What jax version can be installed with CUDA 10.2?

[xmax1]

The suggested install instructions do not appear to work because the cuda 10.2 packages 'don't exist'. What jaxlib version can be installed with cuda10.2? Also, who should I contact to update the install docs it's not clear what versions work with different cuda installs.

pip3 install --upgrade jax jaxlib==0.1.64+cuda102 -f https://storage.googleapis.com/jax-releases/jax_releases.html Collecting jax Downloading https://files.pythonhosted.org/packages/81/6d/2a79504e6295e3b4909ff18af77dc3abd0e33be1f99784822016950b7b28/jax-0.2.11.tar.gz (596kB) 100% |████████████████████████████████| 604kB 1.2MB/s Collecting jaxlib==0.1.64+cuda102 Could not find a version that satisfies the requirement jaxlib==0.1.64+cuda102 (from versions: 0.1, 0.1.1, 0.1.20, 0.1.21, 0.1.22, 0.1.23) No matching distribution found for jaxlib==0.1.64+cuda102

And I tried downgrading the target jaxlib

pip3 install --upgrade jax jaxlib==0.1+cuda102 -f https://storage.googleapis.com/jax-releases/jax_releases.html --no-cache-dir Collecting jax Downloading https://files.pythonhosted.org/packages/81/6d/2a79504e6295e3b4909ff18af77dc3abd0e33be1f99784822016950b7b28/jax-0.2.11.tar.gz (596kB) 100% |████████████████████████████████| 604kB 12.0MB/s Collecting jaxlib==0.1+cuda102 Could not find a version that satisfies the requirement jaxlib==0.1+cuda102 (from versions: 0.1, 0.1.1, 0.1.20, 0.1.21, 0.1.22, 0.1.23) No matching distribution found for jaxlib==0.1+cuda102

Thanks!

[jakevdp]

You can see the available versions here: https://storage.googleapis.com/jax-releases/jax_releases.html

CUDA 10.2 wheels are available for jaxlib 0.1.55 and higher, as long as you're running Python 3.6-3.9 on a linux system.

From the available versions listed in the output of your pip command, I suspect you are running Python 3.5.

[hawkinsp]

Another possibility is your version of pip may be outdated. The versions imply that your pip does not have manylinux2010 support.

[xmax1]

@hawkinsp this was the solution.

I had just setup the environment so I assumed pip was recent.

Thanks.