Repository: google/flax

Issue:
Add nnx.shard_map

Discussion:
review-notebook-app[bot]: Check out this pull request on&nbsp; <a href="https://app.reviewnb.com/google/flax/pull/4490"><img align="absmiddle"  alt="ReviewNB" height="28" class="BotMessageButtonImage" src="https://raw.githubusercontent.com/ReviewNB/support/master/images/button_reviewnb.png"/></a> 

 See visual diffs & provide feedback on Jupyter Notebooks. 

---

 <i>Powered by <a href='https://www.reviewnb.com/?utm_source=gh'>ReviewNB</a></i>

Patch:
diff --git a/docs_nnx/api_reference/flax.nnx/transforms.rst b/docs_nnx/api_reference/flax.nnx/transforms.rst
index 54ba3399..5b4440ed 100644
--- a/docs_nnx/api_reference/flax.nnx/transforms.rst
+++ b/docs_nnx/api_reference/flax.nnx/transforms.rst
@@ -15,6 +15,7 @@ transforms
 
 .. autofunction:: grad
 .. autofunction:: jit
+.. autofunction:: shard_map
 .. autofunction:: remat
 .. autofunction:: scan
 .. autofunction:: value_and_grad
diff --git a/flax/linen/spmd.py b/flax/linen/spmd.py
index 2e8dcec5..5226218e 100644
--- a/flax/linen/spmd.py
+++ b/flax/linen/spmd.py
@@ -63,17 +63,6 @@ class _UnassignedAxis:
 _unassigned_axis = _UnassignedAxis()
 
 
-def is_cpu_platform(mesh: jax.sharding.Mesh | None):
-  if mesh is None:
-    if _global_mesh_defined():
-      device = pxla.thread_resources.env.physical_mesh.devices.reshape(-1)[0]
-    else:
-      device = jax.devices()[0]
-  else:
-    device = mesh.devices.reshape(-1)[0]
-  return device.platform == 'cpu'
-
-
 def _mesh_assignment_free(new_assignment, existing_assignments):
   """Determines if a given mesh axis has already been assigned."""
   new = set(jax.tree_util.tree_leaves(new_assignment))
@@ -208,7 +197,9 @@ def _with_sharding_constraint(
   mesh: jax.sharding.Mesh | None = None,
 ):
   """Wrapper for lax.with_sharding_constraint, no-op on cpu or outside jit."""
-  if is_cpu_platform(mesh) or (not _global_mesh_defined() and mesh is None):
+  if jax.devices()[0].platform == 'cpu' or (
+    not _global_mesh_defined() and mesh is None
+  ):
     return x
   else:
     if mesh is not None and axis_resources is not None:
diff --git a/flax/nnx/__init__.py b/flax/nnx/__init__.py
index f059358b..20bd940e 100644
--- a/flax/nnx/__init__.py
+++ b/flax/nnx/__init__.py
@@ -139,6 +139,7 @@ from .transforms.autodiff import value_and_grad as value_and_grad
 from .transforms.autodiff import custom_vjp as custom_vjp
 from .transforms.autodiff import remat as remat
 from .transforms.compilation import jit as jit
+from .transforms.compilation import shard_map as shard_map
 from .transforms.compilation import StateSharding as StateSharding
 from .transforms.iteration import Carry as Carry
 from .transforms.iteration import scan as scan
diff --git a/flax/nnx/bridge/__init__.py b/flax/nnx/bridge/__init__.py
index 867ffae6..378455da 100644
--- a/flax/nnx/bridge/__init__.py
+++ b/flax/nnx/bridge/__init__.py
@@ -20,7 +20,6 @@ from .wrappers import lazy_init as lazy_init
 from .wrappers import ToLinen as ToLinen
 from .wrappers import to_linen as to_linen
 from .variables import NNXMeta as NNXMeta
-from .variables import with_partitioning as with_partitioning
 from .module import Module as Module
 from .module import Scope as Scope
 from .module import compact as compact
\ No newline at end of file
diff --git a/flax/nnx/bridge/module.py b/flax/nnx/bridge/module.py
index f9b12086..7d4253d9 100644
--- a/flax/nnx/bridge/module.py
+++ b/flax/nnx/bridge/module.py
@@ -24,14 +24,13 @@ import jax
 import typing_extensions as tpe
 
 from flax import errors
-from flax.core import meta
 from flax.core.frozen_dict import FrozenDict
 from flax.nnx import graph, rnglib, statelib, traversals
 from flax.nnx import variablelib
 import flax.nnx.module as nnx_module
 from flax.nnx.object import Object
 from flax.nnx import variablelib
-from flax.nnx.bridge import variables as bridge_variables
+from flax import errors
 import jax.numpy as jnp
 
 A = tp.TypeVar('A')
@@ -40,9 +39,8 @@ F = tp.TypeVar('F', bound=tp.Callable[..., tp.Any])
 
 
 @dataclasses.dataclass
-class ModuleStackEntry:
+class CompactContext:
   module: Module
-  in_compact: bool
   type_counter: defaultdict[type, int] = dataclasses.field(
     default_factory=lambda: defaultdict(int)
   )
@@ -50,14 +48,13 @@ class ModuleStackEntry:
 
 @dataclasses.dataclass
 class ModuleContext(threading.local):
-  module_stack: list[ModuleStackEntry | None] = dataclasses.field(
+  parent_stack: list[tp.Optional[CompactContext]] = dataclasses.field(
     default_factory=lambda: [None]
   )
 
 
 MODULE_CONTEXT = ModuleContext()
 
-
 class ModuleState(statelib.State):
   pass
 
@@ -78,21 +75,27 @@ def has_setup(x: tp.Any) -> tp.TypeGuard[_HasSetup]:
   return hasattr(x, 'setup')
 
 
+class _HasUnbox(tp.Protocol):
+  unbox: bool
+
+
+def has_unbox(x: tp.Any) -> tp.TypeGuard[_HasUnbox]:
+  return hasattr(x, 'unbox')
+
+
 def _maybe_call_setup(module: Module):
   if (
     has_setup(module)
-    and isinstance(module, Module)
+    and isinstance(module, Object)
     and not module._object__state.is_setup
   ):
     # void parent context
-    MODULE_CONTEXT.module_stack.append(
-      ModuleStackEntry(module, in_compact=False)
-    )
+    MODULE_CONTEXT.parent_stack.append(None)
     try:
       module.setup()  # type: ignore[attribute-error]
       module._object__state._is_setup = True
     finally:
-      MODULE_CONTEXT.module_stack.pop()
+      MODULE_CONTEXT.parent_stack.pop()
 
 
 def _bind_module(parent: Module, module: Module) -> Module:
@@ -119,35 +122,32 @@ class ModuleMeta(nnx_module.ModuleMeta):
 
 def _module_meta_call(cls: type[M], *args, **kwargs) -> M:
   # compact behavior
-  parent_ctx = MODULE_CONTEXT.module_stack[-1]
+  parent_ctx = MODULE_CONTEXT.parent_stack[-1]
   parent = None
   module: M
 
-  name = None
   if parent_ctx is not None:
-    if not parent_ctx.in_compact and 'name' in kwargs:
-      raise ValueError(
-        f"'name' can only be set in @compact functions. If in setup(), "
-          "use parent's `self.<attr_name> to set the submodule name.")
+    if 'parent' in kwargs:
+      parent = kwargs.pop('parent')
+      if parent is not None:
+        raise ValueError(
+          f"'parent' can only be set to None, got {type(parent).__name__}"
+        )
+      name = None
+    else:
+      type_index = parent_ctx.type_counter[cls]
+      parent_ctx.type_counter[cls] += 1
 
-    if parent_ctx.in_compact:
-      if 'parent' in kwargs:
-        parent = kwargs.pop('parent')
-        if parent is not None:
-          raise ValueError(
-            f"'parent' can only be set to None, got {type(parent).__name__}"
-          )
+      if 'name' in kwargs:
+        name = kwargs.pop('name')
+        if not isinstance(name, str):
+          raise ValueError(f"'name' must be a 'str', got {type(name).__name__}")
       else:
-        type_index = parent_ctx.type_counter[cls]
-        parent_ctx.type_counter[cls] += 1
+        name = f'{cls.__name__}_{type_index}'
 
-        if 'name' in kwargs:
-          name = kwargs.pop('name')
-          if not isinstance(name, str):
-            raise ValueError(f"'name' must be a 'str', got {type(name).__name__}")
-        else:
-          name = f'{cls.__name__}_{type_index}'
-        parent = parent_ctx.module
+      parent = parent_ctx.module
+  else:
+    name = None
 
   module = nnx_module.ModuleMeta.__call__(cls, *args, **kwargs)
   module.scope = None
@@ -173,8 +173,10 @@ class Module(nnx_module.Module, ModuleBase, metaclass=ModuleMeta):
     cls = dataclasses.dataclass(repr=False)(cls)
     cls.__hash__ = object.__hash__  # type: ignore[method-assign]
 
-  def __getattribute__(self, name: str):
-    return type(self)._getattr(self, name)
+  if not tp.TYPE_CHECKING:
+
+    def __getattribute__(self, name: str):
+      return type(self)._getattr(self, name)
 
   def _getattr(self, name: str) -> tp.Any:
     value = super().__getattribute__(name)
@@ -240,10 +242,16 @@ class Module(nnx_module.Module, ModuleBase, metaclass=ModuleMeta):
         abs_value.raw_value = value
         value = abs_value
 
-      variable = variablelib.Param(value)
+      variable = variablelib.Param(
+        value,
+        unbox=unbox,
+      )
     else:
       value = init_fn(self.make_rng('params'), *init_args, **init_kwargs)
-      variable = variablelib.Param(value)
+      variable = variablelib.Param(
+        value,
+        unbox=unbox,
+      )
 
     setattr(self, name, variable)
     return variable
@@ -291,13 +299,13 @@ class Module(nnx_module.Module, ModuleBase, metaclass=ModuleMeta):
         abs_value.raw_value = value
         value = abs_value
 
-      variable = variable_type(value)
+      variable = variable_type(value, unbox=unbox)
     else:
       if init_fn is None:
         raise ValueError(f"Expected 'init_fn' to be a callable, got None")
 
       value = init_fn(*init_args, **init_kwargs)
-      variable = variable_type(value)
+      variable = variable_type(value, unbox=unbox)
 
     setattr(self, name, variable)
     return variable
@@ -316,13 +324,15 @@ class Module(nnx_module.Module, ModuleBase, metaclass=ModuleMeta):
       if collection not in _variables:
         _variables[collection] = {}
 
-      if (
-        isinstance(variable_state, variablelib.VariableState)
-        and not variable_state._var_metadata
-      ):
-        leaf = variable_state.value
+      if has_unbox(variable_state):
+        assert isinstance(variable_state, variablelib.VariableState)
+        if variable_state.unbox:
+          leaf = variable_state.value
+        else:
+          # TODO: convert to linen box
+          leaf = variable_state
       else:
-        leaf = bridge_variables.to_linen_var(variable_state)
+        leaf = variable_state.value
 
       _variables[collection][path] = leaf
 
@@ -352,17 +362,23 @@ class Module(nnx_module.Module, ModuleBase, metaclass=ModuleMeta):
 
     # create variables
     real_variables = dict(variables)
-    for col_name, linen_collection in variables.items():
+    for collection, state in variables.items():
+      variable_type = variablelib.variable_type_from_name(
+        collection, allow_register=True
+      )
 
       def to_variable(value):
-        return bridge_variables.to_nnx_var(col_name, value)
-
-      linen_collection = jax.tree.map(
-        to_variable,
-        linen_collection,
-        is_leaf=lambda x: isinstance(x, meta.AxisMetadata),
-      )
-      real_variables[col_name] = linen_collection
+        if isinstance(value, variablelib.VariableState):
+          variable = value.to_variable()
+          if not isinstance(variable, variable_type):
+            raise ValueError(
+              f'Expected variable to be of type {variable_type}, got {type(variable).__name__}'
+            )
+          return variable
+        return variable_type(value)
+
+      state = jax.tree.map(to_variable, state)
+      real_variables[collection] = state
 
     states = ({},) if not real_variables else real_variables.values()
     state = ModuleState.merge(*states)
@@ -402,15 +418,12 @@ class Module(nnx_module.Module, ModuleBase, metaclass=ModuleMeta):
         value._object__state._initializing = _initialize
       if isinstance(value, Module):
         value.scope = Scope(rngs)
-        _maybe_call_setup(value)
+        if has_setup(value):
+          value.setup()
 
-    MODULE_CONTEXT.module_stack.append(
-      ModuleStackEntry(module, in_compact=False)
-    )
     try:
       out = _method(module, *args, **kwargs)
     finally:
-      MODULE_CONTEXT.module_stack.pop()
       # reset temporary state
       for _, value in graph.iter_graph(module):
         if isinstance(value, Object):
@@ -462,24 +475,21 @@ class Module(nnx_module.Module, ModuleBase, metaclass=ModuleMeta):
       **kwargs,
     )
 
-  def is_initializing(self) -> bool:
-    return self._object__state._initializing
-
 
 def compact(f: F) -> F:
   @functools.wraps(f)
   def compact_wrapper(self, *args, **kwargs):
     if not isinstance(self, Module):
       raise ValueError(
-        f"Expected 'self' to be a nnx.bridge.Module, got {type(self).__name__}"
+        f"Expected 'self' to be a nnx.compat.Module, got {type(self).__name__}"
       )
 
-    MODULE_CONTEXT.module_stack.append(ModuleStackEntry(self, in_compact=True))
+    MODULE_CONTEXT.parent_stack.append(CompactContext(self))
 
     try:
       return f(self, *args, **kwargs)
     finally:
-      MODULE_CONTEXT.module_stack.pop()
+      MODULE_CONTEXT.parent_stack.pop()
 
   return compact_wrapper  # type: ignore
 
diff --git a/flax/nnx/bridge/variables.py b/flax/nnx/bridge/variables.py
index 7f3a7aea..d515dcf8 100644
--- a/flax/nnx/bridge/variables.py
+++ b/flax/nnx/bridge/variables.py
@@ -14,16 +14,14 @@
 
 from collections import defaultdict
 from typing import Any, TypeVar
-import typing as tp
 
 import jax
 from flax import struct
 from flax.core import meta
-from flax.nnx import graph
 from flax.nnx import spmd
 from flax.nnx import traversals
 from flax.nnx import variablelib
-from flax.typing import LogicalNames
+import typing as tp
 
 
 A = TypeVar('A')
@@ -154,21 +152,8 @@ def nnx_attrs_to_linen_vars(nnx_attrs: dict) -> dict:
     elif isinstance(v, variablelib.VariableState):
       col_name = variablelib.variable_name_from_type(v.type)
       v = to_linen_var(v)
-    elif isinstance(v, graph.NodeDef) or isinstance(v, graph.NodeRef):
-      col_name = 'nnx'  # an nnx.GraphDef for some ToLinen submodule
     else:
-      raise ValueError(f'Cannot infer collection name from value: {v}')
+      col_name = 'nnx'  # it must be an nnx.GraphDef, for some ToLinen submodule
     linen_structured[(col_name, *kp)] = v
   variables = traversals.unflatten_mapping(linen_structured)
   return variables
-
-
-
-def with_partitioning(
-    fn: tp.Callable[..., tp.Any],
-    names: LogicalNames,
-    mesh: jax.sharding.Mesh | None = None,
-) -> tp.Callable[..., meta.Partitioned[tp.Any]]:
-  """Same interface as Linen, but calls NNX `with_partitioning` within."""
-  return spmd.with_partitioning(fn, names, mesh,
-                                linen_meta_type=meta.Partitioned)
\ No newline at end of file
diff --git a/flax/nnx/training/optimizer.py b/flax/nnx/training/optimizer.py
index cad0e7e8..5c017dae 100644
--- a/flax/nnx/training/optimizer.py
+++ b/flax/nnx/training/optimizer.py
@@ -13,8 +13,6 @@
 # limitations under the License.
 from __future__ import annotations
 
-import typing as tp
-
 import jax
 import jax.numpy as jnp
 import optax
@@ -25,8 +23,6 @@ from flax.nnx import variablelib
 from flax.nnx.object import Object
 from flax.nnx.variablelib import Variable, VariableState
 
-M = tp.TypeVar('M', bound=nnx.Module)
-
 # TODO: add tests and docstrings
 
 
@@ -105,7 +101,7 @@ def _update_opt_state(opt_state, updates):
   return jax.tree.map(optimizer_update_variables, opt_state, updates)
 
 
-class Optimizer(Object, tp.Generic[M]):
+class Optimizer(Object):
   """Simple train state for the common case with a single Optax optimizer.
 
   Example usage::
@@ -172,7 +168,7 @@ class Optimizer(Object, tp.Generic[M]):
 
   def __init__(
     self,
-    model: M,
+    model: nnx.Module,
     tx: optax.GradientTransformation,
     wrt: filterlib.Filter = nnx.Param,
   ):
diff --git a/flax/nnx/transforms/compilation.py b/flax/nnx/transforms/compilation.py
index 99d757ae..c064427c 100644
--- a/flax/nnx/transforms/compilation.py
+++ b/flax/nnx/transforms/compilation.py
@@ -17,6 +17,10 @@ import dataclasses
 import functools
 import typing as tp
 
+import jax.experimental
+import jax.experimental.shard_map
+from jax.sharding import PartitionSpec
+
 from flax.nnx import (
   extract,
   filterlib,
@@ -27,11 +31,13 @@ from flax.nnx import (
 import jax
 import jax.core
 import jax.stages
+from jax._src.mesh import Mesh, AbstractMesh
 
 from flax.typing import Missing
 
 F = tp.TypeVar('F', bound=tp.Callable[..., tp.Any])
-
+Specs = tp.Any
+AxisName = tp.Hashable
 
 # -------------------------------
 # jit
@@ -341,7 +347,6 @@ def jit(
         check_aliasing=in_shardings is not None or kwarg_shardings is not None,
         ctxtag=jit_wrapper,
       )
-      jax_in_shardings, kwarg_shardings, jax_out_shardings
       pure_args_out, pure_kwargs_out, pure_out = jitted_fn(
         *pure_args, **pure_kwargs
       )
@@ -371,3 +376,291 @@ def jit(
   jit_wrapper.inner = jitted_fn  # type: ignore
 
   return jit_wrapper  # type: ignore
+
+# -------------------------------
+# shard_map
+# -------------------------------
+
+# TODO: create StateSpec and consider enabling a mode that does
+# not use filters during split for performance. Overall there might
+# be performance limitations for using shard_map at a top-level
+
+@dataclasses.dataclass(eq=False)
+class ShardMapFn:
+  f: tp.Callable[..., tp.Any]
+  in_specs: tp.Any
+  out_specs: tp.Any
+  kwarg_specs: tp.Any
+  ctxtag: tp.Hashable
+
+  def __post_init__(self):
+    functools.update_wrapper(self, self.f)
+
+  def __call__(self, *pure_args, **pure_kwargs):
+    args, kwargs = extract.from_tree(
+      (pure_args, pure_kwargs),
+      merge_fn=_jit_merge_fn,
+      ctxtag=self.ctxtag,
+      is_inner=True,
+    )
+
+    out = self.f(*args, **kwargs)
+
+    args_out, kwargs_out = extract.clear_non_graph_nodes((args, kwargs))
+    pure_args_out, pure_kwargs_out, pure_out = extract.to_tree(
+      (args_out, kwargs_out, out),
+      prefix=(self.in_specs, self.kwarg_specs, self.out_specs),
+      ctxtag=self.ctxtag,
+      split_fn=_jit_split_fn,
+    )
+
+    return pure_args_out, pure_kwargs_out, pure_out
+
+
+@tp.overload
+def shard_map(
+  f: F,
+  *,
+  mesh: Mesh | AbstractMesh,
+  in_specs: Specs,
+  out_specs: Specs,
+  check_rep: bool = True,
+  auto: frozenset[AxisName] = frozenset(),
+) -> F: ...
+@tp.overload
+def shard_map(
+  *,
+  mesh: Mesh | AbstractMesh,
+  in_specs: Specs,
+  out_specs: Specs,
+  check_rep: bool = True,
+  auto: frozenset[AxisName] = frozenset(),
+) -> tp.Callable[[F], F]: ...
+def shard_map(
+  f: F | type[Missing] = Missing,
+  *,
+  mesh: Mesh | AbstractMesh,
+  in_specs: Specs,
+  out_specs: Specs,
+  check_rep: bool = True,
+  auto: frozenset[AxisName] = frozenset(),
+) -> F | tp.Callable[[F], F]:
+  """
+  Lifted version of
+  `jax.experimental.shard_map.shard_map <https://docs.jax.dev/en/latest/_autosummary/jax.experimental.shard_map.shard_map.html>`_
+  that can handle Modules / graph nodes as arguments.
+
+  Simple data parallel example::
+
+    import jax
+    import jax.numpy as jnp
+    from flax import nnx
+    from jax.sharding import PartitionSpec as P
+
+    mesh = jax.sharding.Mesh(jax.local_devices(), ('data',))
+
+    m = nnx.Linear(2, 3, rngs=nnx.Rngs(0))
+    x = jnp.ones((32, 2))
+
+    @nnx.shard_map(
+      mesh=mesh, in_specs=(P(None), P('data')), out_specs=P('data')
+    )
+    def f(m, x):
+      return m(x)
+
+    y = f(m, x)
+
+    jax.debug.visualize_array_sharding(y)
+
+  Notice that here we simply used some ``PartitionSpec`` to define the spec
+  the the whole model and data. This works for simple cases but if we need
+  to assign different ``PartitionSpec`` to different parts of the model we
+  need to use ``StateSharding`` and create some filters that allow us to target
+  specific parts of the model. Here's an example of how to do tensor parallelism
+  for a simple MLP block using ``StateSharding`` and filters::
+
+    mesh = jax.sharding.Mesh(jax.local_devices(), ('model',))
+
+    class MLP(nnx.Module):
+      def __init__(self, din, dhidden, dout, *, rngs: nnx.Rngs):
+        self.linear1 = nnx.Linear(din, dhidden, use_bias=False, rngs=rngs)
+        self.linear2 = nnx.Linear(dhidden, dout, use_bias=False, rngs=rngs)
+
+      def __call__(self, x):
+        return self.linear2(jax.nn.relu(self.linear1(x)))
+
+    m = MLP(2, 64, 3, rngs=nnx.Rngs(0))
+    x = jnp.ones((32, 2))
+
+    def path_ends_with(*path_suffix): # custom filter
+      return lambda path, value: path[-len(path_suffix):] == path_suffix
+
+    model_spec = nnx.StateSharding({
+      path_ends_with('linear1', 'kernel'): P(None, 'model'),
+      path_ends_with('linear2', 'kernel'): P('model', None),
+    })
+
+    @nnx.shard_map(mesh=mesh, in_specs=(model_spec, P(None)), out_specs=P(None))
+    def f(m, x):
+      y = m(x)
+      return jax.lax.psum(y, 'model')
+
+    y = f(m, x)
+
+    jax.debug.visualize_array_sharding(m.linear1.kernel.value)
+    jax.debug.visualize_array_sharding(m.linear2.kernel.value)
+
+
+  Alternatively, a ``State`` object with the exact PartitionSpec for each
+  state then you can be passed to ``StateSharding``::
+
+    mesh = jax.sharding.Mesh(jax.local_devices(), ('model',))
+
+    class MLP(nnx.Module):
+      def __init__(self, din, dhidden, dout, *, rngs: nnx.Rngs):
+        self.linear1 = nnx.Linear(din, dhidden, use_bias=False, rngs=rngs)
+        self.linear2 = nnx.Linear(dhidden, dout, use_bias=False, rngs=rngs)
+
+      def __call__(self, x):
+        return self.linear2(jax.nn.relu(self.linear1(x)))
+
+    m = MLP(2, 64, 3, rngs=nnx.Rngs(0))
+    x = jnp.ones((32, 2))
+
+    model_spec = nnx.State(
+      {
+        'linear1': {'kernel': P(None, 'model')},
+        'linear2': {'kernel': P('model', None)},
+      }
+    )
+
+    @nnx.shard_map(
+      mesh=mesh,
+      in_specs=(nnx.StateSharding(model_spec), P(None)),
+      out_specs=P(None),
+    )
+    def f(m, x):
+      y = m(x)
+      return jax.lax.psum(y, 'model')
+
+    y = f(m, x)
+
+    jax.debug.visualize_array_sharding(m.linear1.kernel.value)
+    jax.debug.visualize_array_sharding(m.linear2.kernel.value)
+
+  Here ``model_spec`` was created manually but you can also automate
+  this process by using ``nnx.get_partition_spec`` to automatically
+  create it for you (see
+  `Scale up on multiple devices <https://flax.readthedocs.io/en/latest/guides/flax_gspmd.html>`_
+  ).
+
+  Args:
+    f: callable to be mapped. Each application of ``f``, or "instance" of ``f``,
+      takes as input a shard of the mapped-over arguments and produces a shard
+      of the output.
+    mesh: a ``jax.sharding.Mesh`` representing the array of devices over which
+      to shard the data and on which to execute instances of ``f``. The names of
+      the ``Mesh`` can be used in collective communication operations in ``f``.
+      This is typically created by a utility function like
+      :func:`jax.experimental.mesh_utils.create_device_mesh`.
+    in_specs: a pytree with ``jax.sharding.PartitionSpec``or ``nnx.StateSharding``
+      (mapping substates to ``PartitionSpec``s) instances as leaves,
+      with a tree structure that is a tree prefix of the
+      args tuple to be mapped over. Similar to ``jax.sharding.NamedSharding``,
+      each ``PartitionSpec`` represents how the corresponding argument (or subtree
+      of arguments) should be sharded along the named axes of ``mesh``. In each
+      ``PartitionSpec``, mentioning a ``mesh`` axis name at a position expresses sharding
+      the corresponding argument array axis along that positional axis; not
+      mentioning an axis name expresses replication. If an argument, or argument
+      subtree, has a corresponding spec of None, that argument is not sharded.
+    out_specs: a pytree with ``jax.sharding.PartitionSpec`` or ``nnx.StateSharding``
+      (mapping substates to ``PartitionSpec``s) instances as leaves, with a tree structure
+      that is a tree prefix of the output of ``f``.
+      Each ``PartitionSpec`` represents how the corresponding output shards should be
+      concatenated. In each ``PartitionSpec``, metioning a ``mesh`` axis name at
+      a position expresses concatenation of that mesh axis's shards along the
+      corresponding positional axis. Not mentioning a ``mesh`` axis name
+      expresses a promise that the output values are equal along that mesh axis,
+      and that rather than concatenating only a single value should be produced.
+    check_rep: If True (default) enable additional validity checks and automatic
+      differentiation optimizations. The validity checks concern whether any mesh
+      axis names not mentioned in ``out_specs`` are consistent with how the outputs
+      of ``f`` are replicated. Must be set False if using a Pallas kernel in ``f``.
+    auto: (experimental) an optional set of axis names from ``mesh`` over which we
+      do not shard the data or map the function, but rather we allow the
+      compiler to control sharding. These names cannot be used in ``in_specs``,
+      ``out_specs``, or in communication collectives in ``f``.
+
+  Returns:
+    A callable that applies the input function ``f`` across data sharded according to
+    the ``mesh`` and ``in_specs``.
+  """
+  if f is Missing:
+    return functools.partial(
+      shard_map,
+      mesh=mesh,
+      in_specs=in_specs,
+      out_specs=out_specs,
+      check_rep=check_rep,
+      auto=auto,
+    )  # type: ignore[return-value]
+  assert not isinstance(f, type)
+
+  kwarg_specs = PartitionSpec()
+  jax_in_specs = jax.tree.map(
+    lambda x: extract.NodeStates(
+      _graphdef=PartitionSpec(),  # type: ignore[arg-type]
+      states=x.shardings,
+      metadata=x,
+    )
+    if isinstance(x, StateSharding)
+    else x,
+    in_specs,
+  )
+  jax_out_specs = jax.tree.map(
+    lambda x: extract.NodeStates(
+      _graphdef=PartitionSpec(),  # type: ignore[arg-type]
+      states=x.shardings,
+      metadata=x,
+    )
+    if isinstance(x, StateSharding)
+    else x,
+    out_specs,
+  )
+
+  @functools.wraps(f)
+  def shard_map_wrapper(*args, **kwargs):
+    # run dynamic_cache_context before update_context
+    with graph.update_context(shard_map_wrapper):
+      pure_args, pure_kwargs = extract.to_tree(
+        (args, kwargs),
+        prefix=(in_specs, kwarg_specs)
+        if in_specs is not None or kwarg_specs is not None
+        else None,
+        split_fn=_jit_split_fn,
+        check_aliasing=in_specs is not None or kwarg_specs is not None,
+        ctxtag=shard_map_wrapper,
+      )
+      pure_args_out, pure_kwargs_out, pure_out = shard_map_fn(
+        *pure_args, **pure_kwargs
+      )
+      _args_out, _kwargs_out, out = extract.from_tree(
+        (pure_args_out, pure_kwargs_out, pure_out),
+        merge_fn=_jit_merge_fn,
+        is_inner=False,
+        ctxtag=shard_map_wrapper,
+      )
+    return out
+
+  shard_map_fn = jax.experimental.shard_map.shard_map(
+    ShardMapFn(f, in_specs, out_specs, kwarg_specs, shard_map_wrapper),
+    mesh=mesh,
+    in_specs=jax_in_specs,
+    out_specs=(jax_in_specs, kwarg_specs, jax_out_specs),  # type: ignore
+    check_rep=check_rep,
+    auto=auto,
+  )
+
+  shard_map_wrapper.inner = shard_map_fn  # type: ignore
+
+  return shard_map_wrapper  # type: ignore
\ No newline at end of file
diff --git a/uv.lock b/uv.lock
index 21c49710..3aaa836d 100644
--- a/uv.lock
+++ b/uv.lock
@@ -2258,7 +2258,7 @@ wheels = [
 
 [[package]]
 name = "orbax-checkpoint"
-version = "0.11.5"
+version = "0.11.6"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
     { name = "absl-py" },
@@ -2276,9 +2276,9 @@ dependencies = [
     { name = "tensorstore" },
     { name = "typing-extensions" },
 ]
-sdist = { url = "https://files.pythonhosted.org/packages/7e/b4/262439a3fe00064b53a5182c0149447304f5a2d5a328a92d64390c18d189/orbax_checkpoint-0.11.5.tar.gz", hash = "sha256:8331ff594980a241ba43eb59dd683e5b590b339cff32a7b72d78cb5a350030b4", size = 249258 }
+sdist = { url = "https://files.pythonhosted.org/packages/a8/66/1499f770885b6b42ed4ac839ac4482e681b3a60f96dcd935a32221c165b3/orbax_checkpoint-0.11.6.tar.gz", hash = "sha256:e16a8bbabe7bc0c94f611d115b2b7790183e6847152804a261048160b81b9628", size = 258729 }
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/66/80/e659696b5b1c2ced427efedd2d9d29c1bc31d841ac8a031215aa38f6b2ae/orbax_checkpoint-0.11.5-py3-none-any.whl", hash = "sha256:b55a7a254ea0ab18237e8234a6ca8bf5522f589fcc2ac698cf6893d5e7ae3500", size = 342800 },
+    { url = "https://files.pythonhosted.org/packages/82/f5/6707aabc6d4928bdac3eadee00972616f41a83362a5b0441d63a93d81b75/orbax_checkpoint-0.11.6-py3-none-any.whl", hash = "sha256:fb208012e5d3601ee37b1100fe4331f9982b814df89f572749be9094fa499e1f", size = 361764 },
 ]
 
 [[package]]


Test patch:
diff --git a/tests/nnx/bridge/wrappers_test.py b/tests/nnx/bridge/wrappers_test.py
index d308c8c3..591442dd 100644
--- a/tests/nnx/bridge/wrappers_test.py
+++ b/tests/nnx/bridge/wrappers_test.py
@@ -16,19 +16,16 @@ import os
 from typing import Any
 
 from flax.linen.dtypes import promote_dtype
+os.environ["XLA_FLAGS"] = '--xla_force_host_platform_device_count=4'
 
-os.environ['XLA_FLAGS'] = '--xla_force_host_platform_device_count=4'
-
-import jax
-import jax.numpy as jnp
-import numpy as np
 from absl.testing import absltest
-
 import flax
 from flax import linen as nn
 from flax import nnx
 from flax.nnx import bridge
-from flax.nnx.bridge.module import MODULE_CONTEXT
+import jax
+import jax.numpy as jnp
+import numpy as np
 
 
 class TestCompatibility(absltest.TestCase):
@@ -502,25 +499,8 @@ class TestCompatModule(absltest.TestCase):
     state = {'b': {'c': nnx.Param(jnp.array(2))}}
     nnx.update(foo, state)
 
-  def test_module_stack(self):
-    """Test that apply set the module stack correctly."""
-    test = self
-
-    class Foo(bridge.Module):
-      def setup(self):
-        current_ctx = MODULE_CONTEXT.module_stack[-1]
-        test.assertIs(current_ctx.module, self)
-        test.assertFalse(current_ctx.in_compact)
-
-      def __call__(self):
-        current_ctx = MODULE_CONTEXT.module_stack[-1]
-        test.assertIs(current_ctx.module, self)
-        test.assertFalse(current_ctx.in_compact)
-
-    foo = Foo()
-    foo.apply({})
-
   def test_compact_basic(self):
+    test = self
     class Linear(bridge.Module):
       dout: int
 
@@ -540,11 +520,20 @@ class TestCompatModule(absltest.TestCase):
         din = x.shape[-1]
         self.linear = Linear(self.dout)
         x = self.linear(x)
+
+        # NNX
+        graphdef, state = nnx.split(self)
+        test.assertIn('Linear_0', state)
+        test.assertIn('w', state['Linear_0'])
+        test.assertIn('b', state['Linear_0'])
+
         return x
 
     foo = Foo(5)
     x = jnp.ones((3, 2))
 
+    self.assertIsInstance(foo, nnx.Module)
+
     variables = foo.init(0, x)
     params = variables['params']
 
@@ -640,14 +629,6 @@ class TestCompatModule(absltest.TestCase):
     y = bar.apply(variables, x)
     self.assertEqual(y.shape, (1, 5))
 
-    with self.assertRaises(ValueError):
-      class SetupBar(bridge.Module):
-        def setup(self):
-          self.xyz = Foo(5, name='xyz')
-        def __call__(self, x):
-          return self.xyz(x)
-      SetupBar().init(0, x)
-
   def test_dense_port(self):
     class Dense(bridge.Module):
       features: int
@@ -727,35 +708,6 @@ class TestCompatModule(absltest.TestCase):
 
     params = train_step(params, x, y)
 
-  def test_metadata(self):
-    class Linear(bridge.Module):
-      dout: int
-
-      @bridge.compact
-      def __call__(self, x):
-        w = self.param(
-          'w', bridge.with_partitioning(nnx.initializers.uniform(), ('in', 'out')),
-          (x.shape[-1], self.dout)
-        )
-        b = self.param('b', nnx.initializers.zeros_init(), (self.dout,))
-        return x @ w + b[None]
-
-    foo = Linear(5)
-    x = jnp.ones((3, 2))
-
-    variables = foo.init(0, x)
-    params = variables['params']
-    self.assertIsInstance(params['w'], nn.Partitioned)
-    self.assertEqual(params['w'].value.shape, (2, 5))
-    self.assertEqual(params['w'].names, ('in', 'out'))
-    self.assertEqual(nn.get_partition_spec(variables)['params']['w'],
-                     jax.sharding.PartitionSpec('in', 'out'))
-    self.assertIsInstance(params['b'], jax.Array)
-    self.assertEqual(params['b'].shape, (5,))
-
-    y: jax.Array = foo.apply(variables, x)
-    self.assertEqual(y.shape, (3, 5))
-
 
 if __name__ == '__main__':
   absltest.main()
diff --git a/tests/nnx/transforms_test.py b/tests/nnx/transforms_test.py
index 41208e88..d166e986 100644
--- a/tests/nnx/transforms_test.py
+++ b/tests/nnx/transforms_test.py
@@ -416,6 +416,129 @@ class TestJIT(absltest.TestCase):
     cached_m2 = cached_f(m)
     self.assertIs(cached_m, cached_m2)
 
+class TestShardMap(absltest.TestCase):
+  def test_basic_shardmap(self):
+    n_devices = jax.local_device_count()
+    devices = mesh_utils.create_device_mesh((n_devices,))
+    mesh = jax.sharding.Mesh(devices, ('a',))
+    PS = jax.sharding.PartitionSpec
+
+    state_sharding = nnx.StateSharding(
+      {
+        nnx.PathContains('kernel'): PS(None, 'a'),
+        nnx.PathContains('bias'): PS(),
+      }
+    )
+
+    m = nnx.Linear(16, 32, rngs=nnx.Rngs(0))
+
+    self.assertNotIsInstance(
+      m.kernel.value.sharding, jax.sharding.NamedSharding
+    )
+
+    @nnx.shard_map(mesh=mesh, in_specs=(state_sharding,), out_specs=None)
+    def f(m: nnx.Linear):
+      self.assertEqual(
+        m.kernel.value.shape, (m.in_features, m.out_features // n_devices)
+      )
+      self.assertEqual(m.bias.shape, (m.out_features,))
+
+    f(m)
+
+    self.assertIsInstance(m.kernel.value.sharding, jax.sharding.NamedSharding)
+
+  def test_from_state(self):
+    n_devices = jax.local_device_count()
+    devices = mesh_utils.create_device_mesh((n_devices,))
+    mesh = jax.sharding.Mesh(devices, ('a',))
+    PS = jax.sharding.PartitionSpec
+
+    state_spec = nnx.State(
+      {
+        'kernel': PS(None, 'a'),
+        'bias': PS(),
+      }
+    )
+    state_sharding = nnx.StateSharding(state_spec)
+
+    m = nnx.Linear(16, 32, rngs=nnx.Rngs(0))
+
+    self.assertNotIsInstance(
+      m.kernel.value.sharding, jax.sharding.NamedSharding
+    )
+
+    @nnx.shard_map(mesh=mesh, in_specs=(state_sharding,), out_specs=None)
+    def f(m: nnx.Linear):
+      self.assertEqual(
+        m.kernel.value.shape, (m.in_features, m.out_features // n_devices)
+      )
+      self.assertEqual(m.bias.shape, (m.out_features,))
+
+    f(m)
+
+    self.assertIsInstance(m.kernel.value.sharding, jax.sharding.NamedSharding)
+    self.assertIsInstance(m.bias.value.sharding, jax.sharding.NamedSharding)
+
+  def test_simple_data_parallel(self):
+    P = jax.sharding.PartitionSpec
+    n_devices = jax.local_device_count()
+
+    mesh = jax.sharding.Mesh(jax.local_devices(), ('data',))
+
+    m = nnx.Linear(2, 3, rngs=nnx.Rngs(0))
+    x = jnp.ones((32, 2))
+
+    @nnx.shard_map(
+      mesh=mesh, in_specs=(P(None), P('data')), out_specs=P('data')
+    )
+    def f(m, x):
+      self.assertEqual(x.shape, (32 // n_devices, 2))
+      return m(x)
+
+    y = f(m, x)
+
+    self.assertEqual(y.shape, (32, 3))
+    self.assertIsInstance(y.sharding, jax.sharding.NamedSharding)
+    self.assertIsInstance(m.kernel.value.sharding, jax.sharding.NamedSharding)
+    self.assertIsInstance(m.bias.value.sharding, jax.sharding.NamedSharding)
+
+  def test_simple_tensor_parallel(self):
+    P = jax.sharding.PartitionSpec
+
+    mesh = jax.sharding.Mesh(jax.local_devices(), ('model',))
+
+    class MLP(nnx.Module):
+      def __init__(self, din, dhidden, dout, *, rngs: nnx.Rngs):
+        self.linear1 = nnx.Linear(din, dhidden, use_bias=False, rngs=rngs)
+        self.linear2 = nnx.Linear(dhidden, dout, use_bias=False, rngs=rngs)
+
+      def __call__(self, x):
+        return self.linear2(jax.nn.relu(self.linear1(x)))
+
+    m = MLP(2, 64, 3, rngs=nnx.Rngs(0))
+    x = jnp.ones((32, 2))
+
+    def path_ends_with(path_suffix):
+      return lambda path, value: path[-len(path_suffix) :] == path_suffix
+
+    model_sharding = nnx.StateSharding(
+      {
+        path_ends_with(('linear1', 'kernel')): P(None, 'model'),
+        path_ends_with(('linear2', 'kernel')): P('model', None),
+      }
+    )
+
+    @nnx.shard_map(
+      mesh=mesh, in_specs=(model_sharding, P(None)), out_specs=P(None)
+    )
+    def f(m, x):
+      y = m(x)
+      return jax.lax.psum(y, 'model')
+
+    y = f(m, x)
+
+    jax.debug.visualize_array_sharding(m.linear1.kernel.value)
+
 
 class TestGrad(parameterized.TestCase):
   def test_grad(self):