WIP: faster low level modules

equinox/_enum.py

@@ -154,7 +154,7 @@ class EnumerationItem(Module):
         def __init__(self, x):
             pass
 
-    def __eq__(self, other) -> Bool[Array, ""]:
+    def __eq__(self, other) -> Bool[Array, ""]:  # pyright: ignore[reportIncompatibleMethodOverride]
         if isinstance(other, EnumerationItem):
             if self._enumeration is other._enumeration:
                 with jax.ensure_compile_time_eval():

equinox/_filters.py

@@ -23,7 +23,7 @@
 _ARRAY_TYPES = (*_NDARRAY_TYPES, jax.Array)  # JAX < 0.7.2
 
 try:  # JAX 0.7.2
-    from jax._src.literals import (
+    from jax._src.literals import (  # pyright: ignore[reportMissingImports]
         LiteralArray,  # pyright: ignore[reportAttributeAccessIssue]
     )
 
@@ -32,7 +32,7 @@
     pass
 
 try:  # JAX > 0.7.2
-    from jax._src.literals import (
+    from jax._src.literals import (  # pyright: ignore[reportMissingImports]
         TypedNdArray,  # pyright: ignore[reportAttributeAccessIssue]
     )
 

equinox/_module/_module.py

@@ -48,6 +48,10 @@ class _ModuleInfo(NamedTuple):
 
     names_tuple: tuple[str, ...]
     names_set: frozenset[str]
+    converter_fields: tuple[tuple[str, Any], ...]  # (name, converter) pairs
+    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
 
 
 MSG_METHOD_IN_INIT: Final = """Cannot assign methods in __init__.
@@ -270,7 +274,7 @@ def call(module, x):
 # This deliberately does not pass `frozen_default=True`, as that clashes with custom
 # `__init__` methods.
 @dataclass_transform(field_specifiers=(dataclasses.field, field))
-class _ModuleMeta(BetterABCMeta):
+class _AbstractModuleMeta(BetterABCMeta):
     def __new__(
         mcs,
         name: str,
@@ -327,7 +331,7 @@ def __new__(
                 else:
                     break
             else:
-                assert name == "Module"
+                assert name in ("Module", "_AbstractModule")
                 has_dataclass_init = True  # eqx.Module itself
         _has_dataclass_init[cls] = has_dataclass_init
         # Check before `dataclass` adds an `__init__` method.
@@ -364,11 +368,29 @@ def __new__(
         if has_dataclass_init:
             cls.__init__.__doc__ = init_doc  # pyright: ignore[reportPossiblyUnboundVariable]
 
-        # Cache the field names for later use.
+        # Cache the field names and per-instantiation metadata for later use.
         names = tuple(f.name for f in fields)
+        converter_fields = tuple(
+            (f.name, converter)
+            for f in fields
+            if (converter := f.metadata.get("converter")) is not None
+        )
+        static_field_names = tuple(
+            f.name for f in fields if f.metadata.get("static", False)
+        )
+        non_init_field_names = tuple(f.name for f in fields if not f.init)
+        check_init_methods = tuple(
+            parent_cls.__dict__["__check_init__"]
+            for parent_cls in cls.__mro__  # pyright: ignore[reportGeneralTypeIssues]
+            if "__check_init__" in parent_cls.__dict__
+        )
         _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,
         )
 
         # Generate optimized flatten/unflatten functions
@@ -385,6 +407,26 @@ def __new__(
 
         return cls
 
+
+class _AbstractModule(Hashable, metaclass=_AbstractModuleMeta):
+    """Lower-level base class for modules, primarily for internal use."""
+
+    def __repr__(self) -> str:
+        return tree_pformat(self)
+
+    def __hash__(self) -> int:
+        return hash(
+            tuple((k, getattr(self, k)) for k in _module_info[type(self)].names_tuple)
+        )
+
+    def __eq__(self, other: object, /) -> bool | np.bool_ | Bool[Array, ""]:  # pyright: ignore
+        return tree_equal(self, other)
+
+
+# This deliberately does not pass `frozen_default=True`, as that clashes with custom
+# `__init__` methods.
+@dataclass_transform(field_specifiers=(dataclasses.field, field))
+class _ModuleMeta(_AbstractModuleMeta):
     def __call__(cls, *args: object, **kwargs: object):  # noqa: N805
         __tracebackhide__ = True
         if cls in _abstract_module_registry:
@@ -403,35 +445,36 @@ def __call__(cls, *args: object, **kwargs: object):  # noqa: N805
             del tryself
         assert not is_abstract_module(cls)  # pyright: ignore[reportArgumentType]
 
-        fields = dataclasses.fields(cls)  # pyright: ignore[reportArgumentType]
+        info = _module_info[cls]
 
-        for f in fields:
-            # Check the field was initialized.
+        # Check all fields were initialized.
+        for name in info.names_tuple:
             try:
-                val = getattr(self, f.name)
+                getattr(self, name)
             except AttributeError as err:
-                raise TypeError(f"Field {f.name!r} was not initialized.") from err
-
-            if (converter := f.metadata.get("converter")) is not None:
-                object.__setattr__(self, f.name, converter(getattr(self, f.name)))
-            if f.metadata.get("static", False):
-                if any(jtu.tree_map(is_array, jtu.tree_leaves(getattr(self, f.name)))):
-                    warnings.warn(
-                        "A JAX array is being set as static! This can result "
-                        "in unexpected behavior and is usually a mistake to do.",
-                        stacklevel=2,
-                    )
-            if not f.init:
-                if any(jtu.tree_map(is_inexact_array_like, jtu.tree_leaves(val))):
-                    warnings.warn(_MSG_FIELD_INIT_FALSE, stacklevel=2)
-
-        for parent_cls in cls.__mro__:
-            try:
-                check_init = parent_cls.__dict__["__check_init__"]
-            except KeyError:  # noqa: PERF203
-                pass
-            else:
-                check_init(self)
+                raise TypeError(f"Field {name!r} was not initialized.") from err
+
+        # Warn about init=False fields containing inexact arrays (pre-converter values).
+        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))):
+                warnings.warn(_MSG_FIELD_INIT_FALSE, stacklevel=2)
+
+        # Apply converters.
+        for name, converter in info.converter_fields:
+            object.__setattr__(self, name, converter(getattr(self, name)))
+
+        # 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)))):
+                warnings.warn(
+                    "A JAX array is being set as static! This can result "
+                    "in unexpected behavior and is usually a mistake to do.",
+                    stacklevel=2,
+                )
+
+        for check_init in info.check_init_methods:
+            check_init(self)
 
         return self
 
@@ -443,7 +486,7 @@ def __signature__(cls):  # noqa: N805
         return sig.replace(parameters=params)
 
 
-class Module(Hashable, metaclass=_ModuleMeta):
+class Module(_AbstractModule, metaclass=_ModuleMeta):
     """Base class. Create your model by inheriting from this.
 
     This will make your model a
@@ -547,17 +590,6 @@ def __new__(cls: "type[_ModuleT]", *args: object, **kwargs: object) -> "_ModuleT
         _currently_initialising.add(self)
         return self
 
-    def __repr__(self) -> str:
-        return tree_pformat(self)
-
-    def __hash__(self) -> int:
-        return hash(
-            tuple((k, getattr(self, k)) for k in _module_info[type(self)].names_tuple)
-        )
-
-    def __eq__(self, other: object, /) -> bool | np.bool_ | Bool[Array, ""]:  # pyright: ignore
-        return tree_equal(self, other)
-
     if not TYPE_CHECKING:
 
         def __setattr__(self, name: str, value: Any) -> None:

equinox/_module/_prebuilt.py

@@ -7,18 +7,18 @@
 
 from ._field import field
 from ._flatten import WRAPPER_FIELD_NAMES
-from ._module import Module
+from ._module import _AbstractModule
 
 
 # Not using `jax.tree_util.Partial` as it doesn't implement __eq__ very well. See #480.
-class BoundMethod(Module):
+class BoundMethod(_AbstractModule):
     """Just like a normal Python bound method... except that this one is a PyTree!
 
     This stores `__self__` as a subnode.
     """
 
     __func__: types.FunctionType = field(static=True)
-    __self__: Module
+    __self__: _AbstractModule
 
     def __post_init__(self):
         for field_name in WRAPPER_FIELD_NAMES:
@@ -29,7 +29,7 @@ def __post_init__(self):
             else:
                 setattr(self, field_name, value)
 
-    def __call__(self, *args, **kwargs):
+    def __call__(self, *args: Any, **kwargs: Any):
         __tracebackhide__ = True
         return self.__func__(self.__self__, *args, **kwargs)
 
@@ -41,7 +41,7 @@ def __wrapped__(self):
 _Return = TypeVar("_Return")
 
 
-class Partial(Module, Generic[_Return]):
+class Partial(_AbstractModule, Generic[_Return]):
     """Like `functools.partial`, but treats the wrapped function, and partially-applied
     args and kwargs, as a PyTree.
 
@@ -85,7 +85,7 @@ def __call__(self, *args: Any, **kwargs: Any) -> _Return:
 _Value = TypeVar("_Value")
 
 
-class Static(Module, Generic[_Value]):
+class Static(_AbstractModule, Generic[_Value]):
     """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

@@ -414,7 +414,7 @@ def f(self, b):
             return self.a + b
 
     m = MyModule(13)
-    assert isinstance(m.f, eqx.Module)
+    assert isinstance(m.f, eqx._module._prebuilt.BoundMethod)  # TODO: better
     flat, treedef = jtu.tree_flatten(m.f)
     assert len(flat) == 1
     assert flat[0] == 13