Add tree-mode support to iter_graph

Add a graph: bool | None parameter to `iter_graph` and other downstream functions such as `iter_modules`, `split_rngs`, `fork_rngs`, `backup_keys`, and `reseed`.

PiperOrigin-RevId: 873211727

Author: Cristian Garcia

flax/nnx/bridge/module.py

@@ -104,7 +104,7 @@ def _maybe_call_setup(module: Module):
 def _bind_module(parent: Module, module: Module) -> Module:
   assert parent.scope is not None
 
-  for _, value in reversed(list(graph.iter_graph(module))):
+  for _, value in reversed(list(graph.iter_graph(module, graph=True))):
     if isinstance(value, Module):
       if module.scope is None:
         value.scope = parent.scope.copy()  # type: ignore[attribute-error]
@@ -471,7 +471,7 @@ def to_variable(value):
     _method = _get_unbound_fn(_method)
 
     # set temporary state
-    for _, value in graph.iter_graph(module):
+    for _, value in graph.iter_graph(module, graph=True):
       if isinstance(value, Pytree):
         value._pytree__state._initializing = _initialize
       if isinstance(value, Module):
@@ -486,7 +486,7 @@ def to_variable(value):
     finally:
       MODULE_CONTEXT.module_stack.pop()
       # reset temporary state
-      for _, value in graph.iter_graph(module):
+      for _, value in graph.iter_graph(module, graph=True):
         if isinstance(value, Pytree):
           value._pytree__state._initializing = False
         if isinstance(value, Module):

flax/nnx/bridge/wrappers.py

@@ -66,7 +66,7 @@ def _functional_constructor(*args: tp.Any, **kwargs: tp.Any) -> Functional[M]:
 
 
 def _set_initializing(module: Module, initializing: bool):
-  for _, value in graph.iter_graph(module):
+  for _, value in graph.iter_graph(module, graph=True):
     if isinstance(value, Pytree):
       value._pytree__state._initializing = initializing
 

flax/nnx/extract.py

@@ -57,7 +57,7 @@ def check_consistent_aliasing(
   node_id_to_variable: dict[int, tp.Any] = {}
 
   # collect all paths and prefixes for each node
-  for path, value in graph.iter_graph(node):
+  for path, value in graph.iter_graph(node, graph=True):
     if graph.is_graph_node(value) or isinstance(value, graph.Variable):
       if isinstance(value, Pytree):
         value._check_valid_context(

flax/nnx/graph.py

@@ -2834,7 +2834,9 @@ def _set_metadata(path: PathParts, variable: V) -> None:
   map_state(_set_metadata, state(node, only))
 
 
-def iter_graph(node: tp.Any, /) -> tp.Iterator[tuple[PathParts, tp.Any]]:
+def iter_graph(
+  node: tp.Any, /, *, graph: bool | None = None,
+) -> tp.Iterator[tuple[PathParts, tp.Any]]:
   """Iterates over all nested nodes and leaves of the given graph node, including the current node.
 
   ``iter_graph`` creates a generator that yields path and value pairs, where the
@@ -2866,22 +2868,35 @@ def iter_graph(node: tp.Any, /) -> tp.Iterator[tuple[PathParts, tp.Any]]:
     (0, 'w') Param
     (0,) Linear
     () list
+
+  Args:
+    node: A graph node object.
+    graph: If ``True`` (default), uses graph-mode which supports the full
+      NNX feature set including shared references. If ``False``, uses
+      tree-mode which treats Modules as regular JAX pytrees, avoiding
+      the overhead of the graph protocol.
   """
+  if graph is None:
+    graph = config.flax_nnx_graph_mode
+  if graph:
+    return _iter_graph(node)
+  else:
+    return _iter_tree(node)
+
+
+def _iter_graph(node: tp.Any, /) -> tp.Iterator[tuple[PathParts, tp.Any]]:
   visited: set[int] = set()
   stack: list[tuple[PathParts, tp.Any, bool]] = [((), node, False)]
   while stack:
-    # Yield if the node is either a leaf or has been traversed already.
     path_parts, node, traversed = stack.pop(-1)
     if traversed or not (is_node(node) or isinstance(node, Variable)):
       yield path_parts, node
       continue
 
-    # Skip if the node has been visited already.
     if id(node) in visited:
       continue
     visited.add(id(node))
 
-    # Traverse the node.
     if (node_impl := get_node_impl(node)) is None:
       yield path_parts, node
       continue
@@ -2891,6 +2906,46 @@ def iter_graph(node: tp.Any, /) -> tp.Iterator[tuple[PathParts, tp.Any]]:
       stack.append(((*path_parts, key), child, False))
 
 
+def _iter_tree(node: tp.Any, /) -> tp.Iterator[tuple[PathParts, tp.Any]]:
+  seen_ids: set[int] = set()
+  stack: list[tuple[PathParts, tp.Any, bool]] = [((), node, False)]
+  while stack:
+    path, current, traversed = stack.pop()
+
+    if traversed:
+      yield path, current
+      continue
+
+    if not is_pytree_node(current, check_graph_registry=False):
+      if isinstance(current, Variable) or variablelib.is_array_ref(current):
+        obj_id = id(current)
+        if obj_id in seen_ids:
+          raise ValueError(
+            f'Found duplicate Variable or Ref at path '
+            f'"{"/".join(map(str, path))}". '
+            'Shared references are not supported with graph=False.'
+          )
+        seen_ids.add(obj_id)
+      yield path, current
+      continue
+
+    obj_id = id(current)
+    if obj_id in seen_ids:
+      raise ValueError(
+        f'Found cycle at path "{"/".join(map(str, path))}". '
+        'Cycles are not supported with graph=False.'
+      )
+    seen_ids.add(obj_id)
+
+    stack.append((path, current, True))
+    children, _ = jax.tree_util.tree_flatten_with_path(
+      current, is_leaf=lambda x: x is not current
+    )
+    for jax_key_path, child in reversed(children):
+      key = _key_path_to_key(jax_key_path[0])
+      stack.append(((*path, key), child, False))
+
+
 def recursive_map(f: tp.Callable[[PathParts, tp.Any], tp.Any], node: tp.Any, /):
   """Recursively applies a function to all nodes and leaves of the given graph node.
 
@@ -3063,8 +3118,10 @@ class GenericPytree: ...
 from jax._src.tree_util import _registry as JAX_PYTREE_REGISTRY
 
 
-def is_pytree_node(x: tp.Any) -> bool:
-  if type(x) in GRAPH_REGISTRY:
+def is_pytree_node(
+  x: tp.Any, *, check_graph_registry: bool = True,
+) -> bool:
+  if check_graph_registry and type(x) in GRAPH_REGISTRY:
     return False
   elif isinstance(x, Variable):
     return False

flax/nnx/module.py

@@ -22,7 +22,7 @@
 
 from flax.nnx import (
   filterlib,
-  graph,
+  graph as graphlib,
 )
 from flax.nnx import variablelib as variableslib
 from flax.nnx.pytreelib import Pytree, PytreeMeta
@@ -477,7 +477,7 @@ def _set_mode_fn(path, node):
         counts[k] += 1
     return node
 
-  out = graph.recursive_map(_set_mode_fn, node)
+  out = graphlib.recursive_map(_set_mode_fn, node)
 
   if raise_if_not_found:
     set_mode_calls = counts.pop("_set_mode_calls")
@@ -563,7 +563,7 @@ def _set_mode_info_fn(path, node):
       classes.add(node.__class__)
     return node
 
-  graph.recursive_map(_set_mode_info_fn, node)
+  graphlib.recursive_map(_set_mode_info_fn, node)
 
   class_list = sorted(list(classes), key=lambda x: x.__qualname__)
   out_str = []
@@ -613,7 +613,9 @@ def first_from(*args: tp.Optional[A], error_msg: str) -> A:
       return arg
   raise ValueError(error_msg)
 
-def iter_modules(module: Module) -> tp.Iterator[tuple[PathParts, Module]]:
+def iter_modules(
+  module: Module, /, *, graph: bool | None = None,
+) -> tp.Iterator[tuple[PathParts, Module]]:
   """Recursively iterates over all nested :class:`Module`'s of the given Module, including
   the argument.
 
@@ -648,8 +650,15 @@ def iter_modules(module: Module) -> tp.Iterator[tuple[PathParts, Module]]:
     ('submodule', 'linear2') Linear
     ('submodule',) SubModule
     () Block
+
+  Args:
+    module: A :class:`Module` object.
+    graph: If ``True`` (default), uses graph-mode which supports the full
+      NNX feature set including shared references. If ``False``, uses
+      tree-mode which treats Modules as regular JAX pytrees, avoiding
+      the overhead of the graph protocol.
   """
-  for path, value in graph.iter_graph(module):
+  for path, value in graphlib.iter_graph(module, graph=graph):
     if isinstance(value, Module):
       yield path, value
 
@@ -687,7 +696,7 @@ def iter_children(module: Module) -> tp.Iterator[tuple[Key, Module]]:
     linear Linear
     submodule SubModule
   """
-  node_impl = graph.get_node_impl(module)
+  node_impl = graphlib.get_node_impl(module)
   assert node_impl is not None
   node_dict = node_impl.node_dict(module)
   for key, value in node_dict.items():