Install old JAX version (0.2.8) with CUDA and CUDNN

[hmonkey17]

Greetings everyone! I'm currently in the process of installing JAX==0.2.8 and jaxlib==0.1.59 with GPU acceleration to run some old codes in a two-year-old paper.

In a newly created conda environment, executing the command
pip install jax[cuda11_pip] -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
works seamlessly, resulting in the installation of a very recent JAX version.

However, attempting to install version 0.2.8 using the command
pip install jax[cuda11_pip]==0.2.8 -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
fails and raises an error stating "jax 0.2.8 does not provide the extra 'cuda11-pip'".

On the other hand, using the command
pip install jaxlib==0.1.59+cuda111 jax==0.2.8 -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
succeeds, but the CUDA dependencies are incomplete.

I'm seeking guidance on the correct method to install the dependencies. Your assistance is greatly appreciated! Thank you in advance!

[jakevdp]

cuda11_pip was only added in JAX v0.4.11 (pip-installable CUDA binaries are a relatively recent development).

You can look at old versions of the documentation to see the installation instructions from that era: https://github.com/google/jax/tree/jax-v0.2.8#pip-installation. This is what it recomments, which is close to the third command in your question:

pip install jax==0.2.8
pip install jaxlib==0.1.57+cuda110 -f https://storage.googleapis.com/jax-releases/jax_releases.html

jaxlib v0.1.59 may also be compatible with jax v0.2.8: you'll have to double-check jax/jaxlib version compatibility, that was before we had a predictable version numbering scheme between the two distributions.

This requires Python 3.8 or older, so unfortunately at the moment I currently have no easy way to check it myself.

[hmonkey17]

Thank you for the information!