jax

1e3b0958 - [ROCm]: Updates for container and build script

[ROCm]: Updates for container and build script -Updated dockerfile.ms -Updated build script to switch building against XLA repo -Update CI script -Update jaxlib setup.py to add rocm version