Broken import with latest jax version

[michaelosthege]

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:

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.

[fehiepsi]

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