How to install a specific version of jax[cuda12]

[YousefElbrolosy]

I want to install jax[cuda12]==0.5.2, because I am using the version of jax 0.5.2 and I want to use the gpu compatible version with that, however I cannot seem to be able to specify the exact version of jax[cuda12] I want, could someone please direct me to how I can restrict the version?

[jakevdp]

Have you tried something like pip install -U "jax[cuda12]==0.5.2" ?

[YousefElbrolosy]

Yes this worked thank you!

[Tong-RuiLiu]

@jakevdp
Hi Jake, thanks for providing this approach. May I ask another question. I am using the Python environment for installation. Is it possible to assign the specific cuda and CUDNN version such as Cuda 12.6 rather than Cuda 12.9 once installing the JAX and Jaxlib? I just found that currently all of the versions of jax are only supported by Cuda 12.9.

[hawkinsp]

You are welcome to install an older version of CUDA (just pip install older package versions for the nvidia- packages). However, you need to make sure to choose versions that match, e.g. all packages from CUDA 12.6. The minimum CUDA version is 12.1

You also cannot downgrade CUDNN past 9.8; the minimum is 9.8.

See the version constraints here: https://docs.jax.dev/en/latest/installation.html#pip-installation-nvidia-gpu-cuda-installed-locally-harder

[Tong-RuiLiu]

@hawkinsp
Thanks for your help! Now I successfully made it!