wip: Pathways proxy and jax client shared memory

[samos123]

Steps to build custom wheel:

git clone https://github.com/samos123/xla --branch v0.5.3-pw-shared-memory
# a typo 30 vs 3 but just use 30 for now
git clone https://github.com/samos123/jax --branch v0.5.30-pw-shared-memory

cd jax

docker build -f Dockerfile.build -t jax-builder .

docker run --rm -it -v $(pwd)/dist:/build/dist -v $(pwd):/build/jax -v $HOME/xla:/build/xla jax-builder bash
cd jax
python build/build.py build --wheels=jaxlib --local_xla_path=../xla --python_version=3.10 --clang_path=/usr/bin/clang-18
exit

ls dist/

[samos123]

note this image won't be accessible, will be updated once Shaurya provides it.