Repository: patrick-kidger/equinox
Pull request #1043 (2025-07-07 19:55:05)

0.13.0


Patch:
diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md
index 6882318..33c3713 100644
--- a/CONTRIBUTING.md
+++ b/CONTRIBUTING.md
@@ -54,9 +54,7 @@ Make your changes. You can then build the documentation by doing
 
 ```bash
 pip install -e '.[docs]'
-mkdocs build
 mkdocs serve
 ```
-(The two back-to-back `mkdocs` commands are needed to get the names of functions correct in the documentation.)
 
 You can then see your local copy of the documentation by navigating to `localhost:8000` in a web browser.
diff --git a/README.md b/README.md
index e57b93a..eebfb3f 100644
--- a/README.md
+++ b/README.md
@@ -101,4 +101,5 @@ If you found this library to be useful in academic work, then please cite: ([arX
 [PySR](https://github.com/milesCranmer/PySR): symbolic regression. (Non-JAX honourable mention!)  
 
 **Awesome JAX**  
-[Awesome JAX](https://github.com/n2cholas/awesome-jax): a longer list of other JAX projects.  
+[Awesome Equinox](https://docs.kidger.site/equinox/awesome-list/)  
+[Awesome JAX](https://github.com/lockwo/awesome-jax): a longer list of other JAX projects.  
diff --git a/docs/api/module/advanced_fields.md b/docs/api/module/advanced_fields.md
index f4b54b8..1604c99 100644
--- a/docs/api/module/advanced_fields.md
+++ b/docs/api/module/advanced_fields.md
@@ -84,39 +84,3 @@ This method has three key differences compared to the `__post_init__` provided b
 ## Creating wrapper modules
 
 ::: equinox.module_update_wrapper
-
-<!--
-## Strict modules
-
-Equinox supports an entirely optional "strict mode", for validating that you follow the abstract/final design pattern as discussed in [this style guide](../../../pattern/).
-
-When enabled via
-```python
-class Foo(eqx.Module, strict=True):
-    ...
-```
-then the following things are checked when you define your class (an error is raised if they fail).
-
-- That all base classes are also strict `eqx.Module`s.
-- That concrete classes are final.
-- The `__init__` method and all fields are all defined on a single class.
-- That abstract classes have names beginning with `"Abstract"`.
-- That no concrete method is overridden. For example, this will raise an error:
-    ```python
-    class Foo(eqx.Module):
-        def f(self): ...
-
-    class Bar(Foo, strict=True):
-        def f(self): ...
-    ```
-    but this is allowed:
-    ```python
-    class Abstract(eqx.Module):
-        @abc.abstractmethod
-        def f(self): ...
-
-    class Concrete(Abstract, strict=True):
-        def f(self): ...
-    ```
-
-Just the strict `Module` is checked. Subclasses will not become strict unless they also opt-in. This makes it possible to safely enable strict modules in a library, without affecting any downstream users.-->
diff --git a/docs/awesome-list.md b/docs/awesome-list.md
index 95179e1..385f39c 100644
--- a/docs/awesome-list.md
+++ b/docs/awesome-list.md
@@ -78,6 +78,9 @@ Coordinates (like `astropy.coordinates`).
 **galax**: [GitHub](https://github.com/GalacticDynamics/galax)  
 Galactic and gravitational dynamics.
 
+**Diffraxtra**: [GitHub](https://github.com/GalacticDynamics/diffraxtra)  
+Extras for `diffrax`: an alternative API on to `diffeqsolve` and dense evaluation.
+
 **FlowJAX**: [GitHub](https://github.com/danielward27/flowjax)  
 Normalizing flows and probabilistic machine learning.
 
@@ -86,3 +89,6 @@ Converts PyTorch `state_dict`s into Equinox pytrees.
 
 **boltz-binder-design**: [GitHub](https://github.com/escalante-bio/boltz-binder-design)   
 Protein design by multi-objective hallucination (~ColabDesign in Equinox). 
+
+**ksim**: [GitHub](https://github.com/kscalelabs/ksim)  
+MJX robotics simulation library.
diff --git a/docs/index.md b/docs/index.md
index 808310d..09b610e 100644
--- a/docs/index.md
+++ b/docs/index.md
@@ -90,4 +90,5 @@ If this quick start has got you interested, then have a read of [All of Equinox]
 [PySR](https://github.com/milesCranmer/PySR): symbolic regression. (Non-JAX honourable mention!)  
 
 **Awesome JAX**  
-[Awesome JAX](https://github.com/n2cholas/awesome-jax): a longer list of other JAX projects.  
+[Awesome Equinox](https://docs.kidger.site/equinox/awesome-list/)  
+[Awesome JAX](https://github.com/lockwo/awesome-jax): a longer list of other JAX projects.  
diff --git a/docs/pattern.md b/docs/pattern.md
index 17f520d..428fa6f 100644
--- a/docs/pattern.md
+++ b/docs/pattern.md
@@ -15,18 +15,6 @@ The following is a useful design pattern. It's not mandatory, but it does come r
 
 Collectively, these rules serve to avoid all the worst ways in which you might spaghetti your code. Abstract classes are used to define interfaces or partial implementations. Final classes are what get passed around at runtime. No overriding (either of methods or of non-abstract classes) means there are no surprises with exactly what is being called. Keeping all fields together means that initialisation is readable.
 
-Equinox will enforce these rules when subclassing with `strict=True`, e.g.
-```python
-class Foo(eqx.Module, strict=True):
-    x: int
-# ...so far so good...
-
-class Bar(Foo, strict=True):
-    y: int
-# ...error raised here: can't define fields in two different classes.
-```
-(In addition, `strict=True` checks two other things: that all abstract classes have a name that starts with `"Abstract"`, and that modules aren't mixed with non-module classes.)
-
 !!! faq "FAQ"
 
     Some quick FAQs:
@@ -150,7 +138,7 @@ class CubicInterpolation(AbstractPolynomialInterpolation):
         coeffs = ...  # some implementation
         super().__init__(coeffs)
 ```
-but this is now much less readable: we've split up initialisation across two different classes. This is a reliable source of bugs. Thus we arrive at the rule that all fields, and the `__init__` method, should all be defined together. Using `strict=True` when subclassing (e.g. `class CubicInterpolation(AbstractPolynomialInterpolation, strict=True)`) will mean that Equinox checks for this, and raises an error if need be.
+but this is now much less readable: we've split up initialisation across two different classes. This is a reliable source of bugs. Thus we arrive at the rule that all fields, and the `__init__` method, should all be defined together.
 
 ## Level 3: implement methods precisely once, and concrete-means-final
 
@@ -161,7 +149,7 @@ In practice, we argue that's a good idea! This rule means that when you see code
 def foo(interp: AbstractPolynomialInterpolation)
     ... = interp.degree()
 ```
-you know that it is calling precisely `AbstractPolynomialInterpolation.degree`, and not an override in some subclass. This is excellent for code readability. Thus we get the rule that no method should be overridden. (And this rule will also be checked via the `strict=True` flag.)
+you know that it is calling precisely `AbstractPolynomialInterpolation.degree`, and not an override in some subclass. This is excellent for code readability. Thus we get the rule that no method should be overridden.
 
 If we assume this, then we now find ourselves arriving at a conclusion: concrete means final. That is, once we have a concrete class (every abstract method/attribute defined in our ABCs is now overridden with an implementation, so we can instantiate this class), then it is now final (we're not allowed to re-override things, so subclassing is pointless). This is how we arrive at the abstract-or-final rule itself!
 
@@ -261,7 +249,7 @@ class ConcreteArrayTwo(ConcreteArray):
     def another_method(self):
         pass
 ```
-We didn't discuss it above, but we do ban things like this as well. (And `strict=True` will prevent this.) The reason is to simplify things when writing something like:
+We didn't discuss it above, but we do ban things like this as well. The reason is to simplify things when writing something like:
 ```python
 def add(x: ConcreteArray, y: ConcreteArray) -> ConcreteArray:
     ...
diff --git a/equinox/__init__.py b/equinox/__init__.py
index 84d128f..ba20453 100644
--- a/equinox/__init__.py
+++ b/equinox/__init__.py
@@ -14,10 +14,6 @@ from ._ad import (
     filter_value_and_grad as filter_value_and_grad,
     filter_vjp as filter_vjp,
 )
-from ._better_abstract import (
-    AbstractClassVar as AbstractClassVar,
-    AbstractVar as AbstractVar,
-)
 from ._caches import clear_caches as clear_caches
 from ._callback import filter_pure_callback as filter_pure_callback
 from ._enum import Enumeration as Enumeration
@@ -39,12 +35,14 @@ from ._filters import (
 from ._jit import EquinoxRuntimeError as EquinoxRuntimeError, filter_jit as filter_jit
 from ._make_jaxpr import filter_make_jaxpr as filter_make_jaxpr
 from ._module import (
+    AbstractClassVar as AbstractClassVar,
+    AbstractVar as AbstractVar,
     field as field,
     Module as Module,
     module_update_wrapper as module_update_wrapper,
     Partial as Partial,
-    static_field as static_field,
-    StrictConfig as StrictConfig,
+    static_field as static_field,  # deprecated
+    StrictConfig as StrictConfig,  # deprecated
 )
 from ._pretty_print import tree_pformat as tree_pformat, tree_pprint as tree_pprint
 from ._serialisation import (
diff --git a/equinox/_filters.py b/equinox/_filters.py
index 6e6e843..7ebe0da 100644
--- a/equinox/_filters.py
+++ b/equinox/_filters.py
@@ -2,6 +2,7 @@ from collections.abc import Callable
 from typing import Any, overload, TypeVar
 
 import jax
+import jax.core
 import jax.numpy as jnp
 import jax.tree_util as jtu
 import numpy as np
@@ -65,6 +66,8 @@ def is_inexact_array_like(element: Any) -> bool:
 
 def _make_filter_tree(is_leaf):
     def _filter_tree(mask: AxisSpec, arg: Any) -> PyTree[bool]:
+        if isinstance(mask, jax.core.Tracer):
+            raise ValueError("`filter_spec` leaf values cannot be traced arrays.")
         if isinstance(mask, bool):
             return jtu.tree_map(lambda _: mask, arg, is_leaf=is_leaf)
         elif callable(mask):
diff --git a/equinox/_module.py b/equinox/_module.py
deleted file mode 100644
index d2af0f9..0000000
--- a/equinox/_module.py
+++ /dev/null
@@ -1,1223 +0,0 @@
-"""Implements the core Module abstraction.
-
-This includes both the core dataclass+pytree combo, plus various related pieces of
-functionality.
-"""
-
-import abc
-import dataclasses
-import functools as ft
-import inspect
-import itertools
-import types
-import typing
-import warnings
-import weakref
-from collections.abc import Callable
-from typing import Any, cast, Final, Protocol, TYPE_CHECKING, TypeVar
-from typing_extensions import dataclass_transform, ParamSpec
-
-import jax
-import jax._src.traceback_util as traceback_util
-import jax.tree_util as jtu
-import numpy as np
-from jaxtyping import Array, Bool, PyTree, PyTreeDef
-
-from ._better_abstract import ABCMeta, dataclass
-from ._caches import cache_clears
-from ._doc_utils import doc_repr
-from ._filters import is_array, is_array_like
-from ._pretty_print import tree_pformat
-from ._tree import tree_equal
-
-
-traceback_util.register_exclusion(__file__)
-
-
-_P = ParamSpec("_P")
-_T = TypeVar("_T")
-
-
-#
-# Part 1: fields. We extend `dataclasses.field` with a wrapper supporting extra features
-# like type converters, and marking a field as static metadata.
-#
-
-
-def static_field(**kwargs):
-    """Deprecated in favour of [`equinox.field`][], i.e. `eqx.field(static=True)`."""
-    return field(**kwargs, static=True)
-
-
-_converter_sentinel: Any = doc_repr(object(), "lambda x: x")
-
-
-def field(
-    *,
-    converter: Callable[[Any], Any] = _converter_sentinel,
-    static: bool = False,
-    **kwargs,
-):
-    """Equinox supports extra functionality on top of the default dataclasses.
-
-    **Arguments:**
-
-    - `converter`: a function to call on this field when the model is initialised. For
-        example, `field(converter=jax.numpy.asarray)` to convert
-        `bool`/`int`/`float`/`complex` values to JAX arrays. This is ran after the
-        `__init__` method (i.e. when using a user-provided `__init__`), and after
-        `__post_init__` (i.e. when using the default dataclass initialisation).
-    - `static`: whether the field should not interact with any JAX transform at all (by
-        making it part of the PyTree structure rather than a leaf).
-    - `**kwargs`: All other keyword arguments are passed on to `dataclass.field`.
-
-    !!! example "Example for `converter`"
-
-        ```python
-        class MyModule(eqx.Module):
-            foo: Array = eqx.field(converter=jax.numpy.asarray)
-
-        mymodule = MyModule(1.0)
-        assert isinstance(mymodule.foo, jax.Array)
-        ```
-
-    !!! example "Example for `static`"
-
-        ```python
-        class MyModule(eqx.Module):
-            normal_field: int
-            static_field: int = eqx.field(static=True)
-
-        mymodule = MyModule("normal", "static")
-        leaves, treedef = jax.tree_util.tree_flatten(mymodule)
-        assert leaves == ["normal"]
-        assert "static" in str(treedef)
-        ```
-
-    `static=True` means that this field is not a node of the PyTree, so it does not
-    interact with any JAX transforms, like JIT or grad. This means that it is usually a
-    bug to make JAX arrays be static fields. `static=True` should very rarely be used.
-    It is preferred to just filter out each field with `eqx.partition` whenever you need
-    to select only some fields.
-    """
-    try:
-        metadata = dict(kwargs.pop("metadata"))  # safety copy
-    except KeyError:
-        metadata = {}
-    if "converter" in metadata:
-        raise ValueError("Cannot use metadata with `static` already set.")
-    if "static" in metadata:
-        raise ValueError("Cannot use metadata with `static` already set.")
-    # We don't just use `lambda x: x` as the default, so that this works:
-    # ```
-    # class Abstract(eqx.Module):
-    #     x: int = eqx.field()
-    #
-    # class Concrete(Abstract):
-    #    @property
-    #    def x(self):
-    #        pass
-    # ```
-    # otherwise we try to call the default converter on a property without a setter,
-    # and an error is raised.
-    # Oddities like the above are to be discouraged, of course, but in particular
-    # `field(init=False)` was sometimes used to denote an abstract field (prior to the
-    # introduction of `AbstractVar`), so we do want to support this.
-    if converter is not _converter_sentinel:
-        metadata["converter"] = converter
-    if static:
-        metadata["static"] = True
-    return dataclasses.field(metadata=metadata, **kwargs)
-
-
-#
-# Part 2: Modules!
-# This is the core of Equinox.
-#
-
-_flatten_sentinel: Final = doc_repr(object(), "-")
-
-
-@dataclass(frozen=True)
-class StrictConfig:
-    """Used to configure strict Equinox modules.
-
-    Passing this as the `strict` parameter of a `Module` marks that the `Module` is
-    strict, and additionally provides some extra custom behaviour.
-
-    Example usage is:
-    ```python
-    class Foo(eqx.Module, strict=StrictConfig(force_abstract=True)):
-        pass
-    ```
-    """
-
-    force_abstract: bool = False
-    allow_abstract_name: bool = False
-    allow_method_override: bool = False
-
-
-StrictConfig.__init__.__doc__ = """**Arguments:**
-
-- `force_abstract`: marks that a class without abstract methods or attributes should
-    still be treated as abstract. Useful for e.g.
-    ```python
-    class AbstractFoo(eqx.Module, strict=True):
-        pass
-    ```
-    which would otherwise be treated as concrete.
-- `allow_abstract_name`: grants this class an exemption from the rule that abstract
-    classes must begin their names with 'Abstract' or '_Abstract', and that concrete
-    classes must not begin their names with 'Abstract' or '_Abstract'.
-- `allow_method_override`: grants this class an exemption from the rule that it cannot
-    override concrete methods.
-
-Both `allow_*` options should be used with care! They exist only to make it easier to
-transition a codebase from non-strict to strict `Module`s in a backward-compatible
-manner. If you are starting a new codebase you should not have need of them.
-"""
-
-
-# Inherits from ABCMeta to support `eqx.{AbstractVar, AbstractClassVar}` and
-# `abc.abstractmethod`.
-class _ActualModuleMeta(ABCMeta):
-    # This method is called whenever you definite a module: `class Foo(eqx.Module): ...`
-    def __new__(
-        mcs,
-        name,
-        bases,
-        dict_,
-        /,
-        strict: bool | StrictConfig = False,
-        **kwargs,
-    ):
-        if isinstance(strict, bool):
-            strict_config = StrictConfig()
-        elif isinstance(strict, StrictConfig):
-            strict_config = strict
-            strict = True
-        else:
-            raise TypeError(
-                "The `strict` parameter in `class Foo(equinox.Module, strict=...)` "
-                "must either be a `bool` or `equinox.StrictConfig` object."
-            )
-
-        # [Step 1] Create the class as normal.
-        cls = super().__new__(mcs, name, bases, dict_, **kwargs)
-        # [Step 2] Arrange for bound methods to be treated as PyTrees as well. This
-        # ensures that
-        # ```
-        # @jax.jit
-        # def run(fn):
-        #     ...
-        # run(SomeModule().some_method)
-        # ```
-        # works.
-        for k, v in cls.__dict__.items():
-            if _not_magic(k) and inspect.isfunction(v):
-                setattr(cls, k, _wrap_method(v))
-
-        # [Step 3] Handle initialisation.
-        #
-        # For context, with any Python dataclass, there are three possible scenarios for
-        # initialisation:
-        # (a) a user-provided `__init__` method is supplied;
-        # (b) dataclasses creates `__init__` , without a user-provided `__post_init__`
-        # (c) dataclasses creates `__init__` , with a user-provided `__post_init__`
-
-        # Create a dataclass `__init__` method if a user method isn't provided.
-        # If a user passed one on this class, then we definitely have a custom __init__.
-        # Else just use whatever our superclass does. Note that this is different to
-        # default dataclass behaviour. Given
-        # ```
-        # @dataclass
-        # class Foo: def __init__(...): ...
-        # @dataclass
-        # class Bar(Foo): pass
-        # ```
-        # then `Bar` will end up with a dataclass-provided `__init__`. That ends up
-        # being ergonomically very annoying, so we disable it.
-        added_custom_init = "__init__" in cls.__dict__
-        if added_custom_init:
-            has_dataclass_init = False
-        else:
-            for kls in cls.__mro__[1:-1]:
-                try:
-                    has_dataclass_init = _has_dataclass_init[kls]
-                except KeyError:
-                    # Non-Module superclasses.
-                    if kls.__init__ is not object.__init__:
-                        has_dataclass_init = False
-                        break
-                else:
-                    break
-            else:
-                assert name == "Module"
-                has_dataclass_init = True  # eqx.Module itself
-
-        # Check for a common error. (Check for `_Initable` to avoid duplicate warnings.)
-        if (
-            not has_dataclass_init
-            and hasattr(cls, "__post_init__")
-            and not issubclass(cls, _Initable)
-        ):
-            warnings.warn(
-                f"Class `{cls.__module__}.{cls.__qualname__}` has both an "
-                "`__init__` method and a `__post_init__` method. This means that "
-                "the `__post_init__` method will not be run!\n"
-                "The reason for this is that `__post_init__` is intended to be "
-                "used with the automatically-generated `__init__` method provided "
-                "by Python dataclasses, which are generated of the form:\n"
-                "```\n"
-                "def __init__(self, field1, field2)\n"
-                "    self.field1 = field1\n"
-                "    self.field2 = field2\n"
-                "    self.__post_init__()\n"
-                "```\n"
-                "and as such a user-provided `__init__` overrides both the setting "
-                "of fields, and the calling of `__post_init__`.\n"
-                "The above is how Python dataclasses work, and has nothing "
-                "to do with Equinox!\n"
-                "If you are using `__post_init__` to check that certain invariants "
-                "hold, then consider using `__check_init__` instead. This is an "
-                "Equinox-specific extension that is always ran. See here for more "
-                "details: "
-                "https://docs.kidger.site/equinox/api/module/advanced_fields/#checking-invariants",  # noqa: E501
-                stacklevel=2,
-            )
-
-        # Fairly common to write `Superclass.__init__.__doc__ = "..."` with
-        # dataclass-provided inits; here we look through the class hierarchy and will
-        # copy this doc forward.
-        if has_dataclass_init:
-            init_doc = cls.__init__.__doc__
-
-        # [Step 4] Register as a dataclass.
-        cls = dataclass(eq=False, repr=False, frozen=True, init=has_dataclass_init)(
-            cls  # pyright: ignore
-        )
-        # [Step 3b] -- finish off building `__init__` methods. Until we'd done
-        # dataclass'ification then we didn't necessarily have our `__init__` method.
-
-        # Set annotation to the converter input. This is useful for runtime type
-        # checkers.
-        # Note that mutating the `__init__.__annotations__` is okay, as it was created
-        # by the dataclass decorator on the previous line, so nothing else owns it.
-        for f in dataclasses.fields(cls):
-            if f.name not in cls.__init__.__annotations__:
-                continue  # Odd behaviour, so skip.
-            try:
-                converter = f.metadata["converter"]
-            except KeyError:
-                pass
-            else:
-                try:
-                    signature = inspect.signature(converter)
-                except ValueError:
-                    # e.g. `inspect.signature(str)` fails
-                    converter_annotation = Any
-                else:
-                    parameters = list(signature.parameters.values())
-                    if len(parameters) == 0:
-                        # No idea what happened, but play it safe.
-                        converter_annotation = Any
-                    else:
-                        converter_annotation = parameters[0].annotation
-                        if converter_annotation is inspect.Signature.empty:
-                            converter_annotation = Any
-                cls.__init__.__annotations__[f.name] = converter_annotation
-
-        # Registering here records that the `dataclass(...)` call has happened.
-        _has_dataclass_init[cls] = has_dataclass_init
-
-        # Assign `__doc__` in case it has been manually overridden:
-        # ```
-        # class Foo(eqx.Module):
-        #     x: int
-        #
-        # Foo.__init__.__doc__ = "Foo should be called with with an integer `x`."
-        #
-        # class Bar(Foo):
-        #     pass
-        #
-        # # Now we try to access `Bar.__init__.__doc__`. (E.g. during docgen.)
-        # ```
-        if has_dataclass_init:
-            cls.__init__.__doc__ = init_doc  # pyright: ignore
-            # TODO: is this next line still necessary?
-            cls.__init__.__module__ = cls.__module__
-
-        # [Step 5] We support an optional `strict` mode for Rust-like strictness in the
-        # type checking.
-        # In practice this is probably too much for your average user, but it's a great
-        # way to build robust libraries.
-        _is_force_abstract[cls] = strict_config.force_abstract
-        _is_strict[cls] = strict
-        if strict:
-            for base in bases:
-                if base is Module:
-                    # We shouldn't check `Module` itself: we can't make it strict as to
-                    # be able to inherit from it we would need to make it abstract. To
-                    # be able to make it abstract its name must start with `Abstract`.
-                    # That we cannot do for backward compatibility.
-                    continue
-                if _is_special_form(base):
-                    # Skip `typing.Generic` etc.
-                    continue
-                # Invariant: all base classes are also strict modules.
-                if not issubclass(base, Module):
-                    raise TypeError(
-                        "Strict `eqx.Module`s must only inherit from other strict "
-                        f"`eqx.Module`s. `{cls.__module__}.{cls.__qualname__}` is a "
-                        "strict Module inheriting from "
-                        f"`{base.__module__}.{base.__qualname__}`, which is not a "
-                        "`Module` at all."
-                    )
-                if not _is_strict[base]:
-                    raise TypeError(
-                        "Strict `eqx.Module`s must only inherit from other strict "
-                        f"`eqx.Module`s. `{cls.__module__}.{cls.__qualname__}` is a "
-                        "strict Module inheriting from "
-                        f"`{base.__module__}.{base.__qualname__}`, which is a "
-                        "`Module`, but not a strict `Module`."
-                    )
-                # Invariant: concrete means final.
-                if not _is_abstract(base):
-                    raise TypeError(
-                        "Every strict `eqx.Module` must be either abstract or final. "
-                        "This means that it is not possible to inherit from a concrete "
-                        f"strict `eqx.Module`. `{cls.__module__}.{cls.__qualname__}` "
-                        "is a strict Module inheriting from "
-                        f"`{base.__module__}.{base.__qualname__}`, which is a concrete "
-                        "strict `Module`."
-                    )
-                # Invariant: field definitions and __init__ methods must all appear on
-                # just one class in a hierarchy.
-                base_num_fields = len(dataclasses.fields(base))
-                if (base_num_fields > 0) or (not _has_dataclass_init[base]):
-                    # If we've added more fields, or added a custom init method, then
-                    # error.
-                    if len(dataclasses.fields(cls)) != base_num_fields:
-                        raise TypeError(
-                            "For readability, any custom `__init__` method, and all "
-                            "fields, must all be defined on the same strict Module. "
-                            f"{cls.__module__}.{cls.__qualname__} is a strict Module "
-                            "that is attempting to add fields, and inherit from "
-                            f"{base.__module__}.{base.__qualname__}. However the "
-                            "latter already has fields or a custom `__init__` defined."
-                        )
-                    if added_custom_init:
-                        raise TypeError(
-                            "For readability, any custom `__init__` method, and all "
-                            "fields, must all be defined on the same strict Module. "
-                            f"{cls.__module__}.{cls.__qualname__} is a strict Module "
-                            "that is attempting to define a custom `__init__` method, "
-                            f"and inherit from {base.__module__}.{base.__qualname__}. "
-                            "However the latter already has fields "
-                            "or a custom `__init__` defined."
-                        )
-            if not strict_config.allow_abstract_name:
-                has_abstract_name = cls.__name__.startswith(
-                    "Abstract"
-                ) or cls.__name__.startswith("_Abstract")
-                if _is_abstract(cls):
-                    if not has_abstract_name:
-                        # Invariant: abstract classes have names beginning with
-                        # `Abstract`.
-                        main = (
-                            "Abstract strict `eqx.Module`s must be named starting "
-                            f"with 'Abstract' or '_Abstract'. Got {name} when defining "
-                            f"{cls.__module__}.{cls.__qualname__}."
-                        )
-                        if _is_force_abstract[cls]:
-                            raise TypeError(main)
-                        inner = []
-                        if len(cls.__abstractmethods__) > 0:
-                            inner.append(
-                                f"abstract methods: {list(cls.__abstractmethods__)}"
-                            )
-                        if len(cls.__abstractvars__) > 0:
-                            inner.append(
-                                f"abstract variables: {list(cls.__abstractvars__)}"
-                            )
-                        if len(cls.__abstractclassvars__) > 0:
-                            inner.append(
-                                "abstract class variables: "
-                                f"{list(cls.__abstractclassvars__)}"
-                            )
-                        inner = ", ".join(inner)
-                        inner = " " + inner + "."
-                        raise TypeError(main + inner)
-                else:
-                    if has_abstract_name:
-                        # Invariant: concrete classes do not have names beginning with
-                        # `Abstract`.
-                        raise TypeError(
-                            "Concrete strict `eqx.Module`s should not have names "
-                            f"starting with 'Abstract' or '_Abstract'. Got '{name}' "
-                            f"when defining {cls.__module__}.{cls.__qualname__}.",
-                        )
-            for k, v in cls.__dict__.items():
-                if isinstance(v, _wrap_method):
-                    v = v.method
-                    # Invariant: concrete methods are not overridden.
-                    if not getattr(v, "__isabstractmethod__", False):
-                        for base in bases:
-                            old_v = getattr(base, k, _dummy_abstract)
-                            if not inspect.isfunction(old_v):
-                                raise TypeError(
-                                    "Strict `eqx.Module`s cannot override non-methods "
-                                    "with methods. "
-                                    f"`{cls.__module__}.{cls.__qualname__}.{k}` is "
-                                    "attempting to override "
-                                    f"`{base.__module__}.{base.__qualname__}.{k}`."
-                                )
-                            if not strict_config.allow_method_override and not getattr(
-                                old_v, "__isabstractmethod__", False
-                            ):
-                                raise TypeError(
-                                    "Strict `eqx.Module`s cannot override concrete "
-                                    "methods. "
-                                    f"`{cls.__module__}.{cls.__qualname__}.{k}` is "
-                                    "attempting to override "
-                                    f"`{base.__module__}.{base.__qualname__}.{k}`."
-                                )
-
-        # [Step 6] Register as a pytree.
-        data_fs = tuple(
-            f.name
-            for f in dataclasses.fields(cls)
-            if not f.metadata.get("static", False)
-        )
-        static_fs = tuple(
-            f.name for f in dataclasses.fields(cls) if f.metadata.get("static", False)
-        )
-
-        sntl = _flatten_sentinel
-
-        def flatten_module(
-            obj: "Module", /
-        ) -> tuple[tuple[PyTree, ...], _FlattenedData]:
-            """Flattens a module into a flattened representation."""
-            d = obj.__dict__
-
-            # Build dynamic keys, values in single pass
-            dynamic_fs = []
-            dynamic_vs = []
-            for k in data_fs:
-                v = d.get(k, sntl)
-                if v is sntl:
-                    continue
-                dynamic_fs.append(k)
-                dynamic_vs.append(v)
-
-            aux = _FlattenedData(
-                tuple(dynamic_fs),
-                tuple([(k, d.get(k, sntl)) for k in static_fs]),
-                tuple([(k, d.get(k, sntl)) for k in _wrapper_field_names]),
-            )
-            return tuple(dynamic_vs), aux
-
-        def flatten_module_with_keys(
-            obj: "Module", /
-        ) -> tuple[tuple[PyTree, ...], _FlattenedData]:
-            """Flattens a module into a flattened representation."""
-            d = obj.__dict__
-
-            # Build dynamic keys, values in single pass
-            dynamic_fs = []
-            dynamic_vs = []
-            for k in data_fs:
-                v = d.get(k, sntl)
-                if v is sntl:
-                    continue
-                dynamic_fs.append(k)
-                dynamic_vs.append((jtu.GetAttrKey(k), v))
-
-            aux = _FlattenedData(
-                tuple(dynamic_fs),
-                tuple([(k, d.get(k, sntl)) for k in static_fs]),
-                tuple([(k, d.get(k, sntl)) for k in _wrapper_field_names]),
-            )
-            return tuple(dynamic_vs), aux
-
-        jtu.register_pytree_with_keys(
-            cls,
-            flatten_with_keys=flatten_module_with_keys,  # pyright: ignore
-            flatten_func=flatten_module,  # pyright: ignore
-            unflatten_func=ft.partial(_unflatten_module, cls),  # pyright: ignore
-        )
-        # Done!
-        return cls
-
-    @property
-    def __signature__(cls):
-        # Use signature of __init__ method for non-callable equinox modules
-        sig = inspect.signature(cls.__init__)
-        params = list(sig.parameters.values())[1:]  # Remove self parameter
-        return sig.replace(parameters=params)
-
-    # This method is called whenever you initialise a module: `MyModule(...)`
-    def __call__(cls, *args, **kwargs):
-        __tracebackhide__ = True
-        if _is_force_abstract[cls]:
-            # Any other is-abstract checks will be handled in super().__call__.
-            raise TypeError("Cannot instantiate abstract `equinox.Module`.")
-        if _has_dataclass_init[cls]:
-            for x in jtu.tree_leaves((args, kwargs)):
-                _warn_jax_transformed_function(cls, x)
-            # else it's handled in __setattr__, but that isn't called here.
-        # [Step 1] Modules are immutable -- except during construction. So defreeze
-        # before init.
-        initable_cls = _make_initable_wrapper(cls)
-        # [Step 2] Instantiate the class as normal.
-        self = super(_ActualModuleMeta, initable_cls).__call__(*args, **kwargs)
-        assert not _is_abstract(cls)
-        # [Step 3] Run converters
-        for field in dataclasses.fields(cls):
-            try:
-                converter = field.metadata["converter"]
-            except KeyError:
-                pass
-            else:
-                try:
-                    value = getattr(self, field.name)
-                except AttributeError:
-                    # Let the all-fields-are-filled check handle the error.
-                    pass
-                else:
-                    setattr(self, field.name, converter(value))
-        # [Step 4] Check that all fields are occupied.
-        missing_names = {
-            field.name
-            for field in dataclasses.fields(cls)  # pyright: ignore
-            # Not `vars` or `__dict__`, to allow for `property`s overwriting a field.
-            # Not recommended, but allowable for backward compatibility.
-            if field.name not in dir(self)
-        }
-        if len(missing_names):
-            raise ValueError(
-                f"The following fields were not initialised during __init__: "
-                f"{missing_names}"
-            )
-        # [Step 5] Prevent arrays from being marked as static
-        for field in dataclasses.fields(self):
-            if field.metadata.get("static", False):
-                if any(
-                    jtu.tree_map(
-                        is_array, jtu.tree_flatten(getattr(self, field.name))[0]
-                    )
-                ):
-                    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,
-                    )
-        # Freeze.
-        object.__setattr__(self, "__class__", cls)
-        # [Step 6] Run any custom validators. (After freezing; as they run
-        # unconditionally across the whole MRO, they aren't allowed to mutate.)
-        for kls in cls.__mro__:
-            try:
-                check = kls.__dict__["__check_init__"]
-            except KeyError:
-                pass
-            else:
-                check(self)
-        return self
-
-    # This bit is kind of sneaky. This is the reason that `_has_dataclass_init` is done
-    # after dataclass registration -- we sneakily treat it as a marker for whether
-    # dataclass registration has happened yet. And if it hasn't, we hide any
-    # `__wrapped__` attribute. We want such attributes for the sake of
-    # `module_update_wrapper`, but if `dataclass` sees it then it tries to follow it.
-    def __getattribute__(cls, item):
-        value = super().__getattribute__(item)
-        if (
-            item == "__wrapped__"
-            and isinstance(value, property)
-            and cls not in _has_dataclass_init
-        ):
-            raise AttributeError
-        else:
-            return value
-
-    def __setattr__(cls, item, value):
-        if _not_magic(item) and inspect.isfunction(value):
-            value = _wrap_method(value)
-        super().__setattr__(item, value)
-
-
-if TYPE_CHECKING:
-
-    @dataclass_transform(field_specifiers=(dataclasses.field, field, static_field))
-    class _ModuleMeta(abc.ABCMeta):
-        __abstractvars__: frozenset[str]
-        __abstractclassvars__: frozenset[str]
-else:
-    _ModuleMeta = _ActualModuleMeta
-
-
-def _is_special_form(cls):
-    # This function is basically a heuristic hack.
-    # If you're getting spurious warnings from this, and think you have another kind of
-    # class that should be excluded from Equinox's checks, then please open a GitHub
-    # issue: https://github.com/patrick-kidger/equinox/issues
-    if cls is _Initable:
-        return True
-    if cls.__module__ in ("typing", "typing_extensions", "collections.abc"):
-        return True
-    if Protocol in cls.__bases__:
-        return True
-    return False
-
-
-def _not_magic(k: str) -> bool:
-    return not (k.startswith("__") and k.endswith("__"))
-
-
-class _wrap_method:
-    def __init__(self, method):
-        self.method = method
-        if getattr(self.method, "__isabstractmethod__", False):
-            self.__isabstractmethod__ = self.method.__isabstractmethod__
-
-    if not hasattr(typing, "GENERATING_DOCUMENTATION"):
-        # Fixes https://github.com/patrick-kidger/equinox/issues/1016
-        @property
-        def __module__(self):
-            return self.method.__module__
-
-        @property
-        def __doc__(self):
-            return self.method.__doc__
-
-    def __get__(self, instance, owner):
-        del owner
-        if instance is None:
-            return self.method
-        else:
-            # Why `inplace=True`?
-            # This is safe because the `BoundMethod` was only instantiated here.
-            # This is necessary so that `_method.__self__ is instance`, which is used
-            # as part of a no-cycle check in `_make_initable`.
-            _method = _module_update_wrapper(
-                BoundMethod(self.method, instance), None, inplace=True
-            )
-            return _method
-
-
-_dummy_abstract = abc.abstractmethod(lambda self: 1)
-_is_force_abstract = weakref.WeakKeyDictionary()
-_is_strict = weakref.WeakKeyDictionary()
-_has_dataclass_init = weakref.WeakKeyDictionary()
-
-
-def _is_abstract(cls):
-    return (
-        _is_force_abstract[cls]
-        or len(cls.__abstractmethods__) > 0
-        or len(cls.__abstractvars__) > 0
-        or len(cls.__abstractclassvars__) > 0
-    )
-
-
-_wrapper_field_names = {
-    "__module__",
-    "__name__",
-    "__qualname__",
-    "__doc__",
-    "__annotations__",
-}
-
-
-class _Initable:
-    # Prevent `__init_subclass__` from triggering when creating initable versions of
-    # classes.
-    def __init_subclass__(cls, **kwargs):
-        del kwargs
-
-
-_transform_types = {
-    type(transform(lambda x: x))
-    for transform in (
-        jax.jit,
-        jax.grad,
-        jax.vmap,
-        jax.value_and_grad,
-        jax.jacfwd,
-        jax.jacrev,
-        jax.hessian,
-        jax.custom_jvp,
-        jax.custom_vjp,
-        jax.checkpoint,  # pyright: ignore
-        jax.pmap,
-    )
-}
-
-
-def _warn_jax_transformed_function(cls, x):
-    # not `isinstance`, just in case JAX every tries to override `__instancecheck__`.
-    if type(x) in _transform_types:
-
-        class _JaxTransformException(Exception):
-            pass
-
-        def _is_array_like(x):
-            if is_array_like(x):
-                raise _JaxTransformException
-
-        while True:
-            try:
-                x = x.__wrapped__
-            except AttributeError:
-                break
-            try:
-                jtu.tree_map(_is_array_like, x)
-            except _JaxTransformException:
-                warnings.warn(
-                    f"""
-Possibly assigning a JAX-transformed callable as an attribute on
-{cls.__module__}.{cls.__qualname__}. This will not have any of its parameters updated.
-
-For example, the following code is buggy:
-```python
-class MyModule(eqx.Module):
-vmap_linear: Callable
-
-def __init__(self, ...):
-    self.vmap_linear = jax.vmap(eqx.nn.Linear(...))
-
-def __call__(self, ...):
-    ... = self.vmap_linear(...)
-```
-This is because the callable returned from `jax.vmap` is *not* a PyTree. This means that
-the parameters inside the `eqx.nn.Linear` layer will not receive gradient updates.
-
-You can most easily fix this either by applying the wrapper at `__call__` time:
-```python
-class MyModule(eqx.Module):
-linear: Callable
-
-def __init__(self, ...):
-    self.linear = eqx.nn.Linear(...)
-
-def __call__(self, ...):
-    ... = jax.vmap(self.linear)(...)
-```
-or by using `eqx.filter_vmap` instead (which *does* return a PyTree):
-```python
-class MyModule(eqx.Module):
-vmap_linear: Callable
-
-def __init__(self, ...):
-    self.vmap_linear = eqx.filter_vmap(eqx.nn.Linear(...))
-
-def __call__(self, ...):
-    ... = self.vmap_linear(...)
-```
-""",
-                    stacklevel=3,
-                )
-                break
-
-
-def _make_initable_wrapper(cls: _ActualModuleMeta) -> _ActualModuleMeta:
-    post_init = getattr(cls, "__post_init__", None)
-    return _make_initable(cls, cls.__init__, post_init, wraps=False)
-
-
-@ft.lru_cache(maxsize=128)
-def _make_initable(
-    cls: _ActualModuleMeta, init, post_init, wraps: bool
-) -> _ActualModuleMeta:
-    # Used as part of the key. Don't cache if these have changed.
-    # In practice, monkey-patching these on the class -- after you've already
-    # instantiated it somewhere! -- is an *ahem*, adventurous, thing to do. But never
-    # let it be said that Equinox doesn't support you in your questionable life choices!
-    del init, post_init
-
-    if wraps:
-        field_names = _wrapper_field_names
-    else:
-        field_names = {field.name for field in dataclasses.fields(cls)}
-
-    class _InitableModule(_Initable, cls):
-        pass
-
-    def __setattr__(self, name, value):
-        if name in field_names:
-            if isinstance(value, BoundMethod) and value.__self__ is self:
-                raise ValueError(
-                    """Cannot assign methods in __init__.
-
-That is, something like the following is not allowed:
-```
-class MyModule(eqx.Module):
-    foo: Callable
-
-    def __init__(self):
-        self.foo = self.bar
-
-    def bar(self):
-        ...
-```
-this is because Equinox modules are PyTrees -- but the above does not have a tree
-structure! `self.foo.__self__` is a cycle that brings us back to `self`.
-
-In the above example, you probably want something like this instead:
-```
-class MyModule(eqx.Module):
-    @property
-    def foo(self):
-        return self.bar
-
-    def bar(self):
-        ...
-```
-so that you can still use `self.foo`, but it is not stored in the PyTree structure.
-
-This is a check that was introduced in Equinox v0.11.0. Before this, the above error
-went uncaught, possibly leading to silently wrong behaviour.
-"""
-                )
-            else:
-                _warn_jax_transformed_function(type(self), value)
-                object.__setattr__(self, name, value)
-        else:
-            raise AttributeError(f"Cannot set attribute {name}")
-
-    # Done like this to avoid dataclasses complaining about overriding setattr on a
-    # frozen class.
-    _InitableModule.__setattr__ = __setattr__
-    # Make beartype happy
-    _InitableModule.__init__ = cls.__init__  # pyright: ignore
-    # Appears in AbstractVar error messages
-    _InitableModule.__name__ = cls.__name__
-    _InitableModule.__qualname__ = cls.__qualname__
-    # I don't have a specific use-case for this but it's probably good practice.
-    _InitableModule.__module__ = cls.__module__
-
-    return _InitableModule
-
-
-cache_clears.append(_make_initable.cache_clear)
-
-
-# Used to provide a pretty repr when doing `jtu.tree_structure(SomeModule(...))`.
-@dataclass(slots=True)
-class _FlattenedData:
-    dynamic_field_names: tuple
-    static_fields: tuple[tuple[str, Any], ...]
-    wrapper_fields: tuple[tuple[str, Any], ...]
-
-    def __repr__(self):
-        return repr((self.dynamic_field_names, self.static_fields))[1:-1]
-
-
-def _unflatten_module(cls: type["Module"], aux: _FlattenedData, dynamic_field_values):
-    # This doesn't go via `__init__`. A user may have done something nontrivial there,
-    # and the field values may be dummy values as used in various places throughout JAX.
-    # See also
-    # https://jax.readthedocs.io/en/latest/pytrees.html#custom-pytrees-and-initialization,
-    # which was (I believe) inspired by Equinox's approach here.
-    module = object.__new__(cls)
-    for name, value in zip(aux.dynamic_field_names, dynamic_field_values):
-        object.__setattr__(module, name, value)
-    for name, value in itertools.chain(aux.static_fields, aux.wrapper_fields):
-        if value is not _flatten_sentinel:
-            object.__setattr__(module, name, value)
-    return module
-
-
-class Module(metaclass=_ModuleMeta):
-    """Base class. Create your model by inheriting from this.
-
-    This will make your model a
-    [dataclass](https://docs.python.org/3/library/dataclasses.html) and a
-    [pytree](https://jax.readthedocs.io/en/latest/pytrees.html).
-
-    **Fields**
-
-    Specify all its fields at the class level (identical to
-    [dataclasses](https://docs.python.org/3/library/dataclasses.html)). This defines
-    its children as a PyTree.
-
-    ```python
-    class MyModule(equinox.Module):
-        weight: jax.Array
-        bias: jax.Array
-        submodule: equinox.Module
-    ```
-
-    **Initialisation**
-
-    A default `__init__` is automatically provided, which just fills in fields with
-    the arguments passed. For example `MyModule(weight, bias, submodule)`.
-
-    Alternatively (quite commonly) you can provide an `__init__` method yourself:
-
-    ```python
-    class MyModule(equinox.Module):
-        weight: jax.Array
-        bias: jax.Array
-        submodule: equinox.Module
-
-        def __init__(self, in_size, out_size, key):
-            wkey, bkey, skey = jax.random.split(key, 3)
-            self.weight = jax.random.normal(wkey, (out_size, in_size))
-            self.bias = jax.random.normal(bkey, (out_size,))
-            self.submodule = equinox.nn.Linear(in_size, out_size, key=skey)
-    ```
-
-    **Methods**
-
-    It is common to create some methods on the class -- for example to define the
-    forward pass of a model.
-
-    ```python
-    class MyModule(equinox.Module):
-        ...  # as above
-
-        def __call__(self, x):
-            return self.submodule(x) + self.weight @ x + self.bias
-    ```
-
-    !!! tip
-
-        You don't have to define `__call__`:
-
-        - You can define other methods if you want.
-        - You can define multiple methods if you want.
-        - You can define no methods if you want. (And just use `equinox.Module` as a
-            nice syntax for custom PyTrees.)
-
-        No method is special-cased.
-
-    **Usage**
-
-    After you have defined your model, then you can use it just like any other
-    PyTree -- that just happens to have some methods attached. In particular you can
-    pass it around across `jax.jit`, `jax.grad` etc. in exactly the way that you're
-    used to.
-
-    !!! example
-
-        If you wanted to, then it would be completely safe to do
-
-        ```python
-        class MyModule(equinox.Module):
-            ...
-
-            @jax.jit
-            def __call__(self, x):
-                ...
-        ```
-
-        because `self` is just a PyTree. Unlike most other neural network libraries,
-        you can mix Equinox and native JAX without any difficulties at all.
-
-    !!! tip "For fans of strong typing."
-
-        Equinox modules are all [ABCs](https://docs.python.org/3/library/abc.html)
-        by default. This means you can use
-        [`abc.abstractmethod`](https://docs.python.org/3/library/abc.html#abc.abstractmethod).
-        You can also create abstract instance attributes or abstract class
-        attributes, see [`equinox.AbstractVar`][] and
-        [`equinox.AbstractClassVar`][].
-    """  # noqa: E501
-
-    #     [`equinox.AbstractClassVar`][]. Finally, some optional strict type-checking
-    #     may be enabled by passing `strict=True`, e.g.
-    #     `class Foo(eqx.Module, strict=True)`; see
-    #     [strict modules](../advanced_fields/#strict-modules) for details.
-    # """  # noqa: E501
-
-    def __hash__(self):
-        return hash(tuple(jtu.tree_leaves(self)))
-
-    def __eq__(  # pyright: ignore
-        self, other
-    ) -> bool | np.bool_ | Bool[Array, ""]:
-        return tree_equal(self, other)
-
-    def __repr__(self):
-        return tree_pformat(self)
-
-
-# Not using `jax.tree_util.Partial` as it doesn't implement __eq__ very well. See #480.
-class BoundMethod(Module):
-    """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
-
-    def __call__(self, *args, **kwargs):
-        __tracebackhide__ = True
-        return self.__func__(self.__self__, *args, **kwargs)
-
-    @property
-    def __wrapped__(self):
-        return self.__func__.__get__(  # pyright: ignore
-            self.__self__, type(self.__self__)
-        )
-
-    # This should be unnecessary in principle. In practice something goes wrong on
-    # Python 3.9 and it returns the wrong thing.
-    @property
-    def __signature__(self):
-        return inspect.signature(self.__wrapped__)
-
-
-#
-# Part 3: some downstream pieces. These don't actually affect the core `Module`
-# abstraction, but this file is a convenient place to put them.
-#
-
-
-def module_update_wrapper(
-    wrapper: Module, wrapped: Callable[_P, _T] | None = None
-) -> Callable[_P, _T]:
-    """Like `functools.update_wrapper` (or its better-known cousin, `functools.wraps`),
-    but acts on [`equinox.Module`][]s, and does not modify its input (it returns the
-    updated module instead).
-
-    !!! Example
-
-        ```python
-        class Wrapper(eqx.Module):
-            fn: Callable
-
-            def __call__(self, *args, **kwargs):
-                return self.fn(*args, **kwargs)
-
-            @property
-            def __wrapped__(self):
-                return self.fn
-
-        def make_wrapper(fn):
-            return eqx.module_update_wrapper(Wrapper(fn))
-        ```
-
-    For example, [`equinox.filter_jit`][] returns a module representing the JIT'd
-    computation. `module_update_wrapper` is used on this module to indicate that this
-    JIT'd computation wraps the original one. (Just like how `functools.wraps` is used.)
-
-    Note that as in the above example, the wrapper class must supply a `__wrapped__`
-    property, which redirects to the wrapped object.
-
-    **Arguments:**
-
-    - `wrapper`: the instance of the wrapper.
-    - `wrapped`: optional, the callable that is being wrapped. If omitted then
-        `wrapper.__wrapped__` will be used.
-
-    **Returns:**
-
-    A copy of `wrapper`, with the attributes `__module__`, `__name__`, `__qualname__`,
-    `__doc__`, and `__annotations__` copied over from the wrapped function.
-    """
-    return _module_update_wrapper(wrapper, wrapped, inplace=False)
-
-
-def _module_update_wrapper(
-    wrapper: Module, wrapped: Callable[_P, _T] | None, inplace: bool
-) -> Callable[_P, _T]:
-    cls = wrapper.__class__
-    if not isinstance(getattr(cls, "__wrapped__", None), property):
-        raise ValueError("Wrapper module must supply `__wrapped__` as a property.")
-
-    if wrapped is None:
-        wrapped = wrapper.__wrapped__  # pyright: ignore
-
-    if not inplace:
-        # Make a clone, to avoid mutating the original input.
-        leaves, treedef = jtu.tree_flatten(wrapper)
-        wrapper = jtu.tree_unflatten(treedef, leaves)
-
-    initable_cls = _make_initable(cls, None, None, wraps=True)
-    object.__setattr__(wrapper, "__class__", initable_cls)
-    try:
-        # Like `ft.update_wrapper(wrapper, wrapped, updated=())`.
-        # We don't update __dict__ as it's common/possible for wrapper and wrapped to
-        # both be classes implementing __call__, in which case copying __dict__ over
-        # basically just breaks the wrapper class.
-        # We don't set __wrapped__, and instead demand that the wrapper class tell us
-        # how to redirect to the wrapped object. This is avoid duplicating part of the
-        # PyTree.
-        for attr in _wrapper_field_names:
-            try:
-                value = getattr(wrapped, attr)
-            except AttributeError:
-                pass
-            else:
-                setattr(wrapper, attr, value)
-    finally:
-        object.__setattr__(wrapper, "__class__", cls)
-    return cast(Callable[_P, _T], wrapper)
-
-
-class Partial(Module):
-    """Like `functools.partial`, but treats the wrapped function, and partially-applied
-    args and kwargs, as a PyTree.
-
-    This is very much like `jax.tree_util.Partial`. The difference is that the JAX
-    version requires that `func` be specifically a *function* -- and will silently
-    misbehave if given any non-function callable, e.g. [`equinox.nn.MLP`][]. In contrast
-    the Equinox version allows for arbitrary callables.
-    """
-
-    func: Callable
-    args: tuple[Any, ...]
-    keywords: dict[str, Any]
-
-    def __init__(self, func, /, *args, **kwargs):
-        """**Arguments:**
-
-        - `func`: the callable to partially apply.
-        - `*args`: any positional arguments to apply.
-        - `**kwargs`: any keyword arguments to apply.
-        """
-        self.func = func
-        self.args = args
-        self.keywords = kwargs
-
-    def __call__(self, *args, **kwargs):
-        """Call the wrapped `self.func`.
-
-        **Arguments:**
-
-        - `*args`: the arguments to apply. Passed after those arguments passed during
-            `__init__`.
-        - `**kwargs`: any keyword arguments to apply.
-
-        **Returns:**
-
-        The result of the wrapped function.
-        """
-        return self.func(*self.args, *args, **kwargs, **self.keywords)
-
-
-class Static(Module):
-    """Wraps a value into a `eqx.field(static=True)`.
-
-    This is useful to treat something as just static metadata with respect to a JAX
-    transformation; for example this is used to return non-arrays from a filtered
-    transform.
-    """
-
-    _leaves: list[Any] = field(static=True)
-    _treedef: PyTreeDef = field(static=True)  # pyright: ignore
-
-    def __init__(self, value: Any):
-        # By flattening, we handle pytrees without `__eq__` methods.
-        # When comparing static metadata for equality, this means we never actually
-        # call `value.__eq__`.
-        self._leaves, self._treedef = jtu.tree_flatten(value)
-
-    @property
-    def value(self):
-        return jtu.tree_unflatten(self._treedef, self._leaves)
diff --git a/equinox/_module/README.md b/equinox/_module/README.md
new file mode 100644
index 0000000..6e03c15
--- /dev/null
+++ b/equinox/_module/README.md
@@ -0,0 +1,35 @@
+# `equinox.Module`
+
+This is the core of the library: subclassing `equinox.Module` ensures that your class is both a PyTree and a frozen dataclass.
+
+The above would be just a couple of lines of code, of course (and this idea exists in numerous libraries, from both before and after Equinox's inception). What else is going on here? Or in other words, why do we think Equinox's implementation is probably the most complete amongst this class of ideas?
+
+## Safety / no footguns:
+
+- Equinox has a lot of nice error messages for every kind of edge case we can think of:
+  - Accidentally marking arrays as static fields;
+  - Accidentally creating cycles via method assignment;
+  - Accidentally closing over arrays;
+  - Accidentally ignoring `__post_init__` because an `__init__` method already exists.
+- Good behaviour around edge cases:
+  - Bound methods like `foo = SomeModule().some_method` are also pytrees, so that things like `jax.jit(...)(foo)` work correctly.
+  - Subclassing propagates `__init__` methods, `__doc__`strings.
+  - Mutating attributes during `__init__` is totally fine because we don't do any `__setattr__` magic.
+- The class is frozen *except* during initialisation, giving both the safety of a frozen dataclass with the nice initialisation syntax of `def __init__(self, ...): self.x = x`.
+
+## Speed:
+
+Flattening/unflattening is pretty fast: we've optimised any overhead down to microseconds.
+
+## A few useful features:
+
+- `equinox.{AbstractVar, AbstractClassVar}`: abstract attributes, analogues for `abc.abstractmethod`.
+- `equinox.field(..., converter=..., static=...)` for conversion of fields on assignment, and for marking that a field should be excluded from the PyTree structure.
+- `equinox.Module.__check_init__`: called at all levels of the class hierarchy, to assert that invariants are true after initialisation.
+- Pretty printing: just like you'd see if you printed it out via Black/`ruff format`, as we use the [Wadler-Lindig](https://github.com/patrick-kidger/wadler_lindig) library.
+
+## Easy to understand:
+
+Equinox Modules should feel familiar from the get-go: they're just frozen dataclasses/pytrees. You can reason about how your code works.
+
+We implement all of the above in a fairly concise, readable amount of code. If you want to understand how Equinox works under the hood, then you should find the code in this directory to be pretty easy to read.
diff --git a/equinox/_module/__init__.py b/equinox/_module/__init__.py
new file mode 100644
index 0000000..af511b0
--- /dev/null
+++ b/equinox/_module/__init__.py
@@ -0,0 +1,14 @@
+from ._better_abstract import (
+    AbstractClassVar as AbstractClassVar,
+    AbstractVar as AbstractVar,
+)
+from ._field import (
+    field as field,
+    static_field as static_field,  # deprecated
+)
+from ._module import (
+    Module as Module,
+    module_update_wrapper as module_update_wrapper,
+    StrictConfig as StrictConfig,  # deprecated
+)
+from ._prebuilt import Partial as Partial, Static as Static
diff --git a/equinox/_better_abstract.py b/equinox/_module/_better_abstract.py
similarity index 98%
rename from equinox/_better_abstract.py
rename to equinox/_module/_better_abstract.py
index 40a0ab1..2ff34f5 100644
--- a/equinox/_better_abstract.py
+++ b/equinox/_module/_better_abstract.py
@@ -171,8 +171,12 @@ def _process_annotation(annotation):
         return is_abstract, is_class
 
 
-class ABCMeta(abc.ABCMeta):
+class BetterABCMeta(abc.ABCMeta):
+    __abstractvars__: frozenset[str]
+    __abstractclassvars__: frozenset[str]
+
     def register(cls, subclass):
+        del subclass
         raise ValueError
 
     def __new__(mcs, name, bases, namespace, /, **kwargs):
@@ -236,7 +240,7 @@ class ABCMeta(abc.ABCMeta):
 
 
 @dataclass_transform()
-def dataclass(**kwargs):
+def better_dataclass(**kwargs):
     def make_dataclass(cls):
         try:
             annotations = cls.__dict__["__annotations__"]
diff --git a/equinox/_module/_field.py b/equinox/_module/_field.py
new file mode 100644
index 0000000..2598f5d
--- /dev/null
+++ b/equinox/_module/_field.py
@@ -0,0 +1,92 @@
+import dataclasses
+import warnings
+from collections.abc import Callable
+from typing import Any
+
+
+def field(
+    *,
+    converter: Callable[[Any], Any] | None = None,
+    static: bool = False,
+    **kwargs,
+):
+    """Equinox supports extra functionality on top of the default dataclasses.
+
+    **Arguments:**
+
+    - `converter`: a function to call on this field when the model is initialised. For
+        example, `field(converter=jax.numpy.asarray)` to convert
+        `bool`/`int`/`float`/`complex` values to JAX arrays. This is ran after the
+        `__init__` method (i.e. when using a user-provided `__init__`), and after
+        `__post_init__` (i.e. when using the default dataclass initialisation).
+        If `converter` is `None`, then no converter is registered.
+    - `static`: whether the field should not interact with any JAX transform at all (by
+        making it part of the PyTree structure rather than a leaf).
+    - `**kwargs`: All other keyword arguments are passed on to `dataclass.field`.
+
+    !!! example "Example for `converter`"
+
+        ```python
+        class MyModule(eqx.Module):
+            foo: Array = eqx.field(converter=jax.numpy.asarray)
+
+        mymodule = MyModule(1.0)
+        assert isinstance(mymodule.foo, jax.Array)
+        ```
+
+    !!! example "Example for `static`"
+
+        ```python
+        class MyModule(eqx.Module):
+            normal_field: int
+            static_field: int = eqx.field(static=True)
+
+        mymodule = MyModule("normal", "static")
+        leaves, treedef = jax.tree_util.tree_flatten(mymodule)
+        assert leaves == ["normal"]
+        assert "static" in str(treedef)
+        ```
+
+    `static=True` means that this field is not a node of the PyTree, so it does not
+    interact with any JAX transforms, like JIT or grad. This means that it is usually a
+    bug to make JAX arrays be static fields. `static=True` should very rarely be used.
+    It is preferred to just filter out each field with `eqx.partition` whenever you need
+    to select only some fields.
+    """
+    try:
+        metadata = dict(kwargs.pop("metadata"))  # safety copy
+    except KeyError:
+        metadata = {}
+    if "converter" in metadata:
+        raise ValueError("Cannot use metadata with `converter` already set.")
+    if "static" in metadata:
+        raise ValueError("Cannot use metadata with `static` already set.")
+    # We don't just use `lambda x: x` as the default, so that this works:
+    # ```
+    # class Abstract(eqx.Module):
+    #     x: int = eqx.field()
+    #
+    # class Concrete(Abstract):
+    #    @property
+    #    def x(self):
+    #        pass
+    # ```
+    # otherwise we try to call the default converter on a property without a setter,
+    # and an error is raised.
+    # Oddities like the above are to be discouraged, of course, but in particular
+    # `field(init=False)` was sometimes used to denote an abstract field (prior to the
+    # introduction of `AbstractVar`), so we do want to support this.
+    if converter is not None:
+        metadata["converter"] = converter
+    if static:
+        metadata["static"] = True
+    return dataclasses.field(metadata=metadata, **kwargs)
+
+
+def static_field(**kwargs):
+    warnings.warn(
+        "`equinox.static_field` is deprecated in favour of "
+        "`equinox.field(static=True)`",
+        stacklevel=2,
+    )
+    return field(**kwargs, static=True)
diff --git a/equinox/_module/_module.py b/equinox/_module/_module.py
new file mode 100644
index 0000000..90285b8
--- /dev/null
+++ b/equinox/_module/_module.py
@@ -0,0 +1,661 @@
+import dataclasses
+import functools as ft
+import inspect
+import itertools
+import types
+import warnings
+import weakref
+from collections.abc import Callable, Hashable
+from typing import Any, cast, Final, ParamSpec, TypeVar
+from typing_extensions import dataclass_transform
+
+import jax
+import jax.tree_util as jtu
+import numpy as np
+from jaxtyping import Array, Bool, PyTree
+
+from .._filters import is_array, is_array_like, is_inexact_array_like
+from .._pretty_print import tree_pformat
+from .._tree import tree_equal
+from ._better_abstract import better_dataclass, BetterABCMeta
+from ._field import field
+
+
+# Legacy compatibibility API, passed to `strict` below.
+def StrictConfig(force_abstact: bool = False, **kwargs):
+    del kwargs
+    if force_abstact:
+        return None
+    else:
+        return False
+
+
+wrapper_field_names: Final = {
+    "__module__",
+    "__name__",
+    "__qualname__",
+    "__doc__",
+    "__annotations__",
+}
+
+
+_abstract_module_registry = weakref.WeakSet()
+_has_dataclass_init = weakref.WeakKeyDictionary()
+_flatten_sentinel = object()
+
+
+# Used to provide a pretty repr when doing `jtu.tree_structure(SomeModule(...))`.
+@dataclasses.dataclass(slots=True)
+class _FlattenedData:
+    dynamic_field_names: tuple
+    static_fields: tuple[tuple[str, Any], ...]
+    wrapper_fields: tuple[tuple[str, Any], ...]
+
+    def __repr__(self):
+        return repr((self.dynamic_field_names, self.static_fields))[1:-1]
+
+
+class _ModuleFlattener:
+    __slots__: tuple[str, str] = ("dynamic_fs", "static_fs")
+    dynamic_fs: tuple[str, ...]
+    static_fs: tuple[str, ...]
+
+    def __init__(self, fields: tuple[dataclasses.Field[Any], ...]):
+        dynamic_fs = []
+        static_fs = []
+        for f in fields:
+            if f.metadata.get("static", False):
+                static_fs.append(f.name)
+            else:
+                dynamic_fs.append(f.name)
+        self.dynamic_fs = tuple(dynamic_fs)
+        self.static_fs = tuple(static_fs)
+
+    def flatten(self, obj: "Module") -> tuple[tuple[PyTree, ...], _FlattenedData]:
+        get = obj.__dict__.get
+        dynamic_fs = []
+        dynamic_vs = []
+        for k in self.dynamic_fs:
+            v = get(k, _flatten_sentinel)
+            if v is _flatten_sentinel:
+                continue
+            dynamic_fs.append(k)
+            dynamic_vs.append(v)
+        aux = _FlattenedData(
+            tuple(dynamic_fs),
+            tuple([(k, get(k, _flatten_sentinel)) for k in self.static_fs]),
+            tuple([(k, get(k, _flatten_sentinel)) for k in wrapper_field_names]),
+        )
+        return tuple(dynamic_vs), aux
+
+    def flatten_with_keys(
+        self, obj: "Module"
+    ) -> tuple[tuple[tuple[Any, PyTree], ...], _FlattenedData]:
+        get = obj.__dict__.get
+        dynamic_fs = []
+        dynamic_vs = []
+        for k in self.dynamic_fs:
+            v = get(k, _flatten_sentinel)
+            if v is _flatten_sentinel:
+                continue
+            dynamic_fs.append(k)
+            dynamic_vs.append((jtu.GetAttrKey(k), v))
+        aux = _FlattenedData(
+            tuple(dynamic_fs),
+            tuple([(k, get(k, _flatten_sentinel)) for k in self.static_fs]),
+            tuple([(k, get(k, _flatten_sentinel)) for k in wrapper_field_names]),
+        )
+        return tuple(dynamic_vs), aux
+
+    @staticmethod
+    def unflatten_with_cls(
+        module_cls: type["Module"],
+        aux: _FlattenedData,
+        dynamic_field_values: tuple[PyTree, ...],
+    ) -> "Module":
+        # This doesn't go via `__init__`. A user may have done something
+        # nontrivial there, and the field values may be dummy values as used in
+        # various places throughout JAX. See also
+        # https://jax.readthedocs.io/en/latest/pytrees.html#custom-pytrees-and-initialization,
+        # which was (I believe) inspired by Equinox's approach here.
+        module = object.__new__(module_cls)
+        for name, value in zip(aux.dynamic_field_names, dynamic_field_values):
+            object.__setattr__(module, name, value)
+        for name, value in itertools.chain(aux.static_fields, aux.wrapper_fields):
+            if value is not _flatten_sentinel:
+                object.__setattr__(module, name, value)
+        return module
+
+
+def _error_method_assignment(self, value):
+    if isinstance(value, BoundMethod) and value.__self__ is self:
+        raise ValueError(
+            """Cannot assign methods in __init__.
+
+That is, something like the following is not allowed:
+```
+class MyModule(eqx.Module):
+    foo: Callable
+
+    def __init__(self):
+        self.foo = self.bar
+
+    def bar(self):
+        ...
+```
+this is because Equinox modules are PyTrees -- but the above does not have a tree
+structure! `self.foo.__self__` is a cycle that brings us back to `self`.
+
+In the above example, you probably want something like this instead:
+```
+class MyModule(eqx.Module):
+    @property
+    def foo(self):
+        return self.bar
+
+    def bar(self):
+        ...
+```
+so that you can still use `self.foo`, but it is not stored in the PyTree structure.
+
+This is a check that was introduced in Equinox v0.11.0. Before this, the above error
+went uncaught, possibly leading to silently wrong behaviour.
+"""
+        )
+
+
+_transform_types = {
+    type(transform(lambda x: x))
+    for transform in (
+        jax.jit,
+        jax.grad,
+        jax.vmap,
+        jax.value_and_grad,
+        jax.jacfwd,
+        jax.jacrev,
+        jax.hessian,
+        jax.custom_jvp,
+        jax.custom_vjp,
+        jax.checkpoint,  # pyright: ignore[reportPrivateImportUsage]
+        jax.pmap,
+    )
+}
+
+
+class _JaxTransformException(Exception):
+    pass
+
+
+def _is_array_like(x):
+    if is_array_like(x):
+        raise _JaxTransformException
+
+
+def _warn_jax_transformed_function(cls, x):
+    # not `isinstance`, just in case JAX every tries to override `__instancecheck__`.
+    if type(x) in _transform_types:
+        while True:
+            try:
+                x = x.__wrapped__
+            except AttributeError:
+                break
+            try:
+                jtu.tree_map(_is_array_like, x)
+            except _JaxTransformException:
+                warnings.warn(
+                    f"""
+Possibly assigning a JAX-transformed callable as an attribute on
+{cls.__module__}.{cls.__qualname__}. This will not have any of its parameters updated.
+
+For example, the following code is buggy:
+```python
+class MyModule(eqx.Module):
+vmap_linear: Callable
+
+def __init__(self, ...):
+    self.vmap_linear = jax.vmap(eqx.nn.Linear(...))
+
+def __call__(self, ...):
+    ... = self.vmap_linear(...)
+```
+This is because the callable returned from `jax.vmap` is *not* a PyTree. This means that
+the parameters inside the `eqx.nn.Linear` layer will not receive gradient updates.
+
+You can most easily fix this either by applying the wrapper at `__call__` time:
+```python
+class MyModule(eqx.Module):
+linear: Callable
+
+def __init__(self, ...):
+    self.linear = eqx.nn.Linear(...)
+
+def __call__(self, ...):
+    ... = jax.vmap(self.linear)(...)
+```
+or by using `eqx.filter_vmap` instead (which *does* return a PyTree):
+```python
+class MyModule(eqx.Module):
+vmap_linear: Callable
+
+def __init__(self, ...):
+    self.vmap_linear = eqx.filter_vmap(eqx.nn.Linear(...))
+
+def __call__(self, ...):
+    ... = self.vmap_linear(...)
+```
+""",
+                    stacklevel=3,
+                )
+                break
+
+
+class _IdSet:
+    __slots__ = ("_dict",)
+
+    def __init__(self):
+        self._dict = {}
+
+    def __contains__(self, key):
+        return id(key) in self._dict.keys()
+
+    def add(self, key):
+        if key not in self:
+            id_key = id(key)
+            # Hold on to `key` to be sure that `id(key)` does not get reallocated.
+            self._dict[id_key] = key
+
+    def remove(self, key):
+        del self._dict[id(key)]
+
+
+_currently_initialising = _IdSet()
+
+
+# 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):
+    def __new__(
+        mcs,
+        name,
+        bases,
+        namespace,
+        *,
+        is_abstract: bool = False,
+        strict: None | bool = False,
+        **kwargs,
+    ):
+        if strict is None:
+            # Legacy compatibility API. Checking that this has the desired behaviour:
+            #
+            # - No argument => not abstract. Still true now, the default `strict=False`
+            #    does not take this branch.
+            # - `strict=False` => not abstract. Still true now as above.
+            # - `strict=True` => not abstract. Still true now as above.
+            # - `strict=StrictConfig(force_abstract=False)` => not abstract as
+            #    `StrictConfig` returns `False` in this case.
+            # - `strict=StrictConfig(force_abstract=True)` => abstract. And indeed then
+            #    we take this branch as `StrictConfig` returns `None` in this case.
+            is_abstract = True
+        del strict
+
+        cls = super().__new__(mcs, name, bases, namespace, **kwargs)
+
+        if is_abstract:
+            _abstract_module_registry.add(cls)
+
+        # Create a dataclass `__init__` method if a user method isn't provided.
+        # If a user passed one on this class, then we definitely have a custom __init__.
+        # Else just use whatever our superclass does. Note that this is different to
+        # default dataclass behaviour. Given
+        # ```
+        # @dataclass
+        # class Foo: def __init__(...): ...
+        # @dataclass
+        # class Bar(Foo): pass
+        # ```
+        # then `Bar` will end up with a dataclass-provided `__init__`. That ends up
+        # being ergonomically very annoying, so we disable it.
+        added_custom_init = "__init__" in cls.__dict__
+        if added_custom_init:
+            has_dataclass_init = False
+        else:
+            for kls in cls.__mro__[1:-1]:
+                try:
+                    has_dataclass_init = _has_dataclass_init[kls]
+                except KeyError:
+                    # Non-Module superclasses.
+                    if kls.__init__ is not object.__init__:
+                        has_dataclass_init = False
+                        break
+                else:
+                    break
+            else:
+                assert name == "Module"
+                has_dataclass_init = True  # eqx.Module itself
+        _has_dataclass_init[cls] = has_dataclass_init
+        # Check before `dataclass` adds an `__init__` method.
+        if not has_dataclass_init and hasattr(cls, "__post_init__"):
+            warnings.warn(
+                f"Class `{cls.__module__}.{cls.__qualname__}` has both an "
+                "`__init__` method and a `__post_init__` method. This means that "
+                "the `__post_init__` method will not be run!\n"
+                "The reason for this is that `__post_init__` is intended to be "
+                "used with the automatically-generated `__init__` method provided "
+                "by Python dataclasses, which are generated of the form:\n"
+                "```\n"
+                "def __init__(self, field1, field2)\n"
+                "    self.field1 = field1\n"
+                "    self.field2 = field2\n"
+                "    self.__post_init__()\n"
+                "```\n"
+                "and as such a user-provided `__init__` overrides both the setting "
+                "of fields, and the calling of `__post_init__`.\n"
+                "The above is how Python dataclasses work, and has nothing "
+                "to do with Equinox!\n"
+                "If you are using `__post_init__` to check that certain invariants "
+                "hold, then consider using `__check_init__` instead. This is an "
+                "Equinox-specific extension that is always ran. See here for more "
+                "details: "
+                "https://docs.kidger.site/equinox/api/module/advanced_fields/#checking-invariants",  # noqa: E501
+                stacklevel=2,
+            )
+        if has_dataclass_init:
+            init_doc = cls.__init__.__doc__
+
+        cls = better_dataclass(eq=False, repr=False, init=has_dataclass_init)(cls)
+        for f in dataclasses.fields(cls):  # pyright: ignore[reportArgumentType]
+            if f.name not in cls.__init__.__annotations__:
+                continue  # Odd behaviour, so skip.
+            try:
+                converter = f.metadata["converter"]
+            except KeyError:
+                pass
+            else:
+                try:
+                    signature = inspect.signature(converter)
+                except ValueError:
+                    # e.g. `inspect.signature(str)` fails
+                    converter_annotation = Any
+                else:
+                    parameters = list(signature.parameters.values())
+                    if len(parameters) == 0:
+                        # No idea what happened, but play it safe.
+                        converter_annotation = Any
+                    else:
+                        converter_annotation = parameters[0].annotation
+                        if converter_annotation is inspect.Signature.empty:
+                            converter_annotation = Any
+                cls.__init__.__annotations__[f.name] = converter_annotation
+        if has_dataclass_init:
+            cls.__init__.__doc__ = init_doc  # pyright: ignore[reportPossiblyUnboundVariable]
+
+        flattener = _ModuleFlattener(dataclasses.fields(cls))  # pyright: ignore[reportArgumentType]
+        jtu.register_pytree_with_keys(
+            cls,
+            flatten_with_keys=flattener.flatten_with_keys,  # pyright: ignore
+            flatten_func=flattener.flatten,  # pyright: ignore
+            unflatten_func=ft.partial(flattener.unflatten_with_cls, cls),  # pyright: ignore
+        )
+
+        return cls
+
+    def __call__(cls, *args, **kwargs):  # noqa: N805
+        __tracebackhide__ = True
+        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]:
+            for x in jtu.tree_leaves((args, kwargs)):
+                _warn_jax_transformed_function(cls, x)
+
+        self = None
+        try:
+            self = super().__call__(*args, **kwargs)  # pyright: ignore[reportAttributeAccessIssue]
+        finally:
+            if self is not None:
+                _currently_initialising.remove(self)
+        assert not is_abstract_module(cls)  # pyright: ignore[reportArgumentType]
+
+        fields = dataclasses.fields(cls)  # pyright: ignore[reportArgumentType]
+        # Not `vars` or `__dict__`, to allow for `property`s overwriting a field.
+        # Not recommended, but allowable for backward compatibility.
+        dir_self = dir(self)
+        missing_names = {f.name for f in fields if f.name not in dir_self}
+        if len(missing_names) > 0:
+            raise TypeError(
+                f"The following fields were not initialised during __init__: "
+                f"{missing_names}"
+            )
+
+        for f in fields:
+            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(getattr(self, f.name))
+                    )
+                ):
+                    warnings.warn(
+                        "Using `field(init=False)` on `equinox.Module` can lead to "
+                        "surprising behaviour when used around `jax.grad`. In the "
+                        "following example, observe how JAX computes gradients with "
+                        "respect to the `.len` attribute (which is a PyTree leaf "
+                        "passed across the `jax.grad` boundary) and that there are no "
+                        "gradients with respect to `.a` or `.b`:\n"
+                        "\n"
+                        "```\n"
+                        "import equinox as eqx\n"
+                        "import jax\n"
+                        "import jax.numpy as jnp\n"
+                        "\n"
+                        "class Foo(eqx.Module):\n"
+                        "    a: jax.Array\n"
+                        "    b: jax.Array\n"
+                        "    len: jax.Array = eqx.field(init=False)\n"
+                        "\n"
+                        "    def __post_init__(self):\n"
+                        "        self.len = jnp.sqrt(self.a**2 + self.b**2)\n"
+                        "\n"
+                        "    def __call__(self, x):\n"
+                        "        return self.len * x\n"
+                        "\n"
+                        "@jax.jit\n"
+                        "@jax.grad\n"
+                        "def call(module, x):\n"
+                        "    return module(x)\n"
+                        "\n"
+                        "grads = call(Foo(jnp.array(3.0), jnp.array(4.0)), 5)\n"
+                        "# Foo(\n"
+                        "#   a=Array(0., dtype=float32, weak_type=True),\n"
+                        "#   b=Array(0., dtype=float32, weak_type=True),\n"
+                        "#   len=Array(5., dtype=float32, weak_type=True)\n"
+                        "# )\n"
+                        "```",
+                        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)
+
+        return self
+
+    # Ensure that `help(FooModule)` still works, even though we've overriden `__call__`.
+    @property
+    def __signature__(cls):  # noqa: N805
+        sig = inspect.signature(cls.__init__)
+        params = list(sig.parameters.values())[1:]  # Remove self parameter
+        return sig.replace(parameters=params)
+
+
+class Module(Hashable, metaclass=_ModuleMeta):
+    def __new__(cls, *args, **kwargs):
+        del args, kwargs
+        self = super().__new__(cls)
+        # We record currently-initialising modules
+        _currently_initialising.add(self)
+        return self
+
+    def __repr__(self):
+        return tree_pformat(self)
+
+    def __hash__(self) -> int:
+        return hash(
+            tuple(
+                (f.name, getattr(self, f.name)) for f in dataclasses.fields(type(self))
+            )
+        )
+
+    def __eq__(self, other) -> bool | np.bool_ | Bool[Array, ""]:  # pyright: ignore
+        return tree_equal(self, other)
+
+    def __setattr__(self, name: str, value: Any):
+        if self in _currently_initialising:
+            allowed_names = frozenset({f.name for f in dataclasses.fields(self)}).union(
+                wrapper_field_names
+            )
+            if name in allowed_names:
+                _error_method_assignment(self, value)
+                _warn_jax_transformed_function(type(self), value)
+                object.__setattr__(self, name, value)
+                return
+        # Allow:
+        # ```
+        # class SomeModule(eqx.Module, Generic[T]): ...
+        # x = SomeModule[int]()
+        # x.__orig_class__ # SomeModule[int]
+        # ```
+        # This attribute is set after instantiation here:
+        # https://github.com/python/cpython/blob/7b3ab5921fa25ed8b97b6296f97c5c78aacf5447/Lib/typing.py#L728
+        # So without special-casing it's incompatible with frozen dataclasses.
+        if name == "__orig_class__":
+            object.__setattr__(self, name, value)
+        raise dataclasses.FrozenInstanceError(f"cannot assign to field '{name}'")
+
+    def __getattribute__(self, name: str, /) -> Any:
+        out = super().__getattribute__(name)
+        # Arrange for bound methods to be treated as PyTrees as well. This
+        # ensures that
+        # ```
+        # @jax.jit
+        # def run(fn):
+        #     ...
+        # run(SomeModule().some_method)
+        # ```
+        # works.
+        if (
+            not _is_magic(name)
+            and isinstance(out, types.MethodType)
+            and out.__self__ is self
+        ):
+            out = BoundMethod(object.__getattribute__(out, "__func__"), self)
+        return out
+
+
+def _is_magic(k: str) -> bool:
+    return (k.startswith("__") and k.endswith("__")) or (k == "_abc_impl")
+
+
+def is_abstract_module(cls: type[Module]) -> bool:
+    if not issubclass(cls, Module):
+        raise TypeError(f"{cls} is not a subclass of `Module`.")
+    return (
+        (len(cls.__abstractmethods__) > 0)
+        or (len(cls.__abstractvars__) > 0)
+        or (len(cls.__abstractclassvars__) > 0)
+        or (cls in _abstract_module_registry)
+    )
+
+
+_P = ParamSpec("_P")
+_T = TypeVar("_T")
+
+
+def module_update_wrapper(
+    wrapper: Module, wrapped: Callable[_P, _T] | None = None
+) -> Callable[_P, _T]:
+    """Like `functools.update_wrapper` (or its better-known cousin, `functools.wraps`),
+    but acts on [`equinox.Module`][]s, and does not modify its input (it returns the
+    updated module instead).
+
+    !!! Example
+
+        ```python
+        class Wrapper(eqx.Module):
+            fn: Callable
+
+            def __call__(self, *args, **kwargs):
+                return self.fn(*args, **kwargs)
+
+            @property
+            def __wrapped__(self):
+                return self.fn
+
+        def make_wrapper(fn):
+            return eqx.module_update_wrapper(Wrapper(fn))
+        ```
+
+    For example, [`equinox.filter_jit`][] returns a module representing the JIT'd
+    computation. `module_update_wrapper` is used on this module to indicate that this
+    JIT'd computation wraps the original one. (Just like how `functools.wraps` is used.)
+
+    Note that as in the above example, the wrapper class must supply a `__wrapped__`
+    property, which redirects to the wrapped object.
+
+    **Arguments:**
+
+    - `wrapper`: the instance of the wrapper.
+    - `wrapped`: optional, the callable that is being wrapped. If omitted then
+        `wrapper.__wrapped__` will be used.
+
+    **Returns:**
+
+    A copy of `wrapper`, with the attributes `__module__`, `__name__`, `__qualname__`,
+    `__doc__`, and `__annotations__` copied over from the wrapped function.
+    """
+    cls = wrapper.__class__
+    if not isinstance(getattr(cls, "__wrapped__", None), property):
+        raise ValueError("Wrapper module must supply `__wrapped__` as a property.")
+
+    if wrapped is None:
+        wrapped = wrapper.__wrapped__  # pyright: ignore
+
+    # Make a clone, to avoid mutating the original input.
+    leaves, treedef = jtu.tree_flatten(wrapper)
+    wrapper = jtu.tree_unflatten(treedef, leaves)
+
+    # Like `ft.update_wrapper(wrapper, wrapped, updated=())`.
+    # We don't update __dict__ as it's common/possible for wrapper and wrapped to
+    # both be classes implementing __call__, in which case copying __dict__ over
+    # basically just breaks the wrapper class.
+    # We don't set __wrapped__, and instead demand that the wrapper class tell us
+    # how to redirect to the wrapped object. This is avoid duplicating part of the
+    # PyTree.
+    _currently_initialising.add(wrapper)
+    try:
+        for field_name in wrapper_field_names:
+            try:
+                value = getattr(wrapped, field_name)
+            except AttributeError:
+                pass
+            else:
+                setattr(wrapper, field_name, value)
+    finally:
+        _currently_initialising.remove(wrapper)
+    return cast(Callable[_P, _T], wrapper)
+
+
+from ._prebuilt import BoundMethod  # After Module is defined.
diff --git a/equinox/_module/_prebuilt.py b/equinox/_module/_prebuilt.py
new file mode 100644
index 0000000..4007820
--- /dev/null
+++ b/equinox/_module/_prebuilt.py
@@ -0,0 +1,103 @@
+import types
+from collections.abc import Callable
+from typing import Any, Generic, TypeVar
+
+import jax.tree_util as jtu
+from jaxtyping import PyTreeDef
+
+from ._field import field
+from ._module import Module, wrapper_field_names
+
+
+# Not using `jax.tree_util.Partial` as it doesn't implement __eq__ very well. See #480.
+class BoundMethod(Module):
+    """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
+
+    def __post_init__(self):
+        for field_name in wrapper_field_names:
+            try:
+                value = getattr(self.__func__, field_name)
+            except AttributeError:
+                pass
+            else:
+                setattr(self, field_name, value)
+
+    def __call__(self, *args, **kwargs):
+        __tracebackhide__ = True
+        return self.__func__(self.__self__, *args, **kwargs)
+
+    @property
+    def __wrapped__(self):
+        return self.__func__.__get__(self.__self__, type(self.__self__))  # pyright: ignore[reportAttributeAccessIssue]
+
+
+_Return = TypeVar("_Return")
+
+
+class Partial(Module, Generic[_Return]):
+    """Like `functools.partial`, but treats the wrapped function, and partially-applied
+    args and kwargs, as a PyTree.
+
+    This is very much like `jax.tree_util.Partial`. The difference is that the JAX
+    version requires that `func` be specifically a *function* -- and will silently
+    misbehave if given any non-function callable, e.g. [`equinox.nn.MLP`][]. In contrast
+    the Equinox version allows for arbitrary callables.
+    """
+
+    func: Callable[..., _Return]
+    args: tuple[Any, ...]
+    keywords: dict[str, Any]
+
+    def __init__(self, func: Callable[..., _Return], /, *args: Any, **kwargs: Any):
+        """**Arguments:**
+
+        - `func`: the callable to partially apply.
+        - `*args`: any positional arguments to apply.
+        - `**kwargs`: any keyword arguments to apply.
+        """
+        self.func = func
+        self.args = args
+        self.keywords = kwargs
+
+    def __call__(self, *args: Any, **kwargs: Any) -> _Return:
+        """Call the wrapped `self.func`.
+
+        **Arguments:**
+
+        - `*args`: the arguments to apply. Passed after those arguments passed during
+            `__init__`.
+        - `**kwargs`: any keyword arguments to apply.
+
+        **Returns:**
+
+        The result of the wrapped function.
+        """
+        return self.func(*self.args, *args, **kwargs, **self.keywords)
+
+
+class Static(Module):
+    """Wraps a value into a `eqx.field(static=True)`.
+
+    This is useful to treat something as just static metadata with respect to a JAX
+    transformation; for example this is used to return non-arrays from a filtered
+    transform.
+    """
+
+    _leaves: list[Any] = field(static=True)
+    _treedef: PyTreeDef = field(static=True)  # pyright: ignore
+
+    def __init__(self, value: Any):
+        # By flattening, we handle pytrees without `__eq__` methods.
+        # When comparing static metadata for equality, this means we never actually
+        # call `value.__eq__`.
+        self._leaves, self._treedef = jtu.tree_flatten(value)
+
+    @property
+    def value(self):
+        return jtu.tree_unflatten(self._treedef, self._leaves)
diff --git a/equinox/_pretty_print.py b/equinox/_pretty_print.py
index 8f48ac8..6a4a660 100644
--- a/equinox/_pretty_print.py
+++ b/equinox/_pretty_print.py
@@ -23,29 +23,33 @@ def _false(_):
     return False
 
 
-def tree_pformat(
-    pytree: Any,
-    *,
-    width: int = 80,
-    indent: int = 2,
-    short_arrays: bool = True,
-    struct_as_array: bool = False,
-    truncate_leaf: Callable[[Any], bool] = _false,
-) -> str:
-    """Pretty-formats a PyTree as a string, whilst abbreviating JAX arrays.
+@dataclasses.dataclass(slots=True)
+class TreeWLCustom:
+    """Wadler-Lindig customisation for PyTree pretty-printing.
 
-    (This is the function used in `__repr__` of [`equinox.Module`][].)
+    **Arguments:**
+    - `truncate_leaf`: A function `Any -> bool`. Applied to all nodes in the
+        PyTree; all truthy nodes will be truncated to just "<class_name>(...)".
+    - `short_arrays`: If `True`, arrays are abbreviated to their shape and
+      dtype (with weak dtype info prepended to the dtype name).
+    - `struct_as_array`: If `True`, JAX ShapeDtypeStructs are abbreviated to
+      their shape and dtype.
 
-    As [`equinox.tree_pprint`][], but returns the string instead of printing it.
     """
 
-    def custom(obj):
-        if truncate_leaf(obj):
+    truncate_leaf: Callable[[Any], bool]
+
+    short_arrays: bool
+
+    struct_as_array: bool
+
+    def __call__(self, obj: Any) -> wl.AbstractDoc | None:
+        if self.truncate_leaf(obj):
             return wl.TextDoc(f"{type(obj).__name__}(...)")
 
-        if short_arrays:
+        if self.short_arrays:
             if isinstance(obj, jax.Array) or (
-                struct_as_array and isinstance(obj, jax.ShapeDtypeStruct)
+                self.struct_as_array and isinstance(obj, jax.ShapeDtypeStruct)
             ):
                 dtype = obj.dtype.name
                 # Added in JAX 0.4.32 to `ShapeDtypeStruct`
@@ -60,8 +64,32 @@ def tree_pformat(
             obj = _Partial(obj.func, obj.args, obj.keywords)
             return wl.pdoc(obj)
 
+
+def tree_pformat(
+    pytree: Any,
+    *,
+    width: int = 80,
+    indent: int = 2,
+    short_arrays: bool = True,
+    struct_as_array: bool = False,
+    truncate_leaf: Callable[[Any], bool] = _false,
+) -> str:
+    """Pretty-formats a PyTree as a string, whilst abbreviating JAX arrays.
+
+    (This is the function used in `__repr__` of [`equinox.Module`][].)
+
+    As [`equinox.tree_pprint`][], but returns the string instead of printing it.
+    """
     return wl.pformat(
-        pytree, width=width, indent=indent, short_arrays=short_arrays, custom=custom
+        pytree,
+        width=width,
+        indent=indent,
+        short_arrays=short_arrays,
+        custom=TreeWLCustom(
+            truncate_leaf=truncate_leaf,
+            short_arrays=short_arrays,
+            struct_as_array=struct_as_array,
+        ),
     )
 
 
diff --git a/equinox/_sharding.py b/equinox/_sharding.py
index fe75829..5d57c00 100644
--- a/equinox/_sharding.py
+++ b/equinox/_sharding.py
@@ -8,7 +8,8 @@ from ._filters import combine, is_array, partition
 
 
 def filter_shard(
-    x: PyTree[Any], device_or_shardings: jax.Device | jax.sharding.Sharding
+    x: PyTree[Any],
+    device_or_shardings: jax.Device | jax.sharding.Sharding,  # pyright: ignore[reportInvalidTypeForm]
 ):
     """Filtered transform combining `jax.lax.with_sharding_constraint`
     and `jax.device_put`.
diff --git a/equinox/_tree.py b/equinox/_tree.py
index 5b062d1..2ddf4b1 100644
--- a/equinox/_tree.py
+++ b/equinox/_tree.py
@@ -351,7 +351,7 @@ def tree_equal(
 
 def tree_flatten_one_level(
     pytree: PyTree,
-) -> tuple[list[PyTree], PyTreeDef]:
+) -> tuple[list[PyTree], PyTreeDef]:  # pyright: ignore[reportInvalidTypeForm]
     """Returns the immediate subnodes of a PyTree node. If called on a leaf node then it
     will return just that leaf.
 
diff --git a/equinox/internal/__init__.py b/equinox/internal/__init__.py
index 497949f..6ee1b8d 100644
--- a/equinox/internal/__init__.py
+++ b/equinox/internal/__init__.py
@@ -1,8 +1,4 @@
 # Backward compatibility: expose via `equinox.internal`. Now available under `equinox`.
-from .._better_abstract import (
-    AbstractClassVar as AbstractClassVar,
-    AbstractVar as AbstractVar,
-)
 from .._compile_utils import (
     hashable_combine as hashable_combine,
     hashable_partition as hashable_partition,
@@ -20,7 +16,12 @@ from .._errors import (
 )
 from .._eval_shape import cached_filter_eval_shape as cached_filter_eval_shape
 from .._misc import left_broadcast_to as left_broadcast_to
-from .._module import Static as Static
+from .._module import (
+    AbstractClassVar as AbstractClassVar,
+    AbstractVar as AbstractVar,
+    Static as Static,
+)
+from .._pretty_print import TreeWLCustom as TreeWLCustom
 from .._unvmap import (
     unvmap_all as unvmap_all,
     unvmap_all_p as unvmap_all_p,
diff --git a/equinox/internal/_loop/common.py b/equinox/internal/_loop/common.py
index c158846..c4ea451 100644
--- a/equinox/internal/_loop/common.py
+++ b/equinox/internal/_loop/common.py
@@ -141,8 +141,9 @@ def _maybe_set_impl(
     pred, xs, x, *i_dynamic_leaves, i_static, i_treedef, kwargs, makes_false_steps
 ):
     i = combine(i_static, jtu.tree_unflatten(i_treedef, i_dynamic_leaves))
-    x = _select_if_vmap(pred, x, xs.at[i].get(**kwargs), makes_false_steps)
-    return [xs.at[i].set(x, **kwargs)]
+    kwargs_dict = dict(kwargs)
+    x = _select_if_vmap(pred, x, xs.at[i].get(**kwargs_dict), makes_false_steps)
+    return [xs.at[i].set(x, **kwargs_dict)]
 
 
 def _maybe_set_abstract(
@@ -157,7 +158,9 @@ def _maybe_set_jvp(
     pred, xs, x, *i_dynamic_leaves = primals
     _, t_xs, t_x, *_ = tangents
     i = combine(i_static, jtu.tree_unflatten(i_treedef, i_dynamic_leaves))
-    out = _maybe_set(pred, xs, x, i, kwargs=kwargs, makes_false_steps=makes_false_steps)
+    out = _maybe_set(
+        pred, xs, x, i, kwargs=dict(kwargs), makes_false_steps=makes_false_steps
+    )
     if type(t_x) is ad.Zero and type(t_xs) is ad.Zero:
         t_out = t_xs
     else:
@@ -166,7 +169,7 @@ def _maybe_set_jvp(
         if type(t_xs) is ad.Zero:
             t_xs = jnp.zeros(t_xs.aval.shape, t_xs.aval.dtype)  # pyright: ignore
         t_out = _maybe_set(
-            pred, t_xs, t_x, i, kwargs=kwargs, makes_false_steps=makes_false_steps
+            pred, t_xs, t_x, i, kwargs=dict(kwargs), makes_false_steps=makes_false_steps
         )
     return [out], [t_out]
 
@@ -202,7 +205,7 @@ def _maybe_set_transpose(
         if type(ct_out) is ad.Zero:
             ct_x = None
         else:
-            ct_x = ct_out.at[i].get(**kwargs)
+            ct_x = ct_out.at[i].get(**dict(kwargs))
             ct_x = _select_if_vmap(pred, ct_x, jnp.zeros_like(ct_x), makes_false_steps)
     else:
         ct_x = None
@@ -252,7 +255,7 @@ def _maybe_set(pred, xs, x, i, *, kwargs, makes_false_steps):
         *i_dynamic_leaves,
         i_static=i_static,
         i_treedef=i_treedef,
-        kwargs=kwargs,
+        kwargs=tuple(kwargs.items()),
         makes_false_steps=makes_false_steps,
     )
     return out
diff --git a/equinox/internal/_primitive.py b/equinox/internal/_primitive.py
index 9eb0ff9..aab7265 100644
--- a/equinox/internal/_primitive.py
+++ b/equinox/internal/_primitive.py
@@ -318,7 +318,7 @@ def create_vprim(name: str, impl, abstract_eval, jvp, transpose):
                 __axis_size=axis_size,
                 __axis_name=axis_name,
                 __batch_axes=batch_axes,
-                params=params,
+                params=tuple(params.items()),
             )
             batch_axes_out = jtu.tree_map(lambda _: 0, out)
         return out, batch_axes_out
@@ -337,7 +337,7 @@ def create_vprim(name: str, impl, abstract_eval, jvp, transpose):
 
 
 def _vprim_impl(*inputs, prim, __axis_size, __axis_name, __batch_axes, params):
-    impl = ft.partial(_vprim_impl_registry[prim], **params)
+    impl = ft.partial(_vprim_impl_registry[prim], **dict(params))
     impl = jax.vmap(
         impl, in_axes=__batch_axes, axis_size=__axis_size, axis_name=__axis_name
     )
@@ -360,7 +360,7 @@ def _vprim_abstract_eval(*inputs, prim, __axis_size, __axis_name, __batch_axes,
         jax.core.mapped_aval(__axis_size, b, x) for x, b in zip(inputs, __batch_axes)
     ]
     abstract_eval = _vprim_abstract_eval_registry[prim]
-    outs = abstract_eval(*inputs, **params)
+    outs = abstract_eval(*inputs, **dict(params))
     outs = [_unmapped_aval(__axis_size, __axis_name, 0, x) for x in outs]
     return outs
 
@@ -393,7 +393,7 @@ def _vprim_jvp(
     assert len(tangents) == len(__batch_axes)
     tangents = [_resolve_zeros_t(t, b) for t, b in zip(tangents, __batch_axes)]
     batch_axes_t = [_resolve_zeros_b(t, b) for t, b in zip(tangents, __batch_axes)]
-    jvp = ft.partial(_vprim_jvp_registry[prim], **params)
+    jvp = ft.partial(_vprim_jvp_registry[prim], **dict(params))
     jvp = jax.vmap(
         jvp,
         in_axes=(__batch_axes, batch_axes_t),
@@ -431,7 +431,7 @@ def _vprim_transpose(
     batch_axes = [_resolve_undefined_b(i, b) for i, b in zip(inputs, __batch_axes)]
 
     def _transpose(*_inputs):
-        _outputs = _vprim_transpose_registry[prim](*_inputs, **params)
+        _outputs = _vprim_transpose_registry[prim](*_inputs, **dict(params))
         # `Zero` is not a JAX type -- it's an internal AD thing -- so we shouldn't pass
         # it across the `vmap` boundary. In particular JAX won't apply the out batch
         # axis to it.
diff --git a/equinox/nn/_activations.py b/equinox/nn/_activations.py
index 054a900..e265fbf 100644
--- a/equinox/nn/_activations.py
+++ b/equinox/nn/_activations.py
@@ -6,7 +6,7 @@ from .._module import Module
 from ._misc import named_scope
 
 
-class PReLU(Module, strict=True):
+class PReLU(Module):
     """PReLU activation function.
 
     This is the elementwise function `x -> max(x, 0) + α * min(x, 0)`.
diff --git a/equinox/nn/_attention.py b/equinox/nn/_attention.py
index 5a50cea..a51d838 100644
--- a/equinox/nn/_attention.py
+++ b/equinox/nn/_attention.py
@@ -78,7 +78,7 @@ else:
     _Mask = Bool[Array, "q_seq kv_seq"] | Bool[Array, "num_heads q_seq kv_seq"]
 
 
-class MultiheadAttention(Module, strict=True):
+class MultiheadAttention(Module):
     r"""
     Computes
 
diff --git a/equinox/nn/_batch_norm.py b/equinox/nn/_batch_norm.py
index eec94a7..56e0ce3 100644
--- a/equinox/nn/_batch_norm.py
+++ b/equinox/nn/_batch_norm.py
@@ -1,9 +1,11 @@
+import warnings
 from collections.abc import Hashable, Sequence
+from typing import Literal
 
 import jax
 import jax.lax as lax
 import jax.numpy as jnp
-from jaxtyping import Array, Bool, Float, PRNGKeyArray
+from jaxtyping import Array, Bool, Float, Int, PRNGKeyArray
 
 from .._misc import default_floating_dtype
 from .._module import field
@@ -12,7 +14,7 @@ from ._sequential import StatefulLayer
 from ._stateful import State, StateIndex
 
 
-class BatchNorm(StatefulLayer, strict=True):
+class BatchNorm(StatefulLayer):
     r"""Computes a mean and standard deviation over the batch and spatial
     dimensions of an array, and uses these to normalise the whole array. Optionally
     applies a channelwise affine transformation afterwards.
@@ -40,20 +42,57 @@ class BatchNorm(StatefulLayer, strict=True):
     statistics updated. During inference then just the running statistics are used.
     Whether the model is in training or inference mode should be toggled using
     [`equinox.nn.inference_mode`][].
+
+    With `mode = "batch"` during training the batch mean and variance are used
+    for normalization.  For inference the exponential running mean and unbiased
+    variance are used for normalization. This is in line with how other machine
+    learning packages (e.g. PyTorch, flax, haiku) implement batch norm.
+
+    With `mode = "ema"` exponential running means and variances are kept.  During
+    training the batch statistics are used to fill in the running statistics until
+    they are populated.  During inference the running statistics are used for
+    normalization.
+
+    ??? cite
+
+        [Batch Normalization: Accelerating Deep Network Training by Reducing
+         Internal Covariate Shift](https://arxiv.org/abs/1502.03167)
+
+        ```bibtex
+        @article{DBLP:journals/corr/IoffeS15,
+            author = {Sergey Ioffe and Christian Szegedy},
+            title = {Batch Normalization: Accelerating Deep Network Training by Reducing
+                     Internal Covariate Shift},
+            journal = {CoRR},
+            volume = {abs/1502.03167},
+            year = {2015},
+            url = {http://arxiv.org/abs/1502.03167},
+            eprinttype = {arXiv},
+            eprint = {1502.03167},
+            timestamp = {Mon, 13 Aug 2018 16:47:06 +0200},
+            biburl = {https://dblp.org/rec/journals/corr/IoffeS15.bib},
+            bibsource = {dblp computer science bibliography, https://dblp.org}
+        }
+        ```
     """  # noqa: E501
 
     weight: Float[Array, "input_size"] | None
     bias: Float[Array, "input_size"] | None
-    first_time_index: StateIndex[Bool[Array, ""]]
-    state_index: StateIndex[
-        tuple[Float[Array, "input_size"], Float[Array, "input_size"]]
-    ]
+    ema_first_time_index: None | StateIndex[Bool[Array, ""]]
+    ema_state_index: (
+        None | StateIndex[tuple[Float[Array, "input_size"], Float[Array, "input_size"]]]
+    )
+    batch_counter: None | StateIndex[Int[Array, ""]]
+    batch_state_index: (
+        None | StateIndex[tuple[Float[Array, "input_size"], Float[Array, "input_size"]]]
+    )
     axis_name: Hashable | Sequence[Hashable]
     inference: bool
     input_size: int = field(static=True)
     eps: float = field(static=True)
     channelwise_affine: bool = field(static=True)
     momentum: float = field(static=True)
+    mode: Literal["ema", "batch"] = field(static=True)
 
     def __init__(
         self,
@@ -64,6 +103,7 @@ class BatchNorm(StatefulLayer, strict=True):
         momentum: float = 0.99,
         inference: bool = False,
         dtype=None,
+        mode: Literal["ema", "batch", "legacy"] = "legacy",
     ):
         """**Arguments:**
 
@@ -85,7 +125,19 @@ class BatchNorm(StatefulLayer, strict=True):
             if `channelwise_affine` is `True`. Defaults to either
             `jax.numpy.float32` or `jax.numpy.float64` depending on whether JAX is in
             64-bit mode.
+        - `mode`: The variant of batch norm to use, either 'ema' or 'batch'.
         """
+        if mode == "legacy":
+            mode = "ema"
+            warnings.warn(
+                "When `eqx.nn.BatchNorm(..., mode=...)` is unspecified it defaults to "
+                "'ema', for backward compatibility. This typically has a performance "
+                "impact, and for new code the user is encouraged to use 'batch' "
+                "instead. See `https://github.com/patrick-kidger/equinox/issues/659`."
+            )
+        if mode not in {"ema", "batch"}:
+            raise ValueError("Invalid mode, must be 'ema' or 'batch'.")
+        self.mode = mode
         dtype = default_floating_dtype() if dtype is None else dtype
         if channelwise_affine:
             self.weight = jnp.ones((input_size,), dtype=dtype)
@@ -93,12 +145,24 @@ class BatchNorm(StatefulLayer, strict=True):
         else:
             self.weight = None
             self.bias = None
-        self.first_time_index = StateIndex(jnp.array(True))
-        init_buffers = (
-            jnp.empty((input_size,), dtype=dtype),
-            jnp.empty((input_size,), dtype=dtype),
-        )
-        self.state_index = StateIndex(init_buffers)
+        if mode == "ema":
+            self.ema_first_time_index = StateIndex(jnp.array(True))
+            init_buffers = (
+                jnp.empty((input_size,), dtype=dtype),
+                jnp.empty((input_size,), dtype=dtype),
+            )
+            self.ema_state_index = StateIndex(init_buffers)
+            self.batch_counter = None
+            self.batch_state_index = None
+        else:
+            self.batch_counter = StateIndex(jnp.array(0))
+            init_hidden = (
+                jnp.zeros((input_size,), dtype=dtype),
+                jnp.ones((input_size,), dtype=dtype),
+            )
+            self.batch_state_index = StateIndex(init_hidden)
+            self.ema_first_time_index = None
+            self.ema_state_index = None
         self.inference = inference
         self.axis_name = axis_name
         self.input_size = input_size
@@ -138,32 +202,18 @@ class BatchNorm(StatefulLayer, strict=True):
         A `NameError` if no `vmap`s are placed around this operation, or if this vmap
         does not have a matching `axis_name`.
         """
+        del key
 
         if inference is None:
             inference = self.inference
-        if inference:
-            running_mean, running_var = state.get(self.state_index)
-        else:
 
-            def _stats(y):
-                mean = jnp.mean(y)
-                mean = lax.pmean(mean, self.axis_name)
-                var = jnp.mean((y - mean) * jnp.conj(y - mean))
-                var = lax.pmean(var, self.axis_name)
-                var = jnp.maximum(0.0, var)
-                return mean, var
-
-            first_time = state.get(self.first_time_index)
-            state = state.set(self.first_time_index, jnp.array(False))
-
-            batch_mean, batch_var = jax.vmap(_stats)(x)
-            running_mean, running_var = state.get(self.state_index)
-            momentum = self.momentum
-            running_mean = (1 - momentum) * batch_mean + momentum * running_mean
-            running_var = (1 - momentum) * batch_var + momentum * running_var
-            running_mean = lax.select(first_time, batch_mean, running_mean)
-            running_var = lax.select(first_time, batch_var, running_var)
-            state = state.set(self.state_index, (running_mean, running_var))
+        def _stats(y):
+            mean = jnp.mean(y)
+            mean = lax.pmean(mean, self.axis_name)
+            var = jnp.mean((y - mean) * jnp.conj(y - mean))
+            var = lax.pmean(var, self.axis_name)
+            var = jnp.maximum(0.0, var)
+            return mean, var
 
         def _norm(y, m, v, w, b):
             out = (y - m) / jnp.sqrt(v + self.eps)
@@ -171,5 +221,45 @@ class BatchNorm(StatefulLayer, strict=True):
                 out = out * w + b
             return out
 
-        out = jax.vmap(_norm)(x, running_mean, running_var, self.weight, self.bias)
+        if self.mode == "ema":
+            assert self.ema_first_time_index is not None
+            assert self.ema_state_index is not None
+            if inference:
+                mean, var = state.get(self.ema_state_index)
+            else:
+                first_time = state.get(self.ema_first_time_index)
+                state = state.set(self.ema_first_time_index, jnp.array(False))
+                batch_mean, batch_var = jax.vmap(_stats)(x)
+                running_mean, running_var = state.get(self.ema_state_index)
+                momentum = self.momentum
+                mean = (1 - momentum) * batch_mean + momentum * running_mean
+                var = (1 - momentum) * batch_var + momentum * running_var
+                # since jnp.array(0) == False
+                mean = lax.select(first_time, batch_mean, mean)
+                var = lax.select(first_time, batch_var, var)
+                state = state.set(self.ema_state_index, (mean, var))
+        else:
+            assert self.batch_state_index is not None
+            assert self.batch_counter is not None
+            counter = state.get(self.batch_counter)
+            hidden_mean, hidden_var = state.get(self.batch_state_index)
+            if inference:
+                # Zero-debias approach: mean = hidden_mean / (1 - momentum^counter)
+                # For simplicity we do the minimal version here (no warmup).
+                scale = jnp.where(counter == 0, 1, 1 - self.momentum**counter)
+                mean = hidden_mean / scale
+                var = hidden_var / scale
+            else:
+                mean, var = jax.vmap(_stats)(x)
+                new_counter = counter + 1
+                new_hidden_mean = hidden_mean * self.momentum + mean * (
+                    1 - self.momentum
+                )
+                new_hidden_var = hidden_var * self.momentum + var * (1 - self.momentum)
+                state = state.set(self.batch_counter, new_counter)
+                state = state.set(
+                    self.batch_state_index, (new_hidden_mean, new_hidden_var)
+                )
+
+        out = jax.vmap(_norm)(x, mean, var, self.weight, self.bias)
         return out, state
diff --git a/equinox/nn/_conv.py b/equinox/nn/_conv.py
index 8fe3edb..9e8be40 100644
--- a/equinox/nn/_conv.py
+++ b/equinox/nn/_conv.py
@@ -67,7 +67,7 @@ def _padding_mode_init(padding_mode: str) -> str:
     return padding_mode
 
 
-class Conv(Module, strict=True):
+class Conv(Module):
     """General N-dimensional convolution."""
 
     num_spatial_dims: int = field(static=True)
@@ -352,7 +352,7 @@ class Conv3d(Conv):
         )
 
 
-class ConvTranspose(Module, strict=True):
+class ConvTranspose(Module):
     """General N-dimensional transposed convolution."""
 
     num_spatial_dims: int = field(static=True)
@@ -576,7 +576,7 @@ class ConvTranspose(Module, strict=True):
         stride = np.expand_dims(self.stride, axis=1)
         pad_width = np.insert(padding_t // stride, 0, 0, axis=0)
         x = jnp.pad(x, pad_width, mode="wrap")
-        padding_t = tuple((p[0].item(), p[1].item()) for p in padding_t % stride)
+        padding_t = tuple((p[0].item(), p[1].item()) for p in padding_t % stride)  # pyright: ignore[reportIndexIssue]
         return x, padding_t
 
     @named_scope("eqx.nn.ConvTranspose")
diff --git a/equinox/nn/_dropout.py b/equinox/nn/_dropout.py
index 13c3011..3f582f5 100644
--- a/equinox/nn/_dropout.py
+++ b/equinox/nn/_dropout.py
@@ -9,7 +9,7 @@ from .._module import Module
 from ._misc import named_scope
 
 
-class Dropout(Module, strict=True):
+class Dropout(Module):
     """Applies dropout.
 
     Note that this layer behaves differently during training and inference. During
diff --git a/equinox/nn/_embedding.py b/equinox/nn/_embedding.py
index adde22b..dcb8564 100644
--- a/equinox/nn/_embedding.py
+++ b/equinox/nn/_embedding.py
@@ -17,7 +17,7 @@ internal_rope_embedding_cache: dict[tuple[int, Any], tuple[Array, Array]] = {}
 cache_clears.append(internal_rope_embedding_cache.clear)
 
 
-class Embedding(Module, strict=True):
+class Embedding(Module):
     """A simple lookup table that stores embeddings of a fixed size."""
 
     num_embeddings: int = field(static=True)
@@ -106,7 +106,7 @@ class Embedding(Module, strict=True):
             )
 
 
-class RotaryPositionalEmbedding(Module, strict=True):
+class RotaryPositionalEmbedding(Module):
     """A rotary positional encoding module, as described in the paper
     "RoFormer: Enhanced Transformer with Rotary Position Embedding". While this module
     can be used in any context, it is particularly useful for providing positional
diff --git a/equinox/nn/_linear.py b/equinox/nn/_linear.py
index 2c2d70e..d3a196d 100644
--- a/equinox/nn/_linear.py
+++ b/equinox/nn/_linear.py
@@ -10,7 +10,7 @@ from .._module import field, Module
 from ._misc import default_init, named_scope
 
 
-class Linear(Module, strict=True):
+class Linear(Module):
     """Performs a linear transformation."""
 
     weight: Array
@@ -106,7 +106,7 @@ class Linear(Module, strict=True):
 _T = TypeVar("_T")
 
 
-class Identity(Module, strict=True):
+class Identity(Module):
     """Identity operation that does nothing. Sometimes useful as a placeholder for
     another Module.
     """
diff --git a/equinox/nn/_mlp.py b/equinox/nn/_mlp.py
index aab22dc..1b7f5ed 100644
--- a/equinox/nn/_mlp.py
+++ b/equinox/nn/_mlp.py
@@ -21,7 +21,7 @@ _identity = doc_repr(lambda x: x, "lambda x: x")
 _relu = doc_repr(jnn.relu, "<function relu>")
 
 
-class MLP(Module, strict=True):
+class MLP(Module):
     """Standard Multi-Layer Perceptron; also known as a feed-forward network.
 
     !!! faq
diff --git a/equinox/nn/_normalisation.py b/equinox/nn/_normalisation.py
index a96d522..dbfd4b5 100644
--- a/equinox/nn/_normalisation.py
+++ b/equinox/nn/_normalisation.py
@@ -14,7 +14,7 @@ from ._misc import named_scope
 from ._stateful import State
 
 
-class LayerNorm(Module, strict=True):
+class LayerNorm(Module):
     r"""
     Computes a mean and standard deviation over the whole input array, and uses these
     to normalise the whole array. Optionally applies an elementwise affine
@@ -161,7 +161,7 @@ class LayerNorm(Module, strict=True):
             return out.astype(orig_dtype), state
 
 
-class GroupNorm(Module, strict=True):
+class GroupNorm(Module):
     r"""
     Splits the first dimension ("channels") into groups of fixed size. Computes a mean
     and standard deviation over the contents of each group, and uses these to normalise
@@ -283,7 +283,7 @@ class GroupNorm(Module, strict=True):
             return out.astype(orig_dtype), state
 
 
-class RMSNorm(Module, strict=True):
+class RMSNorm(Module):
     r"""
     A simplified version of LayerNorm which rescales the inputs, but does not center
     them. Optionally applies a learned reweighting of the transformed array afterward.
diff --git a/equinox/nn/_pool.py b/equinox/nn/_pool.py
index efcc66e..d00855f 100644
--- a/equinox/nn/_pool.py
+++ b/equinox/nn/_pool.py
@@ -11,7 +11,7 @@ from .._module import field, Module
 from ._misc import all_sequences, named_scope
 
 
-class Pool(Module, strict=True):
+class Pool(Module):
     """General N-dimensional downsampling over a sliding window."""
 
     init: int | float | Array
@@ -470,7 +470,7 @@ def _adaptive_pool1d(
     return outputs
 
 
-class AdaptivePool(Module, strict=True):
+class AdaptivePool(Module):
     """General N dimensional adaptive downsampling to a target shape."""
 
     target_shape: Sequence[int] = field(static=True)
diff --git a/equinox/nn/_rnn.py b/equinox/nn/_rnn.py
index 94c6700..0ccc0c0 100644
--- a/equinox/nn/_rnn.py
+++ b/equinox/nn/_rnn.py
@@ -10,7 +10,7 @@ from .._module import field, Module
 from ._misc import default_init, named_scope
 
 
-class GRUCell(Module, strict=True):
+class GRUCell(Module):
     """A single step of a Gated Recurrent Unit (GRU).
 
     !!! example
@@ -108,7 +108,7 @@ class GRUCell(Module, strict=True):
         return new + inp * (hidden - new)
 
 
-class LSTMCell(Module, strict=True):
+class LSTMCell(Module):
     """A single step of a Long-Short Term Memory unit (LSTM).
 
     !!! example
diff --git a/equinox/nn/_sequential.py b/equinox/nn/_sequential.py
index c97dfa2..328c2c9 100644
--- a/equinox/nn/_sequential.py
+++ b/equinox/nn/_sequential.py
@@ -4,14 +4,13 @@ from typing import Any, overload
 import jax.random as jr
 from jaxtyping import Array, PRNGKeyArray
 
-from .._better_abstract import AbstractClassVar
 from .._custom_types import sentinel
-from .._module import Module, StrictConfig
+from .._module import AbstractClassVar, Module
 from ._misc import named_scope
 from ._stateful import State
 
 
-class StatefulLayer(Module, strict=StrictConfig(allow_abstract_name=True)):
+class StatefulLayer(Module):
     """An abstract base class, used by [`equinox.nn.Sequential`][], to mark that a
     layer might be stateful. If `Sequential` sees that a layer inherits from
     `StatefulLayer`, then it will call `layer.is_stateful()` to check whether to
@@ -40,7 +39,7 @@ class StatefulLayer(Module, strict=StrictConfig(allow_abstract_name=True)):
     __call__: AbstractClassVar[Callable]
 
 
-class Sequential(StatefulLayer, strict=StrictConfig(allow_method_override=True)):
+class Sequential(StatefulLayer):
     """A sequence of [`equinox.Module`][]s applied in order.
 
     !!! note
@@ -124,7 +123,7 @@ class Sequential(StatefulLayer, strict=StrictConfig(allow_method_override=True))
         return len(self.layers)
 
 
-class Lambda(Module, strict=True):
+class Lambda(Module):
     """Wraps a callable (e.g. an activation function) for use with
     [`equinox.nn.Sequential`][].
 
diff --git a/equinox/nn/_shared.py b/equinox/nn/_shared.py
index 33d253a..4f8b52f 100644
--- a/equinox/nn/_shared.py
+++ b/equinox/nn/_shared.py
@@ -14,7 +14,7 @@ class SharedNode:
         return "SharedNode"
 
 
-class Shared(Module, strict=True):
+class Shared(Module):
     """Used to tie together multiple nodes across a PyTree.
 
     Note that Equinox modules are Py**Trees** -- so the same layer, appearing in two
diff --git a/equinox/nn/_spectral_norm.py b/equinox/nn/_spectral_norm.py
index 9b3722a..ba65dc2 100644
--- a/equinox/nn/_spectral_norm.py
+++ b/equinox/nn/_spectral_norm.py
@@ -27,7 +27,7 @@ def _power_iteration(weight, u, v, eps):
 _Layer = TypeVar("_Layer")
 
 
-class SpectralNorm(StatefulLayer, Generic[_Layer], strict=True):
+class SpectralNorm(StatefulLayer, Generic[_Layer]):
     """Applies spectral normalisation to a given parameter.
 
     Given a weight matrix $W$, and letting $σ(W)$ denote (an approximation to) its
diff --git a/equinox/nn/_stateful.py b/equinox/nn/_stateful.py
index f7066bd..71afb19 100644
--- a/equinox/nn/_stateful.py
+++ b/equinox/nn/_stateful.py
@@ -30,7 +30,7 @@ class _Sentinel(Module):
 _sentinel = _Sentinel()
 
 
-class StateIndex(Module, Generic[_Value], strict=True):
+class StateIndex(Module, Generic[_Value]):
     """This wraps together (a) a unique dictionary key used for looking up a stateful
     value, and (b) how that stateful value should be initialised.
 
diff --git a/equinox/nn/_weight_norm.py b/equinox/nn/_weight_norm.py
index 8e3142a..d46803b 100644
--- a/equinox/nn/_weight_norm.py
+++ b/equinox/nn/_weight_norm.py
@@ -21,7 +21,7 @@ def _norm_except_axis(v: Array, norm: Callable[[Array], Scalar], axis: int | Non
         return jax.vmap(norm, in_axes=axis, out_axes=axis)(v)
 
 
-class WeightNorm(Module, Generic[_Layer], strict=True):
+class WeightNorm(Module, Generic[_Layer]):
     r"""Applies weight normalisation to a given parameter.
 
     Given the 2D weight matrix
diff --git a/pyproject.toml b/pyproject.toml
index 796be6a..0c440ce 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -27,7 +27,7 @@ name = "equinox"
 readme = "README.md"
 requires-python = ">=3.10"
 urls = {repository = "https://github.com/patrick-kidger/equinox"}
-version = "0.12.2"
+version = "0.13.0"
 
 [project.optional-dependencies]
 docs = [


Test patch:
diff --git a/tests/test_abstract.py b/tests/test_abstract.py
index 03cfa1f..d0b89b6 100644
--- a/tests/test_abstract.py
+++ b/tests/test_abstract.py
@@ -20,10 +20,10 @@ def test_abstract_method():
         pass
 
     with pytest.raises(TypeError, match="abstract method"):
-        A()
+        A()  # pyright: ignore[reportAbstractUsage]
 
     with pytest.raises(TypeError, match="abstract method"):
-        B()
+        B()  # pyright: ignore[reportAbstractUsage]
 
 
 def test_abstract_attribute_stringified():
diff --git a/tests/test_module.py b/tests/test_module.py
index 680341d..4b209cb 100644
--- a/tests/test_module.py
+++ b/tests/test_module.py
@@ -1,4 +1,4 @@
-import abc
+import contextlib
 import dataclasses
 import doctest
 import functools as ft
@@ -6,7 +6,7 @@ import inspect
 import pickle
 from collections.abc import Callable
 from dataclasses import InitVar
-from typing import Any
+from typing import Any, Generic, TypeVar
 
 import equinox as eqx
 import equinox.internal as eqxi
@@ -32,7 +32,7 @@ def test_module_not_enough_attributes():
         def __init__(self) -> None:
             pass
 
-    with pytest.raises(ValueError):
+    with pytest.raises(TypeError):
         MyModule2()
     with pytest.raises(TypeError):
         MyModule2(1)  # pyright: ignore
@@ -210,22 +210,26 @@ def test_method_access_during_init():
 def test_static_field(new):
     if new:
         static_field = ft.partial(eqx.field, static=True)
+        ctx = contextlib.nullcontext()
     else:
         static_field = eqx.static_field
+        ctx = pytest.warns(match="deprecated")
 
-    class MyModule(eqx.Module):
-        field1: int
-        field2: int = static_field()
-        field3: int = static_field(default=3)
+    with ctx:
 
-    m = MyModule(1, 2)
-    flat, treedef = jtu.tree_flatten(m)
-    assert len(flat) == 1
-    assert flat[0] == 1
-    rm = jtu.tree_unflatten(treedef, flat)
-    assert rm.field1 == 1
-    assert rm.field2 == 2
-    assert rm.field3 == 3
+        class MyModule(eqx.Module):
+            field1: int
+            field2: int = static_field()
+            field3: int = static_field(default=3)
+
+        m = MyModule(1, 2)
+        flat, treedef = jtu.tree_flatten(m)
+        assert len(flat) == 1
+        assert flat[0] == 1
+        rm = jtu.tree_unflatten(treedef, flat)
+        assert rm.field1 == 1
+        assert rm.field2 == 2
+        assert rm.field3 == 3
 
 
 def test_converter():
@@ -636,243 +640,6 @@ def test_check_init_no_assignment():
         A(1)
 
 
-def test_strict_noerrors():
-    class Abstract(eqx.Module, strict=True):
-        @abc.abstractmethod
-        def foo(self, x):
-            pass
-
-    class Concrete1(Abstract, strict=True):
-        def foo(self, x):
-            return x + 1
-
-    class Concrete2(Abstract):
-        def foo(self, x):
-            return x + 2
-
-    class Concrete3(Concrete2):
-        def foo(self, x):
-            return x + 3
-
-    class Abstract1(eqx.Module, strict=True):
-        bar: eqx.AbstractVar[int]
-
-        @abc.abstractmethod
-        def foo(self):
-            pass
-
-    class Abstract2(eqx.Module, strict=True):
-        bar: eqx.AbstractClassVar[int]
-
-        @abc.abstractmethod
-        def foo(self):
-            pass
-
-    class Abstract3(eqx.Module, strict=True):
-        bar: int
-
-        @abc.abstractmethod
-        def foo(self):
-            pass
-
-
-def test_strict_non_module_base():
-    class NotAModule:
-        pass
-
-    with pytest.raises(
-        TypeError, match="Strict `eqx.Module`s must only inherit from other strict"
-    ):
-
-        class MyModule(NotAModule, eqx.Module, strict=True):
-            pass
-
-    class NotAStrictModule(eqx.Module):
-        pass
-
-    with pytest.raises(
-        TypeError, match="Strict `eqx.Module`s must only inherit from other strict"
-    ):
-
-        class MyModule2(NotAStrictModule, eqx.Module, strict=True):
-            pass
-
-
-def test_strict_concrete_is_final():
-    class Concrete(eqx.Module, strict=True):
-        pass
-
-    with pytest.raises(
-        TypeError, match="very strict `eqx.Module` must be either abstract or final"
-    ):
-
-        class Concrete2(Concrete, strict=True):
-            pass
-
-
-@pytest.mark.parametrize("super_init_or_attr", ("init", "attr"))
-@pytest.mark.parametrize("sub_init_or_attr", ("init", "attr"))
-@pytest.mark.parametrize("abstract_flag", (False, True))
-def test_strict_init(super_init_or_attr, sub_init_or_attr, abstract_flag):
-    class Abstract(eqx.Module, strict=eqx.StrictConfig(force_abstract=abstract_flag)):
-        if super_init_or_attr == "init":
-
-            def __init__(self):
-                pass
-
-        else:
-            x: int
-
-        if not abstract_flag:
-
-            @abc.abstractmethod
-            def foo(self):
-                pass
-
-    with pytest.raises(
-        TypeError, match="For readability, any custom `__init__` method, and all fields"
-    ):
-
-        class Concrete1(Abstract, strict=True):
-            if sub_init_or_attr == "init":
-
-                def __init__(self):
-                    pass
-
-            else:
-                y: str
-
-            def foo(self):
-                pass
-
-
-def test_strict_init_in_abstract():
-    class AbstractA(eqx.Module):
-        a: int
-        b: int
-
-        def __init__(self, x):
-            self.a = x
-            self.b = x
-
-        @abc.abstractmethod
-        def foo(self):
-            pass
-
-    class ConcreteB(AbstractA):
-        def foo(self):
-            pass
-
-    # No way to teach pyright about Equinox's different behaviour for `__init__` than
-    # is provided by dataclasses.
-    instance = ConcreteB(x=1)  # pyright: ignore
-    assert instance.a == 1
-    assert instance.b == 1
-
-
-def test_strict_init_transitive():
-    class AbstractA(eqx.Module, strict=eqx.StrictConfig(force_abstract=True)):
-        x: int
-
-    class AbstractB(AbstractA, strict=eqx.StrictConfig(force_abstract=True)):
-        pass
-
-    with pytest.raises(
-        TypeError, match="For readability, any custom `__init__` method, and all fields"
-    ):
-
-        class C(AbstractB, strict=True):
-            y: str
-
-
-def test_strict_abstract_name():
-    class Abstract(eqx.Module, strict=eqx.StrictConfig(force_abstract=True)):
-        pass
-
-    class _Abstract(eqx.Module, strict=eqx.StrictConfig(force_abstract=True)):
-        pass
-
-    with pytest.raises(
-        TypeError, match="Abstract strict `eqx.Module`s must be named starting with"
-    ):
-
-        class NotAbstract(eqx.Module, strict=eqx.StrictConfig(force_abstract=True)):
-            pass
-
-
-def test_strict_method_reoverride():
-    class AbstractA(eqx.Module, strict=True):
-        @abc.abstractmethod
-        def foo(self, x):
-            pass
-
-    class AbstractB(AbstractA, strict=True):
-        def foo(self, x):
-            pass
-
-        @abc.abstractmethod
-        def bar(self, x):
-            pass
-
-    with pytest.raises(
-        TypeError, match="Strict `eqx.Module`s cannot override concrete"
-    ):
-
-        class C(AbstractB, strict=True):
-            def foo(self, x):
-                pass
-
-            def bar(self, x):
-                pass
-
-    class AbstractD(eqx.Module, strict=eqx.StrictConfig(force_abstract=True)):
-        foo = 4
-
-    class AbstractD2(AbstractD, strict=eqx.StrictConfig(force_abstract=True)):
-        pass
-
-    with pytest.raises(
-        TypeError, match="Strict `eqx.Module`s cannot override non-methods"
-    ):
-
-        class E(AbstractD, strict=True):
-            def foo(self):
-                pass
-
-    with pytest.raises(
-        TypeError, match="Strict `eqx.Module`s cannot override non-methods"
-    ):
-
-        class E2(AbstractD2, strict=True):
-            def foo(self):
-                pass
-
-
-def test_strict_default():
-    class AbstractA(eqx.Module, strict=eqx.StrictConfig(force_abstract=True)):
-        def foo(self) -> int:
-            return 4
-
-    class AbstractB(
-        AbstractA,
-        strict=eqx.StrictConfig(force_abstract=True, allow_method_override=True),
-    ):
-        def foo(self) -> int:
-            return 5
-
-    with pytest.raises(
-        TypeError, match="Strict `eqx.Module`s cannot override concrete"
-    ):
-
-        class C1(AbstractB, strict=True):
-            def foo(self):
-                return 6
-
-    class C2(AbstractB, strict=eqx.StrictConfig(allow_method_override=True)):
-        def foo(self):
-            return 7
-
-
 def test_post_init_warning():
     class A(eqx.Module):
         called = False
@@ -962,7 +729,7 @@ def test_init_fields():
     class A(eqx.Module):
         x: int = eqx.field(init=False)
 
-    with pytest.raises(ValueError, match="The following fields were not initialised"):
+    with pytest.raises(TypeError, match="The following fields were not initialised"):
         A()
 
     class B(eqx.Module):
@@ -971,7 +738,7 @@ def test_init_fields():
         def __post_init__(self):
             pass
 
-    with pytest.raises(ValueError, match="The following fields were not initialised"):
+    with pytest.raises(TypeError, match="The following fields were not initialised"):
         B()
 
     class C(eqx.Module):
@@ -983,7 +750,7 @@ def test_init_fields():
                 self.x = 1
 
     C(flag=True)
-    with pytest.raises(ValueError, match="The following fields were not initialised"):
+    with pytest.raises(TypeError, match="The following fields were not initialised"):
         C(flag=False)
 
 
@@ -1086,8 +853,6 @@ def test_module_setattr():
     Foo.g = g  # pyright: ignore
     assert Foo.f is f2
     assert Foo.g is g  # pyright: ignore
-    assert type(Foo.__dict__["f"]).__name__ == "_wrap_method"
-    assert type(Foo.__dict__["g"]).__name__ == "_wrap_method"
 
 
 # See https://github.com/patrick-kidger/equinox/issues/206
@@ -1316,3 +1081,101 @@ def test_doctest():
     assert any("baz" in str(test) for test in tests)
     assert any("biz" in str(test) for test in tests)
     assert any("buz" in str(test) for test in tests)
+
+
+def test_reassign():
+    class A(eqx.Module):
+        x: int
+
+        def __init__(self):
+            self.x = 3
+            self.x = 4
+
+    A()
+
+
+def test_not_weakrefable():
+    class A(eqx.Module):
+        __slots__ = ("x",)
+        x: int
+
+        def __init__(self):
+            self.x = 3
+            self.x = 4
+
+    A()
+
+
+def test_classmethod():
+    class A(eqx.Module):
+        x: int
+
+        @classmethod
+        def make(cls, x):
+            return cls(x)
+
+    a = A(x=1).make(x=2)
+    assert a.x == 2
+
+
+def test_staticmethod():
+    class A(eqx.Module):
+        x: int
+
+        @staticmethod
+        def add_five(x: int):
+            return x + 5
+
+    assert A(x=1).add_five(3) == 8
+
+
+def test_property():
+    class A(eqx.Module):
+        x: int
+
+        @property
+        def add_five(self):
+            return self.x + 5
+
+    assert A(x=1).add_five == 6
+
+
+def test_another_classes_method():
+    class A(eqx.Module):
+        x: int
+
+        def add_five(self):
+            return self.x + 5
+
+    class B(eqx.Module):
+        add_five = A(3).add_five
+
+    assert B.add_five() == 8
+    assert B().add_five() == 8
+
+
+def test_another_classes_classmethod():
+    class A(eqx.Module):
+        x: int
+
+        @classmethod
+        def make(cls):
+            assert cls is A
+            return cls(x=3)
+
+    class B(eqx.Module):
+        make = A.make
+
+    assert B.make() == A(3)
+    assert B().make() == A(3)
+
+
+_T = TypeVar("_T")
+
+
+def test_orig_class():
+    class A(eqx.Module, Generic[_T]):
+        pass
+
+    a = A[int]()
+    assert a.__orig_class__ == A[int]
diff --git a/tests/test_nn.py b/tests/test_nn.py
index 9c873c9..dc1a8e6 100644
--- a/tests/test_nn.py
+++ b/tests/test_nn.py
@@ -142,7 +142,7 @@ def test_sequential(getkey):
         [
             eqx.nn.Linear(2, 4, key=getkey()),
             eqx.nn.Linear(4, 1, key=getkey()),
-            eqx.nn.BatchNorm(1, axis_name="batch"),
+            eqx.nn.BatchNorm(1, axis_name="batch", mode="ema"),
             eqx.nn.Linear(1, 3, key=getkey()),
         ]
     )
@@ -176,7 +176,7 @@ def test_nested_sequential(inner_stateful, outer_stateful, getkey):
         inner_seq = eqx.nn.Sequential(
             [
                 eqx.nn.Linear(2, 4, key=getkey()),
-                eqx.nn.BatchNorm(4, axis_name="batch")
+                eqx.nn.BatchNorm(4, axis_name="batch", mode="ema")
                 if inner_stateful
                 else eqx.nn.Identity(),
                 eqx.nn.Linear(4, 3, key=getkey()),
@@ -186,7 +186,7 @@ def test_nested_sequential(inner_stateful, outer_stateful, getkey):
             [
                 eqx.nn.Linear(5, 2, key=getkey()),
                 inner_seq,
-                eqx.nn.BatchNorm(3, axis_name="batch")
+                eqx.nn.BatchNorm(3, axis_name="batch", mode="ema")
                 if outer_stateful
                 else eqx.nn.Identity(),
                 eqx.nn.Linear(3, 6, key=getkey()),
@@ -949,7 +949,8 @@ def test_group_norm(getkey):
         gn = eqx.nn.GroupNorm(groups=4, channels=None, channelwise_affine=True)
 
 
-def test_batch_norm(getkey):
+@pytest.mark.parametrize("mode", ("ema", "batch"))
+def test_batch_norm(getkey, mode):
     x0 = jrandom.uniform(getkey(), (5,))
     x1 = jrandom.uniform(getkey(), (10, 5))
     x2 = jrandom.uniform(getkey(), (10, 5, 6))
@@ -957,14 +958,19 @@ def test_batch_norm(getkey):
 
     # Test that it works with a single vmap'd axis_name
 
-    bn = eqx.nn.BatchNorm(5, "batch")
+    bn = eqx.nn.BatchNorm(5, "batch", mode=mode)
     state = eqx.nn.State(bn)
     vbn = jax.vmap(bn, axis_name="batch", in_axes=(0, None), out_axes=(0, None))
 
     for x in (x1, x2, x3):
         out, state = vbn(x, state)
         assert out.shape == x.shape
-        running_mean, running_var = state.get(bn.state_index)
+        if mode == "ema":
+            assert bn.ema_state_index is not None
+            running_mean, running_var = state.get(bn.ema_state_index)
+        else:
+            assert bn.batch_state_index is not None
+            running_mean, running_var = state.get(bn.batch_state_index)
         assert running_mean.shape == (5,)
         assert running_var.shape == (5,)
 
@@ -985,13 +991,18 @@ def test_batch_norm(getkey):
         in_axes=(0, None),
     )(x2, state)
     assert out.shape == x2.shape
-    running_mean, running_var = state.get(bn.state_index)
+    if mode == "ema":
+        assert bn.ema_state_index is not None
+        running_mean, running_var = state.get(bn.ema_state_index)
+    else:
+        assert bn.batch_state_index is not None
+        running_mean, running_var = state.get(bn.batch_state_index)
     assert running_mean.shape == (10, 5)
     assert running_var.shape == (10, 5)
 
     # Test that it handles multiple axis_names
 
-    vvbn = eqx.nn.BatchNorm(6, ("batch1", "batch2"))
+    vvbn = eqx.nn.BatchNorm(6, ("batch1", "batch2"), mode=mode)
     vvstate = eqx.nn.State(vvbn)
     for axis_name in ("batch1", "batch2"):
         vvbn = jax.vmap(
@@ -999,14 +1010,19 @@ def test_batch_norm(getkey):
         )
     out, out_vvstate = vvbn(x2, vvstate)
     assert out.shape == x2.shape
-    running_mean, running_var = out_vvstate.get(vvbn.state_index)
+    if mode == "ema":
+        assert vvbn.ema_state_index is not None
+        running_mean, running_var = out_vvstate.get(vvbn.ema_state_index)
+    else:
+        assert vvbn.batch_state_index is not None
+        running_mean, running_var = out_vvstate.get(vvbn.batch_state_index)
     assert running_mean.shape == (6,)
     assert running_var.shape == (6,)
 
     # Test that it normalises
 
     x1alt = jrandom.normal(jrandom.PRNGKey(5678), (10, 5))  # avoid flakey test
-    bn = eqx.nn.BatchNorm(5, "batch", channelwise_affine=False)
+    bn = eqx.nn.BatchNorm(5, "batch", channelwise_affine=False, mode=mode)
     state = eqx.nn.State(bn)
     vbn = jax.vmap(bn, axis_name="batch", in_axes=(0, None), out_axes=(0, None))
     out, state = vbn(x1alt, state)
@@ -1017,9 +1033,19 @@ def test_batch_norm(getkey):
 
     # Test that the statistics update during training
     out, state = vbn(x1, state)
-    running_mean, running_var = state.get(bn.state_index)
+    if mode == "ema":
+        assert bn.ema_state_index is not None
+        running_mean, running_var = state.get(bn.ema_state_index)
+    else:
+        assert bn.batch_state_index is not None
+        running_mean, running_var = state.get(bn.batch_state_index)
     out, state = vbn(3 * x1 + 10, state)
-    running_mean2, running_var2 = state.get(bn.state_index)
+    if mode == "ema":
+        assert bn.ema_state_index is not None
+        running_mean2, running_var2 = state.get(bn.ema_state_index)
+    else:
+        assert bn.batch_state_index is not None
+        running_mean2, running_var2 = state.get(bn.batch_state_index)
     assert not jnp.allclose(running_mean, running_mean2)
     assert not jnp.allclose(running_var, running_var2)
 
@@ -1028,7 +1054,12 @@ def test_batch_norm(getkey):
     ibn = eqx.nn.inference_mode(bn, value=True)
     vibn = jax.vmap(ibn, axis_name="batch", in_axes=(0, None), out_axes=(0, None))
     out, state = vibn(4 * x1 + 20, state)
-    running_mean3, running_var3 = state.get(bn.state_index)
+    if mode == "ema":
+        assert bn.ema_state_index is not None
+        running_mean3, running_var3 = state.get(bn.ema_state_index)
+    else:
+        assert bn.batch_state_index is not None
+        running_mean3, running_var3 = state.get(bn.batch_state_index)
     assert jnp.array_equal(running_mean2, running_mean3)
     assert jnp.array_equal(running_var2, running_var3)
 
@@ -1042,6 +1073,16 @@ def test_batch_norm(getkey):
     f(jrandom.normal(getkey(), (1, 5)))
 
 
+# https://github.com/patrick-kidger/equinox/issues/1010
+def test_batch_norm_zero_training_steps(getkey):
+    bn, state = eqx.nn.make_with_state(eqx.nn.BatchNorm)(
+        input_size=3, axis_name="foo", mode="batch", inference=True
+    )
+    x = jrandom.normal(getkey(), (5, 3))
+    out, _ = jax.vmap(bn, in_axes=(0, None), axis_name="foo")(x, state)
+    assert jnp.allclose(x, out)
+
+
 def test_spectral_norm(getkey):
     def λ1():
         u, v = state.get(spectral.uv_index)
diff --git a/tests/test_serialisation.py b/tests/test_serialisation.py
index e89c243..44b214b 100644
--- a/tests/test_serialisation.py
+++ b/tests/test_serialisation.py
@@ -239,12 +239,16 @@ def test_stateful(tmp_path):
         norm1: eqx.nn.BatchNorm
         norm2: eqx.nn.BatchNorm
 
-    model = Model(eqx.nn.BatchNorm(3, "hi"), eqx.nn.BatchNorm(4, "bye"))
+    model = Model(
+        eqx.nn.BatchNorm(3, "hi", mode="ema"), eqx.nn.BatchNorm(4, "bye", mode="ema")
+    )
     state = eqx.nn.State(model)
 
     eqx.tree_serialise_leaves(tmp_path, (model, state))
 
-    model2 = Model(eqx.nn.BatchNorm(3, "hi"), eqx.nn.BatchNorm(4, "bye"))
+    model2 = Model(
+        eqx.nn.BatchNorm(3, "hi", mode="ema"), eqx.nn.BatchNorm(4, "bye", mode="ema")
+    )
     state2 = eqx.nn.State(model2)
 
     eqx.tree_deserialise_leaves(tmp_path, (model2, state2))
diff --git a/tests/test_stateful.py b/tests/test_stateful.py
index d7bc632..cf57b5c 100644
--- a/tests/test_stateful.py
+++ b/tests/test_stateful.py
@@ -7,7 +7,7 @@ import pytest
 
 
 def test_delete_init_state():
-    model = eqx.nn.BatchNorm(3, "batch")
+    model = eqx.nn.BatchNorm(3, "batch", mode="ema")
     eqx.nn.State(model)
     model2 = eqx.nn.delete_init_state(model)