Repository: patrick-kidger/equinox (language: python)

Issue:
Odd behavior for jax.tree_util.Partial and interactions with eqx.Module
Here's something pretty odd about `jax.tree_util.Partial`: it mimics the equality behavior of `functools.partial`, which returns that two partial calls on the same function on the same arguments are not equal, unless they are the same object.

This is a problem, because equinox wraps its method calls to return a `jax.tree_util.Partial`, giving you the following very unexpected behavior:
```
import equinox as eqx

class M(eqx.Module):
  def f(self, x):
    return x

m = M()
M.f == M.f # True
m.f == m.f # False
```
The same instance method isn't the same!  This is because they come from separate calls to the wrapper that turns the
class's method into `jax.tree_util.Partial`.

One possible fix is to switch the wrapper to use equinox's own `Partial`, which behaves sensibly with equality.  Except, you can't use `functools.wraps` on an equinox `Partial` because it's a frozen dataclass and `wraps` tries to change attributes on the method.  The only simple fix is to drop the `wraps` call altogether but that seems like a bad choice too.  Maybe the `Partial` could be unfrozen while it is wrapped?

Probably the better solution is for jax to change how `__eq__` works on `jax.tree_util.Partial`, because it leads to more unusual behavior.  The choice for equality may be a fine decision for `functools.partial`, but `jax.tree_util.Partial` is PyTree compatible and so can be flattened.  I would expect that if two flattened PyTrees are equal that the unflattened PyTrees would be equal.  That doesn't hold for `jax.tree_util.Partial`, which flattens to the function and its arguments, which are equal for two equivalent `jax.tree_util.Partial` objects.

Equinox's `Partial` doesn't have any of these problems.  See below, that only `jax.tree_util.Partial` is inconsistent between equality of its flattened and unflattened representations.

```
import functools as ft
import jax.tree_util as jtu
import equinox as eqx

def f(x):
  return x

print('functools: ', 
      ft.partial(f, 3) == ft.partial(f, 3),
      jtu.tree_flatten(ft.partial(f, 3)) == jtu.tree_flatten(ft.partial(f, 3)))
print('jax.tree_utils: ',
      jtu.Partial(f, 3) == jtu.Partial(f, 3),
      jtu.tree_flatten(jtu.Partial(f, 3)) == jtu.tree_flatten(jtu.Partial(f, 3)))
print('equinox: ',
      eqx.Partial(f, 3) == eqx.Partial(f, 3),
      jtu.tree_flatten(eqx.Partial(f, 3)) == jtu.tree_flatten(eqx.Partial(f, 3)))
```
gives the output:
```
functools:  False False
jax.tree_utils:  False True
equinox:  True True
```



Discussion / hints:

Fix patch:
diff --git a/.gitignore b/.gitignore
index 589331ee..ebeaa306 100644
--- a/.gitignore
+++ b/.gitignore
@@ -5,6 +5,7 @@ build/
 dist/
 site/
 examples/data
+examples/CIFAR
 .all_objects.cache
 .pymon
 .idea
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index abd30f2d..52f1f122 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -12,7 +12,7 @@ repos:
     rev: v1.1.315
     hooks:
     - id: pyright
-      additional_dependencies: [beartype, einops, jax, jaxtyping, pytest, tensorflow, tf2onnx, typing_extensions]
+      additional_dependencies: [beartype, einops, jax, jaxtyping, optax, pytest, tensorflow, tf2onnx, typing_extensions]
   - repo: https://github.com/nbQA-dev/nbQA
     rev: 1.6.3
     hooks:
diff --git a/docs/api/filtering/transformations.md b/docs/api/filtering/transformations.md
index dda82586..c50457f9 100644
--- a/docs/api/filtering/transformations.md
+++ b/docs/api/filtering/transformations.md
@@ -38,6 +38,10 @@ Likewise, `eqx.filter_grad` will automatically differentiate all floating-point
 
 ---
 
+::: equinox.filter_checkpoint
+
+---
+
 ::: equinox.filter_custom_jvp
 
 ---
diff --git a/docs/api/manipulation.md b/docs/api/manipulation.md
index 06a023d1..bab25133 100644
--- a/docs/api/manipulation.md
+++ b/docs/api/manipulation.md
@@ -20,10 +20,6 @@
 
 ---
 
-::: equinox.tree_inference
-
----
-
 ::: equinox.tree_flatten_one_level
 
 ---
diff --git a/docs/api/module/advanced_fields.md b/docs/api/module/advanced_fields.md
index 2c18549a..473b7f99 100644
--- a/docs/api/module/advanced_fields.md
+++ b/docs/api/module/advanced_fields.md
@@ -18,6 +18,105 @@ Equinox modules can be used as [abstract base classes](https://docs.python.org/3
     selection:
         members: false
 
+## Checking invariants
+
+Equinox extends dataclasses with a `__check_init__` method, which is automatically ran after initialisation. This can be used to check invariants like so:
+
+```python
+class Positive(eqx.Module):
+    x: int
+
+    def __check_init__(self):
+        if self.x <= 0:
+            raise ValueError("Oh no!")
+```
+
+This method has three key differences compared to the `__post_init__` provided by dataclasses:
+
+- It is not overridden by an `__init__` method of a subclass. In contrast, the following code has a bug (Equinox will raise a warning if you do this):
+    
+    ```python
+    class Parent(eqx.Module):
+        x: int
+
+        def __post_init__(self):
+            if self.x <= 0:
+                raise ValueError("Oh no!")
+
+    class Child(Parent):
+        x_as_str: str
+
+        def __init__(self, x):
+            self.x = x
+            self.x_as_str = str(x)
+
+    Child(-1)  # No error!
+    ```
+
+- It is automatically called for parent classes; `super().__check_init__()` is not required:
+
+    ```python
+    class Parent(eqx.Module):
+        def __check_init__(self):
+            print("Parent")
+
+    class Child(Parent):
+        def __check_init__(self):
+            print("Child")
+
+    Child()  # prints out both Child and Parent
+    ```
+
+    As with the previous bullet point, this is to prevent child classes accidentally failing to check that the invariants of their parent hold.
+
+- Assignment is not allowed:
+    
+    ```python
+    class MyModule(eqx.Module):
+        foo: int
+
+        def __check_init__(self):
+            self.foo = 1  # will raise an error
+    ```
+
+    This is to prevent `__check_init__` from doing anything too surprising: as the name suggests, it's meant to be used for checking invariants.
+
 ## 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/api/nn/inference.md b/docs/api/nn/inference.md
new file mode 100644
index 00000000..1c743385
--- /dev/null
+++ b/docs/api/nn/inference.md
@@ -0,0 +1,3 @@
+# Training/Inference
+
+::: equinox.nn.inference_mode
diff --git a/docs/api/nn/sequential.md b/docs/api/nn/sequential.md
index 41c57610..df7b84b3 100644
--- a/docs/api/nn/sequential.md
+++ b/docs/api/nn/sequential.md
@@ -25,4 +25,4 @@ These are useful when building fairly straightforward models. But for anything n
 ::: equinox.nn.StatefulLayer
     selection:
         members:
-            - __call__
+            - is_stateful
diff --git a/docs/api/nn/shared.md b/docs/api/nn/shared.md
new file mode 100644
index 00000000..d52e4ed7
--- /dev/null
+++ b/docs/api/nn/shared.md
@@ -0,0 +1,7 @@
+# Sharing layers
+
+::: equinox.nn.Shared
+    selection:
+        members:
+            - __init__
+            - __call__
diff --git a/docs/api/nn/stateful.md b/docs/api/nn/stateful.md
index 51ad63bb..339ca0bc 100644
--- a/docs/api/nn/stateful.md
+++ b/docs/api/nn/stateful.md
@@ -1,17 +1,113 @@
 # Stateful operations
 
-These are the tools that underly stateful operations, like [`equinox.nn.BatchNorm`][] or [`equinox.nn.SpectralNorm`][].
+These are the tools that underlie stateful operations, like [`equinox.nn.BatchNorm`][] or [`equinox.nn.SpectralNorm`][]. These are fairly unusual layers, so most users will not need this part of the API.
 
-See the [stateful example](../../examples/stateful.ipynb) for an example of working with stateful operations.
+!!! Example
 
-::: equinox.nn.State
+    The [stateful example](../../examples/stateful.ipynb) is a good reference for the typical workflow for stateful layers.
+
+---
+
+::: equinox.nn.make_with_state
+
+## Extra features
+
+Let's explain how this works under the hood. First of all, all stateful layers (`BatchNorm` etc.) include an "index". This is basically just a unique hashable value (used later as a dictionary key), and an initial value for the state:
+
+::: equinox.nn.StateIndex
     selection:
         members:
             - __init__
 
 ---
 
-::: equinox.nn.StateIndex
+This `State` object that's being passed around is essentially just a dictionary, mapping from `StateIndex`s to PyTrees-of-arrays. Correspondingly this has `.get` and `.set` methods to read and write values to it.
+
+::: equinox.nn.State
     selection:
         members:
-            - __init__
+            - get
+            - set
+            - substate
+            - update
+
+## Custom stateful layers
+
+Let's use [`equinox.nn.StateIndex`][] to create a custom stateful layer.
+
+```python
+import equinox as eqx
+import jax.numpy as jnp
+from jaxtyping import Array
+
+class Counter(eqx.Module):
+    index: eqx.nn.StateIndex
+
+    def __init__(self):
+        init_state = jnp.array(0)
+        self.index = eqx.nn.StateIndex(init_state)
+
+    def __call__(self, x: Array, state: eqx.nn.State) -> tuple[Array, eqx.nn.State]:
+        value = state.get(self.index)
+        new_x = x + value
+        new_state = state.set(self.index, value + 1)
+        return new_x, new_state
+
+counter, state = eqx.nn.make_with_state(Counter)()
+x = jnp.array(2.3)
+
+num_calls = state.get(counter.index)
+print(f"Called {num_calls} times.")  # 0
+
+_, state = counter(x, state)
+num_calls = state.get(counter.index)
+print(f"Called {num_calls} times.")  # 1
+
+_, state = counter(x, state)
+num_calls = state.get(counter.index)
+print(f"Called {num_calls} times.")  # 2
+```
+
+## Vmap'd stateful layers
+
+This is an advanced thing to do! Here we'll build on [the ensembling guide](../../../tricks/#ensembling), and see how how we can create vmap'd stateful layers.
+
+This follows on from the previous example, in which we define `Counter`.
+```python
+import jax.random as jr
+
+class Model(eqx.Module):
+    linear: eqx.nn.Linear
+    counter: Counter
+    v_counter: Counter
+
+    def __init__(self, key):
+        # Not-stateful layer
+        self.linear = eqx.nn.Linear(2, 2, key=key)
+        # Stateful layer.
+        self.counter = Counter()
+        # Vmap'd stateful layer. (Whose initial state will include a batch dimension.)
+        self.v_counter = eqx.filter_vmap(Counter, axis_size=2)()
+
+    def __call__(self, x: Array, state: eqx.nn.State) -> tuple[Array, eqx.nn.State]:
+        # This bit happens as normal.
+        assert x.shape == (2,)
+        x = self.linear(x)
+        x, state = self.counter(x, state)
+
+        # For the vmap, we have to restrict our state to just those states we want to
+        # vmap, and then update the overall state again afterwards.
+        #
+        # After all, the state for `self.counter` isn't expecting to be batched, so we
+        # have to remove that.
+        substate = state.substate(self.v_counter)
+        x, substate = eqx.filter_vmap(self.v_counter)(x, substate)
+        state = state.update(substate)
+
+        return x, state
+
+key = jr.PRNGKey(0)
+model, state = eqx.nn.make_with_state(Model)(key)
+x = jnp.array([5.0, -1.0])
+model(x, state)
+```
diff --git a/docs/faq.md b/docs/faq.md
index 8180ce01..ccaeca00 100644
--- a/docs/faq.md
+++ b/docs/faq.md
@@ -31,17 +31,13 @@ class Module(eqx.Module):
         self.linear1 = shared_linear
         self.linear2 = shared_linear
 ```
-in which the same object is saved multiple times in the model.
+in which the same object is saved multiple times in the model. However, after making some gradient updates you'll find that `self.linear1` and `self.linear2` are now different.
 
-Don't do this!
+This is intended. In Equinox+JAX, models are Py*Trees*, not DAGs. (Directed acyclic graphs.) This is basically just an arbitrary choice JAX made a long time ago in its design, but it does generally make reasoning about your code fairly easy. (You never need to track if an object is used in multiple places.)
 
-After making some gradient updates you'll find that `self.linear1` and `self.linear2` are now different.
+That said, it can sometimes happen that you really do want to tie together multiple nodes in your PyTree. If this is the case, then use [`equinox.nn.Shared`][], which provides this behaviour. (It stores things as a tree, and then inserts a reference to each node into the right place whenever you need it.)
 
-Recall that in Equinox, models are PyTrees. Meanwhile, JAX treats all PyTrees as *trees*: that is, the same object does not appear more in the tree than once. (If it did, then it would be a *directed acyclic graph* instead.) If JAX ever encounters the same object multiple times then it will unwittingly make independent copies of the object whenever it transforms the overall PyTree.
-
-The resolution is simple: just don't store the same object in multiple places in the PyTree.
-
-You can check for whether you have duplicate nodes by using the [`equinox.tree_check`][] function.
+You can also check for whether you have duplicate nodes by using the [`equinox.tree_check`][] function.
 
 ## How do I input higher-order tensors (e.g. with batch dimensions) into my model?
 
@@ -160,3 +156,85 @@ class Model(eqx.Module):
     def __call__(self, x):
         return self.param * x + jax.lax.stop_gradient(self.buffer)
 ```
+
+## I think my function is being recompiled each time it is run.
+
+You can check each time your function is compiled by adding a print statement:
+```python
+@eqx.filter_jit
+def your_function(x, y, z):
+    print("Compiling!")
+    ...  # rest of your code here
+```
+JAX calls your function each time it needs to compile it. Afterwards, it never actually calls it -- indeed it doesn't use Python at all! Instead, it uses its compiled copy of your function, which only performs array operations. Thus, a print statement is an easy way to check each time JAX is compiling your function.
+
+A function will be recompiled every time the shape or dtype of its arrays changes, or if any of its static (non-array) inputs change (as measured by `__eq__`).
+
+If you want to check which argument is causing an undesired recompilation, then this can be done by checking each argument in turn:
+```python
+@eqx.filter_jit
+def check_arg(arg, name):
+    print(f"Argument {name} is triggering a compile.")
+
+
+for step, (x, y, z) in enumerate(...):  # e.g. a training loop
+    print(f"Step is {step}")
+    check_arg(x, "x")
+    check_arg(y, "y")
+    check_arg(z, "z")
+    your_function(x, y, z)
+```
+for which you'll often see output like
+```
+Step is 0
+Argument x is triggering a compile.
+Argument y is triggering a compile.
+Argument z is triggering a compile.
+Step is 1
+Argument y is triggering a compile.
+Step is 2
+Argument y is triggering a compile.
+...
+```
+On the very first step, none of the arguments have been seen before, so they all trigger a compile. On later steps, just the problematic argument will trigger a recompilation of `check_arg` -- this will be the one that is triggering a recompilation of `your_function` as well!
+
+## How does Equinox compare to...?
+
+#### ...PyTorch?
+
+JAX+Equinox is usually faster than PyTorch (a stronger JIT compiler), and more featureful (e.g. supporting jit-of-vmap, forward-mode autolinearisation, and autoparallelism).
+
+For those doing scientific computing or scientific ML, then JAX+Equinox also has a much stronger ecosystem. For example, PyTorch no longer has a library for solving differential equations (torchdiffeq is unmaintained). Meanwhile, JAX has [Diffrax](https://github.com/patrick-kidger/diffrax).
+
+Both JAX+Equinox and PyTorch are roughly equally easy to use. PyTorch tends to be a easier for new users (e.g. it's closer to being "Python as normal", and there's less functional programming), whilst JAX+Equinox generally supports advanced use-cases more cleanly (e.g. PyTorch has multiple JIT compilers each with their own quirks -- `torch.{fx, jit.script, jit.trace, compile, _dynamo, ...}` -- whilst JAX+Equinox just has the one).
+
+PyTorch is older, and enjoys broader adoption -- it's generally easier to find developers for PyTorch, or off-the-shelf model implementations using it.
+
+#### ...Keras?
+
+These are two very different libraries, with very different target audiences. Keras is great for plug-and-play building of models -- it's often compared to using Lego. This makes it a convenient framework for standing up neural networks quickly. Equinox is much lower level: it tries to support more general use-cases (e.g. its downstream scientific ecosystem), but usually requires greater proficiency with numerical computing / software development / machine learning.
+
+#### ...Flax?
+
+- Flax introduces multiple new abstractions (`flax.linen.Module`, `flax.linen.Variable`, `Module.setup` vs `flax.linen.compact`, `flax.struct.dataclass`, etc.). Equinox tries to avoid adding new abstractions to core JAX; everything is always just a PyTree.
+- Flax is a DSL: it is generally incompatible with non-Flax code, and requires using wrapped `flax.linen.{vmap, scan, ...}` rather than the native `jax.{vmap, ...}`. In contrast, Equinox allows you to use native JAX operations and aims to be compatible with arbitrary JAX code.
+- Bound methods of `eqx.Module` are just PyTrees. In Flax this isn't the case -- passing around bound methods will either result in errors or recompilations, depending what you do. Likewise, `eqx.Module` handles inheritance correctly, including propagating metadata like docstrings. The equivalent `flax.struct.dataclass` silently misbehaves. Overall Equinox seems to have fewer footguns.
+- Equinox offers several advanced features (like [runtime errors](../api/errors/) or [PyTree manipulation](../api/manipulation/#equinox.tree_at)) not found in other libraries.
+
+See also the [Equinox paper](https://arxiv.org/abs/2111.00254).
+
+#### ...Julia?
+
+The Julia ecosystem has [historically been buggy](https://kidger.site/thoughts/jax-vs-julia/).
+
+At time of writing, Julia does not yet have a robust autodifferentiation system. For example, it has multiple competing implementations -- both Diffractor.jl and ForwardDiff.jl for forward-mode autodifferentiation, and all of Tracker.jl, Zygote.jl, Enzyme.jl, ReverseDiff.jl for reverse-mode autodifferentiation. It does not yet support higher-order autodifferentiation robustly. In contrast, JAX+Equinox use a single strong autodifferentiation system.
+
+However, note that JAX+Equinox don't try to offer a completely general programming model: they are optimised for arrays and linear algebra. (Essentially, the sorts of things you use NumPy for.) They're not designed for e.g. a branch-and-bound combinatorial optimisation algorithm, and for these purposes Julia will be superior.
+
+Julia is often a small amount faster on microbenchmarks on CPUs. JAX+Equinox supports running on TPUs, whilst Julia generally does not.
+
+**You're obviously biased! Are the above comparisons fair?**
+
+Seriously, we think they're fair! Nonetheless all of the above approaches have their adherents, so it seems like all of these approaches are doing something right. So if you're already happily using one of them for your current project... then keep using them. (Don't rewrite things for no reason.) But conversely, we'd invite you to try Equinox for your next project. :)
+
+For what it's worth, if you have the time to learn (e.g. you're a grad student), then we'd strongly recommend trying all of the above. All of these libraries have made substantial innovations, and have all made substantially moved the numerical computing space forward. Equinox deliberately takes inspiration from them. For example Julia has an excellent type system, and this has strongly informed [this Equinox design pattern](../pattern/).
diff --git a/docs/pattern.md b/docs/pattern.md
new file mode 100644
index 00000000..d5c03feb
--- /dev/null
+++ b/docs/pattern.md
@@ -0,0 +1,273 @@
+# The abstract/final pattern
+
+The following is a useful design pattern. It's not mandatory, but it does come recommended, as it's designed to produce very readable code. This is also the pattern used throughout the Equinox ecosystem -- [Lineax](https://github.com/google/lineax), [Diffrax](https://github.com/patrick-kidger/diffrax) etc.
+
+!!! Tip "The abstract/final design pattern"
+
+    Every subclass of `eqx.Module` must be either
+
+    (a) abstract (it can be subclassed, but not instantiated); or  
+    (b) final (it can be instantiated, but not subclassed).
+
+    Unless they're abstract, then methods and attributes must never be overridden. Once they've been implemented, that's it.
+
+    The `__init__` method, and all dataclass fields, must all be defined in one class. No defining fields in multiple parts of a class hierarchy.
+
+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 (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:
+
+    - If we want to subclass to override just one method, just to tweak it a little bit -- this can be done by wrapping the original class, not subclassing it.
+    - In practice we typically define all fields, and the `__init__` method, on the final class that we instantiate. Every so often we could do them on an abstract class though, in which case the final subclasses will only be implementing methods.
+    - You should never use the `hasattr` builtin. You probably meant to declare this attribute/method on an abstract base class.
+    - To access an attribute on an abstract class, it can be declared using [`equinox.AbstractVar`][]; the eventual final subclass must provide this as a field or property.
+
+    And for the CS nerds:
+
+    - This is going all-in on nominal subtyping, not structural/duck subtyping.
+    - Yes, this design pattern looks a lot like Rust/Julia/something. For example the abstract-or-final rule is inspired by Julia's type system, and the no-overriding is inspired by how many languages (including Rust) approach the orphan problem.
+    - This has nothing to do with object-oriented (OO) programming, and everything to do with type theory. OO is about mutating classes through their methods, and as Equinox modules are immutable then we never do that here. (It is true that a lot of OO codebases could benefit from rules like these, though.)
+
+---
+
+The above is the really important bit. For those who are curious to know more, here's how we arrive at the above design pattern.
+
+## Level 1: Abstract base classes (ABCs) as interfaces
+
+Let's start off with something very standard: using ABCs to define interfaces. Here's an example.
+
+```python
+class AbstractOptimiser(eqx.Module):
+    @abc.abstractmethod
+    def init(self, params):
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def update(self, params, grads, state):
+        raise NotImplementedError
+
+class Adam(AbstractOptimiser):
+    learning_rate: float
+    beta1: float = 0.9
+    beta2: float = 0.999
+
+    def init(self, params):
+        ...  # some implementation
+        return initial_state
+
+    def update(self, params, grads, state):
+        ...  # some implementation
+        return new_params, new_state
+
+@eqx.filter_jit
+def make_step(params, data, opt_state, optimiser: AbstractOptimiser):
+    grads = eqx.filter_grad(compute_loss)(params, data)
+    new_params, new_opt_state = optimiser.update(params, grads, opt_state)
+    return new_params, new_opt_state
+
+def train(params, dataloader, optimiser: AbstractOptimiser):
+    opt_state = optimiser.init(params)
+    for data in dataloader:
+        params, opt_state = make_step(params, data, opt_state, optimiser)
+    return params
+
+params = ...  # some model
+dataloader = ...  # some dataloader
+optimiser = Adam(learning_rate=3e-4)
+train(params, dataloader, optimiser)
+```
+
+Hopefully the above is indeed easy to read. The `AbstractOptimiser` defines an interface using `init(params)` and `update(params, grad, state)`. (And we should also add type annotations to their arguments, actually.)
+
+Subsequently, the `train` and `make_step` functions can be written without needing to know exactly which optimiser has been passed. (We can later implement some other optimiser and use that in the same place.)
+
+For readability, it's worth following the convention that all abstract classes begin with the word "Abstract".
+
+The above is very common. Indeed Python has a whole module, [abc](https://docs.python.org/3/library/abc.html), for declaring such `abc.abstractmethod`s. This is also our first example of Equinox making this approach easy: Equinox modules automatically inherit from `abc.ABC`, so you don't need to do that yourself.
+
+## Level 2: intermediate ABCS, abstract attributes, and `__init__`-only-once
+
+Now let's move on to a natural extension to the above: intermediate ABCs, that introduce partial implementations.
+
+```python
+class AbstractInterpolation(eqx.Module):
+    @abc.abstractmethod
+    def __call__(self, x: Array) -> Array:
+        raise NotImplementedError
+
+
+class AbstractPolynomialInterpolation(AbstractInterpolation)
+    coeffs: eqx.AbstractVar[Array]
+
+    def degree(self) -> int:
+        return len(self.coeffs)
+
+    def __call__(self, x: Array) -> Array:
+        return jnp.polyval(self.coeffs, x)
+
+
+class CubicInterpolation(AbstractPolynomialInterpolation):
+    coeffs: Array
+
+    def __init__(self, ts: Array, xs: Array):
+        self.coeffs = ...  # some implementation
+```
+in this case, the intermediate ABC `AbstractPolynomialInterpolation` implements the `__call__` method. However, it isn't yet a concrete (non-abstract) class, as it introduces a new abstract variable `coeffs` -- we need to wait until `CubicInterpolation` for that to be defined.
+
+Using an abstract attribute ([`equinox.AbstractVar`][]) here means that we can write `self.coeffs` inside `degree` and `__call__`, and know that this is safe. Unless all abstract attributes are defined then Equinox won't allow us to instantiate the class.
+
+Why didn't we just define `AbstractPolynomialInterpolation.coeffs` as a concrete field? (Just `coeffs: Array`.) Indeed we could have written this:
+```python
+class AbstractPolynomialInterpolation(AbstractInterpolation)
+    coeffs: Array
+
+    def __init__(self, coeffs: Array):
+        self.coeffs = coeffs
+
+    def degree(self) -> int:
+        return len(self.coeffs)
+
+    def __call__(self, x: Array) -> Array:
+        return jnp.polyval(self.coeffs, x)
+
+
+class CubicInterpolation(AbstractPolynomialInterpolation):
+    def __init__(self, ts: Array, xs: Array):
+        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.
+
+## Level 3: implement methods precisely once, and concrete-means-final
+
+Our "`__init__` only once" rule means that `__init__` is defined precisely once, is never overridden, and we never call `super().__init__`. Why stop there: perhaps we should enforce that we never override *any* method?
+
+In practice, we argue that's a good idea! This rule means that when you see code like:
+```python
+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 overriden. (And this rule will also be checked via the `strict=True` flag.)
+
+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 overriden 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!
+
+What about when you have an existing concrete class that you want to tweak just-a-little-bit? In this case, prefer composition over inheritance. Write a wrapper that forwards each method as appropriate. This is just as expressive, and means we keep these readable type-safe rules.
+
+## Level 4: `__check_init__`
+
+It's pretty common to want to validate that certain invariants hold, even in abstract base classes. For this, we have the `__check_init__` method:
+```python
+class AbstractPolynomialInterpolation(AbstractInterpolation)
+    coeffs: AbstractVar[Array]
+
+    def __check_init__(self):
+        if not jnp.issubdtype(self.coeffs.dtype, jnp.floating):
+            raise ValueError("Coefficients must be floating-point!")
+
+    ...
+```
+This method is something that Equinox will look for, and if present it will be ran after initialisation. This is an Equinox-specific extension designed to support this design pattern.
+
+See [checking invariants](../api/module/advanced_fields/#checking-invariants) for more details.
+
+## Extensions and FAQ
+
+**Does this pattern work with multiple inheritance?**
+
+Yes. For example, here's a diamond inheritance pattern (for building a differential equation solver):
+```python
+class AbstractSolver(eqx.Module):
+    @abc.abstractmethod
+    def step(...):
+        raise NotImplementedError
+
+class AbstractAdaptiveSolver(AbstractSolver):
+    tolerance: eqx.AbstractVar[float]
+
+class AbstractImplicitSolver(AbstractSolver):
+    root_finder: eqx.AbstractVar[AbstractRootFinder]
+
+class ImplicitEuler(AbstractAdaptiveSolver, AbstractImplicitSolver):
+    tolerance: float
+    root_finder: AbstractRootFinder = Newton()
+
+    def step(...):
+        ...  # some implementation
+
+solver = ImplicitEuler(tolerance=1e-3)  # this can be instantiated
+```
+
+**That's a lot of `Abstract`s**
+
+Yes.
+
+**Does `super()` ever get used at all?**
+
+No. This design pattern means that you should never need to write `super()` at all.
+
+**What about co-operative multiple inheritance?**
+
+If you're a Python nerd, you'll now be wondering about co-operative multiple inheritance, which specifies using `super()` ubiquitously.
+
+The TL;DR of this is that almost no-one ever uses this properly, and the abstract+final pattern is intended as a direct alternative. One sees a lot of code that looks like this:
+```python
+class A:
+    def __init__(self, x):
+        self.x = x
+        # Not calling super().__init__, because the superclass is just `object`, right?
+
+class AA:
+    def __init__(...):
+        super().__init__(...)  # Being a good citizen.
+        ...  # Do anything else that needs to happen.
+
+class B(A, AA):
+    pass
+
+B()  # AA.__init__ is not called.
+```
+In this case `B()` calls `A.__init__` and this then fails to call `AA.__init__`. Co-operative multiple inheritance only works if everyone, well, co-operates.
+
+Even if everyone wants to do their best, there is another issue. When writing `super().__init__`, it isn't actually know what method is being called -- as above, `super()` could be pointing at almost any class at all. This actually means that it's not possible to know what arguments to pass to `super().__init__`! "Only use keyword arguments" is the closest to a resolution that this issue has, and it's still fragile.
+
+In contrast, our no-overriding and abstract-or-final rules means that we never come across this scenario. We always know precisely what is being called.
+
+**Hang on, I don't buy the abstract-or-concrete part. Can't a concrete subclass add a new method to a concrete superclass?**
+
+You're thinking of something like this this:
+```python
+# This is clearly something we can instantiate: it has no abstract methods/attributes.
+class ConcreteArray(eqx.Module):
+    def some_method(self):
+        pass
+
+# This is clearly also without abstract methods/attributes, and also doesn't break the
+# rule about overriding methods.
+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:
+```python
+def add(x: ConcreteArray, y: ConcreteArray) -> ConcreteArray:
+    ...
+```
+so that there are never any questions about what the return type should be. (If we passed in `ConcreteArrayTwo` in to `x` or `y`, maybe we should try to return a `ConcreteArrayTwo` instead? What if `x = ConcreteArrayTwo()` but `y = ConcreteArrayThree()`, and these two types don't know about each other? Better to avoid the question in the first place.)
+
+**These ideas have appeared in &lt;XYZ language&gt;?**
+
+Yup! Variants of this design pattern are very common, especially in modern languages like Julia/Rust/etc., or in older languages with a strong emphasis on typing.
diff --git a/equinox/__init__.py b/equinox/__init__.py
index 12522da5..a39b0406 100644
--- a/equinox/__init__.py
+++ b/equinox/__init__.py
@@ -2,6 +2,7 @@
 
 from . import debug as debug, internal as internal, nn as nn
 from ._ad import (
+    filter_checkpoint as filter_checkpoint,
     filter_closure_convert as filter_closure_convert,
     filter_custom_jvp as filter_custom_jvp,
     filter_custom_vjp as filter_custom_vjp,
@@ -53,7 +54,6 @@
     tree_check as tree_check,
     tree_equal as tree_equal,
     tree_flatten_one_level as tree_flatten_one_level,
-    tree_inference as tree_inference,
 )
 from ._update import apply_updates as apply_updates
 from ._vmap_pmap import (
@@ -61,6 +61,7 @@
     filter_vmap as filter_vmap,
     if_array as if_array,
 )
+from .nn import inference_mode as tree_inference  # noqa: F401 - backward compatibility
 
 
 __version__ = importlib.metadata.version("equinox")
diff --git a/equinox/_ad.py b/equinox/_ad.py
index a048acf3..e71bd4dc 100644
--- a/equinox/_ad.py
+++ b/equinox/_ad.py
@@ -26,6 +26,7 @@
 from ._custom_types import sentinel
 from ._deprecate import deprecated_0_10
 from ._doc_utils import doc_remove_args
+from ._eval_shape import cached_filter_eval_shape
 from ._filters import (
     combine,
     is_array,
@@ -54,12 +55,19 @@ class _ValueAndGradWrapper(Module):
     def __wrapped__(self):
         return self._fun
 
-    def __call__(self, x, /, *args, **kwargs):
+    def __call__(self, *args, **kwargs):
         @ft.partial(jax.value_and_grad, has_aux=self._has_aux, **self._gradkwargs)
         def fun_value_and_grad(_diff_x, _nondiff_x, *_args, **_kwargs):
             _x = combine(_diff_x, _nondiff_x)
             return self._fun(_x, *_args, **_kwargs)
 
+        if len(args) == 0:
+            raise TypeError(
+                "Functions wrapped with `equinox.filter_{grad, value_and_grad}` must "
+                "have their first argument passed by position, not keyword. (This is "
+                "the argument that is differentiated.)"
+            )
+        x, *args = args
         diff_x, nondiff_x = partition(x, is_inexact_array)
         return fun_value_and_grad(diff_x, nondiff_x, *args, **kwargs)
 
@@ -75,7 +83,7 @@ class _GradWrapper(Module):
 
     @property
     def __wrapped__(self):
-        return self._fun_value_and_grad
+        return self._fun_value_and_grad.__wrapped__  # pyright: ignore
 
     def __call__(self, /, *args, **kwargs):
         value, grad = self._fun_value_and_grad(*args, **kwargs)
@@ -171,7 +179,7 @@ def filter_value_and_grad(
             "as the first argument."
         )
 
-    return module_update_wrapper(_ValueAndGradWrapper(fun, has_aux, gradkwargs), fun)
+    return module_update_wrapper(_ValueAndGradWrapper(fun, has_aux, gradkwargs))
 
 
 @overload
@@ -261,7 +269,7 @@ def grad_func(x__y):
 
     fun_value_and_grad = filter_value_and_grad(fun, has_aux=has_aux, **gradkwargs)
     fun_value_and_grad = cast(_ValueAndGradWrapper, fun_value_and_grad)
-    return module_update_wrapper(_GradWrapper(fun_value_and_grad, has_aux), fun)
+    return module_update_wrapper(_GradWrapper(fun_value_and_grad, has_aux))
 
 
 def _is_none(x):
@@ -416,12 +424,53 @@ def _unflatten(flat_pytree):
 _FlatPyTree = tuple[list[_T], PyTreeDef]
 
 
+def _check_closure_convert_input(self, args, kwargs):
+    self_in_dynamic_struct = _unflatten(self.in_dynamic_struct)
+    self_in_static = _unflatten(self.in_static)
+    in_dynamic, in_static = partition((args, kwargs), is_array)
+    in_dynamic_struct = jax.eval_shape(lambda: in_dynamic)
+    # `is` because `tree_equal` may return a tracer
+    if tree_equal(in_dynamic_struct, self_in_dynamic_struct) is not True:
+        raise ValueError(
+            "Closure-converted function called with different dynamic arguments to "
+            "the example arguments provided."
+        )
+    if tree_equal(in_static, self_in_static) is not True:
+        raise ValueError(
+            "Closure-converted function called with different static arguments to "
+            "the example arguments provided."
+        )
+    return in_dynamic
+
+
+class _TrivialClosureConvert(Module):
+    fn: types.FunctionType
+    in_dynamic_struct: _FlatPyTree[jax.ShapeDtypeStruct] = field(static=True)
+    in_static: _FlatPyTree[Any] = field(static=True)
+
+    @property
+    def in_struct(self):
+        dynamic = _unflatten(self.in_dynamic_struct)
+        static = _unflatten(self.in_static)
+        return combine(dynamic, static)
+
+    @property
+    def out_struct(self):
+        args, kwargs = self.in_struct
+        return cached_filter_eval_shape(self.fn, *args, **kwargs)
+
+    def __call__(self, *args, **kwargs):
+        # unused output
+        _ = _check_closure_convert_input(self, args, kwargs)
+        return self.fn(*args, **kwargs)
+
+
 class _ClosureConvert(Module):
     # Important that `jaxpr` be a leaf (and not static), so that it is a tuple element
     # when passing through `filter_primitive_bind` and thus visible to
     # `jax.core.subjaxprs`
     jaxpr: jax.core.Jaxpr
-    consts: PyTree[Array]  # Captured in the PyTree structure of _ClosureConvert
+    consts: PyTree[ArrayLike]  # Captured in the PyTree structure of _ClosureConvert
     in_dynamic_struct: _FlatPyTree[jax.ShapeDtypeStruct] = field(static=True)
     out_dynamic_struct: _FlatPyTree[jax.ShapeDtypeStruct] = field(static=True)
     in_static: _FlatPyTree[Any] = field(static=True)
@@ -440,27 +489,13 @@ def out_struct(self):
         return combine(dynamic, static)
 
     def __call__(self, *args, **kwargs):
-        self_in_dynamic_struct = _unflatten(self.in_dynamic_struct)
-        self_out_dynamic_struct = _unflatten(self.out_dynamic_struct)
-        self_in_static = _unflatten(self.in_static)
-        self_out_static = _unflatten(self.out_static)
-        in_dynamic, in_static = partition((args, kwargs), is_array)
-        in_dynamic_struct = jax.eval_shape(lambda: in_dynamic)
-        # `is` because `tree_equal` may return a tracer
-        if tree_equal(in_dynamic_struct, self_in_dynamic_struct) is not True:
-            raise ValueError(
-                "Closure-converted function called with different dynamic arguments to "
-                "the example arguments provided."
-            )
-        if tree_equal(in_static, self_in_static) is not True:
-            raise ValueError(
-                "Closure-converted function called with different static arguments to "
-                "the example arguments provided."
-            )
+        in_dynamic = _check_closure_convert_input(self, args, kwargs)
         in_dynamic_flat = jtu.tree_leaves(in_dynamic)
         out_dynamic_flat = jax.core.eval_jaxpr(
             self.jaxpr, self.consts, *in_dynamic_flat
         )
+        self_out_dynamic_struct = _unflatten(self.out_dynamic_struct)
+        self_out_static = _unflatten(self.out_static)
         out_dynamic_struct_flat, out_dynamic_treedef = jtu.tree_flatten(
             self_out_dynamic_struct
         )
@@ -509,24 +544,25 @@ def f(x, y):
         f(1., 1.)
         ```
     """
-    if isinstance(fn, types.FunctionType) and fn.__closure__ is None:
-        # In this case, it's not possible to have any closed-over tracers.
-        # Skip jaxpr tracing for efficiency.
-        return fn
-    closed_jaxpr, out_dynamic_struct, out_static = filter_make_jaxpr(fn)(
-        *args, **kwargs
-    )  # pyright: ignore
     in_dynamic, in_static = partition((args, kwargs), _is_struct)
     in_dynamic_struct = jax.eval_shape(lambda: in_dynamic)
-    jaxpr = closed_jaxpr.jaxpr
-    consts = closed_jaxpr.consts
     in_dynamic_struct = jtu.tree_flatten(in_dynamic_struct)
-    out_dynamic_struct = jtu.tree_flatten(out_dynamic_struct)
     in_static = jtu.tree_flatten(in_static)
-    out_static = jtu.tree_flatten(out_static)
-    closure_converted = _ClosureConvert(
-        jaxpr, consts, in_dynamic_struct, out_dynamic_struct, in_static, out_static
-    )
+    if isinstance(fn, types.FunctionType) and fn.__closure__ is None:
+        # In this case, it's not possible to have any closed-over tracers.
+        # Skip jaxpr tracing for efficiency.
+        closure_converted = _TrivialClosureConvert(fn, in_dynamic_struct, in_static)
+    else:
+        closed_jaxpr, out_dynamic_struct, out_static = filter_make_jaxpr(fn)(
+            *args, **kwargs
+        )  # pyright: ignore
+        jaxpr = closed_jaxpr.jaxpr
+        consts = closed_jaxpr.consts
+        out_dynamic_struct = jtu.tree_flatten(out_dynamic_struct)
+        out_static = jtu.tree_flatten(out_static)
+        closure_converted = _ClosureConvert(
+            jaxpr, consts, in_dynamic_struct, out_dynamic_struct, in_static, out_static
+        )
     closure_converted = cast(Callable[_P, _T], closure_converted)
     return closure_converted
 
@@ -940,3 +976,82 @@ def filter_custom_vjp(fn):
 
     filter_custom_jvp.__doc__ = _filter_custom_jvp_doc
     filter_custom_vjp.__doc__ = _filter_custom_vjp_doc
+
+
+def filter_checkpoint(
+    fun: Callable[_P, _T] = sentinel,
+    *,
+    prevent_cse: bool = True,
+    policy: Optional[Callable[..., bool]] = None,
+) -> Callable[_P, _T]:
+    """Filtered version of `jax.checkpoint`.
+
+    Gradient checkpointing is a technique for reducing memory usage during
+    backpropagation, especially when used with reverse mode automatic differentiation
+    (e.g., `jax.grad` or `equinox.filter_grad`).
+
+    **Arguments:**
+
+    - `fun`: The function to be checkpointed. Will be called as `fun(*args, **kwargs)`.
+        Can return an arbitrary PyTree.
+    - `prevent_cse`: If `True` (the default), then JAX will not perform common
+    subexpression elimination. Please see the documentation for `jax.checkpoint
+    ` for more details.
+    - `policy`: Callable for controlling which intermediate values should be
+    rematerialized. It should be one of the attributes of `jax.checkpoint_policies`.
+    """
+
+    if fun is sentinel:
+        return ft.partial(  # pyright: ignore
+            filter_checkpoint,
+            prevent_cse=prevent_cse,
+            policy=policy,
+        )
+
+    return module_update_wrapper(
+        _CheckpointWrapper(fun, prevent_cse=prevent_cse, policy=policy)
+    )
+
+
+class _CheckpointWrapper(Module):
+    _fun: Callable
+    _prevent_cse: bool
+    _policy: Optional[Callable[..., bool]]
+
+    def __init__(
+        self,
+        fun: Callable,
+        *,
+        prevent_cse: bool = True,
+        policy: Optional[Callable[..., bool]] = None,
+    ):
+        self._fun = fun
+        self._prevent_cse = prevent_cse
+        self._policy = policy
+
+    @property
+    def __wrapped__(self):
+        return self._fun
+
+    def __call__(self, *args, **kwargs):
+        @ft.partial(
+            jax.checkpoint,  # pyright: ignore
+            prevent_cse=self._prevent_cse,
+            policy=self._policy,
+            static_argnums=(0,),
+        )
+        def fun_checkpoint(_static, _dynamic):
+            _args, _kwargs = combine(_static, _dynamic)
+            out = self._fun(*_args, **_kwargs)
+            _dynamic_out, _static_out = partition(out, is_array)
+            return _dynamic_out, Static(_static_out)
+
+        dynamic, static = partition((args, kwargs), is_array)
+        dynamic_out, static_out = fun_checkpoint(static, dynamic)
+
+        return combine(dynamic_out, static_out.value)
+
+    def __get__(self, instance, owner):
+        if instance is None:
+            return self
+        return Partial(self, instance)
diff --git a/equinox/_better_abstract.py b/equinox/_better_abstract.py
index 6971571b..447c0ef6 100644
--- a/equinox/_better_abstract.py
+++ b/equinox/_better_abstract.py
@@ -214,6 +214,7 @@ def __new__(mcs, name, bases, namespace, /, **kwargs):
         cls = super().__new__(mcs, name, bases, namespace, **kwargs)
         abstract_vars = dict()
         abstract_class_vars = dict()
+        cls_annotations = cls.__dict__.get("__annotations__", {})
         for attr, group in [
             ("__abstractvars__", abstract_vars),
             ("__abstractclassvars__", abstract_class_vars),
@@ -233,38 +234,34 @@ def __new__(mcs, name, bases, namespace, /, **kwargs):
                                 "Base classes have mismatched type annotations for "
                                 f"{name}"
                             )
-                    if "__annotations__" in cls.__dict__:
-                        try:
-                            new_annotation = cls.__annotations__[name]
-                        except KeyError:
-                            pass
-                        else:
-                            if not _is_concretisation(new_annotation, annotation):
-                                raise TypeError(
-                                    "Base class and derived class have mismatched type "
-                                    f"annotations for {name}"
-                                )
+                    try:
+                        new_annotation = cls_annotations[name]
+                    except KeyError:
+                        pass
+                    else:
+                        if not _is_concretisation(new_annotation, annotation):
+                            raise TypeError(
+                                "Base class and derived class have mismatched type "
+                                f"annotations for {name}"
+                            )
                     # Not just `if name not in namespace`, as `cls.__dict__` may be
                     # slightly bigger from `__init_subclass__`.
-                    if name not in cls.__dict__:
+                    if name not in cls.__dict__ and name not in cls_annotations:
                         group[name] = annotation
-        if "__annotations__" in cls.__dict__:
-            for name, annotation in cls.__annotations__.items():
-                is_abstract, is_class = _process_annotation(annotation)
-                if is_abstract:
-                    if name in namespace:
-                        if is_class:
-                            raise TypeError(
-                                f"Abstract class attribute {name} cannot have value"
-                            )
-                        else:
-                            raise TypeError(
-                                f"Abstract attribute {name} cannot have value"
-                            )
+        for name, annotation in cls_annotations.items():
+            is_abstract, is_class = _process_annotation(annotation)
+            if is_abstract:
+                if name in namespace:
                     if is_class:
-                        abstract_class_vars[name] = annotation
+                        raise TypeError(
+                            f"Abstract class attribute {name} cannot have value"
+                        )
                     else:
-                        abstract_vars[name] = annotation
+                        raise TypeError(f"Abstract attribute {name} cannot have value")
+                if is_class:
+                    abstract_class_vars[name] = annotation
+                else:
+                    abstract_vars[name] = annotation
         cls.__abstractvars__ = abstract_vars  # pyright: ignore
         cls.__abstractclassvars__ = abstract_class_vars  # pyright: ignore
         return cls
@@ -277,17 +274,11 @@ def __call__(cls, *args, **kwargs):
                 f"attributes {abstract_class_vars}"
             )
         self = super().__call__(*args, **kwargs)
-        abstract_vars = set()
-        for name in cls.__abstractvars__:  # pyright: ignore
-            # Deliberately not doing `if name in self.__dict__` to allow for use of
-            # properties (which are actually class attributes) to override abstract
-            # instance variables.
-            if getattr(self, name, _sentinel) is _sentinel:
-                abstract_vars.add(name)
-        if len(abstract_vars) > 0:
+        if len(cls.__abstractvars__) > 0:  # pyright: ignore
+            abstract_class_vars = set(cls.__abstractvars__)  # pyright: ignore
             raise TypeError(
                 f"Can't instantiate abstract class {cls.__name__} with abstract "
-                f"attributes {abstract_vars}"
+                f"attributes {abstract_class_vars}"
             )
         return self
 
diff --git a/equinox/_config.py b/equinox/_config.py
index 5eb9a571..d1ba6522 100644
--- a/equinox/_config.py
+++ b/equinox/_config.py
@@ -1,4 +1,5 @@
 import os
+import warnings
 from typing import Literal
 
 
@@ -10,3 +11,8 @@
         "Unrecognised value for `EQX_ON_ERROR`. Valid values are `EQX_ON_ERROR=raise`, "
         "`EQX_ON_ERROR=breakpoint`, and `EQX_ON_ERROR=nan`."
     )
+if EQX_ON_ERROR == "breakpoint":
+    warnings.warn(
+        "The environment variable `EQX_ON_ERROR=breakpoint` is currently set. Note "
+        "that this should only be used for debugging, as it slows down runtime speed."
+    )
diff --git a/equinox/_eval_shape.py b/equinox/_eval_shape.py
index 5012eafd..bfd36061 100644
--- a/equinox/_eval_shape.py
+++ b/equinox/_eval_shape.py
@@ -4,8 +4,10 @@
 
 import jax
 import jax._src.traceback_util as traceback_util
+import jax.tree_util as jtu
 from jaxtyping import PyTree
 
+from ._caches import internal_lru_caches
 from ._filters import combine, is_array, partition
 from ._module import Static
 
@@ -35,3 +37,26 @@ def _fn(_static, _dynamic):
     dynamic, static = partition((fun, args, kwargs), _filter)
     dynamic_out, static_out = jax.eval_shape(ft.partial(_fn, static), dynamic)
     return combine(dynamic_out, static_out.value)
+
+
+def _to_struct(x):
+    if is_array(x):
+        return jax.ShapeDtypeStruct(x.shape, x.dtype)
+    else:
+        return x
+
+
+@ft.lru_cache(maxsize=None)
+def _cached_filter_eval_shape(leaves, treedef):
+    fn, args, kwargs = jtu.tree_unflatten(treedef, leaves)
+    return filter_eval_shape(fn, *args, **kwargs)
+
+
+internal_lru_caches.append(_cached_filter_eval_shape)
+
+
+def cached_filter_eval_shape(fn, *args, **kwargs):
+    tree = jtu.tree_map(_to_struct, (fn, args, kwargs))
+    leaves, treedef = jtu.tree_flatten(tree)
+    leaves = tuple(leaves)
+    return _cached_filter_eval_shape(leaves, treedef)
diff --git a/equinox/_jit.py b/equinox/_jit.py
index 633fe99a..b347285c 100644
--- a/equinox/_jit.py
+++ b/equinox/_jit.py
@@ -92,8 +92,9 @@ class XlaRuntimeError(Exception):
 def _modify_traceback(e: Exception):
     # Remove JAX's UnfilteredStackTrace, with its huge error messages.
     e.__cause__ = None
-    # Remove _JitWrapper.__call__ and _JitWrapper._call from the traceback
-    tb = e.__traceback__ = e.__traceback__.tb_next.tb_next  # pyright: ignore
+    # Remove _JitWrapper.__call__ and _JitWrapper._call and Method.__call__ from the
+    # traceback
+    tb = e.__traceback__ = e.__traceback__.tb_next.tb_next.tb_next  # pyright: ignore
     try:
         # See https://github.com/google/jax/blob/69cd3ebe99ce12a9f22e50009c00803a095737c7/jax/_src/traceback_util.py#L190  # noqa: E501
         jax.lib.xla_extension.replace_thread_exc_traceback(tb)  # pyright: ignore
@@ -293,4 +294,4 @@ def f(x, y):  # both args traced if arrays, static if non-arrays
         _cached=cached,
         _filter_warning=filter_warning,
     )
-    return module_update_wrapper(jit_wrapper, fun)
+    return module_update_wrapper(jit_wrapper)
diff --git a/equinox/_make_jaxpr.py b/equinox/_make_jaxpr.py
index ce6f9e55..0cbf2ae3 100644
--- a/equinox/_make_jaxpr.py
+++ b/equinox/_make_jaxpr.py
@@ -76,4 +76,4 @@ def filter_make_jaxpr(
     `int`, `float`, `complex`) are treated as static inputs; wrap them in JAX/NumPy
     arrays if you would like them to be traced.
     """
-    return module_update_wrapper(_MakeJaxpr(fun), fun)
+    return module_update_wrapper(_MakeJaxpr(fun))
diff --git a/equinox/_module.py b/equinox/_module.py
index ed04ec91..23a7182e 100644
--- a/equinox/_module.py
+++ b/equinox/_module.py
@@ -1,9 +1,17 @@
+"""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 types
+import warnings
 import weakref
 from collections.abc import Callable
-from typing import Any, cast, TYPE_CHECKING, TypeVar, Union
+from typing import Any, cast, Optional, Protocol, TYPE_CHECKING, TypeVar, Union
 from typing_extensions import dataclass_transform, ParamSpec
 
 import jax.tree_util as jtu
@@ -21,16 +29,25 @@
 _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)
 
 
-_identity = doc_repr(lambda x: x, "lambda x: x")
+_converter_sentinel: Any = doc_repr(object(), "lambda x: x")
 
 
 def field(
-    *, converter: Callable[[Any], Any] = _identity, static: bool = False, **kwargs
+    *,
+    converter: Callable[[Any], Any] = _converter_sentinel,
+    static: bool = False,
+    **kwargs,
 ):
     """Equinox supports extra functionality on top of the default dataclasses.
 
@@ -38,14 +55,14 @@ def field(
 
     - `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.
+        `bool`/`int`/`float`/`complex` values to JAX arrays. This is ran after the
+        `__init__` method (i.e. when using a user-provided `__init__`), and before
+        `__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`.
 
-    **Converter**
-
-    !!! example
+    !!! example "Example for `converter`"
 
         ```python
         class MyModule(eqx.Module):
@@ -55,9 +72,7 @@ class MyModule(eqx.Module):
         assert isinstance(mymodule.foo, jax.Array)
         ```
 
-    **Static**
-
-    !!! example
+    !!! example "Example for `static`"
 
         ```python
         class MyModule(eqx.Module):
@@ -70,91 +85,335 @@ class MyModule(eqx.Module):
         assert "static" in str(treedef)
         ```
 
-    This means that it does not interact with any JAX transforms, like JIT or grad.
-    This means that it's usually a bug to make JAX arrays be static fields.
-
-    This is an advanced feature that 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.
+    `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["metadata"])
+        metadata = dict(kwargs.pop("metadata"))  # safety copy
     except KeyError:
-        metadata = kwargs["metadata"] = {}
+        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.")
-    metadata["converter"] = converter
+    # 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(**kwargs)
-
-
-class _wrap_method:
-    def __init__(self, method):
-        self.method = method
-        if getattr(self.method, "__isabstractmethod__", False):
-            self.__isabstractmethod__ = self.method.__isabstractmethod__
-
-    def __get__(self, instance, owner):
-        if instance is None:
-            return self.method
-        _method = ft.wraps(self.method)(jtu.Partial(self.method, instance))
-        delattr(_method, "__wrapped__")
-        return _method
-
+    return dataclasses.field(metadata=metadata, **kwargs)
 
-def _not_magic(k: str) -> bool:
-    return not (k.startswith("__") and k.endswith("__"))
 
+#
+# Part 2: Modules!
+# This is the core of Equinox.
+#
 
-_has_dataclass_init = weakref.WeakKeyDictionary()
 
-
-# Inherits from ABCMeta as a convenience for a common use-case.
-# It's not a feature we use ourselves.
+# Inherits from ABCMeta to support `eqx.{AbstractVar, AbstractClassVar}` and
+# `abc.abstractmethod`.
 class _ModuleMeta(ABCMeta):  # pyright: ignore
-    def __new__(mcs, name, bases, dict_):  # pyright: ignore
-        dict_ = {
-            k: _wrap_method(v) if _not_magic(k) and inspect.isfunction(v) else v
-            for k, v in dict_.items()
-        }
-        cls = super().__new__(mcs, name, bases, dict_)
-        # Do override subclasses' dataclass-__init__-s. (None of which call super, so
-        # they must be overriden.)
-        # Don't override custom __init__'s, which leads to poor ergonomics:
-        # e.g. if `B` has a custom init then `class A(B): pass` would otherwise set a
-        # dataclass init that overrides the custom __init__.
-        if "__init__" in cls.__dict__:
-            _init = False
+    # This method is called whenever you definite a module: `class Foo(eqx.Module): ...`
+    def __new__(
+        mcs,
+        name,
+        bases,
+        dict_,
+        /,
+        strict: bool = False,
+        abstract: bool = False,
+        **kwargs,
+    ):
+        # [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__:
+            for kls in cls.__mro__[1:-1]:
                 try:
-                    _init = _has_dataclass_init[kls]
+                    has_dataclass_init = _has_dataclass_init[kls]
                 except KeyError:
-                    pass
+                    # Non-Module superclasses.
+                    if kls.__init__ is not object.__init__:
+                        has_dataclass_init = False
+                        break
                 else:
                     break
             else:
                 assert name == "Module"
-                _init = True  # eqx.Module itself
-        if _init:
+                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
+            )
+
+        # Add support for `eqx.field(converter=...)` when using `__post_init__`.
+        # (Scenario (c) above. Scenarios (a) and (b) are handled later.)
+        if has_dataclass_init and hasattr(cls, "__post_init__"):
+            post_init = cls.__post_init__
+
+            @ft.wraps(post_init)  # pyright: ignore
+            def __post_init__(self, *args, **kwargs):
+                # This `if` is to handle `super()` correctly.
+                # We want to only convert once, at the top level.
+                #
+                # This check is basically testing whether or not the function we're in
+                # now (`cls.__post_init__`) is at the top level
+                # (`self.__class__.__post_init__`). If we are, do conversion. If we're
+                # not, it's presumably because someone is calling us via `super()` in
+                # the middle of their own `__post_init__`. No conversion then; their own
+                # version of this wrapper will do it at the appropriate time instead.
+                #
+                # One small foible: we write `cls.__post_init__`, rather than just
+                # `__post_init__`, to refer to this function. This allows someone else
+                # to also monkey-patch `cls.__post_init__` if they wish, and this won't
+                # remove conversion. (Conversion is a at-the-top-level thing, not a
+                # this-particular-function thing.)
+                #
+                # This top-level business means that this is very nearly the same as
+                # doing conversion in `_ModuleMeta.__call__`. The differences are that
+                # (a) that wouldn't allow us to convert fields before the user-provided
+                # `__post_init__`, and (b) it allows other libraries (i.e. jaxtyping)
+                # to later monkey-patch `__init__`, and we have our converter run before
+                # their own monkey-patched-in code.
+                if self.__class__.__post_init__ is cls.__post_init__:
+                    # Convert all fields currently available.
+                    _convert_fields(self, init=True)
+                post_init(self, *args, **kwargs)  # pyright: ignore
+                if self.__class__.__post_init__ is cls.__post_init__:
+                    # Convert all the fields filled in by `__post_init__` as well.
+                    _convert_fields(self, init=False)
+
+            cls.__post_init__ = __post_init__  # pyright: ignore
+        else:
+            post_init = None
+
+        # 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__
-        cls = dataclass(eq=False, repr=False, frozen=True, init=_init)(
+
+        # [Step 4] Register as a dataclass.
+        cls = dataclass(eq=False, repr=False, frozen=True, init=has_dataclass_init)(
             cls  # pyright: ignore
         )
-        # must happen after `dataclass(...)` we use this in `__getattribute__` to avoid
-        # making any property(def __wrapped__) visible until then. We want to be able to
-        # support property(def __wrapped__) for the sake of classes whose instances are
-        # wrappers (i.e. via `module_update_wrapper`), but this means that a
-        # `__wrapped__` object is also visible on the class object itself, despite not
-        # pointing to a callable. `dataclass(...)` spots this and tries to use it form
-        # a signature.
-        _has_dataclass_init[cls] = _init
-        if _init:
+        # [Step 3b] -- finish off building `__init__` methods. Until we'd done
+        # dataclass'ification then we didn't necessarily have our `__init__` method.
+
+        # Registering here records that the `dataclass(...)` call has happened.
+        _has_dataclass_init[cls] = has_dataclass_init
+
+        # Now handle conversion for cases (a) and (b) above, in which there is no
+        # `__post_init__`.
+        if post_init is None:
+            init = cls.__init__
+
+            @ft.wraps(init)  # pyright: ignore
+            def __init__(self, *args, **kwargs):
+                init(self, *args, **kwargs)
+                # Same `if` trick as with `__post_init__`.
+                if self.__class__.__init__ is cls.__init__:
+                    _convert_fields(self, init=True)
+                    _convert_fields(self, init=False)
+
+            cls.__init__ = __init__
+
+        # Assign `__doc__` in case it has been manually overriden:
+        # ```
+        # 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] = 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 _is_abstract(cls):
+                # Invariant: abstract classes have names beginning with `Abstract`.
+                if not cls.__name__.startswith("Abstract"):
+                    raise TypeError(
+                        "Abstract strict `eqx.Module`s must be named starting with "
+                        f"`Abstract`. Got {name} when defining "
+                        f"{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 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.
         jtu.register_pytree_with_keys(
             cls,
             flatten_with_keys=ft.partial(
@@ -165,53 +424,124 @@ def __new__(mcs, name, bases, dict_):  # pyright: ignore
             ),  # pyright: ignore
             unflatten_func=ft.partial(_unflatten_module, cls),  # pyright: ignore
         )
+        # Done!
         return cls
 
-    # See note above for `_has_dataclass_init`.
-    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
-
+    # This method is called whenever you initialise a module: `MyModule(...)`
     def __call__(cls, *args, **kwargs):
-        # Defreeze it during __init__
-        initable_cls = _make_initable(cls, wraps=False)
+        if _is_force_abstract[cls]:
+            # Any other is-abstract checks will be handled in super().__call__.
+            raise TypeError("Cannot instantiate abstract `equinox.Module`.")
+        # [Step 1] Modules are immutable -- except during construction. So defreeze
+        # before init.
+        post_init = getattr(cls, "__post_init__", None)
+        initable_cls = _make_initable(cls, cls.__init__, post_init, wraps=False)
+        # [Step 2] Instantiate the class as normal.
         self = super(_ModuleMeta, initable_cls).__call__(*args, **kwargs)
+        assert not _is_abstract(cls)
+        # [Step 3] Check that all fields are occupied.
         missing_names = {
             field.name
             for field in dataclasses.fields(cls)  # pyright: ignore
-            if field.init and field.name not in dir(self)
+            # 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}"
             )
-        for field in dataclasses.fields(self):
+        # Freeze.
+        object.__setattr__(self, "__class__", cls)
+        # [Step 4] 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:
-                converter = field.metadata["converter"]
+                check = kls.__dict__["__check_init__"]
             except KeyError:
                 pass
             else:
-                setattr(self, field.name, converter(getattr(self, field.name)))
-        object.__setattr__(self, "__class__", cls)
+                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
+
 
 if TYPE_CHECKING:
-    import abc
 
     @dataclass_transform(field_specifiers=(dataclasses.field, field, static_field))
     class _ModuleMeta(abc.ABCMeta):
         pass
 
 
+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__
+
+    def __get__(self, instance, 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__",
@@ -221,8 +551,18 @@ class _ModuleMeta(abc.ABCMeta):
 }
 
 
+class _Initable:
+    pass
+
+
 @ft.lru_cache(maxsize=128)
-def _make_initable(cls: _ModuleMeta, wraps: bool) -> _ModuleMeta:
+def _make_initable(cls: _ModuleMeta, init, post_init, wraps: bool) -> _ModuleMeta:
+    # 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:
@@ -230,17 +570,52 @@ def _make_initable(cls: _ModuleMeta, wraps: bool) -> _ModuleMeta:
             field.name for field in dataclasses.fields(cls)  # pyright: ignore
         }
 
-    class _InitableModule(cls):  # pyright: ignore
+    class _InitableModule(cls, _Initable):  # pyright: ignore
         pass
 
-    # Done like this to avoid dataclasses complaining about overriding setattr on a
-    # frozen class.
     def __setattr__(self, name, value):
         if name in field_names:
-            object.__setattr__(self, name, 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.
+"""
+                )
+            else:
+                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
@@ -254,8 +629,24 @@ def __setattr__(self, name, value):
 internal_lru_caches.append(_make_initable)
 
 
-# This class exists primarily just to hide the wrapper fields from the PyTreeDef repr.
-# Stuff like the docstring can be pretty long and entirely unhelpful to print out.
+def _convert_fields(module, init: bool):
+    for field in dataclasses.fields(module):
+        if field.init is init:
+            try:
+                converter = field.metadata["converter"]
+            except KeyError:
+                pass
+            else:
+                try:
+                    value = getattr(module, field.name)
+                except AttributeError:
+                    # Let the all-fields-are-filled check handle the error.
+                    pass
+                else:
+                    setattr(module, field.name, converter(value))
+
+
+# Used to provide a pretty repr when doing `jtu.tree_structure(SomeModule(...))`.
 @dataclass()
 class _FlattenedData:
     dynamic_field_names: tuple
@@ -274,17 +665,24 @@ def __repr__(self):
 
 
 def _flatten_module(module: "Module", with_keys: bool):
+    # Subnodes in the PyTree
     dynamic_field_names = []
     dynamic_field_values = []
+    # Static metadata, placed in aux.
     static_field_names = []
     static_field_values = []
+    # Python metadata like `__doc__` and `__module__`.
     wrapper_field_names = []
     wrapper_field_values = []
+
     for field_ in dataclasses.fields(module):
         name = field_.name
         try:
+            # Not `getattr` so that we don't pick up `property`s.
             value = module.__dict__[name]
         except KeyError:
+            # Uninitialised values during `__init__`, or when `property`s overwrite a
+            # field.
             continue
         if field_.metadata.get("static", False):
             static_field_names.append(name)
@@ -312,6 +710,11 @@ def _flatten_module(module: "Module", with_keys: bool):
 
 
 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)
@@ -409,7 +812,7 @@ 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
+    !!! 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
@@ -419,6 +822,12 @@ def __call__(self, x):
         [`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)))
 
@@ -431,12 +840,38 @@ def __repr__(self):
         return tree_pformat(self)
 
 
-# Modifies in-place, just like functools.update_wrapper
+# 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):
+        return self.__func__(self.__self__, *args, **kwargs)
+
+    @property
+    def __wrapped__(self):
+        return self.__func__.__get__(  # pyright: ignore
+            self.__self__, type(self.__self__)
+        )
+
+
+#
+# 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]
+    wrapper: Module, wrapped: Optional[Callable[_P, _T]] = None
 ) -> Callable[_P, _T]:
     """Like `functools.update_wrapper` (or its better-known cousin, `functools.wraps`),
-    but can be used on [`equinox.Module`][]s. (Which are normally immutable.)
+    but acts on [`equinox.Module`][]s, and does not modify its input (it returns the
+    updated module instead).
 
     !!! Example
 
@@ -452,7 +887,7 @@ def __wrapped__(self):
                 return self.fn
 
         def make_wrapper(fn):
-            return eqx.module_update_wrapper(Wrapper(fn), fn)
+            return eqx.module_update_wrapper(Wrapper(fn))
         ```
 
     For example, [`equinox.filter_jit`][] returns a module representing the JIT'd
@@ -461,11 +896,37 @@ def make_wrapper(fn):
 
     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: Optional[Callable[_P, _T]], 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.")
-    initable_cls = _make_initable(cls, wraps=True)
+
+    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=())`.
@@ -487,19 +948,6 @@ def make_wrapper(fn):
     return cast(Callable[_P, _T], wrapper)
 
 
-class Static(Module):
-    _leaves: list[Any] = field(static=True)
-    _treedef: PyTreeDef = field(static=True)  # pyright: ignore
-
-    def __init__(self, value: Any):
-        # Handle pytrees without `__eq__` methods.
-        self._leaves, self._treedef = jtu.tree_flatten(value)
-
-    @property
-    def value(self):
-        return jtu.tree_unflatten(self._treedef, self._leaves)
-
-
 class Partial(Module):
     """Like `functools.partial`, but treats the wrapped function, and partially-applied
     value: Any = field(static=True)
@@ -540,3 +988,25 @@ def __call__(self, *args, **kwargs):
         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 76841c81..d922f921 100644
--- a/equinox/_pretty_print.py
+++ b/equinox/_pretty_print.py
@@ -10,6 +10,8 @@
 import numpy as np
 from jaxtyping import PyTree
 
+from ._doc_utils import WithRepr
+
 
 Dataclass = Any
 NamedTuple = Any  # workaround typeguard bug
@@ -140,7 +142,7 @@ def _pformat_dataclass(obj, **kwargs) -> pp.Doc:
     # values to the module so the repr fails.
     objs = named_objs(
         [
-            (field.name, getattr(obj, field.name, "<uninitialised>"))
+            (field.name, getattr(obj, field.name, WithRepr(None, "<uninitialised>")))
             for field in dataclasses.fields(obj)
             if field.repr
         ],
diff --git a/equinox/_serialisation.py b/equinox/_serialisation.py
index c9fec56e..534337d8 100644
--- a/equinox/_serialisation.py
+++ b/equinox/_serialisation.py
@@ -1,3 +1,4 @@
+import functools as ft
 import pathlib
 from collections.abc import Callable
 from contextlib import contextmanager
@@ -9,6 +10,12 @@
 import numpy as np
 from jaxtyping import PyTree
 
+from ._filters import is_array_like
+
+
+class TreePathError(RuntimeError):
+    path: tuple
+
 
 def _ordered_tree_map(
     f: Callable[..., Any],
@@ -20,9 +27,24 @@ def _ordered_tree_map(
     in fixed order. (Namely depth-first left-to-right.)
     """
     # Discussion: https://github.com/patrick-kidger/equinox/issues/136
-    leaves, treedef = jtu.tree_flatten(tree, is_leaf)
-    all_leaves = [leaves] + [treedef.flatten_up_to(r) for r in rest]
-    return treedef.unflatten(f(*xs) for xs in zip(*all_leaves))
+    paths_and_leaves, treedef = jtu.tree_flatten_with_path(tree, is_leaf)
+    all_leaves = list(zip(*paths_and_leaves)) + [treedef.flatten_up_to(r) for r in rest]
+
+    @ft.wraps(f)
+    def _f(path, *xs):
+        try:
+            return f(*xs)
+        except TreePathError as e:
+            combo_path = path + e.path
+            exc = TreePathError(f"Error at leaf with path {combo_path}")
+            exc.path = combo_path
+            raise exc from e
+        except Exception as e:
+            exc = TreePathError(f"Error at leaf with path {path}")
+            exc.path = path
+            raise exc from e
+
+    return treedef.unflatten(_f(*xs) for xs in zip(*all_leaves))
 
 
 def default_serialise_filter_spec(f: BinaryIO, x: Any) -> None:
@@ -60,9 +82,11 @@ def default_serialise_filter_spec(f: BinaryIO, x: Any) -> None:
     if isinstance(x, jax.Array):
         jnp.save(f, x)
     elif isinstance(x, np.ndarray):
+        # Important to use `np` here so that we don't cast NumPy arrays to JAX arrays.
         np.save(f, x)
-    elif isinstance(x, (bool, float, complex, int)):
-        np.save(f, x)
+    elif is_array_like(x):
+        # Important to use `jnp` here to handle `bfloat16`.
+        jnp.save(f, x)
     else:
         pass
 
@@ -102,9 +126,16 @@ def default_deserialise_filter_spec(f: BinaryIO, x: Any) -> Any:
     if isinstance(x, jax.Array):
         return jnp.load(f)
     elif isinstance(x, np.ndarray):
+        # Important to use `np` here to avoid promoting NumPy arrays to JAX.
         return np.load(f)
-    elif isinstance(x, (bool, float, complex, int)):
-        return np.load(f).item()
+    elif is_array_like(x):
+        # np.generic gets deserialised directly as an array, so convert back to a scalar
+        # type here.
+        # See also https://github.com/google/jax/issues/17858
+        out = np.load(f)
+        if isinstance(x, jax.dtypes.bfloat16):
+            out = out.view(jax.dtypes.bfloat16)
+        return type(x)(out.item())
     else:
         return x
 
@@ -131,22 +162,22 @@ def _maybe_open(path_or_file: Union[str, pathlib.Path, BinaryIO], mode: str):
         yield path_or_file
 
 
-def _assert_same(new, old):
+def _assert_same(path, new, old):
     if type(new) is not type(old):
         raise RuntimeError(
-            f"Deserialised leaf has changed type from {type(old)} in `like` to "
-            f"{type(new)} on disk."
+            f"Deserialised leaf at path {path} has changed type from "
+            f"{type(old)} in `like` to {type(new)} on disk."
         )
     if isinstance(new, (np.ndarray, jax.Array)):
         if new.shape != old.shape:
             raise RuntimeError(
-                f"Deserialised leaf has changed shape from {old.shape} in `like` to "
-                f"{new.shape} on disk."
+                f"Deserialised leaf at path {path} has changed shape from "
+                f"{old.shape} in `like` to {new.shape} on disk."
             )
         if new.dtype != old.dtype:
             raise RuntimeError(
-                f"Deserialised leaf has changed dtype from {old.dtype} in `like` to "
-                f"{new.dtype} on disk."
+                f"Deserialised leaf at path {path} has changed dtype from "
+                f"{old.dtype} in `like` to {new.dtype} on disk."
             )
 
 
@@ -268,5 +299,5 @@ def __deserialise(y):
             return _ordered_tree_map(__deserialise, x, is_leaf=is_leaf)
 
         out = _ordered_tree_map(_deserialise, filter_spec, like)
-    jtu.tree_map(_assert_same, out, like, is_leaf=is_leaf)
+    jtu.tree_map_with_path(_assert_same, out, like, is_leaf=is_leaf)
     return out
diff --git a/equinox/_tree.py b/equinox/_tree.py
index 2ca38663..8dfc26d2 100644
--- a/equinox/_tree.py
+++ b/equinox/_tree.py
@@ -1,6 +1,7 @@
 from collections.abc import Callable, Sequence
 from typing import Any, Optional, TYPE_CHECKING, Union
 
+import jax.numpy as jnp
 import jax.tree_util as jtu
 import numpy as np
 from jaxtyping import Array, Bool, PyTree, PyTreeDef
@@ -230,11 +231,11 @@ def _make_replacement(x: _Node) -> Any:
 
 def tree_equal(*pytrees: PyTree) -> Union[bool, np.bool_, Bool[Array, ""]]:
     """Returns `True` if all input PyTrees are equal. Every PyTree must have the same
-    structure. Any JAX or NumPy arrays (as leaves) must have the same shape, dtype, and
-    values to be considered equal. JAX arrays and NumPy arrays are considered equal
-    to each other.
+    structure, and all leaves must be equal. For JAX arrays, NumPy arrays, or NumPy
+    scalars: they must have the same shape, dtype, and values to be considered equal.
+    JAX and NumPy are considered equal to each other.
 
-    If used under JIT then this may return a tracer.
+    If used under JIT, and any arrays are present, then this may return a tracer.
 
     **Arguments:**
 
@@ -255,9 +256,13 @@ def tree_equal(*pytrees: PyTree) -> Union[bool, np.bool_, Bool[Array, ""]]:
                 if is_array(elem_):
                     if (elem.shape != elem_.shape) or (elem.dtype != elem_.dtype):
                         return False
-                    allsame = (elem == elem_).all()
-                    if allsame is False:
-                        return False
+                    if isinstance(elem, (np.ndarray, np.generic)) and isinstance(
+                        elem_, (np.ndarray, np.generic)
+                    ):
+                        # Avoid promoting to a tracer if possible.
+                        allsame = np.all(elem == elem_)
+                    else:
+                        allsame = jnp.all(elem == elem_)
                     out = out & allsame
                 else:
                     return False
@@ -270,75 +275,6 @@ def tree_equal(*pytrees: PyTree) -> Union[bool, np.bool_, Bool[Array, ""]]:
     return out
 
 
-def _inferences(pytree):
-    is_leaf = lambda x: hasattr(x, "inference") and x is not pytree
-
-    out = [pytree.inference] if hasattr(pytree, "inference") else []
-
-    leaves = [x for x in jtu.tree_leaves(pytree, is_leaf=is_leaf) if is_leaf(x)]
-    # Nodes with an inference flag might have sub-nodes with an inference flag.
-
-    for x in leaves:
-        out.extend(_inferences(x))
-    return out
-
-
-def tree_inference(pytree: PyTree, value: bool) -> PyTree:
-    """Convenience function for setting all `inference` attributes on a PyTree.
-
-    `inference` flags are used to toggle the behaviour of a number of the pre-built
-    neural network layers, such as [`equinox.nn.Dropout`][] or
-    [`equinox.nn.BatchNorm`][].
-
-    !!! Example
-
-        ```python
-        class Model(eqx.Module):
-            norm: eqx.nn.BatchNorm
-            dropout: eqx.nn.Dropout
-            linear: eqx.nn.Linear
-
-            def __init__(self, key):
-                key1, key2 = jax.random.split(key)
-                self.norm = eqx.nn.BatchNorm(3, "batch", key=key1)
-                self.dropout = eqx.nn.Dropout(0.4)
-                self.linear = eqx.nn.Linear(3, 1, key=key2)
-
-            def __call__(self, x, ctx, *, key):
-                x, ctx = self.norm(x, ctx)
-                x = self.dropout(x, key=key)
-                x = self.linear(x)
-                return x, ctx
-
-        training_model = Model(jax.random.PRNGKey(0))
-        inference_model = eqx.tree_inference(training_model, value=True)
-        training_model_again = eqx.tree_inference(inference_model, value=False)
-        ```
-
-    Equivalent to:
-    ```python
-    has_inference = lambda leaf: hasattr(leaf, "inference")
-
-    def where(pytree):
-        return tuple(x.inference
-                     for x in jtu.tree_leaves(pytree, is_leaf=has_inference)
-                     if has_inference(x))
-
-    equinox.tree_at(where, pytree, replace_fn=lambda _: value)
-    ```
-
-    **Arguments:**
-
-    - `pytree`: the PyTree to modify.
-    - `value`: the value to set all `inference` attributes to.
-
-    **Returns:**
-
-    A copy of `pytree` with all `inference` flags set to `value`.
-    """
-    return tree_at(_inferences, pytree, replace_fn=lambda _: value)
-
-
 def tree_flatten_one_level(
     pytree: PyTree,
 ) -> tuple[list[PyTree], PyTreeDef]:  # pyright: ignore
@@ -408,8 +344,42 @@ def is_leaf(node):
 
 
 def tree_check(pytree: Any) -> None:
-    """Checks if the PyTree is well-formed: does it have no repeated nodes, and does it
-    have no self-references.
+    """Checks if the PyTree is well-formed: does it have no self-references, and does
+    it have no duplicate layers.
+
+    Precisely, a "duplicate layer" is any PyTree node with at least one child node.
+
+    !!! info
+
+        This is automatically called when creating an `eqx.Module` instance, to help
+        avoid bugs from duplicating layers.
+
+    !!! Example
+
+        ```python
+        a = 1
+        eqx.tree_check([a, a])  # passes, duplicate is a leaf
+
+        b = eqx.nn.Linear(...)
+        eqx.tree_check([b, b])  # fails, duplicate is nontrivial!
+
+        c = []  # empty list
+        eqx.tree_check([c, c])  # passes, duplicate is trivial
+
+        d = eqx.Module()
+        eqx.tree_check([d, d])  # passes, duplicate is trivial
+
+        eqx.tree_check([None, None])  # passes, duplicate is trivial
+
+        e = [1]
+        eqx.tree_check([e, e])  # fails, duplicate is nontrivial!
+
+        eqx.tree_check([[1], [1]])  # passes, not actually a duplicate: each `[1]`
+                                    # has the same structure, but they're different.
+
+        # passes, not actually a duplicate: each Linear layer is a separate layer.
+        eqx.tree_check([eqx.nn.Linear(...), eqx.nn.Linear(...)])
+        ```
 
     **Arguments:**
 
@@ -427,38 +397,37 @@ def tree_check(pytree: Any) -> None:
     _tree_check(pytree, all_nodes)
 
 
-_trivial_treedef = jtu.tree_structure(0)
+_leaf_treedef = jtu.tree_structure(0)
 
 
 def _tree_check(node, all_nodes):
-    try:
-        self_referential, type_string = all_nodes[id(node)]
-    except KeyError:
-        pass
-    else:
-        if self_referential:
-            raise ValueError(
-                f"PyTree node of type `{type_string}` is self-referential; that is to "
-                "say it appears somewhere within its own PyTree structure. This is "
-                "not allowed."
-            )
-        else:
-            raise ValueError(
-                f"PyTree node of type `{type_string}` appears in the PyTree multiple "
-                "times. This is almost always an error, as these nodes will turn into "
-                "two duplicate copies after flattening/unflattening, e.g. when "
-                "crossing a JIT boundary."
-            )
-    try:
-        type_string = type(node).__name__
-    except AttributeError:
-        # AttributeError: in case we cannot get __name__ for some weird reason.
-        type_string = "<unknown type>"
-    all_nodes[id(node)] = (True, type_string)
     subnodes, treedef = tree_flatten_one_level(node)
-    if treedef != _trivial_treedef:
-        # This does mean that leaves can appear multiple times. This is valid, e.g.
-        # [4, 4].
+    # We allow duplicate leaves and empty containers, so don't raise an error with those
+    if treedef != _leaf_treedef and treedef.num_leaves > 0:
+        try:
+            self_referential, type_string = all_nodes[id(node)]
+        except KeyError:
+            pass
+        else:
+            if self_referential:
+                raise ValueError(
+                    f"PyTree node of type `{type_string}` is self-referential; that is "
+                    "to say it appears somewhere within its own PyTree structure. This "
+                    "is not allowed."
+                )
+            else:
+                raise ValueError(
+                    f"PyTree node of type `{type_string}` appears in the PyTree "
+                    "multiple times. This is almost always an error, as these nodes "
+                    "will turn into two duplicate copies after "
+                    "flattening/unflattening, e.g. when crossing a JIT boundary."
+                )
+        try:
+            type_string = type(node).__name__
+        except AttributeError:
+            # AttributeError: in case we cannot get __name__ for some weird reason.
+            type_string = "<unknown type>"
+        all_nodes[id(node)] = (True, type_string)
         for subnode in subnodes:
             _tree_check(subnode, all_nodes)
-    all_nodes[id(node)] = (False, type_string)
+        all_nodes[id(node)] = (False, type_string)
diff --git a/equinox/_vmap_pmap.py b/equinox/_vmap_pmap.py
index 9653fb0e..d487ce82 100644
--- a/equinox/_vmap_pmap.py
+++ b/equinox/_vmap_pmap.py
@@ -406,7 +406,7 @@ def evaluate_per_ensemble(model, x):
         _axis_size=axis_size,
         _vmapkwargs=vmapkwargs,
     )
-    return module_update_wrapper(vmap_wrapper, fun)
+    return module_update_wrapper(vmap_wrapper)
 
 
 @compile_cache
@@ -722,4 +722,4 @@ def g(x, y):
         _filter_warning=filter_warning,
         _pmapkwargs=pmapkwargs,
     )
-    return module_update_wrapper(pmap_wrapper, fun)
+    return module_update_wrapper(pmap_wrapper)
diff --git a/equinox/internal/README.md b/equinox/internal/README.md
index a15cb566..0779ebb4 100644
--- a/equinox/internal/README.md
+++ b/equinox/internal/README.md
@@ -8,8 +8,6 @@ A few of the highlights available here are:
 
 **_loop/** A while loop which may be backpropagated through, using an online checkpointing scheme. Also implements a recursive checkpointed ("treeverse") scan.
 
-**../_errors.py:** Runtime errors. For verifying arguments or for checking that a computation succeeded. Now working on CPUs+GPUs+TPUs! (The file exists in the main namespace for dependency reasons but is (semi-publicly) exposed only in `equinox.internal`.)
-
 **_noinline.py:** MLIR sub-graphs. Can reduce compile times by removing the inlining of a function called repeatedly. Also makes it possible to iteratively recompile: change and recompile just the sub/super-computation, without needing to recompile the whole computation graph. Only tested on CPU.
 
 **_omega.py**: Neat syntax for tree-mapping arithmetic. E.g. `(x**ω + y**ω).ω == jtu.tree_map(operator.add, x, y)`. See also [tree-math](https://github.com/google/tree-math), for a similar idea but without the neat syntax.
diff --git a/equinox/internal/__init__.py b/equinox/internal/__init__.py
index 48a29731..73c5e6cf 100644
--- a/equinox/internal/__init__.py
+++ b/equinox/internal/__init__.py
@@ -18,6 +18,7 @@
     branched_error_if as branched_error_if,
     error_if as error_if,
 )
+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 .._pretty_print import tree_pp as tree_pp
@@ -41,6 +42,7 @@
     store_dce as store_dce,
 )
 from ..debug._announce_transform import announce_jaxpr_p as announce_jaxpr_p
+from ._closure_to_pytree import closure_to_pytree as closure_to_pytree
 from ._finalise_jaxpr import (
     finalise_eval_jaxpr as finalise_eval_jaxpr,
     finalise_fn as finalise_fn,
diff --git a/equinox/internal/_closure_to_pytree.py b/equinox/internal/_closure_to_pytree.py
new file mode 100644
index 00000000..c0e764d4
--- /dev/null
+++ b/equinox/internal/_closure_to_pytree.py
@@ -0,0 +1,134 @@
+# This is some mildly unpleasant code.
+#
+# Basically, Optax make the decision *not* to register their optimisers as PyTrees.
+# This means that we often end up with spurious recompilation, just because a learning
+# rate changed. That results in a new optimiser instance, which is just a function and
+# is treated statically.
+#
+# So here we simply replace all function closures with pytrees, with each of their cell
+# contents as their subnodes.
+
+import inspect
+import types
+from typing import Any, Optional
+
+import jax.tree_util as jtu
+
+from .._module import Module
+
+
+def _make_cell(val):
+    fn = lambda: val
+    return fn.__closure__[0]  # pyright: ignore
+
+
+def _adjust_function_closure(fn, closure):
+    out = types.FunctionType(
+        code=fn.__code__,
+        globals=fn.__globals__,
+        name=fn.__name__,
+        argdefs=fn.__defaults__,
+        closure=closure,
+    )
+    out.__module__ = fn.__module__
+    out.__qualname__ = fn.__qualname__
+    out.__doc__ = fn.__doc__
+    out.__annotations__.update(fn.__annotations__)
+    if fn.__kwdefaults__ is not None:
+        out.__kwdefaults__ = fn.__kwdefaults__.copy()
+    return out
+
+
+# Not a pytree.
+# Used so that two different local functions, with different identities, can still
+# compare equal. This is needed as these leaves are compared statically when
+# filter-jit'ing.
+class _FunctionWithEquality:
+    def __init__(self, fn: types.FunctionType):
+        self.fn = fn
+
+    def information(self):
+        try:
+            # `fn` not defined in REPL.
+            source = inspect.getsource(self.fn)
+        except OSError:
+            # `fn` defined in REPL. In practice this will lead to recompilations based
+            # on function identity, but correctness >> speed.
+            return self.fn
+        else:
+            return self.fn.__qualname__, self.fn.__module__, source
+
+    def __hash__(self):
+        return hash(self.information())
+
+    def __eq__(self, other):
+        return type(self) == type(other) and self.information() == other.information()
+
+
+class _Closure(Module):
+    fn: _FunctionWithEquality
+    contents: Optional[tuple[Any, ...]]
+
+    def __init__(self, fn: types.FunctionType):
+        self.fn = _FunctionWithEquality(fn)
+        if fn.__closure__ is None:
+            contents = None
+        else:
+            contents = tuple(
+                closure_to_pytree(cell.cell_contents) for cell in fn.__closure__
+            )
+        self.contents = contents
+
+    def __call__(self, *args, **kwargs):
+        if self.contents is None:
+            closure = None
+        else:
+            closure = tuple(_make_cell(contents) for contents in self.contents)
+        fn = _adjust_function_closure(self.fn.fn, closure)
+        return fn(*args, **kwargs)
+
+
+def _fixup_closure(leaf):
+    if isinstance(leaf, types.FunctionType):
+        return _Closure(leaf)
+    else:
+        return leaf
+
+
+def closure_to_pytree(tree):
+    """Convert all function closures into pytree nodes.
+
+    **Arguments:**
+
+    - `tree`: Any pytree.
+
+    **Returns:**
+
+    A copy of `tree`, where all function closures have been replaced by a new object
+    that is (a) callable like the original function, but (b) iterates over its
+    `__closure__` as subnodes in the pytree.
+
+    !!! Example
+
+        ```python
+        def some_fn():
+            a = jnp.array(1.)
+
+            @closure_to_pytree
+            def f(x):
+                return x + a
+
+            print(jax.tree_util.tree_leaves(f))  # prints out `a`
+        ```
+
+    !!! Warning
+
+        One annoying technical detail in the above example: we had to wrap the whole lot
+        in a `some_fn`, so that we're in a local scope. Python treats functions at the
+        global scope differently, and this conversion won't result in any global
+        variable being treated as part of the pytree.
+
+        In practice, the intended use case of this function is to fix Optax, which
+        always uses local functions.
+    """
+    return jtu.tree_map(_fixup_closure, tree)
diff --git a/equinox/internal/_noinline.py b/equinox/internal/_noinline.py
index 940b11b1..0514394d 100644
--- a/equinox/internal/_noinline.py
+++ b/equinox/internal/_noinline.py
@@ -506,4 +506,4 @@ def abstract_fn(__dynamic_fn, *args, **kwargs):
         _index_to_fn.append(static_fn)
     dynamic_index = jnp.array(dynamic_index)
     noinline_fn = _NoInlineWrapper(dynamic_index, abstract_fn, dynamic_fn)
-    return module_update_wrapper(noinline_fn, abstract_fn)
+    return module_update_wrapper(noinline_fn)
diff --git a/equinox/nn/__init__.py b/equinox/nn/__init__.py
index 2f70587c..a5fb3a3d 100644
--- a/equinox/nn/__init__.py
+++ b/equinox/nn/__init__.py
@@ -13,6 +13,7 @@
 )
 from ._dropout import Dropout as Dropout
 from ._embedding import Embedding as Embedding
+from ._inference import inference_mode as inference_mode
 from ._linear import Identity as Identity, Linear as Linear
 from ._mlp import MLP as MLP
 from ._normalisation import GroupNorm as GroupNorm, LayerNorm as LayerNorm
@@ -38,5 +39,11 @@
     Sequential as Sequential,
     StatefulLayer as StatefulLayer,
 )
+from ._shared import Shared as Shared
 from ._spectral_norm import SpectralNorm as SpectralNorm
-from ._stateful import State as State, StateIndex as StateIndex
+from ._stateful import (
+    delete_init_state as delete_init_state,
+    make_with_state as make_with_state,
+    State as State,
+    StateIndex as StateIndex,
+)
diff --git a/equinox/nn/_attention.py b/equinox/nn/_attention.py
index 435f949f..d4b4d5b6 100644
--- a/equinox/nn/_attention.py
+++ b/equinox/nn/_attention.py
@@ -168,7 +168,7 @@ def __init__(
         - `dropout_p`: Dropout probability on attention weights.
         - `inference`: Whether to actually apply dropout at all. If `True` then dropout
             is not applied. If `False` then dropout is applied. This may be toggled
-            with [`equinox.tree_inference`][] or overridden during
+            with [`equinox.nn.inference_mode`][] or overridden during
             [`equinox.nn.MultiheadAttention.__call__`][].
         - `key`: A `jax.random.PRNGKey` used to provide randomness for parameter
             initialisation. (Keyword only argument.)
diff --git a/equinox/nn/_batch_norm.py b/equinox/nn/_batch_norm.py
index 66c9d338..dca0353c 100644
--- a/equinox/nn/_batch_norm.py
+++ b/equinox/nn/_batch_norm.py
@@ -38,7 +38,7 @@ class BatchNorm(StatefulLayer):
     training then statistics are computed using the input data, and the running
     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.tree_inference`][].
+    [`equinox.nn.inference_mode`][].
     """  # noqa: E501
 
     weight: Optional[Float[Array, "input_size"]]
@@ -79,7 +79,7 @@ def __init__(
         - `inference`: If `False` then the batch means and variances will be calculated
             and used to update the running statistics. If `True` then the running
             statistics are directly used for normalisation. This may be toggled with
-            [`equinox.tree_inference`][] or overridden during
+            [`equinox.nn.inference_mode`][] or overridden during
             [`equinox.nn.BatchNorm.__call__`][].
         - `dtype`: The dtype of the input array.
         """
@@ -92,12 +92,12 @@ def __init__(
         else:
             self.weight = None
             self.bias = None
-        self.first_time_index = StateIndex(lambda **_: jnp.array(True))
-        make_buffers = lambda **_: (
+        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(make_buffers)
+        self.state_index = StateIndex(init_buffers)
         self.inference = inference
         self.axis_name = axis_name
         self.input_size = input_size
diff --git a/equinox/nn/_dropout.py b/equinox/nn/_dropout.py
index e3b82960..023a00b6 100644
--- a/equinox/nn/_dropout.py
+++ b/equinox/nn/_dropout.py
@@ -16,7 +16,7 @@ class Dropout(Module):
     Note that this layer behaves differently during training and inference. During
     training then dropout is randomly applied; during inference this layer does nothing.
     Whether the model is in training or inference mode should be toggled using
-    [`equinox.tree_inference`][].
+    [`equinox.nn.inference_mode`][].
     """
 
     # Not static fields as it makes sense to want to modify them via equinox.tree_at.
@@ -35,7 +35,7 @@ def __init__(
         - `p`: The fraction of entries to set to zero. (On average.)
         - `inference`: Whether to actually apply dropout at all. If `True` then dropout
             is *not* applied. If `False` then dropout is applied. This may be toggled
-            with [`equinox.tree_inference`][] or overridden during
+            with [`equinox.nn.inference_mode`][] or overridden during
             [`equinox.nn.Dropout.__call__`][].
         - `deterministic`: Deprecated alternative to `inference`.
         """
diff --git a/equinox/nn/_embedding.py b/equinox/nn/_embedding.py
index acda0bb4..7780bf2b 100644
--- a/equinox/nn/_embedding.py
+++ b/equinox/nn/_embedding.py
@@ -1,9 +1,11 @@
 from typing import Optional
 
 import jax
+import jax.numpy as jnp
 import jax.random as jrandom
-from jaxtyping import Array, Float, PRNGKeyArray
+from jaxtyping import Array, ArrayLike, Float, Int, PRNGKeyArray
 
+from .._filters import is_array_like
 from .._module import field, Module
 
 
@@ -16,28 +18,49 @@ class Embedding(Module):
 
     def __init__(
         self,
-        num_embeddings: int,
-        embedding_size: int,
+        num_embeddings: Optional[int] = None,  # pyright: ignore
+        embedding_size: Optional[int] = None,  # pyright: ignore
         weight: Optional[Float[Array, "num_embeddings embedding_size"]] = None,
         *,
-        key: PRNGKeyArray,
+        key: Optional[PRNGKeyArray] = None,
         **kwargs,
     ):
         """**Arguments:**
 
+        `Embedding` should be initialised with either:
+
         - `num_embeddings`: Size of embedding dictionary. Must be non-negative.
         - `embedding_size`: Size of each embedding vector. Must be non-negative.
-        - `weight`: If given, the embedding lookup table. Will be generated randomly
-            if not provided.
-        - `key`: A `jax.random.PRNGKey` used to provide randomness for parameter
-            initialisation. (Keyword only argument.)
+        - `key`: A `jax.random.PRNGKey` used to provide randomness for initialisation
+            of the embedding lookup table. (Keyword only argument.)
+
+        Or:
+
+        - `weight`: The embedding lookup table, of shape
+            `(num_embeddings, embedding_size)`.
         """
         super().__init__(**kwargs)
         if weight is None:
-            assert num_embeddings >= 0, "num_embeddings must not be negative."
-            assert embedding_size >= 0, "embedding_size must not be negative."
+            if num_embeddings is None or embedding_size is None or key is None:
+                raise ValueError(
+                    "Must provide `eqx.nn.Embedding(num_embeddings=..., "
+                    "embedding_size=..., key=...)` if not providing the weight "
+                    "directly."
+                )
+            if num_embeddings < 0:
+                raise ValueError("num_embeddings must not be negative.")
+            if embedding_size < 0:
+                raise ValueError("embedding_size must not be negative.")
             self.weight = jrandom.normal(key, (num_embeddings, embedding_size))
         else:
+            if weight.ndim != 2:
+                raise ValueError(
+                    "weight must have shape (num_embeddings, embedding_size)."
+                )
+            if num_embeddings is None:
+                num_embeddings: int = weight.shape[0]
+            if embedding_size is None:
+                embedding_size: int = weight.shape[1]
             if weight.shape != (num_embeddings, embedding_size):
                 raise ValueError(
                     "weight must have shape (num_embeddings, embedding_size)."
@@ -47,10 +70,12 @@ def __init__(
         self.embedding_size = embedding_size
 
     @jax.named_scope("eqx.nn.Embedding")
-    def __call__(self, x: Array, *, key: Optional[PRNGKeyArray] = None) -> Array:
+    def __call__(
+        self, x: Int[ArrayLike, ""], *, key: Optional[PRNGKeyArray] = None
+    ) -> Array:
         """**Arguments:**
 
-        - `x`: The table index.
+        - `x`: The table index. Should be a scalar integer array.
         - `key`: Ignored; provided for compatibility with the rest of the Equinox API.
             (Keyword only argument.)
 
@@ -59,4 +84,10 @@ def __call__(self, x: Array, *, key: Optional[PRNGKeyArray] = None) -> Array:
         A JAX array of shape `(embedding_size,)`, from the x-th index of the embedding
         table.
         """
-        return self.weight[x]
+        if is_array_like(x) and jnp.shape(x) == ():
+            return self.weight[x]
+        else:
+            raise ValueError(
+                "`eqx.nn.Embedding()(x)` should be called with a scalar index `x`. "
+                "Use `jax.vmap` if you would like to index with multiple values."
+            )
diff --git a/equinox/nn/_inference.py b/equinox/nn/_inference.py
new file mode 100644
index 00000000..e37f57f3
--- /dev/null
+++ b/equinox/nn/_inference.py
@@ -0,0 +1,74 @@
+import jax.tree_util as jtu
+from jaxtyping import PyTree
+
+from .._tree import tree_at
+
+
+def _inferences(pytree):
+    is_leaf = lambda x: hasattr(x, "inference") and x is not pytree
+
+    out = [pytree.inference] if hasattr(pytree, "inference") else []
+
+    leaves = [x for x in jtu.tree_leaves(pytree, is_leaf=is_leaf) if is_leaf(x)]
+    # Nodes with an inference flag might have sub-nodes with an inference flag.
+
+    for x in leaves:
+        out.extend(_inferences(x))
+    return out
+
+
+def inference_mode(pytree: PyTree, value: bool = True) -> PyTree:
+    """Convenience function for setting all `inference` attributes.
+
+    `inference` flags are used to toggle the behaviour of a number of the pre-built
+    neural network layers, such as [`equinox.nn.Dropout`][] or
+    [`equinox.nn.BatchNorm`][].
+
+    !!! Example
+
+        ```python
+        class Model(eqx.Module):
+            norm: eqx.nn.BatchNorm
+            dropout: eqx.nn.Dropout
+            linear: eqx.nn.Linear
+
+            def __init__(self, key):
+                key1, key2 = jax.random.split(key)
+                self.norm = eqx.nn.BatchNorm(3, "batch", key=key1)
+                self.dropout = eqx.nn.Dropout(0.4)
+                self.linear = eqx.nn.Linear(3, 1, key=key2)
+
+            def __call__(self, x, ctx, *, key):
+                x, ctx = self.norm(x, ctx)
+                x = self.dropout(x, key=key)
+                x = self.linear(x)
+                return x, ctx
+
+        training_model = Model(jax.random.PRNGKey(0))
+        inference_model = eqx.nn.inference_mode(training_model)
+        training_model_again = eqx.nn.inference_mode(inference_model, value=False)
+        ```
+
+    This function is essentially equivalent to:
+    ```python
+    has_inference = lambda leaf: hasattr(leaf, "inference")
+
+    def where(pytree):
+        return tuple(x.inference
+                     for x in jtu.tree_leaves(pytree, is_leaf=has_inference)
+                     if has_inference(x))
+
+    inference_pytree = equinox.tree_at(where, pytree, replace_fn=lambda _: value)
+    ```
+
+    **Arguments:**
+
+    - `pytree`: the PyTree to modify.
+    - `value`: the value to set all `inference` attributes to. Defaults to `True`, i.e.
+        inference mode.
+
+    **Returns:**
+
+    A copy of `pytree` with all `inference` flags set to `value`.
+    """
+    return tree_at(_inferences, pytree, replace_fn=lambda _: value)
diff --git a/equinox/nn/_normalisation.py b/equinox/nn/_normalisation.py
index d15276ef..88865e24 100644
--- a/equinox/nn/_normalisation.py
+++ b/equinox/nn/_normalisation.py
@@ -5,7 +5,7 @@
 
 import jax
 import jax.numpy as jnp
-from jaxtyping import Array, PRNGKeyArray
+from jaxtyping import Array, Float, PRNGKeyArray
 
 from .._custom_types import sentinel
 from .._misc import left_broadcast_to
@@ -50,16 +50,16 @@ class LayerNorm(Module):
 
     """
 
-    shape: Union[None, int, Sequence[int]] = field(static=True)
+    shape: tuple[int] = field(static=True)
     eps: float = field(static=True)
     use_weight: bool = field(static=True)
     use_bias: bool = field(static=True)
-    weight: Optional[Array]
-    bias: Optional[Array]
+    weight: Optional[Float[Array, "*shape"]]
+    bias: Optional[Float[Array, "*shape"]]
 
     def __init__(
         self,
-        shape: Union[None, int, Sequence[int]],
+        shape: Union[int, Sequence[int]],
         eps: float = 1e-5,
         use_weight: bool = True,
         use_bias: bool = True,
@@ -69,14 +69,17 @@ def __init__(
     ):
         """**Arguments:**
 
-        - `shape`: Input shape. May be left unspecified (e.g. just `None`) if
-            `elementwise_affine=False`.
+        - `shape`: Shape of the input.
         - `eps`: Value added to denominator for numerical stability.
         - `use_weight`: Whether the module has learnable affine weights.
         - `use_bias`: Whether the module has learnable affine biases.
         - `elementwise_affine`: Deprecated alternative to `use_weight` and `use_bias`.
         """
         super().__init__(**kwargs)
+        if isinstance(shape, int):
+            shape = (shape,)
+        else:
+            shape = tuple(shape)
         self.shape = shape
         self.eps = eps
         if elementwise_affine is not None:
@@ -103,11 +106,15 @@ def __call__(
 
     @jax.named_scope("eqx.nn.LayerNorm")
     def __call__(
-        self, x: Array, state: State = sentinel, *, key: Optional[PRNGKeyArray] = None
+        self,
+        x: Float[Array, "*shape"],
+        state: State = sentinel,
+        *,
+        key: Optional[PRNGKeyArray] = None,
     ) -> Union[Array, tuple[Array, State]]:
         """**Arguments:**
 
-        - `x`: A JAX array.
+        - `x`: A JAX array, with the same shape as the `shape` passed to `__init__`.
         - `state`: Ignored; provided for interchangability with the
             [`equinox.nn.BatchNorm`][] API.
         - `key`: Ignored; provided for compatibility with the rest of the Equinox API.
@@ -121,6 +128,17 @@ def __call__(
         is passed through unchanged. If `state` is not passed, then just the output is
         returned.
         """
+        if x.shape != self.shape:
+            raise ValueError(
+                "`LayerNorm(shape)(x)` must satisfy the invariant `shape == x.shape`"
+                f"Received `shape={self.shape} and `x.shape={x.shape}`. You might need "
+                "to replace `layer_norm(x)` with `jax.vmap(layer_norm)(x)`.\n"
+                "\n"
+                "If this is a new error for you, it might be because this became "
+                "stricter in Equinox v0.11.0. Previously all that was required is that "
+                "`x.shape` ended with `shape`. However, this turned out to be a "
+                "frequent source of bugs, so we made the check stricter!"
+            )
         mean = jnp.mean(x, keepdims=True)
         variance = jnp.var(x, keepdims=True)
         variance = jnp.maximum(0.0, variance)
diff --git a/equinox/nn/_sequential.py b/equinox/nn/_sequential.py
index c335251f..cc11e84d 100644
--- a/equinox/nn/_sequential.py
+++ b/equinox/nn/_sequential.py
@@ -1,6 +1,5 @@
-import abc
 from collections.abc import Callable, Sequence
-from typing import Any, Optional, overload, Union
+from typing import Any, ClassVar, Optional, overload, Union
 
 import jax
 import jax.random as jr
@@ -11,7 +10,36 @@
 from ._stateful import State
 
 
-class Sequential(Module):
+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 call the layer as
+    `new_x = layer(x)` or `(new_x, new_state) = layer(x, state)`.
+    """
+
+    def is_stateful(self) -> bool:
+        """Indicates whether this layer should be considered stateful.
+
+        The default implementation just returns True, but subclasses may override this
+        to provide custom logic if the layer is only "maybe stateful". (E.g. if they
+        optioanlly use stateful sublayers themselves.)
+
+        **Arguments:**
+
+        None
+
+        **Returns:**
+
+        A boolean. `True` indicates that the layer should be called as
+        `(new_x, new_state) = layer(x, state)`. `False` indicates that the layer should
+        be called as `new_x = layer(x)`.
+        """
+        return True
+
+    __call__: ClassVar[Callable]
+
+
+class Sequential(StatefulLayer):
     """A sequence of [`equinox.Module`][]s applied in order.
 
     !!! note
@@ -24,6 +52,11 @@ class Sequential(Module):
     def __init__(self, layers: Sequence[Callable]):
         self.layers = tuple(layers)
 
+    def is_stateful(self) -> bool:
+        return any(
+            isinstance(x, StatefulLayer) and x.is_stateful() for x in self.layers
+        )
+
     @overload
     def __call__(self, x: Array, *, key: Optional[PRNGKeyArray] = None) -> Array:
         ...
@@ -62,7 +95,7 @@ def __call__(
         else:
             keys = jr.split(key, len(self.layers))
         for layer, key in zip(self.layers, keys):
-            if isinstance(layer, StatefulLayer):
+            if isinstance(layer, StatefulLayer) and layer.is_stateful():
                 x, state = layer(x, state=state, key=key)
             else:
                 x = layer(x, key=key)
@@ -92,30 +125,6 @@ def __len__(self):
 """
 
 
-class StatefulLayer(Module):
-    """An abstract base class, used to mark a stateful layer for the sake of
-    [`equinox.nn.Sequential`][]. If `Sequential` sees that a layer inherits
-    from `StatefulLayer`, then it will know to pass in `state` as well as the
-    piped data `x`.
-
-    Subclasses must implement the `__call__` method that takes input data and the
-    current state as arguments and returns the output data and updated state.
-    """
-
-    @abc.abstractmethod
-    def __call__(
-        self,
-        x: Array,
-        state: State,
-        *,
-        key: Optional[PRNGKeyArray],
-    ) -> tuple[Array, State]:
-        """The function signature that stateful layers should conform to, to be
-        compatible with [`equinox.nn.Sequential`][].
-        """
-        raise NotImplementedError("Subclasses must implement the __call__ method.")
-
-
 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
new file mode 100644
index 00000000..4f8b52fe
--- /dev/null
+++ b/equinox/nn/_shared.py
@@ -0,0 +1,158 @@
+from collections.abc import Callable
+
+from jaxtyping import PyTree
+
+from .._eval_shape import filter_eval_shape
+from .._module import Module
+from .._tree import tree_at, tree_equal
+
+
+class SharedNode:
+    """Placeholder value for nodes that have been removed by `eqx.nn.Shared`."""
+
+    def __repr__(self):
+        return "SharedNode"
+
+
+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
+    difference parts of the tree, will be treated as two copies of this layer. For
+    example,
+    ```python
+    class SubModel(eqx.Module):
+        linear: eqx.nn.Linear
+
+    class Model(eqx.Module):
+        linear: eqx.nn.Linear
+        submodel: SubModel
+
+        def __init__(self):
+            linear = eqx.nn.Linear(...)
+            self.linear = linear
+            self.submodel = SubModel(linear)
+    ```
+    is used to declare `model.linear` and `model.submodel.linear` as two separate
+    layers. They will start with the same initial parameter values, and then update
+    independently during training.
+
+    For when we really do want to share layers or weights across different parts of a
+    model, then `eqx.nn.Shared` exists as a way to easily express this in the PyTree
+    paradigm.
+
+    !!! Example
+
+        It is common in many language models to have an initial embedding matrix at the
+        start, and then to reuse this as the weight of the final linear transformation.
+
+        ```python
+        import equinox as eqx
+        import jax.numpy as jnp
+        from jaxtyping import Array, Int
+
+        class LanguageModel(eqx.Module):
+            shared: eqx.nn.Shared
+
+            def __init__(self):
+                embedding = eqx.nn.Embedding(...)
+                linear = eqx.nn.Linear(...)
+                # These two weights will now be tied together.
+                where = lambda embed_and_lin: embed_and_lin[1].weight
+                get = lambda embed_and_lin: embed_and_lin[0].weight
+                self.shared = eqx.nn.Shared((embedding, linear), where, get)
+
+            def __call__(self, tokens: Int[Array, "sequence"]):
+                # Expand back out so we can evaluate these layers.
+                embedding, linear = self.shared()
+                assert embedding.weight is linear.weight  # same parameter!
+                # Now go ahead and evaluate your language model.
+                values = jax.vmap(embedding)(tokens)
+                ...  # other layers, probably
+                return jax.vmap(linear)(values)
+        ```
+
+        _(Side note: you will sometimes see some authors referring to transposing
+        the embedding matrix prior to the final linear layer. This is because some
+        other libraries store the weight matrices of linear layers the other way
+        around. If that had been necessary here then we could have done it with
+        `get = lambda embed_and_lin: jnp.transpose(embed_and_lin[0].weight)`.)_
+
+    """
+
+    pytree: PyTree
+    where: Callable
+    get: Callable
+
+    def __init__(self, pytree: PyTree, where: Callable, get: Callable):
+        """**Arguments:**
+
+        - `pytree`: The PyTree to share some nodes across.
+        - `where`: a function specifying either a single node, or a sequence of nodes,
+            as with `eqx.tree_at(where, pytree, ...)`.
+        - `get`: a function, which when evaluated on `pytree`, returns either a single
+            value (if `where` does), or a sequence of values (if `where` does, and in
+            this case this must be a sequence of the same length as `where`).
+
+        The node(s) of `get(pytree)` and the corresponding value(s) of `where(pytree)`
+        will be tied together.
+
+        !!! info
+
+            To explain how this works. The implementation is just:
+            ```python
+            class Shared(eqx.Module):
+                pytree: PyTree
+                where: Callable
+                get: Callable
+
+                def __init__(self, pytree, where, get):
+                    # `0` is just some dummy value
+                    self.pytree = eqx.tree_at(where, pytree, replace_fn=lambda _: 0)
+                    self.where = where
+                    self.get = get
+
+                def __call__(self):
+                    return eqx.tree_at(self.where, self.pytree, self.get(self.pytree))
+            ```
+            so that at `__init__` time, the duplicate nodes specified in `where` are
+            removed from the PyTree. We no longer have a separate copy updating during
+            training.
+
+            And then at `__call__` time, references to the values returned by
+            `get(pytree)` are put in their place. We end up with a pytree of the same
+            structure as what we started with, which we can now use (evaluate as a
+            layer etc.) as normal.
+
+        !!! tip
+
+            If you need to apply any transform (e.g. transposing a matrix), then this
+            can be done as part of `get`. For example,
+            `get = lambda pair: jnp.transpose(pair[1].weight)`.
+        """
+
+        source_struct = filter_eval_shape(get, pytree)
+        dest_struct = filter_eval_shape(where, pytree)
+        if tree_equal(source_struct, dest_struct) is not True:
+            raise ValueError(
+                "Every node being shared together must have the same pytree "
+                "structure, shape+dtype of arrays, etc., as each other. Got:\n"
+                f"{source_struct}\n"
+                "and\n"
+                f"{dest_struct}"
+            )
+        self.pytree = tree_at(where, pytree, replace_fn=lambda _: SharedNode())
+        self.where = where
+        self.get = get
+
+    def __call__(self):
+        """**Arguments:**
+
+        None.
+
+        **Returns:**
+
+        A PyTree of the same structure as the original `pytree`, with `get(pytree)` in
+        the place of the nodes at `where(pytree)`.
+        """
+        return tree_at(self.where, self.pytree, self.get(self.pytree))
diff --git a/equinox/nn/_spectral_norm.py b/equinox/nn/_spectral_norm.py
index 4998c607..618df99f 100644
--- a/equinox/nn/_spectral_norm.py
+++ b/equinox/nn/_spectral_norm.py
@@ -49,7 +49,7 @@ class SpectralNorm(StatefulLayer, Generic[_Layer]):
     Note that this layer behaves differently during training and inference. During
     training then power iterations are updated; during inference they are fixed.
     Whether the model is in training or inference mode should be toggled using
-    [`equinox.tree_inference`][].
+    [`equinox.nn.inference_mode`][].
     """  # noqa: E501
 
     layer: _Layer
@@ -81,7 +81,7 @@ def __init__(
         - `eps`: Epsilon for numerical stability when calculating norms.
         - `inference`: Whether this is in inference mode, at which time no power
             iterations are performed.  This may be toggled with
-            [`equinox.tree_inference`][].
+            [`equinox.nn.inference_mode`][].
         - `key`: A `jax.random.PRNGKey` used to provide randomness for initialisation.
             (Keyword only argument.)
         """
@@ -103,7 +103,7 @@ def __init__(
         v0 = jr.normal(vkey, (v_len,))
         for _ in range(15):
             u0, v0 = _power_iteration(weight, u0, v0, eps)
-        self.uv_index = StateIndex(lambda **_: (u0, v0))
+        self.uv_index = StateIndex((u0, v0))
 
     @jax.named_scope("eqx.nn.SpectralNorm")
     def __call__(
diff --git a/equinox/nn/_stateful.py b/equinox/nn/_stateful.py
index b1e5a5aa..fc9b9ab0 100644
--- a/equinox/nn/_stateful.py
+++ b/equinox/nn/_stateful.py
@@ -1,39 +1,67 @@
 import types
 from collections.abc import Callable
-from typing import Any, Generic, TypeVar
+from typing import Any, Generic, TYPE_CHECKING, TypeVar, Union
+from typing_extensions import ParamSpec
 
 import jax
 import jax.numpy as jnp
 import jax.tree_util as jtu
 from jaxtyping import PyTree
 
-from .._module import Module
+from .._module import field, Module
 from .._pretty_print import bracketed, named_objs, text, tree_pformat
+from .._tree import tree_at
 
 
 _Value = TypeVar("_Value")
+_P = ParamSpec("_P")
+_T = TypeVar("_T")
 
 
 class StateIndex(Module, Generic[_Value]):
-    """This is an advanced feature, used when creating custom stateful layers.
+    """This wraps together (a) a unique dictionary key used for looking up a stateful
+    value, and (b) how that stateful value should be initialised.
 
-    This wraps a dictionary key used for looking up the stateful value.
+    !!! Example
 
-    See the source code of [`equinox.nn.BatchNorm`][] for an example.
-    """
+        ```python
+        class MyStatefulLayer(eqx.Module):
+            index: eqx.nn.StateIndex
+
+            def __init__(self):
+                init_state = jnp.array(0)
+                self.index = StateIndex(init_state)
+
+            def __call__(self, x: Array, state: eqx.nn.State) -> tuple[Array, eqx.nn.State]:
+                current_state = state.get(self.index)
+                new_x = x + current_state
+                new_state = state.set(current_state + 1)
+                return new_x, new_state
+        ```
 
-    marker: object
-    init: types.FunctionType
+    See also e.g. the source code of built-in stateful layers like
+    [`equinox.nn.BatchNorm`][] for further reference.
+    """  # noqa: E501
 
-    def __init__(self, init: Callable[..., _Value]):
+    # Starts off as an `object` when initialised; later replaced with an `int` inside
+    # `make_with_state`.
+    marker: Union[object, int] = field(static=True)
+    init: _Value
+
+    def __init__(self, init: _Value):
         """**Arguments:**
 
-        - `init`: A function that is used to initialise the state of the layer. Should
-            be a function that returns a PyTree of JAX arrays. It will be called with
-            all keyword arguments passed to [`equinox.nn.State.__init__`].
+        - `init`: The initial value for the state.
         """
-        if not isinstance(init, types.FunctionType):
-            raise TypeError("`StateIndex(init=...)` must be a function")
+        if isinstance(init, types.FunctionType):
+            # Technically a function is valid here, since we could allow any pytree.
+            # In practice that's weird / kind of useless, so better to explicitly raise
+            # the deprecation error.
+            raise ValueError(
+                "As of Equinox v0.11.0, `eqx.nn.StateIndex` now accepts the value "
+                "of the initial state directly. (Not a function that creates the "
+                "initial state.)"
+            )
         self.marker = object()
         self.init = init
 
@@ -42,37 +70,46 @@ def _is_index(x: Any) -> bool:
     return isinstance(x, StateIndex)
 
 
+# Used as a sentinel in two ways: keeping track of updated `State`s, and keeping track
+# of deleted initial states.
 _sentinel = object()
 
 
 _state_error = """
 Attempted to use old state. Probably you have done something like:
-
+```
 x, state2 = layer1(x, state1)
 x, state3 = layer1(x, state1)  # bug! Re-used state1 instead of using state2.
+```
+
+If you have done this intentionally, because you want to use an old state, then you can
+avoid this error by making a clone of the state:
+```
+leaves, treedef = jax.tree_util.tree_flatten(state)
+state_clone = jax.tree_util.tree_unflatten(treedef, leaves)
+```
 """.strip()
 
 
-# Basically just a dictionary which (a) works only with Markers, and which (b) works
-# around a JAX bug that prevents flattening dicts with `object()` keys, and which (c)
-# does error-checking that you're using the most up-to-date version of it.
+# Basically just a dictionary which (a) works only with StateIndex-s, and which (b)
+# works around a JAX bug that prevents flattening dicts with `object()` keys, and which
+# (c) does error-checking that you're using the most up-to-date version of it.
 @jtu.register_pytree_node_class
 class State:
     """Stores the state of a model. For example, the running statistics of all
     [`equinox.nn.BatchNorm`][] layers in the model.
 
-    Most models won't need this. (As most models don't have any stateful layers.)
-    If used, the state will be passed to each layer at call time; see the
-    [stateful example](../../examples/stateful.ipynb).
+    This is essentially a dictionary mapping from [`equinox.nn.StateIndex`][]s to
+    PyTrees of arrays.
+
+    This class should be initialised via [`equinox.nn.make_with_state`][].
     """
 
-    def __init__(self, model: PyTree, **kwargs):
+    def __init__(self, model: PyTree):
         """**Arguments:**
 
         - `model`: any PyTree. All stateful layers (e.g. [`equinox.nn.BatchNorm`][])
-            will have their state initialised and stored inside the `State` object.
-        - `**kwargs`: all keyword arguments are forwarded to the `init` function of
-            `equinox.nn.StateIndex(init=...)`  (used inside each stateful layer).
+            will have their initial state stored inside the `State` object.
         """
         # Note that de/serialisation depends on the ordered-ness of this dictionary,
         # between serialisation and deserialisation.
@@ -80,30 +117,115 @@ def __init__(self, model: PyTree, **kwargs):
         leaves = jtu.tree_leaves(model, is_leaf=_is_index)
         for leaf in leaves:
             if _is_index(leaf):
-                value = leaf.init(**kwargs)
-                value = jtu.tree_map(jnp.asarray, value)
-                state[leaf.marker] = value
+                if leaf.init is _sentinel:
+                    raise ValueError(
+                        "Cannot call `eqx.nn.State(eqx.nn.delete_init_state(model))`. "
+                        "You should call `eqx.nn.State(model)`, using the original "
+                        "model."
+                    )
+                state[leaf.marker] = jtu.tree_map(jnp.asarray, leaf.init)
         self._state = state
 
     def get(self, item: StateIndex[_Value]) -> _Value:
+        """Given an [`equinox.nn.StateIndex`][], returns the value of its state.
+
+        **Arguments:**
+
+        - `item`: an [`equinox.nn.StateIndex`][].
+
+        **Returns:**
+
+        The current state associated with that index.
+        """
         if self._state is _sentinel:
             raise ValueError(_state_error)
         if type(item) is not StateIndex:
-            raise ValueError("Can only use `eqx.nn.Marker`s as state keys.")
+            raise ValueError("Can only use `eqx.nn.StateIndex`s as state keys.")
         return self._state[item.marker]  # pyright: ignore
 
     def set(self, item: StateIndex[_Value], value: _Value) -> "State":
+        """Sets a new value for an [`equinox.nn.StateIndex`][], **and returns the
+        updated state**.
+
+        **Arguments:**
+
+        - `item`: an [`equinox.nn.StateIndex`][].
+        - `value`: the new value associated with that index.
+
+        **Returns:**
+
+        A new [`equinox.nn.State`][] object, with the update.
+
+        As a safety guard against accidentally writing `state.set(item, value)` without
+        assigning it to a new value, then the old object (`self`) will become invalid.
+        """
         if self._state is _sentinel:
             raise ValueError(_state_error)
         if type(item) is not StateIndex:
-            raise ValueError("Can only use `eqx.nn.Marker`s as state keys.")
+            raise ValueError("Can only use `eqx.nn.StateIndex`s as state keys.")
         old_value = self._state[item.marker]  # pyright: ignore
         value = jtu.tree_map(jnp.asarray, value)
         if jax.eval_shape(lambda: old_value) != jax.eval_shape(lambda: value):
             raise ValueError("Old and new values have different structures.")
-        self._state[item.marker] = value  # pyright: ignore
+        state = self._state.copy()  # pyright: ignore
+        state[item.marker] = value
         new_self = object.__new__(State)
-        new_self._state = self._state
+        new_self._state = state
+        self._state = _sentinel
+        return new_self
+
+    def substate(self, pytree: PyTree) -> "State":
+        """Creates a smaller `State` object, that tracks only the states of some smaller
+        part of the overall model.
+
+        **Arguments:**
+
+        - `pytree`: any PyTree. It will be iterated over to check for
+            [`equinox.nn.StateIndex`]s.
+
+        **Returns:**
+
+        A new [`equinox.nn.State`][] object, which tracks only some of the overall
+        states.
+        """
+        if self._state is _sentinel:
+            raise ValueError(_state_error)
+        leaves = jtu.tree_leaves(pytree, is_leaf=_is_index)
+        markers = [x.marker for x in leaves if _is_index(x)]
+        substate = {k: self._state[k] for k in markers}  # pyright: ignore
+        subself = object.__new__(State)
+        subself._state = substate
+        return subself
+
+    def update(self, substate: "State") -> "State":
+        """Takes a smaller `State` object (typically produces via `.substate`), and
+        updates states by using all of its values.
+
+        **Arguments:**
+
+        - `substate`: a `State` object whose keys are a subset of the keys of `self`.
+
+        **Returns:**
+
+        A new [`equinox.nn.State`][] object, containing all of the updated values.
+
+        As a safety guard against accidentally writing `state.set(item, value)` without
+        assigning it to a new value, then the old object (`self`) will become invalid.
+        """
+        if self._state is _sentinel:
+            raise ValueError(_state_error)
+        if type(substate) is not State:
+            raise ValueError("Can only use `eqx.nn.State`s in `update`.")
+        state = self._state.copy()  # pyright: ignore
+        for key, value in substate._state.items():  # pyright: ignore
+            if key not in state:
+                raise ValueError(
+                    "Cannot call `state1.update(state2)` unless `state2` is a substate "
+                    "of `state1`."
+                )
+            state[key] = value
+        new_self = object.__new__(State)
+        new_self._state = state
         self._state = _sentinel
         return new_self
 
@@ -145,3 +267,113 @@ def tree_unflatten(cls, keys, values):
             state[key] = value
         self._state = state
         return self
+
+
+def _delete_init_state(x):
+    if _is_index(x):
+        return tree_at(lambda y: y.init, x, _sentinel)
+    else:
+        return x
+
+
+def delete_init_state(model: PyTree) -> PyTree:
+    """For memory efficiency, this deletes the initial state stored within a model.
+
+    Every stateful layer keeps a copy of the initial value of its state. This is then
+    collected by [`equinox.nn.State`][], when it is called on the model. However, this
+    means that the model must keep a copy of the initial state around, in case
+    `eqx.nn.State` is called on it again. This extra copy consumes extra memory.
+
+    But in practice, it is quite common to only need to initialise the state once. In
+    this case, we can use this function to delete this extra copy, and in doing so save
+    some memory.
+
+    !!! Example
+
+        Here is the typical pattern in which this is used:
+        ```python
+        model_and_state = eqx.nn.BatchNorm(...)
+        state = eqx.nn.State(model_and_state)
+        model = eqx.nn.delete_init_state(model)
+        del model_and_state  # ensure this goes out of scope and is garbage collected
+        ```
+        Indeed the above is precisely what [`equinox.nn.make_with_state`][] does.
+
+    **Arguments:**
+
+    - `model`: any PyTree.
+
+    **Returns:**
+
+    A copy of `model`, with all the initial states stripped out. (As in the exampels
+    above, you should then dispose of the original `model` object.)
+    """
+    return jtu.tree_map(_delete_init_state, model, is_leaf=_is_index)
+
+
+def make_with_state(make_model: Callable[_P, _T]) -> Callable[_P, tuple[_T, State]]:
+    """This function is the most common API for working with stateful models. This
+    initialises both the parameters and the state of a stateful model.
+
+    `eqx.nn.make_with_state(Model)(*args, **kwargs)` simply calls
+    `model_with_state = Model(*args, **kwargs)`, and then partitions the resulting
+    PyTree into two pieces: the parameters, and the state.
+
+    **Arguments:**
+
+    - `make_model`: some callable returning a PyTree.
+
+    **Returns:**
+
+    A callable, which when evaluated returns a 2-tuple of `(model, state)`, where
+    `model` is the result of `make_model(*args, **kwargs)` but with all of the initial
+    states stripped out, and `state` is an [`equinox.nn.State`][] object encapsulating
+    the initial states.
+
+    !!! Example
+
+        See [the stateful example](../../examples/stateful.ipynb) for a runnable
+        example.
+
+        ```python
+        class Model(eqx.Module):
+            def __init__(self, foo, bar):
+                ...
+
+            ...
+
+        model, state = eqx.nn.make_with_state(Model)(foo=3, bar=4)
+        ```
+    """
+
+    # _P.{args, kwargs} not supported by beartype
+    if TYPE_CHECKING:
+
+        def make_with_state_impl(
+            *args: _P.args, **kwargs: _P.kwargs
+        ) -> tuple[_T, State]:
+            ...
+
+    else:
+
+        def make_with_state_impl(*args, **kwargs) -> tuple[_T, State]:
+            model = make_model(*args, **kwargs)
+
+            # Replace all markers with `int`s. This is needed to ensure that two calls
+            # to `make_with_state` produce compatible models and states.
+            leaves, treedef = jtu.tree_flatten(model, is_leaf=_is_index)
+            counter = 0
+            new_leaves = []
+            for leaf in leaves:
+                if _is_index(leaf):
+                    leaf = StateIndex(leaf.init)
+                    object.__setattr__(leaf, "marker", counter)
+                    counter += 1
+                new_leaves.append(leaf)
+            model = jtu.tree_unflatten(treedef, new_leaves)
+
+            state = State(model)
+            model = delete_init_state(model)
+            return model, state
+
+    return make_with_state_impl
diff --git a/examples/deep_convolutional_gan.ipynb b/examples/deep_convolutional_gan.ipynb
index 9f9b4f9b..c56aa286 100644
--- a/examples/deep_convolutional_gan.ipynb
+++ b/examples/deep_convolutional_gan.ipynb
@@ -640,12 +640,12 @@
     "    return out\n",
     "\n",
     "\n",
-    "inference_gen = eqx.tree_inference(generator, value=True)\n",
+    "inference_gen = eqx.nn.inference_mode(generator)\n",
     "inference_gen = eqx.Partial(inference_gen, state=generator_state)\n",
     "\n",
     "generated_images = evaluate(inference_gen, noise)\n",
     "\n",
-    "inference_discriminator = eqx.tree_inference(discriminator, value=True)\n",
+    "inference_discriminator = eqx.nn.inference_mode(discriminator)\n",
     "inference_discriminator = eqx.Partial(\n",
     "    inference_discriminator, state=discriminator_state\n",
     ")\n",
diff --git a/examples/stateful.ipynb b/examples/stateful.ipynb
index 38044b8c..ef872523 100644
--- a/examples/stateful.ipynb
+++ b/examples/stateful.ipynb
@@ -11,6 +11,8 @@
     "\n",
     "This just means that we need to plumb an extra input and output through our models. This example demonstrates both [`equinox.nn.BatchNorm`][] and [`equinox.nn.SpectralNorm`][].\n",
     "\n",
+    "See also the [stateful API reference](../api/nn/stateful.md).\n",
+    "\n",
     "This example is available as a Jupyter notebook [here](https://github.com/patrick-kidger/equinox/blob/main/examples/stateful.ipynb)."
    ]
   },
@@ -37,8 +39,8 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "# This model is just a weird mish-mash of layers for demonstration purposes, it isn't\n",
-    "# doing any clever.\n",
+    "# This model is just a weird mish-mash of stateful and non-stateful layers for\n",
+    "# demonstration purposes, it isn't doing any clever.\n",
     "class Model(eqx.Module):\n",
     "    norm1: eqx.nn.BatchNorm\n",
     "    spectral_linear: eqx.nn.SpectralNorm[eqx.nn.Linear]\n",
@@ -69,6 +71,24 @@
     "        return x, state"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "id": "2b931647-01ff-4eb9-972f-55d3d50771c6",
+   "metadata": {},
+   "source": [
+    "We see from the above that we just define our models like normal. As advertised, we just need to thread the additional `state` object in and out of every call. An updated state object is returned.\n",
+    "\n",
+    "There's really nothing special here about stateful layers. Equinox isn't special-casing them in any way. We thread `state` in and out, just like we're thread `x` in and out. In fact calling it \"state\" is really just a matter of how it's advertised!\n",
+    "\n",
+    "---\n",
+    "\n",
+    "Alright, now let's see how we might train this model. This is also much like normal.\n",
+    "\n",
+    "Note the use of `in_axes` and `out_axes`: our data is batched, but our model state isn't batched -- just like how our model parameters isn't batched.\n",
+    "\n",
+    "Note how the `axis_name` argment matches the `axis_name` argument that the `BatchNorm` layers were initialised with. This tells `BatchNorm` which vmap'd axis it should compute statistics over. (This is a detail specific to `BatchNorm`, and is unrelated to stateful operations in general.)"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": 3,
@@ -77,9 +97,6 @@
    "outputs": [],
    "source": [
     "def compute_loss(model, state, xs, ys):\n",
-    "    # The `axis_name` argument is needed specifically for `BatchNorm`: so it knows\n",
-    "    #     what axis to compute batch statistics over.\n",
-    "    # The `in_axes` and `out_axes` are needed so that `state` isn't batched.\n",
     "    batch_model = jax.vmap(\n",
     "        model, axis_name=\"batch\", in_axes=(0, None), out_axes=(0, None)\n",
     "    )\n",
@@ -96,6 +113,16 @@
     "    return model, state, opt_state"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "id": "d85113ae-7d63-4051-83bd-e9e921dfcae1",
+   "metadata": {},
+   "source": [
+    "---\n",
+    "\n",
+    "And now, let's see how we initialise this model, and initialise its state."
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": 4,
@@ -112,15 +139,16 @@
     "\n",
     "key = jr.PRNGKey(seed)\n",
     "mkey, xkey, xkey2 = jr.split(key, 3)\n",
-    "model = Model(mkey)\n",
-    "state = eqx.nn.State(model)\n",
+    "\n",
+    "model, state = eqx.nn.make_with_state(Model)(mkey)\n",
+    "\n",
     "xs = jr.normal(xkey, (dataset_size, 3))\n",
     "ys = jnp.sin(xs) + 1\n",
     "optim = optax.adam(learning_rate)\n",
     "opt_state = optim.init(eqx.filter(model, eqx.is_inexact_array))\n",
     "\n",
-    "# Full-batch gradient descent in this simple example.\n",
     "for _ in range(steps):\n",
+    "    # Full-batch gradient descent in this simple example.\n",
     "    model, state, opt_state = make_step(model, state, opt_state, xs, ys)"
    ]
   },
@@ -129,15 +157,19 @@
    "id": "9c2135ba-07dc-4ce7-89c1-105bd699e78d",
    "metadata": {},
    "source": [
-    "Overall, we see that this should be relatively straightforward!\n",
+    "!!! Info \"What is `eqx.nn.make_with_state` doing?\"\n",
     "\n",
-    "When calling `state = eqx.nn.State(model)`, then the model PyTree is iterated over, and any stateful layers store their initial states in the resulting `state` object. The `state` object is itself also a PyTree, so it can just be passed around in the usual way.\n",
+    "    Here we come to the only interesting bit about using stateful layers!\n",
     "\n",
-    "In this example, `state` will store the running statistics for `BatchNorm`, and U-V power iterations for `SpectralNorm`.\n",
+    "    When we initialise the model -- e.g. if we were to call `Model(mkey)` directly -- then the model PyTree would be initialised containing both (a) the initial parameters, and (b) the initial state. So `make_with_state` simply calls this, and then separates these two things. The returned `model` is a PyTree holding all the initial parameters (just like any other model), and `state` is a PyTree holding the initial state.\n",
     "\n",
-    "Subsequently, we just need to thread the `state` object in-and-out of every call. Each time a new state object is returned. (And the old state object should not be reused.)\n",
+    "---\n",
     "\n",
-    "Finally, let's use our trained model to perform inference:"
+    "Finally, let's use our trained model to perform inference.\n",
+    "\n",
+    "Remember to set the inference flag! Some layers have different behaviour between training and inference, and `BatchNorm` is one of these. (This is a detail specific to layers like `BatchNorm` and [`equinox.nn.Dropout`][], and is unrelated to stateful operations in general.)\n",
+    "\n",
+    "We also fix the final state in the model, using [`equinox.Partial`][]. The resulting `inference_model` is a PyTree (specifically, an `equinox.Partial`) containing both `model` and `state`."
    ]
   },
   {
@@ -147,7 +179,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "inference_model = eqx.tree_inference(model, value=True)\n",
+    "inference_model = eqx.nn.inference_mode(model)\n",
     "inference_model = eqx.Partial(inference_model, state=state)\n",
     "\n",
     "\n",
@@ -168,17 +200,15 @@
    "id": "d22e4395-9006-465e-ba0b-a4b0d8131bd7",
    "metadata": {},
    "source": [
-    "Here, we don't need the updated state object that is output, so we just discard it.\n",
-    "\n",
-    "(Also, don't forget to set the `inference` flags.)"
+    "Here, we don't need the updated state object that is produced, so we just discard it."
    ]
   }
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "jax38",
+   "display_name": "py311",
    "language": "python",
-   "name": "jax38"
+   "name": "py311"
   },
   "language_info": {
    "codemirror_mode": {
@@ -190,7 +220,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.8.16"
+   "version": "3.11.3"
   }
  },
  "nbformat": 4,
diff --git a/examples/unet.ipynb b/examples/unet.ipynb
index 609a1b05..2a0b8932 100644
--- a/examples/unet.ipynb
+++ b/examples/unet.ipynb
@@ -9,7 +9,7 @@
     "\n",
     "This is an advanced example, providing an implementation of a U-Net architecture.\n",
     "\n",
-    "This version is intended for use in a [score-based diffusion](./score_based_diffusion/), so it accepts a `t` argument.\n",
+    "This version is intended for use in a [score-based diffusion](../score_based_diffusion/), so it accepts a `t` argument.\n",
     "\n",
     "This example is available as a Jupyter notebook [here](https://github.com/patrick-kidger/equinox/blob/main/examples/unet.ipynb)\n",
     "\n",
diff --git a/examples/vision_transformer.ipynb b/examples/vision_transformer.ipynb
new file mode 100644
index 00000000..e61f2981
--- /dev/null
+++ b/examples/vision_transformer.ipynb
@@ -0,0 +1,662 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "x96qLutZxINh"
+   },
+   "source": [
+    "# Vision Transformer (ViT)\n",
+    "\n",
+    "This example builds a vision transformer model using Equinox, an implementation based on the paper: _An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale_.\n",
+    "\n",
+    "In addition to this tutorial example, you may also like the ViT implementation [available here in Eqxvision here](https://eqxvision.readthedocs.io/en/latest/api/models/classification/vit/).\n",
+    "\n",
+    "!!! warning\n",
+    "\n",
+    "    This example will take a short while to run on a GPU.\n",
+    "\n",
+    "!!! cite \"Reference\"\n",
+    "\n",
+    "    [arXiv link](https://arxiv.org/abs/2010.11929)\n",
+    "\n",
+    "    ```bibtex\n",
+    "    @inproceedings{dosovitskiy2021an,\n",
+    "        title={An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale},\n",
+    "        author={Alexey Dosovitskiy and Lucas Beyer and Alexander Kolesnikov and Dirk Weissenborn\n",
+    "                and Xiaohua Zhai and Thomas Unterthiner and Mostafa Dehghani and Matthias Minderer\n",
+    "                and Georg Heigold and Sylvain Gelly and Jakob Uszkoreit and Neil Houlsby},\n",
+    "        booktitle={International Conference on Learning Representations},\n",
+    "        year={2021},\n",
+    "    }\n",
+    "    ```"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {
+    "id": "3-NddIhhxINj"
+   },
+   "outputs": [],
+   "source": [
+    "import functools\n",
+    "\n",
+    "import einops  # https://github.com/arogozhnikov/einops\n",
+    "import jax\n",
+    "import jax.numpy as jnp\n",
+    "import jax.random as jr\n",
+    "import numpy as np\n",
+    "import optax  # https://github.com/deepmind/optax\n",
+    "from jaxtyping import PRNGKeyArray, Array, Float\n",
+    "\n",
+    "# We'll use PyTorch to load the dataset.\n",
+    "import torch\n",
+    "import torchvision\n",
+    "import torchvision.transforms as transforms\n",
+    "\n",
+    "import equinox as eqx"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {
+    "id": "bYi-XlXRxINl"
+   },
+   "outputs": [],
+   "source": [
+    "# Hyperparameters\n",
+    "lr = 0.0001\n",
+    "dropout_rate = 0.1\n",
+    "beta1 = 0.9\n",
+    "beta2 = 0.999\n",
+    "batch_size = 64\n",
+    "patch_size = 4\n",
+    "num_patches = 64\n",
+    "num_steps = 100000\n",
+    "image_size = (32, 32, 3)\n",
+    "embedding_dim = 512\n",
+    "hidden_dim = 256\n",
+    "num_heads = 8\n",
+    "num_layers = 6\n",
+    "height, width, channels = image_size\n",
+    "num_classes = 10"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "ClIEhf1dBa8x"
+   },
+   "source": [
+    "Let's first load the CIFAR10 dataset using torchvision"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {
+    "colab": {
+     "base_uri": "https://localhost:8080/"
+    },
+    "id": "Vcwi4un6CMu_",
+    "outputId": "fad94424-789b-46b2-f3df-41f004ddfaf7"
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Files already downloaded and verified\n",
+      "Files already downloaded and verified\n"
+     ]
+    }
+   ],
+   "source": [
+    "transform_train = transforms.Compose(\n",
+    "    [\n",
+    "        transforms.RandomCrop(32, padding=4),\n",
+    "        transforms.Resize((height, width)),\n",
+    "        transforms.RandomHorizontalFlip(),\n",
+    "        transforms.ToTensor(),\n",
+    "        transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),\n",
+    "    ]\n",
+    ")\n",
+    "\n",
+    "transform_test = transforms.Compose(\n",
+    "    [\n",
+    "        transforms.Resize((height, width)),\n",
+    "        transforms.ToTensor(),\n",
+    "        transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),\n",
+    "    ]\n",
+    ")\n",
+    "\n",
+    "train_dataset = torchvision.datasets.CIFAR10(\n",
+    "    \"CIFAR\",\n",
+    "    train=True,\n",
+    "    download=True,\n",
+    "    transform=transform_train,\n",
+    ")\n",
+    "\n",
+    "test_dataset = torchvision.datasets.CIFAR10(\n",
+    "    \"CIFAR\",\n",
+    "    train=False,\n",
+    "    download=True,\n",
+    "    transform=transform_test,\n",
+    ")\n",
+    "\n",
+    "trainloader = torch.utils.data.DataLoader(\n",
+    "    train_dataset, batch_size=batch_size, shuffle=True, drop_last=True\n",
+    ")\n",
+    "\n",
+    "testloader = torch.utils.data.DataLoader(\n",
+    "    test_dataset, batch_size=batch_size, shuffle=True, drop_last=True\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "h-Q4A5H8OQLs"
+   },
+   "source": [
+    "Now Let's start by making the patch embeddings layer that will turn images into embedded patches to be processed then by the attention layers."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "metadata": {
+    "id": "SFo1GzZvxINl"
+   },
+   "outputs": [],
+   "source": [
+    "class PatchEmbedding(eqx.Module):\n",
+    "    linear: eqx.nn.Embedding\n",
+    "    patch_size: int\n",
+    "\n",
+    "    def __init__(\n",
+    "        self,\n",
+    "        input_channels: int,\n",
+    "        output_shape: int,\n",
+    "        patch_size: int,\n",
+    "        key: PRNGKeyArray,\n",
+    "    ):\n",
+    "        self.patch_size = patch_size\n",
+    "\n",
+    "        self.linear = eqx.nn.Linear(\n",
+    "            self.patch_size**2 * input_channels,\n",
+    "            output_shape,\n",
+    "            key=key,\n",
+    "        )\n",
+    "\n",
+    "    def __call__(\n",
+    "        self, x: Float[Array, \"channels height width\"]\n",
+    "    ) -> Float[Array, \"num_patches embedding_dim\"]:\n",
+    "        x = einops.rearrange(\n",
+    "            x,\n",
+    "            \"c (h ph) (w pw) -> (h w) (c ph pw)\",\n",
+    "            ph=self.patch_size,\n",
+    "            pw=self.patch_size,\n",
+    "        )\n",
+    "        x = jax.vmap(self.linear)(x)\n",
+    "\n",
+    "        return x"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "dOJp3mZNOjJW"
+   },
+   "source": [
+    "After that, we implement the attention block which is the core of the transformer architecture."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "metadata": {
+    "id": "mDg-L_9ixINm"
+   },
+   "outputs": [],
+   "source": [
+    "class AttentionBlock(eqx.Module):\n",
+    "    layer_norm1: eqx.nn.LayerNorm\n",
+    "    layer_norm2: eqx.nn.LayerNorm\n",
+    "    attention: eqx.nn.MultiheadAttention\n",
+    "    linear1: eqx.nn.Linear\n",
+    "    linear2: eqx.nn.Linear\n",
+    "    dropout1: eqx.nn.Dropout\n",
+    "    dropout2: eqx.nn.Dropout\n",
+    "\n",
+    "    def __init__(\n",
+    "        self,\n",
+    "        input_shape: int,\n",
+    "        hidden_dim: int,\n",
+    "        num_heads: int,\n",
+    "        dropout_rate: float,\n",
+    "        key: PRNGKeyArray,\n",
+    "    ):\n",
+    "        key1, key2, key3 = jr.split(key, 3)\n",
+    "\n",
+    "        self.layer_norm1 = eqx.nn.LayerNorm(input_shape)\n",
+    "        self.layer_norm2 = eqx.nn.LayerNorm(input_shape)\n",
+    "        self.attention = eqx.nn.MultiheadAttention(num_heads, input_shape, key=key1)\n",
+    "\n",
+    "        self.linear1 = eqx.nn.Linear(input_shape, hidden_dim, key=key2)\n",
+    "        self.linear2 = eqx.nn.Linear(hidden_dim, input_shape, key=key3)\n",
+    "        self.dropout1 = eqx.nn.Dropout(dropout_rate)\n",
+    "        self.dropout2 = eqx.nn.Dropout(dropout_rate)\n",
+    "\n",
+    "    def __call__(\n",
+    "        self,\n",
+    "        x: Float[Array, \"num_patches embedding_dim\"],\n",
+    "        enable_dropout: bool,\n",
+    "        key: PRNGKeyArray,\n",
+    "    ) -> Float[Array, \"num_patches embedding_dim\"]:\n",
+    "        input_x = jax.vmap(self.layer_norm1)(x)\n",
+    "        x = x + self.attention(input_x, input_x, input_x)\n",
+    "\n",
+    "        input_x = jax.vmap(self.layer_norm2)(x)\n",
+    "        input_x = jax.vmap(self.linear1)(input_x)\n",
+    "        input_x = jax.nn.gelu(input_x)\n",
+    "\n",
+    "        key1, key2 = jr.split(key, num=2)\n",
+    "\n",
+    "        input_x = self.dropout1(input_x, inference=not enable_dropout, key=key1)\n",
+    "        input_x = jax.vmap(self.linear2)(input_x)\n",
+    "        input_x = self.dropout2(input_x, inference=not enable_dropout, key=key2)\n",
+    "\n",
+    "        x = x + input_x\n",
+    "\n",
+    "        return x"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "t_RB1ip0PEk4"
+   },
+   "source": [
+    "Lastly, we build the full Vision Transformer model, which is composed of embeddings layers, a series of transformer blocks, and a classification head."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "metadata": {
+    "id": "nG6fLPhyQEBx"
+   },
+   "outputs": [],
+   "source": [
+    "class VisionTransformer(eqx.Module):\n",
+    "    patch_embedding: PatchEmbedding\n",
+    "    positional_embedding: jnp.ndarray\n",
+    "    cls_token: jnp.ndarray\n",
+    "    attention_blocks: list[AttentionBlock]\n",
+    "    dropout: eqx.nn.Dropout\n",
+    "    mlp: eqx.nn.Sequential\n",
+    "    num_layers: int\n",
+    "\n",
+    "    def __init__(\n",
+    "        self,\n",
+    "        embedding_dim: int,\n",
+    "        hidden_dim: int,\n",
+    "        num_heads: int,\n",
+    "        num_layers: int,\n",
+    "        dropout_rate: float,\n",
+    "        patch_size: int,\n",
+    "        num_patches: int,\n",
+    "        num_classes: int,\n",
+    "        key: PRNGKeyArray,\n",
+    "    ):\n",
+    "        key1, key2, key3, key4, key5 = jr.split(key, 5)\n",
+    "\n",
+    "        self.patch_embedding = PatchEmbedding(channels, embedding_dim, patch_size, key1)\n",
+    "\n",
+    "        self.positional_embedding = jr.normal(key2, (num_patches + 1, embedding_dim))\n",
+    "\n",
+    "        self.cls_token = jr.normal(key3, (1, embedding_dim))\n",
+    "\n",
+    "        self.num_layers = num_layers\n",
+    "\n",
+    "        self.attention_blocks = [\n",
+    "            AttentionBlock(embedding_dim, hidden_dim, num_heads, dropout_rate, key4)\n",
+    "            for _ in range(self.num_layers)\n",
+    "        ]\n",
+    "\n",
+    "        self.dropout = eqx.nn.Dropout(dropout_rate)\n",
+    "\n",
+    "        self.mlp = eqx.nn.Sequential(\n",
+    "            [\n",
+    "                eqx.nn.LayerNorm(embedding_dim),\n",
+    "                eqx.nn.Linear(embedding_dim, num_classes, key=key5),\n",
+    "            ]\n",
+    "        )\n",
+    "\n",
+    "    def __call__(\n",
+    "        self,\n",
+    "        x: Float[Array, \"channels height width\"],\n",
+    "        enable_dropout: bool,\n",
+    "        key: PRNGKeyArray,\n",
+    "    ) -> Float[Array, \"num_classes\"]:\n",
+    "        x = self.patch_embedding(x)\n",
+    "\n",
+    "        x = jnp.concatenate((self.cls_token, x), axis=0)\n",
+    "\n",
+    "        x += self.positional_embedding[\n",
+    "            : x.shape[0]\n",
+    "        ]  # Slice to the same length as x, as the positional embedding may be longer.\n",
+    "\n",
+    "        dropout_key, *attention_keys = jr.split(key, num=self.num_layers + 1)\n",
+    "\n",
+    "        x = self.dropout(x, inference=not enable_dropout, key=dropout_key)\n",
+    "\n",
+    "        for block, attention_key in zip(self.attention_blocks, attention_keys):\n",
+    "            x = block(x, enable_dropout, key=attention_key)\n",
+    "\n",
+    "        x = x[0]  # Select the CLS token.\n",
+    "        x = self.mlp(x)\n",
+    "\n",
+    "        return x"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "metadata": {
+    "id": "agBSRsXVxINn"
+   },
+   "outputs": [],
+   "source": [
+    "@eqx.filter_value_and_grad\n",
+    "def compute_grads(\n",
+    "    model: VisionTransformer, images: jnp.ndarray, labels: jnp.ndarray, key\n",
+    "):\n",
+    "    logits = jax.vmap(model, in_axes=(0, None, 0))(images, True, key)\n",
+    "    loss = optax.softmax_cross_entropy_with_integer_labels(logits, labels)\n",
+    "\n",
+    "    return jnp.mean(loss)\n",
+    "\n",
+    "\n",
+    "@eqx.filter_jit\n",
+    "def step_model(\n",
+    "    model: VisionTransformer,\n",
+    "    optimizer: optax.GradientTransformation,\n",
+    "    state: optax.OptState,\n",
+    "    images: jnp.ndarray,\n",
+    "    labels: jnp.ndarray,\n",
+    "    key,\n",
+    "):\n",
+    "    loss, grads = compute_grads(model, images, labels, key)\n",
+    "    updates, new_state = optimizer.update(grads, state, model)\n",
+    "\n",
+    "    model = eqx.apply_updates(model, updates)\n",
+    "\n",
+    "    return model, new_state, loss\n",
+    "\n",
+    "\n",
+    "def train(\n",
+    "    model: VisionTransformer,\n",
+    "    optimizer: optax.GradientTransformation,\n",
+    "    state: optax.OptState,\n",
+    "    data_loader: torch.utils.data.DataLoader,\n",
+    "    num_steps: int,\n",
+    "    print_every: int = 1000,\n",
+    "    key=None,\n",
+    "):\n",
+    "    losses = []\n",
+    "\n",
+    "    def infinite_trainloader():\n",
+    "        while True:\n",
+    "            yield from data_loader\n",
+    "\n",
+    "    for step, batch in zip(range(num_steps), infinite_trainloader()):\n",
+    "        images, labels = batch\n",
+    "\n",
+    "        images = images.numpy()\n",
+    "        labels = labels.numpy()\n",
+    "\n",
+    "        key, *subkeys = jr.split(key, num=batch_size + 1)\n",
+    "        subkeys = jnp.array(subkeys)\n",
+    "\n",
+    "        (model, state, loss) = step_model(\n",
+    "            model, optimizer, state, images, labels, subkeys\n",
+    "        )\n",
+    "\n",
+    "        losses.append(loss)\n",
+    "\n",
+    "        if (step % print_every) == 0 or step == num_steps - 1:\n",
+    "            print(f\"Step: {step}/{num_steps}, Loss: {loss}.\")\n",
+    "\n",
+    "    return model, state, losses"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "metadata": {
+    "id": "y3Bm_Xln-rSp"
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Step: 0/100000, Loss: 2.5608019828796387.\n",
+      "Step: 1000/100000, Loss: 1.711548089981079.\n",
+      "Step: 2000/100000, Loss: 1.4029508829116821.\n",
+      "Step: 3000/100000, Loss: 1.405516505241394.\n",
+      "Step: 4000/100000, Loss: 1.1661641597747803.\n",
+      "Step: 5000/100000, Loss: 1.1351711750030518.\n",
+      "Step: 6000/100000, Loss: 1.11599600315094.\n",
+      "Step: 7000/100000, Loss: 0.796968936920166.\n",
+      "Step: 8000/100000, Loss: 0.6870157718658447.\n",
+      "Step: 9000/100000, Loss: 1.0474591255187988.\n",
+      "Step: 10000/100000, Loss: 0.9413787722587585.\n",
+      "Step: 11000/100000, Loss: 0.8514565229415894.\n",
+      "Step: 12000/100000, Loss: 0.6746965646743774.\n",
+      "Step: 13000/100000, Loss: 0.7895829677581787.\n",
+      "Step: 14000/100000, Loss: 0.6844460964202881.\n",
+      "Step: 15000/100000, Loss: 0.6571178436279297.\n",
+      "Step: 16000/100000, Loss: 0.5611618757247925.\n",
+      "Step: 17000/100000, Loss: 0.610838770866394.\n",
+      "Step: 18000/100000, Loss: 0.7180566787719727.\n",
+      "Step: 19000/100000, Loss: 0.6528561115264893.\n",
+      "Step: 20000/100000, Loss: 0.5517654418945312.\n",
+      "Step: 21000/100000, Loss: 0.6301887035369873.\n",
+      "Step: 22000/100000, Loss: 0.5667067766189575.\n",
+      "Step: 23000/100000, Loss: 0.43517759442329407.\n",
+      "Step: 24000/100000, Loss: 0.5348870754241943.\n",
+      "Step: 25000/100000, Loss: 0.44732385873794556.\n",
+      "Step: 26000/100000, Loss: 0.49118855595588684.\n",
+      "Step: 27000/100000, Loss: 0.5242345929145813.\n",
+      "Step: 28000/100000, Loss: 0.44588926434516907.\n",
+      "Step: 29000/100000, Loss: 0.23619337379932404.\n",
+      "Step: 30000/100000, Loss: 0.4560542702674866.\n",
+      "Step: 31000/100000, Loss: 0.3148268163204193.\n",
+      "Step: 32000/100000, Loss: 0.4813237488269806.\n",
+      "Step: 33000/100000, Loss: 0.40532559156417847.\n",
+      "Step: 34000/100000, Loss: 0.2517223358154297.\n",
+      "Step: 35000/100000, Loss: 0.322698712348938.\n",
+      "Step: 36000/100000, Loss: 0.3052283525466919.\n",
+      "Step: 37000/100000, Loss: 0.37322986125946045.\n",
+      "Step: 38000/100000, Loss: 0.27499520778656006.\n",
+      "Step: 39000/100000, Loss: 0.2547920346260071.\n",
+      "Step: 40000/100000, Loss: 0.27322614192962646.\n",
+      "Step: 41000/100000, Loss: 0.6049947738647461.\n",
+      "Step: 42000/100000, Loss: 0.28800976276397705.\n",
+      "Step: 43000/100000, Loss: 0.2901820242404938.\n",
+      "Step: 44000/100000, Loss: 0.3800655007362366.\n",
+      "Step: 45000/100000, Loss: 0.15261484682559967.\n",
+      "Step: 46000/100000, Loss: 0.17970965802669525.\n",
+      "Step: 47000/100000, Loss: 0.23651015758514404.\n",
+      "Step: 48000/100000, Loss: 0.3813527822494507.\n",
+      "Step: 49000/100000, Loss: 0.35252541303634644.\n",
+      "Step: 50000/100000, Loss: 0.16249465942382812.\n",
+      "Step: 51000/100000, Loss: 0.10218428075313568.\n",
+      "Step: 52000/100000, Loss: 0.2192973792552948.\n",
+      "Step: 53000/100000, Loss: 0.1880446970462799.\n",
+      "Step: 54000/100000, Loss: 0.14270251989364624.\n",
+      "Step: 55000/100000, Loss: 0.1278090476989746.\n",
+      "Step: 56000/100000, Loss: 0.0856819674372673.\n",
+      "Step: 57000/100000, Loss: 0.16201086342334747.\n",
+      "Step: 58000/100000, Loss: 0.20575015246868134.\n",
+      "Step: 59000/100000, Loss: 0.20935538411140442.\n",
+      "Step: 60000/100000, Loss: 0.09025183320045471.\n",
+      "Step: 61000/100000, Loss: 0.21367806196212769.\n",
+      "Step: 62000/100000, Loss: 0.06895419955253601.\n",
+      "Step: 63000/100000, Loss: 0.14567255973815918.\n",
+      "Step: 64000/100000, Loss: 0.18438486754894257.\n",
+      "Step: 65000/100000, Loss: 0.11639232933521271.\n",
+      "Step: 66000/100000, Loss: 0.06631053984165192.\n",
+      "Step: 67000/100000, Loss: 0.11763929575681686.\n",
+      "Step: 68000/100000, Loss: 0.046494871377944946.\n",
+      "Step: 69000/100000, Loss: 0.14044761657714844.\n",
+      "Step: 70000/100000, Loss: 0.1277393102645874.\n",
+      "Step: 71000/100000, Loss: 0.154437854886055.\n",
+      "Step: 72000/100000, Loss: 0.15087449550628662.\n",
+      "Step: 73000/100000, Loss: 0.05043340474367142.\n",
+      "Step: 74000/100000, Loss: 0.3183276355266571.\n",
+      "Step: 75000/100000, Loss: 0.15685151517391205.\n",
+      "Step: 76000/100000, Loss: 0.13796621561050415.\n",
+      "Step: 77000/100000, Loss: 0.1036764532327652.\n",
+      "Step: 78000/100000, Loss: 0.08222786337137222.\n",
+      "Step: 79000/100000, Loss: 0.1525675356388092.\n",
+      "Step: 80000/100000, Loss: 0.06328584253787994.\n",
+      "Step: 81000/100000, Loss: 0.1235610619187355.\n",
+      "Step: 82000/100000, Loss: 0.03093503788113594.\n",
+      "Step: 83000/100000, Loss: 0.07480041682720184.\n",
+      "Step: 84000/100000, Loss: 0.016707731410861015.\n",
+      "Step: 85000/100000, Loss: 0.0491723008453846.\n",
+      "Step: 86000/100000, Loss: 0.0650872215628624.\n",
+      "Step: 87000/100000, Loss: 0.08738622069358826.\n",
+      "Step: 88000/100000, Loss: 0.10671466588973999.\n",
+      "Step: 89000/100000, Loss: 0.11922930181026459.\n",
+      "Step: 90000/100000, Loss: 0.1234014481306076.\n",
+      "Step: 91000/100000, Loss: 0.08588997274637222.\n",
+      "Step: 92000/100000, Loss: 0.036773063242435455.\n",
+      "Step: 93000/100000, Loss: 0.03425668179988861.\n",
+      "Step: 94000/100000, Loss: 0.21202465891838074.\n",
+      "Step: 95000/100000, Loss: 0.26020047068595886.\n",
+      "Step: 96000/100000, Loss: 0.154791459441185.\n",
+      "Step: 97000/100000, Loss: 0.1340092271566391.\n",
+      "Step: 98000/100000, Loss: 0.11398129910230637.\n",
+      "Step: 99000/100000, Loss: 0.16246598958969116.\n",
+      "Step: 99999/100000, Loss: 0.04668630287051201.\n"
+     ]
+    }
+   ],
+   "source": [
+    "key = jr.PRNGKey(2003)\n",
+    "\n",
+    "model = VisionTransformer(\n",
+    "    embedding_dim=embedding_dim,\n",
+    "    hidden_dim=hidden_dim,\n",
+    "    num_heads=num_heads,\n",
+    "    num_layers=num_layers,\n",
+    "    dropout_rate=dropout_rate,\n",
+    "    patch_size=patch_size,\n",
+    "    num_patches=num_patches,\n",
+    "    num_classes=num_classes,\n",
+    "    key=key,\n",
+    ")\n",
+    "\n",
+    "optimizer = optax.adamw(\n",
+    "    learning_rate=lr,\n",
+    "    b1=beta1,\n",
+    "    b2=beta2,\n",
+    ")\n",
+    "\n",
+    "state = optimizer.init(eqx.filter(model, eqx.is_inexact_array))\n",
+    "\n",
+    "model, state, losses = train(model, optimizer, state, trainloader, num_steps, key=key)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "X4GPbpuMQEB1"
+   },
+   "source": [
+    "And now let's see how the vision transformer performs on the CIFAR10 dataset."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 10,
+   "metadata": {
+    "colab": {
+     "base_uri": "https://localhost:8080/"
+    },
+    "id": "pZu5pMRZW3tF",
+    "outputId": "d74d06b7-0340-4e4f-b723-8e5d532aecb5",
+    "scrolled": true
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Accuracy: 79.13661858974359%\n"
+     ]
+    }
+   ],
+   "source": [
+    "accuracies = []\n",
+    "\n",
+    "for batch in range(len(test_dataset) // batch_size):\n",
+    "    images, labels = next(iter(testloader))\n",
+    "\n",
+    "    logits = jax.vmap(functools.partial(model, enable_dropout=False))(\n",
+    "        images.numpy(), key=jax.random.split(key, num=batch_size)\n",
+    "    )\n",
+    "\n",
+    "    predictions = jnp.argmax(logits, axis=-1)\n",
+    "\n",
+    "    accuracy = jnp.mean(predictions == labels.numpy())\n",
+    "\n",
+    "    accuracies.append(accuracy)\n",
+    "\n",
+    "print(f\"Accuracy: {np.sum(accuracies) / len(accuracies) * 100}%\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "K6i8KZl4Ba87"
+   },
+   "source": [
+    "Of course this is not the best accuracy you can get on CIFAR10, but with more training and hyperparameter tuning, you can get better results using the vision transformer."
+   ]
+  }
+ ],
+ "metadata": {
+  "accelerator": "GPU",
+  "colab": {
+   "gpuType": "T4",
+   "provenance": []
+  },
+  "kernelspec": {
+   "display_name": "jax",
+   "language": "python",
+   "name": "jax"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.11.4"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 4
+}
diff --git a/mkdocs.yml b/mkdocs.yml
index 8f363f65..b1eb4dd0 100644
--- a/mkdocs.yml
+++ b/mkdocs.yml
@@ -104,6 +104,7 @@ nav:
             - Generative score-based diffusion: 'examples/score_based_diffusion.ipynb'
             - BERT language model: 'examples/bert.ipynb'
             - U-Net implementation: 'examples/unet.ipynb'
+            - Vision transformer: 'examples/vision_transformer.ipynb'
             - Image GAN: 'examples/deep_convolutional_gan.ipynb'
         - Features:
             - Freezing parameters: 'examples/frozen_layer.ipynb'
@@ -127,6 +128,8 @@ nav:
             - 'api/nn/embedding.md'
             - 'api/nn/mlp.md'
             - 'api/nn/sequential.md'
+            - 'api/nn/inference.md'
+            - 'api/nn/shared.md'
             - 'api/nn/stateful.md'
         - Filtering:
             - 'api/filtering/partition-combine.md'
@@ -140,6 +143,7 @@ nav:
         - 'api/pretty-printing.md'
         - 'api/serialisation.md'
     - Misc:
+        - 'faq.md'
         - 'tricks.md'
+        # - 'pattern.md'
         - 'citation.md'
-        - 'faq.md'
diff --git a/pyproject.toml b/pyproject.toml
index 0ec31729..622cea25 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,6 +1,6 @@
 [project]
 name = "equinox"
-version = "0.10.11"
+version = "0.11.0"
 description = "Elegant easy-to-use neural networks in JAX."
 readme = "README.md"
 requires-python ="~=3.9"


Test patch:
diff --git a/tests/requirements.txt b/tests/requirements.txt
index 6147e392..946ff033 100644
--- a/tests/requirements.txt
+++ b/tests/requirements.txt
@@ -1,3 +1,4 @@
 jaxlib
+optax
 pytest
 beartype
diff --git a/tests/test_ad.py b/tests/test_ad.py
index fd8e1c76..409a3706 100644
--- a/tests/test_ad.py
+++ b/tests/test_ad.py
@@ -340,6 +340,15 @@ def f(x, y):
     f(1.0, 1.0)
 
 
+def test_closure_convert_trivial():
+    def f(a):
+        return a + 1
+
+    f2 = eqx.filter_closure_convert(f, 1)
+    f2.out_struct
+    assert type(f2).__name__ == "_TrivialClosureConvert"
+
+
 def test_closure_convert_custom_jvp():
     @eqx.filter_custom_jvp
     def call(f, x):
@@ -575,3 +584,11 @@ def f_bwd(token, ct, _, x):
         return 2 * ct
 
     jax.grad(f)(1.0)
+
+
+def test_positional_first_argument():
+    x = jnp.array(1.0)
+    with pytest.raises(TypeError, match="Functions wrapped with"):
+        eqx.filter_grad(lambda x: x + 1)(x=x)
+    with pytest.raises(TypeError, match="Functions wrapped with"):
+        eqx.filter_value_and_grad(lambda x, y: x + y)(x=x, y=x)
diff --git a/tests/test_checkpoint.py b/tests/test_checkpoint.py
new file mode 100644
index 00000000..f62394e6
--- /dev/null
+++ b/tests/test_checkpoint.py
@@ -0,0 +1,25 @@
+from functools import partial
+
+import jax.numpy as jnp
+import jax.tree_util as jtu
+
+import equinox as eqx
+
+
+def test_checkpoint(getkey):
+    mlp = eqx.nn.MLP(1, 2, 5, 2, key=getkey())
+
+    def fun(mlp, x, checkpoint):
+        if checkpoint:
+            mlp = eqx.filter_checkpoint(mlp)
+        return jnp.sum(mlp(x))
+
+    x = jnp.array([1.0])
+
+    grad_f_check = eqx.filter_grad(partial(fun, checkpoint=True))
+    grad_f_nocheck = eqx.filter_grad(partial(fun, checkpoint=False))
+
+    for l1, l2 in zip(
+        jtu.tree_leaves(grad_f_check(mlp, x)), jtu.tree_leaves(grad_f_nocheck(mlp, x))
+    ):
+        assert jnp.allclose(l1, l2)
diff --git a/tests/test_closure_to_pytree.py b/tests/test_closure_to_pytree.py
new file mode 100644
index 00000000..ecf1d4e2
--- /dev/null
+++ b/tests/test_closure_to_pytree.py
@@ -0,0 +1,73 @@
+import jax.numpy as jnp
+import jax.tree_util as jtu
+import optax
+
+import equinox as eqx
+import equinox.internal as eqxi
+
+
+def test_fixup_optax():
+    lr = jnp.array(1e-3)
+    optim = optax.chain(
+        optax.adam(lr),
+        optax.scale_by_schedule(optax.piecewise_constant_schedule(1, {200: 0.1})),
+    )
+    optim = eqxi.closure_to_pytree(optim)
+
+    for leaf in jtu.tree_leaves(optim):
+        if eqx.is_array(leaf) and leaf == -lr:
+            break
+    else:
+        assert False
+
+    # Check that we can still init and update as normal.
+    grads = params = {"foo": jnp.array(1.0)}
+    state = optim.init(params)
+    optim.update(grads, state)
+
+    lr = jnp.array(1e-2)
+    optim2 = optax.chain(
+        optax.adam(lr),
+        optax.scale_by_schedule(optax.piecewise_constant_schedule(1, {200: 0.1})),
+    )
+    optim2 = eqxi.closure_to_pytree(optim2)
+
+    compiling = 0
+
+    @eqx.filter_jit
+    def f(x):
+        nonlocal compiling
+        compiling += 1
+
+    f(optim)
+    assert compiling == 1
+    f(optim2)
+    assert compiling == 1
+
+
+def test_closure_same_name():
+    def f(flag):
+        if flag:
+
+            def g(y):
+                return 1 + y
+
+        else:
+
+            def g(y):
+                return 2 + y
+
+        def h(y):
+            return g(y)
+
+        return h
+
+    h1 = eqxi.closure_to_pytree(f(True))
+    h2 = eqxi.closure_to_pytree(f(False))
+
+    @eqx.filter_jit
+    def run(f, y):
+        return f(y)
+
+    assert run(h1, jnp.array(1.0)) == 2
+    assert run(h2, jnp.array(1.0)) == 3
diff --git a/tests/test_enum.py b/tests/test_enum.py
index 71e37af6..305ea9d6 100644
--- a/tests/test_enum.py
+++ b/tests/test_enum.py
@@ -2,6 +2,7 @@
 import jax.numpy as jnp
 import pytest
 
+import equinox as eqx
 import equinox.internal as eqxi
 
 from .helpers import shaped_allclose
@@ -245,11 +246,11 @@ class A(eqxi.Enumeration):
 
     token = jnp.array(True)
     A.a.error_if(token, False)
-    jax.jit(A.a.error_if)(token, False)
+    eqx.filter_jit(A.a.error_if)(token, False)
     with pytest.raises(Exception):
         A.a.error_if(token, True)
     with pytest.raises(Exception):
-        jax.jit(A.a.error_if)(token, True)
+        eqx.filter_jit(A.a.error_if)(token, True)
 
 
 def test_compile_time_eval():
diff --git a/tests/test_errors.py b/tests/test_errors.py
index 18354973..031f2cd7 100644
--- a/tests/test_errors.py
+++ b/tests/test_errors.py
@@ -181,11 +181,13 @@ def _raises():
         while tb is not None:
             code_stack.append(tb.tb_frame.f_code)
             tb = tb.tb_next
-        assert len(code_stack) == 3
-        one, two, three = code_stack
+        assert len(code_stack) == 4
+        one, two, three, four = code_stack
         assert one.co_filename.endswith("test_errors.py")
         assert one.co_name == "test_traceback_runtime_custom"
         assert two.co_filename.endswith("equinox/_jit.py")
         assert two.co_name == "__call__"
-        assert three.co_filename.endswith("equinox/_jit.py")
-        assert three.co_name == "_call"
+        assert three.co_filename.endswith("equinox/_module.py")
+        assert three.co_name == "__call__"
+        assert four.co_filename.endswith("equinox/_jit.py")
+        assert four.co_name == "_call"
diff --git a/tests/test_module.py b/tests/test_module.py
index 0f426d31..3840a3d9 100644
--- a/tests/test_module.py
+++ b/tests/test_module.py
@@ -1,5 +1,9 @@
+import abc
+import dataclasses
 import functools as ft
-from typing import Any
+from collections.abc import Callable
+from dataclasses import InitVar
+from typing import Any, Optional
 
 import jax
 import jax.numpy as jnp
@@ -7,6 +11,7 @@
 import pytest
 
 import equinox as eqx
+import equinox.internal as eqxi
 
 from .helpers import shaped_allclose
 
@@ -138,6 +143,66 @@ def __init__(self, value7: Any, **kwargs):
     assert m.weight7 == 2
 
 
+# This used to be allowed, but was frankly very sketchy -- it introduced cycles into
+# things. We now explicitly try to catch this as an error.
+def test_method_assignment():
+    class A(eqx.Module):
+        foo: Callable
+        x: jax.Array
+
+        def __init__(self, x):
+            # foo is assigned before x! We have that `self.bar.__self__` is a copy of
+            # `self`, but for which self.bar.__self__.x` doesn't exist yet. Then later
+            # calling `self.foo()` would raise an error.
+            self.foo = self.bar
+            self.x = x
+
+        def bar(self):
+            return self.x + 1
+
+    x = jnp.array(3)
+    with pytest.raises(ValueError, match="Cannot assign methods in __init__"):
+        A(x)
+
+
+def test_method_assignment2():
+    # From https://github.com/patrick-kidger/equinox/issues/327
+    class Component(eqx.Module):
+        transform: Callable[[float], float]
+        validator: Callable[[Callable], Callable]
+
+        def __init__(self, transform=lambda x: x, validator=lambda f: f) -> None:
+            self.transform = transform
+            self.validator = validator
+
+        def __call__(self, x):
+            return self.validator(self.transform)(x)
+
+    class SubComponent(Component):
+        test: Callable[[float], float]
+
+        def __init__(self, test=lambda x: 2 * x) -> None:
+            self.test = test
+            super().__init__(self._transform)
+
+        def _transform(self, x):
+            return self.test(x)
+
+    with pytest.raises(ValueError, match="Cannot assign methods in __init__"):
+        SubComponent()
+
+
+def test_method_access_during_init():
+    class Foo(eqx.Module):
+        def __init__(self):
+            self.method()
+
+        def method(self):
+            pass
+
+    Foo()
+
+
 @pytest.mark.parametrize("new", (False, True))
 def test_static_field(new):
     if new:
@@ -171,8 +236,165 @@ class MyModuleWithInit(eqx.Module):
 
         def __init__(self, a):
             self.field = a
+            assert type(self.field) is float
 
-    assert shaped_allclose(MyModule(1.0).field, jnp.array(1.0))  # pyright: ignore
+    assert shaped_allclose(MyModuleWithInit(1.0).field, jnp.array(1.0))
+
+    called = False
+
+    class MyModuleWithPostInit(eqx.Module):
+        field: jax.Array = eqx.field(converter=jnp.asarray)
+
+        def __post_init__(self):
+            nonlocal called
+            assert not called
+            called = True
+            assert shaped_allclose(self.field, jnp.array(1.0))
+
+    MyModuleWithPostInit(1.0)  # pyright: ignore
+    assert called
+
+
+def test_converter_monkeypatched_init():
+    class Foo(eqx.Module):
+        field: jax.Array = eqx.field(converter=jnp.asarray)
+
+    assert shaped_allclose(Foo(1.0).field, jnp.array(1.0))  # pyright: ignore
+
+    called = False
+    init = Foo.__init__
+
+    def __init__(self, *args, **kwargs):
+        nonlocal called
+        assert not called
+        called = True
+        init(self, *args, **kwargs)
+
+    Foo.__init__ = __init__
+    assert shaped_allclose(Foo(1.0).field, jnp.array(1.0))  # pyright: ignore
+    assert called
+
+
+# Note that `Foo` had to start with a `__post_init__` method for this to work.
+# Dataclasses check for the presence of a `__post_init__` method when the class is
+# created, and at that time creates a flag declaring whether to run `__post_init__` at
+# initialisation time.
+def test_converter_monkeypatched_postinit():
+    called1 = False
+
+    class Foo(eqx.Module):
+        field: jax.Array = eqx.field(converter=jnp.asarray)
+
+        def __post_init__(self):
+            nonlocal called1
+            assert not called1
+            called1 = True
+            assert shaped_allclose(self.field, jnp.array(1.0))
+
+    assert shaped_allclose(Foo(1.0).field, jnp.array(1.0))  # pyright: ignore
+    assert called1
+
+    called2 = False
+    post_init = Foo.__post_init__
+
+    def __post_init__(self):
+        nonlocal called1
+        nonlocal called2
+        assert not called2
+        called1 = False
+        called2 = True
+        post_init(self)
+
+    Foo.__post_init__ = __post_init__  # pyright: ignore
+    assert shaped_allclose(Foo(1.0).field, jnp.array(1.0))  # pyright: ignore
+    assert called2
+    assert called1
+
+
+@pytest.mark.parametrize("base_is_module", (False, True))
+def test_converter_init_hierarchy(base_is_module):
+    class A(eqx.Module if base_is_module else object):  # pyright: ignore
+        def __init__(self, x):
+            nonlocal called
+            assert not called
+            called = True
+            self.x = x
+
+    class B(eqx.Module):
+        x: jax.Array = eqx.field(converter=jnp.asarray)
+
+    class C(A, B):
+        # Use `A.__init__`
+        pass
+
+    class D(B, A):
+        # Use the autogenerated `B.__init__`
+        pass
+
+    # In either case, conversion should happen.
+
+    called = False
+    assert shaped_allclose(C(1).x, jnp.array(1))
+    assert called
+
+    assert shaped_allclose(D(1).x, jnp.array(1))  # pyright: ignore
+    # No `called` check, we're not using `A.__init__`.
+
+
+@pytest.mark.parametrize("base_is_module", (False, True))
+def test_converter_post_init_hierarchy(base_is_module):
+    class A(eqx.Module if base_is_module else object):  # pyright: ignore
+        def __post_init__(self):
+            nonlocal called
+            assert not called
+            called = True
+
+    class B(eqx.Module):
+        x: jax.Array = eqx.field(converter=jnp.asarray)
+
+    class C(A, B):
+        pass
+
+    class D(B, A):
+        pass
+
+    called = False
+    assert shaped_allclose(C(1).x, jnp.array(1))  # pyright: ignore
+    assert called
+
+    called = False
+    assert shaped_allclose(D(1).x, jnp.array(1))  # pyright: ignore
+    assert called
+
+
+def test_init_and_postinit():
+    class Foo(eqx.Module):
+        field: jax.Array = eqx.field(converter=jnp.asarray)
+
+        def __post_init__(self):
+            assert False
+
+    with pytest.warns(UserWarning, match="__init__` method and a `__post_init__"):
+
+        class Bar(Foo):
+            def __init__(self):
+                self.field = 1  # pyright: ignore
+
+    assert shaped_allclose(Bar().field, jnp.array(1))
+
+    class Qux(eqx.Module):
+        field: jax.Array = eqx.field(converter=jnp.asarray)
+
+        def __init__(self):
+            self.field = 1  # pyright: ignore
+
+    with pytest.warns(UserWarning, match="__init__` method and a `__post_init__"):
+
+        class Quux(Qux):
+            def __post_init__(self):
+                assert False
+
+    assert shaped_allclose(Quux().field, jnp.array(1))  # pyright: ignore
 
 
 def test_wrap_method():
@@ -183,13 +405,29 @@ def f(self, b):
             return self.a + b
 
     m = MyModule(13)
-    assert isinstance(m.f, jtu.Partial)
+    assert isinstance(m.f, eqx.Module)
     flat, treedef = jtu.tree_flatten(m.f)
     assert len(flat) == 1
     assert flat[0] == 13
     assert jtu.tree_unflatten(treedef, flat)(2) == 15
 
 
+def test_eq_method():
+    # Expected behaviour from non-Module methods
+    class A:
+        def f(self):
+            pass
+
+    a = A()
+    assert a.f == a.f
+
+    class B(eqx.Module):
+        def f(self):
+            pass
+
+    assert B().f == B().f
+
+
 def test_init_subclass():
     ran = []
 
@@ -205,17 +443,40 @@ class AnotherModule(MyModule):
 
 
 def test_wrapper_attributes():
-    def f(x):
-        pass
-
-    fjit = eqx.filter_jit(f)
-    # Gets __name__ attribute from module_update_wrapper
-
-    @eqx.filter_jit  # Flattens and unflattens
-    def g(k):
-        k.__name__
-
-    g(fjit)
+    def f(x: int) -> str:
+        """some doc"""
+        return "hi"
+
+    def h(x: int) -> str:
+        """some other doc"""
+        return "bye"
+
+    noinline = lambda x: eqxi.noinline(h, abstract_fn=x)
+
+    for wrapper in (
+        eqx.filter_jit,
+        eqx.filter_grad,
+        eqx.filter_value_and_grad,
+        eqx.filter_vmap,
+        eqx.filter_pmap,
+        noinline,
+    ):
+        f_wrap = wrapper(f)
+        # Gets __name__ attribute from module_update_wrapper
+
+        called = False
+
+        @eqx.filter_jit  # Flattens and unflattens
+        def g(k):
+            nonlocal called
+            called = True
+            assert k.__name__ == "f"
+            assert k.__doc__ == "some doc"
+            assert k.__qualname__ == "test_wrapper_attributes.<locals>.f"
+            assert k.__annotations__ == {"x": int, "return": str}
+
+        g(f_wrap)
+        assert called
 
 
 # https://github.com/patrick-kidger/equinox/issues/337
@@ -252,3 +513,480 @@ def test_wrapped(getkey):
     y = eqx.filter_vmap(eqx.nn.Linear(2, 2, key=getkey()))
     x, y = eqx.filter((x, y), eqx.is_array)
     jtu.tree_map(lambda x, y: x + y, x, y)
+
+
+def test_class_creation_kwargs():
+    called = False
+
+    class A(eqx.Module):
+        def __init_subclass__(cls, foo, **kwargs) -> None:
+            nonlocal called
+            assert not called
+            called = True
+            assert foo
+            assert len(kwargs) == 0
+
+    class B(A, foo=True):
+        pass
+
+    assert called
+
+
+def test_check_init():
+    class FooException(Exception):
+        pass
+
+    called_a = False
+    called_b = False
+
+    class A(eqx.Module):
+        a: int
+
+        def __check_init__(self):
+            nonlocal called_a
+            called_a = True
+            if self.a >= 0:
+                raise FooException
+
+    class B(A):
+        def __check_init__(self):
+            nonlocal called_b
+            called_b = True
+
+    class C(A):
+        pass
+
+    assert not called_a
+    assert not called_b
+    A(-1)
+    assert called_a
+    assert not called_b
+
+    called_a = False
+    with pytest.raises(FooException):
+        A(1)
+    assert called_a
+    assert not called_b
+
+    called_a = False
+    B(-1)
+    assert called_a
+    assert called_b
+
+    called_a = False
+    called_b = False
+    with pytest.raises(FooException):
+        B(1)
+    assert called_a
+    assert called_b  # B.__check_init__ is called before A.__check_init__
+
+    called_a = False
+    called_b = False
+    C(-1)
+    assert called_a
+    assert not called_b
+
+    called_a = False
+    with pytest.raises(FooException):
+        C(1)
+    assert called_a
+    assert not called_b
+
+
+def test_check_init_order():
+    called_a = False
+    called_b = False
+    called_c = False
+
+    class A(eqx.Module):
+        def __check_init__(self):
+            nonlocal called_a
+            called_a = True
+
+    class B(A):
+        def __check_init__(self):
+            nonlocal called_b
+            called_b = True
+            raise ValueError
+
+    class C(B):
+        def __check_init__(self):  # pyright: ignore
+            nonlocal called_c
+            called_c = True
+
+    with pytest.raises(ValueError):
+        C()
+
+    assert called_c
+    assert called_b
+    assert not called_a
+
+
+def test_check_init_no_assignment():
+    class A(eqx.Module):
+        x: int
+
+        def __check_init__(self):
+            self.x = 4
+
+    with pytest.raises(dataclasses.FrozenInstanceError):
+        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=True, 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=True, abstract=True):
+        x: int
+
+    class AbstractB(AbstractA, strict=True, 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():
+    with pytest.raises(
+        TypeError, match="Abstract strict `eqx.Module`s must be named starting with"
+    ):
+
+        class NotAbstract(eqx.Module, strict=True, 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=True, abstract=True):
+        foo = 4
+
+    class AbstractD2(AbstractD, strict=True, 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_post_init_warning():
+    class A(eqx.Module):
+        called = False
+
+        def __post_init__(self):
+            type(self).called = True
+
+    with pytest.warns(
+        UserWarning, match="test_module.test_post_init_warning.<locals>.B"
+    ):
+
+        class B(A):
+            def __init__(self):
+                pass
+
+    with pytest.warns(
+        UserWarning, match="test_module.test_post_init_warning.<locals>.C"
+    ):
+
+        class C(B):
+            pass
+
+    B()
+    C()
+    assert not A.called
+
+
+def test_inherit_doc():
+    # This is not what dataclasses do by default -- they would set `B.__init__.__doc__`
+    # to `None`
+
+    class A(eqx.Module):
+        pass
+
+    A.__init__.__doc__ = "Hey there!"
+
+    class B(A):
+        pass
+
+    assert B.__init__.__doc__ == "Hey there!"
+
+
+@pytest.mark.parametrize("a_post_init", (False, True))
+@pytest.mark.parametrize("b_post_init", (False, True))
+@pytest.mark.parametrize("c_post_init", (False, True))
+def test_conversion_once(a_post_init, b_post_init, c_post_init):
+    def converter(x):
+        nonlocal called
+        assert not called
+        called = True
+        return x
+
+    class A(eqx.Module):
+        x: int = eqx.field(converter=converter)
+
+        if a_post_init:
+
+            def __post_init__(self):
+                pass
+
+    class B(A):
+        if b_post_init:
+
+            def __post_init__(self):
+                pass
+
+    class C(B):
+        if c_post_init:
+
+            def __post_init__(self):
+                pass
+
+    called = False
+    A(1)
+    assert called
+
+    called = False
+    B(1)
+    assert called
+
+    called = False
+    C(1)
+    assert called
+
+
+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"):
+        A()
+
+    class B(eqx.Module):
+        x: int = eqx.field(init=False)
+
+        def __post_init__(self):
+            pass
+
+    with pytest.raises(ValueError, match="The following fields were not initialised"):
+        B()
+
+    class C(eqx.Module):
+        x: int = eqx.field(init=False)
+        flag: InitVar[bool]
+
+        def __post_init__(self, flag):
+            if flag:
+                self.x = 1
+
+    C(flag=True)
+    with pytest.raises(ValueError, match="The following fields were not initialised"):
+        C(flag=False)
+
+
+@pytest.mark.parametrize("field", (dataclasses.field, eqx.field))
+def test_init_as_abstract(field):
+    # Before the introduction of AbstractVar, it was possible to sort-of get the same
+    # behaviour by marking it as an `init=False` field. Here we check that we don't
+    # break that, in particular when it's overridden by a non-field.
+
+    class Abstract(eqx.Module):
+        foo: int = field(init=False)
+
+    class Concrete(Abstract):
+        @property
+        def foo(self):
+            return 1
+
+    x = Concrete()
+    leaves, treedef = jtu.tree_flatten(x)
+    assert len(leaves) == 0
+    y = jtu.tree_unflatten(treedef, leaves)
+    assert y.foo == 1
+
+
+# https://github.com/patrick-kidger/equinox/issues/522
+def test_custom_field():
+    def my_field(*, foo: Optional[bool] = None, **kwargs: Any):
+        metadata = kwargs.pop("metadata", {})
+        if foo is not None:
+            metadata["foo"] = foo
+        return eqx.field(metadata=metadata, **kwargs)
+
+    class ExampleModel(eqx.Module):
+        dynamic: jax.Array = my_field(foo=True)
+        static: int = my_field(foo=False, static=True)
+
+    model = ExampleModel(dynamic=jnp.array(1), static=1)
+    dynamic_field, static_field = dataclasses.fields(model)
+    assert dynamic_field.metadata == dict(foo=True)
+    assert static_field.metadata == dict(foo=False, static=True)
diff --git a/tests/test_nn.py b/tests/test_nn.py
index 2ce46abd..034892d3 100644
--- a/tests/test_nn.py
+++ b/tests/test_nn.py
@@ -152,6 +152,45 @@ def test_sequential(getkey):
     assert eqx.nn.Sequential(list(seq)) == seq
 
 
+@pytest.mark.parametrize("inner_stateful", (False, True))
+@pytest.mark.parametrize("outer_stateful", (False, True))
+def test_nested_sequential(inner_stateful, outer_stateful, getkey):
+    def make():
+        inner_seq = eqx.nn.Sequential(
+            [
+                eqx.nn.Linear(2, 4, key=getkey()),
+                eqx.nn.BatchNorm(4, axis_name="batch")
+                if inner_stateful
+                else eqx.nn.Identity(),
+                eqx.nn.Linear(4, 3, key=getkey()),
+            ]
+        )
+        outer_seq = eqx.nn.Sequential(
+            [
+                eqx.nn.Linear(5, 2, key=getkey()),
+                inner_seq,
+                eqx.nn.BatchNorm(3, axis_name="batch")
+                if outer_stateful
+                else eqx.nn.Identity(),
+                eqx.nn.Linear(3, 6, key=getkey()),
+            ]
+        )
+        return outer_seq
+
+    x = jrandom.normal(getkey(), (1, 5))
+    if inner_stateful or outer_stateful:
+        model, state = eqx.nn.make_with_state(make)()
+        out, state = jax.vmap(
+            model, axis_name="batch", in_axes=(0, None), out_axes=(0, None)
+        )(x, state)
+        assert isinstance(state, eqx.nn.State)
+    else:
+        model = make()
+        out = jax.vmap(model, axis_name="batch")(x)
+    assert isinstance(out, jax.Array)
+    assert out.shape == (1, 6)
+
+
 def test_mlp(getkey):
     mlp = eqx.nn.MLP(2, 3, 8, 2, key=getkey())
     x = jrandom.normal(getkey(), (2,))
@@ -671,17 +710,21 @@ def test_multihead_attention_deterministic(getkey):
 
 def test_embedding(getkey):
     emb = eqx.nn.Embedding(100, 512, key=getkey())
-    x = jnp.array([1])
-    assert emb(x).shape == (1, 512)
+    x = jnp.array(1)
+    assert emb(x).shape == (512,)
 
     emb = eqx.nn.Embedding(num_embeddings=10, embedding_size=20, key=getkey())
-    x = jnp.array([0])
-    assert emb(x).shape == (1, 20)
+    x = jnp.array(0)
+    assert emb(x).shape == (20,)
 
     emb = eqx.nn.Embedding(
         10, 10, weight=jnp.linspace(0.1, 10, 100).reshape(10, 10), key=getkey()
     )
-    x = jnp.array([-1])
+    x = jnp.array(-1)
+    assert jnp.allclose(emb(x), jnp.linspace(9.1, 10.0, 10))
+
+    emb = eqx.nn.Embedding(weight=jnp.linspace(0.1, 10, 100).reshape(10, 10))
+    x = jnp.array(-1)
     assert jnp.allclose(emb(x), jnp.linspace(9.1, 10.0, 10))
 
 
@@ -823,7 +866,7 @@ def test_batch_norm(getkey):
 
     # Test that the statistics don't update at inference
 
-    ibn = eqx.tree_inference(bn, value=True)
+    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)
@@ -844,7 +887,7 @@ def test_spectral_norm(getkey):
     def λ1():
         u, v = state.get(spectral.uv_index)
         σ = jnp.einsum("i,ij,j->", u, spectral.layer.weight, v)
-        _, s, _ = jnp.linalg.svd(spectral.layer.weight / σ)
+        _, s, _ = jnp.linalg.svd(spectral.layer.weight / σ)  # pyright: ignore
         return s[0]
 
     x = jrandom.normal(getkey(), (5,))
@@ -869,7 +912,7 @@ def λ1():
     spectral = eqx.tree_at(
         lambda s: s.layer.weight, spectral, spectral.layer.weight + 1
     )
-    spectral = eqx.tree_inference(spectral, value=True)
+    spectral = eqx.nn.inference_mode(spectral, value=True)
     assert not jnp.allclose(λ1(), 1)
     for _ in range(100):
         _, state = spectral(x, state)
diff --git a/tests/test_serialisation.py b/tests/test_serialisation.py
index bb571742..c5158b6c 100644
--- a/tests/test_serialisation.py
+++ b/tests/test_serialisation.py
@@ -3,6 +3,8 @@
 import jax
 import jax.numpy as jnp
 import numpy as np
+import pytest
+from jax.dtypes import bfloat16
 
 import equinox as eqx
 
@@ -79,6 +81,84 @@ def test_leaf_serialisation_file(getkey, tmp_path):
     assert tree_loaded_obj is like_obj
 
 
+def test_helpful_errors(getkey, tmp_path):
+    # Test that we get a helpful error message when the loading itself fails
+    tree = jnp.array(1), {}
+    eqx.tree_serialise_leaves(tmp_path, tree)
+    bad_like_tree = jnp.array(2), {
+        "a": jnp.array(2),
+        "b": jnp.array(2),
+    }
+    with pytest.raises(
+        RuntimeError,
+        match=r"Error at leaf with path \(SequenceKey\(idx=1\), DictKey\(key='a'\)\)",
+    ):
+        _ = eqx.tree_deserialise_leaves(tmp_path, bad_like_tree)
+
+    # Test that we get a helpful error message when the types don't match
+    tree = jnp.array(1), {
+        "a": jnp.array(1),
+        "b": jnp.array(1),
+    }
+    eqx.tree_serialise_leaves(tmp_path, tree)
+    bad_like_tree = jnp.array(2), {
+        "a": jnp.array(2, dtype=jnp.float32),
+        "b": jnp.array(2),
+    }
+    with pytest.raises(
+        RuntimeError,
+        match=r"Deserialised leaf at path \(SequenceKey\(idx=1\), DictKey\(key='a'\)\)",
+    ):
+        _ = eqx.tree_deserialise_leaves(tmp_path, bad_like_tree)
+
+
+def test_generic_dtype_serialisation(getkey, tmp_path):
+    # Ensure we can round trip when we start with an array
+    jax_array = jnp.array(bfloat16(1))
+    eqx.tree_serialise_leaves(tmp_path, jax_array)
+    like_jax_array = jnp.array(bfloat16(2))
+    loaded_jax_array = eqx.tree_deserialise_leaves(tmp_path, like_jax_array)
+    assert jax_array.item() == loaded_jax_array.item()
+
+    tree = (
+        jnp.array(1e-8),
+        bfloat16(1e-8),
+        np.float32(1e-8),
+        jnp.array(1e-8),
+        np.float64(1e-8),
+    )
+    like_tree = (
+        jnp.array(2.0),
+        bfloat16(2),
+        np.float32(2),
+        jnp.array(2.0),
+        np.float64(2.0),
+    )
+
+    # Ensure we can round trip when we start with a scalar
+    eqx.tree_serialise_leaves(tmp_path, tree)
+    loaded_tree = eqx.tree_deserialise_leaves(tmp_path, like_tree)
+    assert len(loaded_tree) == len(tree)
+    for a, b in zip(loaded_tree, tree):
+        assert type(a) is type(b)
+        assert a.item() == b.item()
+
+    # Ensure we can round trip when we start with a scalar that we've JAX JITed
+    # `[:-1]` to skip float64, as JIT turns it into float32
+    eqx.tree_serialise_leaves(tmp_path, jax.jit(lambda x: x)(tree[:-1]))
+    loaded_tree = eqx.tree_deserialise_leaves(tmp_path, like_tree[:-1])
+    assert len(loaded_tree) == len(tree[:-1])
+    for a, b in zip(loaded_tree, tree):
+        assert type(a) is type(b)
+        assert a.item() == b.item()
+
+
+def test_python_scalar(tmp_path):
+    eqx.tree_serialise_leaves(tmp_path, 1e-8)
+    out = eqx.tree_deserialise_leaves(tmp_path, 0.0)
+    assert out == 1e-8
+
+
 def test_custom_leaf_serialisation(getkey, tmp_path):
     jax_array1 = jnp.array(1)
     jax_array2 = jnp.array([1.0, 2.0])
diff --git a/tests/test_shared.py b/tests/test_shared.py
new file mode 100644
index 00000000..5fa332bb
--- /dev/null
+++ b/tests/test_shared.py
@@ -0,0 +1,125 @@
+import jax
+import jax.numpy as jnp
+import jax.random as jr
+import pytest
+from jaxtyping import Array, Float, Int
+
+import equinox as eqx
+
+
+def test_shared_array(getkey):
+    class MyModule(eqx.Module):
+        shared: eqx.nn.Shared
+
+        def __init__(self):
+            embedding = eqx.nn.Embedding(
+                num_embeddings=3, embedding_size=4, key=getkey()
+            )
+            head = eqx.nn.Linear(4, 3, key=getkey())
+            where = lambda pair: pair[1].weight
+            get = lambda pair: pair[0].weight
+            self.shared = eqx.nn.Shared((embedding, head), where, get)
+
+        def __call__(self, token: Int[Array, ""]):
+            nonlocal called
+            called = True
+            embedding, head = self.shared()
+            assert embedding.weight is head.weight
+            return head(embedding(token))
+
+    called = False
+    module = MyModule()
+    module(jnp.array(0))
+    assert called
+
+
+# We share a non-leaf node
+def test_shared_node(getkey):
+    class MyModule(eqx.Module):
+        shared: eqx.nn.Shared
+
+        def __init__(self):
+            attention = eqx.nn.MultiheadAttention(
+                num_heads=3, query_size=12, key=getkey()
+            )
+            my_proj = eqx.nn.Linear(12, 12, use_bias=False, key=getkey())
+            where = lambda pair: pair[1].key_proj
+            get = lambda pair: pair[0]
+            self.shared = eqx.nn.Shared((my_proj, attention), where, get)
+
+        def __call__(self, x: Float[Array, "seq 12"]):
+            nonlocal called
+            called = True
+            my_proj, attention = self.shared()
+            eq = eqx.tree_equal(my_proj, attention.key_proj)
+            x = attention(x, x, x)
+            out = jax.vmap(my_proj)(x)
+            return out, eq
+
+    called = False
+    module = MyModule()
+    x = jr.normal(getkey(), (5, 12))
+
+    @eqx.filter_jit
+    @eqx.filter_grad(has_aux=True)
+    def f(module, x):
+        out, eq = module(x)
+        return jnp.sum(out), eq
+
+    d_module, eq = f(module, x)
+    assert called
+    assert eq
+    module = eqx.apply_updates(module, d_module)
+    d_module, eq = f(module, x)
+    assert eq
+    module = eqx.apply_updates(module, d_module)
+
+
+def test_mismatched_structure(getkey):
+    x = jr.normal(getkey(), (3, 4))
+    y = jr.normal(getkey(), (4, 3))
+    with pytest.raises(ValueError, match="Every node being shared together"):
+        eqx.nn.Shared((x, y), lambda pair: pair[0], lambda pair: pair[1])
+
+
+def test_multi_shared(getkey):
+    class MyModule(eqx.Module):
+        shared: eqx.nn.Shared
+
+        def __init__(self):
+            my_proj = eqx.nn.Linear(12, 12, use_bias=False, key=getkey())
+            attention = eqx.nn.MultiheadAttention(
+                num_heads=3, query_size=12, key=getkey()
+            )
+            where = lambda pair: (pair[1].key_proj, pair[1].query_proj.weight)
+            get = lambda pair: (pair[0], pair[0].weight + 1)
+            self.shared = eqx.nn.Shared((my_proj, attention), where, get)
+
+        def __call__(self, x: Float[Array, "seq 12"]):
+            nonlocal called
+            called = True
+            my_proj, attention = self.shared()
+            eq1 = eqx.tree_equal(my_proj, attention.key_proj)
+            eq2 = (my_proj.weight + 1 == attention.query_proj.weight).all()
+            x = attention(x, x, x)
+            out = jax.vmap(my_proj)(x)
+            eq = eq1 & eq2
+            return out, eq
+
+    called = False
+    module = MyModule()
+    x = jr.normal(getkey(), (5, 12))
+
+    @eqx.filter_jit
+    @eqx.filter_grad(has_aux=True)
+    def f(module, x):
+        out, eq = module(x)
+        return jnp.sum(out), eq
+
+    d_module, eq = f(module, x)
+    assert called
+    assert eq
+    module = eqx.apply_updates(module, d_module)
+    d_module, eq = f(module, x)
+    assert eq
+    module = eqx.apply_updates(module, d_module)
diff --git a/tests/test_stateful.py b/tests/test_stateful.py
new file mode 100644
index 00000000..dc97c13e
--- /dev/null
+++ b/tests/test_stateful.py
@@ -0,0 +1,74 @@
+import jax
+import jax.numpy as jnp
+import jax.random as jr
+import jax.tree_util as jtu
+import pytest
+
+import equinox as eqx
+
+
+def test_delete_init_state():
+    model = eqx.nn.BatchNorm(3, "batch")
+    eqx.nn.State(model)
+    model2 = eqx.nn.delete_init_state(model)
+
+    eqx.nn.State(model)
+    with pytest.raises(ValueError):
+        eqx.nn.State(model2)
+
+    leaves = [x for x in jtu.tree_leaves(model) if eqx.is_array(x)]
+    leaves2 = [x for x in jtu.tree_leaves(model2) if eqx.is_array(x)]
+    assert len(leaves) == len(leaves2) + 3
+
+
+def test_double_state():
+    # From https://github.com/patrick-kidger/equinox/issues/450#issuecomment-1714501666
+
+    class Counter(eqx.Module):
+        index: eqx.nn.StateIndex
+
+        def __init__(self):
+            init_state = jnp.array(0)
+            self.index = eqx.nn.StateIndex(init_state)
+
+        def __call__(self, x, state):
+            value = state.get(self.index)
+            new_x = x + value
+            new_state = state.set(self.index, value + 1)
+            return new_x, new_state
+
+    class Model(eqx.Module):
+        linear: eqx.nn.Linear
+        counter: Counter
+        v_counter: Counter
+
+        def __init__(self, key):
+            # Not-stateful layer
+            self.linear = eqx.nn.Linear(2, 2, key=key)
+            # Stateful layer.
+            self.counter = Counter()
+            # Vmap'd stateful layer. (Whose initial state will include a batch
+            # dimension.)
+            self.v_counter = eqx.filter_vmap(Counter, axis_size=2)()
+
+        def __call__(self, x, state):
+            assert x.shape == (2,)
+            x = self.linear(x)
+            x, state = self.counter(x, state)
+            substate = state.substate(self.v_counter)
+            x, substate = eqx.filter_vmap(self.v_counter)(x, substate)
+            state = state.update(substate)
+            return x, state
+
+    key = jr.PRNGKey(0)
+    model, state = eqx.nn.make_with_state(Model)(key)
+    x = jnp.array([5.0, -1.0])
+    model(x, state)
+
+    @jax.jit
+    def make_state(key):
+        _, state = eqx.nn.make_with_state(Model)(key)
+        return state
+
+    new_state = make_state(jr.PRNGKey(1))
+    model(x, new_state)
diff --git a/tests/test_tree.py b/tests/test_tree.py
index 1f40cc07..c5f92623 100644
--- a/tests/test_tree.py
+++ b/tests/test_tree.py
@@ -1,14 +1,16 @@
-from typing import Callable
-
 import jax
+import jax.core
 import jax.nn as jnn
 import jax.numpy as jnp
 import jax.random as jrandom
 import jax.tree_util as jtu
+import numpy as np
 import pytest
 
 import equinox as eqx
 
+from .helpers import shaped_allclose
+
 
 def test_tree_at_replace(getkey):
     key = getkey()
@@ -155,10 +157,45 @@ def run3(x, y):
     assert not run3(a, 1)
 
 
-def test_tree_inference(getkey):
+def test_tree_equal_numpy():
+    x = np.array([1, 2], dtype=np.float32)
+    y = jnp.array([1, 2], dtype=jnp.float32)
+    z = jnp.array([1, 2], dtype=jnp.float16)
+    assert shaped_allclose(eqx.tree_equal(x, x), np.bool_(True))
+    assert shaped_allclose(eqx.tree_equal(x, y), jnp.array(True))
+    assert shaped_allclose(eqx.tree_equal(y, y), jnp.array(True))
+
+    assert shaped_allclose(eqx.tree_equal(x, z), False)
+
+    @jax.jit
+    def f():
+        assert shaped_allclose(eqx.tree_equal(x, x), np.bool_(True))
+        out = eqx.tree_equal(x, y)
+        assert isinstance(out, jax.core.Tracer)
+        return out
+
+    assert shaped_allclose(f(), jnp.array(True))
+
+
+def test_tree_equal_scalars():
+    x = np.float32(1)
+    y = np.array(1, dtype=np.float32)
+    z = np.array(1, dtype=np.float16)
+    # scalar-ness does not matter
+    assert shaped_allclose(eqx.tree_equal(x, y), np.bool_(True))
+    # dtype does matter
+    assert shaped_allclose(eqx.tree_equal(x, z), False)
+
+    z = jax.dtypes.bfloat16(1)
+    w = jax.dtypes.bfloat16(2)
+    assert shaped_allclose(eqx.tree_equal(z, z), np.bool_(True))
+    assert shaped_allclose(eqx.tree_equal(z, w), np.bool_(False))
+
+
+def test_inference_mode(getkey):
     attention = eqx.nn.MultiheadAttention(2, 4, key=getkey())
     assert attention.dropout.inference is False
-    attention2 = eqx.tree_inference(attention, True)
+    attention2 = eqx.nn.inference_mode(attention)
     assert attention.dropout.inference is False
     assert attention2.dropout.inference is True
 
@@ -188,6 +225,14 @@ def test_tree_flatten_one_level():
         eqx.tree_flatten_one_level(x)
 
 
+# This matches the behaviour of `jax._src.tree_util.flatten_one_level`
+def test_tree_flatten_one_level_special():
+    x = [None, None, eqx.Module(), 1, 2]
+    leaves, treedef = eqx.tree_flatten_one_level(x)
+    assert leaves == [None, None, eqx.Module(), 1, 2]
+    assert treedef == jtu.tree_structure([0, 0, 0, 0, 0])
+
+
 def test_tree_check():
     x = []
     y = []
@@ -201,28 +246,15 @@ def test_tree_check():
     with pytest.raises(ValueError):
         eqx.tree_check(x)
 
-    # From https://github.com/patrick-kidger/equinox/issues/327
-    class Component(eqx.Module):
-        transform: Callable[[float], float]
-        validator: Callable[[Callable], Callable]
-
-        def __init__(self, transform=lambda x: x, validator=lambda f: f) -> None:
-            self.transform = transform
-            self.validator = validator
 
-        def __call__(self, x):
-            return self.validator(self.transform)(x)
+def test_tree_check_none():
+    eqx.tree_check([None, None])
 
-    class SubComponent(Component):
-        test: Callable[[float], float]
 
-        def __init__(self, test=lambda x: 2 * x) -> None:
-            self.test = test
-            super().__init__(self._transform)
+def test_tree_check_integer():
+    eqx.tree_check([0, 0])
 
-        def _transform(self, x):
-            return self.test(x)
 
-    a = SubComponent()
-    with pytest.raises(ValueError):
-        eqx.tree_check(a)
+def test_tree_check_module():
+    a = eqx.Module()  # same `id(...)` for both entries passed to `tree_check`.
+    eqx.tree_check([a, a])