Mark xfail for leaked tracer tests (#1997)

* Mark xfail for leaked tracer tests

* test_handlers::test_plate is renamed to test_jit_trace

* remove tracer leak xfail for the current passing tests

* remove all global jax live arrays

* add issues for each tracer leak test

* fix failing tests

* fix jax core deprecation in provenance

Author: fehiepsi

.github/workflows/ci.yml

@@ -78,6 +78,20 @@ jobs:
     - name: Test x64
       run: |
         JAX_ENABLE_X64=1 pytest -vs test/test_distributions.py -k powerLaw
+    - name: Test tracer leak
+      if: matrix.python-version == '3.10'
+      env:
+        JAX_CHECK_TRACER_LEAKS: 1
+      run: |
+        pytest -vs test/contrib/einstein/test_steinvi.py::test_run_smoke -k ASVGD
+        pytest -vs test/contrib/test_infer_discrete.py::test_scan_hmm_smoke
+        pytest -vs test/infer/test_mcmc.py::test_chain_inside_jit
+        pytest -vs test/infer/test_mcmc.py::test_chain_jit_args_smoke
+        pytest -vs test/infer/test_mcmc.py::test_reuse_mcmc_run
+        pytest -vs test/infer/test_mcmc.py::test_model_with_multiple_exec_paths
+        pytest -vs test/infer/test_svi.py::test_mutable_state
+        pytest -vs test/test_distributions.py::test_mean_var -k Gompertz
+
     - name: Coveralls
       if: github.repository == 'pyro-ppl/numpyro' && matrix.python-version == '3.10'
       uses: coverallsapp/github-action@v2

numpyro/distributions/discrete.py

@@ -944,7 +944,7 @@ def entropy(self):
         logq = -jax.nn.softplus(self.logits)
         logp = -jax.nn.softplus(-self.logits)
         p = jax.scipy.special.expit(self.logits)
-        p_clip = jnp.clip(p, min=jnp.finfo(p).tiny)
+        p_clip = jnp.clip(p, jnp.finfo(p).tiny)
         return -(1 - p) * logq / p_clip - logp
 
 

numpyro/ops/provenance.py

@@ -3,7 +3,6 @@
 
 import jax
 from jax.api_util import flatten_fun, shaped_abstractify
-import jax.core as core
 from jax.experimental.pjit import pjit_p
 import jax.util as util
 
@@ -12,11 +11,21 @@
 except ImportError:
     import jax.linear_util as lu
 
+try:
+    from jax.extend.core import Literal
+except ImportError:
+    from jax.core import Literal
+
 try:
     from jax.extend.core.primitives import call_p, closed_call_p
 except ImportError:
     from jax.core import call_p, closed_call_p
 
+try:
+    from jax.api_util import debug_info
+except ImportError:
+    debug_info = None
+
 from jax.interpreters.partial_eval import trace_to_jaxpr_dynamic
 from jax.interpreters.pxla import xla_pmap_p
 
@@ -44,14 +53,29 @@ def eval_provenance(fn, **kwargs):
     """
     # Flatten the function and its arguments
     args, in_tree = jax.tree.flatten(((), kwargs))
-    wrapped_fun, out_tree = flatten_fun(lu.wrap_init(fn), in_tree)
+    fn_info = (
+        dict(debug_info=debug_info("eval_provenance fn", fn, (), kwargs))
+        if debug_info is not None
+        else {}
+    )
+    wrapped_fun, out_tree = flatten_fun(lu.wrap_init(fn, **fn_info), in_tree)
     # Abstract eval to get output pytree
     avals = util.safe_map(shaped_abstractify, args)
     # XXX: we split out the process of abstract evaluation and provenance tracking
     # for simplicity. In principle, they can be merged so that we only need to walk
     # through the equations once.
+
+    wrapped_info = (
+        dict(
+            debug_info=debug_info(
+                "eval_provenance wrapped", wrapped_fun.call_wrapped, args, {}
+            )
+        )
+        if debug_info is not None
+        else {}
+    )
     jaxpr, avals_out, _ = trace_to_jaxpr_dynamic(
-        lu.wrap_init(wrapped_fun.call_wrapped, {}), avals
+        lu.wrap_init(wrapped_fun.call_wrapped, {}, **wrapped_info), avals
     )
 
     # get provenances of flatten kwargs
@@ -69,12 +93,12 @@ def track_deps_jaxpr(jaxpr, provenance_inputs):
     env = {}
 
     def read(v):
-        if isinstance(v, core.Literal):
+        if isinstance(v, Literal):
             return frozenset()
         return env.get(v, frozenset())
 
     def write(v, p):
-        if isinstance(v, core.Literal):
+        if isinstance(v, Literal):
             return
         env[v] = read(v) | p
 

test/conftest.py

@@ -3,14 +3,22 @@
 
 import os
 
+import jax
 from jax import config
 
 from numpyro.util import set_rng_seed
 
 config.update("jax_platform_name", "cpu")  # noqa: E702
 
 
+SETUP_STATE = {"is_first_test": True}
+
+
 def pytest_runtest_setup(item):
+    if SETUP_STATE["is_first_test"]:
+        SETUP_STATE["is_first_test"] = False
+        assert len(jax.live_arrays()) == 0
+
     if "JAX_ENABLE_X64" in os.environ:
         config.update("jax_enable_x64", True)
     set_rng_seed(0)

test/contrib/einstein/test_steinvi.py

@@ -3,6 +3,7 @@
 
 from collections import namedtuple
 from functools import partial
+import os
 import string
 
 import numpy as np
@@ -119,6 +120,9 @@ def model(features, labels):
 @pytest.mark.parametrize("kernel", KERNELS)
 @pytest.mark.parametrize("problem", (uniform_normal, regression))
 @pytest.mark.parametrize("method", ("ASVGD", "SVGD", "SteinVI"))
+@pytest.mark.xfail(
+    os.getenv("JAX_CHECK_TRACER_LEAKS") == "1", reason="Expected tracer leak"
+)
 def test_run_smoke(kernel, problem, method):
     true_coefs, data, model = problem()
     if method == "ASVGD":

test/contrib/hsgp/test_laplacian.py

@@ -131,12 +131,12 @@ def test_eigenfunctions(x: ArrayLike, ell: float | int, m: int | list[int]):
         (1, 2, False),
         ([1, 1], 2, False),
         (np.array([1, 1])[..., None], 2, False),
-        (jnp.array([1, 1])[..., None], 2, False),
+        (np.array([1, 1])[..., None], 2, False),
+        (np.array([1, 1]), 2, True),
         (np.array([1, 1]), 2, True),
-        (jnp.array([1, 1]), 2, True),
         ([1, 1], 1, True),
         (np.array([1, 1]), 1, True),
-        (jnp.array([1, 1]), 1, True),
+        (np.array([1, 1]), 1, True),
     ],
     ids=[
         "ell-float",

test/contrib/stochastic_support/test_dcc.py

@@ -19,14 +19,14 @@
 
 @pytest.mark.parametrize(
     "branch_dist",
-    [dist.Normal(0, 1), dist.Gamma(1, 1)],
+    [lambda: dist.Normal(0, 1), lambda: dist.Gamma(1, 1)],
 )
 @pytest.mark.xfail(raises=RuntimeError)
 def test_continuous_branching(branch_dist):
     rng_key = random.PRNGKey(0)
 
     def model():
-        model1 = numpyro.sample("model1", branch_dist, infer={"branching": True})
+        model1 = numpyro.sample("model1", branch_dist(), infer={"branching": True})
         mean = 1.0 if model1 == 0 else 2.0
         numpyro.sample("obs", dist.Normal(mean, 1.0), obs=0.2)
 

test/contrib/test_enum_elbo.py

@@ -1247,7 +1247,6 @@ def test_elbo_enumerate_plates_6(scale):
 
     @config_enumerate
     @handlers.scale(scale=scale)
-    @handlers.trace
     def model_iplate_iplate(data, params):
         probs_a = pyro.param(
             "probs_a", params["probs_a"], constraint=constraints.simplex
@@ -1305,7 +1304,6 @@ def model_iplate_plate(data, params):
 
     @config_enumerate
     @handlers.scale(scale=scale)
-    @handlers.trace
     def model_plate_iplate(data, params):
         probs_a = pyro.param(
             "probs_a", params["probs_a"], constraint=constraints.simplex
@@ -1423,7 +1421,6 @@ def test_elbo_enumerate_plates_7(scale):
 
     @config_enumerate
     @handlers.scale(scale=scale)
-    @handlers.trace
     def model_iplate_iplate(data, params):
         probs_a = pyro.param(
             "probs_a", params["probs_a"], constraint=constraints.simplex
@@ -1489,7 +1486,6 @@ def model_iplate_plate(data, params):
 
     @config_enumerate
     @handlers.scale(scale=scale)
-    @handlers.trace
     def model_plate_iplate(data, params):
         probs_a = pyro.param(
             "probs_a", params["probs_a"], constraint=constraints.simplex

test/contrib/test_infer_discrete.py

@@ -2,6 +2,7 @@
 # SPDX-License-Identifier: Apache-2.0
 
 import logging
+import os
 
 import numpy as np
 from numpy.testing import assert_allclose
@@ -95,6 +96,10 @@ def hmm(data, hidden_dim=10):
     ],
 )
 @pytest.mark.parametrize("temperature", [0, 1])
+@pytest.mark.xfail(
+    os.getenv("JAX_CHECK_TRACER_LEAKS") == "1",
+    reason="Expected tracer leak: https://github.com/pyro-ppl/numpyro/issues/1998",
+)
 def test_scan_hmm_smoke(length, temperature):
     # This should match the example in the infer_discrete docstring.
     def hmm(data, hidden_dim=10):

test/contrib/test_tfp.py

@@ -179,7 +179,7 @@ def model(data):
         (
             "ReplicaExchangeMC",
             dict(
-                inverse_temperatures=0.5 ** jnp.arange(4), make_kernel_fn=make_kernel_fn
+                inverse_temperatures=0.5 ** np.arange(4), make_kernel_fn=make_kernel_fn
             ),
         ),
     ],