revert is_leaf logic in _check_carry_same_references

PiperOrigin-RevId: 799324513

Author: Cristian Garcia

flax/nnx/transforms/iteration.py

@@ -652,7 +652,11 @@ def check_carry_same_references(key_path, arg, out):
       )
 
   jax.tree_util.tree_map_with_path(
-    check_carry_same_references, carry_arg, carry_arg_out, is_leaf=graph.is_graph_node
+      check_carry_same_references,
+      carry_arg,
+      carry_arg_out,
+      is_leaf=lambda x: graph.is_graph_node(x)
+      and not isinstance(x, variablelib.Variable),
   )
 
 def _extract_graphdefs(
@@ -1308,9 +1312,6 @@ def scan_wrapper(*args, **kwargs):
   return scan_wrapper  # type: ignore
 
 
-
-
-
 # -------------------------------
 # while_loop
 # -------------------------------