diff --git a/flax/nnx/extract.py b/flax/nnx/extract.py
index ae8b71ae8..ff8d97f2b 100644
--- a/flax/nnx/extract.py
+++ b/flax/nnx/extract.py
@@ -17,6 +17,7 @@
 import dataclasses
 import functools
 import inspect
+import types
 import typing as tp
 
 from flax import struct
@@ -1101,6 +1102,135 @@ def to_masked(tree, all_updates: list[Updates]):
       is_leaf=lambda x: x is None
   )
 
+
+def find_captured_nodes(f) -> list[object]:
+  """Find nnx Variables and graph nodes in f's closure chain.
+
+  Recursively searches through ``__closure__`` attributes at each level
+  and follows ``__wrapped__`` and ``functools.partial`` chains so that
+  Variables captured by decorators are also discovered.
+  """
+  seen: set[int] = set()
+  captured_nodes: list[object] = []
+
+  def _collect_from_pytree(obj):
+    for _, value in graphlib.iter_graph(obj, graph=True):
+      if isinstance(value, variablelib.Variable):
+        obj_id = id(value)
+        if obj_id not in seen:
+          seen.add(obj_id)
+          captured_nodes.append(value)
+
+  def _search(func):
+    if isinstance(func, functools.partial):
+      _collect_from_pytree((func.args, func.keywords))
+      _search(func.func)
+      return
+    if hasattr(func, '__closure__') and func.__closure__ is not None:
+      for cell in func.__closure__:
+        obj = cell.cell_contents
+        _collect_from_pytree(obj)
+    if hasattr(func, '__wrapped__'):
+      _search(func.__wrapped__)
+
+  _search(f)
+  return captured_nodes
+
+
+def _make_cell(val):
+  """Create a closure cell containing val."""
+  x = val
+  return (lambda: x).__closure__[0]  # type: ignore
+
+
+def replace_closure_cells(
+    f: tp.Callable,
+    captured_info: list[object],
+    replacements: tuple,
+) -> tp.Callable:
+  """Return a copy of f with captured closure cells replaced by replacements.
+
+  Uses identity matching to replace Variables at every level of the
+  ``__wrapped__`` / ``functools.partial`` chain, so that Variables captured
+  by intermediate decorators are also replaced.
+  """
+  if not captured_info:
+    return f
+  id_to_replacement = {
+      id(node): new_val
+      for node, new_val in zip(captured_info, replacements)
+  }
+  return _replace_closure_cells(f, id_to_replacement)
+
+
+def _replace_variables_in_pytree(
+    obj: tp.Any,
+    id_to_replacement: dict[int, tp.Any],
+) -> tp.Any:
+  """Replace Variables nested inside a pytree, returning the original if unchanged."""
+  is_leaf = lambda x: isinstance(x, variablelib.Variable)
+  def _replace(x):
+    if isinstance(x, variablelib.Variable) and id(x) in id_to_replacement:
+      return id_to_replacement[id(x)]
+    return x
+  return jax.tree.map(_replace, obj, is_leaf=is_leaf)
+
+
+def _replace_closure_cells(
+    f: tp.Callable,
+    id_to_replacement: dict[int, tp.Any],
+) -> tp.Callable:
+  if isinstance(f, functools.partial):
+    new_func = _replace_closure_cells(f.func, id_to_replacement)
+    new_args = _replace_variables_in_pytree(f.args, id_to_replacement)
+    new_keywords = _replace_variables_in_pytree(f.keywords, id_to_replacement)
+    if new_func is f.func and new_args is f.args and new_keywords is f.keywords:
+      return f
+    return functools.partial(new_func, *new_args, **new_keywords)
+
+  cells: list | None = None
+
+  # Replace captured Variables in this function's closure.
+  if hasattr(f, '__closure__') and f.__closure__ is not None:
+    cells = list(f.__closure__)
+    for i, cell in enumerate(cells):
+      try:
+        obj = cell.cell_contents
+      except ValueError:
+        continue
+      if id(obj) in id_to_replacement:
+        cells[i] = _make_cell(id_to_replacement[id(obj)])
+      else:
+        # Search inside pytree cell contents for nested Variables.
+        new_obj = _replace_variables_in_pytree(obj, id_to_replacement)
+        if new_obj is not obj:
+          cells[i] = _make_cell(new_obj)
+
+  # Recurse into __wrapped__ and update the cell referencing the old inner.
+  if hasattr(f, '__wrapped__'):
+    old_inner = f.__wrapped__
+    new_inner = _replace_closure_cells(old_inner, id_to_replacement)
+    if new_inner is not old_inner:
+      if cells is None and hasattr(f, '__closure__') and f.__closure__ is not None:
+        cells = list(f.__closure__)
+      if cells is not None:
+        for i, cell in enumerate(cells):
+          try:
+            if cell.cell_contents is old_inner:
+              cells[i] = _make_cell(new_inner)
+              break
+          except ValueError:
+            continue
+
+  new_f = types.FunctionType(
+    f.__code__, f.__globals__, f.__name__,
+    f.__defaults__, tuple(cells) if cells else f.__closure__,
+  )
+  if hasattr(f, '__wrapped__'):
+    new_f.__wrapped__ = new_inner
+  return new_f
+
+
 def filter_kwargs(f, **kwargs):
   sig = inspect.signature(f)
   has_var_keyword = any(
diff --git a/flax/nnx/transforms/autodiff.py b/flax/nnx/transforms/autodiff.py
index 222f9e439..bca98cc23 100644
--- a/flax/nnx/transforms/autodiff.py
+++ b/flax/nnx/transforms/autodiff.py
@@ -72,16 +72,24 @@ class SimpleGradFn:
   f: tp.Callable[..., tp.Any]
   has_aux: bool
   graph: bool
+  captured_info: list[tp.Any]
 
   def __post_init__(self):
     functools.update_wrapper(self, self.f, updated=())
 
   @extract.treemap_copy_args
   def __call__(self, *args, **kwargs):
-    current, snapshot = extract.snapshot(labeled(args=args, kwargs=kwargs))
+    captured_args = kwargs.pop('__captures__', None)
+    current, snapshot = extract.snapshot(
+        labeled(args=args, kwargs=kwargs, captured_args=captured_args))
     if self.graph:
       args, kwargs = extract.from_tree2((args, kwargs))
-    out = self.f(*args, **kwargs)
+    if captured_args is not None:
+      f = extract.replace_closure_cells(
+          self.f, self.captured_info, captured_args)
+    else:
+      f = self.f
+    out = f(*args, **kwargs)
     if self.graph:
       out = extract.to_tree2(out)
     extract.check_no_aliases('grad', **current, out=out, check=['out'])
@@ -151,8 +159,10 @@ def _grad_general(
 
   if not graph or not graph_updates:
 
+    captured_info = extract.find_captured_nodes(f)
+
     gradded_fn = transform(
-        SimpleGradFn(f, has_aux, graph=graph),
+        SimpleGradFn(f, has_aux, graph=graph, captured_info=captured_info),
         argnums=argnums,  # type: ignore[arg-type]
         has_aux=True,
         holomorphic=holomorphic,
@@ -169,7 +179,11 @@ def tree_grad_wrapper(*args, **kwargs):
           (args, kwargs), prefix=(args_prefix, False),
         )
 
-      variables = extract.check_no_aliases('grad', args=args, kwargs=kwargs)
+      variables = extract.check_no_aliases(
+          'grad', args=args, kwargs=kwargs, captured_args=captured_info)
+
+      if captured_info:
+        kwargs = {**kwargs, '__captures__': captured_info}
 
       fn_out = gradded_fn(*args, **kwargs)
 
diff --git a/flax/nnx/transforms/compilation.py b/flax/nnx/transforms/compilation.py
index 1b55009a6..e9f8fd550 100644
--- a/flax/nnx/transforms/compilation.py
+++ b/flax/nnx/transforms/compilation.py
@@ -354,6 +354,41 @@ def jit(
       support shared ``Variable`` references or returning mutable array
       references from the jitted function.
 
+  .. note::
+    **Captured Variables.** When ``fun`` closes over NNX ``Variable``
+    objects (e.g. references from an enclosing scope), ``nnx.jit``
+    automatically discovers them by inspecting the function's
+    ``__closure__``, ``functools.partial`` args, and ``__wrapped__``
+    chains. These captured Variables are extracted before tracing and
+    re-injected afterward so that state updates are propagated correctly.
+
+    However, **Variables referenced via module-level (global) names are
+    not detected**. Python stores global references in
+    ``f.__globals__`` (the entire module namespace), and there is no
+    cheap way to know which globals a function actually reads without
+    bytecode analysis. If you need to use a module-level Variable
+    inside a jitted function, pass it as an explicit argument instead
+    of relying on the global reference::
+
+      # Bad – global Variable is invisible to nnx.jit:
+      model = nnx.Linear(2, 3, rngs=nnx.Rngs(0))
+
+      @nnx.jit
+      def f(x):
+        return model(x)  # model is a global, not captured
+
+      # Good – pass as an argument:
+      @nnx.jit
+      def f(model, x):
+        return model(x)
+
+      # Also good – capture via a closure:
+      def make_f(model):
+        @nnx.jit
+        def f(x):
+          return model(x)  # model is in __closure__
+        return f
+
   Returns:
     A wrapped version of ``fun``, set up for just-in-time compilation.
   """
@@ -426,18 +461,32 @@ class SimpleJitFn:
   donate_argnames: frozenset[str]
   graph: bool
   update_shardings: tuple[tp.Any, ...]
+  captured_info: list[tp.Any]
 
   def __post_init__(self):
     functools.update_wrapper(self, self.f, updated=())
+    # The captures list is always prepended as an extra first argument,
+    # so adjust the signature so inspect.signature (used by JAX for
+    # validation of static_argnums / donate_argnums) sees the physical
+    # call signature instead of following __wrapped__ to the original.
+    sig = inspect.signature(self.f)
+    captures_param = inspect.Parameter(
+        '__captures__', inspect.Parameter.POSITIONAL_ONLY)
+    self.__signature__ = sig.replace(
+        parameters=[captures_param, *sig.parameters.values()])
 
   @extract.treemap_copy_args
-  def __call__(self, *args, **kwargs):
+  def __call__(self, captured_args, *args, **kwargs):
+    captured_shardings = self.in_shardings[0] if isinstance(self.in_shardings, tuple) else None
+    user_shardings = self.in_shardings[1:] if isinstance(self.in_shardings, tuple) else self.in_shardings
     current, snapshot = extract.snapshot(
-        labeled(args=args, kwargs=kwargs)
+        labeled(captured_args=captured_args, args=args, kwargs=kwargs)
     )
     if self.graph:
-      args, kwargs = extract.from_tree2((args, kwargs))
-    out = self.f(*args, **kwargs)
+      captured_args, args, kwargs = extract.from_tree2((captured_args, args, kwargs))
+    f = extract.replace_closure_cells(
+          self.f, self.captured_info, captured_args)
+    out = f(*args, **kwargs)
     if self.graph:
       out = extract.to_tree2(out, prefix=self.out_shardings)
     extract.check_no_aliases('jit', **current, out=out, check=['out'])
@@ -445,12 +494,15 @@ def keep_fn(jax_path, prefix, c, s):
       if extract.variable_changed(c, s):
         return True
       arg_type, arg_key, *_ = graphlib.jax_to_nnx_path(jax_path)
+      if arg_type == 'captured_args':
+        return True
       if arg_type == 'args':
         return arg_key in self.donate_argnums
       else:  # arg_type == 'kwargs':
         return arg_key in self.donate_argnames
+    prefix = labeled(captured_args=captured_shardings, args=user_shardings, kwargs=None)
     updates = extract.get_updates(
-        current, snapshot, prefix=labeled(args=self.in_shardings, kwargs=None),
+        current, snapshot, prefix=prefix,
         known_prefixes=self.update_shardings, keep_fn=keep_fn
     )
     return out, updates
@@ -533,6 +585,9 @@ def __init__(
     self.partial_args = partial_args
     self.graph = graph
 
+    # Capture closure nodes once at construction time
+    self._captured_info = extract.find_captured_nodes(fun)
+
     resolved = _resolve_argnums(fun, static_argnums, static_argnames)
     if isinstance(in_shardings, (tuple, list)) and resolved:
       expanded = list(in_shardings)
@@ -542,6 +597,16 @@ def __init__(
     else:
       self.in_shardings = in_shardings
 
+    # Prepend None shardings for captured nodes (passed as a list).
+    # The captures list is always the first argument, even when empty.
+    jit_in_shardings: tp.Any
+    if isinstance(in_shardings, (tuple, list)):
+      jit_in_shardings = (None,) + tuple(in_shardings)
+    elif in_shardings is not None:
+      raise ValueError("Broadcasting a sharding across arguments is not supported in NNX.")
+    else:
+      jit_in_shardings = in_shardings
+
     donate_argnums_set = frozenset(
         (donate_argnums,) if isinstance(donate_argnums, int)
         else donate_argnums or ()
@@ -550,21 +615,48 @@ def __init__(
         (donate_argnames,) if isinstance(donate_argnames, str)
         else donate_argnames or ()
     )
+
+    # The captures list is always the first argument, so offset
+    # index-based jax.jit parameters so they still refer to the correct
+    # physical arguments.
+    jit_static_argnums: int | tp.Sequence[int] | None = _offset_argnums(static_argnums, 1)
+    # Always donate the captured variables (arg 0) since we reassign
+    # the updated values back after the jit call.
+    offset_donate = _offset_argnums(donate_argnums, 1)
+    if offset_donate is None:
+      jit_donate_argnums: int | tp.Sequence[int] | None = (0,)
+    elif isinstance(offset_donate, int):
+      jit_donate_argnums = (0, offset_donate)
+    else:
+      jit_donate_argnums = (0, *offset_donate)
+
+    # Build the in_shardings for SimpleJitFn (used internally for
+    # snapshot/update prefixes). This needs both the captures prefix
+    # AND static arg expansion, unlike jit_in_shardings which only
+    # needs the captures prefix (JAX handles statics itself).
+    if isinstance(self.in_shardings, (tuple, list)):
+      fn_in_shardings: tp.Any = (None,) + tuple(self.in_shardings)
+    elif self.in_shardings is not None:
+      raise ValueError("Broadcasting in shardings across arguments is not supported in NNX.")
+    else: # self.in_shardings == None
+      fn_in_shardings = self.in_shardings
+
     self.jitted_fn = jax.jit(
         SimpleJitFn(
             fun,
-            self.in_shardings,
+            fn_in_shardings,
             out_shardings,
             donate_argnums_set,
             donate_argnames_set,
             graph,
             tuple(update_shardings),
+            self._captured_info,
         ),
-        in_shardings=in_shardings,
+        in_shardings=jit_in_shardings,
         out_shardings=(out_shardings, update_shardings),
-        static_argnums=static_argnums,
+        static_argnums=jit_static_argnums,
         static_argnames=static_argnames,
-        donate_argnums=donate_argnums,
+        donate_argnums=jit_donate_argnums,
         donate_argnames=donate_argnames,
         keep_unused=keep_unused,
         device=device,
@@ -575,7 +667,10 @@ def __init__(
   def _maybe_to_tree(self, args, kwargs):
     if self.graph:
       if self.in_shardings is not None and isinstance(self.in_shardings, (tuple, list)):
-        runtime_prefix = self.in_shardings[len(self.partial_args):]
+        # args[0] is always the captures list; skip partial_args in the
+        # user portion of shardings.
+        user_prefix = self.in_shardings[len(self.partial_args):]
+        runtime_prefix = (None, *user_prefix)
       else:
         runtime_prefix = self.in_shardings
 
@@ -594,10 +689,11 @@ def _maybe_from_tree(self, out):
     return out
 
   def __call__(self, *args: P.args, **kwargs: P.kwargs) -> R:
-    args = (*self.partial_args, *args)  # type: ignore[assignment]
-    args, kwargs = self._maybe_to_tree(args, kwargs)
-    variables = extract.check_no_aliases('jit', args=args, kwargs=kwargs)
-    out, updates = self.jitted_fn(*args, **kwargs)
+    all_args: tuple[tp.Any, ...] = (self._captured_info, *self.partial_args, *args)
+    all_args, kwargs = self._maybe_to_tree(all_args, kwargs)
+    variables = extract.check_no_aliases(
+        'jit', captured_args=all_args[0], args=all_args[1:], kwargs=kwargs)
+    out, updates = self.jitted_fn(*all_args, **kwargs)
     extract.apply_updates(variables, updates)
     return self._maybe_from_tree(out)
 
@@ -607,26 +703,29 @@ def __get__(self, obj, objtype=None):
     return functools.partial(self, obj)
 
   def eval_shape(self, *args, **kwargs):
-    args = (*self.partial_args, *args)
+    args = (self._captured_info, *self.partial_args, *args)
     args, kwargs = self._maybe_to_tree(args, kwargs)
     if not self.graph:
-      extract.check_no_aliases('jit', args=args, kwargs=kwargs)
+      extract.check_no_aliases(
+          'jit', captured_args=args[0], args=args[1:], kwargs=kwargs)
     out, _ = self.jitted_fn.eval_shape(*args, **kwargs)
     return self._maybe_from_tree(out)
 
   def trace(self, *args, **kwargs):
-    args = (*self.partial_args, *args)
+    args = (self._captured_info, *self.partial_args, *args)
     args, kwargs = self._maybe_to_tree(args, kwargs)
     if not self.graph:
-      extract.check_no_aliases('jit', args=args, kwargs=kwargs)
+      extract.check_no_aliases(
+          'jit', captured_args=args[0], args=args[1:], kwargs=kwargs)
     traced = self.jitted_fn.trace(*args, **kwargs)
     return SimpleTraced(traced, self)
 
   def lower(self, *args, **kwargs):
-    args = (*self.partial_args, *args)
+    args = (self._captured_info, *self.partial_args, *args)
     args, kwargs = self._maybe_to_tree(args, kwargs)
     if not self.graph:
-      extract.check_no_aliases('jit', args=args, kwargs=kwargs)
+      extract.check_no_aliases(
+          'jit', captured_args=args[0], args=args[1:], kwargs=kwargs)
     lowered = self.jitted_fn.lower(*args, **kwargs)
     return SimpleLowered(lowered, self)
 
@@ -1333,9 +1432,10 @@ def call(*args, **kwargs):
     raise NotImplementedError
 
   def __call__(self, *args, **kwargs):
-    args = (*self.jit_wrapped.partial_args, *args)
+    args = (self.jit_wrapped._captured_info, *self.jit_wrapped.partial_args, *args)
     args, kwargs = self.jit_wrapped._maybe_to_tree(args, kwargs)
-    variables = extract.check_no_aliases('jit', args=args, kwargs=kwargs)
+    variables = extract.check_no_aliases(
+        'jit', captured_args=args[0], args=args[1:], kwargs=kwargs)
     out, updates = self.compiled(*args, **kwargs)
     extract.apply_updates(variables, updates)
     return self.jit_wrapped._maybe_from_tree(out)
@@ -1916,6 +2016,18 @@ def shard_map_wrapper(*args, **kwargs):
   return shard_map_wrapper  # type: ignore
 
 
+def _offset_argnums(
+    argnums: int | tp.Sequence[int] | None,
+    offset: int,
+) -> int | tuple[int, ...] | None:
+  """Shift argnum indices by ``offset`` (e.g. when prepending captures)."""
+  if argnums is None:
+    return None
+  if isinstance(argnums, int):
+    return argnums + offset
+  return tuple(i + offset for i in argnums)
+
+
 # We can't use private methods from jax._src.api_util
 # We copy the function: api_util.fun_signature
 def _fun_signature(fun: tp.Callable) -> inspect.Signature | None:
diff --git a/flax/nnx/transforms/iteration.py b/flax/nnx/transforms/iteration.py
index 791265a11..d44f3b937 100644
--- a/flax/nnx/transforms/iteration.py
+++ b/flax/nnx/transforms/iteration.py
@@ -1412,13 +1412,14 @@ class SimpleScanFn:
   carry_idx: int | None
   carry_out_idx: int | None
   update_axes: tuple[tp.Any, ...]
+  captured_info: list[tp.Any]
 
   def __post_init__(self):
     functools.update_wrapper(self, self.f, updated=())
 
   @extract.treemap_copy_args
   def __call__(self, full_carry: tp.Any, args: tp.Any):
-    carry, broadcasts = full_carry
+    carry, broadcasts, captured_args = full_carry
     carry_in = extract.copy_var_structure(carry)
     args = extract.insert(args, broadcasts)
     current, snapshot = extract.snapshot(labeled(args=args))
@@ -1427,7 +1428,9 @@ def __call__(self, full_carry: tp.Any, args: tp.Any):
       carry = extract.from_tree2(carry)
     args = extract.replace_at(args, self.carry_idx, carry)
 
-    out = self.f(*args)
+    f = extract.replace_closure_cells(
+        self.f, self.captured_info, captured_args)
+    out = f(*args)
 
     if self.carry_idx is None:  # has carry
       carry_out = None
@@ -1454,7 +1457,7 @@ def __call__(self, full_carry: tp.Any, args: tp.Any):
         and extract.variable_changed(cur, snap),
     )
 
-    return (carry_out, broadcasts), (ys, updates)
+    return (carry_out, broadcasts, captured_args), (ys, updates)
 
 
 @tp.overload
@@ -1677,6 +1680,8 @@ def _simple_scan(
   carry_idx = extract.find(in_axes, Carry)
   carry_out_idx = extract.find(out_axes, Carry)
 
+  captured_info = extract.find_captured_nodes(f_unbound)
+
   simple_scan_fn = SimpleScanFn(
       f_unbound, graph=graph,
       in_axes=in_axes, out_axes=out_axes,
@@ -1684,6 +1689,7 @@ def _simple_scan(
       carry_idx=carry_idx,
       carry_out_idx=carry_out_idx,
       update_axes=tuple(updates_axes),
+      captured_info=captured_info,
   )
 
   @functools.wraps(f)
@@ -1700,7 +1706,8 @@ def simple_scan_wrapper(*args):
     if graph:
       args = extract.to_tree2(args, prefix=in_axes)
       carry = extract.to_tree2(carry)
-    variables = extract.check_no_aliases('scan', args=args, carry=carry)
+    variables = extract.check_no_aliases(
+        'scan', args=args, carry=carry, captured_args=captured_info)
     args, broadcasts = extract.extract(
         lambda _, axes, x: axes is None,
         in_axes, args,
@@ -1711,9 +1718,9 @@ def simple_scan_wrapper(*args):
         lambda ax, leaf: jnp.moveaxis(leaf, ax, 0),
         in_axes, args,
     )
-    (carry_out, final_broadcasts), (ys, updates) = jax.lax.scan(
+    (carry_out, final_broadcasts, captured_out), (ys, updates) = jax.lax.scan(
         simple_scan_fn,
-        (carry, broadcasts),
+        (carry, broadcasts, captured_info),
         args_t,
         length=length,
         reverse=reverse,
@@ -1731,6 +1738,13 @@ def simple_scan_wrapper(*args):
       if isinstance(broadcast, variablelib.Variable):
         broadcast.update_from_state(update)
 
+    # Write back captured variable values from the scan output.
+    # captured_out contains unflattened Variables (rebuilt by JAX pytree),
+    # so extract their raw values and write back to the originals.
+    for node, new_var in zip(captured_info, captured_out):
+      if isinstance(node, variablelib.Variable):
+        node.set_value(new_var.get_value())
+
     if graph:
       ys = extract.from_tree2(ys)
       carry = extract.from_tree2(carry)
diff --git a/tests/nnx/spmd_test.py b/tests/nnx/spmd_test.py
index 5f075ce77..5c7fde199 100644
--- a/tests/nnx/spmd_test.py
+++ b/tests/nnx/spmd_test.py
@@ -605,6 +605,30 @@ def test_get_abstract_no_sharding_metadata(self):
         getattr(abs_model.kernel.get_value(), 'sharding', None)
     )
 
+  def test_jit_closure_partition_spec_in_shardings(self):
+    """A PartitionSpec in_shardings applies to the arg but not captures."""
+    mesh = jax.make_mesh((4,), ('data',))
+    data_sharding = NamedSharding(mesh, P('data'))
+
+    with jax.set_mesh(mesh):
+      model = nnx.Linear(2, 3, rngs=nnx.Rngs(0))
+      observed_specs = []
+
+      def callback(sharding):
+        observed_specs.append(sharding.spec)
+
+      @nnx.jit(in_shardings=(P('data'),), graph=False)
+      def forward(x):
+        jax.debug.inspect_array_sharding(model.kernel[...], callback=callback)
+        jax.debug.inspect_array_sharding(x, callback=callback)
+        return model(x)
+
+      x = jax.device_put(jnp.ones((4, 2)), data_sharding)
+      forward(x)
+      self.assertEqual(observed_specs[0], P())
+      self.assertEqual(observed_specs[1], P('data'))
+
+
 def has_sharding_spec(array):
     sharding = array.sharding
     if hasattr(sharding, 'spec'):
diff --git a/tests/nnx/transforms_test.py b/tests/nnx/transforms_test.py
index b1d75bc73..36177f67a 100644
--- a/tests/nnx/transforms_test.py
+++ b/tests/nnx/transforms_test.py
@@ -16,7 +16,7 @@
 os.environ['XLA_FLAGS'] = '--xla_force_host_platform_device_count=4'
 
 import dataclasses
-from functools import partial
+from functools import partial, wraps
 import typing as tp
 
 from absl.testing import absltest
@@ -8017,5 +8017,277 @@ def forward_block(rng_state, rest_state, x):
     assert not jnp.allclose(y, y2)
 
 
+class TestClosureCapture(parameterized.TestCase):
+  """Tests for auto-capture of closure Variables in tree-mode transforms."""
+
+  def test_jit_closure_no_recompilation(self):
+    """nnx.jit with a captured variable does not recompile on repeated calls."""
+    model = nnx.Linear(2, 3, rngs=nnx.Rngs(0))
+
+    @nnx.jit(graph=False)
+    def forward(x):
+      return model(x)
+
+    x = jnp.ones((4, 2))
+    forward(x)
+    cache_size_after_first = forward.jitted_fn._cache_size()
+    model.kernel[...] = jnp.zeros((2,3))
+    forward(x)
+    cache_size_after_second = forward.jitted_fn._cache_size()
+
+    self.assertEqual(cache_size_after_first, cache_size_after_second,
+                     'nnx.jit recompiled on the second call with a closure capture')
+
+  def test_jit_closure_static_argnums(self):
+    """static_argnums indices refer to user args, not shifted by captures."""
+    model = nnx.Linear(2, 3, rngs=nnx.Rngs(0))
+
+    @nnx.jit(graph=False, static_argnums=1)
+    def forward(x, mode):
+      if mode == 'train':
+        return model(x)
+      return model(x) * 0
+
+    self.assertTrue(forward._captured_info, 'expected non-empty captures')
+
+    x = jnp.ones((4, 2))
+    y_train = forward(x, 'train')
+    y_eval = forward(x, 'eval')
+    np.testing.assert_array_equal(y_eval, jnp.zeros_like(y_eval))
+    self.assertFalse(jnp.allclose(y_train, jnp.zeros_like(y_train)))
+
+  def test_jit_closure_static_argnums_sequence(self):
+    """static_argnums as a sequence is offset correctly with captures."""
+    model = nnx.Linear(2, 3, rngs=nnx.Rngs(0))
+
+    def make_forward():
+      @nnx.jit(graph=False, static_argnums=(1, 2))
+      def forward(x, mode, scale):
+        out = model(x)
+        if mode == 'scale':
+          return out * scale
+        return out
+      return forward
+
+    forward = make_forward()
+    self.assertTrue(forward._captured_info, 'expected non-empty captures')
+
+    x = jnp.ones((4, 2))
+    y1 = forward(x, 'scale', 2.0)
+    y2 = forward(x, 'noscale', 1.0)
+    np.testing.assert_allclose(y1, y2 * 2.0)
+
+  def test_jit_closure_donates_captured_variables(self):
+    """Captured nnx.Variable buffers are donated to the jit call."""
+    model = nnx.Linear(2, 3, rngs=nnx.Rngs(0))
+
+    @nnx.jit(graph=False)
+    def forward(x):
+      return model(x)
+
+    self.assertTrue(forward._captured_info, 'expected non-empty captures')
+
+    x = jnp.ones((4, 2))
+
+    # Grab references to the raw JAX arrays backing the captured variables
+    old_kernel = model.kernel.get_value()
+    old_bias = model.bias.get_value()
+
+    _ = forward(x)
+
+    # After the jit call the original buffers should have been donated
+    # (invalidated), and the variables should hold new buffers.
+    self.assertTrue(old_kernel.is_deleted(),
+                    'captured kernel buffer was not donated')
+    self.assertTrue(old_bias.is_deleted(),
+                    'captured bias buffer was not donated')
+
+  @parameterized.parameters('jit', 'scan', 'grad')
+  def test_closure_state_update(self, transform):
+    """Mutations to captured Variables propagate back."""
+    count = nnx.Variable(jnp.array(0))
+
+    def body(x):
+      count[...] += 1
+      return x
+
+    if transform == 'jit':
+      forward = nnx.jit(body, graph=False)
+      forward(jnp.array(1.0))
+      self.assertEqual(count[...], 1)
+      forward(jnp.array(1.0))
+      self.assertEqual(count[...], 2)
+    elif transform == 'scan':
+      forward = nnx.scan(body, in_axes=0, out_axes=0, graph=False)
+      forward(jnp.ones((3,)))
+      self.assertEqual(count[...], 3)
+    elif transform == 'grad':
+      def loss(x):
+        count[...] += 1
+        return jnp.sum(x ** 2)
+      grad_fn = nnx.grad(loss, graph_updates=False)
+      grad_fn(jnp.array([1.0]))
+      self.assertEqual(count[...], 1)
+      grad_fn(jnp.array([1.0]))
+      self.assertEqual(count[...], 2)
+
+  @parameterized.parameters('jit', 'scan', 'grad')
+  def test_closure_pytree_state_update(self, transform):
+    """Mutations to captured pytrees of Variables propagate back."""
+    count = [nnx.Variable(jnp.array(0))]
+
+    def body(x):
+      count[0][...] += 1
+      return x * 2
+
+    if transform == 'jit':
+      forward = nnx.jit(body, graph=False)
+      forward(jnp.array(1.0))
+      self.assertEqual(count[0][...], 1)
+      forward(jnp.array(1.0))
+      self.assertEqual(count[0][...], 2)
+    elif transform == 'scan':
+      forward = nnx.scan(body, in_axes=0, out_axes=0, graph=False)
+      forward(jnp.ones((3,)))
+      self.assertEqual(count[0][...], 3)
+    elif transform == 'grad':
+      def loss(x):
+        count[0][...] += 1
+        return jnp.sum(x ** 2)
+      grad_fn = nnx.grad(loss, graph_updates=False)
+      grad_fn(jnp.array([1.0]))
+      self.assertEqual(count[0][...], 1)
+      grad_fn(jnp.array([1.0]))
+      self.assertEqual(count[0][...], 2)
+
+  @parameterized.parameters('jit', 'scan', 'grad')
+  def test_double_closure(self, transform):
+    x = nnx.Variable(jnp.array(0))
+
+    def f(y):
+      def g():
+        x[...] += 1
+      g()
+    if transform == 'jit':
+      nnx.jit(f, graph_updates=False)(jnp.array(1.0))
+      self.assertEqual(x[...], 1)
+    elif transform == 'scan':
+      nnx.scan(f, in_axes=0, out_axes=0, graph=False)(jnp.ones((4,)))
+      self.assertEqual(x[...], 4)
+    elif transform == 'grad':
+      def loss(y):
+        def g():
+          x[...] += 1
+        g()
+        return jnp.sum(y ** 2)
+      nnx.grad(loss, graph_updates=False)(jnp.array([1.0]))
+      self.assertEqual(x[...], 1)
+
+  @parameterized.parameters('jit', 'scan', 'grad')
+  def test_closure_nested(self, transform):
+    """Captures through functools.partial work."""
+    count = nnx.Variable(jnp.array(0))
+
+    def forward(scale, x):
+      count[...] += 1
+      return x * scale
+
+    f = partial(forward, 2)
+    if transform == 'jit':
+      nnx.jit(f, graph_updates=False)(jnp.array(1.0))
+      self.assertEqual(count[...], 1)
+    elif transform == 'scan':
+      nnx.scan(f, in_axes=0, out_axes=0, graph=False)(jnp.ones((4,)))
+      self.assertEqual(count[...], 4)
+    elif transform == 'grad':
+      def loss(scale, x):
+        count[...] += 1
+        return jnp.sum(x * scale)
+      g = partial(loss, 2.0)
+      nnx.grad(g, graph_updates=False)(jnp.array([1.0]))
+      self.assertEqual(count[...], 1)
+
+  @parameterized.parameters('jit', 'scan', 'grad')
+  def test_closure_in_decorator(self, transform):
+    """Captures through @wraps decorator chains work."""
+    count = nnx.Variable(jnp.array(0))
+
+    def my_decorator(f):
+      @wraps(f)
+      def wrapper(x):
+        count[...] += 1
+        return f(x)
+      return wrapper
+
+    def silly(x):
+      count[...] += 1
+      return jnp.sum(x ** 2)
+
+    decorated = my_decorator(silly)
+    if transform == 'jit':
+      nnx.jit(decorated, graph_updates=False)(jnp.array(1.0))
+      self.assertEqual(count[...], 2)
+    elif transform == 'scan':
+      nnx.scan(decorated, in_axes=0, out_axes=0, graph=False)(jnp.ones((3,)))
+      self.assertEqual(count[...], 6)
+    elif transform == 'grad':
+      nnx.grad(decorated, graph_updates=False)(jnp.array([1.0]))
+      self.assertEqual(count[...], 2)
+
+  @parameterized.parameters('jit', 'scan', 'grad')
+  def test_closure_duplicate_error_mentions_captured_args(self, transform):
+    """Error for duplicate Variable mentions 'captured_args' in the path."""
+    count = nnx.Variable(jnp.array(0))
+
+    def body(x):
+      count[...] += 1
+      return x
+
+    if transform == 'jit':
+      forward = nnx.jit(body, graph=False)
+    elif transform == 'scan':
+      forward = nnx.scan(body, in_axes=0, out_axes=0, graph=False)
+    elif transform == 'grad':
+      def loss(x):
+        count[...] += 1
+        return jnp.sum(x ** 2)
+      forward = nnx.grad(loss, graph_updates=False)
+
+    with self.assertRaisesRegex(ValueError, 'captured_args'):
+      forward(count)
+
+  @parameterized.parameters('jit', 'scan')
+  def test_closure_duplicate_error_no_captured_args_without_captures(self, transform):
+    """Without captures, duplicate error shows 'args' paths, not 'captured_args'."""
+    v = nnx.Variable(jnp.array(0))
+
+    def body(x, y):
+      return x[...] + y[...]
+
+    if transform == 'jit':
+      forward = nnx.jit(body, graph=False)
+    elif transform == 'scan':
+      forward = nnx.scan(body, in_axes=(0, 0), out_axes=0, graph=False)
+
+    with self.assertRaisesRegex(ValueError, r"args\[") as cm:
+      forward(v, v)
+    self.assertNotIn('captured_args', str(cm.exception))
+
+  def test_double_transform_closure(self):
+     x = nnx.Variable(jnp.array(0))
+
+     @nnx.jit(graph_updates=False)
+     def f():
+       x[...] += 1
+       @nnx.jit(graph_updates=False)
+       def g():
+         x[...] += 1
+       g()
+     f()
+     self.assertEqual(x[...], 2)
+     f()
+     self.assertEqual(x[...], 4)
+
+
 if __name__ == '__main__':
   absltest.main()