Jax installation error: Jax requires jaxlib to be installed (Multiple cuda installations)

[tonygracious]

In my system there are multiple versions of cuda installation

 /usr/local$ ls
bin  cuda  cuda-10.0  cuda-10.1  cuda-11.2  dcgm  etc  games  include  lib  man  sbin  share  src

nvcc --version
nvcc: NVIDIA (R) Cuda compiler driver
Copyright (c) 2005-2020 NVIDIA Corporation
Built on Mon_Nov_30_19:08:53_PST_2020
Cuda compilation tools, release 11.2, V11.2.67
Build cuda_11.2.r11.2/compiler.29373293_0

System have Python 3.7.3. I have installed jax, but when I am importing jax I am getting following error,

ModuleNotFoundError: jax requires jaxlib to be installed.

Can anyone help me in fixing this error?

cudnn_version.h file in /usr/local/cuda-11.2/include have the following headers,
`#ifndef CUDNN_VERSION_H_
#define CUDNN_VERSION_H_

#define CUDNN_MAJOR 8
#define CUDNN_MINOR 2
#define CUDNN_PATCHLEVEL 1

#define CUDNN_VERSION (CUDNN_MAJOR * 1000 + CUDNN_MINOR * 100 + CUDNN_PATCHLEVEL)

#endif /* CUDNN_VERSION_H */`

[YouJiacheng]

Did you install jax by

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

[tonygracious]

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

It gave me the following error. "No matching distribution found for jaxlib==0.3.10+cuda11.cudnn82; extra == "cuda" (from jax[cuda])"

Looking in links: https://storage.googleapis.com/jax-releases/jax_releases.html
Collecting jax[cuda]
Requirement already satisfied, skipping upgrade: numpy>=1.19 in ./anaconda3/lib/python3.7/site-packages (from jax[cuda]) (1.19.5)
Requirement already satisfied, skipping upgrade: absl-py in ./anaconda3/lib/python3.7/site-packages (from jax[cuda]) (0.13.0)
Requirement already satisfied, skipping upgrade: scipy>=1.2.1 in ./anaconda3/lib/python3.7/site-packages (from jax[cuda]) (1.7.3)
Requirement already satisfied, skipping upgrade: opt-einsum in ./anaconda3/lib/python3.7/site-packages (from jax[cuda]) (3.3.0)
Requirement already satisfied, skipping upgrade: typing-extensions in ./anaconda3/lib/python3.7/site-packages (from jax[cuda]) (3.7.4.3)
Collecting jaxlib==0.3.10+cuda11.cudnn82; extra == "cuda" (from jax[cuda])
ERROR: Could not find a version that satisfies the requirement jaxlib==0.3.10+cuda11.cudnn82; extra == "cuda" (from jax[cuda]) (from versions: 0.1.23, 0.1.29, 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.52+cuda100, 0.1.52+cuda101, 0.1.52+cuda102, 0.1.52+cuda110, 0.1.55, 0.1.55+cuda100, 0.1.55+cuda101, 0.1.55+cuda102, 0.1.55+cuda110, 0.1.56, 0.1.56+cuda100, 0.1.56+cuda101, 0.1.56+cuda102, 0.1.56+cuda110, 0.1.56+cuda111, 0.1.57, 0.1.57+cuda101, 0.1.57+cuda102, 0.1.57+cuda110, 0.1.57+cuda111, 0.1.58, 0.1.58+cuda101, 0.1.58+cuda102, 0.1.58+cuda110, 0.1.58+cuda111, 0.1.59, 0.1.59+cuda101, 0.1.59+cuda102, 0.1.59+cuda110, 0.1.59+cuda111, 0.1.60, 0.1.60+cuda101, 0.1.60+cuda102, 0.1.60+cuda110, 0.1.60+cuda111, 0.1.61, 0.1.61+cuda101, 0.1.61+cuda102, 0.1.61+cuda110, 0.1.61+cuda111, 0.1.61+cuda112, 0.1.62, 0.1.62+cuda101, 0.1.62+cuda102, 0.1.62+cuda110, 0.1.62+cuda111, 0.1.62+cuda112, 0.1.63, 0.1.63+cuda101, 0.1.63+cuda102, 0.1.63+cuda110, 0.1.63+cuda111, 0.1.63+cuda112, 0.1.64, 0.1.64+cuda101, 0.1.64+cuda102, 0.1.64+cuda110, 0.1.64+cuda111, 0.1.64+cuda112, 0.1.65, 0.1.65+cuda101, 0.1.65+cuda102, 0.1.65+cuda110, 0.1.65+cuda111, 0.1.65+cuda112, 0.1.66, 0.1.66+cuda101, 0.1.66+cuda102, 0.1.66+cuda110, 0.1.66+cuda111, 0.1.67, 0.1.67+cuda101, 0.1.67+cuda102, 0.1.67+cuda110, 0.1.67+cuda111, 0.1.68, 0.1.68+cuda101, 0.1.68+cuda102, 0.1.68+cuda110, 0.1.68+cuda111, 0.1.69, 0.1.69+cuda101, 0.1.69+cuda102, 0.1.69+cuda110, 0.1.69+cuda111, 0.1.70, 0.1.70+cuda101, 0.1.70+cuda102, 0.1.70+cuda110, 0.1.70+cuda111, 0.1.71, 0.1.71+cuda102, 0.1.71+cuda110, 0.1.71+cuda111, 0.1.72, 0.1.72+cuda111, 0.1.73, 0.1.73+cuda11.cudnn805, 0.1.73+cuda11.cudnn82, 0.1.74, 0.1.74+cuda11.cudnn805, 0.1.74+cuda11.cudnn82, 0.1.75, 0.1.75+cuda11.cudnn805, 0.1.75+cuda11.cudnn82, 0.1.76, 0.1.76+cuda11.cudnn805, 0.1.76+cuda11.cudnn82, 0.3.0, 0.3.0+cuda11.cudnn805, 0.3.0+cuda11.cudnn82, 0.3.2, 0.3.2+cuda11.cudnn805, 0.3.2+cuda11.cudnn82, 0.3.5, 0.3.5+cuda11.cudnn805, 0.3.5+cuda11.cudnn82)
ERROR: No matching distribution found for jaxlib==0.3.10+cuda11.cudnn82; extra == "cuda" (from jax[cuda])

[YouJiacheng]

Hi - could you try upgrading your python and pip first?
I think your python and pip version is only compatible with manylinux2010 platform tag, but JAX only provide wheels with manylinux2014 platform tag for jaxlib-0.3.10.
manylinux2014 requires python 3.7.8+/3.8.4+/3.9.0+ and pip >= 19.3.

[tonygracious]

Thanks it worked

[kbirger]

My dockerfile looks like this:

RUN apt-get update && apt-get install -y \
  git \
  python3 \
  python3-pip \
  && rm -rf /var/lib/apt/lists/*

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

should be latest python and pip. I'm still gettings the same error

[YouJiacheng]

@kbirger You can open another question, and provide a complete error log.