TheRock/external-builds/jax/README.md

Build JAX with ROCm support

This directory provides tooling for building JAX with ROCm Python wheels.

Tip

If you want to install our prebuilt JAX packages instead of building them from source, see RELEASES.md instead.

Table of contents:

• Support status
• Build instructions
• Test instructions
• Nightly releases

For JAX development references, see:

• jax-ml/jax - upstream JAX
• ROCm/jax - ROCm's downstream JAX fork, used to build the ROCm JAX wheels
• ROCm/rocm-jax - infrastructure (e.g. Dockerfiles), no longer used to build the JAX wheels
• JAX developer documentation

Support status

Project and feature support status

Project / feature	Linux support	Windows support
jaxlib	✅ Supported	❌ Not supported
jax_rocm<major>_pjrt	✅ Supported	❌ Not supported
jax_rocm<major>_plugin	✅ Supported	❌ Not supported

Important

The plugin and PJRT package names embed the major version of the ROCm release they were built against, so the name changes across ROCm releases:

ROCm release	Package names
7.x	jax_rocm7_plugin, jax_rocm7_pjrt
10.x	jax_rocm10_plugin, jax_rocm10_pjrt

Wheels published from earlier ROCm releases keep their original name, so look for jax_rocm7_* when installing a historical ROCm 7.x release.

This document writes jax_rocm<major>_* where either spelling applies.

Supported JAX versions

Support for JAX is provided via stable release branches from ROCm/jax, built with the manylinux flow.

JAX version	Linux	Windows
0.11.0	✅ Supported via ROCm/jax rocm-jaxlib-v0.11.0	❌ Not supported
0.10.2	✅ Supported via ROCm/jax rocm-jaxlib-v0.10.2	❌ Not supported
0.10.1	✅ Supported via ROCm/jax rocm-jaxlib-v0.10.1	❌ Not supported
0.10.0	✅ Supported via ROCm/jax rocm-jaxlib-v0.10.0	❌ Not supported

Note

Python 3.11 is not supported for JAX 0.11.0 and later (dropped upstream).

See also:

• Workflow source code: multi_arch_build_linux_jax_wheels.yml

Build instructions

This repository builds the following ROCm-enabled JAX artifacts:

• jax_rocm<major>_pjrt (PJRT runtime for ROCm)
• jax_rocm<major>_plugin (JAX runtime plugin for ROCm)

Note

jaxlib is not built for supported JAX versions (0.10.0+); it is installed from upstream PyPI (e.g. pip install jaxlib==0.11.0). Only jax_rocm<major>_pjrt and jax_rocm<major>_plugin are built.

How building with TheRock differs from upstream

The downstream ROCm/jax build instructions assume that a stable ROCm version is already installed on the system.

Supported JAX versions (0.10.0+) build against ROCm Python packages installed from the TheRock multi-arch Python package index (the manylinux flow).

Prerequisites

• OS: Linux (supported distributions with ROCm)
• Python: 3.12 recommended (Python 3.11 is not supported for JAX 0.11.0+)
• Compiler: Clang provided by the manylinux build environment
• ROCm: ROCm Python packages from the TheRock multi-arch package index

Steps

1. Checkout the source repository for your JAX version (replace the ref with the version you want to build):

   git clone https://github.com/ROCm/jax.git
   
   pushd jax
   git checkout rocm-jaxlib-v0.11.0
   popd
   

2. Choose your configuration:

   • JAX version: e.g. 0.10.0, 0.10.1, 0.10.2, or 0.11.0
   • Python version: e.g. 3.12
   • Package index: the TheRock multi-arch Python package index.

3. Build the JAX wheels (multi-arch package flow).

   From the ROCm/jax checkout, build the ROCm plugin and PJRT wheels:

   python build/build.py build --wheels=jax-rocm-plugin,jax-rocm-pjrt \
     --python_version=3.12 \
     --bazel_startup_options=--bazelrc=build/rocm/rocm.bazelrc \
     --bazel_options=--config=rocm_release_wheel \
     --bazel_options=--repo_env=ROCM_PATH=$(rocm-sdk path --root) \
     --bazel_options=--repo_env=ML_WHEEL_TYPE=release \
     --bazel_options=--//jaxlib/tools:jaxlib_git_hash=$(git rev-parse HEAD) \
     --verbose --detailed_timestamped_log --output_path=$(pwd)/dist
   

   This is the same flow the GitHub Actions workflow uses:

   • .github/workflows/multi_arch_build_linux_jax_wheels.yml

   The workflow installs ROCm Python packages from the configured TheRock multi-arch package index before building jax_rocm<major>_plugin and jax_rocm<major>_pjrt, where <major> is the ROCm major version being built against.

4. Locate built wheels.

   After a successful build, wheels will be available in:

   jax/dist/
   

For more detailed build options, see the ROCm/jax repository and the .github/workflows/multi_arch_build_linux_jax_wheels.yml workflow in TheRock.

Test instructions

Prerequisites

• AMD GPU matching the target amdgpu_family
• Python environment with pip
• Access to the JAX wheel package index

Testing JAX wheels

1. Checkout the JAX test repo:

   git clone https://github.com/ROCm/jax.git jax_tests
   pushd jax_tests
   git checkout rocm-jaxlib-v<JAX_VERSION>
   popd
   

2. Create a virtual environment:

   python3 -m venv jax_test_env
   source jax_test_env/bin/activate
   

3. Install requirements:

   cd jax
   pip install -r build/test-requirements.txt
   pip install pytest-html pytest-csv uv pytest-json-report
   

4. Install ROCm Python packages:

   pip install \
   --index-url <package_index_url> \
   "rocm[libraries,device-<gfx_arch>]==<rocm_version>"
   

5. Install JAX wheels from the package index:

   # Replace <rocm_major> with the ROCm major version of the index you install
   # from, e.g. 7 for a ROCm 7.x release or 10 for a ROCm 10.x release.
   pip install \
   --index-url <package_index_url> \
   jax_rocm<rocm_major>_plugin \
   jax_rocm<rocm_major>_pjrt
   
   # Install jax from PyPI to match the version
   pip install jax==<JAX_VERSION>
   

6. Run JAX tests:

   cd jax
   # Create a dist directory (required to run the pytest-rocm.sh script).
   mkdir -p dist
   sh ci/run_pytest_rocm.sh
   

Nightly releases

Gating releases with JAX tests

Successful builds publish JAX wheels to the nightly multi-arch Python package index:

https://rocm.nightlies.amd.com/whl-multi-arch/

The published wheels are validated by the JAX test workflow as part of the nightly release process before being made available for use.