How to build jax[cuda12]==0.4.34 with local xla repo?

[TTTNCTTT]

Background

I have some applications running on jax installed via pip install jax[cuda12]==0.4.34 jaxlib==0.4.34. Recently, I encountered this issue when enabling both cuDNN and persistent caching simultaneously: String field 'xla.gpu.CompilationResultProto.DnnCompiledGraphsEntry.value' contains invalid UTF-8 data.
I noticed that this bug was fixed in the commit to the xla repository on January 9, 2025, and after that, the latest version of jax was 0.5.0. So, I reinstalled the latest jax using pip install jax[cuda12]==0.5.0 jaxlib==0.5.0 after removing the old jax. The aforementioned issue was resolved!
However, my application is not compatible with jax 0.5.0, as it runs slower and encounters some NaN errors. I decided to build jaxlib and jax[cuda12] using a local xla repository.

What did I do

I pulled the jax and xla repositories and switched to specific commits using the following commands:

git clone --recurse-submodules https://github.com/jax-ml/jax.git
git clone --recurse-submodules https://github.com/openxla/xla.git
cd jax
# This corresponds to the [jax v0.4.34 release](https://github.com/jax-ml/jax/tree/jax-v0.4.34) version.
git checkout affba367c5533df8900e32cbc3d31ca92dd1c1ea 
git submodule update --init --recursive
cd ..
cd xla
# This is the XLA version defined in the [/jax/third_party/xla/workspace.bzl](https://github.com/jax-ml/jax/commit/aa9ee7abfab2344ce56483af29266a31ca7b7708) file from the `jax v0.4.34 release`.
git checkout cd6e808c59f53b40a99df1f1b860db9a3e598bff 
git submodule update --init --recursive

After modifying the XLA source code to fix the bug, I read developer.md and used the following command to build jaxlib:

python3 build/build.py \
    --python_version=3.11 \
    --enable_cuda \
    --cuda_version=12.6.1 \
    --cudnn_version=9.4.0 \
    --bazel_options=--override_repository=xla=/home/thomas/xla \
    --verbose

After a long wait, I received the following output:

Target //jaxlib/tools:build_wheel up-to-date:
  bazel-bin/jaxlib/tools/build_wheel
INFO: Elapsed time: 1928.511s, Critical Path: 287.80s
INFO: 3390 processes: 24 internal, 3366 local.
INFO: Build completed successfully, 3390 total actions
INFO: Running command line: bazel-bin/jaxlib/tools/build_wheel '--output_path=/home/thomas/jax/dist' '--jaxlib_git_hash=affba367c5533df8900e32cbc3d31ca92dd1c1ea' '--cpu=x86_64'
Output wheel: /home/thomas/jax/dist/jaxlib-0.4.34.dev20250218-cp311-cp311-manylinux2014_x86_64.whl
To install the newly-built jaxlib wheel on system Python, run:
  pip install /home/thomas/jax/dist/jaxlib-0.4.34.dev20250218-cp311-cp311-manylinux2014_x86_64.whl --force-reinstall
To install the newly-built jaxlib wheel on hermetic Python, run:
  echo -e "\n/home/thomas/jax/dist/jaxlib-0.4.34.dev20250218-cp311-cp311-manylinux2014_x86_64.whl" >> build/requirements.in
  bazel run //build:requirements.update --repo_env=HERMETIC_PYTHON_VERSION=3.11

Question

I successfully built jaxlib. However, the question arises: how do I build jax[cuda12]? I did not find any instructions in the jax documentation for building a specific wheel with the [cuda12] tag. Thank you for your response!

[L-Dandan]

Excuse me, is the jax installed by this method the CPU version or the cuda 12 version?

[hawkinsp]

Please see the docs: https://docs.jax.dev/en/latest/developer.html#building-jaxlib-from-source

The [cuda12] tag includes the jax-cuda12-pjrt and jax-cuda12-plugin wheels.