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

Diff for: .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

Diff for: 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
 
 

Diff for: 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
 

Diff for: 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)

Diff for: 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":

Diff for: 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",

Diff for: 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)
 

Diff for: 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

Diff for: 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):

Diff for: 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
             ),
         ),
     ],

Diff for: test/infer/test_ensemble_mcmc.py

@@ -1,6 +1,7 @@
 # Copyright Contributors to the Pyro project.
 # SPDX-License-Identifier: Apache-2.0
 
+import numpy as np
 import pytest
 
 import jax.numpy as jnp
@@ -15,10 +16,13 @@
 # reused for all smoke-tests
 N, dim = 3000, 3
 
-data = random.normal(random.PRNGKey(0), (N, dim))
-true_coefs = jnp.arange(1.0, dim + 1.0)
-logits = jnp.sum(true_coefs * data, axis=-1)
-labels = dist.Bernoulli(logits=logits).sample(random.PRNGKey(1))
+data = np.random.default_rng(0).normal(N, dim)
+true_coefs = np.arange(1.0, dim + 1.0)
+logits = np.sum(true_coefs * data, axis=-1)
+
+
+def labels_maker():
+    return dist.Bernoulli(logits=logits).sample(random.PRNGKey(1))
 
 
 def model(labels):
@@ -54,7 +58,7 @@ def test_chain_smoke(kernel_cls, n_chain, method):
     )
 
     with pytest.raises(AssertionError, match="chain_method"):
-        mcmc.run(random.PRNGKey(2), labels)
+        mcmc.run(random.PRNGKey(2), labels_maker())
 
 
 @pytest.mark.parametrize("kernel_cls", [AIES, ESS])
@@ -70,7 +74,7 @@ def test_out_shape_smoke(kernel_cls):
         num_chains=n_chains,
         chain_method="vectorized",
     )
-    mcmc.run(random.PRNGKey(2), labels)
+    mcmc.run(random.PRNGKey(2), labels_maker())
 
     assert mcmc.get_samples(group_by_chain=True)["coefs"].shape[0] == n_chains
 
@@ -94,6 +98,7 @@ def test_multirun(kernel_cls):
         num_chains=n_chains,
         chain_method="vectorized",
     )
+    labels = labels_maker()
     mcmc.run(random.PRNGKey(2), labels)
     mcmc.run(random.PRNGKey(3), labels)
 
@@ -111,5 +116,6 @@ def test_warmup(kernel_cls):
         num_chains=n_chains,
         chain_method="vectorized",
     )
+    labels = labels_maker()
     mcmc.warmup(random.PRNGKey(2), labels)
     mcmc.run(random.PRNGKey(3), labels)

Diff for: test/infer/test_gradient.py

@@ -37,7 +37,7 @@ def guide_0(data, params):
         pyro.sample("z", dist.Categorical(probs))
 
 
-params_0 = {"probs": jnp.array([[0.4, 0.6], [0.5, 0.5]])}
+params_0 = {"probs": np.array([[0.4, 0.6], [0.5, 0.5]])}
 
 
 def model_1(data, params):
@@ -57,8 +57,8 @@ def guide_1(data, params):
 
 
 params_1 = {
-    "probs_a": jnp.array([0.5, 0.5]),
-    "probs_b": jnp.array([[[0.5, 0.5], [0.6, 0.4]], [[0.4, 0.6], [0.35, 0.65]]]),
+    "probs_a": np.array([0.5, 0.5]),
+    "probs_b": np.array([[[0.5, 0.5], [0.6, 0.4]], [[0.4, 0.6], [0.35, 0.65]]]),
 }
 
 
@@ -88,19 +88,19 @@ def guide_2(data, params):
 
 
 params_2 = {
-    "probs_a": jnp.array([0.5, 0.5]),
-    "probs_b": jnp.array([[0.4, 0.6], [0.3, 0.7]]),
-    "probs_c": jnp.array([[[0.3, 0.7], [0.8, 0.2]], [[0.2, 0.8], [0.5, 0.5]]]),
-    "probs_d": jnp.array([[[0.2, 0.8], [0.9, 0.1]], [[0.1, 0.9], [0.4, 0.6]]]),
+    "probs_a": np.array([0.5, 0.5]),
+    "probs_b": np.array([[0.4, 0.6], [0.3, 0.7]]),
+    "probs_c": np.array([[[0.3, 0.7], [0.8, 0.2]], [[0.2, 0.8], [0.5, 0.5]]]),
+    "probs_d": np.array([[[0.2, 0.8], [0.9, 0.1]], [[0.1, 0.9], [0.4, 0.6]]]),
 }
 
 
 @pytest.mark.parametrize(
     "model,guide,params,data",
     [
-        (model_0, guide_0, params_0, jnp.array([-0.5, 2.0])),
-        (model_1, guide_1, params_1, jnp.array([-0.5, 2.0])),
-        (model_2, guide_2, params_2, jnp.array([0, 1])),
+        (model_0, guide_0, params_0, np.array([-0.5, 2.0])),
+        (model_1, guide_1, params_1, np.array([-0.5, 2.0])),
+        (model_2, guide_2, params_2, np.array([0, 1])),
     ],
 )
 def test_gradient(model, guide, params, data):

Diff for: test/infer/test_hmc_util.py

@@ -125,8 +125,8 @@ def register_fn(model):
             num_steps=100,
             q_i={"x": 0.0},
             p_i={"x": 1.0},
-            q_f={"x": jnp.sin(1.0)},
-            p_f={"x": jnp.cos(1.0)},
+            q_f={"x": np.sin(1.0)},
+            p_f={"x": np.cos(1.0)},
             m_inv=np.array([1.0]),
             prec=1e-4,
         )