rocm-jax-v0.7.1

JAX 0.7.1 Installation Instructions

Python 3.11

jaxlib

python3 -m pip install https://github.com/ROCm/jax/releases/download/rocm-jax-v0.7.1/jaxlib-0.7.1-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl

JAX ROCm Plugin

python3 -m pip install https://github.com/ROCm/jax/releases/download/rocm-jax-v0.7.1/jax_rocm7_pjrt-0.7.1-py3-none-manylinux_2_28_x86_64.whl  https://github.com/ROCm/jax/releases/download/rocm-jax-v0.7.1/jax_rocm7_plugin-0.7.1-cp311-cp311-manylinux_2_28_x86_64.whl

Python 3.12

jaxlib

python3 -m pip install https://github.com/ROCm/jax/releases/download/rocm-jax-v0.7.1/jaxlib-0.7.1-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl

JAX ROCm Plugin

python3 -m pip install https://github.com/ROCm/jax/releases/download/rocm-jax-v0.7.1/jax_rocm7_pjrt-0.7.1-py3-none-manylinux_2_28_x86_64.whl  https://github.com/ROCm/jax/releases/download/rocm-jax-v0.7.1/jax_rocm7_plugin-0.7.1-cp312-cp312-manylinux_2_28_x86_64.whl

Python 3.13

jaxlib

python3 -m pip install https://github.com/ROCm/jax/releases/download/rocm-jax-v0.7.1/jaxlib-0.7.1-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl

JAX ROCm Plugin

python3 -m pip install https://github.com/ROCm/jax/releases/download/rocm-jax-v0.7.1/jax_rocm7_pjrt-0.7.1-py3-none-manylinux_2_28_x86_64.whl https://github.com/ROCm/jax/releases/download/rocm-jax-v0.7.1/jax_rocm7_plugin-0.7.1-cp313-cp313-manylinux_2_28_x86_64.whl

What's Changed

• applied rocm-jax jax build fixes patches by @zahiqbal in #541
• Porting v0.6.0 changes to v0.7.1 by @zahiqbal in #549
• The ROCm plugin should fail gracefully and allow for other devices to… by @JehandadKhan in #550
• Skip test_dot_precision() with TF32 for gfx11xx and gfx12xx architecture (#520) by @AleksaArsic in #551

New Contributors

• @AleksaArsic made their first contribution in #551

Full Changelog: rocm-jax-v0.6.0...rocm-jax-v0.7.1