We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Leak when running examples with JAX_CHECK_TRACER_LEAKS=1
Hi
when I run ar2.py from yours examples (or bnn.py did not try with others) with the environment variable JAX_CHECK_TRACER_LEAKS=1 they fail. (I had to use it to try to find an issue with a function I had written)

the exception raised is 

Exception: Leaked level MainTrace(1,DynamicJaxprTrace). Leaked tracer(s): [Traced<ShapedArray(int32[])>with<DynamicJaxprTrace(level=1/0)>].

here the complete log:

Traceback (most recent call last):
  File "/home/ffranco/Downloads/ar2.py", line 138, in <module>
    main(args)
  File "/home/ffranco/Downloads/ar2.py", line 117, in main
    run_inference(model, args, rng_key, y)
  File "/home/ffranco/Downloads/ar2.py", line 96, in run_inference
    mcmc.run(rng_key, y=y)
  File "/home/ffranco/numpyro08/lib/python3.7/site-packages/numpyro/infer/mcmc.py", line 593, in run
    states_flat, last_state = partial_map_fn(map_args)
  File "/home/ffranco/numpyro08/lib/python3.7/site-packages/numpyro/infer/mcmc.py", line 386, in _single_chain_mcmc
    model_kwargs=kwargs,
  File "/home/ffranco/numpyro08/lib/python3.7/site-packages/numpyro/infer/hmc.py", line 707, in init
    rng_key_init_model, model_args, model_kwargs, init_params
  File "/home/ffranco/numpyro08/lib/python3.7/site-packages/numpyro/infer/hmc.py", line 659, in _init_state
    forward_mode_differentiation=self._forward_mode_differentiation,
  File "/home/ffranco/numpyro08/lib/python3.7/site-packages/numpyro/infer/util.py", line 606, in initialize_model
    ) = _get_model_transforms(substituted_model, model_args, model_kwargs)
  File "/home/ffranco/numpyro08/lib/python3.7/site-packages/numpyro/infer/util.py", line 404, in _get_model_transforms
    model_trace = trace(model).get_trace(*model_args, **model_kwargs)
  File "/home/ffranco/numpyro08/lib/python3.7/site-packages/numpyro/handlers.py", line 171, in get_trace
    self(*args, **kwargs)
  File "/home/ffranco/numpyro08/lib/python3.7/site-packages/numpyro/primitives.py", line 105, in __call__
    return self.fn(*args, **kwargs)
  File "/home/ffranco/numpyro08/lib/python3.7/site-packages/numpyro/primitives.py", line 105, in __call__
    return self.fn(*args, **kwargs)
  File "/home/ffranco/numpyro08/lib/python3.7/site-packages/numpyro/primitives.py", line 105, in __call__
    return self.fn(*args, **kwargs)
  File "/home/ffranco/Downloads/ar2.py", line 67, in ar2_scan
    scan(transition, init, timesteps)
  File "/home/ffranco/numpyro08/lib/python3.7/site-packages/numpyro/contrib/control_flow/scan.py", line 438, in scan
    msg = apply_stack(initial_msg)
  File "/home/ffranco/numpyro08/lib/python3.7/site-packages/numpyro/primitives.py", line 53, in apply_stack
    default_process_message(msg)
  File "/home/ffranco/numpyro08/lib/python3.7/site-packages/numpyro/primitives.py", line 28, in default_process_message
    msg["value"] = msg["fn"](*msg["args"], **msg["kwargs"])
  File "/home/ffranco/numpyro08/lib/python3.7/site-packages/numpyro/contrib/control_flow/scan.py", line 306, in scan_wrapper
    body_fn, wrapped_carry, xs, length=length, reverse=reverse
  File "/home/ffranco/numpyro08/lib/python3.7/site-packages/jax/_src/traceback_util.py", line 162, in reraise_with_filtered_traceback
    return fun(*args, **kwargs)
  File "/home/ffranco/numpyro08/lib/python3.7/site-packages/jax/_src/lax/control_flow.py", line 1345, in scan
    init_flat, carry_avals, carry_avals_out, init_tree, *rest = _create_jaxpr(init)
  File "/home/ffranco/numpyro08/lib/python3.7/site-packages/jax/_src/lax/control_flow.py", line 1332, in _create_jaxpr
    f, in_tree, carry_avals + x_avals, "scan")
  File "/home/ffranco/numpyro08/lib/python3.7/site-packages/jax/_src/util.py", line 185, in wrapper
    return f(*args, **kwargs)
  File "/home/ffranco/numpyro08/lib/python3.7/site-packages/jax/_src/lax/control_flow.py", line 78, in _initial_style_jaxpr
    fun, in_tree, in_avals, primitive_name)
  File "/home/ffranco/numpyro08/lib/python3.7/site-packages/jax/_src/util.py", line 185, in wrapper
    return f(*args, **kwargs)
  File "/home/ffranco/numpyro08/lib/python3.7/site-packages/jax/_src/lax/control_flow.py", line 71, in _initial_style_open_jaxpr
    jaxpr, _, consts = pe.trace_to_jaxpr_dynamic(wrapped_fun, in_avals, debug)
  File "/home/ffranco/numpyro08/lib/python3.7/site-packages/jax/interpreters/partial_eval.py", line 1511, in trace_to_jaxpr_dynamic
    del main, fun
  File "/home/ffranco/anaconda3/lib/python3.7/contextlib.py", line 119, in __exit__
    next(self.gen)
  File "/home/ffranco/numpyro08/lib/python3.7/site-packages/jax/core.py", line 810, in new_main
    raise Exception(f'Leaked level {t()}. Leaked tracer(s): {leaked_tracers}.')
jax._src.traceback_util.UnfilteredStackTrace: Exception: Leaked level MainTrace(1,DynamicJaxprTrace). Leaked tracer(s): [Traced<ShapedArray(int32[])>with<DynamicJaxprTrace(level=1/0)>].

The stack trace below excludes JAX-internal frames.
The preceding is the original exception that occurred, unmodified.

--------------------

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/home/ffranco/Downloads/ar2.py", line 138, in <module>
    main(args)
  File "/home/ffranco/Downloads/ar2.py", line 117, in main
    run_inference(model, args, rng_key, y)
  File "/home/ffranco/Downloads/ar2.py", line 96, in run_inference
    mcmc.run(rng_key, y=y)
  File "/home/ffranco/numpyro08/lib/python3.7/site-packages/numpyro/infer/mcmc.py", line 593, in run
    states_flat, last_state = partial_map_fn(map_args)
  File "/home/ffranco/numpyro08/lib/python3.7/site-packages/numpyro/infer/mcmc.py", line 386, in _single_chain_mcmc
    model_kwargs=kwargs,
  File "/home/ffranco/numpyro08/lib/python3.7/site-packages/numpyro/infer/hmc.py", line 707, in init
    rng_key_init_model, model_args, model_kwargs, init_params
  File "/home/ffranco/numpyro08/lib/python3.7/site-packages/numpyro/infer/hmc.py", line 659, in _init_state
    forward_mode_differentiation=self._forward_mode_differentiation,
  File "/home/ffranco/numpyro08/lib/python3.7/site-packages/numpyro/infer/util.py", line 606, in initialize_model
    ) = _get_model_transforms(substituted_model, model_args, model_kwargs)
  File "/home/ffranco/numpyro08/lib/python3.7/site-packages/numpyro/infer/util.py", line 404, in _get_model_transforms
    model_trace = trace(model).get_trace(*model_args, **model_kwargs)
  File "/home/ffranco/numpyro08/lib/python3.7/site-packages/numpyro/handlers.py", line 171, in get_trace
    self(*args, **kwargs)
  File "/home/ffranco/numpyro08/lib/python3.7/site-packages/numpyro/primitives.py", line 105, in __call__
    return self.fn(*args, **kwargs)
  File "/home/ffranco/numpyro08/lib/python3.7/site-packages/numpyro/primitives.py", line 105, in __call__
    return self.fn(*args, **kwargs)
  File "/home/ffranco/numpyro08/lib/python3.7/site-packages/numpyro/primitives.py", line 105, in __call__
    return self.fn(*args, **kwargs)
  File "/home/ffranco/Downloads/ar2.py", line 67, in ar2_scan
    scan(transition, init, timesteps)
  File "/home/ffranco/numpyro08/lib/python3.7/site-packages/numpyro/contrib/control_flow/scan.py", line 438, in scan
    msg = apply_stack(initial_msg)
  File "/home/ffranco/numpyro08/lib/python3.7/site-packages/numpyro/primitives.py", line 53, in apply_stack
    default_process_message(msg)
  File "/home/ffranco/numpyro08/lib/python3.7/site-packages/numpyro/primitives.py", line 28, in default_process_message
    msg["value"] = msg["fn"](*msg["args"], **msg["kwargs"])
  File "/home/ffranco/numpyro08/lib/python3.7/site-packages/numpyro/contrib/control_flow/scan.py", line 306, in scan_wrapper
    body_fn, wrapped_carry, xs, length=length, reverse=reverse
  File "/home/ffranco/anaconda3/lib/python3.7/contextlib.py", line 119, in __exit__
    next(self.gen)
Exception: Leaked level MainTrace(1,DynamicJaxprTrace). Leaked tracer(s): [Traced<ShapedArray(int32[])>with<DynamicJaxprTrace(level=1/0)>].

--- 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/ops/pytree.py`
Content:
```
1 # Copyright Contributors to the Pyro project.
2 # SPDX-License-Identifier: Apache-2.0
3 
4 from jax.tree_util import register_pytree_node_class
5 
6 
7 @register_pytree_node_class
8 class PytreeTrace:
9     def __init__(self, trace):
10         self.trace = trace
11 
12     def tree_flatten(self):
13         trace, aux_trace = {}, {}
14         for name, site in self.trace.items():
15             if site["type"] in ["sample", "deterministic", "plate", "param"]:
16                 trace[name], aux_trace[name] = {}, {"_control_flow_done": True}
17                 for key in site:
18                     if key == "fn":
19                         if site["type"] == "sample":
20                             trace[name][key] = site[key]
21                         elif site["type"] == "plate":
22                             aux_trace[name][key] = site[key]
23                     elif key in ["args", "value", "intermediates"]:
24                         trace[name][key] = site[key]
25                     # scanned sites have stop field because we trace them inside a block handler
26                     elif key != "stop":
27                         if key == "kwargs":
28                             kwargs = site["kwargs"].copy()
29                             if "rng_key" in kwargs:
30                                 # rng_key is not traced else it is collected by the
31                                 # scan primitive which doesn't make sense
32                                 # set to None to avoid leaks during tracing by JAX
33                                 kwargs["rng_key"] = None
34                             aux_trace[name][key] = kwargs
35                         else:
36                             aux_trace[name][key] = site[key]
37         # keep the site order information because in JAX, flatten and unflatten do not preserve
38         # the order of keys in a dict
39         site_names = list(trace.keys())
40         return (trace,), (aux_trace, site_names)
41 
42     @classmethod
43     def tree_unflatten(cls, aux_data, children):
44         aux_trace, site_names = aux_data
45         (trace,) = children
46         trace_with_aux = {}
47         for name in site_names:
48             trace[name].update(aux_trace[name])
49             trace_with_aux[name] = trace[name]
50         return cls(trace_with_aux)
51
```
--- 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/ops/pytree.py b/numpyro/ops/pytree.py
--- a/numpyro/ops/pytree.py
+++ b/numpyro/ops/pytree.py
@@ -32,6 +32,12 @@
                                 # set to None to avoid leaks during tracing by JAX
                                 kwargs["rng_key"] = None
                             aux_trace[name][key] = kwargs
+                        elif key == "infer":
+                            kwargs = site["infer"].copy()
+                            if "_scan_current_index" in kwargs:
+                                # set to None to avoid leaks during tracing by JAX
+                                kwargs["_scan_current_index"] = None
+                            aux_trace[name][key] = kwargs
                         else:
                             aux_trace[name][key] = site[key]
         # keep the site order information because in JAX, flatten and unflatten do not preserve

{"golden_diff": "diff --git a/numpyro/ops/pytree.py b/numpyro/ops/pytree.py\n--- a/numpyro/ops/pytree.py\n+++ b/numpyro/ops/pytree.py\n@@ -32,6 +32,12 @@\n                                 # set to None to avoid leaks during tracing by JAX\n                                 kwargs[\"rng_key\"] = None\n                             aux_trace[name][key] = kwargs\n+                        elif key == \"infer\":\n+                            kwargs = site[\"infer\"].copy()\n+                            if \"_scan_current_index\" in kwargs:\n+                                # set to None to avoid leaks during tracing by JAX\n+                                kwargs[\"_scan_current_index\"] = None\n+                            aux_trace[name][key] = kwargs\n                         else:\n                             aux_trace[name][key] = site[key]\n         # keep the site order information because in JAX, flatten and unflatten do not preserve\n", "issue": "Leak when running examples with JAX_CHECK_TRACER_LEAKS=1\nHi\r\nwhen I run ar2.py from yours examples (or bnn.py did not try with others) with the environment variable JAX_CHECK_TRACER_LEAKS=1 they fail. (I had to use it to try to find an issue with a function I had written)\r\n\r\nthe exception raised is \r\n\r\nException: Leaked level MainTrace(1,DynamicJaxprTrace). Leaked tracer(s): [Traced<ShapedArray(int32[])>with<DynamicJaxprTrace(level=1/0)>].\r\n\r\nhere the complete log:\r\n\r\nTraceback (most recent call last):\r\n  File \"/home/ffranco/Downloads/ar2.py\", line 138, in <module>\r\n    main(args)\r\n  File \"/home/ffranco/Downloads/ar2.py\", line 117, in main\r\n    run_inference(model, args, rng_key, y)\r\n  File \"/home/ffranco/Downloads/ar2.py\", line 96, in run_inference\r\n    mcmc.run(rng_key, y=y)\r\n  File \"/home/ffranco/numpyro08/lib/python3.7/site-packages/numpyro/infer/mcmc.py\", line 593, in run\r\n    states_flat, last_state = partial_map_fn(map_args)\r\n  File \"/home/ffranco/numpyro08/lib/python3.7/site-packages/numpyro/infer/mcmc.py\", line 386, in _single_chain_mcmc\r\n    model_kwargs=kwargs,\r\n  File \"/home/ffranco/numpyro08/lib/python3.7/site-packages/numpyro/infer/hmc.py\", line 707, in init\r\n    rng_key_init_model, model_args, model_kwargs, init_params\r\n  File \"/home/ffranco/numpyro08/lib/python3.7/site-packages/numpyro/infer/hmc.py\", line 659, in _init_state\r\n    forward_mode_differentiation=self._forward_mode_differentiation,\r\n  File \"/home/ffranco/numpyro08/lib/python3.7/site-packages/numpyro/infer/util.py\", line 606, in initialize_model\r\n    ) = _get_model_transforms(substituted_model, model_args, model_kwargs)\r\n  File \"/home/ffranco/numpyro08/lib/python3.7/site-packages/numpyro/infer/util.py\", line 404, in _get_model_transforms\r\n    model_trace = trace(model).get_trace(*model_args, **model_kwargs)\r\n  File \"/home/ffranco/numpyro08/lib/python3.7/site-packages/numpyro/handlers.py\", line 171, in get_trace\r\n    self(*args, **kwargs)\r\n  File \"/home/ffranco/numpyro08/lib/python3.7/site-packages/numpyro/primitives.py\", line 105, in __call__\r\n    return self.fn(*args, **kwargs)\r\n  File \"/home/ffranco/numpyro08/lib/python3.7/site-packages/numpyro/primitives.py\", line 105, in __call__\r\n    return self.fn(*args, **kwargs)\r\n  File \"/home/ffranco/numpyro08/lib/python3.7/site-packages/numpyro/primitives.py\", line 105, in __call__\r\n    return self.fn(*args, **kwargs)\r\n  File \"/home/ffranco/Downloads/ar2.py\", line 67, in ar2_scan\r\n    scan(transition, init, timesteps)\r\n  File \"/home/ffranco/numpyro08/lib/python3.7/site-packages/numpyro/contrib/control_flow/scan.py\", line 438, in scan\r\n    msg = apply_stack(initial_msg)\r\n  File \"/home/ffranco/numpyro08/lib/python3.7/site-packages/numpyro/primitives.py\", line 53, in apply_stack\r\n    default_process_message(msg)\r\n  File \"/home/ffranco/numpyro08/lib/python3.7/site-packages/numpyro/primitives.py\", line 28, in default_process_message\r\n    msg[\"value\"] = msg[\"fn\"](*msg[\"args\"], **msg[\"kwargs\"])\r\n  File \"/home/ffranco/numpyro08/lib/python3.7/site-packages/numpyro/contrib/control_flow/scan.py\", line 306, in scan_wrapper\r\n    body_fn, wrapped_carry, xs, length=length, reverse=reverse\r\n  File \"/home/ffranco/numpyro08/lib/python3.7/site-packages/jax/_src/traceback_util.py\", line 162, in reraise_with_filtered_traceback\r\n    return fun(*args, **kwargs)\r\n  File \"/home/ffranco/numpyro08/lib/python3.7/site-packages/jax/_src/lax/control_flow.py\", line 1345, in scan\r\n    init_flat, carry_avals, carry_avals_out, init_tree, *rest = _create_jaxpr(init)\r\n  File \"/home/ffranco/numpyro08/lib/python3.7/site-packages/jax/_src/lax/control_flow.py\", line 1332, in _create_jaxpr\r\n    f, in_tree, carry_avals + x_avals, \"scan\")\r\n  File \"/home/ffranco/numpyro08/lib/python3.7/site-packages/jax/_src/util.py\", line 185, in wrapper\r\n    return f(*args, **kwargs)\r\n  File \"/home/ffranco/numpyro08/lib/python3.7/site-packages/jax/_src/lax/control_flow.py\", line 78, in _initial_style_jaxpr\r\n    fun, in_tree, in_avals, primitive_name)\r\n  File \"/home/ffranco/numpyro08/lib/python3.7/site-packages/jax/_src/util.py\", line 185, in wrapper\r\n    return f(*args, **kwargs)\r\n  File \"/home/ffranco/numpyro08/lib/python3.7/site-packages/jax/_src/lax/control_flow.py\", line 71, in _initial_style_open_jaxpr\r\n    jaxpr, _, consts = pe.trace_to_jaxpr_dynamic(wrapped_fun, in_avals, debug)\r\n  File \"/home/ffranco/numpyro08/lib/python3.7/site-packages/jax/interpreters/partial_eval.py\", line 1511, in trace_to_jaxpr_dynamic\r\n    del main, fun\r\n  File \"/home/ffranco/anaconda3/lib/python3.7/contextlib.py\", line 119, in __exit__\r\n    next(self.gen)\r\n  File \"/home/ffranco/numpyro08/lib/python3.7/site-packages/jax/core.py\", line 810, in new_main\r\n    raise Exception(f'Leaked level {t()}. Leaked tracer(s): {leaked_tracers}.')\r\njax._src.traceback_util.UnfilteredStackTrace: Exception: Leaked level MainTrace(1,DynamicJaxprTrace). Leaked tracer(s): [Traced<ShapedArray(int32[])>with<DynamicJaxprTrace(level=1/0)>].\r\n\r\nThe stack trace below excludes JAX-internal frames.\r\nThe preceding is the original exception that occurred, unmodified.\r\n\r\n--------------------\r\n\r\nThe above exception was the direct cause of the following exception:\r\n\r\nTraceback (most recent call last):\r\n  File \"/home/ffranco/Downloads/ar2.py\", line 138, in <module>\r\n    main(args)\r\n  File \"/home/ffranco/Downloads/ar2.py\", line 117, in main\r\n    run_inference(model, args, rng_key, y)\r\n  File \"/home/ffranco/Downloads/ar2.py\", line 96, in run_inference\r\n    mcmc.run(rng_key, y=y)\r\n  File \"/home/ffranco/numpyro08/lib/python3.7/site-packages/numpyro/infer/mcmc.py\", line 593, in run\r\n    states_flat, last_state = partial_map_fn(map_args)\r\n  File \"/home/ffranco/numpyro08/lib/python3.7/site-packages/numpyro/infer/mcmc.py\", line 386, in _single_chain_mcmc\r\n    model_kwargs=kwargs,\r\n  File \"/home/ffranco/numpyro08/lib/python3.7/site-packages/numpyro/infer/hmc.py\", line 707, in init\r\n    rng_key_init_model, model_args, model_kwargs, init_params\r\n  File \"/home/ffranco/numpyro08/lib/python3.7/site-packages/numpyro/infer/hmc.py\", line 659, in _init_state\r\n    forward_mode_differentiation=self._forward_mode_differentiation,\r\n  File \"/home/ffranco/numpyro08/lib/python3.7/site-packages/numpyro/infer/util.py\", line 606, in initialize_model\r\n    ) = _get_model_transforms(substituted_model, model_args, model_kwargs)\r\n  File \"/home/ffranco/numpyro08/lib/python3.7/site-packages/numpyro/infer/util.py\", line 404, in _get_model_transforms\r\n    model_trace = trace(model).get_trace(*model_args, **model_kwargs)\r\n  File \"/home/ffranco/numpyro08/lib/python3.7/site-packages/numpyro/handlers.py\", line 171, in get_trace\r\n    self(*args, **kwargs)\r\n  File \"/home/ffranco/numpyro08/lib/python3.7/site-packages/numpyro/primitives.py\", line 105, in __call__\r\n    return self.fn(*args, **kwargs)\r\n  File \"/home/ffranco/numpyro08/lib/python3.7/site-packages/numpyro/primitives.py\", line 105, in __call__\r\n    return self.fn(*args, **kwargs)\r\n  File \"/home/ffranco/numpyro08/lib/python3.7/site-packages/numpyro/primitives.py\", line 105, in __call__\r\n    return self.fn(*args, **kwargs)\r\n  File \"/home/ffranco/Downloads/ar2.py\", line 67, in ar2_scan\r\n    scan(transition, init, timesteps)\r\n  File \"/home/ffranco/numpyro08/lib/python3.7/site-packages/numpyro/contrib/control_flow/scan.py\", line 438, in scan\r\n    msg = apply_stack(initial_msg)\r\n  File \"/home/ffranco/numpyro08/lib/python3.7/site-packages/numpyro/primitives.py\", line 53, in apply_stack\r\n    default_process_message(msg)\r\n  File \"/home/ffranco/numpyro08/lib/python3.7/site-packages/numpyro/primitives.py\", line 28, in default_process_message\r\n    msg[\"value\"] = msg[\"fn\"](*msg[\"args\"], **msg[\"kwargs\"])\r\n  File \"/home/ffranco/numpyro08/lib/python3.7/site-packages/numpyro/contrib/control_flow/scan.py\", line 306, in scan_wrapper\r\n    body_fn, wrapped_carry, xs, length=length, reverse=reverse\r\n  File \"/home/ffranco/anaconda3/lib/python3.7/contextlib.py\", line 119, in __exit__\r\n    next(self.gen)\r\nException: Leaked level MainTrace(1,DynamicJaxprTrace). Leaked tracer(s): [Traced<ShapedArray(int32[])>with<DynamicJaxprTrace(level=1/0)>].\n", "before_files": [{"content": "# Copyright Contributors to the Pyro project.\n# SPDX-License-Identifier: Apache-2.0\n\nfrom jax.tree_util import register_pytree_node_class\n\n\n@register_pytree_node_class\nclass PytreeTrace:\n    def __init__(self, trace):\n        self.trace = trace\n\n    def tree_flatten(self):\n        trace, aux_trace = {}, {}\n        for name, site in self.trace.items():\n            if site[\"type\"] in [\"sample\", \"deterministic\", \"plate\", \"param\"]:\n                trace[name], aux_trace[name] = {}, {\"_control_flow_done\": True}\n                for key in site:\n                    if key == \"fn\":\n                        if site[\"type\"] == \"sample\":\n                            trace[name][key] = site[key]\n                        elif site[\"type\"] == \"plate\":\n                            aux_trace[name][key] = site[key]\n                    elif key in [\"args\", \"value\", \"intermediates\"]:\n                        trace[name][key] = site[key]\n                    # scanned sites have stop field because we trace them inside a block handler\n                    elif key != \"stop\":\n                        if key == \"kwargs\":\n                            kwargs = site[\"kwargs\"].copy()\n                            if \"rng_key\" in kwargs:\n                                # rng_key is not traced else it is collected by the\n                                # scan primitive which doesn't make sense\n                                # set to None to avoid leaks during tracing by JAX\n                                kwargs[\"rng_key\"] = None\n                            aux_trace[name][key] = kwargs\n                        else:\n                            aux_trace[name][key] = site[key]\n        # keep the site order information because in JAX, flatten and unflatten do not preserve\n        # the order of keys in a dict\n        site_names = list(trace.keys())\n        return (trace,), (aux_trace, site_names)\n\n    @classmethod\n    def tree_unflatten(cls, aux_data, children):\n        aux_trace, site_names = aux_data\n        (trace,) = children\n        trace_with_aux = {}\n        for name in site_names:\n            trace[name].update(aux_trace[name])\n            trace_with_aux[name] = trace[name]\n        return cls(trace_with_aux)\n", "path": "numpyro/ops/pytree.py"}], "after_files": [{"content": "# Copyright Contributors to the Pyro project.\n# SPDX-License-Identifier: Apache-2.0\n\nfrom jax.tree_util import register_pytree_node_class\n\n\n@register_pytree_node_class\nclass PytreeTrace:\n    def __init__(self, trace):\n        self.trace = trace\n\n    def tree_flatten(self):\n        trace, aux_trace = {}, {}\n        for name, site in self.trace.items():\n            if site[\"type\"] in [\"sample\", \"deterministic\", \"plate\", \"param\"]:\n                trace[name], aux_trace[name] = {}, {\"_control_flow_done\": True}\n                for key in site:\n                    if key == \"fn\":\n                        if site[\"type\"] == \"sample\":\n                            trace[name][key] = site[key]\n                        elif site[\"type\"] == \"plate\":\n                            aux_trace[name][key] = site[key]\n                    elif key in [\"args\", \"value\", \"intermediates\"]:\n                        trace[name][key] = site[key]\n                    # scanned sites have stop field because we trace them inside a block handler\n                    elif key != \"stop\":\n                        if key == \"kwargs\":\n                            kwargs = site[\"kwargs\"].copy()\n                            if \"rng_key\" in kwargs:\n                                # rng_key is not traced else it is collected by the\n                                # scan primitive which doesn't make sense\n                                # set to None to avoid leaks during tracing by JAX\n                                kwargs[\"rng_key\"] = None\n                            aux_trace[name][key] = kwargs\n                        elif key == \"infer\":\n                            kwargs = site[\"infer\"].copy()\n                            if \"_scan_current_index\" in kwargs:\n                                # set to None to avoid leaks during tracing by JAX\n                                kwargs[\"_scan_current_index\"] = None\n                            aux_trace[name][key] = kwargs\n                        else:\n                            aux_trace[name][key] = site[key]\n        # keep the site order information because in JAX, flatten and unflatten do not preserve\n        # the order of keys in a dict\n        site_names = list(trace.keys())\n        return (trace,), (aux_trace, site_names)\n\n    @classmethod\n    def tree_unflatten(cls, aux_data, children):\n        aux_trace, site_names = aux_data\n        (trace,) = children\n        trace_with_aux = {}\n        for name in site_names:\n            trace[name].update(aux_trace[name])\n            trace_with_aux[name] = trace[name]\n        return cls(trace_with_aux)\n", "path": "numpyro/ops/pytree.py"}]}