Install and/or package jax matching cuda+cudnn versions

[olivier-peltre]

How should one install jax + jaxlib with the most recent cuda versions on ubuntu? As of today, this is

• cuda: 12.4
• cudnn: 9.4

Currently, the jax_cuda_releases page only holds versions up to 0.4.29, while pip install jax[cuda] yields 0.4.33 which is not compatible with my current system (expecting cuda 12.6)

I'm currently trying pip install "jax[cuda]==0.4.30", then will try 0.4.31, etc, until I hopefully get compatible jaxlib + cudnn version.
Would you have any better installation instructions?
I guess uploading the "3 months old" jaxlibs to jax_cuda_releases with the same cuda+cudnn tags might already help finding the version that fits a shoe!

A follow-up question that I've been already reading a lot of internet about, is how to register jax/jaxlib as a flexible enough dependency within a platform-agnostic library (not just an application repo)... I'm happy to drop poetry for uv or plain old setup-tools, but I doubt there is an easy solution for this right now. Any advice welcome!

Cheers 🙏

---

EDIT: jax[cuda12] == 0.4.30 seems to match cudnn 9.4, but the next question would follow in a month ^^

---

I've been through e.g.

• Please provide PEP 503 compliant indices for CUDA versions of packages #5410
• https://github.com/jorenham/jax_pep503?tab=readme-ov-file

[olivier-peltre]

Update: running

pip install "jax[cuda_12]" "nvidia-cuda-nvcc-cu12>=12.4,<12.5

seemed to provide satsifactory versions. Is this robust or recommended?

[hawkinsp]

JAX is compatible with CUDA 12.4, see https://jax.readthedocs.io/en/latest/installation.html#pip-installation-nvidia-gpu-cuda-installed-locally-harder.

However, I'm trying to understand what went wrong when you used CUDA 12.6 and just let JAX install it.

[hawkinsp]

I'll also note that jax_cuda_releases will never be updated again: we now ship wheels on pypi.

[olivier-peltre]

Thanks for the up-to-date link and precious info! 🙏