How to install compatible versions of Jax and Jaxlib for a particular CUDA version?

[davidrpugh]

I am having trouble getting both jax and jaxlib installed with compatible versions. For example, how would I get the most recent versions of Jax and Jaxlib compatible with CUDA 10.1 and cuDNN 7.6? At the moment I am just trying combinations to find one that appears to work. Here would be a typical requirements.txt file.

fedjax
jax=0.2.16  
jaxlib==0.1.65+cuda101
tensorflow==2.3
tensorflow-federated==0.17

--find-links https://storage.googleapis.com/jax-releases/jax_releases.html

Is there a better way?

[jakevdp]

JAX dropped support for CUDA 10.X in jaxlib version 0.1.72 (See https://github.com/google/jax/blob/main/CHANGELOG.md#jaxlib-0172-oct-12-2021), and looking at the changelog the current JAX version at that release was v 0.2.16. So you should be able to use something like

pip install jax[cuda101]==0.2.16 -f https://storage.googleapis.com/jax-releases/jax_releases.html

Note that if you already have a version of JAX installed, pip uses the local version to parse extras (which feels like a bug to me), so you have to work around it by first installing the appropriate JAX version without extras:

pip install jax==0.2.16
pip install jax[cuda101]==0.2.16 -f https://storage.googleapis.com/jax-releases/jax_releases.html

I think you can also use the jax[cuda101]==0.2.16 syntax in your requirements file.

[davidrpugh]

Excellent! Thanks. I have started using this syntax in my Jax project template repo.

[hawkinsp]

CUDA 10.1 is buggy and causes wrong outputs in some cases, which is the key reason we dropped support for it. I really do mean that CUDA (ptxas) itself is buggy, not JAX: this is not something we can work around. I recommend upgrading to CUDA 11.1 or newer.

[davidrpugh]

Indeed I would upgrade but this particular application uses FedJax which in turn installs TensorFlow 2.3 which requires CUDA 10.1.

[sumowi]

From: https://pypi.org/project/jax/

pip install --upgrade pip

CUDA 12 installation

• Note: wheels only available on linux.

pip install --upgrade "jax[cuda12_pip]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

CUDA 11 installation

• Note: wheels only available on linux.

pip install --upgrade "jax[cuda11_pip]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html