Repository: pyro-ppl/numpyro

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)>].


Discussion:

Resolving patch:
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 @@ def tree_flatten(self):
                                 # 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


Test patch: