fix nnx_wrapper.py comment

Author: Sharon Yu

src/maxtext/layers/nnx_wrappers.py

@@ -26,7 +26,6 @@
 from flax.core import FrozenDict
 from flax.core import meta
 from flax.nnx import graph
-from flax.nnx import tracers as nnx_tracers
 from flax.nnx import variablelib
 from flax.nnx.bridge import module as bdg_module
 from flax.nnx.module import Module
@@ -180,19 +179,6 @@ def is_linen_initializing() -> bool:
   return False
 
 
-def _refresh_variable_trace_state(module: Module) -> None:
-  """Resets stale ``_trace_state`` on Variables to unblock downstream ``nnx.split``.
-
-  ``nnx.update`` called with JAX tracer values uses ``_unsafe_bypass_check=True``,
-  which leaves Variables with a stale ``_trace_state`` from the outer Python
-  context and breaks ``nnx.split`` with "Cannot extract graph node from different
-  trace level". Resets ``_trace_state`` on any Variable whose ``_can_update`` is False.
-  """
-  for _, v in nnx.graph.iter_graph(module):
-    if isinstance(v, variablelib.Variable) and not v._can_update:  # pylint: disable=protected-access
-      object.__setattr__(v, "_trace_state", nnx_tracers.TraceState())
-
-
 class ToNNX(Module):
   """A wrapper to turn any Linen module into an NNX module.
 
@@ -505,8 +491,11 @@ def maybe_unbox(x):
       filtered_state_flat = {k: v for k, v in new_state_flat.items() if k not in unknown_state_flat}
       new_state = nnx.State(nnx.traversals.unflatten_mapping(filtered_state_flat))
 
-    nnx.update(module, new_state)
-    _refresh_variable_trace_state(module)
+    # Use split and merge to create a new module bound to the current trace level
+    # instead of using nnx.update which can leave stale tracers.
+    _, graphdef = module.split()
+    module = graphdef.merge(new_state)
+
     _fix_for_qwix_quantization(module)
     method_fn = _get_module_method(module, nnx_method)
     out = method_fn(module, *args, **kwargs)