Handle captures in jax transforms

Author: samanklesaria

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(

flax/nnx/transforms/compilation.py

@@ -426,31 +426,62 @@ class SimpleJitFn:
   donate_argnames: frozenset[str]
   graph: bool
   update_shardings: tuple[tp.Any, ...]
+  captured_info: tp.Optional[list[tp.Any]]
 
   def __post_init__(self):
     functools.update_wrapper(self, self.f, updated=())
+    # When captures are prepended as an extra argument, the original
+    # function's signature no longer matches the physical call signature.
+    # Remove __wrapped__ so inspect.signature (used by JAX for validation
+    # of static_argnums / donate_argnums) sees *args instead.
+    if self.captured_info:
+      delattr(self, '__wrapped__')
 
   @extract.treemap_copy_args
   def __call__(self, *args, **kwargs):
-    current, snapshot = extract.snapshot(
-        labeled(args=args, kwargs=kwargs)
-    )
+    if self.captured_info:
+      captured_args = args[0]
+      user_args = args[1:]
+      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(captured_args=captured_args, args=user_args, kwargs=kwargs)
+      )
+    else:
+      user_args = args
+      current, snapshot = extract.snapshot(
+          labeled(args=user_args, kwargs=kwargs)
+      )
     if self.graph:
       args, kwargs = extract.from_tree2((args, kwargs))
-    out = self.f(*args, **kwargs)
+    if self.captured_info:
+      captured_args = args[0]
+      user_args = args[1:]
+      f = extract.replace_closure_cells(
+          self.f, self.captured_info, captured_args)
+    else:
+      user_args = args
+      f = self.f
+    out = f(*user_args, **kwargs)
     if self.graph:
       out = extract.to_tree2(out, prefix=self.out_shardings)
     extract.check_no_aliases('jit', **current, out=out, check=['out'])
     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
+    if self.captured_info:
+      prefix = labeled(captured_args=captured_shardings, args=user_shardings, kwargs=None)
+    else:
+      prefix = labeled(args=self.in_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 +564,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 +576,25 @@ def __init__(
     else:
       self.in_shardings = in_shardings
 
+    # Prepend None shardings for captured nodes (passed as a list)
+    n_captured = len(self._captured_info)
+    jit_in_shardings: tp.Any
+    if n_captured > 0:
+      if isinstance(in_shardings, (tuple, list)):
+        jit_in_shardings = ([None] * n_captured,) + tuple(in_shardings)
+      elif in_shardings is not None:
+        # Expand scalar in_shardings to match the number of user parameters
+        # so it doesn't broadcast over the captures.
+        sig = _fun_signature(fun)
+        n_params = len(sig.parameters) if sig is not None else 1
+        jit_in_shardings = ([None] * n_captured,) + tuple(
+            in_shardings for _ in range(n_params)
+        )
+      else:
+        jit_in_shardings = in_shardings
+    else:
+      jit_in_shardings = in_shardings
+
     donate_argnums_set = frozenset(
         (donate_argnums,) if isinstance(donate_argnums, int)
         else donate_argnums or ()
@@ -550,21 +603,33 @@ def __init__(
         (donate_argnames,) if isinstance(donate_argnames, str)
         else donate_argnames or ()
     )
+
+    # When captures are prepended as an extra first argument, offset
+    # index-based jax.jit parameters so they still refer to the correct
+    # physical arguments.
+    if n_captured > 0:
+      jit_static_argnums: int | tp.Sequence[int] | None = _offset_argnums(static_argnums, 1)
+      jit_donate_argnums: int | tp.Sequence[int] | None = _offset_argnums(donate_argnums, 1)
+    else:
+      jit_static_argnums = static_argnums
+      jit_donate_argnums = donate_argnums
+
     self.jitted_fn = jax.jit(
         SimpleJitFn(
             fun,
-            self.in_shardings,
+            jit_in_shardings if n_captured > 0 else self.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,
@@ -594,10 +659,16 @@ 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.partial_args, *args)
+    if self._captured_info:
+      all_args = (self._captured_info, *all_args)
+    all_args, kwargs = self._maybe_to_tree(all_args, kwargs)
+    if self._captured_info:
+      variables = extract.check_no_aliases(
+          'jit', captured_args=all_args[0], args=all_args[1:], kwargs=kwargs)
+    else:
+      variables = extract.check_no_aliases('jit', args=all_args, kwargs=kwargs)
+    out, updates = self.jitted_fn(*all_args, **kwargs)
     extract.apply_updates(variables, updates)
     return self._maybe_from_tree(out)
 
@@ -608,22 +679,33 @@ def __get__(self, obj, objtype=None):
 
   def eval_shape(self, *args, **kwargs):
     args = (*self.partial_args, *args)
+    args = (self._captured_info, *args) if self._captured_info else args
     args, kwargs = self._maybe_to_tree(args, kwargs)
     if not self.graph:
-      extract.check_no_aliases('jit', args=args, kwargs=kwargs)
+      if self._captured_info:
+        extract.check_no_aliases(
+            'jit', captured_args=args[0], args=args[1:], kwargs=kwargs)
+      else:
+        extract.check_no_aliases('jit', args=args, 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, *args) if self._captured_info else args
     args, kwargs = self._maybe_to_tree(args, kwargs)
     if not self.graph:
-      extract.check_no_aliases('jit', args=args, kwargs=kwargs)
+      if self._captured_info:
+        extract.check_no_aliases(
+            'jit', captured_args=args[0], args=args[1:], kwargs=kwargs)
+      else:
+        extract.check_no_aliases('jit', args=args, 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, *args) if self._captured_info else args
     args, kwargs = self._maybe_to_tree(args, kwargs)
     if not self.graph:
       extract.check_no_aliases('jit', args=args, kwargs=kwargs)
@@ -1916,6 +1998,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: