Installing jax[cuda101]

[LUOBO123LUOBO123]

I try to install jax with the following command.
pip install --upgrade pip
pip install --upgrade "jax[cuda101]" -f https://storage.googleapis.com/jax-releases/jax_releases.html

The following error occurred during installation.
ERROR: Packages installed from PyPI cannot depend on packages which are not also hosted on PyPI.
jax depends on jaxlib@ https://storage.googleapis.com/jax-releases/cuda101/jaxlib-0.1.67+cuda101-cp38-none-manylinux2010_x86_64.whl

[lumip]

This is a problem with the 0.2.15 release, which seems to be what is getting installed - can you verify that is correct? If so, can you try forcing it to install v 0.2.16 (or 0.2.14):

pip install --upgrade "jax[cuda101]==0.2.16" -f https://storage.googleapis.com/jax-releases/jax_releases.html

[jakevdp]

Agreed with @lumip – I suspect the issue is that you have 0.2.15 installed locally, which has some problematic code for the cuda wheel path that was fixed in 0.2.16. You can either specify the version as recommended above, or first upgrade jax alone before installing the cuda101 wheel:

pip install --upgrade jax
pip install "jax[cuda101]" -f https://storage.googleapis.com/jax-releases/jax_releases.html

[LUOBO123LUOBO123]

This is a problem with the 0.2.15 release, which seems to be what is getting installed - can you verify that is correct? If so, can you try forcing it to install v 0.2.16 (or 0.2.14):

pip install --upgrade "jax[cuda101]==0.2.16" -f https://storage.googleapis.com/jax-releases/jax_releases.html

Thank you for your reply.!This method does not work. I install jax from source.

[LUOBO123LUOBO123]

This is a problem with the 0.2.15 release, which seems to be what is getting installed - can you verify that is correct? If so, can you try forcing it to install v 0.2.16 (or 0.2.14):

pip install --upgrade "jax[cuda101]==0.2.16" -f https://storage.googleapis.com/jax-releases/jax_releases.html

Thank you for your reply.!This method does not work. I install jax from source.

Thank you for your reply.!This method does not work. I install jax from source.
According to your method, the following error will appear.

[lumip]

You are having a connection issue with the jax release package index there. The final error message seem to come from pip falling back to the default (PyPI), after the (remaining) retry count hit 0. That one does not have cuda versions.
I'm not sure how to fix it, but maybe it was a temporary issue and works now if you try again. Or maybe @jakevdp knows something about this.

[jakevdp]

Is it possible you have a firewall that is preventing your system from accessing https://storage.googleapis.com/jax-releases/jax_releases.html?

[PeterZhizhin]

I'm having the same problem:

PS C:\Users\peter> python -m pip install --user "jax[cuda111]"  -f https://storage.googleapis.com/jax-releases/jax_releases.html
Looking in links: https://storage.googleapis.com/jax-releases/jax_releases.html
Collecting jax[cuda111]
  Using cached jax-0.2.18.tar.gz (708 kB)
Requirement already satisfied: numpy>=1.17 in c:\users\peter\python3\lib\site-packages (from jax[cuda111]) (1.18.5)
Requirement already satisfied: absl-py in c:\users\peter\python3\lib\site-packages (from jax[cuda111]) (0.10.0)
Requirement already satisfied: opt_einsum in c:\users\peter\python3\lib\site-packages (from jax[cuda111]) (3.3.0)
  Using cached jax-0.2.17.tar.gz (693 kB)
  Using cached jax-0.2.16.tar.gz (680 kB)
  Using cached jax-0.2.15.tar.gz (681 kB)
ERROR: Packages installed from PyPI cannot depend on packages which are not also hosted on PyPI.
jax depends on jaxlib@ https://storage.googleapis.com/jax-releases/cuda111/jaxlib-0.1.67+cuda111-cp38-none-manylinux2010_x86_64.whl

I can access both URLs: https://storage.googleapis.com/jax-releases/jax_releases.html and https://storage.googleapis.com/jax-releases/cuda111/jaxlib-0.1.67+cuda111-cp38-none-manylinux2010_x86_64.whl

Maybe, this has something to do with the fact that I'm installing on Windows. I see that Windows PyPi packages are not yet provided.

[jakevdp]

We do not currently provide jaxlib wheels for windows; see #5795.