We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Make HMCState picklable
Currently, we use `laxtuple` to make HMCState. This way however make it not possible to pickle HMCState (unless we transform it to a dict/tuple/namedtuple).

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `numpyro/mcmc.py`
Content:
```
1 import math
2 import os
3 
4 import tqdm
5 
6 import jax.numpy as np
7 from jax import jit, partial, random
8 from jax.flatten_util import ravel_pytree
9 from jax.random import PRNGKey
10 
11 import numpyro.distributions as dist
12 from numpyro.hmc_util import IntegratorState, build_tree, find_reasonable_step_size, velocity_verlet, warmup_adapter
13 from numpyro.util import cond, fori_loop, laxtuple
14 
15 HMCState = laxtuple('HMCState', ['z', 'z_grad', 'potential_energy', 'num_steps', 'accept_prob',
16                                  'step_size', 'inverse_mass_matrix', 'rng'])
17 
18 
19 def _get_num_steps(step_size, trajectory_length):
20     num_steps = np.array(trajectory_length / step_size, dtype=np.int32)
21     return np.where(num_steps < 1, np.array(1, dtype=np.int32), num_steps)
22 
23 
24 def _sample_momentum(unpack_fn, inverse_mass_matrix, rng):
25     if inverse_mass_matrix.ndim == 1:
26         r = dist.Normal(0., np.sqrt(np.reciprocal(inverse_mass_matrix))).sample(rng)
27         return unpack_fn(r)
28     elif inverse_mass_matrix.ndim == 2:
29         raise NotImplementedError
30 
31 
32 def _euclidean_ke(inverse_mass_matrix, r):
33     r, _ = ravel_pytree(r)
34 
35     if inverse_mass_matrix.ndim == 2:
36         v = np.matmul(inverse_mass_matrix, r)
37     elif inverse_mass_matrix.ndim == 1:
38         v = np.multiply(inverse_mass_matrix, r)
39 
40     return 0.5 * np.dot(v, r)
41 
42 
43 def hmc(potential_fn, kinetic_fn=None, algo='NUTS'):
44     r"""
45     Hamiltonian Monte Carlo inference, using either fixed number of
46     steps or the No U-Turn Sampler (NUTS) with adaptive path length.
47 
48     **References:**
49 
50     1. *MCMC Using Hamiltonian Dynamics*, Radford M. Neal
51     2. *The No-U-turn sampler: adaptively setting path lengths in Hamiltonian Monte Carlo*,
52        Matthew D. Hoffman, and Andrew Gelman.
53 
54     :param potential_fn: Python callable that computes the potential energy
55         given input parameters. The input parameters to `potential_fn` can be
56         any python collection type, provided that `init_samples` argument to
57         `init_kernel` has the same type.
58     :param kinetic_fn: Python callable that returns the kinetic energy given
59         inverse mass matrix and momentum. If not provided, the default is
60         euclidean kinetic energy.
61     :param str algo: Whether to run ``HMC`` with fixed number of steps or ``NUTS``
62         with adaptive path length. Default is ``NUTS``.
63     :return init_kernel, sample_kernel: Returns a tuple of callables, the first
64         one to initialize the sampler, and the second one to generate samples
65         given an existing one.
66 
67     **Example**
68 
69     .. testsetup::
70 
71         import jax
72         from jax import random
73         import jax.numpy as np
74         import numpyro.distributions as dist
75         from numpyro.handlers import sample
76         from numpyro.hmc_util import initialize_model
77         from numpyro.mcmc import hmc
78         from numpyro.util import fori_collect
79 
80     .. doctest::
81 
82         >>> true_coefs = np.array([1., 2., 3.])
83         >>> data = random.normal(random.PRNGKey(2), (2000, 3))
84         >>> dim = 3
85         >>> labels = dist.Bernoulli(logits=(true_coefs * data).sum(-1)).sample(random.PRNGKey(3))
86         >>>
87         >>> def model(data, labels):
88         ...     coefs_mean = np.zeros(dim)
89         ...     coefs = sample('beta', dist.Normal(coefs_mean, np.ones(3)))
90         ...     return sample('y', dist.Bernoulli(logits=(coefs * data).sum(-1)), obs=labels)
91         >>>
92         >>> init_params, potential_fn, constrain_fn = initialize_model(random.PRNGKey(0), model, data, labels)
93         >>> init_kernel, sample_kernel = hmc(potential_fn, algo='NUTS')
94         >>> hmc_state = init_kernel(init_params,
95         ...                         trajectory_length=10,
96         ...                         num_warmup_steps=300)
97         >>> hmc_states = fori_collect(500, sample_kernel, hmc_state,
98         ...                           transform=lambda x: constrain_fn(x.z))
99         >>> print(np.mean(hmc_states['beta'], axis=0))
100         [0.9153987 2.0754058 2.9621222]
101     """
102     if kinetic_fn is None:
103         kinetic_fn = _euclidean_ke
104     vv_init, vv_update = velocity_verlet(potential_fn, kinetic_fn)
105     trajectory_len = None
106     max_treedepth = None
107     momentum_generator = None
108     wa_update = None
109 
110     def init_kernel(init_samples,
111                     num_warmup,
112                     step_size=1.0,
113                     adapt_step_size=True,
114                     adapt_mass_matrix=True,
115                     diag_mass=True,
116                     target_accept_prob=0.8,
117                     trajectory_length=2*math.pi,
118                     max_tree_depth=10,
119                     run_warmup=True,
120                     progbar=True,
121                     rng=PRNGKey(0)):
122         """
123         Initializes the HMC sampler.
124 
125         :param init_samples: Initial parameters to begin sampling. The type can
126             must be consistent with the input type to `potential_fn`.
127         :param int num_warmup_steps: Number of warmup steps; samples generated
128             during warmup are discarded.
129         :param float step_size: Determines the size of a single step taken by the
130             verlet integrator while computing the trajectory using Hamiltonian
131             dynamics. If not specified, it will be set to 1.
132         :param bool adapt_step_size: A flag to decide if we want to adapt step_size
133             during warm-up phase using Dual Averaging scheme.
134         :param bool adapt_mass_matrix: A flag to decide if we want to adapt mass
135             matrix during warm-up phase using Welford scheme.
136         :param bool diag_mass: A flag to decide if mass matrix is diagonal (default)
137             or dense (if set to ``False``).
138         :param float target_accept_prob: Target acceptance probability for step size
139             adaptation using Dual Averaging. Increasing this value will lead to a smaller
140             step size, hence the sampling will be slower but more robust. Default to 0.8.
141         :param float trajectory_length: Length of a MCMC trajectory for HMC. Default
142             value is :math:`2\\pi`.
143         :param int max_tree_depth: Max depth of the binary tree created during the doubling
144             scheme of NUTS sampler. Defaults to 10.
145         :param bool run_warmup: Flag to decide whether warmup is run. If ``True``,
146             `init_kernel` returns an initial :data:`HMCState` that can be used to
147             generate samples using MCMC. Else, returns the arguments and callable
148             that does the initial adaptation.
149         :param bool progbar: Whether to enable progress bar updates. Defaults to
150             ``True``.
151         :param bool heuristic_step_size: If ``True``, a coarse grained adjustment of
152             step size is done at the beginning of each adaptation window to achieve
153             `target_acceptance_prob`.
154         :param jax.random.PRNGKey rng: random key to be used as the source of
155             randomness.
156         """
157         step_size = float(step_size)
158         nonlocal momentum_generator, wa_update, trajectory_len, max_treedepth
159         trajectory_len = float(trajectory_length)
160         max_treedepth = max_tree_depth
161         z = init_samples
162         z_flat, unravel_fn = ravel_pytree(z)
163         momentum_generator = partial(_sample_momentum, unravel_fn)
164 
165         find_reasonable_ss = partial(find_reasonable_step_size,
166                                      potential_fn, kinetic_fn,
167                                      momentum_generator)
168 
169         wa_init, wa_update = warmup_adapter(num_warmup,
170                                             adapt_step_size=adapt_step_size,
171                                             adapt_mass_matrix=adapt_mass_matrix,
172                                             diag_mass=diag_mass,
173                                             target_accept_prob=target_accept_prob,
174                                             find_reasonable_step_size=find_reasonable_ss)
175 
176         rng_hmc, rng_wa = random.split(rng)
177         wa_state = wa_init(z, rng_wa, step_size, mass_matrix_size=np.size(z_flat))
178         r = momentum_generator(wa_state.inverse_mass_matrix, rng)
179         vv_state = vv_init(z, r)
180         hmc_state = HMCState(vv_state.z, vv_state.z_grad, vv_state.potential_energy, 0, 0.,
181                              wa_state.step_size, wa_state.inverse_mass_matrix, rng_hmc)
182 
183         wa_update = jit(wa_update)
184         if run_warmup:
185             # JIT if progress bar updates not required
186             if not progbar:
187                 hmc_state, _ = jit(fori_loop, static_argnums=(2,))(0, num_warmup,
188                                                                    warmup_update,
189                                                                    (hmc_state, wa_state))
190             else:
191                 for i in tqdm.trange(num_warmup):
192                     hmc_state, wa_state = warmup_update(i, (hmc_state, wa_state))
193             return hmc_state
194         else:
195             return hmc_state, wa_state, warmup_update
196 
197     def warmup_update(t, states):
198         hmc_state, wa_state = states
199         hmc_state = sample_kernel(hmc_state)
200         wa_state = wa_update(t, hmc_state.accept_prob, hmc_state.z, wa_state)
201         hmc_state = hmc_state.update(step_size=wa_state.step_size,
202                                      inverse_mass_matrix=wa_state.inverse_mass_matrix)
203         return hmc_state, wa_state
204 
205     def _hmc_next(step_size, inverse_mass_matrix, vv_state, rng):
206         num_steps = _get_num_steps(step_size, trajectory_len)
207         vv_state_new = fori_loop(0, num_steps,
208                                  lambda i, val: vv_update(step_size, inverse_mass_matrix, val),
209                                  vv_state)
210         energy_old = vv_state.potential_energy + kinetic_fn(inverse_mass_matrix, vv_state.r)
211         energy_new = vv_state_new.potential_energy + kinetic_fn(inverse_mass_matrix, vv_state_new.r)
212         delta_energy = energy_new - energy_old
213         delta_energy = np.where(np.isnan(delta_energy), np.inf, delta_energy)
214         accept_prob = np.clip(np.exp(-delta_energy), a_max=1.0)
215         transition = random.bernoulli(rng, accept_prob)
216         vv_state = cond(transition,
217                         vv_state_new, lambda state: state,
218                         vv_state, lambda state: state)
219         return vv_state, num_steps, accept_prob
220 
221     def _nuts_next(step_size, inverse_mass_matrix, vv_state, rng):
222         binary_tree = build_tree(vv_update, kinetic_fn, vv_state,
223                                  inverse_mass_matrix, step_size, rng,
224                                  max_tree_depth=max_treedepth)
225         accept_prob = binary_tree.sum_accept_probs / binary_tree.num_proposals
226         num_steps = binary_tree.num_proposals
227         vv_state = vv_state.update(z=binary_tree.z_proposal,
228                                    potential_energy=binary_tree.z_proposal_pe,
229                                    z_grad=binary_tree.z_proposal_grad)
230         return vv_state, num_steps, accept_prob
231 
232     _next = _nuts_next if algo == 'NUTS' else _hmc_next
233 
234     @jit
235     def sample_kernel(hmc_state):
236         """
237         Given an existing :data:`HMCState`, run HMC with fixed (possibly adapted)
238         step size and return a new :data:`HMCState`.
239 
240         :param hmc_state: Current sample (and associated state).
241         :return: new proposed :data:`HMCState` from simulating
242             Hamiltonian dynamics given existing state.
243         """
244         rng, rng_momentum, rng_transition = random.split(hmc_state.rng, 3)
245         r = momentum_generator(hmc_state.inverse_mass_matrix, rng_momentum)
246         vv_state = IntegratorState(hmc_state.z, r, hmc_state.potential_energy, hmc_state.z_grad)
247         vv_state, num_steps, accept_prob = _next(hmc_state.step_size,
248                                                  hmc_state.inverse_mass_matrix,
249                                                  vv_state, rng_transition)
250         return HMCState(vv_state.z, vv_state.z_grad, vv_state.potential_energy, num_steps,
251                         accept_prob, hmc_state.step_size, hmc_state.inverse_mass_matrix, rng)
252 
253     # Make `init_kernel` and `sample_kernel` visible from the global scope once
254     # `hmc` is called for sphinx doc generation.
255     if 'SPHINX_BUILD' in os.environ:
256         hmc.init_kernel = init_kernel
257         hmc.sample_kernel = sample_kernel
258 
259     return init_kernel, sample_kernel
260
```
--- 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/numpyro/mcmc.py b/numpyro/mcmc.py
--- a/numpyro/mcmc.py
+++ b/numpyro/mcmc.py
@@ -1,5 +1,6 @@
 import math
 import os
+from collections import namedtuple
 
 import tqdm
 
@@ -7,13 +8,24 @@
 from jax import jit, partial, random
 from jax.flatten_util import ravel_pytree
 from jax.random import PRNGKey
+from jax.tree_util import register_pytree_node
 
 import numpyro.distributions as dist
 from numpyro.hmc_util import IntegratorState, build_tree, find_reasonable_step_size, velocity_verlet, warmup_adapter
-from numpyro.util import cond, fori_loop, laxtuple
+from numpyro.util import cond, fori_loop
 
-HMCState = laxtuple('HMCState', ['z', 'z_grad', 'potential_energy', 'num_steps', 'accept_prob',
-                                 'step_size', 'inverse_mass_matrix', 'rng'])
+HMCState = namedtuple('HMCState', ['z', 'z_grad', 'potential_energy', 'num_steps', 'accept_prob',
+                                   'step_size', 'inverse_mass_matrix', 'rng'])
+
+
+register_pytree_node(
+    HMCState,
+    lambda xs: (tuple(xs), None),
+    lambda _, xs: HMCState(*xs)
+)
+
+
+HMCState.update = HMCState._replace
 
 
 def _get_num_steps(step_size, trajectory_length):

{"golden_diff": "diff --git a/numpyro/mcmc.py b/numpyro/mcmc.py\n--- a/numpyro/mcmc.py\n+++ b/numpyro/mcmc.py\n@@ -1,5 +1,6 @@\n import math\n import os\n+from collections import namedtuple\n \n import tqdm\n \n@@ -7,13 +8,24 @@\n from jax import jit, partial, random\n from jax.flatten_util import ravel_pytree\n from jax.random import PRNGKey\n+from jax.tree_util import register_pytree_node\n \n import numpyro.distributions as dist\n from numpyro.hmc_util import IntegratorState, build_tree, find_reasonable_step_size, velocity_verlet, warmup_adapter\n-from numpyro.util import cond, fori_loop, laxtuple\n+from numpyro.util import cond, fori_loop\n \n-HMCState = laxtuple('HMCState', ['z', 'z_grad', 'potential_energy', 'num_steps', 'accept_prob',\n-                                 'step_size', 'inverse_mass_matrix', 'rng'])\n+HMCState = namedtuple('HMCState', ['z', 'z_grad', 'potential_energy', 'num_steps', 'accept_prob',\n+                                   'step_size', 'inverse_mass_matrix', 'rng'])\n+\n+\n+register_pytree_node(\n+    HMCState,\n+    lambda xs: (tuple(xs), None),\n+    lambda _, xs: HMCState(*xs)\n+)\n+\n+\n+HMCState.update = HMCState._replace\n \n \n def _get_num_steps(step_size, trajectory_length):\n", "issue": "Make HMCState picklable\nCurrently, we use `laxtuple` to make HMCState. This way however make it not possible to pickle HMCState (unless we transform it to a dict/tuple/namedtuple).\n", "before_files": [{"content": "import math\nimport os\n\nimport tqdm\n\nimport jax.numpy as np\nfrom jax import jit, partial, random\nfrom jax.flatten_util import ravel_pytree\nfrom jax.random import PRNGKey\n\nimport numpyro.distributions as dist\nfrom numpyro.hmc_util import IntegratorState, build_tree, find_reasonable_step_size, velocity_verlet, warmup_adapter\nfrom numpyro.util import cond, fori_loop, laxtuple\n\nHMCState = laxtuple('HMCState', ['z', 'z_grad', 'potential_energy', 'num_steps', 'accept_prob',\n                                 'step_size', 'inverse_mass_matrix', 'rng'])\n\n\ndef _get_num_steps(step_size, trajectory_length):\n    num_steps = np.array(trajectory_length / step_size, dtype=np.int32)\n    return np.where(num_steps < 1, np.array(1, dtype=np.int32), num_steps)\n\n\ndef _sample_momentum(unpack_fn, inverse_mass_matrix, rng):\n    if inverse_mass_matrix.ndim == 1:\n        r = dist.Normal(0., np.sqrt(np.reciprocal(inverse_mass_matrix))).sample(rng)\n        return unpack_fn(r)\n    elif inverse_mass_matrix.ndim == 2:\n        raise NotImplementedError\n\n\ndef _euclidean_ke(inverse_mass_matrix, r):\n    r, _ = ravel_pytree(r)\n\n    if inverse_mass_matrix.ndim == 2:\n        v = np.matmul(inverse_mass_matrix, r)\n    elif inverse_mass_matrix.ndim == 1:\n        v = np.multiply(inverse_mass_matrix, r)\n\n    return 0.5 * np.dot(v, r)\n\n\ndef hmc(potential_fn, kinetic_fn=None, algo='NUTS'):\n    r\"\"\"\n    Hamiltonian Monte Carlo inference, using either fixed number of\n    steps or the No U-Turn Sampler (NUTS) with adaptive path length.\n\n    **References:**\n\n    1. *MCMC Using Hamiltonian Dynamics*, Radford M. Neal\n    2. *The No-U-turn sampler: adaptively setting path lengths in Hamiltonian Monte Carlo*,\n       Matthew D. Hoffman, and Andrew Gelman.\n\n    :param potential_fn: Python callable that computes the potential energy\n        given input parameters. The input parameters to `potential_fn` can be\n        any python collection type, provided that `init_samples` argument to\n        `init_kernel` has the same type.\n    :param kinetic_fn: Python callable that returns the kinetic energy given\n        inverse mass matrix and momentum. If not provided, the default is\n        euclidean kinetic energy.\n    :param str algo: Whether to run ``HMC`` with fixed number of steps or ``NUTS``\n        with adaptive path length. Default is ``NUTS``.\n    :return init_kernel, sample_kernel: Returns a tuple of callables, the first\n        one to initialize the sampler, and the second one to generate samples\n        given an existing one.\n\n    **Example**\n\n    .. testsetup::\n\n        import jax\n        from jax import random\n        import jax.numpy as np\n        import numpyro.distributions as dist\n        from numpyro.handlers import sample\n        from numpyro.hmc_util import initialize_model\n        from numpyro.mcmc import hmc\n        from numpyro.util import fori_collect\n\n    .. doctest::\n\n        >>> true_coefs = np.array([1., 2., 3.])\n        >>> data = random.normal(random.PRNGKey(2), (2000, 3))\n        >>> dim = 3\n        >>> labels = dist.Bernoulli(logits=(true_coefs * data).sum(-1)).sample(random.PRNGKey(3))\n        >>>\n        >>> def model(data, labels):\n        ...     coefs_mean = np.zeros(dim)\n        ...     coefs = sample('beta', dist.Normal(coefs_mean, np.ones(3)))\n        ...     return sample('y', dist.Bernoulli(logits=(coefs * data).sum(-1)), obs=labels)\n        >>>\n        >>> init_params, potential_fn, constrain_fn = initialize_model(random.PRNGKey(0), model, data, labels)\n        >>> init_kernel, sample_kernel = hmc(potential_fn, algo='NUTS')\n        >>> hmc_state = init_kernel(init_params,\n        ...                         trajectory_length=10,\n        ...                         num_warmup_steps=300)\n        >>> hmc_states = fori_collect(500, sample_kernel, hmc_state,\n        ...                           transform=lambda x: constrain_fn(x.z))\n        >>> print(np.mean(hmc_states['beta'], axis=0))\n        [0.9153987 2.0754058 2.9621222]\n    \"\"\"\n    if kinetic_fn is None:\n        kinetic_fn = _euclidean_ke\n    vv_init, vv_update = velocity_verlet(potential_fn, kinetic_fn)\n    trajectory_len = None\n    max_treedepth = None\n    momentum_generator = None\n    wa_update = None\n\n    def init_kernel(init_samples,\n                    num_warmup,\n                    step_size=1.0,\n                    adapt_step_size=True,\n                    adapt_mass_matrix=True,\n                    diag_mass=True,\n                    target_accept_prob=0.8,\n                    trajectory_length=2*math.pi,\n                    max_tree_depth=10,\n                    run_warmup=True,\n                    progbar=True,\n                    rng=PRNGKey(0)):\n        \"\"\"\n        Initializes the HMC sampler.\n\n        :param init_samples: Initial parameters to begin sampling. The type can\n            must be consistent with the input type to `potential_fn`.\n        :param int num_warmup_steps: Number of warmup steps; samples generated\n            during warmup are discarded.\n        :param float step_size: Determines the size of a single step taken by the\n            verlet integrator while computing the trajectory using Hamiltonian\n            dynamics. If not specified, it will be set to 1.\n        :param bool adapt_step_size: A flag to decide if we want to adapt step_size\n            during warm-up phase using Dual Averaging scheme.\n        :param bool adapt_mass_matrix: A flag to decide if we want to adapt mass\n            matrix during warm-up phase using Welford scheme.\n        :param bool diag_mass: A flag to decide if mass matrix is diagonal (default)\n            or dense (if set to ``False``).\n        :param float target_accept_prob: Target acceptance probability for step size\n            adaptation using Dual Averaging. Increasing this value will lead to a smaller\n            step size, hence the sampling will be slower but more robust. Default to 0.8.\n        :param float trajectory_length: Length of a MCMC trajectory for HMC. Default\n            value is :math:`2\\\\pi`.\n        :param int max_tree_depth: Max depth of the binary tree created during the doubling\n            scheme of NUTS sampler. Defaults to 10.\n        :param bool run_warmup: Flag to decide whether warmup is run. If ``True``,\n            `init_kernel` returns an initial :data:`HMCState` that can be used to\n            generate samples using MCMC. Else, returns the arguments and callable\n            that does the initial adaptation.\n        :param bool progbar: Whether to enable progress bar updates. Defaults to\n            ``True``.\n        :param bool heuristic_step_size: If ``True``, a coarse grained adjustment of\n            step size is done at the beginning of each adaptation window to achieve\n            `target_acceptance_prob`.\n        :param jax.random.PRNGKey rng: random key to be used as the source of\n            randomness.\n        \"\"\"\n        step_size = float(step_size)\n        nonlocal momentum_generator, wa_update, trajectory_len, max_treedepth\n        trajectory_len = float(trajectory_length)\n        max_treedepth = max_tree_depth\n        z = init_samples\n        z_flat, unravel_fn = ravel_pytree(z)\n        momentum_generator = partial(_sample_momentum, unravel_fn)\n\n        find_reasonable_ss = partial(find_reasonable_step_size,\n                                     potential_fn, kinetic_fn,\n                                     momentum_generator)\n\n        wa_init, wa_update = warmup_adapter(num_warmup,\n                                            adapt_step_size=adapt_step_size,\n                                            adapt_mass_matrix=adapt_mass_matrix,\n                                            diag_mass=diag_mass,\n                                            target_accept_prob=target_accept_prob,\n                                            find_reasonable_step_size=find_reasonable_ss)\n\n        rng_hmc, rng_wa = random.split(rng)\n        wa_state = wa_init(z, rng_wa, step_size, mass_matrix_size=np.size(z_flat))\n        r = momentum_generator(wa_state.inverse_mass_matrix, rng)\n        vv_state = vv_init(z, r)\n        hmc_state = HMCState(vv_state.z, vv_state.z_grad, vv_state.potential_energy, 0, 0.,\n                             wa_state.step_size, wa_state.inverse_mass_matrix, rng_hmc)\n\n        wa_update = jit(wa_update)\n        if run_warmup:\n            # JIT if progress bar updates not required\n            if not progbar:\n                hmc_state, _ = jit(fori_loop, static_argnums=(2,))(0, num_warmup,\n                                                                   warmup_update,\n                                                                   (hmc_state, wa_state))\n            else:\n                for i in tqdm.trange(num_warmup):\n                    hmc_state, wa_state = warmup_update(i, (hmc_state, wa_state))\n            return hmc_state\n        else:\n            return hmc_state, wa_state, warmup_update\n\n    def warmup_update(t, states):\n        hmc_state, wa_state = states\n        hmc_state = sample_kernel(hmc_state)\n        wa_state = wa_update(t, hmc_state.accept_prob, hmc_state.z, wa_state)\n        hmc_state = hmc_state.update(step_size=wa_state.step_size,\n                                     inverse_mass_matrix=wa_state.inverse_mass_matrix)\n        return hmc_state, wa_state\n\n    def _hmc_next(step_size, inverse_mass_matrix, vv_state, rng):\n        num_steps = _get_num_steps(step_size, trajectory_len)\n        vv_state_new = fori_loop(0, num_steps,\n                                 lambda i, val: vv_update(step_size, inverse_mass_matrix, val),\n                                 vv_state)\n        energy_old = vv_state.potential_energy + kinetic_fn(inverse_mass_matrix, vv_state.r)\n        energy_new = vv_state_new.potential_energy + kinetic_fn(inverse_mass_matrix, vv_state_new.r)\n        delta_energy = energy_new - energy_old\n        delta_energy = np.where(np.isnan(delta_energy), np.inf, delta_energy)\n        accept_prob = np.clip(np.exp(-delta_energy), a_max=1.0)\n        transition = random.bernoulli(rng, accept_prob)\n        vv_state = cond(transition,\n                        vv_state_new, lambda state: state,\n                        vv_state, lambda state: state)\n        return vv_state, num_steps, accept_prob\n\n    def _nuts_next(step_size, inverse_mass_matrix, vv_state, rng):\n        binary_tree = build_tree(vv_update, kinetic_fn, vv_state,\n                                 inverse_mass_matrix, step_size, rng,\n                                 max_tree_depth=max_treedepth)\n        accept_prob = binary_tree.sum_accept_probs / binary_tree.num_proposals\n        num_steps = binary_tree.num_proposals\n        vv_state = vv_state.update(z=binary_tree.z_proposal,\n                                   potential_energy=binary_tree.z_proposal_pe,\n                                   z_grad=binary_tree.z_proposal_grad)\n        return vv_state, num_steps, accept_prob\n\n    _next = _nuts_next if algo == 'NUTS' else _hmc_next\n\n    @jit\n    def sample_kernel(hmc_state):\n        \"\"\"\n        Given an existing :data:`HMCState`, run HMC with fixed (possibly adapted)\n        step size and return a new :data:`HMCState`.\n\n        :param hmc_state: Current sample (and associated state).\n        :return: new proposed :data:`HMCState` from simulating\n            Hamiltonian dynamics given existing state.\n        \"\"\"\n        rng, rng_momentum, rng_transition = random.split(hmc_state.rng, 3)\n        r = momentum_generator(hmc_state.inverse_mass_matrix, rng_momentum)\n        vv_state = IntegratorState(hmc_state.z, r, hmc_state.potential_energy, hmc_state.z_grad)\n        vv_state, num_steps, accept_prob = _next(hmc_state.step_size,\n                                                 hmc_state.inverse_mass_matrix,\n                                                 vv_state, rng_transition)\n        return HMCState(vv_state.z, vv_state.z_grad, vv_state.potential_energy, num_steps,\n                        accept_prob, hmc_state.step_size, hmc_state.inverse_mass_matrix, rng)\n\n    # Make `init_kernel` and `sample_kernel` visible from the global scope once\n    # `hmc` is called for sphinx doc generation.\n    if 'SPHINX_BUILD' in os.environ:\n        hmc.init_kernel = init_kernel\n        hmc.sample_kernel = sample_kernel\n\n    return init_kernel, sample_kernel\n", "path": "numpyro/mcmc.py"}], "after_files": [{"content": "import math\nimport os\nfrom collections import namedtuple\n\nimport tqdm\n\nimport jax.numpy as np\nfrom jax import jit, partial, random\nfrom jax.flatten_util import ravel_pytree\nfrom jax.random import PRNGKey\nfrom jax.tree_util import register_pytree_node\n\nimport numpyro.distributions as dist\nfrom numpyro.hmc_util import IntegratorState, build_tree, find_reasonable_step_size, velocity_verlet, warmup_adapter\nfrom numpyro.util import cond, fori_loop\n\nHMCState = namedtuple('HMCState', ['z', 'z_grad', 'potential_energy', 'num_steps', 'accept_prob',\n                                   'step_size', 'inverse_mass_matrix', 'rng'])\n\n\nregister_pytree_node(\n    HMCState,\n    lambda xs: (tuple(xs), None),\n    lambda _, xs: HMCState(*xs)\n)\n\n\nHMCState.update = HMCState._replace\n\n\ndef _get_num_steps(step_size, trajectory_length):\n    num_steps = np.array(trajectory_length / step_size, dtype=np.int32)\n    return np.where(num_steps < 1, np.array(1, dtype=np.int32), num_steps)\n\n\ndef _sample_momentum(unpack_fn, inverse_mass_matrix, rng):\n    if inverse_mass_matrix.ndim == 1:\n        r = dist.Normal(0., np.sqrt(np.reciprocal(inverse_mass_matrix))).sample(rng)\n        return unpack_fn(r)\n    elif inverse_mass_matrix.ndim == 2:\n        raise NotImplementedError\n\n\ndef _euclidean_ke(inverse_mass_matrix, r):\n    r, _ = ravel_pytree(r)\n\n    if inverse_mass_matrix.ndim == 2:\n        v = np.matmul(inverse_mass_matrix, r)\n    elif inverse_mass_matrix.ndim == 1:\n        v = np.multiply(inverse_mass_matrix, r)\n\n    return 0.5 * np.dot(v, r)\n\n\ndef hmc(potential_fn, kinetic_fn=None, algo='NUTS'):\n    r\"\"\"\n    Hamiltonian Monte Carlo inference, using either fixed number of\n    steps or the No U-Turn Sampler (NUTS) with adaptive path length.\n\n    **References:**\n\n    1. *MCMC Using Hamiltonian Dynamics*, Radford M. Neal\n    2. *The No-U-turn sampler: adaptively setting path lengths in Hamiltonian Monte Carlo*,\n       Matthew D. Hoffman, and Andrew Gelman.\n\n    :param potential_fn: Python callable that computes the potential energy\n        given input parameters. The input parameters to `potential_fn` can be\n        any python collection type, provided that `init_samples` argument to\n        `init_kernel` has the same type.\n    :param kinetic_fn: Python callable that returns the kinetic energy given\n        inverse mass matrix and momentum. If not provided, the default is\n        euclidean kinetic energy.\n    :param str algo: Whether to run ``HMC`` with fixed number of steps or ``NUTS``\n        with adaptive path length. Default is ``NUTS``.\n    :return init_kernel, sample_kernel: Returns a tuple of callables, the first\n        one to initialize the sampler, and the second one to generate samples\n        given an existing one.\n\n    **Example**\n\n    .. testsetup::\n\n        import jax\n        from jax import random\n        import jax.numpy as np\n        import numpyro.distributions as dist\n        from numpyro.handlers import sample\n        from numpyro.hmc_util import initialize_model\n        from numpyro.mcmc import hmc\n        from numpyro.util import fori_collect\n\n    .. doctest::\n\n        >>> true_coefs = np.array([1., 2., 3.])\n        >>> data = random.normal(random.PRNGKey(2), (2000, 3))\n        >>> dim = 3\n        >>> labels = dist.Bernoulli(logits=(true_coefs * data).sum(-1)).sample(random.PRNGKey(3))\n        >>>\n        >>> def model(data, labels):\n        ...     coefs_mean = np.zeros(dim)\n        ...     coefs = sample('beta', dist.Normal(coefs_mean, np.ones(3)))\n        ...     return sample('y', dist.Bernoulli(logits=(coefs * data).sum(-1)), obs=labels)\n        >>>\n        >>> init_params, potential_fn, constrain_fn = initialize_model(random.PRNGKey(0), model, data, labels)\n        >>> init_kernel, sample_kernel = hmc(potential_fn, algo='NUTS')\n        >>> hmc_state = init_kernel(init_params,\n        ...                         trajectory_length=10,\n        ...                         num_warmup_steps=300)\n        >>> hmc_states = fori_collect(500, sample_kernel, hmc_state,\n        ...                           transform=lambda x: constrain_fn(x.z))\n        >>> print(np.mean(hmc_states['beta'], axis=0))\n        [0.9153987 2.0754058 2.9621222]\n    \"\"\"\n    if kinetic_fn is None:\n        kinetic_fn = _euclidean_ke\n    vv_init, vv_update = velocity_verlet(potential_fn, kinetic_fn)\n    trajectory_len = None\n    max_treedepth = None\n    momentum_generator = None\n    wa_update = None\n\n    def init_kernel(init_samples,\n                    num_warmup,\n                    step_size=1.0,\n                    adapt_step_size=True,\n                    adapt_mass_matrix=True,\n                    diag_mass=True,\n                    target_accept_prob=0.8,\n                    trajectory_length=2*math.pi,\n                    max_tree_depth=10,\n                    run_warmup=True,\n                    progbar=True,\n                    rng=PRNGKey(0)):\n        \"\"\"\n        Initializes the HMC sampler.\n\n        :param init_samples: Initial parameters to begin sampling. The type can\n            must be consistent with the input type to `potential_fn`.\n        :param int num_warmup_steps: Number of warmup steps; samples generated\n            during warmup are discarded.\n        :param float step_size: Determines the size of a single step taken by the\n            verlet integrator while computing the trajectory using Hamiltonian\n            dynamics. If not specified, it will be set to 1.\n        :param bool adapt_step_size: A flag to decide if we want to adapt step_size\n            during warm-up phase using Dual Averaging scheme.\n        :param bool adapt_mass_matrix: A flag to decide if we want to adapt mass\n            matrix during warm-up phase using Welford scheme.\n        :param bool diag_mass: A flag to decide if mass matrix is diagonal (default)\n            or dense (if set to ``False``).\n        :param float target_accept_prob: Target acceptance probability for step size\n            adaptation using Dual Averaging. Increasing this value will lead to a smaller\n            step size, hence the sampling will be slower but more robust. Default to 0.8.\n        :param float trajectory_length: Length of a MCMC trajectory for HMC. Default\n            value is :math:`2\\\\pi`.\n        :param int max_tree_depth: Max depth of the binary tree created during the doubling\n            scheme of NUTS sampler. Defaults to 10.\n        :param bool run_warmup: Flag to decide whether warmup is run. If ``True``,\n            `init_kernel` returns an initial :data:`HMCState` that can be used to\n            generate samples using MCMC. Else, returns the arguments and callable\n            that does the initial adaptation.\n        :param bool progbar: Whether to enable progress bar updates. Defaults to\n            ``True``.\n        :param bool heuristic_step_size: If ``True``, a coarse grained adjustment of\n            step size is done at the beginning of each adaptation window to achieve\n            `target_acceptance_prob`.\n        :param jax.random.PRNGKey rng: random key to be used as the source of\n            randomness.\n        \"\"\"\n        step_size = float(step_size)\n        nonlocal momentum_generator, wa_update, trajectory_len, max_treedepth\n        trajectory_len = float(trajectory_length)\n        max_treedepth = max_tree_depth\n        z = init_samples\n        z_flat, unravel_fn = ravel_pytree(z)\n        momentum_generator = partial(_sample_momentum, unravel_fn)\n\n        find_reasonable_ss = partial(find_reasonable_step_size,\n                                     potential_fn, kinetic_fn,\n                                     momentum_generator)\n\n        wa_init, wa_update = warmup_adapter(num_warmup,\n                                            adapt_step_size=adapt_step_size,\n                                            adapt_mass_matrix=adapt_mass_matrix,\n                                            diag_mass=diag_mass,\n                                            target_accept_prob=target_accept_prob,\n                                            find_reasonable_step_size=find_reasonable_ss)\n\n        rng_hmc, rng_wa = random.split(rng)\n        wa_state = wa_init(z, rng_wa, step_size, mass_matrix_size=np.size(z_flat))\n        r = momentum_generator(wa_state.inverse_mass_matrix, rng)\n        vv_state = vv_init(z, r)\n        hmc_state = HMCState(vv_state.z, vv_state.z_grad, vv_state.potential_energy, 0, 0.,\n                             wa_state.step_size, wa_state.inverse_mass_matrix, rng_hmc)\n\n        wa_update = jit(wa_update)\n        if run_warmup:\n            # JIT if progress bar updates not required\n            if not progbar:\n                hmc_state, _ = jit(fori_loop, static_argnums=(2,))(0, num_warmup,\n                                                                   warmup_update,\n                                                                   (hmc_state, wa_state))\n            else:\n                for i in tqdm.trange(num_warmup):\n                    hmc_state, wa_state = warmup_update(i, (hmc_state, wa_state))\n            return hmc_state\n        else:\n            return hmc_state, wa_state, warmup_update\n\n    def warmup_update(t, states):\n        hmc_state, wa_state = states\n        hmc_state = sample_kernel(hmc_state)\n        wa_state = wa_update(t, hmc_state.accept_prob, hmc_state.z, wa_state)\n        hmc_state = hmc_state.update(step_size=wa_state.step_size,\n                                     inverse_mass_matrix=wa_state.inverse_mass_matrix)\n        return hmc_state, wa_state\n\n    def _hmc_next(step_size, inverse_mass_matrix, vv_state, rng):\n        num_steps = _get_num_steps(step_size, trajectory_len)\n        vv_state_new = fori_loop(0, num_steps,\n                                 lambda i, val: vv_update(step_size, inverse_mass_matrix, val),\n                                 vv_state)\n        energy_old = vv_state.potential_energy + kinetic_fn(inverse_mass_matrix, vv_state.r)\n        energy_new = vv_state_new.potential_energy + kinetic_fn(inverse_mass_matrix, vv_state_new.r)\n        delta_energy = energy_new - energy_old\n        delta_energy = np.where(np.isnan(delta_energy), np.inf, delta_energy)\n        accept_prob = np.clip(np.exp(-delta_energy), a_max=1.0)\n        transition = random.bernoulli(rng, accept_prob)\n        vv_state = cond(transition,\n                        vv_state_new, lambda state: state,\n                        vv_state, lambda state: state)\n        return vv_state, num_steps, accept_prob\n\n    def _nuts_next(step_size, inverse_mass_matrix, vv_state, rng):\n        binary_tree = build_tree(vv_update, kinetic_fn, vv_state,\n                                 inverse_mass_matrix, step_size, rng,\n                                 max_tree_depth=max_treedepth)\n        accept_prob = binary_tree.sum_accept_probs / binary_tree.num_proposals\n        num_steps = binary_tree.num_proposals\n        vv_state = vv_state.update(z=binary_tree.z_proposal,\n                                   potential_energy=binary_tree.z_proposal_pe,\n                                   z_grad=binary_tree.z_proposal_grad)\n        return vv_state, num_steps, accept_prob\n\n    _next = _nuts_next if algo == 'NUTS' else _hmc_next\n\n    @jit\n    def sample_kernel(hmc_state):\n        \"\"\"\n        Given an existing :data:`HMCState`, run HMC with fixed (possibly adapted)\n        step size and return a new :data:`HMCState`.\n\n        :param hmc_state: Current sample (and associated state).\n        :return: new proposed :data:`HMCState` from simulating\n            Hamiltonian dynamics given existing state.\n        \"\"\"\n        rng, rng_momentum, rng_transition = random.split(hmc_state.rng, 3)\n        r = momentum_generator(hmc_state.inverse_mass_matrix, rng_momentum)\n        vv_state = IntegratorState(hmc_state.z, r, hmc_state.potential_energy, hmc_state.z_grad)\n        vv_state, num_steps, accept_prob = _next(hmc_state.step_size,\n                                                 hmc_state.inverse_mass_matrix,\n                                                 vv_state, rng_transition)\n        return HMCState(vv_state.z, vv_state.z_grad, vv_state.potential_energy, num_steps,\n                        accept_prob, hmc_state.step_size, hmc_state.inverse_mass_matrix, rng)\n\n    # Make `init_kernel` and `sample_kernel` visible from the global scope once\n    # `hmc` is called for sphinx doc generation.\n    if 'SPHINX_BUILD' in os.environ:\n        hmc.init_kernel = init_kernel\n        hmc.sample_kernel = sample_kernel\n\n    return init_kernel, sample_kernel\n", "path": "numpyro/mcmc.py"}]}