[Documentation]: Instructions for building Jax binaries using The Rock (as of 3/24/2025)

[gabeweisz]

These are the steps necessary (as of 3/24/2025) to build JAX using The Rock:

Set Environment Variables

# Environment variables;
#  THE_ROCK_DIR points to the location of The Rock
#  JAX_DIR points to the location of the Jax source
#  XLA_DIR points to the location of the XLA source   

export ROCM_PATH=$THE_ROCK_DIR
export PATH=$ROCM_PATH/bin:$ROCM_PATH/llvm/bin:$PATH

create and activate a python venv if needed

pip -m venv .venv-jax
source .venv-jax/bin/activate
# jax needs patchelf to build its wheels
pip install patchelf

Build and install roctracer into The Rock

The Rock comes with device libs but is missing a symlink needed for roctracer to build:

ln -s $ROCM_PATH/lib/llvm/amdgcn/ $ROCM_PATH/amdgcn

Build roctracer as follows:

# An active Python venv is assumed
pip install CppHeaderParser
git clone https://github.com/rocm/roctracer
cd roctracer
export HIP_PLATFORM=amd
# this will avoid building some targets that we don't need and make building faster
sed -i '/make mytest/d' build.sh
sed -i '/make package/d' build.sh
sed -i 's/^make/make -j \`nproc\`/' build.sh

# build and install
bash build.sh
cd build
make install

Build Jax

Jax currently uses the Bazel build system, which (as of version 6.5, the version used to build Jax) fails with a versioned libamdhip64.so.6. This can be solved by editing third_party/tsl/third_party/gpus/rocm/BUILD.tpl inside $XLA_PATH and changing

srcs = glob(["%{rocm_root}/lib/libamdhip*.so.6*"])

to

srcs = glob(["%{rocm_root}/lib/libamdhip*.so.6"])

around line 140

After this, wheels for the Jax binaries can be built with:

# we want a release build
export JAXLIB_RELEASE=1

cd $JAX_DIR
python build/build.py build --wheels jaxlib,jax-rocm-plugin,jax-rocm-pjrt --rocm_path $ROCM_PATH --local_xla_path=$XLA_PATH --rocm_amdgpu_targets gfx942

If on NFS, the build process will be much faster if you configure bazel to store its intermediate files on local storage, by adding

--bazel_startup_options='--output_user_root=/tmp/bazel-jax-build'

to the build command above

To build jax 0.4.35-qa, check out the appropriate branches of jax and xla and then run

export JAXLIB_RELEASE=1
python ./build/build.py --noenable_cuda --enable_rocm --build_gpu_kernel_plugin=rocm --build_gpu_plugin=true --build_gpu_pjrt_plugin=true  --rocm_amdgpu_targets=gfx942 --bazel_options=--override_repository=xla="$XLA_PATH" --rocm_path="$ROCM_PATH"

[gcapodagAMD]

just in case it is useful to other people, building with ROCm 6.4.0 and version 4.35 of https://github.com/ROCm/jax and https://github.com/ROCm/xla with this command:

python3 build/build.py --enable_rocm --rocm_path=$ROCM_PATH \
                                --bazel_options=--override_repository=xla=$XLA_PATH \
                                --rocm_amdgpu_targets=$AMDGPU_GFXMODEL \
                                --build_gpu_plugin --gpu_plugin_rocm_version=60 --build_gpu_kernel_plugin=rocm \
                                --bazel_options=--jobs=128 \
                                --bazel_startup_options=--host_jvm_args=-Xmx512m

I had to modify the .bazelrc file like this:

sed -i '$a build:rocm --copt=-Wno-error=c23-extensions' .bazelrc

to make it work.

[gabeweisz]

Here is a script that will download a nightly from the rock and use it to build Jax

# this script will create a directory called "jax_build" underneath the current directory
# and then use the rock (currently nightly version 20250212)
# to build jax 0.5.0 (with the matching Jax) in a python virtual environment
# called ".venv-jax" underneath "jax_build"
# the rock will be available under "jax_build/therock"
# and the venv can be activated with "source jax_build/.venv-jax/bin/activate"
# If using it, you should use the same rocm from the rock by setting
# export ROCM_DIR=`pwd`/jax_build/therock
# export PATH=$ROCM_PATH/bin:$PATH
# export LD_LIBRARY_PATH=$ROCM_PATH/lib:$LD_LIBRARY_PATH
# this script will also leave a jax wheel in jax_build/jax, and jaxlib wheels in jax_build/jaxlib/dist

mkdir jax_build
cd jax_build

# update this filename as desired
# this version does not work right now because tensile cannot load
#wget https://github.com/ROCm/TheRock/releases/download/nightly-tarball/therock-dist-linux-gfx94X-dcgpu-6.4.0rc20250512.tar.gz -O therock.rgz
wget https://github.com/ROCm/TheRock/releases/download/mainline-snapshot-2025-03-13/therock-dist-gfx94X-dcgpu.tar.gz -O therock.tgz

mkdir therock
pushd therock
echo "extracting The Rock!"
tar xzf ../therock.tgz
ln -s lib/llvm/amdgcn/ amdgcn
popd
export ROCM_PATH=`pwd`/therock
export PATH=$ROCM_PATH/bin:$ROCM_PATH/llvm/bin:$PATH

python -m venv .venv-jax
source .venv-jax/bin/activate
pip install --upgrade pip
pip install patchelf

pip install CppHeaderParser
echo "Building roctracer"
git clone https://github.com/rocm/roctracer
pushd roctracer
export HIP_PLATFORM=amd
# this will avoid building some targets that we don't need and make building faster
sed -i '/make mytest/d' build.sh
sed -i '/make package/d' build.sh
sed -i 's/^make/make -j `nproc`/' build.sh

# build and install
bash build.sh
cd build
make install
popd

echo "checking out xla"
git clone https://github.com/rocm/xla
pushd xla
git checkout rocm-jaxlib-v0.5.0
git cherry-pick 265eb4a21e4184d50aba5ffc26dc2ea100ac5b57
popd

echo "checking out jax"
git clone https://github.com/rocm/jax
cd jax
git checkout rocm-jaxlib-v0.5.0
export XLA_PATH=`pwd`/../xla
export JAXLIB_RELEASE=1
echo "building jax and jaxlib"
python build/build.py build --wheels jaxlib,jax-rocm-plugin,jax-rocm-pjrt --rocm_path $ROCM_PATH --local_xla_path=$XLA_PATH --rocm_amdgpu_targets gfx942 --bazel_startup_options='--output_user_root=/tmp/bazel-jax-build'
pip wheel . --no-deps
echo "installing jax"
pip install *.whl dist/*.whl

[gabeweisz]

As of 7/3/2025:
With the latest TheRock 20250701 and Jax/xla 0.6.0 (jax 7acefb2024dd71f75502a3bc3af696183bc32702, xla 60eb599c959c9720a4a6242ba3c92924e9d886f4)

These 2 manual steps are needed to get Jax to build with therock:

ln -s $THEROCK_PATH/llvm/amdgcn $THEROCK_PATH/amdgcn # this is needed for both compiling and running
pip install patchelf # needed to compile jax

[dmartinovic-htec]

I'm currently trying to build JAX with ROCm support using the build.py script, but I’m running into this error:

ERROR - Incorrect wheel name provided, valid choices are jaxlib, jax-cuda-plugin or cuda-plugin, jax-cuda-pjrt or cuda-pjrt, jax-rocm-plugin or rocm-plugin, jax-rocm-pjrt or rocm-pjrt

The command I’m running is:
python build/build.py build --wheels jaxlib,jax-rocm-plugin,jax-rocm-pjrt --rocm_path $ROCM_PATH --local_xla_path=$XLA_PATH --rocm_amdgpu_targets gfx942

I believe I’m using the correct wheel names (jaxlib,jax-rocm-plugin,jax-rocm-pjrt), but the script still throws this error. Could someone please clarify the expected format or let me know if I'm missing something?

Thanks in advance!

[gabeweisz]

As of 9/26/2025, the preferred way to build Jax with any ROCm, including TheRock, is to use the rocm-jax repo. This changed recently, and I had not updated these instructions yet.

Automatic nightly builds of Jax with TheRock are almost working as well, and you should be able to pull wheels using from the last passing run of the PR to fix the nightly builds.

This flow supports Jax 0.6.0. Support for Jax 0.6.2 and 0.7.0 will come soon. You can use the pjrt and plugin wheels generated by this process with Jax 0.6.2, but need to patch Jaxlib in a few places to remove some erroneous version checks. I think this also works if you use JaxLib 0.6.0 from this build process with the rest of these wheels and Jax 0.6.2.

[gcapodagAMD]

@dmartinovic-htec you could also take a look at this script we have been using: https://github.com/amd/HPCTrainingDock/blob/main/extras/scripts/jax_setup.sh