fix tracer leak in loss function with mutable state (#2001)

Author: fehiepsi

.github/workflows/ci.yml

@@ -39,9 +39,11 @@ jobs:
       run: |
         make lint
     - name: Build documentation
+      if: matrix.python-version != '3.9'
       run: |
         make docs
     - name: Test documentation
+      if: matrix.python-version != '3.9'
       run: |
         make doctest
         python -m doctest -v README.md
@@ -89,7 +91,6 @@ jobs:
         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

numpyro/infer/svi.py

@@ -51,7 +51,7 @@ def _make_loss_fn(
     def loss_fn(params):
         params = constrain_fn(params)
         if mutable_state is not None:
-            params.update(mutable_state)
+            params.update(jax.lax.stop_gradient(mutable_state))
             result = elbo.loss_with_mutable_state(
                 rng_key, params, model, guide, *args, **kwargs, **static_kwargs
             )

test/infer/test_svi.py

@@ -2,7 +2,6 @@
 # SPDX-License-Identifier: Apache-2.0
 
 from functools import partial
-import os
 
 import numpy as np
 from numpy.testing import assert_allclose
@@ -520,9 +519,6 @@ def guide():
 @pytest.mark.parametrize("stable_update", [True, False])
 @pytest.mark.parametrize("num_particles", [1, 10])
 @pytest.mark.parametrize("elbo", [Trace_ELBO, TraceMeanField_ELBO])
-@pytest.mark.xfail(
-    os.getenv("JAX_CHECK_TRACER_LEAKS") == "1", reason="Expected tracer leak"
-)
 def test_mutable_state(stable_update, num_particles, elbo):
     def model():
         x = numpyro.sample("x", dist.Normal(-1, 1))