Fix: Tracer leaks during NUTS sampling on random_nnx_module when NNX module contains mutable state (#2162)

* fix

* rm skip

Author: juanitorduz

numpyro/handlers.py

@@ -971,6 +971,8 @@ def process_message(self, msg: Message) -> None:
             value = self.substitute_fn(msg)
 
         if value is not None:
+            if msg["type"] == "mutable" and isinstance(value, dict):
+                value = value.copy()
             msg["value"] = value
 
 

test/contrib/test_module.py

@@ -435,6 +435,39 @@ def model(data, labels=None):
     assert f"nn{scope_divider}layers.1.bias" in samples
 
 
+@pytest.mark.parametrize("use_deterministic", [True, False])
+def test_random_nnx_module_mcmc_with_mutable_state(use_deterministic):
+    from flax import nnx
+
+    class NNXModel(nnx.Module):
+        def __init__(self):
+            self.linear = nnx.Linear(10, 1, rngs=nnx.Rngs(0))
+            self.mutable = nnx.Variable(0)
+
+        def __call__(self, x):
+            return self.linear(x)
+
+    nn_module = NNXModel()
+
+    def model(x, y=None):
+        random_model = random_nnx_module("model", nn_module, dist.Normal(0, 1))
+        pred = random_model(x)
+        with numpyro.plate("plate", size=x.shape[0]):
+            if use_deterministic:
+                pred = numpyro.deterministic("pred", pred)
+            numpyro.sample("obs", dist.Normal(pred, 1.0).to_event(1), obs=y)
+
+    x = jax.random.uniform(jax.random.key(0), shape=(10, 10))
+    y = jax.random.uniform(jax.random.key(0), shape=(10, 1))
+
+    mcmc = MCMC(NUTS(model), num_warmup=5, num_samples=5, progress_bar=False)
+    with jax.check_tracer_leaks(True):
+        mcmc.run(jax.random.key(0), x, y)
+    samples = mcmc.get_samples()
+    assert "model/linear.kernel" in samples
+    assert "model/linear.bias" in samples
+
+
 @pytest.mark.skipif(sys.version_info[:2] == (3, 9), reason="Skipping on Python 3.9")
 def test_eqx_module():
     import equinox as eqx