jax cuda installation can't find wheel

[kefirski]

Trying to install jax[cuda] inside of docker container with nvidia/cuda:11.4.2-devel-ubuntu20.04 image, but get installation error

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

ERROR: Could not find a version that satisfies the requirement jaxlib==0.3.10+cuda11.cudnn82 (from jax[cuda]) (from versions: 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)
ERROR: No matching distribution found for jaxlib==0.3.10+cuda11.cudnn82 (from jax[cuda])

python3 --version
nvcc --version

Python 3.9.5
nvcc: NVIDIA (R) Cuda compiler driver
Copyright (c) 2005-2021 NVIDIA Corporation
Built on Mon_Oct_11_21:27:02_PDT_2021
Cuda compilation tools, release 11.4, V11.4.152
Build cuda_11.4.r11.4/compiler.30521435_0

[cryptozealot]

pip install --upgrade "jax[cuda]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
inside the container or add it in docker file

[jakevdp]

Yes - the key change is that CUDA wheels are now at jax_cuda_releases.html rather than jax_releases.html.