Jaxlib for Cuda 11.8

[smtabatabaie]

Hi, I'm trying to install jax and jaxlib for snerg (https://github.com/google-research/google-research/tree/master/snerg). Unfortunately when I try to install jax and jaxlib for Cuda 11.8 with the following command :
pip install --upgrade jax jaxlib==0.1.69+cuda118 -f https://storage.googleapis.com/jax-releases/jax_releases.html
I get the following error:

ERROR: Ignored the following versions that require a different python version: 0.4.14 Requires-Python >=3.9
ERROR: Could not find a version that satisfies the requirement jaxlib==0.1.69+cuda118 (from versions: 0.1.32, 0.1.40, 0.1.41, 0.1.42, 0.1.43, 0.1.44, 0.1.46, 0.1.50, 0.1.51, 0.1.52, 0.1.55, 0.1.56, 0.1.57, 0.1.58, 0.1.59, 0.1.60, 0.1.61, 0.1.62, 0.1.63, 0.1.64, 0.1.65, 0.1.66, 0.1.67, 0.1.68, 0.1.69, 0.1.70, 0.1.71, 0.1.72, 0.1.73, 0.1.74, 0.1.75, 0.1.76, 0.3.0, 0.3.2, 0.3.5, 0.3.7, 0.3.8, 0.3.10, 0.3.14, 0.3.15, 0.3.18, 0.3.20, 0.3.22, 0.3.24, 0.3.25, 0.4.0, 0.4.1, 0.4.2, 0.4.3, 0.4.4, 0.4.6, 0.4.7, 0.4.9, 0.4.10, 0.4.11, 0.4.12, 0.4.13)
ERROR: No matching distribution found for jaxlib==0.1.69+cuda118

Would appreciate any help. Thanks

[jakevdp]

jaxlib version 0.1.69 is quite old. It actually pre-dated the release of CUDA 11.8 by over a year, so I would not expect to find pre-built wheels for that version combination. I suspect you'll have more luck using a more recent version of JAX.

If for some reason you absolutely need this old JAX version, your best bet would probably be to downgrade to a CUDA version more contemporary to that JAX release (circa September 2021). From the v0.1.69 installation instructions, it looks like that would be CUDA 11.3 or older.

[ayaka14732]

You can install the cuda-compat-11-3 package.

https://docs.nvidia.com/deploy/cuda-compatibility/