Install and/or package jax matching cuda+cudnn versions #24042
Unanswered
olivier-peltre
asked this question in
Q&A
Replies: 2 comments 1 reply
-
Update: running
seemed to provide satsifactory versions. Is this robust or recommended? |
Beta Was this translation helpful? Give feedback.
0 replies
-
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. |
Beta Was this translation helpful? Give feedback.
1 reply
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
-
How should one install jax + jaxlib with the most recent cuda versions on ubuntu? As of today, this is
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.
Beta Was this translation helpful? Give feedback.
All reactions