Add tree-mode support for nnx.recursive_map, nnx.view, and nnx.view_info

Add a `graph` parameter to `nnx.recursive_map`, `nnx.view`, and
`nnx.view_info`. When `graph=False`, these functions use JAX's native
pytree traversal instead of Flax's graph protocol. Cycles and shared
Variable/Ref references are detected and raise errors in tree mode.
Added parametrized tests for both graph and tree modes.

PiperOrigin-RevId: 873225517

Author: Cristian Garcia

flax/nnx/graph.py

@@ -2907,35 +2907,37 @@ def _iter_graph(node: tp.Any, /) -> tp.Iterator[tuple[PathParts, tp.Any]]:
 
 
 def _iter_tree(node: tp.Any, /) -> tp.Iterator[tuple[PathParts, tp.Any]]:
-  seen_ids: set[int] = set()
+  in_progress: set[int] = set()
+  seen_refs: set[int] = set()
   stack: list[tuple[PathParts, tp.Any, bool]] = [((), node, False)]
   while stack:
     path, current, traversed = stack.pop()
 
     if traversed:
+      in_progress.discard(id(current))
       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:
+        if obj_id in seen_refs:
           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)
+        seen_refs.add(obj_id)
       yield path, current
       continue
 
     obj_id = id(current)
-    if obj_id in seen_ids:
+    if obj_id in in_progress:
       raise ValueError(
         f'Found cycle at path "{"/".join(map(str, path))}". '
         'Cycles are not supported with graph=False.'
       )
-    seen_ids.add(obj_id)
+    in_progress.add(obj_id)
 
     stack.append((path, current, True))
     children, _ = jax.tree_util.tree_flatten_with_path(
@@ -2946,7 +2948,13 @@ def _iter_tree(node: tp.Any, /) -> tp.Iterator[tuple[PathParts, tp.Any]]:
       stack.append(((*path, key), child, False))
 
 
-def recursive_map(f: tp.Callable[[PathParts, tp.Any], tp.Any], node: tp.Any, /):
+def recursive_map(
+  f: tp.Callable[[PathParts, tp.Any], tp.Any],
+  node: tp.Any,
+  /,
+  *,
+  graph: bool | None = None,
+):
   """Recursively applies a function to all nodes and leaves of the given graph node.
 
   Example::
@@ -2969,15 +2977,28 @@ def recursive_map(f: tp.Callable[[PathParts, tp.Any], tp.Any], node: tp.Any, /):
     Path = .conv     Conv
     Path = .lin      Linear
     Path = .         MyModule
+
+  Args:
+    f: A function that takes a path and a node and returns a new node.
+    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.
   """
-  node = clone(node, variables=False)
-  path_parts: PathParts = ()
-  visited: set[int] = set()
-  results: dict[int, tp.Any] = {}
-  return _recursive_map(f, node, path_parts, visited, results)
+  if graph is None:
+    graph = config.flax_nnx_graph_mode
+  if graph:
+    node = clone(node, variables=False)
+    path_parts: PathParts = ()
+    visited: set[int] = set()
+    results: dict[int, tp.Any] = {}
+    return _recursive_map_graph(f, node, path_parts, visited, results)
+  else:
+    return _recursive_map_tree(f, node)
 
 
-def _recursive_map(
+def _recursive_map_graph(
     f: tp.Callable[[PathParts, tp.Any], tp.Any],
     node: tp.Any,
     path: PathParts,
@@ -3002,7 +3023,7 @@ def _recursive_map(
     visited.add(node_id)
   if node_impl is not None:
     for key, value in node_impl.node_dict(node).items():
-      new_value = _recursive_map(f, value, (*path, key), visited, results)
+      new_value = _recursive_map_graph(f, value, (*path, key), visited, results)
       if new_value is not value:
         if node_impl.set_key is not None and value is not new_value:
           node_impl.set_key(node, key, new_value)
@@ -3017,6 +3038,52 @@ def _recursive_map(
   return new_node
 
 
+def _recursive_map_tree(
+    f: tp.Callable[[PathParts, tp.Any], tp.Any],
+    node: tp.Any,
+) -> tp.Any:
+  in_progress: set[int] = set()
+  seen_refs: set[int] = set()
+
+  def _recurse(path: PathParts, current: tp.Any) -> tp.Any:
+    if not is_pytree_node(current, check_graph_registry=False):
+      if isinstance(current, Variable) or is_array_ref(current):
+        obj_id = id(current)
+        if obj_id in seen_refs:
+          raise ValueError(
+            f'Found duplicate Variable or Ref at path '
+            f'"{"/".join(map(str, path))}". '
+            'Shared references are not supported with graph=False.'
+          )
+        seen_refs.add(obj_id)
+      return f(path, current)
+
+    obj_id = id(current)
+    if obj_id in in_progress:
+      raise ValueError(
+        f'Found cycle at path "{"/".join(map(str, path))}". '
+        'Cycles are not supported with graph=False.'
+      )
+    in_progress.add(obj_id)
+
+    children_with_path, treedef = jax.tree_util.tree_flatten_with_path(
+      current, is_leaf=lambda x: x is not current
+    )
+    new_children = []
+    for jax_key_path, child in children_with_path:
+      key = _key_path_to_key(jax_key_path[0])
+      new_child = _recurse((*path, key), child)
+      new_children.append(new_child)
+
+    new_node = treedef.unflatten(new_children)
+    result = f(path, new_node)
+
+    in_progress.discard(obj_id)
+    return result
+
+  return _recurse((), node)
+
+
 def find_duplicates(node: tp.Any, /, *, only: filterlib.Filter = ...) -> list[list[PathParts]]:
   """Finds duplicate nodes or node leaves in the given node.
 

flax/nnx/module.py

@@ -428,7 +428,7 @@ def eval(self, **attributes):
       raise_if_not_found=False,
     )
 
-def view(node: A, /, *, only: filterlib.Filter = ..., raise_if_not_found: bool = True,  **kwargs) -> A:
+def view(node: A, /, *, only: filterlib.Filter = ..., raise_if_not_found: bool = True, graph: bool | None = None, **kwargs) -> A:
   """Creates a new node with static attributes updated according to ``**kwargs``.
 
   The new node contains references to jax arrays in the original node. If a
@@ -462,6 +462,10 @@ def view(node: A, /, *, only: filterlib.Filter = ..., raise_if_not_found: bool =
   Args:
     node: the object to create a copy of.
     only: Filters to select the Modules to set the attributes of.
+    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.
     **kwargs: The attributes to set.
   """
   predicate = filterlib.to_predicate(only)
@@ -477,7 +481,7 @@ def _set_mode_fn(path, node):
         counts[k] += 1
     return node
 
-  out = graphlib.recursive_map(_set_mode_fn, node)
+  out = graphlib.recursive_map(_set_mode_fn, node, graph=graph)
 
   if raise_if_not_found:
     set_mode_calls = counts.pop("_set_mode_calls")
@@ -518,7 +522,7 @@ def _parse_docstring_args(doc_str: str) -> dict[str, str]:
 
 
 
-def view_info(node: Module, /, *, only: filterlib.Filter = ...) -> str:
+def view_info(node: Module, /, *, only: filterlib.Filter = ..., graph: bool | None = None) -> str:
   """Provides information about the ``view`` arguments for a module and all
   submodules. If no docstring is provided for a module's `set_view`, this function
   puts the `set_view` signature below the function.
@@ -554,6 +558,10 @@ def view_info(node: Module, /, *, only: filterlib.Filter = ...) -> str:
   Args:
     node: the object to display ``view`` information for.
     only: Filters to select the Modules to display information for.
+    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.
   """
   predicate = filterlib.to_predicate(only)
   classes: set[Module] = set()
@@ -563,7 +571,7 @@ def _set_mode_info_fn(path, node):
       classes.add(node.__class__)
     return node
 
-  graphlib.recursive_map(_set_mode_info_fn, node)
+  graphlib.recursive_map(_set_mode_info_fn, node, graph=graph)
 
   class_list = sorted(list(classes), key=lambda x: x.__qualname__)
   out_str = []

flax/nnx/rnglib.py

@@ -896,31 +896,31 @@ def _tree_split_rngs(
   only: filterlib.Filter = ...,
   squeeze: bool = False,
 ) -> tp.Any:
-  node = graphlib.clone(node, graph=False)
   predicate = filterlib.to_predicate(only)
-  for path, stream in graphlib.iter_graph(node, graph=False):
+
+  def _split_stream(path, node):
     if (
-      isinstance(stream, RngStream)
-      and predicate((*path, 'key'), stream.key)
-      and predicate((*path, 'count'), stream.count)
+      isinstance(node, RngStream)
+      and predicate((*path, 'key'), node.key)
+      and predicate((*path, 'count'), node.count)
     ):
-      key = stream()
-      key = random.split(key, splits)
+      key = random.split(node(), splits)
       if squeeze:
         key = key[0]
       if squeeze:
-        counts_shape = stream.count.shape
+        counts_shape = node.count.shape
       elif isinstance(splits, int):
-        counts_shape = (splits, *stream.count.shape)
+        counts_shape = (splits, *node.count.shape)
       else:
-        counts_shape = (*splits, *stream.count.shape)
+        counts_shape = (*splits, *node.count.shape)
 
-      stream.key = RngKey(key, tag=stream.tag)
-      stream.count = RngCount(
-        jnp.zeros(counts_shape, dtype=jnp.uint32), tag=stream.tag
+      node.key = RngKey(key, tag=node.tag)
+      node.count = RngCount(
+        jnp.zeros(counts_shape, dtype=jnp.uint32), tag=node.tag
       )
+    return node
 
-  return node
+  return graphlib.recursive_map(_split_stream, node, graph=False)
 
 @tp.overload
 def fork_rngs(

tests/nnx/graph_utils_test.py

@@ -1209,7 +1209,8 @@ def swap(path, node):
     self.assertEqual(bar2[1].d, -20)
     self.assertEqual(n, 2)
 
-  def test_recursive_map_with_list(self):
+  @parameterized.parameters(True, False)
+  def test_recursive_map_with_list(self, graph):
     rngs = nnx.Rngs(0)
     model = nnx.Sequential(nnx.Linear(2, 3, rngs=rngs), nnx.relu, nnx.Linear(3, 4, rngs=rngs))
 
@@ -1218,7 +1219,7 @@ def add_rank2_lora(_, node):
         return nnx.LoRA(node.in_features, 2, node.out_features, base_module=node, rngs=rngs)
       return node
 
-    self.assertEqual(len(nnx.recursive_map(add_rank2_lora, model).layers), 3)
+    self.assertEqual(len(nnx.recursive_map(add_rank2_lora, model, graph=graph).layers), 3)
 
   def test_graphdef_hash_with_sequential(self):
     rngs = nnx.Rngs(0)
@@ -1503,6 +1504,86 @@ def __init__(self, rngs):
     self.assertIn('Dropout', module_types)
     self.assertLen(modules, 3)
 
+  def test_recursive_map_tree_mode(self):
+    class Foo(nnx.Pytree):
+      def __init__(self, d):
+        self.d = d
+
+    foo1 = Foo(10)
+    foo2 = Foo(20)
+    bar = [foo1, foo2]
+    n = 0
+
+    def inc_d(path, node):
+      nonlocal n
+      if isinstance(node, Foo):
+        n += 1
+        node.d += 1
+      return node
+
+    bar2 = nnx.recursive_map(inc_d, bar, graph=False)
+    self.assertEqual(bar2[0].d, 11)
+    self.assertEqual(bar2[1].d, 21)
+    self.assertEqual(n, 2)
+
+  def test_recursive_map_tree_mode_replace(self):
+    class Foo(nnx.Pytree):
+      def __init__(self, d):
+        self.d = d
+
+    foo1 = Foo(10)
+    foo2 = Foo(20)
+    bar = [foo1, foo2]
+    n = 0
+
+    def swap(path, node):
+      nonlocal n
+      if isinstance(node, Foo):
+        n += 1
+        node = Foo(-node.d)
+      return node
+
+    bar2 = nnx.recursive_map(swap, bar, graph=False)
+    self.assertEqual(bar2[0].d, -10)
+    self.assertEqual(bar2[1].d, -20)
+    self.assertEqual(n, 2)
+
+  def test_recursive_map_tree_mode_with_list(self):
+    rngs = nnx.Rngs(0)
+    model = nnx.Sequential(
+      nnx.Linear(2, 3, rngs=rngs), nnx.relu, nnx.Linear(3, 4, rngs=rngs)
+    )
+
+    def add_rank2_lora(_, node):
+      if isinstance(node, nnx.Linear):
+        return nnx.LoRA(
+          node.in_features, 2, node.out_features,
+          base_module=node, rngs=rngs,
+        )
+      return node
+
+    result = nnx.recursive_map(add_rank2_lora, model, graph=False)
+    self.assertLen(result.layers, 3)
+
+  def test_recursive_map_tree_mode_shared_variable_raises(self):
+    v = nnx.Param(jnp.array(1))
+    g = [v, v]
+
+    with self.assertRaisesRegex(
+      ValueError, 'Shared references are not supported with graph=False'
+    ):
+      nnx.recursive_map(lambda path, node: node, g, graph=False)
+
+  def test_recursive_map_tree_mode_cycle_raises(self):
+    a = nnx.List([1])
+    b = nnx.List([2, a])
+    a.append(b)
+
+    with self.assertRaisesRegex(
+      ValueError, 'Cycles are not supported with graph=False'
+    ):
+      nnx.recursive_map(lambda path, node: node, a, graph=False)
+
 
 if __name__ == '__main__':
   absltest.main()