We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
fori_collect does not work with non-float arrays
The issue is ravel_pytree which calls `vjp` and does not work with int32 or bool arrays. As such, `fori_collect` will fail if the default identity transform is used.

--- 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/util.py`
Content:
```
1 from contextlib import contextmanager
2 import random
3 
4 import numpy as onp
5 import tqdm
6 
7 from jax import jit, lax, ops, vmap
8 from jax.flatten_util import ravel_pytree
9 import jax.numpy as np
10 
11 _DATA_TYPES = {}
12 _DISABLE_CONTROL_FLOW_PRIM = False
13 
14 
15 def set_rng_seed(rng_seed):
16     random.seed(rng_seed)
17     onp.random.seed(rng_seed)
18 
19 
20 @contextmanager
21 def optional(condition, context_manager):
22     """
23     Optionally wrap inside `context_manager` if condition is `True`.
24     """
25     if condition:
26         with context_manager:
27             yield
28     else:
29         yield
30 
31 
32 @contextmanager
33 def control_flow_prims_disabled():
34     global _DISABLE_CONTROL_FLOW_PRIM
35     stored_flag = _DISABLE_CONTROL_FLOW_PRIM
36     try:
37         _DISABLE_CONTROL_FLOW_PRIM = True
38         yield
39     finally:
40         _DISABLE_CONTROL_FLOW_PRIM = stored_flag
41 
42 
43 def cond(pred, true_operand, true_fun, false_operand, false_fun):
44     if _DISABLE_CONTROL_FLOW_PRIM:
45         if pred:
46             return true_fun(true_operand)
47         else:
48             return false_fun(false_operand)
49     else:
50         return lax.cond(pred, true_operand, true_fun, false_operand, false_fun)
51 
52 
53 def while_loop(cond_fun, body_fun, init_val):
54     if _DISABLE_CONTROL_FLOW_PRIM:
55         val = init_val
56         while cond_fun(val):
57             val = body_fun(val)
58         return val
59     else:
60         # TODO: consider jitting while_loop similar to fori_loop
61         return lax.while_loop(cond_fun, body_fun, init_val)
62 
63 
64 def fori_loop(lower, upper, body_fun, init_val):
65     if _DISABLE_CONTROL_FLOW_PRIM:
66         val = init_val
67         for i in range(int(lower), int(upper)):
68             val = body_fun(i, val)
69         return val
70     else:
71         return jit(lax.fori_loop, static_argnums=(2,))(lower, upper, body_fun, init_val)
72 
73 
74 def identity(x):
75     return x
76 
77 
78 def fori_collect(lower, upper, body_fun, init_val, transform=identity, progbar=True, **progbar_opts):
79     """
80     This looping construct works like :func:`~jax.lax.fori_loop` but with the additional
81     effect of collecting values from the loop body. In addition, this allows for
82     post-processing of these samples via `transform`, and progress bar updates.
83     Note that, `progbar=False` will be faster, especially when collecting a
84     lot of samples. Refer to example usage in :func:`~numpyro.mcmc.hmc`.
85 
86     :param int lower: the index to start the collective work. In other words,
87         we will skip collecting the first `lower` values.
88     :param int upper: number of times to run the loop body.
89     :param body_fun: a callable that takes a collection of
90         `np.ndarray` and returns a collection with the same shape and
91         `dtype`.
92     :param init_val: initial value to pass as argument to `body_fun`. Can
93         be any Python collection type containing `np.ndarray` objects.
94     :param transform: a callable to post-process the values returned by `body_fn`.
95     :param progbar: whether to post progress bar updates.
96     :param `**progbar_opts`: optional additional progress bar arguments. A
97         `diagnostics_fn` can be supplied which when passed the current value
98         from `body_fun` returns a string that is used to update the progress
99         bar postfix. Also a `progbar_desc` keyword argument can be supplied
100         which is used to label the progress bar.
101     :return: collection with the same type as `init_val` with values
102         collected along the leading axis of `np.ndarray` objects.
103     """
104     assert lower < upper
105     init_val_flat, unravel_fn = ravel_pytree(transform(init_val))
106     ravel_fn = lambda x: ravel_pytree(transform(x))[0]  # noqa: E731
107 
108     if not progbar:
109         collection = np.zeros((upper - lower,) + init_val_flat.shape)
110 
111         def _body_fn(i, vals):
112             val, collection = vals
113             val = body_fun(val)
114             i = np.where(i >= lower, i - lower, 0)
115             collection = ops.index_update(collection, i, ravel_fn(val))
116             return val, collection
117 
118         _, collection = fori_loop(0, upper, _body_fn, (init_val, collection))
119     else:
120         diagnostics_fn = progbar_opts.pop('diagnostics_fn', None)
121         progbar_desc = progbar_opts.pop('progbar_desc', '')
122         collection = []
123 
124         val = init_val
125         with tqdm.trange(upper, desc=progbar_desc) as t:
126             for i in t:
127                 val = body_fun(val)
128                 if i >= lower:
129                     collection.append(jit(ravel_fn)(val))
130                 if diagnostics_fn:
131                     t.set_postfix_str(diagnostics_fn(val), refresh=False)
132 
133         collection = np.stack(collection)
134 
135     return vmap(unravel_fn)(collection)
136 
137 
138 def copy_docs_from(source_class, full_text=False):
139     """
140     Decorator to copy class and method docs from source to destin class.
141     """
142 
143     def decorator(destin_class):
144         # This works only in python 3.3+:
145         # if not destin_class.__doc__:
146         #     destin_class.__doc__ = source_class.__doc__
147         for name in dir(destin_class):
148             if name.startswith('_'):
149                 continue
150             destin_attr = getattr(destin_class, name)
151             destin_attr = getattr(destin_attr, '__func__', destin_attr)
152             source_attr = getattr(source_class, name, None)
153             source_doc = getattr(source_attr, '__doc__', None)
154             if source_doc and not getattr(destin_attr, '__doc__', None):
155                 if full_text or source_doc.startswith('See '):
156                     destin_doc = source_doc
157                 else:
158                     destin_doc = 'See :meth:`{}.{}.{}`'.format(
159                         source_class.__module__, source_class.__name__, name)
160                 if isinstance(destin_attr, property):
161                     # Set docs for object properties.
162                     # Since __doc__ is read-only, we need to reset the property
163                     # with the updated doc.
164                     updated_property = property(destin_attr.fget,
165                                                 destin_attr.fset,
166                                                 destin_attr.fdel,
167                                                 destin_doc)
168                     setattr(destin_class, name, updated_property)
169                 else:
170                     destin_attr.__doc__ = destin_doc
171         return destin_class
172 
173     return decorator
174
```
--- 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/util.py b/numpyro/util.py
--- a/numpyro/util.py
+++ b/numpyro/util.py
@@ -1,3 +1,4 @@
+from collections import namedtuple
 from contextlib import contextmanager
 import random
 
@@ -5,8 +6,8 @@
 import tqdm
 
 from jax import jit, lax, ops, vmap
-from jax.flatten_util import ravel_pytree
 import jax.numpy as np
+from jax.tree_util import tree_flatten, tree_unflatten, tree_map
 
 _DATA_TYPES = {}
 _DISABLE_CONTROL_FLOW_PRIM = False
@@ -171,3 +172,29 @@
         return destin_class
 
     return decorator
+
+
+pytree_metadata = namedtuple('pytree_metadata', ['flat', 'shape', 'size', 'dtype'])
+
+
+def _ravel_list(*leaves):
+    leaves_metadata = tree_map(lambda l: pytree_metadata(np.ravel(l), np.shape(l), np.size(l), lax.dtype(l)),
+                               leaves)
+    leaves_idx = np.cumsum(np.array((0,) + tuple(d.size for d in leaves_metadata)))
+
+    def unravel_list(arr):
+        return [np.reshape(lax.dynamic_slice_in_dim(arr, leaves_idx[i], m.size),
+                           m.shape).astype(m.dtype)
+                for i, m in enumerate(leaves_metadata)]
+
+    return np.concatenate([m.flat for m in leaves_metadata]), unravel_list
+
+
+def ravel_pytree(pytree):
+    leaves, treedef = tree_flatten(pytree)
+    flat, unravel_list = _ravel_list(*leaves)
+
+    def unravel_pytree(arr):
+        return tree_unflatten(treedef, unravel_list(arr))
+
+    return flat, unravel_pytree

{"golden_diff": "diff --git a/numpyro/util.py b/numpyro/util.py\n--- a/numpyro/util.py\n+++ b/numpyro/util.py\n@@ -1,3 +1,4 @@\n+from collections import namedtuple\n from contextlib import contextmanager\n import random\n \n@@ -5,8 +6,8 @@\n import tqdm\n \n from jax import jit, lax, ops, vmap\n-from jax.flatten_util import ravel_pytree\n import jax.numpy as np\n+from jax.tree_util import tree_flatten, tree_unflatten, tree_map\n \n _DATA_TYPES = {}\n _DISABLE_CONTROL_FLOW_PRIM = False\n@@ -171,3 +172,29 @@\n         return destin_class\n \n     return decorator\n+\n+\n+pytree_metadata = namedtuple('pytree_metadata', ['flat', 'shape', 'size', 'dtype'])\n+\n+\n+def _ravel_list(*leaves):\n+    leaves_metadata = tree_map(lambda l: pytree_metadata(np.ravel(l), np.shape(l), np.size(l), lax.dtype(l)),\n+                               leaves)\n+    leaves_idx = np.cumsum(np.array((0,) + tuple(d.size for d in leaves_metadata)))\n+\n+    def unravel_list(arr):\n+        return [np.reshape(lax.dynamic_slice_in_dim(arr, leaves_idx[i], m.size),\n+                           m.shape).astype(m.dtype)\n+                for i, m in enumerate(leaves_metadata)]\n+\n+    return np.concatenate([m.flat for m in leaves_metadata]), unravel_list\n+\n+\n+def ravel_pytree(pytree):\n+    leaves, treedef = tree_flatten(pytree)\n+    flat, unravel_list = _ravel_list(*leaves)\n+\n+    def unravel_pytree(arr):\n+        return tree_unflatten(treedef, unravel_list(arr))\n+\n+    return flat, unravel_pytree\n", "issue": "fori_collect does not work with non-float arrays\nThe issue is ravel_pytree which calls `vjp` and does not work with int32 or bool arrays. As such, `fori_collect` will fail if the default identity transform is used.\n", "before_files": [{"content": "from contextlib import contextmanager\nimport random\n\nimport numpy as onp\nimport tqdm\n\nfrom jax import jit, lax, ops, vmap\nfrom jax.flatten_util import ravel_pytree\nimport jax.numpy as np\n\n_DATA_TYPES = {}\n_DISABLE_CONTROL_FLOW_PRIM = False\n\n\ndef set_rng_seed(rng_seed):\n    random.seed(rng_seed)\n    onp.random.seed(rng_seed)\n\n\n@contextmanager\ndef optional(condition, context_manager):\n    \"\"\"\n    Optionally wrap inside `context_manager` if condition is `True`.\n    \"\"\"\n    if condition:\n        with context_manager:\n            yield\n    else:\n        yield\n\n\n@contextmanager\ndef control_flow_prims_disabled():\n    global _DISABLE_CONTROL_FLOW_PRIM\n    stored_flag = _DISABLE_CONTROL_FLOW_PRIM\n    try:\n        _DISABLE_CONTROL_FLOW_PRIM = True\n        yield\n    finally:\n        _DISABLE_CONTROL_FLOW_PRIM = stored_flag\n\n\ndef cond(pred, true_operand, true_fun, false_operand, false_fun):\n    if _DISABLE_CONTROL_FLOW_PRIM:\n        if pred:\n            return true_fun(true_operand)\n        else:\n            return false_fun(false_operand)\n    else:\n        return lax.cond(pred, true_operand, true_fun, false_operand, false_fun)\n\n\ndef while_loop(cond_fun, body_fun, init_val):\n    if _DISABLE_CONTROL_FLOW_PRIM:\n        val = init_val\n        while cond_fun(val):\n            val = body_fun(val)\n        return val\n    else:\n        # TODO: consider jitting while_loop similar to fori_loop\n        return lax.while_loop(cond_fun, body_fun, init_val)\n\n\ndef fori_loop(lower, upper, body_fun, init_val):\n    if _DISABLE_CONTROL_FLOW_PRIM:\n        val = init_val\n        for i in range(int(lower), int(upper)):\n            val = body_fun(i, val)\n        return val\n    else:\n        return jit(lax.fori_loop, static_argnums=(2,))(lower, upper, body_fun, init_val)\n\n\ndef identity(x):\n    return x\n\n\ndef fori_collect(lower, upper, body_fun, init_val, transform=identity, progbar=True, **progbar_opts):\n    \"\"\"\n    This looping construct works like :func:`~jax.lax.fori_loop` but with the additional\n    effect of collecting values from the loop body. In addition, this allows for\n    post-processing of these samples via `transform`, and progress bar updates.\n    Note that, `progbar=False` will be faster, especially when collecting a\n    lot of samples. Refer to example usage in :func:`~numpyro.mcmc.hmc`.\n\n    :param int lower: the index to start the collective work. In other words,\n        we will skip collecting the first `lower` values.\n    :param int upper: number of times to run the loop body.\n    :param body_fun: a callable that takes a collection of\n        `np.ndarray` and returns a collection with the same shape and\n        `dtype`.\n    :param init_val: initial value to pass as argument to `body_fun`. Can\n        be any Python collection type containing `np.ndarray` objects.\n    :param transform: a callable to post-process the values returned by `body_fn`.\n    :param progbar: whether to post progress bar updates.\n    :param `**progbar_opts`: optional additional progress bar arguments. A\n        `diagnostics_fn` can be supplied which when passed the current value\n        from `body_fun` returns a string that is used to update the progress\n        bar postfix. Also a `progbar_desc` keyword argument can be supplied\n        which is used to label the progress bar.\n    :return: collection with the same type as `init_val` with values\n        collected along the leading axis of `np.ndarray` objects.\n    \"\"\"\n    assert lower < upper\n    init_val_flat, unravel_fn = ravel_pytree(transform(init_val))\n    ravel_fn = lambda x: ravel_pytree(transform(x))[0]  # noqa: E731\n\n    if not progbar:\n        collection = np.zeros((upper - lower,) + init_val_flat.shape)\n\n        def _body_fn(i, vals):\n            val, collection = vals\n            val = body_fun(val)\n            i = np.where(i >= lower, i - lower, 0)\n            collection = ops.index_update(collection, i, ravel_fn(val))\n            return val, collection\n\n        _, collection = fori_loop(0, upper, _body_fn, (init_val, collection))\n    else:\n        diagnostics_fn = progbar_opts.pop('diagnostics_fn', None)\n        progbar_desc = progbar_opts.pop('progbar_desc', '')\n        collection = []\n\n        val = init_val\n        with tqdm.trange(upper, desc=progbar_desc) as t:\n            for i in t:\n                val = body_fun(val)\n                if i >= lower:\n                    collection.append(jit(ravel_fn)(val))\n                if diagnostics_fn:\n                    t.set_postfix_str(diagnostics_fn(val), refresh=False)\n\n        collection = np.stack(collection)\n\n    return vmap(unravel_fn)(collection)\n\n\ndef copy_docs_from(source_class, full_text=False):\n    \"\"\"\n    Decorator to copy class and method docs from source to destin class.\n    \"\"\"\n\n    def decorator(destin_class):\n        # This works only in python 3.3+:\n        # if not destin_class.__doc__:\n        #     destin_class.__doc__ = source_class.__doc__\n        for name in dir(destin_class):\n            if name.startswith('_'):\n                continue\n            destin_attr = getattr(destin_class, name)\n            destin_attr = getattr(destin_attr, '__func__', destin_attr)\n            source_attr = getattr(source_class, name, None)\n            source_doc = getattr(source_attr, '__doc__', None)\n            if source_doc and not getattr(destin_attr, '__doc__', None):\n                if full_text or source_doc.startswith('See '):\n                    destin_doc = source_doc\n                else:\n                    destin_doc = 'See :meth:`{}.{}.{}`'.format(\n                        source_class.__module__, source_class.__name__, name)\n                if isinstance(destin_attr, property):\n                    # Set docs for object properties.\n                    # Since __doc__ is read-only, we need to reset the property\n                    # with the updated doc.\n                    updated_property = property(destin_attr.fget,\n                                                destin_attr.fset,\n                                                destin_attr.fdel,\n                                                destin_doc)\n                    setattr(destin_class, name, updated_property)\n                else:\n                    destin_attr.__doc__ = destin_doc\n        return destin_class\n\n    return decorator\n", "path": "numpyro/util.py"}], "after_files": [{"content": "from collections import namedtuple\nfrom contextlib import contextmanager\nimport random\n\nimport numpy as onp\nimport tqdm\n\nfrom jax import jit, lax, ops, vmap\nimport jax.numpy as np\nfrom jax.tree_util import tree_flatten, tree_unflatten, tree_map\n\n_DATA_TYPES = {}\n_DISABLE_CONTROL_FLOW_PRIM = False\n\n\ndef set_rng_seed(rng_seed):\n    random.seed(rng_seed)\n    onp.random.seed(rng_seed)\n\n\n@contextmanager\ndef optional(condition, context_manager):\n    \"\"\"\n    Optionally wrap inside `context_manager` if condition is `True`.\n    \"\"\"\n    if condition:\n        with context_manager:\n            yield\n    else:\n        yield\n\n\n@contextmanager\ndef control_flow_prims_disabled():\n    global _DISABLE_CONTROL_FLOW_PRIM\n    stored_flag = _DISABLE_CONTROL_FLOW_PRIM\n    try:\n        _DISABLE_CONTROL_FLOW_PRIM = True\n        yield\n    finally:\n        _DISABLE_CONTROL_FLOW_PRIM = stored_flag\n\n\ndef cond(pred, true_operand, true_fun, false_operand, false_fun):\n    if _DISABLE_CONTROL_FLOW_PRIM:\n        if pred:\n            return true_fun(true_operand)\n        else:\n            return false_fun(false_operand)\n    else:\n        return lax.cond(pred, true_operand, true_fun, false_operand, false_fun)\n\n\ndef while_loop(cond_fun, body_fun, init_val):\n    if _DISABLE_CONTROL_FLOW_PRIM:\n        val = init_val\n        while cond_fun(val):\n            val = body_fun(val)\n        return val\n    else:\n        # TODO: consider jitting while_loop similar to fori_loop\n        return lax.while_loop(cond_fun, body_fun, init_val)\n\n\ndef fori_loop(lower, upper, body_fun, init_val):\n    if _DISABLE_CONTROL_FLOW_PRIM:\n        val = init_val\n        for i in range(int(lower), int(upper)):\n            val = body_fun(i, val)\n        return val\n    else:\n        return jit(lax.fori_loop, static_argnums=(2,))(lower, upper, body_fun, init_val)\n\n\ndef identity(x):\n    return x\n\n\ndef fori_collect(lower, upper, body_fun, init_val, transform=identity, progbar=True, **progbar_opts):\n    \"\"\"\n    This looping construct works like :func:`~jax.lax.fori_loop` but with the additional\n    effect of collecting values from the loop body. In addition, this allows for\n    post-processing of these samples via `transform`, and progress bar updates.\n    Note that, `progbar=False` will be faster, especially when collecting a\n    lot of samples. Refer to example usage in :func:`~numpyro.mcmc.hmc`.\n\n    :param int lower: the index to start the collective work. In other words,\n        we will skip collecting the first `lower` values.\n    :param int upper: number of times to run the loop body.\n    :param body_fun: a callable that takes a collection of\n        `np.ndarray` and returns a collection with the same shape and\n        `dtype`.\n    :param init_val: initial value to pass as argument to `body_fun`. Can\n        be any Python collection type containing `np.ndarray` objects.\n    :param transform: a callable to post-process the values returned by `body_fn`.\n    :param progbar: whether to post progress bar updates.\n    :param `**progbar_opts`: optional additional progress bar arguments. A\n        `diagnostics_fn` can be supplied which when passed the current value\n        from `body_fun` returns a string that is used to update the progress\n        bar postfix. Also a `progbar_desc` keyword argument can be supplied\n        which is used to label the progress bar.\n    :return: collection with the same type as `init_val` with values\n        collected along the leading axis of `np.ndarray` objects.\n    \"\"\"\n    assert lower < upper\n    init_val_flat, unravel_fn = ravel_pytree(transform(init_val))\n    ravel_fn = lambda x: ravel_pytree(transform(x))[0]  # noqa: E731\n\n    if not progbar:\n        collection = np.zeros((upper - lower,) + init_val_flat.shape)\n\n        def _body_fn(i, vals):\n            val, collection = vals\n            val = body_fun(val)\n            i = np.where(i >= lower, i - lower, 0)\n            collection = ops.index_update(collection, i, ravel_fn(val))\n            return val, collection\n\n        _, collection = fori_loop(0, upper, _body_fn, (init_val, collection))\n    else:\n        diagnostics_fn = progbar_opts.pop('diagnostics_fn', None)\n        progbar_desc = progbar_opts.pop('progbar_desc', '')\n        collection = []\n\n        val = init_val\n        with tqdm.trange(upper, desc=progbar_desc) as t:\n            for i in t:\n                val = body_fun(val)\n                if i >= lower:\n                    collection.append(jit(ravel_fn)(val))\n                if diagnostics_fn:\n                    t.set_postfix_str(diagnostics_fn(val), refresh=False)\n\n        collection = np.stack(collection)\n\n    return vmap(unravel_fn)(collection)\n\n\ndef copy_docs_from(source_class, full_text=False):\n    \"\"\"\n    Decorator to copy class and method docs from source to destin class.\n    \"\"\"\n\n    def decorator(destin_class):\n        # This works only in python 3.3+:\n        # if not destin_class.__doc__:\n        #     destin_class.__doc__ = source_class.__doc__\n        for name in dir(destin_class):\n            if name.startswith('_'):\n                continue\n            destin_attr = getattr(destin_class, name)\n            destin_attr = getattr(destin_attr, '__func__', destin_attr)\n            source_attr = getattr(source_class, name, None)\n            source_doc = getattr(source_attr, '__doc__', None)\n            if source_doc and not getattr(destin_attr, '__doc__', None):\n                if full_text or source_doc.startswith('See '):\n                    destin_doc = source_doc\n                else:\n                    destin_doc = 'See :meth:`{}.{}.{}`'.format(\n                        source_class.__module__, source_class.__name__, name)\n                if isinstance(destin_attr, property):\n                    # Set docs for object properties.\n                    # Since __doc__ is read-only, we need to reset the property\n                    # with the updated doc.\n                    updated_property = property(destin_attr.fget,\n                                                destin_attr.fset,\n                                                destin_attr.fdel,\n                                                destin_doc)\n                    setattr(destin_class, name, updated_property)\n                else:\n                    destin_attr.__doc__ = destin_doc\n        return destin_class\n\n    return decorator\n\n\npytree_metadata = namedtuple('pytree_metadata', ['flat', 'shape', 'size', 'dtype'])\n\n\ndef _ravel_list(*leaves):\n    leaves_metadata = tree_map(lambda l: pytree_metadata(np.ravel(l), np.shape(l), np.size(l), lax.dtype(l)),\n                               leaves)\n    leaves_idx = np.cumsum(np.array((0,) + tuple(d.size for d in leaves_metadata)))\n\n    def unravel_list(arr):\n        return [np.reshape(lax.dynamic_slice_in_dim(arr, leaves_idx[i], m.size),\n                           m.shape).astype(m.dtype)\n                for i, m in enumerate(leaves_metadata)]\n\n    return np.concatenate([m.flat for m in leaves_metadata]), unravel_list\n\n\ndef ravel_pytree(pytree):\n    leaves, treedef = tree_flatten(pytree)\n    flat, unravel_list = _ravel_list(*leaves)\n\n    def unravel_pytree(arr):\n        return tree_unflatten(treedef, unravel_list(arr))\n\n    return flat, unravel_pytree\n", "path": "numpyro/util.py"}]}