Merge pull request #5055 from mohsinm-dev:fix-bound-method-auto-unbinding

PiperOrigin-RevId: 833378860

Author: Flax Authors

flax/nnx/transforms/autodiff.py

@@ -27,11 +27,13 @@
 )
 from flax.nnx.statelib import State
 import jax
-import jax.core
-import jax.stages
 
 from flax.nnx.transforms import general
-from flax.nnx.transforms.transforms import resolve_kwargs
+from flax.nnx.transforms.transforms import (
+  resolve_kwargs,
+  _resolve_bound_callable,
+  _raise_bound_method_error,
+)
 from flax.typing import MISSING, Missing
 
 
@@ -60,6 +62,7 @@ class DiffState:
   filter: filterlib.Filter
 
 
+
 @dataclasses.dataclass(eq=False)
 class GradFn:
   f: tp.Callable[..., tp.Any]
@@ -312,15 +315,23 @@ def grad(
       holomorphic=holomorphic,
       allow_int=allow_int,
     )
-  return _grad_general(
-    f,
+  # Detect bound nnx.Module methods and raise error.
+  f_unbound, _, was_bound = _resolve_bound_callable(f)
+
+  if was_bound:
+    _raise_bound_method_error('grad')
+
+  grad_fn = _grad_general(
+    f_unbound,
     argnums,
     has_aux,
     holomorphic,
     allow_int,
     return_value=False,
   )
 
+  return grad_fn
+
 
 @tp.overload
 def value_and_grad(
@@ -366,8 +377,14 @@ def value_and_grad(
       holomorphic=holomorphic,
       allow_int=allow_int,
     )
+  # Detect bound nnx.Module methods and raise error.
+  f_unbound, _, was_bound = _resolve_bound_callable(f)
+
+  if was_bound:
+    _raise_bound_method_error('value_and_grad')
+
   return _grad_general(
-    f,
+    f_unbound,
     argnums,
     has_aux,
     holomorphic,
@@ -829,7 +846,13 @@ def custom_vjp(
   """
   if isinstance(fun, Missing):
     return functools.partial(custom_vjp, nondiff_argnums=nondiff_argnums)
-  return CustomVjp(fun, nondiff_argnums)
+
+  # Detect bound nnx.Module methods and raise error.
+  fun_unbound, _, was_bound = _resolve_bound_callable(fun)
+  if was_bound:
+    _raise_bound_method_error('custom_vjp')
+
+  return CustomVjp(fun_unbound, nondiff_argnums)
 
 
 # -------------------------------
@@ -881,11 +904,18 @@ def remat(
       policy=policy,
     )  # type: ignore[return-value]
 
-  return resolve_kwargs()(
+  # Detect bound nnx.Module methods and raise error.
+  f_unbound, _, was_bound = _resolve_bound_callable(f)
+
+  if was_bound:
+    _raise_bound_method_error('remat')
+
+  # Unbound function path: preserve the concise composition used in NNX.
+  return resolve_kwargs()(  # type: ignore[return-value]
     graph.update_context('remat')(
       general.split_inputs(
         jax.checkpoint(
-          general.merge_inputs(f, ctxtag='remat'),
+          general.merge_inputs(f_unbound, ctxtag='remat'),
           prevent_cse=prevent_cse,
           static_argnums=static_argnums,
           policy=policy,

flax/nnx/transforms/compilation.py

@@ -30,6 +30,10 @@
   statelib,
   variablelib,
 )
+from flax.nnx.transforms.transforms import (
+  _resolve_bound_callable,
+  _raise_bound_method_error,
+)
 from flax.typing import MISSING, Missing, PathParts
 
 F = tp.TypeVar('F', bound=tp.Callable[..., tp.Any])
@@ -218,6 +222,11 @@ def jit(
       JAX keeps a weak reference to ``fun`` for use as a compilation cache key,
       so the object ``fun`` must be weakly-referenceable. Most :class:`Callable`
       objects will already satisfy this requirement.
+
+      .. note::
+        Bound methods (e.g., ``module.method``) are not supported. Use the
+        decorator form ``@nnx.jit`` on the method definition or call
+        ``nnx.jit(MyClass.method)(instance, ...)`` with the unbound method.
     in_shardings: Pytree of structure matching that of arguments to ``fun``,
       with all actual arguments replaced by resource assignment specifications.
       It is also valid to specify a pytree prefix (e.g. one value in place of a
@@ -335,8 +344,13 @@ def jit(
       inline=inline,
       abstracted_axes=abstracted_axes,
     )  # type: ignore[return-value]
+  # Detect bound nnx.Module methods and raise error.
+  fun_unbound, _, was_bound = _resolve_bound_callable(fun)
+  if was_bound:
+    _raise_bound_method_error('jit')
+
   return JitWrapped(
-    fun,
+    fun_unbound,
     in_shardings=in_shardings,
     out_shardings=out_shardings,
     static_argnums=static_argnums,
@@ -986,6 +1000,11 @@ def f(m, x):
     )  # type: ignore[return-value]
   assert not isinstance(f, type)
 
+  # Detect bound nnx.Module methods and raise error.
+  f_unbound, _, was_bound = _resolve_bound_callable(f)
+  if was_bound:
+    _raise_bound_method_error('shard_map')
+
   kwarg_specs = PartitionSpec()
   jax_in_specs = jax.tree.map(
     lambda x: extract.NodeStates(
@@ -1033,7 +1052,7 @@ def shard_map_wrapper(*args, **kwargs):
     return out
 
   shard_map_fn = jax.shard_map(
-      ShardMapFn(f, in_specs, out_specs, kwarg_specs, shard_map_wrapper),
+      ShardMapFn(f_unbound, 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

flax/nnx/transforms/iteration.py

@@ -25,7 +25,11 @@
 from flax.nnx import statelib
 from flax.nnx.module import Module
 from flax.nnx.statelib import State
-from flax.nnx.transforms.transforms import resolve_kwargs
+from flax.nnx.transforms.transforms import (
+  resolve_kwargs,
+  _resolve_bound_callable,
+  _raise_bound_method_error,
+)
 from flax.typing import Leaf, Missing, PytreeDeque
 import jax
 import jax.core
@@ -329,6 +333,11 @@ def vmap(
         transform_metadata=transform_metadata,
     )  # type: ignore[return-value]
 
+  # Detect bound nnx.Module methods and raise error.
+  f_unbound, _, was_bound = _resolve_bound_callable(f)
+  if was_bound:
+    _raise_bound_method_error('vmap')
+
   jax_in_axes = jax.tree.map(
     lambda x: extract.NodeStates.from_prefixes(x.axes, metadata=x)
     if isinstance(x, StateAxes)
@@ -342,7 +351,7 @@ def vmap(
     out_axes,
   )
   vmapped_fn = jax.vmap(
-      VmapFn(f, transform_metadata, in_axes, out_axes),
+      VmapFn(f_unbound, transform_metadata, in_axes, out_axes),
       in_axes=jax_in_axes,
       out_axes=(jax_in_axes, jax_out_axes),
       axis_name=axis_name,
@@ -551,6 +560,11 @@ def pmap(
         transform_metadata=transform_metadata,
     )  # type: ignore[return-value]
 
+  # Detect bound nnx.Module methods and raise error.
+  f_unbound, _, was_bound = _resolve_bound_callable(f)
+  if was_bound:
+    _raise_bound_method_error('pmap')
+
   jax_in_axes = jax.tree.map(
     lambda x: extract.NodeStates.from_prefixes(x.axes, metadata=x)
     if isinstance(x, StateAxes)
@@ -564,7 +578,7 @@ def pmap(
     out_axes,
   )
   pmapped_fn = jax.pmap(
-      PmapFn(f, transform_metadata, in_axes, out_axes),
+      PmapFn(f_unbound, transform_metadata, in_axes, out_axes),
       axis_name=axis_name,
       in_axes=jax_in_axes,
       out_axes=(jax_in_axes, jax_out_axes),
@@ -1258,6 +1272,10 @@ def forward(x, model):
         transform_metadata=transform_metadata,
     )  # type: ignore[return-value]
 
+  f_unbound, _, was_bound = _resolve_bound_callable(f)
+  if was_bound:
+    _raise_bound_method_error('scan')
+
   _check_out_axes(out_axes)
 
   input_carry_argnum = _get_carry_argnum(in_axes, is_in_axes=True)
@@ -1272,7 +1290,7 @@ def forward(x, model):
     )
 
   scan_fn = ScanFn(
-    f,
+    f_unbound,
     input_carry_argnum,
     output_carry_argnum,
     in_axes,

flax/nnx/transforms/transforms.py

@@ -94,6 +94,66 @@ def resolve_kwargs_wrapper(*args, **kwargs):
 
 
 
+# -------------------------------
+# helper utilities for bound methods & indices
+# -------------------------------
+
+def _resolve_bound_callable(
+  f: tp.Callable[..., tp.Any],
+) -> tuple[tp.Callable[..., tp.Any], tp.Any | None, bool]:
+  """Detects and extracts bound methods from NNX Module callables.
+
+  This function unwraps functools.partial layers to reach the underlying
+  callable before checking if it's a bound method of an NNX Module.
+
+  Args:
+    f: A callable that may be a bound method of an NNX Module, potentially
+       wrapped in functools.partial.
+
+  Returns:
+    A tuple of (unbound_fn, bound_self, was_bound) where:
+    - unbound_fn: The unbound function (or original if not bound)
+    - bound_self: The Module instance if f was bound, None otherwise
+    - was_bound: True if f was a bound method, False otherwise
+
+  Note:
+    Preserves functools.partial wrappers around the callable and follows
+    the same detection pattern as _get_unbound_fn in bridge/module.py.
+    Detection occurs before any argnum shifting or index normalization.
+  """
+  # Unwrap functools.partial layers to reach the underlying callable.
+  partials: list[tuple[tuple[tp.Any, ...], dict[str, tp.Any] | None]] = []
+  g = f
+  while isinstance(g, functools.partial):  # type: ignore[arg-type]
+    partials.append((g.args or (), g.keywords))  # type: ignore[attr-defined]
+    g = g.func  # type: ignore[attr-defined]
+
+  bound_self = getattr(g, "__self__", None)
+  was_bound = bool(inspect.ismethod(g) and isinstance(bound_self, Module))
+  if was_bound:
+    g = g.__func__  # type: ignore[attr-defined]
+
+  # Reapply partials in reverse unwrap order.
+  for args, kwargs in reversed(partials):
+    kwargs = {} if kwargs is None else kwargs
+    g = functools.partial(g, *args, **kwargs)
+
+  return g, (bound_self if was_bound else None), was_bound
+
+
+def _raise_bound_method_error(transform_name: str):
+  """Raises a standardized error for bound method usage with NNX transforms.
+
+  Args:
+    transform_name: Name of the transform (e.g., 'grad', 'jit', 'remat').
+  """
+  raise ValueError(
+    f"nnx.{transform_name} does not support bound methods. "
+    f"Use the decorator form @nnx.{transform_name} or call "
+    f"nnx.{transform_name}(MyClass.method)(instance, ...) with the unbound method."
+  )
+
+
 class LiftedModule(tp.Generic[M], Module):  # type: ignore[ignored-abstractmethod]
   @abstractmethod
   def _call(self, accessor: DelayedAccessor, *args, **kwargs) -> tp.Any:
@@ -195,12 +255,18 @@ def eval_shape(
     performing any floating point operations (FLOPs) which can be expensive. This can be
     useful for performing shape inference, for example.
   """
+  # Detect bound nnx.Module methods and raise error.
+  f_call, _, was_bound = _resolve_bound_callable(f)
+
+  if was_bound:
+    _raise_bound_method_error('eval_shape')
+
   args, kwargs = extract.to_tree((args, kwargs))
 
   @functools.wraps(f)
   def _eval_shape_fn(*args, **kwargs):
     args, kwargs = extract.from_tree((args, kwargs))
-    out = f(*args, **kwargs)
+    out = f_call(*args, **kwargs)
     return _to_value_metadata(extract.to_tree(out))
 
   out = jax.eval_shape(_eval_shape_fn, *args, **kwargs)
@@ -255,7 +321,13 @@ def checkify(
     >>> print(err)
     Error(nan generated by primitive: sin.)
   """
-  checkify_fn = checkify_lib.checkify(CheckifyFn(f), errors)
+  # Detect bound nnx.Module methods and raise error.
+  f_call, _, was_bound = _resolve_bound_callable(f)
+
+  if was_bound:
+    _raise_bound_method_error('checkify')
+
+  checkify_fn = checkify_lib.checkify(CheckifyFn(f_call), errors)
 
   @functools.wraps(f)
   @graph.update_context('checkify')
@@ -307,4 +379,3 @@ def switch(
     [general.merge_inputs(f, ctxtag='switch') for f in branches],
     *operands,
   )
-