Jax/Jaxlib 0.4.24 release for ROCm
**Jax/Jaxlib 0.4.24 release for ROCm
Install Jaxlib
Python 3.9
python3 -m pip install https://github.com/ROCmSoftwarePlatform/jax/releases/download/jaxlib-v0.4.24/jaxlib-0.4.24+rocm570-cp39-cp39-manylinux2014_x86_64.whl
python3 -m pip install https://github.com/ROCmSoftwarePlatform/jax/releases/download/jaxlib-v0.4.24/jaxlib-0.4.24+rocm600-cp39-cp39-manylinux2014_x86_64.whl
Python 3.10
python3 -m pip install https://github.com/ROCmSoftwarePlatform/jax/releases/download/jaxlib-v0.4.24/jaxlib-0.4.24+rocm570-cp310-cp310-manylinux2014_x86_64.whl
python3 -m pip install https://github.com/ROCmSoftwarePlatform/jax/releases/download/jaxlib-v0.4.24/jaxlib-0.4.24+rocm600-cp310-cp310-manylinux2014_x86_64.whl
Python 3.11
python3 -m pip install https://github.com/ROCmSoftwarePlatform/jax/releases/download/jaxlib-v0.4.24/jaxlib-0.4.24+rocm570-cp311-cp311-manylinux2014_x86_64.whl
python3 -m pip install https://github.com/ROCmSoftwarePlatform/jax/releases/download/jaxlib-v0.4.24/jaxlib-0.4.24+rocm600-cp311-cp311-manylinux2014_x86_64.whl
Install Jax
python3 -m pip install https://github.com/ROCmSoftwarePlatform/jax/archive/refs/tags/jaxlib-v0.4.24.tar.gz
Docker Images
Python 3.9
docker pull rocm/jax:rocm5.7.0-jax0.4.24-py3.9.0
docker pull rocm/jax:rocm6.0.0-jax0.4.24-py3.9.0
Python 3.10
docker pull rocm/jax:rocm5.7.0-jax0.4.24-py3.10.0
docker pull rocm/jax:rocm6.0.0-jax0.4.24-py3.10.0
Python 3.11
docker pull rocm/jax:rocm5.7.0-jax0.4.24-py3.11.0
docker pull rocm/jax:rocm6.0.0-jax0.4.24-py3.11.0