JAXlib 0.1.73 + Cuda do not recognize GPU

[jecampagne]

Hi,
I experienced a problem, using the latest version of Numpyro 0.8.0. Using conda install, it downloads the jaxlib 0.1.73 version. But this version installed with cuda, does not recognize my V100 GPU ( print(jax.devices())).
Moreover if I install the Numpyro previous version, then it downloads jaxlib 0.1.69 but then jax.partial does not exists.

DO you experience the same behavior? In case, you can provide me a sequence of installations for a conda install in a dedicated environment with Numpyro working on GPU, I would be glad. Here is the CUDA version

• cuda_version: '11-3'
• cuda_release: '11.3.0-1'
• cuda_drivers_version: '465.19.01-1'
• cudnn_version: '8.2.0.53'
• nccl_version: '2.9.8'

Thanks

[jecampagne]

After fresh Anaconda3 install as well as conda env python 3.8 creation

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

jax 0.2.24 pyhd8ed1ab_0 conda-forge
jaxlib 0.1.73+cuda11.cudnn82 pypi_0 pypi
numpyro 0.8.0 pyhd8ed1ab_0 conda-forge

solves my problem. Sorry, it was certainly due to pip/conda install of an other package.