Repository: pyro-ppl/numpyro

Issue:
pip install numpyro[cuda111]: pb
Hi, 
I experience a pb with Numpyro install on a GPU (Linux)
```python
 - 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'
```

Using prescription on [this Numpyro webpage](https://pypi.org/project/numpyro/)
```python
pip install numpyro[cuda111] -f https://storage.googleapis.com/jax-releases/jax_releases.html
```

I get:
```
Looking in links: https://storage.googleapis.com/jax-releases/jax_releases.html
Collecting numpyro[cuda111]
  Downloading numpyro-0.8.0-py3-none-any.whl (264 kB)
     |████████████████████████████████| 264 kB 8.2 MB/s 
Collecting jax>=0.2.13
  Using cached jax-0.2.24-py3-none-any.whl
Collecting tqdm
  Downloading tqdm-4.62.3-py2.py3-none-any.whl (76 kB)
     |████████████████████████████████| 76 kB 2.6 MB/s 
Collecting jaxlib>=0.1.65
  Using cached https://storage.googleapis.com/jax-releases/cuda11/jaxlib-0.1.73%2Bcuda11.cudnn82-cp38-none-manylinux2010_x86_64.whl (138.6 MB)
Requirement already satisfied: numpy>=1.18 in anaconda3/envs/jaxtmp/lib/python3.8/site-packages (from jax>=0.2.13->numpyro[cuda111]) (1.21.3)
Collecting typing-extensions
  Using cached typing_extensions-3.10.0.2-py3-none-any.whl (26 kB)
Requirement already satisfied: scipy>=1.2.1 in /anaconda3/envs/jaxtmp/lib/python3.8/site-packages (from jax>=0.2.13->numpyro[cuda111]) (1.7.1)
Requirement already satisfied: absl-py in .local/lib/python3.8/site-packages (from jax>=0.2.13->numpyro[cuda111]) (0.13.0)
Requirement already satisfied: opt-einsum in.local/lib/python3.8/site-packages (from jax>=0.2.13->numpyro[cuda111]) (3.3.0)
WARNING: jax 0.2.24 does not provide the extra 'cuda111'
Requirement already satisfied: flatbuffers<3.0,>=1.12 in /.local/lib/python3.8/site-packages (from jaxlib>=0.1.65->numpyro[cuda111]) (2.0)
Collecting six
  Using cached six-1.16.0-py2.py3-none-any.whl (11 kB)
Installing collected packages: six, typing-extensions, tqdm, jaxlib, jax, numpyro
ERROR: pip's dependency resolver does not currently take into account all the packages that are installed. This behaviour is the source of the following dependency conflicts.
tensorflow-probability 0.13.0 requires cloudpickle>=1.3, which is not installed.
tensorflow-probability 0.13.0 requires decorator, which is not installed.
Successfully installed jax-0.2.24 jaxlib-0.1.73+cuda11.cudnn82 numpyro-0.8.0 six-1.16.0 tqdm-4.62.3 typing-extensions-3.10.0.2
```
Some listing via pip
```python
(jaxtmp) ccwgige016:$ pip list
Package                Version                         Location
---------------------- ------------------------------- ---------------------------------------------
absl-py                0.13.0
appdirs                1.4.4
certifi                2021.10.8
configparser           5.0.2
dill                   0.3.4
dm-tree                0.1.6
flatbuffers            2.0
gast                   0.5.1
jax                    0.2.24
jax-cosmo              0.1rc8.dev27+g3e30b6c.d20211024 /sps/lsstcest/users/campagne/jax_cosmo_forked
jaxlib                 0.1.73+cuda11.cudnn82
numpy                  1.21.3
numpyro                0.8.0
opt-einsum             3.3.0
pip                    21.2.4
scipy                  1.7.1
setuptools             58.0.4
six                    1.16.0
tensorflow-probability 0.13.0
torch                  1.9.0
tqdm                   4.62.3
typing-extensions      3.10.0.2
wheel                  0.37.0
```
idem with conda
```python
# packages in environment at /sps/lsst/users/campagne/anaconda3/envs/jaxtmp:
#
# Name                    Version                   Build  Channel
_libgcc_mutex             0.1                        main  
_openmp_mutex             4.5                       1_gnu  
ca-certificates           2021.9.30            h06a4308_1  
certifi                   2021.10.8        py38h06a4308_0  
jax                       0.2.24                   pypi_0    pypi
jaxlib                    0.1.73+cuda11.cudnn82          pypi_0    pypi
ld_impl_linux-64          2.35.1               h7274673_9  
libffi                    3.3                  he6710b0_2  
libgcc-ng                 9.3.0               h5101ec6_17  
libgomp                   9.3.0               h5101ec6_17  
libstdcxx-ng              9.3.0               hd4cf53a_17  
ncurses                   6.2                  he6710b0_1  
numpy                     1.21.3                   pypi_0    pypi
numpyro                   0.8.0                    pypi_0    pypi
openssl                   1.1.1l               h7f8727e_0  
pip                       21.2.4           py38h06a4308_0  
python                    3.8.12               h12debd9_0  
readline                  8.1                  h27cfd23_0  
scipy                     1.7.1                    pypi_0    pypi
setuptools                58.0.4           py38h06a4308_0  
six                       1.16.0                   pypi_0    pypi
sqlite                    3.36.0               hc218d9a_0  
tk                        8.6.11               h1ccaba5_0  
tqdm                      4.62.3                   pypi_0    pypi
typing-extensions         3.10.0.2                 pypi_0    pypi
wheel                     0.37.0             pyhd3eb1b0_1  
xz                        5.2.5                h7b6447c_0  
zlib                      1.2.11               h7b6447c_3 
```
Now try to do something
```python
(jaxtmp) ccwgige016:$ python
Python 3.8.12 (default, Oct 12 2021, 13:49:34) 
[GCC 7.5.0] :: Anaconda, Inc. on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import jax
ModuleNotFoundError: No module named 'numpy.core'
ImportError: numpy.core._multiarray_umath failed to import
ImportError: numpy.core.umath failed to import
2021-10-27 12:44:41.774617: F external/org_tensorflow/tensorflow/compiler/xla/python/xla.cc:77] Check failed: tensorflow::RegisterNumpyBfloat16() 
Aborted
```
Do you have a solution? Thanks


Discussion:
I guess this is due a recent breaking change of jaxlib. You can try to follow [this guide](https://jax.readthedocs.io/en/latest/changelog.html#jaxlib-0-1-73-oct-18-2021).

Resolving patch:
diff --git a/setup.py b/setup.py
--- a/setup.py
+++ b/setup.py
@@ -9,12 +9,6 @@
 from setuptools import find_packages, setup
 
 PROJECT_PATH = os.path.dirname(os.path.abspath(__file__))
-_available_cuda_versions = [
-    "101",
-    "102",
-    "110",
-    "111",
-]  # TODO: align these with what's available in JAX before release
 _jax_version_constraints = ">=0.2.13"
 _jaxlib_version_constraints = ">=0.1.65"
 
@@ -80,12 +74,9 @@
         ],
         "cpu": f"jax[cpu]{_jax_version_constraints}",
         # TPU and CUDA installations, currently require to add package repository URL, i.e.,
-        # pip install numpyro[cuda101] -f https://storage.googleapis.com/jax-releases/jax_releases.html
+        # pip install numpyro[cuda] -f https://storage.googleapis.com/jax-releases/jax_releases.html
         "tpu": f"jax[tpu]{_jax_version_constraints}",
-        **{
-            f"cuda{version}": f"jax[cuda{version}]{_jax_version_constraints}"
-            for version in _available_cuda_versions
-        },
+        "cuda": f"jax[cuda]{_jax_version_constraints}",
     },
     long_description=long_description,
     long_description_content_type="text/markdown",


Test patch: