Question about JAX working with CUDA 11.2

[dhruvsreenivas]

Hi,

Hope you are doing well! I'm working on a large-scale Jax project, but the cluster I'm using only uses CUDA 11.2, not CUDA ≥ 11.4. I had to install Jax 0.4.2 and Jaxlib 0.4.2 for CUDA recently due to having to use PyTorch dataloading (TensorFlow dataloading ended up being pretty memory heavy so now I currently have a PyTorch setup, which made me upgrade Jax). Previously, the following worked:

jax version == 0.3.16
jaxlib version == 0.3.15 (cuda version, trying to find that now)
dm-haiku version 0.0.7
optax version 0.1.3 I believe
rlax version 0.1.2 I believe

I don't know how I made it work before, but I want to bring that back. How can I install the dependencies for CUDA appropriately? I really like JAX and so want to use it with my available NVIDIA drivers well, but I know on the website for installation it only talks about JAX installation with cuda versions 11.4 or above. I've downgraded PyTorch in the hope that the JAX library will not auto-update upon installation. Thanks!

[dhruvsreenivas]

Alright I think I got it to work... found the wheel online somewhere and now everything seems to work as expected. Terribly sorry for the inconvenience!

[jakevdp]

Glad you got it working! For future reference you can find a list of jaxlib CUDA wheels here: https://storage.googleapis.com/jax-releases/jax_cuda_releases.html, which is the URL listed in the GPU installation instructions.

These can be pip installed using the URL directly, e.g.

$ pip install https://storage.googleapis.com/jax-releases/cuda11/jaxlib-0.3.15+cuda11.cudnn82-cp39-none-manylinux2014_x86_64.whl

The targeted versions are encoded in the filename. The one above is for jaxlib 0.3.15, on linux, with CUDA 11, CUDNN 8,2, and cpython 3.9.