ROCm JAX 0.4.31 Release
ROCm 6.0.3, 6.1.3 and 6.2.3
Python 3.10.15
jaxlib
python3 -m pip install https://github.com/ROCm/jax/releases/download/rocm-jax-v0.4.31/jaxlib-0.4.31-cp310-cp310-manylinux_2_28_x86_64.whl
JAX ROCm Plugin
python3 -m pip install https://github.com/ROCm/jax/releases/download/rocm-jax-v0.4.31/jax_rocm60_pjrt-0.4.31-py3-none-manylinux_2_28_x86_64.whl https://github.com/ROCm/jax/releases/download/rocm-jax-v0.4.31/jax_rocm60_plugin-0.4.31-cp310-cp310-manylinux_2_28_x86_64.whl
Python 3.11.10
jaxlib
python3 -m pip install https://github.com/ROCm/jax/releases/download/rocm-jax-v0.4.31/jaxlib-0.4.31-cp311-cp311-manylinux_2_28_x86_64.whl
JAX ROCm Plugin
python3 -m pip install https://github.com/ROCm/jax/releases/download/rocm-jax-v0.4.31/jax_rocm60_pjrt-0.4.31-py3-none-manylinux_2_28_x86_64.whl https://github.com/ROCm/jax/releases/download/rocm-jax-v0.4.31/jax_rocm60_plugin-0.4.31-cp311-cp311-manylinux_2_28_x86_64.whl
Python 3.12.6
jaxlib
python3 -m pip install https://github.com/ROCm/jax/releases/download/rocm-jax-v0.4.31/jaxlib-0.4.31-cp312-cp312-manylinux_2_28_x86_64.whl
JAX ROCm Plugin
python3 -m pip install https://github.com/ROCm/jax/releases/download/rocm-jax-v0.4.31/jax_rocm60_pjrt-0.4.31-py3-none-manylinux_2_28_x86_64.whl https://github.com/ROCm/jax/releases/download/rocm-jax-v0.4.31/jax_rocm60_plugin-0.4.31-cp312-cp312-manylinux_2_28_x86_64.whl
JAX
python3 -m pip install https://github.com/ROCm/jax/archive/refs/tags/rocm-jaxlib-v0.4.31.tar.gz
Docker Images
ROCm 6.0.3
docker pull rocm/jax-community:rocm6.0.3-jax0.4.31-py3.10.15
docker pull rocm/jax-community:rocm6.0.3-jax0.4.31-py3.11.10
docker pull rocm/jax-community:rocm6.0.3-jax0.4.31-py3.12.6
ROCm 6.1.3
docker pull rocm/jax-community:rocm6.1.3-jax0.4.31-py3.10.15
docker pull rocm/jax-community:rocm6.1.3-jax0.4.31-py3.11.10
docker pull rocm/jax-community:rocm6.1.3-jax0.4.31-py3.12.6
ROCm 6.2.3
docker pull rocm/jax-community:rocm6.2.3-jax0.4.31-py3.10.15
docker pull rocm/jax-community:latest-py310
docker pull rocm/jax-community:rocm6.2.3-jax0.4.31-py3.11.10
docker pull rocm/jax-community:latest-py311
docker pull rocm/jax-community:rocm6.2.3-jax0.4.31-py3.12.6
docker pull rocm/jax-community:latest-py312
docker pull rocm/jax-community:latest
