Install JAX for ROCm 7.14.0#
This page guides you through installing JAX with ROCm support on AMD hardware. It applies to supported AMD GPUs and platforms.
Prerequisites#
Ensure your host system has the AMD GPU Driver (amdgpu) installed. See the ROCm compatibility matrix for driver support information. For installation instructions, see the AMD GPU Driver documentation.
Ensure the host system has Docker Engine installed.
Ensure your system has a supported Python version installed and accessible: 3.11, 3.12, 3.13, or 3.14.
Complete the ROCm Core SDK installation prerequisites. See Prerequisites (Install ROCm) for instructions.
Important
Unlike PyTorch, the JAX packages don’t automatically install ROCm library and device packages as dependencies. The following section includes recommended instructions to install ROCm in a Python virtual environment alongside JAX. See Install ROCm for other installation methods.
Get started#
Pull the ROCm JAX 0.10.0 Docker image.
docker pull rocm/jax:rocm7.14-jax0.10.0-py3.14
docker pull rocm/jax:rocm7.14-jax0.10.0-py3.13
docker pull rocm/jax:rocm7.14-jax0.10.0-py3.12
docker pull rocm/jax:rocm7.14-jax0.10.0-py3.11
Pull the ROCm JAX 0.9.1 Docker image.
docker pull rocm/jax:rocm7.14-jax0.9.1-py3.14
docker pull rocm/jax:rocm7.14-jax0.9.1-py3.13
docker pull rocm/jax:rocm7.14-jax0.9.1-py3.12
docker pull rocm/jax:rocm7.14-jax0.9.1-py3.11
Start the Docker container.
docker run -it --rm \ --device /dev/kfd \ --device /dev/dri \ --network=host \ --ipc=host \ --group-add=video \ --cap-add=SYS_PTRACE \ --security-opt seccomp=unconfined \ rocm/jax:rocm7.14-jax0.10.0-py3.14 \ bash
docker run -it --rm \ --device /dev/kfd \ --device /dev/dri \ --network=host \ --ipc=host \ --group-add=video \ --cap-add=SYS_PTRACE \ --security-opt seccomp=unconfined \ rocm/jax:rocm7.14-jax0.10.0-py3.13 \ bash
docker run -it --rm \ --device /dev/kfd \ --device /dev/dri \ --network=host \ --ipc=host \ --group-add=video \ --cap-add=SYS_PTRACE \ --security-opt seccomp=unconfined \ rocm/jax:rocm7.14-jax0.10.0-py3.12 \ bash
docker run -it --rm \ --device /dev/kfd \ --device /dev/dri \ --network=host \ --ipc=host \ --group-add=video \ --cap-add=SYS_PTRACE \ --security-opt seccomp=unconfined \ rocm/jax:rocm7.14-jax0.10.0-py3.11 \ bash
docker run -it --rm \ --device /dev/kfd \ --device /dev/dri \ --network=host \ --ipc=host \ --group-add=video \ --cap-add=SYS_PTRACE \ --security-opt seccomp=unconfined \ rocm/jax:rocm7.14-jax0.9.1-py3.14 \ bash
docker run -it --rm \ --device /dev/kfd \ --device /dev/dri \ --network=host \ --ipc=host \ --group-add=video \ --cap-add=SYS_PTRACE \ --security-opt seccomp=unconfined \ rocm/jax:rocm7.14-jax0.9.1-py3.13 \ bash
docker run -it --rm \ --device /dev/kfd \ --device /dev/dri \ --network=host \ --ipc=host \ --group-add=video \ --cap-add=SYS_PTRACE \ --security-opt seccomp=unconfined \ rocm/jax:rocm7.14-jax0.9.1-py3.12 \ bash
docker run -it --rm \ --device /dev/kfd \ --device /dev/dri \ --network=host \ --ipc=host \ --group-add=video \ --cap-add=SYS_PTRACE \ --security-opt seccomp=unconfined \ rocm/jax:rocm7.14-jax0.9.1-py3.11 \ bash
Install JAX using pip#
For prerequisite steps and post-installation recommendations, see the ROCm installation instructions.
Set up your Python virtual environment.
python3.14 -m venv .venv
python3.13 -m venv .venv
python3.12 -m venv .venv
python3.11 -m venv .venv
Activate your Python virtual environment.
source .venv/bin/activate
If you don’t have an existing ROCm installation, install ROCm using the following command; otherwise, proceed to installing JAX libraries.
python -m pip install --index-url https://repo.amd.com/rocm/whl-multi-arch/ \ "rocm[libraries,device-all]"
python -m pip install --index-url https://repo.amd.com/rocm/whl-multi-arch/ \ "rocm[libraries,device-gfx950]"
python -m pip install --index-url https://repo.amd.com/rocm/whl-multi-arch/ \ "rocm[libraries,device-gfx942]"
python -m pip install --index-url https://repo.amd.com/rocm/whl-multi-arch/ \ "rocm[libraries,device-gfx90a]"
python -m pip install --index-url https://repo.amd.com/rocm/whl-multi-arch/ \ "rocm[libraries,device-gfx1200]"
python -m pip install --index-url https://repo.amd.com/rocm/whl-multi-arch/ \ "rocm[libraries,device-gfx1201]"
python -m pip install --index-url https://repo.amd.com/rocm/whl-multi-arch/ \ "rocm[libraries,device-gfx1100]"
python -m pip install --index-url https://repo.amd.com/rocm/whl-multi-arch/ \ "rocm[libraries,device-gfx1101]"
python -m pip install --index-url https://repo.amd.com/rocm/whl-multi-arch/ \ "rocm[libraries,device-gfx1102]"
Install the ROCm-enabled JAX libraries.
Note
The
jaxandjaxlibpackages are not published to the AMD package repository. After installing GFX architecture-basedjax_rocm7_pluginandjax_rocm7_pjrtpackages from the AMD repository, installjaxandjaxlibfrom PyPI.python -m pip install --index-url https://repo.amd.com/rocm/whl-multi-arch/ \ "jax_rocm7_plugin==0.10.0+rocm7.14.0" \ "jax_rocm7_pjrt==0.10.0+rocm7.14.0" # Install jax from PyPI python -m pip install \ "jax==0.10.0" \ "jaxlib==0.10.0"
python -m pip install --index-url https://repo.amd.com/rocm/whl-multi-arch/ \ "jax_rocm7_plugin==0.9.1+rocm7.14.0" \ "jax_rocm7_pjrt==0.9.1+rocm7.14.0" # Install jax from PyPI python -m pip install \ "jax==0.9.1" \ "jaxlib==0.9.1"
Verify your JAX installation.
python -c "import jax; print(jax.devices())"
This prints something like
[RocmDevice(id=0)]if JAX and ROCm are installed properly and your AMD GPUs are detected.