Install JAX on ROCm 7.13.0

This topic guides you through installing JAX with ROCm support on AMD hardware. It applies to supported AMD GPUs and platforms.

Device family

AMD Instinct™

AMD Radeon™

AMD Ryzen™

Instinct GPU

AMD Instinct MI355X (gfx950)AMD Instinct MI350X (gfx950)AMD Instinct MI350P (gfx950)AMD Instinct MI325X (gfx942)AMD Instinct MI300X (gfx942)AMD Instinct MI300A (gfx942)AMD Instinct MI250X (gfx90a)AMD Instinct MI250 (gfx90a)AMD Instinct MI210 (gfx90a)AMD Instinct MI100 (gfx908)

Radeon GPU

AMD Radeon AI PRO R9700 (gfx1201)AMD Radeon AI PRO R9600D (gfx1201)AMD Radeon RX 9070 XT (gfx1201)AMD Radeon RX 9070 GRE (gfx1201)AMD Radeon RX 9070 (gfx1201)AMD Radeon RX 9060 XT LP (gfx1200)AMD Radeon RX 9060 XT (gfx1200)AMD Radeon RX 9060 (gfx1200)AMD Radeon PRO W7900 Dual Slot (gfx1100)AMD Radeon PRO W7900 (gfx1100)AMD Radeon PRO W7800 48GB (gfx1100)AMD Radeon PRO W7800 (gfx1100)AMD Radeon RX 7900 XTX (gfx1100)AMD Radeon RX 7900 XT (gfx1100)AMD Radeon RX 7900 GRE (gfx1100)AMD Radeon PRO W7700 (gfx1101)AMD Radeon RX 7800 XT (gfx1101)AMD Radeon RX 7700 XT (gfx1101)AMD Radeon RX 7700 XE (gfx1101)AMD Radeon RX 7700 (gfx1101)AMD Radeon PRO V710 (gfx1101)AMD Radeon RX 7600 (gfx1102)AMD Radeon PRO W6800 (gfx1030)AMD Radeon PRO V620 (gfx1030)

Ryzen APU

AMD Ryzen AI Max+ PRO 395 (gfx1151)AMD Ryzen AI Max PRO 390 (gfx1151)AMD Ryzen AI Max PRO 385 (gfx1151)AMD Ryzen AI Max PRO 380 (gfx1151)AMD Ryzen AI Max+ 395 (gfx1151)AMD Ryzen AI Max+ 392 (gfx1151)AMD Ryzen AI Max+ 388 (gfx1151)AMD Ryzen AI Max 390 (gfx1151)AMD Ryzen AI Max 385 (gfx1151)AMD Ryzen AI 9 HX PRO 475 (gfx1150)AMD Ryzen AI 9 HX PRO 470 (gfx1150)AMD Ryzen AI 9 PRO 465 (gfx1150)AMD Ryzen AI 7 PRO 450 (gfx1152)AMD Ryzen AI 5 PRO 440 (gfx1152)AMD Ryzen AI 9 HX 475 (gfx1150)AMD Ryzen AI 9 HX 470 (gfx1150)AMD Ryzen AI 9 465 (gfx1150)AMD Ryzen AI 7 450 (gfx1152)AMD Ryzen AI 9 HX PRO 375 (gfx1150)AMD Ryzen AI 9 HX PRO 370 (gfx1150)AMD Ryzen AI 7 PRO 350 (gfx1152)AMD Ryzen AI 5 PRO 340 (gfx1152)AMD Ryzen AI 9 HX 375 (gfx1150)AMD Ryzen AI 9 HX 370 (gfx1150)AMD Ryzen AI 9 365 (gfx1150)AMD Ryzen AI 7 350 (gfx1152)AMD Ryzen AI 7 345 (gfx1152)AMD Ryzen AI 5 340 (gfx1152)AMD Ryzen AI 5 330 (gfx1152)AMD Ryzen 7 PRO 250 (gfx1103)AMD Ryzen 5 PRO 230 (gfx1103)AMD Ryzen 5 PRO 220 (gfx1103)AMD Ryzen 5 PRO 215 (gfx1103)AMD Ryzen 3 PRO 210 (gfx1103)AMD Ryzen 9 270 (gfx1103)AMD Ryzen 7 260 (gfx1103)AMD Ryzen 7 250 (gfx1103)AMD Ryzen 5 240 (gfx1103)AMD Ryzen 5 230 (gfx1103)AMD Ryzen 5 220 (gfx1103)AMD Ryzen 3 210 (gfx1103)

Operating system

Linux

JAX version

0.9.1

0.8.2

Prerequisites

• Ensure your system has the AMD GPU Driver (amdgpu) installed. See the ROCm 7.13.0 compatibility matrix for driver support information. For installation instructions, see the AMD GPU Driver documentation.

• Ensure your system has a supported Python version installed and accessible: 3.11, 3.12, 3.13, or 3.14.

• Install the ROCm Core SDK – it’s recommended to use pip to install JAX and the ROCm Core SDK in the same Python virtual environment.

  Important

  Unlike PyTorch, the JAX packages do not automatically install rocm[libraries] as a dependency.

Install JAX using pip

For prerequisite steps and post-installation recommendations, see the ROCm installation instructions.

1. Set up your Python virtual environment. For example, run the following command to create one with Python 3.13:

   python3.13 -m venv .venv
   

2. Activate your Python virtual environment. For example:

   source .venv/bin/activate
   

3. Install the appropriate ROCm-enabled JAX libraries for your operating system and AMD hardware architecture.

   Note

   The jax and jaxlib packages are not published to the AMD package repository. After installing GFX architecture-based jax_rocm7_plugin and jax_rocm7_pjrt packages from the AMD repository, install jax and jaxlib from PyPI.

   Note

   The jax package is not published to the AMD package repository. After installing GFX architecture-based jaxlib, jax_rocm7_plugin and jax_rocm7_pjrt packages from the AMD repository, install jax from PyPI.

   python -m pip install --index-url https://repo.amd.com/rocm/whl/gfx950-dcgpu/ \
     "jax_rocm7_plugin==0.9.1+rocm7.13.0" \
     "jax_rocm7_pjrt==0.9.1+rocm7.13.0"
   
   # Install jax from PyPI
   python -m pip install \
     "jax==0.9.1" \
     "jaxlib==0.9.1"
   

   python -m pip install --index-url https://repo.amd.com/rocm/whl/gfx94X-dcgpu/ \
     "jax_rocm7_plugin==0.9.1+rocm7.13.0" \
     "jax_rocm7_pjrt==0.9.1+rocm7.13.0"
   
   # Install jax from PyPI
   python -m pip install \
     "jax==0.9.1" \
     "jaxlib==0.9.1"
   

   python -m pip install --index-url https://repo.amd.com/rocm/whl/gfx90a/ \
     "jax_rocm7_plugin==0.9.1+rocm7.13.0" \
     "jax_rocm7_pjrt==0.9.1+rocm7.13.0"
   
   # Install jax from PyPI
   python -m pip install \
     "jax==0.9.1" \
     "jaxlib==0.9.1"
   

   python -m pip install --index-url https://repo.amd.com/rocm/whl/gfx908/ \
     "jax_rocm7_plugin==0.9.1+rocm7.13.0" \
     "jax_rocm7_pjrt==0.9.1+rocm7.13.0"
   
   # Install jax from PyPI
   python -m pip install \
     "jax==0.9.1" \
     "jaxlib==0.9.1"
   

   python -m pip install --index-url https://repo.amd.com/rocm/whl/gfx120X-all/ \
     "jax_rocm7_plugin==0.9.1+rocm7.13.0" \
     "jax_rocm7_pjrt==0.9.1+rocm7.13.0"
   
   # Install jax from PyPI
   python -m pip install \
     "jax==0.9.1" \
     "jaxlib==0.9.1"
   

   python -m pip install --index-url https://repo.amd.com/rocm/whl/gfx110X-all/ \
     "jax_rocm7_plugin==0.9.1+rocm7.13.0" \
     "jax_rocm7_pjrt==0.9.1+rocm7.13.0"
   
   # Install jax from PyPI
   python -m pip install \
     "jax==0.9.1" \
     "jaxlib==0.9.1"
   

   python -m pip install --index-url https://repo.amd.com/rocm/whl/gfx103X-all/ \
     "jax_rocm7_plugin==0.9.1+rocm7.13.0" \
     "jax_rocm7_pjrt==0.9.1+rocm7.13.0"
   
   # Install jax from PyPI
   python -m pip install \
     "jax==0.9.1" \
     "jaxlib==0.9.1"
   

   python -m pip install --index-url https://repo.amd.com/rocm/whl/gfx1151/ \
     "jax_rocm7_plugin==0.9.1+rocm7.13.0" \
     "jax_rocm7_pjrt==0.9.1+rocm7.13.0"
   
   # Install jax from PyPI
   python -m pip install \
     "jax==0.9.1" \
     "jaxlib==0.9.1"
   

   python -m pip install --index-url https://repo.amd.com/rocm/whl/gfx1150/ \
     "jax_rocm7_plugin==0.9.1+rocm7.13.0" \
     "jax_rocm7_pjrt==0.9.1+rocm7.13.0"
   
   # Install jax from PyPI
   python -m pip install \
     "jax==0.9.1" \
     "jaxlib==0.9.1"
   

   python -m pip install --index-url https://repo.amd.com/rocm/whl/gfx1152/ \
     "jax_rocm7_plugin==0.9.1+rocm7.13.0" \
     "jax_rocm7_pjrt==0.9.1+rocm7.13.0"
   
   # Install jax from PyPI
   python -m pip install \
     "jax==0.9.1" \
     "jaxlib==0.9.1"
   

   python -m pip install --index-url https://repo.amd.com/rocm/whl/gfx950-dcgpu/ \
     "jaxlib==0.8.2" \
     "jax_rocm7_plugin==0.8.2+rocm7.13.0" \
     "jax_rocm7_pjrt==0.8.2+rocm7.13.0"
   
   # Install jax from PyPI
   python -m pip install "jax==0.8.2"
   

   python -m pip install --index-url https://repo.amd.com/rocm/whl/gfx94X-dcgpu/ \
     "jaxlib==0.8.2" \
     "jax_rocm7_plugin==0.8.2+rocm7.13.0" \
     "jax_rocm7_pjrt==0.8.2+rocm7.13.0"
   
   # Install jax from PyPI
   python -m pip install "jax==0.8.2"
   

   python -m pip install --index-url https://repo.amd.com/rocm/whl/gfx90a/ \
     "jaxlib==0.8.2" \
     "jax_rocm7_plugin==0.8.2+rocm7.13.0" \
     "jax_rocm7_pjrt==0.8.2+rocm7.13.0"
   
   # Install jax from PyPI
   python -m pip install "jax==0.8.2"
   

   python -m pip install --index-url https://repo.amd.com/rocm/whl/gfx908/ \
     "jaxlib==0.8.2" \
     "jax_rocm7_plugin==0.8.2+rocm7.13.0" \
     "jax_rocm7_pjrt==0.8.2+rocm7.13.0"
   
   # Install jax from PyPI
   python -m pip install "jax==0.8.2"
   

   python -m pip install --index-url https://repo.amd.com/rocm/whl/gfx120X-all/ \
     "jaxlib==0.8.2" \
     "jax_rocm7_plugin==0.8.2+rocm7.13.0" \
     "jax_rocm7_pjrt==0.8.2+rocm7.13.0"
   
   # Install jax from PyPI
   python -m pip install "jax==0.8.2"
   

   python -m pip install --index-url https://repo.amd.com/rocm/whl/gfx110X-all/ \
     "jaxlib==0.8.2" \
     "jax_rocm7_plugin==0.8.2+rocm7.13.0" \
     "jax_rocm7_pjrt==0.8.2+rocm7.13.0"
   
   # Install jax from PyPI
   python -m pip install "jax==0.8.2"
   

   python -m pip install --index-url https://repo.amd.com/rocm/whl/gfx103X-all/ \
     "jaxlib==0.8.2" \
     "jax_rocm7_plugin==0.8.2+rocm7.13.0" \
     "jax_rocm7_pjrt==0.8.2+rocm7.13.0"
   
   # Install jax from PyPI
   python -m pip install "jax==0.8.2"
   

   python -m pip install --index-url https://repo.amd.com/rocm/whl/gfx1151/ \
     "jaxlib==0.8.2" \
     "jax_rocm7_plugin==0.8.2+rocm7.13.0" \
     "jax_rocm7_pjrt==0.8.2+rocm7.13.0"
   
   # Install jax from PyPI
   python -m pip install "jax==0.8.2"
   

   python -m pip install --index-url https://repo.amd.com/rocm/whl/gfx1150/ \
     "jaxlib==0.8.2" \
     "jax_rocm7_plugin==0.8.2+rocm7.13.0" \
     "jax_rocm7_pjrt==0.8.2+rocm7.13.0"
   
   # Install jax from PyPI
   python -m pip install "jax==0.8.2"
   

   python -m pip install --index-url https://repo.amd.com/rocm/whl/gfx1152/ \
     "jaxlib==0.8.2" \
     "jax_rocm7_plugin==0.8.2+rocm7.13.0" \
     "jax_rocm7_pjrt==0.8.2+rocm7.13.0"
   
   # Install jax from PyPI
   python -m pip install "jax==0.8.2"
   

4. Set the following environment variable before running JAX as a workaround for a known issue.

   export XLA_FLAGS="--xla_gpu_enable_command_buffer="
   

5. Check your JAX installation.

   python -c "import jax; print(jax.devices())"
   

   This prints something like [RocmDevice(id=0)] if JAX and ROCm are installed properly.

Known issues

These are known issues related to JAX installation on ROCm 7.13.0 and their workarounds.

Segfaults with JAX 0.9.1

JAX 0.9.1 might segfault during execution. To work around this, disable XLA command buffers by setting the following flag before running your script:

export XLA_FLAGS="--xla_gpu_enable_command_buffer="