ensure compatibility

Author: juanitorduz

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)