We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN 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

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `setup.py`
Content:
```
1 # Copyright Contributors to the Pyro project.
2 # SPDX-License-Identifier: Apache-2.0
3 
4 from __future__ import absolute_import, division, print_function
5 
6 import os
7 import sys
8 
9 from setuptools import find_packages, setup
10 
11 PROJECT_PATH = os.path.dirname(os.path.abspath(__file__))
12 _available_cuda_versions = [
13     "101",
14     "102",
15     "110",
16     "111",
17 ]  # TODO: align these with what's available in JAX before release
18 _jax_version_constraints = ">=0.2.13"
19 _jaxlib_version_constraints = ">=0.1.65"
20 
21 # Find version
22 for line in open(os.path.join(PROJECT_PATH, "numpyro", "version.py")):
23     if line.startswith("__version__ = "):
24         version = line.strip().split()[2][1:-1]
25 
26 # READ README.md for long description on PyPi.
27 try:
28     long_description = open("README.md", encoding="utf-8").read()
29 except Exception as e:
30     sys.stderr.write("Failed to read README.md:\n  {}\n".format(e))
31     sys.stderr.flush()
32     long_description = ""
33 
34 setup(
35     name="numpyro",
36     version=version,
37     description="Pyro PPL on NumPy",
38     packages=find_packages(include=["numpyro", "numpyro.*"]),
39     url="https://github.com/pyro-ppl/numpyro",
40     author="Uber AI Labs",
41     install_requires=[
42         f"jax{_jax_version_constraints}",
43         f"jaxlib{_jaxlib_version_constraints}",
44         "tqdm",
45     ],
46     extras_require={
47         "doc": [
48             "ipython",  # sphinx needs this to render codes
49             "nbsphinx>=0.8.5",
50             "readthedocs-sphinx-search==0.1.0",
51             "sphinx",
52             "sphinx_rtd_theme",
53             "sphinx-gallery",
54         ],
55         "test": [
56             "black[jupyter]>=21.8b0",
57             "flake8",
58             "isort>=5.0",
59             "pytest>=4.1",
60             "pyro-api>=0.1.1",
61             "scipy>=1.1",
62         ],
63         "dev": [
64             "dm-haiku",
65             "flax",
66             "funsor==0.4.1",
67             "graphviz",
68             "jaxns==0.0.7",
69             "optax>=0.0.6",
70             "tensorflow_probability>=0.13",
71         ],
72         "examples": [
73             "arviz",
74             "jupyter",
75             "matplotlib",
76             "pandas",
77             "seaborn",
78             "scikit-learn",
79             "wordcloud",
80         ],
81         "cpu": f"jax[cpu]{_jax_version_constraints}",
82         # TPU and CUDA installations, currently require to add package repository URL, i.e.,
83         # pip install numpyro[cuda101] -f https://storage.googleapis.com/jax-releases/jax_releases.html
84         "tpu": f"jax[tpu]{_jax_version_constraints}",
85         **{
86             f"cuda{version}": f"jax[cuda{version}]{_jax_version_constraints}"
87             for version in _available_cuda_versions
88         },
89     },
90     long_description=long_description,
91     long_description_content_type="text/markdown",
92     keywords="probabilistic machine learning bayesian statistics",
93     license="Apache License 2.0",
94     classifiers=[
95         "Intended Audience :: Developers",
96         "Intended Audience :: Education",
97         "Intended Audience :: Science/Research",
98         "License :: OSI Approved :: Apache Software License",
99         "Operating System :: POSIX :: Linux",
100         "Operating System :: MacOS :: MacOS X",
101         "Programming Language :: Python :: 3.6",
102         "Programming Language :: Python :: 3.7",
103         "Programming Language :: Python :: 3.8",
104         "Programming Language :: Python :: 3.9",
105     ],
106 )
107
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

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",

{"golden_diff": "diff --git a/setup.py b/setup.py\n--- a/setup.py\n+++ b/setup.py\n@@ -9,12 +9,6 @@\n from setuptools import find_packages, setup\n \n PROJECT_PATH = os.path.dirname(os.path.abspath(__file__))\n-_available_cuda_versions = [\n-    \"101\",\n-    \"102\",\n-    \"110\",\n-    \"111\",\n-]  # TODO: align these with what's available in JAX before release\n _jax_version_constraints = \">=0.2.13\"\n _jaxlib_version_constraints = \">=0.1.65\"\n \n@@ -80,12 +74,9 @@\n         ],\n         \"cpu\": f\"jax[cpu]{_jax_version_constraints}\",\n         # TPU and CUDA installations, currently require to add package repository URL, i.e.,\n-        # pip install numpyro[cuda101] -f https://storage.googleapis.com/jax-releases/jax_releases.html\n+        # pip install numpyro[cuda] -f https://storage.googleapis.com/jax-releases/jax_releases.html\n         \"tpu\": f\"jax[tpu]{_jax_version_constraints}\",\n-        **{\n-            f\"cuda{version}\": f\"jax[cuda{version}]{_jax_version_constraints}\"\n-            for version in _available_cuda_versions\n-        },\n+        \"cuda\": f\"jax[cuda]{_jax_version_constraints}\",\n     },\n     long_description=long_description,\n     long_description_content_type=\"text/markdown\",\n", "issue": "pip install numpyro[cuda111]: pb\nHi, \r\nI experience a pb with Numpyro install on a GPU (Linux)\r\n```python\r\n - cuda_version: '11-3'\r\n- cuda_release: '11.3.0-1' \r\n- cuda_drivers_version: '465.19.01-1' \r\n- cudnn_version: '8.2.0.53'\r\n- nccl_version: '2.9.8'\r\n```\r\n\r\nUsing prescription on [this Numpyro webpage](https://pypi.org/project/numpyro/)\r\n```python\r\npip install numpyro[cuda111] -f https://storage.googleapis.com/jax-releases/jax_releases.html\r\n```\r\n\r\nI get:\r\n```\r\nLooking in links: https://storage.googleapis.com/jax-releases/jax_releases.html\r\nCollecting numpyro[cuda111]\r\n  Downloading numpyro-0.8.0-py3-none-any.whl (264 kB)\r\n     |\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 264 kB 8.2 MB/s \r\nCollecting jax>=0.2.13\r\n  Using cached jax-0.2.24-py3-none-any.whl\r\nCollecting tqdm\r\n  Downloading tqdm-4.62.3-py2.py3-none-any.whl (76 kB)\r\n     |\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588\u2588| 76 kB 2.6 MB/s \r\nCollecting jaxlib>=0.1.65\r\n  Using cached https://storage.googleapis.com/jax-releases/cuda11/jaxlib-0.1.73%2Bcuda11.cudnn82-cp38-none-manylinux2010_x86_64.whl (138.6 MB)\r\nRequirement already satisfied: numpy>=1.18 in anaconda3/envs/jaxtmp/lib/python3.8/site-packages (from jax>=0.2.13->numpyro[cuda111]) (1.21.3)\r\nCollecting typing-extensions\r\n  Using cached typing_extensions-3.10.0.2-py3-none-any.whl (26 kB)\r\nRequirement 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)\r\nRequirement already satisfied: absl-py in .local/lib/python3.8/site-packages (from jax>=0.2.13->numpyro[cuda111]) (0.13.0)\r\nRequirement already satisfied: opt-einsum in.local/lib/python3.8/site-packages (from jax>=0.2.13->numpyro[cuda111]) (3.3.0)\r\nWARNING: jax 0.2.24 does not provide the extra 'cuda111'\r\nRequirement already satisfied: flatbuffers<3.0,>=1.12 in /.local/lib/python3.8/site-packages (from jaxlib>=0.1.65->numpyro[cuda111]) (2.0)\r\nCollecting six\r\n  Using cached six-1.16.0-py2.py3-none-any.whl (11 kB)\r\nInstalling collected packages: six, typing-extensions, tqdm, jaxlib, jax, numpyro\r\nERROR: 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.\r\ntensorflow-probability 0.13.0 requires cloudpickle>=1.3, which is not installed.\r\ntensorflow-probability 0.13.0 requires decorator, which is not installed.\r\nSuccessfully 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\r\n```\r\nSome listing via pip\r\n```python\r\n(jaxtmp) ccwgige016:$ pip list\r\nPackage                Version                         Location\r\n---------------------- ------------------------------- ---------------------------------------------\r\nabsl-py                0.13.0\r\nappdirs                1.4.4\r\ncertifi                2021.10.8\r\nconfigparser           5.0.2\r\ndill                   0.3.4\r\ndm-tree                0.1.6\r\nflatbuffers            2.0\r\ngast                   0.5.1\r\njax                    0.2.24\r\njax-cosmo              0.1rc8.dev27+g3e30b6c.d20211024 /sps/lsstcest/users/campagne/jax_cosmo_forked\r\njaxlib                 0.1.73+cuda11.cudnn82\r\nnumpy                  1.21.3\r\nnumpyro                0.8.0\r\nopt-einsum             3.3.0\r\npip                    21.2.4\r\nscipy                  1.7.1\r\nsetuptools             58.0.4\r\nsix                    1.16.0\r\ntensorflow-probability 0.13.0\r\ntorch                  1.9.0\r\ntqdm                   4.62.3\r\ntyping-extensions      3.10.0.2\r\nwheel                  0.37.0\r\n```\r\nidem with conda\r\n```python\r\n# packages in environment at /sps/lsst/users/campagne/anaconda3/envs/jaxtmp:\r\n#\r\n# Name                    Version                   Build  Channel\r\n_libgcc_mutex             0.1                        main  \r\n_openmp_mutex             4.5                       1_gnu  \r\nca-certificates           2021.9.30            h06a4308_1  \r\ncertifi                   2021.10.8        py38h06a4308_0  \r\njax                       0.2.24                   pypi_0    pypi\r\njaxlib                    0.1.73+cuda11.cudnn82          pypi_0    pypi\r\nld_impl_linux-64          2.35.1               h7274673_9  \r\nlibffi                    3.3                  he6710b0_2  \r\nlibgcc-ng                 9.3.0               h5101ec6_17  \r\nlibgomp                   9.3.0               h5101ec6_17  \r\nlibstdcxx-ng              9.3.0               hd4cf53a_17  \r\nncurses                   6.2                  he6710b0_1  \r\nnumpy                     1.21.3                   pypi_0    pypi\r\nnumpyro                   0.8.0                    pypi_0    pypi\r\nopenssl                   1.1.1l               h7f8727e_0  \r\npip                       21.2.4           py38h06a4308_0  \r\npython                    3.8.12               h12debd9_0  \r\nreadline                  8.1                  h27cfd23_0  \r\nscipy                     1.7.1                    pypi_0    pypi\r\nsetuptools                58.0.4           py38h06a4308_0  \r\nsix                       1.16.0                   pypi_0    pypi\r\nsqlite                    3.36.0               hc218d9a_0  \r\ntk                        8.6.11               h1ccaba5_0  \r\ntqdm                      4.62.3                   pypi_0    pypi\r\ntyping-extensions         3.10.0.2                 pypi_0    pypi\r\nwheel                     0.37.0             pyhd3eb1b0_1  \r\nxz                        5.2.5                h7b6447c_0  \r\nzlib                      1.2.11               h7b6447c_3 \r\n```\r\nNow try to do something\r\n```python\r\n(jaxtmp) ccwgige016:$ python\r\nPython 3.8.12 (default, Oct 12 2021, 13:49:34) \r\n[GCC 7.5.0] :: Anaconda, Inc. on linux\r\nType \"help\", \"copyright\", \"credits\" or \"license\" for more information.\r\n>>> import jax\r\nModuleNotFoundError: No module named 'numpy.core'\r\nImportError: numpy.core._multiarray_umath failed to import\r\nImportError: numpy.core.umath failed to import\r\n2021-10-27 12:44:41.774617: F external/org_tensorflow/tensorflow/compiler/xla/python/xla.cc:77] Check failed: tensorflow::RegisterNumpyBfloat16() \r\nAborted\r\n```\r\nDo you have a solution? Thanks\n", "before_files": [{"content": "# Copyright Contributors to the Pyro project.\n# SPDX-License-Identifier: Apache-2.0\n\nfrom __future__ import absolute_import, division, print_function\n\nimport os\nimport sys\n\nfrom setuptools import find_packages, setup\n\nPROJECT_PATH = os.path.dirname(os.path.abspath(__file__))\n_available_cuda_versions = [\n    \"101\",\n    \"102\",\n    \"110\",\n    \"111\",\n]  # TODO: align these with what's available in JAX before release\n_jax_version_constraints = \">=0.2.13\"\n_jaxlib_version_constraints = \">=0.1.65\"\n\n# Find version\nfor line in open(os.path.join(PROJECT_PATH, \"numpyro\", \"version.py\")):\n    if line.startswith(\"__version__ = \"):\n        version = line.strip().split()[2][1:-1]\n\n# READ README.md for long description on PyPi.\ntry:\n    long_description = open(\"README.md\", encoding=\"utf-8\").read()\nexcept Exception as e:\n    sys.stderr.write(\"Failed to read README.md:\\n  {}\\n\".format(e))\n    sys.stderr.flush()\n    long_description = \"\"\n\nsetup(\n    name=\"numpyro\",\n    version=version,\n    description=\"Pyro PPL on NumPy\",\n    packages=find_packages(include=[\"numpyro\", \"numpyro.*\"]),\n    url=\"https://github.com/pyro-ppl/numpyro\",\n    author=\"Uber AI Labs\",\n    install_requires=[\n        f\"jax{_jax_version_constraints}\",\n        f\"jaxlib{_jaxlib_version_constraints}\",\n        \"tqdm\",\n    ],\n    extras_require={\n        \"doc\": [\n            \"ipython\",  # sphinx needs this to render codes\n            \"nbsphinx>=0.8.5\",\n            \"readthedocs-sphinx-search==0.1.0\",\n            \"sphinx\",\n            \"sphinx_rtd_theme\",\n            \"sphinx-gallery\",\n        ],\n        \"test\": [\n            \"black[jupyter]>=21.8b0\",\n            \"flake8\",\n            \"isort>=5.0\",\n            \"pytest>=4.1\",\n            \"pyro-api>=0.1.1\",\n            \"scipy>=1.1\",\n        ],\n        \"dev\": [\n            \"dm-haiku\",\n            \"flax\",\n            \"funsor==0.4.1\",\n            \"graphviz\",\n            \"jaxns==0.0.7\",\n            \"optax>=0.0.6\",\n            \"tensorflow_probability>=0.13\",\n        ],\n        \"examples\": [\n            \"arviz\",\n            \"jupyter\",\n            \"matplotlib\",\n            \"pandas\",\n            \"seaborn\",\n            \"scikit-learn\",\n            \"wordcloud\",\n        ],\n        \"cpu\": f\"jax[cpu]{_jax_version_constraints}\",\n        # TPU and CUDA installations, currently require to add package repository URL, i.e.,\n        # pip install numpyro[cuda101] -f https://storage.googleapis.com/jax-releases/jax_releases.html\n        \"tpu\": f\"jax[tpu]{_jax_version_constraints}\",\n        **{\n            f\"cuda{version}\": f\"jax[cuda{version}]{_jax_version_constraints}\"\n            for version in _available_cuda_versions\n        },\n    },\n    long_description=long_description,\n    long_description_content_type=\"text/markdown\",\n    keywords=\"probabilistic machine learning bayesian statistics\",\n    license=\"Apache License 2.0\",\n    classifiers=[\n        \"Intended Audience :: Developers\",\n        \"Intended Audience :: Education\",\n        \"Intended Audience :: Science/Research\",\n        \"License :: OSI Approved :: Apache Software License\",\n        \"Operating System :: POSIX :: Linux\",\n        \"Operating System :: MacOS :: MacOS X\",\n        \"Programming Language :: Python :: 3.6\",\n        \"Programming Language :: Python :: 3.7\",\n        \"Programming Language :: Python :: 3.8\",\n        \"Programming Language :: Python :: 3.9\",\n    ],\n)\n", "path": "setup.py"}], "after_files": [{"content": "# Copyright Contributors to the Pyro project.\n# SPDX-License-Identifier: Apache-2.0\n\nfrom __future__ import absolute_import, division, print_function\n\nimport os\nimport sys\n\nfrom setuptools import find_packages, setup\n\nPROJECT_PATH = os.path.dirname(os.path.abspath(__file__))\n_jax_version_constraints = \">=0.2.13\"\n_jaxlib_version_constraints = \">=0.1.65\"\n\n# Find version\nfor line in open(os.path.join(PROJECT_PATH, \"numpyro\", \"version.py\")):\n    if line.startswith(\"__version__ = \"):\n        version = line.strip().split()[2][1:-1]\n\n# READ README.md for long description on PyPi.\ntry:\n    long_description = open(\"README.md\", encoding=\"utf-8\").read()\nexcept Exception as e:\n    sys.stderr.write(\"Failed to read README.md:\\n  {}\\n\".format(e))\n    sys.stderr.flush()\n    long_description = \"\"\n\nsetup(\n    name=\"numpyro\",\n    version=version,\n    description=\"Pyro PPL on NumPy\",\n    packages=find_packages(include=[\"numpyro\", \"numpyro.*\"]),\n    url=\"https://github.com/pyro-ppl/numpyro\",\n    author=\"Uber AI Labs\",\n    install_requires=[\n        f\"jax{_jax_version_constraints}\",\n        f\"jaxlib{_jaxlib_version_constraints}\",\n        \"tqdm\",\n    ],\n    extras_require={\n        \"doc\": [\n            \"ipython\",  # sphinx needs this to render codes\n            \"nbsphinx>=0.8.5\",\n            \"readthedocs-sphinx-search==0.1.0\",\n            \"sphinx\",\n            \"sphinx_rtd_theme\",\n            \"sphinx-gallery\",\n        ],\n        \"test\": [\n            \"black[jupyter]>=21.8b0\",\n            \"flake8\",\n            \"isort>=5.0\",\n            \"pytest>=4.1\",\n            \"pyro-api>=0.1.1\",\n            \"scipy>=1.1\",\n        ],\n        \"dev\": [\n            \"dm-haiku\",\n            \"flax\",\n            \"funsor==0.4.1\",\n            \"graphviz\",\n            \"jaxns==0.0.7\",\n            \"optax>=0.0.6\",\n            \"tensorflow_probability>=0.13\",\n        ],\n        \"examples\": [\n            \"arviz\",\n            \"jupyter\",\n            \"matplotlib\",\n            \"pandas\",\n            \"seaborn\",\n            \"scikit-learn\",\n            \"wordcloud\",\n        ],\n        \"cpu\": f\"jax[cpu]{_jax_version_constraints}\",\n        # TPU and CUDA installations, currently require to add package repository URL, i.e.,\n        # pip install numpyro[cuda] -f https://storage.googleapis.com/jax-releases/jax_releases.html\n        \"tpu\": f\"jax[tpu]{_jax_version_constraints}\",\n        \"cuda\": f\"jax[cuda]{_jax_version_constraints}\",\n    },\n    long_description=long_description,\n    long_description_content_type=\"text/markdown\",\n    keywords=\"probabilistic machine learning bayesian statistics\",\n    license=\"Apache License 2.0\",\n    classifiers=[\n        \"Intended Audience :: Developers\",\n        \"Intended Audience :: Education\",\n        \"Intended Audience :: Science/Research\",\n        \"License :: OSI Approved :: Apache Software License\",\n        \"Operating System :: POSIX :: Linux\",\n        \"Operating System :: MacOS :: MacOS X\",\n        \"Programming Language :: Python :: 3.6\",\n        \"Programming Language :: Python :: 3.7\",\n        \"Programming Language :: Python :: 3.8\",\n        \"Programming Language :: Python :: 3.9\",\n    ],\n)\n", "path": "setup.py"}]}