perf: add fastpaths in _ModuleMeta.__call__

equinox/_module/_module.py

@@ -9,11 +9,15 @@
     Any,
     cast,
     Final,
+    get_args,
+    get_origin,
+    get_type_hints,
     Literal,
     NamedTuple,
     ParamSpec,
     TYPE_CHECKING,
     TypeVar,
+    Union,
 )
 from typing_extensions import dataclass_transform
 
@@ -52,6 +56,49 @@ class _ModuleInfo(NamedTuple):
     static_field_names: tuple[str, ...]  # names of static=True fields
     non_init_field_names: tuple[str, ...]  # names of init=False fields
     check_init_methods: tuple[Any, ...]  # unbound __check_init__ funcs from MRO
+    # Fastpath: init=False fields that have NO default/default_factory and
+    # therefore are not guaranteed to be set by the dataclass-generated
+    # __init__.  When the class uses the dataclass __init__ we only need to
+    # check these; for all other (init=True) fields the generated __init__
+    # guarantees they are set.
+    unchecked_init_false_names: tuple[str, ...]
+    # Fastpath: False when every init=True field is annotated with a type that
+    # can never hold a JAX-transformed callable. When False, the jtu.tree_leaves
+    # scan in __call__ is skipped entirely.
+    may_receive_callable_args: bool
+    # Mirrors _has_dataclass_init[cls]: cached here to avoid a second
+    # WeakKeyDictionary lookup in the hot __call__ path.
+    has_dataclass_init: bool
+    # When True, _ModuleMeta.__call__ skips all checks/warnings (other than the
+    # abstract check) and only creates the instance, applies converters, and
+    # removes from _currently_initialising.
+    unchecked_init: bool
+
+
+# ---------------------------------------------------------------------------
+# Callable-warning scan helpers
+# ---------------------------------------------------------------------------
+
+_SAFE_FROM_CALLABLE_TYPES: frozenset[type] = frozenset(
+    {int, float, str, bool, bytes, complex, type(None), np.ndarray, jax.Array}
+)
+
+# types.UnionType (``X | Y`` syntax) is available from Python 3.10 onwards.
+_UNION_TYPE: type | None = getattr(types, "UnionType", None)
+
+
+def _is_safe_from_callable_warning(annotation: object, /) -> bool:
+    """True only if no value with this annotation can be a JAX-transformed callable."""
+    if annotation is Any or annotation is object:
+        return False
+    if annotation in _SAFE_FROM_CALLABLE_TYPES:
+        return True
+    origin = get_origin(annotation)
+    if origin is Union or (
+        _UNION_TYPE is not None and isinstance(annotation, _UNION_TYPE)
+    ):
+        return all(_is_safe_from_callable_warning(a) for a in get_args(annotation))
+    return False
 
 
 MSG_METHOD_IN_INIT: Final = """Cannot assign methods in __init__.
@@ -187,11 +234,11 @@ def __init__(self):
         self._dict: dict[int, Module] = {}
 
     def __contains__(self, key: "Module") -> bool:
-        return id(key) in self._dict.keys()
+        return id(key) in self._dict
 
     def add(self, key: "Module") -> None:
-        if key not in self:
-            id_key = id(key)
+        id_key = id(key)
+        if id_key not in self._dict:
             # Hold on to `key` to be sure that `id(key)` does not get reallocated.
             self._dict[id_key] = key
 
@@ -283,6 +330,7 @@ def __new__(
         *,
         is_abstract: bool = False,
         strict: None | bool = False,
+        unchecked_init: bool = False,
         **kwargs: object,
     ):
         if strict is None:
@@ -384,13 +432,40 @@ def __new__(
             for parent_cls in cls.__mro__  # pyright: ignore[reportGeneralTypeIssues]
             if "__check_init__" in parent_cls.__dict__
         )
+        # Precompute fastpath: init=False fields without a default or
+        # default_factory are the only ones NOT guaranteed to be set by the
+        # dataclass-generated __init__, so those are the ones we still have to
+        # check even on the fastpath.
+        unchecked_init_false_names = tuple(
+            f.name
+            for f in fields  # pyright: ignore[reportArgumentType]
+            if not f.init
+            and f.default is dataclasses.MISSING
+            and f.default_factory is dataclasses.MISSING  # pyright: ignore[reportAttributeAccessIssue]
+        )
+        # Precompute whether any init=True field could hold a JAX-transformed
+        # callable. On failure (e.g. unresolvable forward refs) fall back to
+        # True so the scan is never incorrectly skipped.
+        try:
+            hints = get_type_hints(cls)
+        except Exception:
+            hints = {}
+        may_receive_callable_args = not all(
+            _is_safe_from_callable_warning(hints.get(f.name, Any))
+            for f in fields  # pyright: ignore[reportArgumentType]
+            if f.init
+        )
         _module_info[cls] = _ModuleInfo(
             names_tuple=names,
             names_set=frozenset(names),
             converter_fields=converter_fields,
             static_field_names=static_field_names,
             non_init_field_names=non_init_field_names,
             check_init_methods=check_init_methods,
+            unchecked_init_false_names=unchecked_init_false_names,
+            may_receive_callable_args=may_receive_callable_args,
+            has_dataclass_init=has_dataclass_init,
+            unchecked_init=unchecked_init,
         )
 
         # Generate optimized flatten/unflatten functions
@@ -412,7 +487,23 @@ def __call__(cls, *args: object, **kwargs: object):  # noqa: N805
         if cls in _abstract_module_registry:
             # Any other is-abstract checks will be handled in super().__call__.
             raise TypeError("Cannot instantiate abstract `equinox.Module`.")
-        if _has_dataclass_init[cls]:
+
+        info = _module_info[cls]
+
+        if info.unchecked_init:
+            tryself = None
+            try:
+                self = tryself = super().__call__(*args, **kwargs)  # pyright: ignore[reportAttributeAccessIssue]
+            finally:
+                if tryself is not None:
+                    _currently_initialising.remove(tryself)
+                del tryself
+            for name, converter in info.converter_fields:
+                object.__setattr__(self, name, converter(getattr(self, name)))
+            return self
+
+        has_dcls_init = info.has_dataclass_init
+        if has_dcls_init and info.may_receive_callable_args:
             for x in jtu.tree_leaves((args, kwargs)):
                 _warn_jax_transformed_function(cls, x)
 
@@ -425,10 +516,16 @@ def __call__(cls, *args: object, **kwargs: object):  # noqa: N805
             del tryself
         assert not is_abstract_module(cls)  # pyright: ignore[reportArgumentType]
 
-        info = _module_info[cls]
-
         # Check all fields were initialized.
-        for name in info.names_tuple:
+        # Fastpath: when using the dataclass-generated __init__ every init=True
+        # field is guaranteed to be set, and every init=False field that has a
+        # default/default_factory is also set.  Only init=False fields without
+        # any default can be missing, so we only iterate over those.
+        # With a custom __init__ we fall back to checking everything.
+        names_to_check = (
+            info.unchecked_init_false_names if has_dcls_init else info.names_tuple
+        )
+        for name in names_to_check:
             try:
                 getattr(self, name)
             except AttributeError as err:
@@ -441,12 +538,12 @@ def __call__(cls, *args: object, **kwargs: object):  # noqa: N805
         # Warn about init=False fields containing inexact arrays.
         for name in info.non_init_field_names:
             val = getattr(self, name)
-            if any(jtu.tree_map(is_inexact_array_like, jtu.tree_leaves(val))):
+            if any(map(is_inexact_array_like, jtu.tree_leaves(val))):
                 warnings.warn(_MSG_FIELD_INIT_FALSE, stacklevel=2)
 
         # Warn about static fields containing JAX arrays (post-converter values).
         for name in info.static_field_names:
-            if any(jtu.tree_map(is_array, jtu.tree_leaves(getattr(self, name)))):
+            if any(map(is_array, jtu.tree_leaves(getattr(self, name)))):
                 warnings.warn(
                     "A JAX array is being set as static! This can result "
                     "in unexpected behavior and is usually a mistake to do.",
@@ -584,14 +681,18 @@ def __eq__(self, other: object, /) -> bool | np.bool_ | Bool[Array, ""]:  # pyri
     if not TYPE_CHECKING:
 
         def __setattr__(self, name: str, value: Any) -> None:
-            if self in _currently_initialising and (
-                name in _module_info[type(self)].names_set
-                or name in WRAPPER_FIELD_NAMES
-            ):
-                _error_method_assignment(self, value)
-                _warn_jax_transformed_function(type(self), value)
-                object.__setattr__(self, name, value)
-                return
+            if self in _currently_initialising:
+                cls = type(self)
+                if name in _module_info[cls].names_set:
+                    _error_method_assignment(self, value)
+                    _warn_jax_transformed_function(cls, value)
+                    object.__setattr__(self, name, value)
+                    return
+                elif name in WRAPPER_FIELD_NAMES:
+                    # Allow setting wrapper fields without warning, because they
+                    # are ignored by the PyTree machinery.
+                    object.__setattr__(self, name, value)
+                    return
             # Allow:
             # ```
             # class SomeModule(eqx.Module, Generic[T]): ...
@@ -617,9 +718,9 @@ def __getattribute__(self, name: str, /) -> Any:
             # ```
             # works.
             if (
-                not _is_magic(name)
-                and isinstance(out, types.MethodType)
+                isinstance(out, types.MethodType)
                 and out.__self__ is self
+                and not _is_magic(name)
             ):
                 out = BoundMethod(object.__getattribute__(out, "__func__"), self)
             return out

equinox/_module/_prebuilt.py

@@ -23,7 +23,7 @@ def __get__(self, obj: Any, objtype: Any = None) -> Callable[..., _Return_co]: .
 
 
 # Not using `jax.tree_util.Partial` as it doesn't implement __eq__ very well. See #480.
-class BoundMethod(Module, Generic[_Return]):
+class BoundMethod(Module, Generic[_Return], unchecked_init=True):
     """Just like a normal Python bound method... except that this one is a PyTree!
 
     This stores `__self__` as a subnode.
@@ -48,7 +48,7 @@ def __getattr__(self, name: str) -> Any:
         )
 
 
-class Partial(Module, Generic[_Return]):
+class Partial(Module, Generic[_Return], unchecked_init=True):
     """Like `functools.partial`, but treats the wrapped function, and partially-applied
     args and kwargs, as a PyTree.
 
@@ -92,7 +92,7 @@ def __call__(self, *args: Any, **kwargs: Any) -> _Return:
 _Value = TypeVar("_Value")
 
 
-class Static(Module, Generic[_Value]):
+class Static(Module, Generic[_Value], unchecked_init=True):
     """Wraps a value into a `eqx.field(static=True)`.
 
     This is useful to treat something as just static metadata with respect to a JAX

tests/test_module.py

@@ -756,6 +756,134 @@ def __post_init__(self, flag):
         C(flag=False)
 
 
+def test_init_fastpath_metadata():
+    """_ModuleInfo precomputes which init=False fields need checking (RED test)."""
+    from equinox._module._module import _module_info
+
+    class AllInitTrue(eqx.Module):
+        x: int
+        y: float
+
+    # No init=False fields → nothing to check with dataclass-generated __init__
+    assert _module_info[AllInitTrue].unchecked_init_false_names == ()
+
+    class InitFalseNoDefault(eqx.Module):
+        x: int = eqx.field(init=False)
+
+    # init=False without a default → dataclass __init__ does NOT set it; must check
+    assert _module_info[InitFalseNoDefault].unchecked_init_false_names == ("x",)
+
+    class InitFalseWithDefault(eqx.Module):
+        x: int = eqx.field(init=False, default=42)
+
+    # init=False WITH a default → dataclass __init__ sets it; no need to check
+    assert _module_info[InitFalseWithDefault].unchecked_init_false_names == ()
+
+    class InitFalseWithFactory(eqx.Module):
+        x: list = eqx.field(init=False, default_factory=list)
+
+    # init=False WITH a default_factory → also initialized; no need to check
+    assert _module_info[InitFalseWithFactory].unchecked_init_false_names == ()
+
+
+def test_init_fastpath_correctness():
+    """The fastpath still catches errors that custom __init__ leaves fields unset."""
+
+    # Custom __init__ that forgets to set a field → must still raise
+    class CustomForgetsX(eqx.Module):
+        x: int
+
+        def __init__(self):
+            pass  # deliberately forgets self.x
+
+    with pytest.raises(TypeError, match="Field 'x' was not initialized."):
+        CustomForgetsX()
+
+    # Dataclass init + init=False without default → must still raise
+    class DataclassInitFalseNoDefault(eqx.Module):
+        x: int = eqx.field(init=False)
+
+    with pytest.raises(TypeError, match="Field 'x' was not initialized."):
+        DataclassInitFalseNoDefault()
+
+    # Dataclass init + all init=True → no error (common fast path)
+    class DataclassAllInitTrue(eqx.Module):
+        x: int
+
+    m = DataclassAllInitTrue(42)
+    assert m.x == 42
+
+
+def test_callable_warning_scan_metadata():
+    """_ModuleInfo.may_receive_callable_args is False only when all init=True
+    fields have annotations that can never hold a JAX-transformed callable."""
+    from equinox._module._module import _module_info
+
+    class SafeAnnotations(eqx.Module):
+        a: int
+        b: float
+        c: str
+
+    assert _module_info[SafeAnnotations].may_receive_callable_args is False
+
+    class HasAny(eqx.Module):
+        x: Any
+
+    assert _module_info[HasAny].may_receive_callable_args is True
+
+    class HasCallable(eqx.Module):
+        fn: Callable
+
+    assert _module_info[HasCallable].may_receive_callable_args is True
+
+    class HasOptionalInt(eqx.Module):
+        x: int | None
+
+    assert _module_info[HasOptionalInt].may_receive_callable_args is False
+
+    class HasUnionInt(eqx.Module):
+        x: int | None
+
+    assert _module_info[HasUnionInt].may_receive_callable_args is False
+
+    class AllInitFalse(eqx.Module):
+        x: int = eqx.field(init=False, default=0)
+
+    # No init=True fields → nothing passed as constructor args → no scan needed
+    assert _module_info[AllInitFalse].may_receive_callable_args is False
+
+
+def test_callable_warning_scan_correctness(capsys):
+    """When may_receive_callable_args is True the scan still runs; when False it
+    is skipped without changing observable behaviour for correctly-typed modules."""
+    from equinox._module._module import _module_info
+
+    # Safe-typed module: flag must be False (scan is skipped)
+    class SafeModule(eqx.Module):
+        x: int
+        y: float
+
+    assert _module_info[SafeModule].may_receive_callable_args is False
+    # Constructing with correct types must not warn
+    import warnings as _warnings
+
+    with _warnings.catch_warnings():
+        _warnings.simplefilter("error")
+        m = SafeModule(1, 2.0)
+    assert m.x == 1
+
+    # Any-typed module: flag must be True (scan runs as before)
+    class AnyModule(eqx.Module):
+        fn: Any
+
+    assert _module_info[AnyModule].may_receive_callable_args is True
+    # Constructing with a plain int must not warn (scan runs, finds nothing bad)
+    with _warnings.catch_warnings():
+        _warnings.simplefilter("error")
+        m2 = AnyModule(42)
+    assert m2.fn == 42
+
+
 @pytest.mark.parametrize("field", (dataclasses.field, eqx.field))
 def test_init_as_abstract(field):
     # Before the introduction of AbstractVar, it was possible to sort-of get the same