diff --git a/equinox/_enum.py b/equinox/_enum.py index b788c033..db02cdb1 100644 --- a/equinox/_enum.py +++ b/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(): diff --git a/equinox/_filters.py b/equinox/_filters.py index f6b16ebc..8fbcea0c 100644 --- a/equinox/_filters.py +++ b/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] ) diff --git a/equinox/_module/_module.py b/equinox/_module/_module.py index 7af6013e..b5db104b 100644 --- a/equinox/_module/_module.py +++ b/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: diff --git a/equinox/_module/_prebuilt.py b/equinox/_module/_prebuilt.py index ea078a80..6fa52a46 100644 --- a/equinox/_module/_prebuilt.py +++ b/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 diff --git a/tests/test_module.py b/tests/test_module.py index 3116d244..481cbf76 100644 --- a/tests/test_module.py +++ b/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