fix(qwix): Remove explicit leaf assignment

Author: hsuan-lun-chiang

src/maxtext/layers/nnx_wrappers.py

@@ -496,7 +496,7 @@ def maybe_unbox(x):
 
     if unknown_state_flat:
       paths_str = ""
-      for path, value_state in unknown_state_flat.items():
+      for path, _ in unknown_state_flat.items():
         paths_str += f"\n  - {'/'.join(map(str, path))}"
 
         # Dynamically reconstruct the unknown variables
@@ -505,10 +505,8 @@ def maybe_unbox(x):
           if not hasattr(curr, p):
             setattr(curr, p, nnx.Module())
           curr = getattr(curr, p)
-        if not hasattr(curr, path[-1]):
-          setattr(curr, path[-1], value_state.type(value_state.value))
 
-      warnings.warn(f"Found unknown module paths in incoming state:{paths_str}. They have been dynamically reconstructed.")
+      warnings.warn(f"Found unknown module paths in incoming state:{paths_str}. Intermediate modules have been reconstructed.")
 
     nnx.update(module, new_state)
     _refresh_variable_trace_state(module)