Ensure compatibility NNX 0.11.2 (#2067)

juanitorduz · web-flow · commit 68d86dc0535c · 2025-09-05T23:36:39.000-04:00
* ensure compatibility

* rm xfail
diff --git a/numpyro/contrib/module.py b/numpyro/contrib/module.py
@@ -498,7 +498,7 @@ def apply_fn(params, *call_args, **call_kwargs):
         if mutable_holder:
             nnx.replace_by_pure_dict(mutable_state, mutable_holder["state"])
 
-        model = nnx.merge(graph_def, params_state, mutable_state)
+        model = nnx.merge(graph_def, params_state, mutable_state, copy=True)
 
         model_call = model(*call_args, **call_kwargs)
 
diff --git a/test/contrib/test_module.py b/test/contrib/test_module.py
@@ -385,9 +385,6 @@ def nnx_model_eager(x, y):
 @pytest.mark.parametrize(
     argnames="batchnorm", argvalues=[True, False], ids=["batchnorm", "no_batchnorm"]
 )
-@pytest.mark.xfail(
-    reason="Temporary marking to pass CI. Bug fixed in https://github.com/pyro-ppl/numpyro/pull/2067"
-)
 def test_nnx_state_dropout_smoke(dropout, batchnorm):
     from flax import nnx
 

Original file line number	Diff line number	Diff line change
`@@ -385,9 +385,6 @@ def nnx_model_eager(x, y):`
`385`	`385`	`@pytest.mark.parametrize(`
`386`	`386`	`argnames="batchnorm", argvalues=[True, False], ids=["batchnorm", "no_batchnorm"]`
`387`	`387`	`)`
`388`		`-@pytest.mark.xfail(`
`389`		`- reason="Temporary marking to pass CI. Bug fixed in https://github.com/pyro-ppl/numpyro/pull/2067"`
`390`		`-)`
`391`	`388`	`def test_nnx_state_dropout_smoke(dropout, batchnorm):`
`392`	`389`	`from flax import nnx`
`393`	`390`