add nnx.map and nnx.abstract_with_sharding

PiperOrigin-RevId: 886959302

Author: Cristian Garcia

flax/nnx/__init__.py

@@ -67,6 +67,7 @@
 from .graphlib import iter_graph as iter_graph
 from .graphlib import recursive_map as recursive_map
 from .graphlib import find_duplicates as find_duplicates
+from .graphlib import map as map
 from .graphlib import call as call
 from .graphlib import set_metadata as set_metadata
 from .graphlib import SplitContext as SplitContext
@@ -151,6 +152,7 @@
 from .spmd import get_named_sharding as get_named_sharding
 from .spmd import with_partitioning as with_partitioning
 from .spmd import get_abstract_model as get_abstract_model
+from .spmd import abstract_with_sharding as abstract_with_sharding
 from .statelib import FlatState as FlatState
 from .statelib import State as State
 from .statelib import to_flat_state as to_flat_state

flax/nnx/graphlib.py

@@ -19,6 +19,7 @@
 import functools
 import threading
 import typing as tp
+import builtins
 
 import jax.core
 
@@ -2526,6 +2527,48 @@ def state(
 variables = state
 
 
+def map(
+  f: tp.Callable[[tuple, tp.Any], tp.Any],
+  node: A,
+  /,
+  *,
+  graph: bool | None = None,
+) -> A:
+  """Map a function over the state of a graph node.
+
+  ``map`` extracts the state from ``node`` using :func:`split`, applies ``f``
+  to every ``(path, value)`` pair using :func:`map_state`, and returns a
+  new node with the mapped values merged back into the original structure.
+  Note that the leaves in the state are :class:`Variable` objects, so ``f``
+  should handle them accordingly.
+
+  Example usage::
+
+    >>> from flax import nnx
+    >>> import jax.numpy as jnp
+
+    >>> model = nnx.Linear(2, 3, rngs=nnx.Rngs(0))
+    >>> new_model = nnx.map(lambda path, v: v.replace(jnp.zeros_like(v)), model)
+    >>> assert jnp.all(new_model.kernel[...] == 0)
+    >>> assert jnp.all(new_model.bias[...] == 0)
+
+  Args:
+    f: A callable ``(path, value) -> new_value`` applied to each leaf in the
+      state. ``path`` is a tuple of path parts and ``value`` is the
+      corresponding leaf (typically a :class:`Variable`).
+    node: A graph node object.
+    graph: If ``True``, 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.
+  Returns:
+    A :class:`State` with the mapped values.
+  """
+  graphdef, state = split(node, graph=graph)
+  state = statelib.map_state(f, state)
+  return merge(graphdef, state)
+
+
 def graphdef(
   node: tp.Any, /, *, graph: bool | None = None,
 ) -> GraphDef[tp.Any]:
@@ -2695,7 +2738,7 @@ def _different_vars(path, x):
     duplicates_strs = '\n  ---'
     for node_duplicates in all_duplicates:
       for path in node_duplicates:
-        path_str = '/'.join(map(str, path))
+        path_str = '/'.join(builtins.map(str, path))
         duplicates_strs += f'\n  {path_str}'
       duplicates_strs += '\n  ---'
     raise ValueError(f'Found duplicate at paths:{duplicates_strs}')
@@ -2976,10 +3019,10 @@ def _iter_tree(node: tp.Any, /) -> tp.Iterator[tuple[PathParts, tp.Any]]:
       continue
 
     if not is_pytree_node(current, check_graph_registry=False):
-      _check_valid_pytree(current, 'iter_graph', '/'.join(map(str, path)))
+      _check_valid_pytree(current, 'iter_graph', '/'.join(builtins.map(str, path)))
       if isinstance(current, Variable) or variablelib.is_array_ref(current):
         obj_id = id(current)
-        str_path = '/'.join(map(str, path))
+        str_path = '/'.join(builtins.map(str, path))
         if obj_id in seen_refs:
           raise ValueError(
             f'Duplicate {current}\nfound at paths:\n\n'
@@ -2993,7 +3036,7 @@ def _iter_tree(node: tp.Any, /) -> tp.Iterator[tuple[PathParts, tp.Any]]:
       continue
 
     obj_id = id(current)
-    str_path = '/'.join(map(str, path))
+    str_path = '/'.join(builtins.map(str, path))
     if obj_id in in_progress:
       raise ValueError(
         f'Cycle detected for {type(current).__name__}\nfound at paths:\n\n'
@@ -3146,7 +3189,7 @@ def _recursive_map_graph(
   if node_id in visited:
     if node_id in results:
       return results[node_id]
-    path_str = '/'.join(map(str, path))
+    path_str = '/'.join(builtins.map(str, path))
     raise ValueError(
         f"Found cycle in the graph at path '{path_str}'. Node of type"
         f' {type(node)} has already been visited but has not been returned yet.'
@@ -3184,10 +3227,10 @@ def _recursive_map_tree(
 
   def _recurse(path: PathParts, current: tp.Any) -> tp.Any:
     if not is_pytree_node(current, check_graph_registry=False):
-      _check_valid_pytree(current, 'recursive_map', '/'.join(map(str, path)))
+      _check_valid_pytree(current, 'recursive_map', '/'.join(builtins.map(str, path)))
       if isinstance(current, Variable) or is_array_ref(current):
         obj_id = id(current)
-        str_path = '/'.join(map(str, path))
+        str_path = '/'.join(builtins.map(str, path))
         if obj_id in seen_refs:
           raise ValueError(
             f'Duplicate {current}\nfound at paths:\n\n'
@@ -3200,7 +3243,7 @@ def _recurse(path: PathParts, current: tp.Any) -> tp.Any:
       return f(path, current)
 
     obj_id = id(current)
-    str_path = '/'.join(map(str, path))
+    str_path = '/'.join(builtins.map(str, path))
     if obj_id in in_progress:
       raise ValueError(
         f'Cycle detected for {type(current).__name__}\nfound at paths:\n\n'

flax/nnx/spmd.py

@@ -179,4 +179,62 @@ def get_abstract_model(init_fn, mesh, *, graph: bool | None = None):
       lambda a, s: jax.ShapeDtypeStruct(a.shape, a.dtype, sharding=s),
       abs_state, get_named_sharding(abs_state, mesh)
     )
-  return gdef, abs_state
+  return gdef, abs_state
+
+
+def abstract_with_sharding(
+    tree: A, graph: bool | None = None
+) -> A:
+  """Add sharding information to abstract Variables.
+
+  When creating models with :func:`eval_shape`, Variables are abstract
+  (backed by ``jax.ShapeDtypeStruct``) and may not carry sharding
+  information, especially when using meshes with
+  :attr:`jax.sharding.AxisType.Auto` axes. ``abstract_with_sharding`` inspects each
+  Variable in ``tree`` and, if it has ``out_sharding`` metadata but no
+  sharding already set, attaches a :class:`jax.sharding.NamedSharding`
+  derived from the Variable's ``out_sharding`` and either its ``mesh``
+  metadata or the current abstract mesh (``jax.sharding.get_abstract_mesh``).
+
+  Example usage::
+
+    from flax import nnx
+    import jax
+
+    mesh = jax.make_mesh((2, 2), ('a', 'b'),
+      axis_types=(jax.sharding.AxisType.Auto,) * 2)
+    with jax.set_mesh(mesh):
+      abs_model = nnx.eval_shape(
+        lambda: nnx.Linear(4, 8, rngs=nnx.Rngs(0),
+          kernel_metadata={'out_sharding': ('a', 'b')}))
+      abs_model = nnx.abstract_with_sharding(abs_model)
+    assert abs_model.kernel.sharding.spec == jax.P('a', 'b')
+
+  Args:
+    tree: A graph node (e.g. an :class:`nnx.Module`) whose Variables should
+      be annotated with sharding (via ``out_sharding`` metadata).
+    graph: Forwarded to :func:`nnx.map`. If ``True``, uses graph-mode;
+      if ``False``, uses tree-mode.
+  Returns:
+    A tree with sharding-annotated ShapeDtypeStruct values inside Variables.
+  """
+  def add_sharding(_path, x):
+    if (
+        isinstance(x, variablelib.Variable)
+        and hasattr(value := x.get_value(), 'shape')
+        and hasattr(value, 'dtype')
+        and getattr(value, 'sharding', None) is None
+        and x.has_metadata('out_sharding')
+    ):
+      if x.has_metadata('mesh'):
+        mesh = x.get_metadata('mesh')
+      else:
+        mesh = jax.sharding.get_abstract_mesh()
+      specs = get_var_pspec(x)
+      sharding = jax.sharding.NamedSharding(mesh, specs)
+      abs_var = x.replace(
+          jax.ShapeDtypeStruct(value.shape, value.dtype, sharding=sharding)
+      )
+      return abs_var
+    return x
+  return graphlib.map(add_sharding, tree, graph=graph)

tests/nnx/graph_utils_test.py

@@ -1652,7 +1652,53 @@ def __init__(self):
     ):
       nnx.recursive_map(lambda path, node: node, node, graph=False)
 
+  @parameterized.parameters(True, False)
+  def test_map(self, graph):
+    model = nnx.Linear(2, 3, rngs=nnx.Rngs(0))
+    new_model = nnx.map(lambda path, x: jnp.zeros_like(x), model, graph=graph)
+
+    self.assertTrue(hasattr(new_model, 'kernel'))
+    self.assertTrue(hasattr(new_model, 'bias'))
+    np.testing.assert_array_equal(new_model.kernel, jnp.zeros((2, 3)))
+    np.testing.assert_array_equal(new_model.bias, jnp.zeros((3,)))
+
+  def test_map_with_path(self):
+    model = nnx.Linear(2, 3, rngs=nnx.Rngs(0))
+    paths_seen = []
+
+    def record_path(path, x):
+      paths_seen.append(path)
+      return x
+
+    nnx.map(record_path, model)
+    self.assertLen(paths_seen, 2)
+    path_last_parts = sorted(p[-1] for p in paths_seen)
+    self.assertEqual(path_last_parts, ['bias', 'kernel'])
+
+  def test_map_nested(self):
+    class Model(nnx.Module):
+      def __init__(self, rngs):
+        self.linear = nnx.Linear(2, 3, rngs=rngs)
+
+    model = Model(rngs=nnx.Rngs(0))
+    new_model = nnx.map(lambda path, x: jnp.ones_like(x), model)
+
+    self.assertTrue(hasattr(new_model, 'linear'))
+    np.testing.assert_array_equal(new_model.linear.kernel, jnp.ones((2, 3)))
+    np.testing.assert_array_equal(new_model.linear.bias, jnp.ones((3,)))
+
+  def test_map_replace(self):
+    model = nnx.Linear(2, 3, rngs=nnx.Rngs(0))
+    new_model = nnx.map(
+      lambda path, v: v.replace(jnp.zeros_like(v)), model
+    )
+
+    self.assertTrue(hasattr(new_model, 'kernel'))
+    self.assertTrue(hasattr(new_model, 'bias'))
+    self.assertIsInstance(new_model.kernel, nnx.Param)
+    np.testing.assert_array_equal(new_model.kernel[...], jnp.zeros((2, 3)))
+    np.testing.assert_array_equal(new_model.bias[...], jnp.zeros((3,)))
+
 
 if __name__ == '__main__':
   absltest.main()
-

tests/nnx/spmd_test.py

@@ -457,6 +457,73 @@ def test_variable_out_sharding_types(self, axis_type_name):
       v_format = nnx.Variable(value, out_sharding=Format(Layout(major_to_minor=(1, 0)), ns))
       self.assertEqual(v_format.sharding, ns)
 
+  def test_get_abstract_with_abstract_mesh(self):
+    mesh = jax.make_mesh(
+        (2, 2),
+        ('a', 'b'),
+        axis_types=(jax.sharding.AxisType.Auto,) * 2,
+    )
+    with jax.set_mesh(mesh):
+      abs_model = nnx.eval_shape(
+          lambda: nnx.Linear(
+              4,
+              8,
+              rngs=nnx.Rngs(0),
+              kernel_metadata={'out_sharding': ('a', 'b')},
+          )
+      )
+      abs_model = nnx.abstract_with_sharding(abs_model)
+
+    self.assertIsInstance(abs_model.kernel, nnx.Param)
+    self.assertEqual(abs_model.kernel.sharding.spec, P('a', 'b'))
+    self.assertEqual(
+        abs_model.kernel.sharding.mesh.axis_names,
+        mesh.axis_names,
+    )
+
+  def test_get_abstract_with_per_variable_mesh(self):
+    mesh1 = jax.make_mesh(
+        (2, 2),
+        ('a', 'b'),
+        axis_types=(jax.sharding.AxisType.Auto,) * 2,
+    )
+    mesh2 = jax.make_mesh(
+        (1, 4),
+        ('c', 'd'),
+        axis_types=(jax.sharding.AxisType.Auto,) * 2,
+    )
+
+    class Model(nnx.Module):
+      def __init__(self):
+        self.p1 = nnx.Linear(
+            4,
+            8,
+            rngs=nnx.Rngs(0),
+            kernel_metadata={'out_sharding': ('a', 'b'), 'mesh': mesh1},
+        )
+        self.p2 = nnx.Linear(
+            4,
+            8,
+            rngs=nnx.Rngs(0),
+            kernel_metadata={'out_sharding': ('c', 'd'), 'mesh': mesh2},
+        )
+
+    abs_model = nnx.eval_shape(lambda: Model())
+    abs_model = nnx.abstract_with_sharding(abs_model)
+
+    self.assertEqual(abs_model.p1.kernel.sharding.spec, P('a', 'b'))
+    self.assertEqual(abs_model.p1.kernel.sharding.mesh, mesh1)
+    self.assertEqual(abs_model.p2.kernel.sharding.spec, P('c', 'd'))
+    self.assertEqual(abs_model.p2.kernel.sharding.mesh, mesh2)
+
+  def test_get_abstract_no_sharding_metadata(self):
+    abs_model = nnx.eval_shape(lambda: nnx.Linear(4, 8, rngs=nnx.Rngs(0)))
+    abs_model = nnx.abstract_with_sharding(abs_model)
+
+    self.assertIsInstance(abs_model.kernel, nnx.Param)
+    self.assertIsNone(
+        getattr(abs_model.kernel.get_value(), 'sharding', None)
+    )
 
 def has_sharding_spec(array):
     sharding = array.sharding