Repository: pyro-ppl/numpyro

Issue:
Broken import with latest jax version
We're seeing this error in our CI pipeline:
```
/usr/share/miniconda/envs/pymc3-dev-py39/lib/python3.9/site-packages/numpyro/__init__.py:6: in <module>
    from numpyro import compat, diagnostics, distributions, handlers, infer, optim
/usr/share/miniconda/envs/pymc3-dev-py39/lib/python3.9/site-packages/numpyro/infer/__init__.py:12: in <module>
    from numpyro.infer.hmc import HMC, NUTS
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

    from collections import OrderedDict, namedtuple
    import math
    import os
    
>   from jax import device_put, lax, partial, random, vmap
E   ImportError: cannot import name 'partial' from 'jax' (/usr/share/miniconda/envs/pymc3-dev-py39/lib/python3.9/site-packages/jax/__init__.py)

/usr/share/miniconda/envs/pymc3-dev-py39/lib/python3.9/site-packages/numpyro/infer/hmc.py:8: ImportError
```

Quoting the [v0.2.21 release notes](https://github.com/google/jax/releases/tag/jax-v0.2.21):

> jax.partial, jax.lax.partial, and jax.util.partial were accidental
exports that have now been removed. Use functools.partial from the Python
standard library instead.

So it'll be an easy fix for you.


Discussion:
Thanks for the ping, @michaelosthege! I'll update this shortly.

Resolving patch:
diff --git a/numpyro/handlers.py b/numpyro/handlers.py
--- a/numpyro/handlers.py
+++ b/numpyro/handlers.py
@@ -674,11 +674,11 @@ class seed(Messenger):
 
     def __init__(self, fn=None, rng_seed=None):
         if isinstance(rng_seed, int) or (
-            isinstance(rng_seed, jnp.ndarray) and not jnp.shape(rng_seed)
+            isinstance(rng_seed, (np.ndarray, jnp.ndarray)) and not jnp.shape(rng_seed)
         ):
             rng_seed = random.PRNGKey(rng_seed)
         if not (
-            isinstance(rng_seed, jnp.ndarray)
+            isinstance(rng_seed, (np.ndarray, jnp.ndarray))
             and rng_seed.dtype == jnp.uint32
             and rng_seed.shape == (2,)
         ):
diff --git a/numpyro/infer/hmc.py b/numpyro/infer/hmc.py
--- a/numpyro/infer/hmc.py
+++ b/numpyro/infer/hmc.py
@@ -2,10 +2,11 @@
 # SPDX-License-Identifier: Apache-2.0
 
 from collections import OrderedDict, namedtuple
+from functools import partial
 import math
 import os
 
-from jax import device_put, lax, partial, random, vmap
+from jax import device_put, lax, random, vmap
 from jax.flatten_util import ravel_pytree
 import jax.numpy as jnp
 
diff --git a/numpyro/infer/mcmc.py b/numpyro/infer/mcmc.py
--- a/numpyro/infer/mcmc.py
+++ b/numpyro/infer/mcmc.py
@@ -7,6 +7,8 @@
 import os
 import warnings
 
+import numpy as np
+
 from jax import jit, lax, local_device_count, pmap, random, vmap
 from jax.core import Tracer
 from jax.interpreters.xla import DeviceArray
@@ -190,7 +192,7 @@ def _hashable(x):
         return x
     elif isinstance(x, DeviceArray):
         return x.copy().tobytes()
-    elif isinstance(x, jnp.ndarray):
+    elif isinstance(x, (np.ndarray, jnp.ndarray)):
         return x.tobytes()
     return x
 


Test patch:
diff --git a/test/contrib/test_optim.py b/test/contrib/test_optim.py
--- a/test/contrib/test_optim.py
+++ b/test/contrib/test_optim.py
@@ -1,10 +1,12 @@
 # Copyright Contributors to the Pyro project.
 # SPDX-License-Identifier: Apache-2.0
 
+from functools import partial
+
 from numpy.testing import assert_allclose
 import pytest
 
-from jax import grad, jit, partial, random
+from jax import grad, jit, random
 from jax.lax import fori_loop
 import jax.numpy as jnp
 from jax.test_util import check_close
diff --git a/test/infer/test_compute_downstream_costs.py b/test/infer/test_compute_downstream_costs.py
--- a/test/infer/test_compute_downstream_costs.py
+++ b/test/infer/test_compute_downstream_costs.py
@@ -304,7 +304,7 @@ def test_compute_downstream_costs_big_model_guide_pair(
 
     for k in dc:
         assert guide_trace[k]["log_prob"].shape == dc[k].shape
-        assert_allclose(dc[k], dc_brute[k])
+        assert_allclose(dc[k], dc_brute[k], rtol=2e-7)
 
 
 def plate_reuse_model_guide(include_obs=True, dim1=3, dim2=2):
diff --git a/test/test_optimizers.py b/test/test_optimizers.py
--- a/test/test_optimizers.py
+++ b/test/test_optimizers.py
@@ -1,9 +1,11 @@
 # Copyright Contributors to the Pyro project.
 # SPDX-License-Identifier: Apache-2.0
 
+from functools import partial
+
 import pytest
 
-from jax import grad, jit, partial
+from jax import grad, jit
 import jax.numpy as jnp
 
 from numpyro import optim