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

Issue:
Add proper documentation!
Rather than the current README-slog ;) 
Tweak filtering to match JAX?
At the moment we mostly do `isinstance` checks to determine what is an array, or array-like.

However see https://github.com/google/jax/blob/4a17c78605e7fc69a69a999e2f6298db79d3837a/jax/_src/numpy/lax_numpy.py#L542 which is JAX's internal definition of what is array-like.

Mostly it looks like these overlap, but there's a couple of edge cases. (Such as `__jax_array__`.)

Besides this, JAX now supports meaningful `isinstance(..., jnp.ndarray)`, so we should bump the minimum version required and switch to using that.


Discussion / hints:

Fix patch:
diff --git a/.gitignore b/.gitignore
index d05e8374..ec6f56a4 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,4 +1,7 @@
 **/__pycache__
+**/.ipynb_checkpoints
 *.egg-info/
 build/
 dist/
+site/
+.all_objects.cache
diff --git a/.isort.cfg b/.isort.cfg
index bdc8fe4c..64737cc5 100644
--- a/.isort.cfg
+++ b/.isort.cfg
@@ -3,4 +3,3 @@ force_alphabetical_sort_within_sections=true
 lines_after_imports=2
 profile=black
 treat_comments_as_code=true
-extra_standard_library=typing_extensions
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index faa62ec2..07ac779b 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -3,6 +3,12 @@ repos:
     rev: stable
     hooks:
     - id: black
+  - repo: https://github.com/nbQA-dev/nbQA
+    rev: 1.1.1
+    hooks:
+    - id: nbqa-black
+    - id: nbqa-isort 
+    - id: nbqa-flake8
   - repo: https://github.com/PyCQA/isort
     rev: 5.9.3
     hooks:
diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md
index fbe0ae3f..610027c2 100644
--- a/CONTRIBUTING.md
+++ b/CONTRIBUTING.md
@@ -4,6 +4,8 @@ Contributions (pull requests) are very welcome! Here's how to get started.
 
 ---
 
+**Getting started**
+
 First fork the library on GitHub.
 
 Then clone and install the library in development mode:
@@ -23,6 +25,10 @@ pre-commit install
 
 These hooks use Black and isort to format the code, and flake8 to lint it.
 
+---
+
+**If you're making changes to the code:**
+
 Now make your changes. Make sure to include additional tests if necessary.
 
 Next verify the tests all pass:
@@ -39,3 +45,21 @@ git push
 ```
 
 Finally, open a pull request on GitHub!
+
+---
+
+**If you're making changes to the documentation:**
+
+Make your changes. You can then build the documentation by doing
+
+```bash
+pip install docs/requirements.txt
+mkdocs serve
+```
+Then doing `Control-C`, and running:
+```
+mkdocs serve
+```
+(So you run `mkdocs serve` twice.)
+
+You can then see your local copy of the documentation by navigating to `localhost:8000` in a web browser.
diff --git a/README.md b/README.md
index c2a51562..f77e9400 100644
--- a/README.md
+++ b/README.md
@@ -1,151 +1,72 @@
 <h1 align='center'>Equinox</h1>
-<h2 align='center'>Callable PyTrees and filtered JIT/grad transformations<br>=> neural networks in JAX</h2>
 
-Equinox brings more power to your model building in [JAX](https://github.com/google/jax).<br>
-Represent *parameterised functions as data* and use *filtered transformations* for powerful fine-grained control of the model-building process. Equinox demonstrates how to use a PyTorch-like class-based API without compromising on JAX-like functional programming.
+Equinox is a JAX library based around a simple idea: **represent parameterised functions (such as neural networks) as PyTrees.**
 
-Equinox is half tech-demo, half neural network library, and comes with no behind-the-scenes magic, guaranteed.
+In doing so:
+
+- We get a PyTorch-like API...
+- ...that's fully compatible with *native* JAX transformations...
+- ...with no new concepts you have to learn. (It's all just PyTrees.)
 
 The elegance of Equinox is its selling point in a world that already has [Haiku](https://github.com/deepmind/dm-haiku), [Flax](https://github.com/google/flax) and so on.
 
+_(In other words, why should you care? Because Equinox is really simple to learn, and really simple to use.)_
 
-## Quick start
-### Installation
+## Installation
 
-```
+```bash
 pip install equinox
 ```
-Requires Python 3.7+ and JAX 0.2.18+.
-
-### Parameterised functions as data
-
-Equinox represents *parameterised functions as [PyTrees](https://jax.readthedocs.io/en/latest/pytrees.html)*. (For example a neural network is a function parameterised by its weights, biases, etc.) Now you can JIT/grad/etc. a higher-order function (like a loss function) with respect to a parameterised function as its input (like a model).
-
-Previous libraries have introduced a lot of extra complexity to make this work. e.g. custom notions of parameter groups, class-to-functional transformations, or specially-wrapped `library.jit` or `library.grad` to be compatible with JAX's JIT/grad/etc.
-
-In contrast, Equinox makes it elegant:
-
-```python
-import equinox as eqx
-import jax
-import jax.nn as jnn
-import jax.numpy as jnp
-import jax.random as jrandom
-
-class MyModule(eqx.Module):
-    # Specify the module's attributes;
-    layers: list
-    bias: jnp.ndarray
-    
-    # And how to initialise them;
-    def __init__(self, key):
-        key1, key2 = jrandom.split(key)
-        self.layers = [eqx.nn.Linear(2, 8, key=key1),
-                       eqx.nn.Linear(8, 2, key=key2)]
-        self.bias = jnp.ones(2)
 
-    # And the forward pass of the model.
-    def __call__(self, x):
-        for layer in self.layers[:-1]:
-            x = jnn.relu(layer(x))
-        return self.layers[-1](x) + self.bias
-
-@jax.jit
-@jax.grad
-def loss(model, x, y):
-    pred_y = jax.vmap(model)(x)
-    return jnp.mean((y - pred_y) ** 2)
-
-x_key, y_key, model_key = jrandom.split(jrandom.PRNGKey(0), 3)
-x, y = jrandom.normal(x_key, (100, 2)), jrandom.normal(y_key, (100, 2))
-model = MyModule(model_key)
-grads = loss(model, x, y)
-```
+Requires Python 3.8+ and JAX 0.2.18+.
 
-And there's no magic there! All `eqx.Module` really does is register your class with JAX as a PyTree node. (In fact the source code for `eqx.Module` is only about 100 lines long.)
+## Documentation
 
-### Filtering
+Available at [https://docs.kidger.site/equinox](https://docs.kidger.site/equinox).
 
-In the previous example, all of the model attributes were Modules and JAX arrays.
+## Quick example
 
-Arbitrary Python objects are fine too! We just need to handle them appropriately around `jax.jit` and `jax.grad`.
+Models are defined using PyTorch-like syntax:
 
 ```python
 import equinox as eqx
-import functools as ft
 import jax
-import jax.nn as jnn
-import jax.numpy as jnp
-import jax.random as jrandom
 
-class AnotherModule(eqx.Module):
-    layers: list
+class Linear(eqx.Module):
+    weight: jax.numpy.ndarray
+    bias: jax.numpy.ndarray
 
-    def __init__(self, key):
-        key1, key2 = jrandom.split(key)
-        # Model now has `jnn.relu` -- a Python function -- as part of its PyTree.
-        self.layers = [eqx.nn.Linear(2, 8, key=key1),
-                       jnn.relu,
-                       eqx.nn.Linear(8, 2, key=key2)]
+    def __init__(self, in_size, out_size, key):
+        wkey, bkey = jax.random.split(key)
+        self.weight = jax.random.normal(wkey, (out_size, in_size))
+        self.bias = jax.random.normal(bkey, (out_size,))
 
     def __call__(self, x):
-        for layer in self.layers:
-            x = layer(x)
-        return x
-
-x_key, y_key, model_key = jrandom.split(jrandom.PRNGKey(0), 3)
-x, y = jrandom.normal(x_key, (100, 2)), jrandom.normal(y_key, (100, 2))
-model = AnotherModule(model_key)
+        return self.weight @ x + self.bias
+```
 
-# Option 1: explicitly filter out anything that isn't JIT/grad-able.
+and fully compatible with normal JAX operations:
 
-@ft.partial(jax.jit, static_argnums=1)
+```python
+@jax.jit
 @jax.grad
-def loss(params, static, x, y):
-    model = eqx.combine(params, static)
+def loss_fn(model, x, y):
     pred_y = jax.vmap(model)(x)
-    return jnp.mean((y - pred_y) ** 2)
-
-params, static = eqx.partition(model, eqx.is_array)
-loss(params, static, x, y)
+    return jax.numpy.mean((y - pred_y) ** 2)
 
-# Option 2: use filtered transformations, which automates the above process for you.
-# (Can be handy if you want to JIT/grad with respect to different things!)
-
-@eqx.filter_jit
-@eqx.filter_grad
-def loss(model, x, y):
-    pred_y = jax.vmap(model)(x)
-    return jnp.mean((y - pred_y) ** 2)
-
-loss(model, x, y)
+batch_size, in_size, out_size = 32, 2, 3
+model = Linear(in_size, out_size, key=jax.random.PRNGKey(0))
+x = jax.numpy.zeros((batch_size, in_size))
+y = jax.numpy.zeros((batch_size, out_size))
+grads = loss_fn(model, x, y)
 ```
 
-Here, `params` and `static` are actually both instances of `AnotherModule`. `params` keeps just the attributes that are JAX arrays, and `static` keeps everything else. Then `combine` just merges the two PyTrees back together afterwards.
-
-### Integrates smoothly with JAX
-
-And that's it! That's pretty much all of Equinox.
-
-Equinox introduces a powerful yet straightforward way to build neural networks, without introducing lots of new notions or tieing you into a framework.
-
-Equinox is all just regular JAX -- PyTrees and transformations. Together, these two pieces allow us to specify complex models in JAX-friendly ways.
-
-## Examples
-
-- [`build_model.py`](./examples/build_model.py) builds an MLP from scratch, demonstrating the easy parameterised-functions-as-data approach that Equinox introduces. We'll then pass it into higher-order functions like JIT and grad. Overall we produce models using a familiar class-based syntax, that are also functional and integrate directly with JAX's JIT/autograd.
-
-- [`filtered_transformations.py`](./examples/filtered_transformations.py) introduces `equinox.filter_jit` and `equinox.filter_grad`. These will be used to select the parameters of an MLP and train them.
- 
-- [`frozen_layer.py`](./examples/frozen_layer.py) demonstrates how this approach really shines: some of the parameters will be trained, some of them will be frozen, but *all* of them will be efficiently JIT-traced.
-
-- [`train_rnn.py`](./examples/train_rnn.py) trains an RNN on a toy clockwise/anticlockwise spiral classification problem.
-
-- [`modules_to_initapply.py`](./examples/modules_to_initapply.py) demonstrates how to use Equinox in an init/apply-style way, which some JAX libraries have been built around. (e.g. Stax)
+Finally, there's no magic behind the scenes. All `eqx.Module` does is register your class as a PyTree. From that point onwards, JAX already knows how to work with PyTrees.
 
 ## Citation
 
-If you find Equinox useful in academic work then please consider a citation:
+If you found this library to be useful in academic work, then please cite:
+
 ```bibtex
 @article{kidger2021equinox,
     author={Patrick Kidger and Cristian Garcia},
@@ -155,247 +76,4 @@ If you find Equinox useful in academic work then please consider a citation:
 }
 ```
 
-## API
-
-### Full API list
-
-```python
-# Module                         # Neural networks
-equinox.Module                   equinox.nn.Linear
-                                 equinox.nn.Identity
-# Filtering/combining            equinox.nn.Dropout
-equinox.filter                   equinox.nn.GRUCell
-equinox.partition                equinox.nn.LSTMCell
-equinox.combine                  equinox.nn.Sequential
-                                 equinox.nn.MLP
-# Filtered transformations       equinox.nn.Conv
-equinox.filter_jit               equinox.nn.Conv1d
-equinox.filter_grad              equinox.nn.Conv2d
-equinox.filter_value_and_grad    equinox.nn.Conv3d
-                                 
-# Filters                        # Utilities
-equinox.is_array                 equinox.apply_updates
-equinox.is_array_like            equinox.static_field
-equinox.is_inexact_array         equinox.tree_at
-equinox.is_inexact_array_like    equinox.tree_equal
-```
-
-### Module
-
-```python
-equinox.Module
-```
-Base class; create your model by inheriting from this.
-
-Specify all its attributes at the class level (identical to [dataclasses](https://docs.python.org/3/library/dataclasses.html)). This defines its children in the PyTree.
-
-```python
-class MyModule(equinox.Module):
-    weight: typing.Any
-    bias: typing.Any
-    submodule: Module
-```
-
-In this case a default `__init__` method is provided, which just fills in these attributes with the argments passed: `MyModule(weight, bias, submodule)`. Alternatively you can provide an `__init__` method yourself. (For example to specify dimension sizes instead of raw weights.) By the end of `__init__`, every attribute must have been assigned.
-
-```python
-class AnotherModule(equinox.Module):
-    weight: Any
-
-    def __init__(self, input_size, output_size, key):
-        self.weight = jax.random.normal(key, (output_size, input_size))
-```
-
-After initialisation then attributes cannot be modified: models are immutable as per functional programming. (Parameter updates are made by creating a new model, not by mutating parameters in-place; see for example [`train_rnn.py`](./examples/train_rnn.py).)
-
-It is typical to also create some methods on the class. As `self` will be an input parameter -- treated as a PyTree -- then these methods will get access to the attributes of the instance. Defining `__call__` gives an easy way to define a forward pass for a model (although any method can be used, and no methods are special-cased):
-
-```python
-class LinearWithoutBias(equinox.Module):
-    weight: Any
-
-    def __call__(self, x):
-        return self.weight @ x
-```
-
-If defining a method `meth`, then take care not to write `instance = MyModule(...); jax.jit(instance.meth)(...)`. (Or similarly with `jax.grad`, `equinox.filter_jit` etc.) This is because `instance.meth` is not a pure function as it already has the `self` parameter passed implicitly. Instead do
-```python
-@jax.jit
-def func(instance, args):
-    instance.meth(args)
-    # Also use this pattern with instance(args) if you defined `__call__` instead of `meth`.
-```
-
-### Filtering/combining
-
-Filtering can be used to organise the contents of PyTrees.
-
-```python
-equinox.filter(pytree, filter_spec, inverse=False, replace=None)
-```
-Filters out the leaves of a PyTree not satisfying a condition. Those not satisfying the condition are replaced with `replace`.
-
-- `pytree` is any PyTree
-- `filter_spec` is a PyTree whose structure should be a prefix of the structure of `pytree`. Each of its leaves should either be:
-  - `True`, in which case the leaf or subtree is kept;
-  - `False`, in which case the leaf or subtree is replaced with `replace`;
-  - a callable `Leaf -> bool`, in which case this is evaluted on the leaf or mapped over the subtree, and the leaf kept or replaced as appropriate.
-- `inverse` switches the truthy/falsey behaviour: falsey results are kept and truthy results are replaced.
-- `replace` is what to replace any falsey leaves with. Defaults to `None`.
-
-Returns a PyTree of the same structure as `pytree`.
-
-An important special case is something like `equinox.filter(pytree, equinox.is_array)`. Then `equinox.is_array` is evaluted on all of `pytree`'s leaves, and each leaf then kept or replaced.
-
-See also `equinox.combine` to reconstitute the PyTree again.
-
-```python
-equinox.partition(pytree, filter_spec, replace=None)
-```
-Equivalent to `filter(...), filter(..., inverse=True)`, but slightly more efficient.
-
-```python
-equinox.combine(*pytrees)
-```
-
-Every element of `pytrees` must be a PyTree of the same structure. The return value is also a PyTree of the same structure. Each leaf will be the first non-`None` leaf found in the corresponding leaves of `pytrees`, as they are iterated over. The intention is that this be used to undo a call to `equinox.filter` or `equinox.partition`.
-
-### Filtered transformations
-
-It's very common to need to filter just to handle JAX transformations. Equinox provides the following convenience wrappers.<br>
-They're not designed to handle every edge case -- they're just a way to streamline the common cases. Use separate `equinox.filter`+`jax.jit` etc. if you need finer control.
-
-```python
-equinox.filter_jit(fun, *, filter_spec=is_array, **kwargs)
-```
-Wraps `jax.jit`.
-
-- `fun` is a pure function to JIT compile.
-- `filter_spec` is a PyTree whose structure should be a prefix of the structure of the inputs to `fun`. Each of its leaves should either be `True`, `False`, or a callable `Leaf -> bool`. It behaves exactly as the `filter_spec` argument to `equinox.filter`. Truthy values will be traced; falsey values will be held static. Specifically, if calling `fun(*args, **kwargs)`, then `filter_spec` must have a structure which is a prefix for `(args, kwargs)`.
-- `**kwargs` are any other arguments to `jax.jit`.
-
-An important special case is to pass a function as `filter_spec`, which will be applied to every leaf of every input. For example, `equinox.filter_jit(fun, equinox.is_array)`.
-
-See also `equinox.is_array`, which is the default choice of `filter_spec`. This will trace every JAX array, and make every other argument static.
-
-```python
-equinox.filter_grad(fun, *, filter_spec=is_inexact_array, **kwargs)
-```
-Wraps `jax.grad`.
-
-- `fun` is a pure function to differentiate.
-- `filter_spec` is a PyTree whose structure should be a prefix of the structure of the **first** input to `fun`. Each of its leaves should either be `True`, `False`, or a callable `Leaf -> bool`. It behaves exactly as the `filter_spec` argument to `equinox.filter`. Truthy values will be differentiated; falsey values will not. Specifically, if calling `fun(x, *args, **kwargs)`, then `filter_spec` must have a structure which is a prefix for the structure of `x`.
-- `**kwargs` are any other arguments to `jax.grad`.
-
-An important special case is to pass a function as `filter_spec`, which will be applied to every leaf of the first input. For example, `equinox.filter_grad(fun, equinox.is_inexact_array)`.<br>
-
-See also `equinox.is_inexact_array`, which is the default choice of `filter_spec`. This will differentiate all floating-point JAX arrays.
-
-Note that as the returned gradients must have the same structure as the inputs, then all nondifferentiable components of the input PyTree will have gradient `None`. See `equinox.apply_updates` for a convenience to only apply non-`None` updates.
-
-```python
-equinox.filter_value_and_grad(fun, *, filter_spec=is_inexact_array, **kwargs)
-```
-Wraps `jax.value_and_grad`. Arguments are as `equinox.filter_grad`.
-
-### Filters
-
-Any function `Any -> bool` can be used as a filter. We provide some convenient common choices.
-
-```python
-equinox.is_array(element)
-```
-Returns `True` if `element` is a JAX array (not but a NumPy array).
-
-```python
-equinox.is_array_like(element)
-```
-Returns `True` if `element` is a JAX array, NumPy array, or a Python float/int/bool/complex.
-
-```python
-equinox.is_inexact_array(element)
-```
-Returns `True` if `element` is a floating point JAX array (but not a NumPy array).
-
-```python
-equinox.is_inexact_array_like(element)
-```
-Returns `True` if `element` is a floating point JAX array, floating point NumPy array, or a Python float or complex.
-
-### Utilities
-
-```python
-equinox.apply_updates(model, updates)
-```
-Performs a training update to a model.
-- `model` must be a PyTree;
-- `updates` must be a PyTree with the same structure.
-
-It essentially performs `jax.tree_map(lambda m, u: m + u, model, updates)` (or `optax.apply_upates(model, updates)`). However anywhere `updates` is `None` then no update is made at all, so as to handle nondifferentiable parts of `model`.
-
-The returned value is the updated model. (`model` is not mutated in place, as is usual in JAX and functional programming.)
-
-To produce `updates`, it is typical to take the gradients from the loss function, and then adjust them according to any standard optimiser; for example [Optax](https://github.com/deepmind/optax) provides `optax.sgd` or `optax.adam`.
-
-```python
-equinox.static_field(**kwargs)
-```
-This is a relatively advanced feature. Use it to mark one of the fields of a `Module` as being "static": that is, never differentiated, and always a `static_argnum` to JIT. Best used only if you control whatever will be assigned to that field. For example `equinox.nn.MLP` does *not* use this for its activation function, as in principle a learnt activation function could be passed.
-
-Example:
-```python
-class MyModule(equinox.Module):
-    value: list = equinox.static_field()
-```
-
-If any `**kwargs` are passed, then they will be forwarded on to `dataclasses.field`. (Recall that Equinox uses dataclasses as its modules, so general `dataclasses` behaviour should work as normal.)
-
-```python
-equinox.tree_at(where, pytree, replace=_sentinel, replace_fn=_sentinel)
-```
-Modifies an existing tree, and returns the modified tree. (Like `.at` for "in place modifications" of JAX arrays.)
-
-- `where` is a callable `PyTree -> Leaf` or `PyTree -> Tuple[Leaf, ...]`. It should consume a PyTree of the same shape as `pytree`, and return the leaf or leaves that should be replaced. For example `where=lambda mlp: mlp.layers[-1].linear.weight`.
-- `pytree` is the existing PyTree to modify.
-- `replace` should either be a single element, or a tuple of the same length as returned by `where`. This specifies the replacements to make at the locations specified by `where`. Mutually exclusive with `replace_fn`.
-- `replace_fn` should be a function `Leaf -> Any`. It will be called on every leaf replaced using `where`. The return value from `replace_fn` will be used in its place. Mutually exclusive with `replace`.
-
-For example this can be used to help specify the weights of a model to train or not train:
-```python
-trainable = jax.tree_map(lambda _: False, model)
-trainable = equinox.tree_at(lambda mlp: mlp.layers[-1].linear.weight, model, replace=True)
-equinox.filter_grad(..., filter_spec=trainable)
-```
-
-```python
-equinox.tree_equal(*pytrees)
-```
-Returns `True` if all PyTrees in the list are equal. All arrays must have the same shape, dtype, and values. JAX arrays and NumPy arrays are not considered equal.
-
-### Neural network library
-
-Equinox includes a small neural network library, mostly as a tech demo for how the rest of the library can be used. Its API is broadly modelled after PyTorch.
-
-```python
-equinox.nn.Linear(in_features, out_features, use_bias=True, *, key)(input)
-equinox.nn.Identity(*args, **kwargs)(input)  # args and kwargs are ignored
-equinox.nn.Dropout(p=0.5, deterministic=False)(input, *, key=None, deterministic=None)
-equinox.nn.GRUCell(input_size, hidden_size, use_bias=True, *, key)(input, hidden)
-equinox.nn.LSTMCell(input_size, hidden_size, use_bias=True, *, key)(input, hidden)
-equinox.nn.Sequential(layers)(input, *, key=None)
-equinox.nn.MLP(in_size, out_size, width_size, depth,
-               activation=jax.nn.relu, final_activation=lambda x: x, *, key)(input)
-equinox.nn.Conv(num_spatial_dims, in_channels, out_channels, kernel_size, stride=1,
-                padding=0, dilation=1, use_bias=True, *, key)(input)
-equinox.nn.Conv1d(in_channels, out_channels, kernel_size, stride=1, 
-                  padding=0, dilation=1, use_bias=True, * key)(input)
-equinox.nn.Conv2d(in_channels, out_channels, kernel_size, stride=1, 
-                  padding=0, dilation=1, use_bias=True, * key)(input)
-equinox.nn.Conv3d(in_channels, out_channels, kernel_size, stride=1, 
-                  padding=0, dilation=1, use_bias=True, * key)(input)
-```
-These all behave in the way you expect. The `key` arguments are used to generate the random initial weights, or to generate randomness on the forward pass of stochastic layers like `Dropout`.
-
-The `Dropout(deterministic=...)(deterministic=...)` options determines whether to have the layer act as the identity function, as is commonly done with dropout during inference time. The call-time `deterministic` takes precendence if it passed; otherwise the init-time `deterministic` is used. (Note that because models are PyTrees, you can modify the init-time `deterministic` flag using `equinox.tree_at`. This is perfectly fine, and might be handy if it's easier than using the call-time flag.)
-
-The `MLP(final_activation=...)` option determines any final activation function to apply after the last layer. (In some cases it is desirable for this to be different to the activation used in the main part of the network.)
+(Also consider starring the project on GitHub.)
diff --git a/docs/.citation.md b/docs/.citation.md
new file mode 100644
index 00000000..f4f036f1
--- /dev/null
+++ b/docs/.citation.md
@@ -0,0 +1,12 @@
+If you found this library to be useful in academic work, then please cite:
+
+```bibtex
+@article{kidger2021equinox,
+    author={Patrick Kidger and Cristian Garcia},
+    title={{E}quinox: neural networks in {JAX} via callable {P}y{T}rees and filtered transformations},
+    year={2021},
+    journal={Differentiable Programming workshop at Neural Information Processing Systems 2021}
+}
+```
+
+(Also consider starring the project [on GitHub](https://github.com/patrick-kidger/equinox).)
diff --git a/docs/.htaccess b/docs/.htaccess
new file mode 100644
index 00000000..d47c427b
--- /dev/null
+++ b/docs/.htaccess
@@ -0,0 +1,1 @@
+ErrorDocument 404 /equinox/404.html
diff --git a/docs/_overrides/partials/source.html b/docs/_overrides/partials/source.html
new file mode 100644
index 00000000..f3d57935
--- /dev/null
+++ b/docs/_overrides/partials/source.html
@@ -0,0 +1,20 @@
+{% import "partials/language.html" as lang with context %}
+<a href="{{ config.repo_url }}" title="{{ lang.t('source.link.title') }}" class="md-source" data-md-component="source">
+  <div class="md-source__icon md-icon">
+    {% set icon = config.theme.icon.repo or "fontawesome/brands/git-alt" %}
+    {% include ".icons/" ~ icon ~ ".svg" %}
+  </div>
+  <div class="md-source__repository">
+    {{ config.repo_name }}
+  </div>
+</a>
+{% if config.theme.twitter_url %}
+<a href="{{ config.theme.twitter_url }}" title="Go to Twitter" class="md-source">
+  <div class="md-source__icon md-icon">
+    {% include ".icons/fontawesome/brands/twitter.svg" %}
+  </div>
+  <div class="md-source__repository">
+    {{ config.theme.twitter_name }}
+  </div>
+</a>
+{% endif %}
diff --git a/docs/_static/.README.md b/docs/_static/.README.md
new file mode 100644
index 00000000..010a7c92
--- /dev/null
+++ b/docs/_static/.README.md
@@ -0,0 +1,2 @@
+The favicon is `math-integral` from https://materialdesignicons.com, found by way of https://pictogrammers.com.
+(The logo is `math-integral-box`.)
diff --git a/docs/_static/custom_css.css b/docs/_static/custom_css.css
new file mode 100644
index 00000000..37ae5ecc
--- /dev/null
+++ b/docs/_static/custom_css.css
@@ -0,0 +1,128 @@
+/* Fit the Twitter handle alongside the GitHub one in the top right. */
+
+div.md-header__source {
+    width: revert;
+    max-width: revert;
+}
+
+a.md-source {
+    display: inline-block;
+}
+
+.md-source__repository {
+    max-width: 100%;
+}
+
+/* Emphasise sections of nav on left hand side */
+
+nav.md-nav {
+  padding-left: 5px;
+}
+
+nav.md-nav--secondary {
+    border-left: revert !important;
+}
+
+.md-nav__title {
+  font-size: 0.9rem;
+}
+
+.md-nav__item--section > .md-nav__link {
+  font-size: 0.9rem;
+}
+
+/* Indent autogenerated documentation */
+
+div.doc-contents {
+  padding-left: 25px;
+  border-left: 4px solid rgba(230, 230, 230);
+}
+
+/* Increase visibility of splitters "---" */
+
+[data-md-color-scheme="default"] .md-typeset hr {
+    border-bottom-color: rgb(0, 0, 0);
+    border-bottom-width: 1pt;
+}
+
+[data-md-color-scheme="slate"] .md-typeset hr {
+    border-bottom-color: rgb(230, 230, 230);
+}
+
+/* More space at the bottom of the page */
+
+.md-main__inner {
+  margin-bottom: 1.5rem;
+}
+
+/* Remove prev/next footer buttons */
+
+.md-footer__inner {
+    display: none;
+}
+
+/* Change font sizes */
+
+html {
+    /* Decrease font size for overall webpage
+       Down from 137.5% which is the Material default */
+    font-size: 110%;
+}
+
+.md-typeset .admonition {
+    /* Increase font size in admonitions */
+    font-size: 100% !important;
+}
+
+.md-typeset details {
+    /* Increase font size in details */
+    font-size: 100% !important;
+}
+
+.md-typeset h1 {
+    font-size: 1.6rem;
+}
+
+.md-typeset h2 {
+    font-size: 1.5rem;
+}
+
+.md-typeset h3 {
+    font-size: 1.3rem;
+}
+
+.md-typeset h4 {
+    font-size: 1.1rem;
+}
+
+.md-typeset h5 {
+    font-size: 0.9rem;
+}
+
+.md-typeset h6 {
+    font-size: 0.8rem;
+}
+
+/* Bugfix: remove the superfluous parts generated when doing:
+
+??? Blah
+
+    ::: library.something
+*/
+
+.md-typeset details .mkdocstrings > h4 {
+    display: none;
+}
+
+.md-typeset details .mkdocstrings > h5 {
+    display: none;
+}
+
+/* Change default colours for <a> tags */
+
+[data-md-color-scheme="default"] {
+    --md-typeset-a-color: rgb(0, 189, 164) !important;
+}
+[data-md-color-scheme="slate"] {
+    --md-typeset-a-color: rgb(0, 189, 164) !important;
+}
diff --git a/docs/_static/favicon.png b/docs/_static/favicon.png
new file mode 100755
index 00000000..8cf632a5
Binary files /dev/null and b/docs/_static/favicon.png differ
diff --git a/docs/all-of-equinox.md b/docs/all-of-equinox.md
new file mode 100644
index 00000000..0fa34ee8
--- /dev/null
+++ b/docs/all-of-equinox.md
@@ -0,0 +1,150 @@
+# All of Equinox
+
+Equinox is a very small and easy to understand library. (Because it uses JAX abstractions like PyTrees, it doesn't need any real complexity.)
+
+So as the title suggests, this page tells you essentially everything you need to know to use Equinox.
+
+## Parameterised functions as PyTrees
+
+As we saw on the [Getting Started](./index.md) page, Equinox represents parameterised functions as PyTrees.
+
+!!! example
+
+    A neural network is a function parameterised by its weights, biases, etc.
+
+    (But you can use Equinox to represent any kind of parameterised function.)
+
+And now you can JIT/grad/etc. with respect to your model. For example, using a few built-in layers by way of demonstration:
+
+```python
+import equinox as eqx
+import jax
+
+class MyModule(eqx.Module):
+    layers: list
+    bias: jax.numpy.ndarray
+    
+    def __init__(self, key):
+        key1, key2, key3 = jax.random.split(key, 3)
+        self.layers = [eqx.nn.Linear(2, 8, key=key1),
+                       eqx.nn.Linear(8, 8, key=key2),
+                       eqx.nn.Linear(8, 2, key=key3)]
+        self.bias = jax.numpy.ones(2)
+
+    def __call__(self, x):
+        for layer in self.layers[:-1]:
+            x = jax.nn.relu(layer(x))
+        return self.layers[-1](x) + self.bias
+
+@jax.jit
+@jax.grad
+def loss(model, x, y):
+    pred_y = jax.vmap(model)(x)
+    return jax.numpy.mean((y - pred_y) ** 2)
+
+x_key, y_key, model_key = jax.random.split(jax.random.PRNGKey(0), 3)
+x = jax.random.normal(x_key, (100, 2))
+y = jax.random.normal(y_key, (100, 2))
+model = MyModule(model_key)
+grads = loss(model, x, y)
+learning_rate = 0.1
+model = jax.tree_map(lambda m, g: m - learning_rate * g, model, grads)
+```
+
+## Filtering
+
+In the previous example, all of the model attributes were `Module`s and JAX arrays.
+
+But Equinox supports using arbitrary Python objects too! If you choose to include those in your models, then Equinox offers the tools to handle them appropriately around `jax.jit` and `jax.grad`. (Which themselves only really work with JAX arrays.)
+
+!!! example
+
+    You might have a `bool`-ean flag in your model-as-a-PyTree, specifying whether to enable some extra piece of behaviour. You might want to treat that as a `static_argnum` to `jax.jit`.
+
+!!! example
+
+    The activation function in a multi-layer perceptron isn't a JAX array -- it's an arbitrary Python function -- but you may reasonably wish to make that part of your model's parameterisation. (Rather than just baking in a particular choice into your forward pass.)
+
+If you want to do this, then Equinox offers *filtering*, like follows.
+
+```python
+import equinox as eqx
+import functools as ft
+import jax
+
+# Start off by crreating a model just like normal, but with some arbitrary Python
+# objects as part of its parameterisation. In this case, we have `jax.nn.relu`, which
+# is a Python function.
+class AnotherModule(eqx.Module):
+    layers: list
+
+    def __init__(self, key):
+        key1, key2 = jax.random.split(key)
+        self.layers = [eqx.nn.Linear(2, 8, key=key1),
+                       jax.nn.relu,
+                       eqx.nn.Linear(8, 2, key=key2)]
+
+    def __call__(self, x):
+        for layer in self.layers:
+            x = layer(x)
+        return x
+
+x_key, y_key, model_key = jax.random.split(jax.random.PRNGKey(0), 3)
+x, y = jax.random.normal(x_key, (100, 2)), jax.random.normal(y_key, (100, 2))
+model = AnotherModule(model_key)
+
+# Option 1: manually filter out anything that isn't JIT/grad-able.
+
+@ft.partial(jax.jit, static_argnums=1)
+@jax.grad
+def loss(params, static, x, y):
+    model = eqx.combine(params, static)
+    pred_y = jax.vmap(model)(x)
+    return jax.numpy.mean((y - pred_y) ** 2)
+
+params, static = eqx.partition(model, eqx.is_array)
+loss(params, static, x, y)
+
+# Option 2: use filtered transformations, which automates the above process for you.
+
+@eqx.filter_jit
+@eqx.filter_grad
+def loss(model, x, y):
+    pred_y = jax.vmap(model)(x)
+    return jax.numpy.mean((y - pred_y) ** 2)
+
+loss(model, x, y)
+```
+
+Here, `params` and `static` are both instances of `AnotherModule`: `params` keeps just the leaves that are JAX arrays; `static` keeps everything else. Then `combine` merges the two PyTrees back together after crossing the `jax.jit` and `jax.grad` API boundaries.
+
+The choice of `eqx.is_array` is a *filter function*: a boolean function specifying whether each leaf should go into `params` or into `static`. In this case very simply `eqx.is_array(x) == isinstance(x, jax.numpy.ndarray)`.
+
+As a convenience, `eqx.filter_jit` and `eqx.filter_grad` just wrap filtering and transformation together. It turns out to be really common to only need to filter around JAX transformations.
+
+## Integrates smoothly with JAX
+
+Equinox introduces a powerful yet straightforward way to build neural networks, without introducing lots of new notions or tieing you into a framework.
+
+Equinox is all just regular JAX -- PyTrees and transformations. Together, these two pieces allow us to specify complex models in JAX-friendly ways.
+
+## Next steps
+
+And that's it! That's pretty much everything you need to know about Equinox. Everything you've seen so far should be enough to get started with using the library. Also see the [Train RNN](./examples/train_rnn.ipynb) example for a fully worked example.
+
+## Summary
+
+Equinox includes four main things:
+
+- For building models: `equinox.Module`.
+- Prebuilt neural network layers: `equinox.nn.Linear`, `equinox.nn.Conv2d`, etc.
+- Filtering, and filtered transformations: `equinox.filter`, `equinox.filter_jit` etc.
+- Some utilities to help manipulate PyTrees: `equinox.tree_at` etc.
+
+See also the API reference on the left.
+
+!!! faq "FAQ"
+
+    One common question that gets asked: a lot of other libraries introduce custom `library.jit` etc. operations, specifically to work with `library.Module`. What makes the filtered transformations of Equinox different?
+
+    The answer is that filters are tools that apply to *any PyTree*; and models just happen to be PyTrees. There is no special coupling between a filtered transformation and `eqx.Module`. (Not to mention that the purpose of filtered transformations -- being able to include non-JAX-arrays in your model's parameters -- is itself unusual/impossible in many other libraries.)
diff --git a/docs/api/filtering/filter-functions.md b/docs/api/filtering/filter-functions.md
new file mode 100644
index 00000000..850fec6a
--- /dev/null
+++ b/docs/api/filtering/filter-functions.md
@@ -0,0 +1,17 @@
+# Filter functions
+
+These are some built-in common ways of specifying *how* to split a PyTree into pieces. (For example, JAX array vs non-JAX-array.)
+
+::: equinox.is_array
+
+---
+
+::: equinox.is_array_like
+
+---
+
+::: equinox.is_inexact_array
+
+---
+
+::: equinox.is_inexact_array_like
diff --git a/docs/api/filtering/filter-operations.md b/docs/api/filtering/filter-operations.md
new file mode 100644
index 00000000..edb1d3d7
--- /dev/null
+++ b/docs/api/filtering/filter-operations.md
@@ -0,0 +1,13 @@
+# Filter, partition, combine
+
+These are the functions that actually split a PyTree into pieces, or combine several PyTrees back into a PyTree.
+
+::: equinox.filter
+
+---
+
+::: equinox.partition
+
+---
+
+::: equinox.combine
diff --git a/docs/api/filtering/filtered-transformations.md b/docs/api/filtering/filtered-transformations.md
new file mode 100644
index 00000000..22a3e7b3
--- /dev/null
+++ b/docs/api/filtering/filtered-transformations.md
@@ -0,0 +1,21 @@
+# Filtered transformations
+
+These typically combine [`equinox.partition`][], a filter function, and a JAX transformation, all together.
+
+Practically speaking these are usually the only kind of filtering you ever have to use. (But it's good to understand what e.g. [`equinox.partition`][] and [`equinox.is_array`][] are doing under the hood, just so that these don't seem too magical.)
+
+::: equinox.filter_jit
+
+---
+
+::: equinox.filter_grad
+
+---
+
+::: equinox.filter_value_and_grad
+
+---
+
+::: equinox.filter_custom_vjp
+    selection:
+        members: false
diff --git a/docs/api/module.md b/docs/api/module.md
new file mode 100644
index 00000000..478f3625
--- /dev/null
+++ b/docs/api/module.md
@@ -0,0 +1,5 @@
+# Module
+
+::: equinox.Module
+    selection:
+        members: false
diff --git a/docs/api/nn/composed.md b/docs/api/nn/composed.md
new file mode 100644
index 00000000..2fb89644
--- /dev/null
+++ b/docs/api/nn/composed.md
@@ -0,0 +1,15 @@
+# Composites
+
+::: equinox.nn.Sequential
+    selection:
+        members:
+            - __init__
+            - __call__
+
+---
+
+::: equinox.nn.MLP
+    selection:
+        members:
+            - __init__
+            - __call__
diff --git a/docs/api/nn/conv.md b/docs/api/nn/conv.md
new file mode 100644
index 00000000..8f4d90de
--- /dev/null
+++ b/docs/api/nn/conv.md
@@ -0,0 +1,24 @@
+# Convolutional layers
+::: equinox.nn.Conv
+    selection:
+        members:
+            - __init__
+            - __call__
+
+---
+
+::: equinox.nn.Conv1d
+    selection:
+        members: false
+
+---
+
+::: equinox.nn.Conv2d
+    selection:
+        members: false
+
+---
+
+::: equinox.nn.Conv3d
+    selection:
+        members: false
diff --git a/docs/api/nn/dropout.md b/docs/api/nn/dropout.md
new file mode 100644
index 00000000..a5f14a2e
--- /dev/null
+++ b/docs/api/nn/dropout.md
@@ -0,0 +1,7 @@
+# Dropout
+
+::: equinox.nn.Dropout
+    selection:
+        members:
+            - __init__
+            - __call__
diff --git a/docs/api/nn/linear.md b/docs/api/nn/linear.md
new file mode 100644
index 00000000..0c988fee
--- /dev/null
+++ b/docs/api/nn/linear.md
@@ -0,0 +1,15 @@
+# Linear layers
+
+::: equinox.nn.Linear
+    selection:
+        members:
+            - __init__
+            - __call__
+
+---
+
+::: equinox.nn.Identity
+    selection:
+        members:
+            - __init__
+            - __call__
diff --git a/docs/api/nn/rnn.md b/docs/api/nn/rnn.md
new file mode 100644
index 00000000..42e93755
--- /dev/null
+++ b/docs/api/nn/rnn.md
@@ -0,0 +1,16 @@
+# Recurrent layers
+
+::: equinox.nn.GRUCell
+    selection:
+        members:
+            - __init__
+            - __call__
+
+---
+
+::: equinox.nn.LSTMCell
+    selection:
+        members:
+        members:
+            - __init__
+            - __call__
diff --git a/docs/api/utilities.md b/docs/api/utilities.md
new file mode 100644
index 00000000..a79e0901
--- /dev/null
+++ b/docs/api/utilities.md
@@ -0,0 +1,15 @@
+# Utilities
+
+::: equinox.tree_at
+
+---
+
+::: equinox.tree_equal
+
+---
+
+::: equinox.apply_updates
+
+---
+
+::: equinox.static_field
diff --git a/docs/citation.md b/docs/citation.md
new file mode 100644
index 00000000..e3ae3010
--- /dev/null
+++ b/docs/citation.md
@@ -0,0 +1,3 @@
+# Citation
+
+--8<-- ".citation.md"
diff --git a/docs/examples b/docs/examples
new file mode 120000
index 00000000..a6573af9
--- /dev/null
+++ b/docs/examples
@@ -0,0 +1,1 @@
+../examples
\ No newline at end of file
diff --git a/docs/index.md b/docs/index.md
new file mode 100644
index 00000000..0ef8aefa
--- /dev/null
+++ b/docs/index.md
@@ -0,0 +1,68 @@
+# Getting started
+
+Equinox is a JAX library based around a simple idea: **represent parameterised functions (such as neural networks) as PyTrees.**
+
+In doing so:
+
+- We get a PyTorch-like API...
+- ...that's fully compatible with *native* JAX transformations...
+- ...with no new concepts you have to learn. (It's all just PyTrees.)
+
+The elegance of Equinox is its selling point in a world that already has [Haiku](https://github.com/deepmind/dm-haiku), [Flax](https://github.com/google/flax) and so on.
+
+_(In other words, why should you care? Because Equinox is really simple to learn, and really simple to use.)_
+
+## Installation
+
+```bash
+pip install equinox
+```
+
+Requires Python 3.7+ and JAX 0.2.18+.
+
+## Quick example
+
+Models are defined using PyTorch-like syntax:
+
+```python
+import equinox as eqx
+import jax
+
+class Linear(eqx.Module):
+    weight: jax.numpy.ndarray
+    bias: jax.numpy.ndarray
+
+    def __init__(self, in_size, out_size, key):
+        wkey, bkey = jax.random.split(key)
+        self.weight = jax.random.normal(wkey, (out_size, in_size))
+        self.bias = jax.random.normal(bkey, (out_size,))
+
+    def __call__(self, x):
+        return self.weight @ x + self.bias
+```
+
+and fully compatible with normal JAX operations:
+
+```python
+@jax.jit
+@jax.grad
+def loss_fn(model, x, y):
+    pred_y = jax.vmap(model)(x)
+    return jax.numpy.mean((y - pred_y) ** 2)
+
+batch_size, in_size, out_size = 32, 2, 3
+model = Linear(in_size, out_size, key=jax.random.PRNGKey(0))
+x = jax.numpy.zeros((batch_size, in_size))
+y = jax.numpy.zeros((batch_size, out_size))
+grads = loss_fn(model, x, y)
+```
+
+Finally, there's no magic behind the scenes. All `eqx.Module` does is register your class as a PyTree. From that point onwards, JAX already knows how to work with PyTrees.
+
+## Next steps
+
+If this quick start has got you interested, then have a read of [All of Equinox](./all-of-equinox.md), which introduces you to basically everything in Equinox. (Doesn't take very long! Equinox is simple because everything is a PyTree.)
+
+## Citation
+
+--8<-- ".citation.md"
diff --git a/docs/requirements.txt b/docs/requirements.txt
new file mode 100644
index 00000000..88af1fdf
--- /dev/null
+++ b/docs/requirements.txt
@@ -0,0 +1,12 @@
+# Latest versions at time of writing.
+mkdocs==1.2.3            # Main documentation generator.
+mkdocs-material==7.3.6   # Theme
+pymdown-extensions==9.0  # Markdown extensions e.g. to handle LaTeX.
+mkdocstrings==0.16.2     # Autogenerate documentation from docstrings.
+mknotebooks==0.7.1       # Turn Jupyter Lab notebooks into webpages.
+pytkdocs_tweaks==0.0.1   # Tweaks mkdocstrings to improve various aspects
+mkdocs_include_exclude_files==0.0.1  # Tweak which files are included/excluded
+
+# Dependencies of Equinox itself.
+# Always use most up-to-date versions.
+jax[cpu]
diff --git a/equinox/__init__.py b/equinox/__init__.py
index 4e135926..ce8ab2ac 100644
--- a/equinox/__init__.py
+++ b/equinox/__init__.py
@@ -10,11 +10,17 @@
     partition,
     split,
 )
-from .gradf import filter_grad, filter_value_and_grad, gradf, value_and_grad_f
-from .jitf import filter_jit, jitf
+from .grad import (
+    filter_custom_vjp,
+    filter_grad,
+    filter_value_and_grad,
+    gradf,
+    value_and_grad_f,
+)
+from .jit import filter_jit, jitf
 from .module import Module, static_field
 from .tree import tree_at, tree_equal
 from .update import apply_updates
 
 
-__version__ = "0.1.3"
+__version__ = "0.1.4"
diff --git a/equinox/custom_types.py b/equinox/custom_types.py
index 0a955d27..38c67736 100644
--- a/equinox/custom_types.py
+++ b/equinox/custom_types.py
@@ -1,17 +1,15 @@
-from typing import Any, Union
+import typing
+from typing import Any
 
 import jax
 import jax.numpy as jnp
-import jaxlib
-import numpy as np
 
 
-# JAX arrays
-Array = Union[jax.core.Tracer, jaxlib.xla_extension.DeviceArray]
-# JAX arrays + numpy arrays
-# jnp.ndarray counts as a numpy array: `assert isinstance(np.array(1.), jnp.ndarray)`
-MoreArrays = Union[Array, np.ndarray, jnp.ndarray]
+if getattr(typing, "GENERATING_DOCUMENTATION", False):
+    Array = "jax.numpy.ndarray"
+else:
+    Array = jnp.ndarray
 
 PyTree = Any
 
-TreeDef = jaxlib.xla_extension.PyTreeDef
+TreeDef = type(jax.tree_structure(0))
diff --git a/equinox/deprecated.py b/equinox/deprecated.py
index 91857261..b5de16ff 100644
--- a/equinox/deprecated.py
+++ b/equinox/deprecated.py
@@ -3,8 +3,11 @@
 
 
 def deprecated(*, in_favour_of):
+    if not isinstance(in_favour_of, str):
+        in_favour_of = in_favour_of.__name__
+
     def decorator(fn):
-        msg = f"{fn.__name__} is deprecated in favour of {in_favour_of.__name__}"
+        msg = f"{fn.__name__} is deprecated in favour of {in_favour_of}"
 
         @ft.wraps(fn)
         def wrapper(*args, **kwargs):
diff --git a/equinox/filters.py b/equinox/filters.py
index 5e324fac..8ddec1e6 100644
--- a/equinox/filters.py
+++ b/equinox/filters.py
@@ -1,10 +1,10 @@
 from typing import Any, Callable, List, Optional, Tuple, Union
-from typing_extensions import get_args
 
 import jax
 import jax.numpy as jnp
+import numpy as np
 
-from .custom_types import Array, MoreArrays, PyTree, TreeDef
+from .custom_types import PyTree, TreeDef
 from .deprecated import deprecated
 
 
@@ -12,29 +12,37 @@
 # Filter functions
 #
 
-_array_types = get_args(Array)
-_morearray_types = get_args(MoreArrays)
-_arraylike_types = _morearray_types + (int, float, complex, bool)
 
-
-# TODO: not sure if this is the best way to do this? In light of:
-# https://github.com/google/jax/commit/258ae44303b1539eff6253263694ec768b8803f0#diff-de759f969102e9d64b54a299d11d5f0e75cfe3052dc17ffbcd2d43b250719fb0
 def is_array(element: Any) -> bool:
-    return isinstance(element, _array_types)
+    """Returns `True` if `element` is a JAX array (but not a NumPy array)."""
+    return isinstance(element, jnp.ndarray)
 
 
 # Does _not_ do a try/except on jnp.asarray(element) because that's very slow.
+# Chosen to match
+# https://github.com/google/jax/blob/4a17c78605e7fc69a69a999e2f6298db79d3837a/jax/_src/numpy/lax_numpy.py#L542  # noqa: E501
 def is_array_like(element: Any) -> bool:
-    return isinstance(element, _arraylike_types)
+    """Returns `True` if `element` is a JAX array, a NumPy array, or a Python
+    `float`/`complex`/`bool`/`int`.
+    """
+    return isinstance(
+        element, (jnp.ndarray, np.ndarray, float, complex, bool, int)
+    ) or hasattr(element, "__jax_array__")
 
 
 def is_inexact_array(element: Any) -> bool:
+    """Returns `True` if `element` is an inexact (i.e. floating point) JAX array."""
     return is_array(element) and jnp.issubdtype(element.dtype, jnp.inexact)
 
 
 def is_inexact_array_like(element: Any) -> bool:
+    """Returns `True` if `element` is an inexact JAX array, an inexact NumPy array, or
+    a Python `float` or `complex`.
+    """
+    if hasattr(element, "__jax_array__"):
+        element = element.__jax_array__()
     return (
-        isinstance(element, _morearray_types)
+        isinstance(element, (jnp.ndarray, np.ndarray))
         and jnp.issubdtype(element.dtype, jnp.inexact)
     ) or isinstance(element, (float, complex))
 
@@ -56,6 +64,37 @@ def _make_filter_tree(mask: Union[bool, Callable[[Any], bool]], arg: Any) -> boo
 def filter(
     pytree: PyTree, filter_spec: PyTree, inverse: bool = False, replace: Any = None
 ) -> PyTree:
+    """
+    Filters out the leaves of a PyTree not satisfying a condition. Those not satisfying
+    the condition are replaced with `replace`.
+
+    **Arguments:**
+
+    - `pytree` is any PyTree.
+    - `filter_spec` is a PyTree whose structure should be a prefix of the structure of
+        `pytree`. Each of its leaves should either be:
+        - `True`, in which case the leaf or subtree is kept;
+        - `False`, in which case the leaf or subtree is replaced with `replace`;
+        - a callable `Leaf -> bool`, in which case this is evaluted on the leaf or
+            mapped over the subtree, and the leaf kept or replaced as appropriate.
+    - `inverse` switches the truthy/falsey behaviour: falsey results are kept and
+        truthy results are replaced.
+    - `replace` is what to replace any falsey leaves with. Defaults to `None`.
+
+    **Returns:**
+
+    A PyTree of the same structure as `pytree`.
+
+    !!! info
+
+        An common special case is `equinox.filter(pytree, equinox.is_array)`. Then
+        `equinox.is_array` is evaluted on all of `pytree`'s leaves, and each leaf then
+        kept or replaced.
+
+    !!! info
+
+        See also [`equinox.combine`][] to reconstitute the PyTree again.
+    """
 
     inverse = bool(inverse)  # just in case, to make the != trick below work reliably
     filter_tree = jax.tree_map(_make_filter_tree, filter_spec, pytree)
@@ -65,6 +104,9 @@ def filter(
 
 
 def partition(pytree: PyTree, filter_spec: PyTree, replace: Any = None) -> PyTree:
+    """Equivalent to `filter(...), filter(..., inverse=True)`, but slightly more
+    efficient.
+    """
 
     filter_tree = jax.tree_map(_make_filter_tree, filter_spec, pytree)
     left = jax.tree_map(lambda mask, x: x if mask else replace, filter_tree, pytree)
@@ -83,7 +125,33 @@ def _is_none(x):
     return x is None
 
 
-def combine(*pytrees: PyTree):
+def combine(*pytrees: PyTree) -> PyTree:
+    """Combines multiple PyTrees into one PyTree, by replacing `None` leaves.
+
+    !!! example
+
+        ```python
+        pytree1 = [None, 1, 2]
+        pytree2 = [0, None, None]
+        equinox.combine(pytree1, pytree2)  # [0, 1, 2]
+        ```
+
+    !!! tip
+
+        The idea is that `equinox.combine` should be used to undo a call to
+        [`equinox.filter`][] or [`equinox.partition`][].
+
+    **Arguments:**
+
+        - `*pytrees*: a sequence of PyTrees all with the same structure.
+
+    **Returns:**
+
+    A PyTree with the same structure as its inputs. Each leaf will be the first
+    non-`None` leaf found in the corresponding leaves of `pytrees` as they are
+    iterated over.
+    """
+
     return jax.tree_map(_combine, *pytrees, is_leaf=_is_none)
 
 
diff --git a/equinox/grad.py b/equinox/grad.py
new file mode 100644
index 00000000..f4b92b9b
--- /dev/null
+++ b/equinox/grad.py
@@ -0,0 +1,304 @@
+import functools as ft
+import typing
+
+import jax
+
+from .deprecated import deprecated
+from .filters import (
+    combine,
+    is_array,
+    is_inexact_array,
+    merge,
+    partition,
+    split,
+    validate_filters,
+)
+
+
+def filter_value_and_grad(
+    fun, *, filter_spec=is_inexact_array, argnums=None, **gradkwargs
+):
+    """As [`equinox.filter_grad`][], except that it is `jax.value_and_grad` that is
+    wrapped.
+    """
+    if argnums is not None:
+        raise ValueError(
+            "`argnums` should not be passed. If you need to differentiate "
+            "multiple objects then collect them into a tuple and pass that "
+            "as the first argument."
+        )
+
+    @ft.partial(jax.value_and_grad, argnums=0, **gradkwargs)
+    def fun_value_and_grad(diff_x, nondiff_x, *args, **kwargs):
+        x = combine(diff_x, nondiff_x)
+        return fun(x, *args, **kwargs)
+
+    def fun_value_and_grad_wrapper(x, *args, **kwargs):
+        diff_x, nondiff_x = partition(x, filter_spec)
+        return fun_value_and_grad(diff_x, nondiff_x, *args, **kwargs)
+
+    return fun_value_and_grad_wrapper
+
+
+def filter_grad(fun, *, filter_spec=is_inexact_array, **gradkwargs):
+    """Wraps together [`equinox.partition`][] and `jax.grad`.
+
+    **Arguments:**
+
+    - `fun` is a pure function to JIT compile.
+    - `filter_spec` is a PyTree whose structure should be a prefix of the structure of
+        the **first** argument to `fun`. It behaves as the `filter_spec` argument to
+        [`equinox.filter`][]. Truthy values will be differentiated; falsey values will
+        not.
+    - `**gradkwargs` are any other keyword arguments to `jax.grad`.
+
+    **Returns:**
+
+    A function computing the derivative of `fun` with respect to its first input. Any
+    nondifferentiable leaves will have `None` as the gradient. See
+    `equinox.apply_updates` for a convenience function that will only attempt to apply
+    non-`None` updates.
+
+    !!! info
+
+        A very important special case is to trace all inexact (i.e. floating point)
+        JAX arrays and treat all other objects as nondifferentiable.
+
+        This is accomplished with `filter_spec=equinox.is_inexact_array`, which is the
+        default.
+
+    !!! tip
+
+        If you need to differentiate multiple objects, then put them together into a
+        tuple and pass that through the first argument:
+        ```python
+        # We want to differentiate `func` with respect to both `x` and `y`.
+        def func(x, y):
+            ...
+
+        @equinox.filter_grad
+        def grad_func(x__y):
+            x, y = x__y
+            return func(x, y)
+        ```
+    """
+
+    has_aux = gradkwargs.get("has_aux", False)
+
+    fun_value_and_grad = filter_value_and_grad(
+        fun, filter_spec=filter_spec, **gradkwargs
+    )
+
+    def fun_grad(*args, **kwargs):
+        value, grad = fun_value_and_grad(*args, **kwargs)
+        if has_aux:
+            value, aux = value
+            return aux, grad
+        else:
+            return grad
+
+    return fun_grad
+
+
+class filter_custom_vjp:
+    """Provides an easier API for `jax.custom_vjp`, by using filtering.
+
+    Usage is:
+    ```python
+    @equinox.filter_custom_vjp
+    def fn(vjp_arg, *args, **kwargs):
+        # vjp_arg is some PyTree of arbitrary Python objects.
+        # args, kwargs contain arbitrary Python objects.
+        ...
+        return obj  # some PyTree of arbitrary Python objects.
+
+    def fn_fwd(vjp_arg, *args, **kwargs):
+        ...
+        # Should return `obj` as before. `residuals` can be any collection of JAX
+        # arrays you want to keep around for the backward pass.
+        return obj, residuals
+
+    def fn_bwd(residuals, grad_obj, vjp_arg, *args, **kwargs):
+        # grad_obj will have `None` as the gradient for any leaves of `obj` that were
+        # not JAX arrays
+        ...
+        # grad_vjp_arg should have `None` as the gradient for any leaves of `vjp_arg`
+        # that were not JAX arrays.
+        return grad_vjp_arg
+
+    fn.defvjp(fn_fwd, fn_bwd)
+    ```
+
+    The key differences to `jax.custom_vjp` are that:
+
+    - Only the gradient of the first argument, `vjp_arg`, should be computed on the
+        backward pass. Everything else will automatically have zero gradient.
+    - You do not need to distinguish differentiable from nondifferentiable manually.
+        Instead you should return gradients for all inexact JAX arrays in the first
+        argument. (And just put `None` on every other leaf of the PyTree.)
+    - As a convenience, all of the inputs from the forward pass are additionally made
+        available to you on the backward pass.
+
+    !!! tip
+
+        If you need gradients with respect to multiple arguments, then just pack them
+        together as a tuple via the first argument `vjp_arg`. (See also
+        [`equinox.filter_grad`][] for a similar trick.)
+    """
+
+    def __init__(self, fn):
+        self.fn = fn
+        self.fn_wrapped = None
+
+    def defvjp(self, fn_fwd, fn_bwd):
+        def fn_wrapped(
+            nonarray_vjp_arg,
+            nonarray_args_kwargs,
+            diff_array_vjp_arg,
+            nondiff_array_vjp_arg,
+            array_args_kwargs,
+        ):
+            vjp_arg = combine(
+                nonarray_vjp_arg, diff_array_vjp_arg, nondiff_array_vjp_arg
+            )
+            args, kwargs = combine(nonarray_args_kwargs, array_args_kwargs)
+            return self.fn(vjp_arg, *args, **kwargs)
+
+        def fn_fwd_wrapped(
+            nonarray_vjp_arg,
+            nonarray_args_kwargs,
+            diff_array_vjp_arg,
+            nondiff_array_vjp_arg,
+            array_args_kwargs,
+        ):
+            vjp_arg = combine(
+                nonarray_vjp_arg, diff_array_vjp_arg, nondiff_array_vjp_arg
+            )
+            args, kwargs = combine(nonarray_args_kwargs, array_args_kwargs)
+            out, residuals = fn_fwd(vjp_arg, *args, **kwargs)
+            return out, (
+                residuals,
+                diff_array_vjp_arg,
+                nondiff_array_vjp_arg,
+                array_args_kwargs,
+            )
+
+        def fn_bwd_wrapped(nonarray_vjp_arg, nonarray_args_kwargs, residuals, grad_out):
+            (
+                residuals,
+                diff_array_vjp_arg,
+                nondiff_array_vjp_arg,
+                array_args_kwargs,
+            ) = residuals
+            vjp_arg = combine(
+                nonarray_vjp_arg, diff_array_vjp_arg, nondiff_array_vjp_arg
+            )
+            args, kwargs = combine(nonarray_args_kwargs, array_args_kwargs)
+            out = fn_bwd(residuals, grad_out, vjp_arg, *args, **kwargs)
+            if jax.tree_structure(out) != jax.tree_structure(diff_array_vjp_arg):
+                raise RuntimeError(
+                    "custom_vjp gradients must have the same structure as "
+                    "`equinox.filter(vjp_arg, equinox.is_inexact_array)`, where "
+                    "`vjp_arg` is the first argument used in the forward pass."
+                )
+            # None is the gradient through nondiff_array_vjp_arg and array_args_kwargs
+            return out, None, None
+
+        fn_wrapped = jax.custom_vjp(fn_wrapped, nondiff_argnums=(0, 1))
+        fn_wrapped.defvjp(fn_fwd_wrapped, fn_bwd_wrapped)
+        self.fn_wrapped = fn_wrapped
+
+    def __call__(self, vjp_arg, /, *args, **kwargs):
+        if self.fn_wrapped is None:
+            raise RuntimeError(f"defvjp not yet called for {self.fn.__name__}")
+        array_vjp_arg, nonarray_vjp_arg = partition(vjp_arg, is_array)
+        diff_array_vjp_arg, nondiff_array_vjp_arg = partition(
+            array_vjp_arg, is_inexact_array
+        )
+        array_args_kwargs, nonarray_args_kwargs = partition((args, kwargs), is_array)
+        return self.fn_wrapped(
+            nonarray_vjp_arg,
+            nonarray_args_kwargs,
+            diff_array_vjp_arg,
+            nondiff_array_vjp_arg,
+            array_args_kwargs,
+        )
+
+
+if getattr(typing, "GENERATING_DOCUMENTATION", False):
+    _filter_custom_vjp_doc = filter_custom_vjp.__doc__
+
+    def filter_custom_vjp(fn):
+        pass
+
+    filter_custom_vjp.__doc__ = _filter_custom_vjp_doc
+
+
+#
+# Deprecated
+#
+
+
+@deprecated(in_favour_of=filter_value_and_grad)
+def value_and_grad_f(fun, *, filter_fn=None, filter_tree=None, argnums=0, **gradkwargs):
+    if isinstance(argnums, int):
+        unwrap = True
+        argnums = (argnums,)
+        if filter_tree is not None:
+            filter_tree = (filter_tree,)
+    else:
+        unwrap = False
+
+    validate_filters("value_and_grad_f", filter_fn, filter_tree)
+
+    @ft.partial(jax.value_and_grad, argnums=argnums, **gradkwargs)
+    def f_value_and_grad(*args, **kwargs):
+        *args, notes = args
+        args = list(args)
+        for i, (arg_nograd, which, treedef) in notes.items():
+            arg_grad = args[i]
+            arg = merge(arg_grad, arg_nograd, which, treedef)
+            args[i] = arg
+        return fun(*args, **kwargs)
+
+    def f_value_and_grad_wrapper(*args, **kwargs):
+        args = list(args)
+        notes = {}
+        for j, i in enumerate(argnums):
+            arg = args[i]
+            if filter_fn is None:
+                # implies filter_tree is not None
+                arg_grad, arg_nograd, which, treedef = split(
+                    arg, filter_tree=filter_tree[j]
+                )
+            else:
+                arg_grad, arg_nograd, which, treedef = split(arg, filter_fn=filter_fn)
+            args[i] = arg_grad
+            notes[i] = (arg_nograd, which, treedef)
+        value, grad = f_value_and_grad(*args, notes, **kwargs)
+        grad = list(grad)
+        for j, i in enumerate(argnums):
+            g = grad[j]
+            arg_nograd, which, treedef = notes[i]
+            none_grad = [None for _ in arg_nograd]
+            grad[j] = merge(g, none_grad, which, treedef)
+        if unwrap:
+            (grad,) = grad
+        return value, grad
+
+    return f_value_and_grad_wrapper
+
+
+@deprecated(in_favour_of=filter_grad)
+def gradf(fun, *, has_aux=False, **gradkwargs):
+    f_value_and_grad = value_and_grad_f(fun, has_aux=has_aux, **gradkwargs)
+
+    def f_grad(*args, **kwargs):
+        value, grad = f_value_and_grad(*args, **kwargs)
+        if has_aux:
+            value, aux = value
+            return aux, grad
+        else:
+            return grad
+
+    return f_grad
diff --git a/equinox/gradf.py b/equinox/gradf.py
deleted file mode 100644
index 2bf2adb1..00000000
--- a/equinox/gradf.py
+++ /dev/null
@@ -1,121 +0,0 @@
-import functools as ft
-
-import jax
-
-from .deprecated import deprecated
-from .filters import (
-    combine,
-    is_inexact_array,
-    merge,
-    partition,
-    split,
-    validate_filters,
-)
-
-
-def filter_value_and_grad(
-    fun, *, filter_spec=is_inexact_array, argnums=None, **gradkwargs
-):
-    if argnums is not None:
-        raise ValueError(
-            "`argnums` should not be passed. If you need to differentiate "
-            "multiple objects then collect them into a tuple and pass that "
-            "as the first argument."
-        )
-
-    @ft.partial(jax.value_and_grad, argnums=0, **gradkwargs)
-    def fun_value_and_grad(diff_x, nondiff_x, *args, **kwargs):
-        x = combine(diff_x, nondiff_x)
-        return fun(x, *args, **kwargs)
-
-    def fun_value_and_grad_wrapper(x, *args, **kwargs):
-        diff_x, nondiff_x = partition(x, filter_spec)
-        return fun_value_and_grad(diff_x, nondiff_x, *args, **kwargs)
-
-    return fun_value_and_grad_wrapper
-
-
-def filter_grad(fun, *, filter_spec=is_inexact_array, has_aux=False, **gradkwargs):
-    fun_value_and_grad = filter_value_and_grad(
-        fun, filter_spec=filter_spec, has_aux=has_aux, **gradkwargs
-    )
-
-    def fun_grad(*args, **kwargs):
-        value, grad = fun_value_and_grad(*args, **kwargs)
-        if has_aux:
-            value, aux = value
-            return aux, grad
-        else:
-            return grad
-
-    return fun_grad
-
-
-#
-# Deprecated
-#
-
-
-@deprecated(in_favour_of=filter_value_and_grad)
-def value_and_grad_f(fun, *, filter_fn=None, filter_tree=None, argnums=0, **gradkwargs):
-    if isinstance(argnums, int):
-        unwrap = True
-        argnums = (argnums,)
-        if filter_tree is not None:
-            filter_tree = (filter_tree,)
-    else:
-        unwrap = False
-
-    validate_filters("value_and_grad_f", filter_fn, filter_tree)
-
-    @ft.partial(jax.value_and_grad, argnums=argnums, **gradkwargs)
-    def f_value_and_grad(*args, **kwargs):
-        *args, notes = args
-        args = list(args)
-        for i, (arg_nograd, which, treedef) in notes.items():
-            arg_grad = args[i]
-            arg = merge(arg_grad, arg_nograd, which, treedef)
-            args[i] = arg
-        return fun(*args, **kwargs)
-
-    def f_value_and_grad_wrapper(*args, **kwargs):
-        args = list(args)
-        notes = {}
-        for j, i in enumerate(argnums):
-            arg = args[i]
-            if filter_fn is None:
-                # implies filter_tree is not None
-                arg_grad, arg_nograd, which, treedef = split(
-                    arg, filter_tree=filter_tree[j]
-                )
-            else:
-                arg_grad, arg_nograd, which, treedef = split(arg, filter_fn=filter_fn)
-            args[i] = arg_grad
-            notes[i] = (arg_nograd, which, treedef)
-        value, grad = f_value_and_grad(*args, notes, **kwargs)
-        grad = list(grad)
-        for j, i in enumerate(argnums):
-            g = grad[j]
-            arg_nograd, which, treedef = notes[i]
-            none_grad = [None for _ in arg_nograd]
-            grad[j] = merge(g, none_grad, which, treedef)
-        if unwrap:
-            (grad,) = grad
-        return value, grad
-
-    return f_value_and_grad_wrapper
-
-
-@deprecated(in_favour_of=filter_grad)
-def gradf(fun, *, has_aux=False, **gradkwargs):
-    f_value_and_grad = value_and_grad_f(fun, has_aux=has_aux, **gradkwargs)
-
-    def f_grad(*args, **kwargs):
-        value, grad = f_value_and_grad(*args, **kwargs)
-        if has_aux:
-            value, aux = value
-            return aux, grad
-        else:
-            return grad
-
-    return f_grad
diff --git a/equinox/jitf.py b/equinox/jit.py
similarity index 74%
rename from equinox/jitf.py
rename to equinox/jit.py
index d1712377..bea07d60 100644
--- a/equinox/jitf.py
+++ b/equinox/jit.py
@@ -29,30 +29,50 @@ def f_wrapped(
     return f_wrapped
 
 
-def filter_jit(
-    fun,
-    *,
-    filter_spec=is_array,
-    filter_spec_return=is_array,
-    static_argnums=None,
-    static_argnames=None,
-    donate_argnums=None,
-    **jitkwargs
-):
+def filter_jit(fun, *, filter_spec=is_array, filter_spec_return=is_array, **jitkwargs):
+    """Wraps together [`equinox.partition`][] and `jax.jit`.
 
-    if static_argnums is not None:
-        raise ValueError("`static_argnums` should not be passed; use a filter instead.")
-    if static_argnames is not None:
+    **Arguments:**
+
+    - `fun` is a pure function to JIT compile.
+    - `filter_spec` is a PyTree whose structure should be a prefix of the structure of
+        the inputs to `fun`. It behaves as the `filter_spec` argument to
+        [`equinox.filter`][]. Truthy values will be traced; falsey values will be held
+        static.
+    - `filter_spec_return` is a PyTree whose structure should be a prefix of the
+        structure of the outputs of `fun`. It behaves as the `filter_spec` argument to
+        [`equinox.filter`][]. Truthy values should be tracers; falsely values are any
+        (non-tracer) auxiliary information to return.
+    - `**jitkwargs` are any other keyword arguments to `jax.jit`.
+
+        !!! info
+
+            Specifically, if calling `fun(*args, **kwargs)`, then `filter_spec` must
+            have a structure which is a prefix for `(args, kwrgs)`.
+
+    **Returns:**
+
+    The JIT'd version of `fun`.
+
+    !!! info
+
+        A very important special case is to trace all JAX arrays and treat all other
+        objects as static.
+
+        This is accomplished with `filter_spec=equinox.is_array`,
+        `filter_spec_return=equinox.is_array` -- which are the defaults. (It is
+        unusual to need different behaviour to this.)
+    """
+
+    if any(
+        x in jitkwargs for x in ("static_argnums", "static_argnames", "donate_argnums")
+    ):
         raise ValueError(
-            "`static_argnames` should not be passed; use a filter instead."
-        )
-    if donate_argnums is not None:
-        raise NotImplementedError(
-            "`donate_argnums` is not implemented for filter_jit. Manually combine "
-            "`equinox.filter` and `jax.jit` instead.."
+            "`jitkwargs` cannot contain 'static_argnums', 'static_argnames' or "
+            "'donate_argnums'."
         )
 
-    # We choose not to make a distinction between ([arg, ... ,arg], kwargs) and ((arg, ... ,arg), kwargs)
+    # We choose not to make a distinction between ([arg, ..., arg], kwargs) and ((arg, ..., arg), kwargs)
     if (
         isinstance(filter_spec, tuple)
         and len(filter_spec) == 2
diff --git a/equinox/module.py b/equinox/module.py
index 8ae14973..34240abe 100644
--- a/equinox/module.py
+++ b/equinox/module.py
@@ -9,6 +9,29 @@
 
 
 def static_field(**kwargs):
+    """Used for marking that a field should _not_ be treated as part of the PyTree
+    of a [`equinox.Module`][]. (And is instead just treated as extra metadata.)
+
+    !!! example
+
+        ```python
+        class MyModule(equinox.Module):
+            normal_field: int
+            static_field: int = equinox.static_field()
+
+        mymodule = MyModule()
+        leaves, treedef = jax.tree_flatten(mymodule)
+        assert len(leaves) == 1
+        ```
+
+    In practice this should rarely be used; it is usually preferential to just filter
+    out each field with `eqx.filter` whenever you need to select only some fields.
+
+    **Arguments:**
+
+    - `**kwargs`: If any are passed then they are passed on to `datacalss.field`.
+        (Recall that Equinox uses dataclasses for its modules.)
+    """
     try:
         metadata = dict(kwargs["metadata"])
     except KeyError:
@@ -22,6 +45,8 @@ def static_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:
@@ -100,6 +125,87 @@ def __call__(cls, *args, **kwargs):
 
 
 class Module(metaclass=_ModuleMeta):
+    """Base class. Created your model by inheriting from this.
+
+    **Fields**
+
+    Specify all its fields at the class level (identical to
+    [dataclasses](https://docs.python.org/3/library/dataclasses.html)). This defines
+    its children as a PyTree.
+
+    ```python
+    class MyModule(equinox.Module):
+        weight: jax.numpy.ndarray
+        bias: jax.numpy.ndarray
+        submodule: equinox.Module
+    ```
+
+    **Initialisation**
+
+    A default `__init__` is automatically provided, which just fills in fields with the
+    arguments passed. For example `MyModule(weight, bias, submodule)`.
+
+    Alternatively (quite commonly) you can provide an `__init__` method yourself:
+
+    ```python
+    class MyModule(equinox.Module):
+        weight: jax.numpy.ndarray
+        bias: jax.numpy.ndarray
+        submodule: equinox.Module
+
+        def __init__(self, in_size, out_size, key):
+            wkey, bkey, skey = jax.random.split(key, 3)
+            self.weight = jax.random.normal(wkey, (out_size, in_size))
+            self.bias = jax.random.normal(bkey, (out_size,))
+            self.submodule = equinox.nn.Linear(in_size, out_size, key=skey)
+    ```
+
+    **Methods**
+
+    It is common to create some methods on the class -- for example to define the forward pass of a model.
+
+    ```python
+    class MyModule(equinox.Module):
+        ...  # as above
+
+        def __call__(self, x):
+            return self.submodule(x) + self.weight @ x + self.bias
+    ```
+
+    !!! tip
+
+        You don't have to define `__call__`:
+
+        - You can define other methods if you want.
+        - You can define multiple methods if you want.
+        - You can define no methods if you want. (And just use `equinox.Module` as a
+            nice syntax for custom PyTrees.)
+
+        No method is special-cased.
+
+    **Usage**
+
+    After you have defined your model, then you can use it just like any other PyTree
+    -- that just happens to have some methods attached. In particular you can pass it
+    around across `jax.jit`, `jax.grad` etc. in exactly the way that you're used to.
+
+    !!! example
+
+        If you wanted to, then it would be completely safe to do
+
+        ```python
+        class MyModule(equinox.Module):
+            ...
+
+            @jax.jit
+            def __call__(self, x):
+                ...
+        ```
+
+        because `self` is just a PyTree. Unlike most other neural network libaries, you
+        can mix Equinox and native JAX without any difficulties at all.
+    """
+
     _has_dataclass_init = True
 
     def __hash__(self):
diff --git a/equinox/nn/composed.py b/equinox/nn/composed.py
index 76bfbfc7..4fff39a7 100644
--- a/equinox/nn/composed.py
+++ b/equinox/nn/composed.py
@@ -1,16 +1,34 @@
-from typing import List
+import typing
+from typing import Any, Callable, List, Optional, Sequence
 
+import jax
 import jax.nn as jnn
 import jax.random as jrandom
 
+from ..custom_types import Array
 from ..module import Module, static_field
 from .linear import Linear
 
 
+def _identity(x):
+    return x
+
+
+if getattr(typing, "GENERATING_DOCUMENTATION", False):
+
+    def relu(_):
+        pass
+
+    jnn.relu = relu
+    _identity.__qualname__ = "identity"  # Renamed for nicer documentation.
+
+
 class MLP(Module):
+    """Standard Multi-Layer Perceptron; also known as a feed-forward network."""
+
     layers: List[Linear]
-    activation: callable
-    final_activation: callable
+    activation: Callable
+    final_activation: Callable
     in_size: int = static_field()
     out_size: int = static_field()
     width_size: int = static_field()
@@ -18,16 +36,30 @@ class MLP(Module):
 
     def __init__(
         self,
-        in_size,
-        out_size,
-        width_size,
-        depth,
-        activation=jnn.relu,
-        final_activation=lambda x: x,
+        in_size: int,
+        out_size: int,
+        width_size: int,
+        depth: int,
+        activation: Callable = jnn.relu,
+        final_activation: Callable = _identity,
         *,
-        key,
+        key: "jax.random.PRNGKey",
         **kwargs
     ):
+        """**Arguments**:
+
+        - `in_size`: The size of the input layer.
+        - `out_size`: The size of the output layer.
+        - `width_size`: The size of each hidden layer.
+        - `depth`: The number of hidden layers.
+        - `activation`: The activation function after each hidden layer. Defaults to
+            ReLU.
+        - `final_activation`: The activation function after the output layer. Defaults
+            to the identity.
+        - `key`: A `jax.random.PRNGKey` used to provide randomness for parameter
+            initialisation. (Keyword only argument.)
+        """
+
         super().__init__(**kwargs)
         keys = jrandom.split(key, depth + 1)
         layers = []
@@ -46,7 +78,19 @@ def __init__(
         self.activation = activation
         self.final_activation = final_activation
 
-    def __call__(self, x, *, key=None):
+    def __call__(
+        self, x: Array, *, key: Optional["jax.random.PRNGKey"] = None
+    ) -> Array:
+        """**Arguments:**
+
+        - `x`: A JAX array with shape `(in_size,)`.
+        - `key`: Ignored; provided for compatibility with the rest of the Equinox API.
+            (Keyword only argument.)
+
+        **Returns:**
+
+        A JAX array with shape `(out_size,)`.
+        """
         for layer in self.layers[:-1]:
             x = layer(x)
             x = self.activation(x)
@@ -56,9 +100,22 @@ def __call__(self, x, *, key=None):
 
 
 class Sequential(Module):
-    layers: List[Module]
+    """A sequence of [`equinox.Module`][]s applied in order."""
+
+    layers: Sequence[Module]
+
+    def __call__(self, x: Any, *, key: Optional["jax.random.PRNGKey"] = None) -> Any:
+        """**Arguments:**
+
+        - `x`: Argument passed to the first member of the sequence.
+        - `key`: A `jax.random.PRNGKey`, which will be split and passed to every layer
+            to provide any desired randomness. (Optional. Keyword only argument.)
+
+        **Returns:**
+
+        The output of the last member of the sequence.
+        """
 
-    def __call__(self, x, *, key=None):
         if key is None:
             keys = [None] * len(self.layers)
         else:
@@ -66,3 +123,9 @@ def __call__(self, x, *, key=None):
         for layer, key in zip(self.layers, keys):
             x = layer(x, key=key)
         return x
+
+
+Sequential.__init__.__doc__ = """**Arguments:**
+
+- `layers`: A sequence of [`equinox.Module`][]s.
+"""
diff --git a/equinox/nn/conv.py b/equinox/nn/conv.py
index 3c058f53..26e788f3 100644
--- a/equinox/nn/conv.py
+++ b/equinox/nn/conv.py
@@ -1,7 +1,8 @@
 import collections
 from itertools import repeat
-from typing import Optional, Sequence, Union
+from typing import Any, Optional, Sequence, Tuple, Union
 
+import jax
 import jax.numpy as jnp
 import jax.random as jrandom
 import numpy as np
@@ -11,8 +12,8 @@
 from ..module import Module, static_field
 
 
-def _ntuple(n):
-    def parse(x):
+def _ntuple(n: int) -> callable:
+    def parse(x: Any) -> tuple:
         if isinstance(x, collections.abc.Iterable):
             if len(x) == n:
                 return tuple(x)
@@ -26,31 +27,58 @@ def parse(x):
 
 
 class Conv(Module):
+    """General N-dimensional convolution."""
+
     num_spatial_dims: int = static_field()
     weight: Array
     bias: Optional[Array]
     in_channels: int = static_field()
     out_channels: int = static_field()
-    kernel_size: Union[int, Sequence[int]] = static_field()
-    stride: Union[int, Sequence[int]] = static_field()
-    padding: Union[int, Sequence[int]] = static_field()
-    dilation: Union[int, Sequence[int]] = static_field()
+    kernel_size: Tuple[int] = static_field()
+    stride: Tuple[int] = static_field()
+    padding: Tuple[int] = static_field()
+    dilation: Tuple[int] = static_field()
     use_bias: bool = static_field()
 
     def __init__(
         self,
-        num_spatial_dims,
-        in_channels,
-        out_channels,
-        kernel_size,
-        stride=1,
-        padding=0,
-        dilation=1,
-        use_bias=True,
+        num_spatial_dims: int,
+        in_channels: int,
+        out_channels: int,
+        kernel_size: Union[int, Sequence[int]],
+        stride: Union[int, Sequence[int]] = 1,
+        padding: Union[int, Sequence[int]] = 0,
+        dilation: Union[int, Sequence[int]] = 1,
+        use_bias: bool = True,
         *,
-        key,
+        key: "jax.random.PRNGKey",
         **kwargs,
     ):
+        """**Arguments:**
+
+        - `num_spatial_dims`: The number of spatial dimensions. For example traditional
+            convolutions for image processing have this set to `2`.
+        - `in_channels`: The number of input channels.
+        - `out_channels`: The number of output channels.
+        - `kernel_size`: The size of the convolutional kernel.
+        - `stride`: The stride of the convolution.
+        - `padding`: The amount of padding to apply before and after each spatial
+            dimension. The same amount of padding is applied both before and after.
+        - `dilation`: The dilation of the convolution.
+        - `use_bias`: Whether to add on a bias after the convolution.
+        - `key`: A `jax.random.PRNGKey` used to provide randomness for parameter
+            initialisation. (Keyword only argument.)
+
+        !!! info
+
+            All of `kernel_size`, `stride`, `padding`, `dilation` can be either an
+            integer or a sequence of integers. If they are a sequence then the sequence
+            should be of length equal to `num_spatial_dims`, and specify the value of
+            each property down each spatial dimension in turn.. If they are an integer
+            then the same kernel size / stride / padding / dilation will be used along
+            every spatial dimension.
+
+        """
         super().__init__(**kwargs)
         self.num_spatial_dims = num_spatial_dims
         parse = _ntuple(self.num_spatial_dims)
@@ -95,7 +123,21 @@ def __init__(
             )
         self.dilation = parse(dilation)
 
-    def __call__(self, x, *, key=None):
+    def __call__(
+        self, x: Array, *, key: Optional["jax.random.PRNGKey"] = None
+    ) -> Array:
+        """**Arguments:**
+
+        - `x`: The input. Should be a JAX array of shape `(in_channels, dim_1, ..., dim_N)`, where
+            `N = num_spatial_dims`.
+        - `key`: Ignored; provided for compatibility with the rest of the Equinox API.
+            (Keyword only argument.)
+
+        **Returns:**
+
+        A JAX array of shape `(out_channels, new_dim_1, ..., new_dim_N)`.
+        """
+
         unbatched_rank = self.num_spatial_dims + 1
         if x.ndim != unbatched_rank:
             raise ValueError(
@@ -117,6 +159,8 @@ def __call__(self, x, *, key=None):
 
 
 class Conv1d(Conv):
+    """As [`equinox.nn.Conv`][] with `num_spatial_dims=1`."""
+
     def __init__(
         self,
         in_channels,
@@ -145,6 +189,8 @@ def __init__(
 
 
 class Conv2d(Conv):
+    """As [`equinox.nn.Conv`][] with `num_spatial_dims=2`."""
+
     def __init__(
         self,
         in_channels,
@@ -173,6 +219,8 @@ def __init__(
 
 
 class Conv3d(Conv):
+    """As [`equinox.nn.Conv`][] with `num_spatial_dims=3`."""
+
     def __init__(
         self,
         in_channels,
diff --git a/equinox/nn/dropout.py b/equinox/nn/dropout.py
index c1416cab..c3c82d65 100644
--- a/equinox/nn/dropout.py
+++ b/equinox/nn/dropout.py
@@ -1,15 +1,37 @@
+from typing import Optional
+
+import jax
 import jax.numpy as jnp
 import jax.random as jrandom
 
+from ..custom_types import Array
 from ..module import Module
 
 
 class Dropout(Module):
+    """Applies dropout."""
+
     # Not static_fields as it makes sense to want to modify them via equinox.tree_at.
     p: float = 0.5
     deterministic: bool = False
 
-    def __call__(self, x, *, key=None, deterministic=None):
+    def __call__(
+        self,
+        x: Array,
+        *,
+        key: Optional["jax.random.PRNGKey"] = None,
+        deterministic: Optional[bool] = None
+    ) -> Array:
+        """**Arguments:**
+
+        - `x`: An any-dimensional JAX array to dropout.
+        - `key`: A `jax.random.PRNGKey` used to provide randomness for calculating
+            which elements to dropout. (Keyword only argument.)
+        - `deterministic`: As per [`equinox.nn.Dropout.__init__`][]. If `True` or
+            `False` then it will take priority over `self.deterministic`. If `None`
+            then the value from `self.deterministic` will be used.
+        """
+
         if deterministic is None:
             deterministic = self.deterministic
         if deterministic:
@@ -22,3 +44,30 @@ def __call__(self, x, *, key=None, deterministic=None):
             q = 1 - self.p
             mask = jrandom.bernoulli(key, q, x.shape)
             return jnp.where(mask, x / q, 0)
+
+
+Dropout.__init__.__doc__ = """**Arguments:**
+
+- `p`: The fraction of entries to set to zero. (On average.)
+- `deterministic`: Whether to actually apply dropout at all. If `True` then dropout
+    is *not* applied. If `False` then dropout is applied.
+
+!!! info
+
+    The `deterministic` flag is provided as it is common to only apply dropout
+    during training, but not to apply it during inference. If you want to change this
+    flag between training and inference, then you can either:
+
+    - Override it with the `__call__`-time `deterministic` flag, see below.
+    - Modify the `deterministic` flag directly -- possible to do because `Dropout` is
+        just a PyTree. For example this sets all `deterministic` flags to `True`:
+        ```python
+        model = ...  # some model featuring a Dropout layer somewhere
+        is_dropout = lambda x: isinstance(x, Dropout)
+        def find_deterministic(m):
+            return tuple(d.deterministic
+                         for d in jax.tree_flatten(m, is_leaf=is_dropout)[0]
+                         if is_dropout(d))
+        model = eqx.tree_at(find_deterministic, model, replace_fn=lambda _: True)
+        ```
+"""
diff --git a/equinox/nn/linear.py b/equinox/nn/linear.py
index e258de4f..bc965740 100644
--- a/equinox/nn/linear.py
+++ b/equinox/nn/linear.py
@@ -1,6 +1,7 @@
 import math
-from typing import Optional
+from typing import Optional, TypeVar
 
+import jax
 import jax.random as jrandom
 
 from ..custom_types import Array
@@ -8,13 +9,30 @@
 
 
 class Linear(Module):
+    """Performs a linear transformation."""
+
     weight: Array
     bias: Optional[Array]
     in_features: int = static_field()
     out_features: int = static_field()
     use_bias: bool = static_field()
 
-    def __init__(self, in_features, out_features, use_bias=True, *, key):
+    def __init__(
+        self,
+        in_features: int,
+        out_features: int,
+        use_bias: bool = True,
+        *,
+        key: "jax.random.PRNGKey"
+    ):
+        """**Arguments:**
+
+        - `in_features`: The input size.
+        - `out_features`: The output size.
+        - `use_bias`: Whether to add on a bias as well.
+        - `key`: A `jax.random.PRNGKey` used to provide randomness for parameter
+            initialisation. (Keyword only argument.)
+        """
         super().__init__()
         wkey, bkey = jrandom.split(key, 2)
         lim = 1 / math.sqrt(in_features)
@@ -30,17 +48,48 @@ def __init__(self, in_features, out_features, use_bias=True, *, key):
         self.out_features = out_features
         self.use_bias = use_bias
 
-    def __call__(self, x, *, key=None):
+    def __call__(
+        self, x: Array, *, key: Optional["jax.random.PRNGKey"] = None
+    ) -> Array:
+        """**Arguments:**
+
+        - `x`: The input. Should be a JAX array of shape `(in_features,)`.
+        - `key`: Ignored; provided for compatibility with the rest of the Equinox API.
+            (Keyword only argument.)
+
+        **Returns:**
+
+        A JAX array of shape `(out_features,)`
+        """
+
         x = self.weight @ x
         if self.bias is not None:
             x = x + self.bias
         return x
 
 
+_T = TypeVar("T")
+
+
 class Identity(Module):
+    """Identity operation that does nothing. Sometimes useful as a placeholder for
+    another Module.
+    """
+
     def __init__(self, *args, **kwargs):
+        """Consumes arbitrary `*args` and `**kwargs` but ignores them."""
         # Ignores args and kwargs
         super().__init__()
 
-    def __call__(self, x, *, key=None):
+    def __call__(self, x: _T, *, key: Optional["jax.random.PRNGKey"] = None) -> _T:
+        """**Arguments:**
+
+        - `x`: The input, of any type.
+        - `key`: Ignored; provided for compatibility with the rest of the Equinox API.
+            (Keyword only argument.)
+
+        **Returns:**
+
+        The input, unchanged.
+        """
         return x
diff --git a/equinox/nn/rnn.py b/equinox/nn/rnn.py
index d80faf48..b9bda30d 100644
--- a/equinox/nn/rnn.py
+++ b/equinox/nn/rnn.py
@@ -2,6 +2,7 @@
 import warnings
 from typing import Optional
 
+import jax
 import jax.nn as jnn
 import jax.numpy as jnp
 import jax.random as jrandom
@@ -11,6 +12,27 @@
 
 
 class GRUCell(Module):
+    """A single step of a Gated Recurrent Unit (GRU).
+
+    !!! example
+
+        This is often used by wrapping it into a `jax.lax.scan`. For example:
+
+        ```python
+        class Model(Module):
+            cell: GRUCell
+
+            def __init__(self, ...):
+                self.cell = GRUCell(...)
+
+            def __call__(self, xs):
+                scan_fn = lambda state, input: (cell(input, state), None)
+                init_state = jnp.zeros(self.cell.hidden_size)
+                final_state, _ = jax.lax.scan(scan_fn init_state, xs)
+                return final_state
+        ```
+    """
+
     weight_ih: Array
     weight_hh: Array
     bias: Optional[Array]
@@ -20,8 +42,24 @@ class GRUCell(Module):
     use_bias: bool = static_field()
 
     def __init__(
-        self, input_size, hidden_size, use_bias=True, bias=None, *, key, **kwargs
+        self,
+        input_size: int,
+        hidden_size: int,
+        use_bias: bool = True,
+        bias=None,
+        *,
+        key: Optional["jax.random.PRNGKey"],
+        **kwargs
     ):
+        """**Arguments:**
+
+        - `input_size`: The dimensionality of the input vector at each time step.
+        - `hidden_size`: The dimensionality of the hidden state passed along between
+            time steps.
+        - `use_bias`: Whether to add on a bias after each update.
+        - `key`: A `jax.random.PRNGKey` used to provide randomness for parameter
+            initialisation. (Keyword only argument.)
+        """
         super().__init__(**kwargs)
         if bias is not None:
             warnings.warn("`bias` is deprecated in favour of `use_bias`.")
@@ -51,7 +89,21 @@ def __init__(
         self.hidden_size = hidden_size
         self.use_bias = use_bias
 
-    def __call__(self, input, hidden, *, key=None):
+    def __call__(
+        self, input: Array, hidden: Array, *, key: Optional["jax.random.PRNGKey"] = None
+    ):
+        """**Arguments:**
+
+        - `input`: The input, which should be a JAX array of shape `(input_size,)`.
+        - `hidden`: The hidden state, which should be a JAX array of shape
+            `(hidden_size,)`.
+        - `key`: Ignored; provided for compatibility with the rest of the Equinox API.
+            (Keyword only argument.)
+
+        **Returns:**
+
+        The updated hidden state, which is a JAX array of shape `(hidden_size,)`.
+        """
         if self.bias is None:
             bias = 0
             bias_n = 0
@@ -67,6 +119,28 @@ def __call__(self, input, hidden, *, key=None):
 
 
 class LSTMCell(Module):
+    """A single step of a Long-Short Term Memory unit (LSTM).
+
+    !!! example
+
+        This is often used by wrapping it into a `jax.lax.scan`. For example:
+
+        ```python
+        class Model(Module):
+            cell: LSTMCell
+
+            def __init__(self, ...):
+                self.cell = LSTMCell(...)
+
+            def __call__(self, xs):
+                scan_fn = lambda state, input: (cell(input, state), None)
+                init_state = (jnp.zeros(self.cell.hidden_size),
+                              jnp.zeros(self.cell.hidden_size))
+                final_state, _ = jax.lax.scan(scan_fn init_state, xs)
+                return final_state
+        ```
+    """
+
     weight_ih: Array
     weight_hh: Array
     bias: Optional[Array]
@@ -75,8 +149,24 @@ class LSTMCell(Module):
     use_bias: bool = static_field()
 
     def __init__(
-        self, input_size, hidden_size, use_bias=True, bias=None, *, key, **kwargs
+        self,
+        input_size: int,
+        hidden_size: int,
+        use_bias: bool = True,
+        bias=None,
+        *,
+        key: "jax.random.PRNGKey",
+        **kwargs
     ):
+        """**Arguments:**
+
+        - `input_size`: The dimensionality of the input vector at each time step.
+        - `hidden_size`: The dimensionality of the hidden state passed along between
+            time steps.
+        - `use_bias`: Whether to add on a bias after each update.
+        - `key`: A `jax.random.PRNGKey` used to provide randomness for parameter
+            initialisation. (Keyword only argument.)
+        """
         super().__init__(**kwargs)
         if bias is not None:
             warnings.warn("`bias` is deprecated in favour of `use_bias`.")
@@ -103,6 +193,19 @@ def __init__(
         self.use_bias = use_bias
 
     def __call__(self, input, hidden, *, key=None):
+        """**Arguments:**
+
+        - `input`: The input, which should be a JAX array of shape `(input_size,)`.
+        - `hidden`: The hidden state, which should be a 2-tuple of JAX arrays, each of
+            shape `(hidden_size,)`.
+        - `key`: Ignored; provided for compatibility with the rest of the Equinox API.
+            (Keyword only argument.)
+
+        **Returns:**
+
+        The updated hidden state, which is a 2-tuple of JAX arrays, each of shape
+        `(hidden_size,)`.
+        """
         h, c = hidden
         lin = self.weight_ih @ input + self.weight_hh @ h
         if self.bias is not None:
diff --git a/equinox/tree.py b/equinox/tree.py
index 7f37358e..2c99a0cc 100644
--- a/equinox/tree.py
+++ b/equinox/tree.py
@@ -1,20 +1,65 @@
-from typing import Any, Callable, Optional, Tuple, Union
-from typing_extensions import get_args
+from typing import Any, Callable, Sequence, Union
 
 import jax
+import jax.numpy as jnp
+import numpy as np
 
-from .custom_types import MoreArrays, PyTree
+from .custom_types import PyTree
 
 
 _sentinel = object()
 
+_Leaf = Any
+
 
 def tree_at(
-    where: Callable[[PyTree], Union[Any, Tuple[Any, ...]]],
+    where: Callable[[PyTree], Union[_Leaf, Sequence[_Leaf]]],
     pytree: PyTree,
-    replace: Optional[Union[Any, Tuple[Any, ...]]] = _sentinel,
-    replace_fn: Optional[Callable[[Any], Any]] = _sentinel,
+    replace: Union[_Leaf, Sequence[_Leaf]] = _sentinel,
+    replace_fn: Callable[[_Leaf], _Leaf] = _sentinel,
+    is_leaf: Callable[[_Leaf], bool] = None,
 ) -> PyTree:
+    """Updates a PyTree out-of-place; a bit like using `.at[].set()` on a JAX array.
+
+    **Arguments:**
+
+    - `where`: A callable `PyTree -> Leaf` or `PyTree -> Sequence[Leaf]`. It should
+        consume a PyTree with the same structure as `pytree`, and return the leaf or
+        leaves that should be replaced. For example
+        `where = lambda mlp: mlp.layers[-1].linear.weight`.
+    - `pytree`: The PyTree to modify.
+    - `replace`: Either a single element, or a sequence of the same length as returned
+        by `where`. This specifies the replacements to make at the locations specified
+        by `where`. Mutually exclusive with `replace_fn`.
+    - `replace_fn`: A function `Leaf -> Any`. It will be called on every leaf specified
+        by `where`. The return value from `replace_fn` will be used in its place.
+        Mutually exclusive with `replace`.
+
+    **Returns:**
+
+    A copy of the input PyTree, with the appropriate modifications.
+
+    !!! example
+
+        This can be used to help specify the weights of a model to train or not to
+        train:
+
+        ```python
+        model = ...
+        trainable = jax.tree_map(lambda _: False, model)
+        trainable = equinox.tree_at(lambda mlp: mlp.layers[-1].linear.weight, model, replace=True)
+        equinox.filter_grad(..., filter_spec=trainable)
+        ```
+
+    !!! example
+
+        Sub-PyTrees can be replaced by flattening them to leaves first:
+
+        ```python
+        equinox.tree_at(lambda t: jax.tree_leaves(t.subtree), pytree,
+                        jax.tree_leaves(new_subtree))
+        ```
+    """
 
     if (replace is _sentinel and replace_fn is _sentinel) or (
         replace is not _sentinel and replace_fn is not _sentinel
@@ -30,17 +75,32 @@ def tree_at(
         replacer = lambda j, i: replace[j]
 
     # TODO: is there a neater way of accomplishing this?
-    flat, treedef = jax.tree_flatten(pytree)
+    flat, treedef = jax.tree_flatten(pytree, is_leaf=is_leaf)
     flat_indices = list(range(len(flat)))
     index_pytree = jax.tree_unflatten(treedef, flat_indices)
     index = where(index_pytree)
-    # where can return either a single entry, or a tuple.
+    # where can return either a single entry, or a sequence
     if isinstance(index, int):
         index = (index,)
         replace = (replace,)
+    elif isinstance(index, Sequence):
+        for i in index:
+            if not isinstance(i, int):
+                raise ValueError(
+                    r"""`where` must return a sequence of only leaves; not some subtree.
+
+                    If you want to replace all of a subtree, you can do so by replacing
+                    >>> eqx.tree_at(lambda t: t.subtree, tree, new_subtree)  # buggy
+                    with
+                    >>> eqx.tree_at(lambda t: jax.tree_leaves(t.subtree), tree, 
+                    ...             jax.tree_leaves(new_subtree))  # fixed
+                    """
+                )
 
     if replace_passed and len(index) != len(replace):
-        raise ValueError("`where` must return a tuple of the same length as `replace`.")
+        raise ValueError(
+            "`where` must return a sequence of leaves of the same length as `replace`."
+        )
     for j, i in enumerate(index):
         flat[i] = replacer(j, i)
 
@@ -48,8 +108,21 @@ def tree_at(
 
 
 def tree_equal(*pytrees: PyTree) -> bool:
+    """Returns `True` if all input PyTrees are equal. Every PyTrees 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 not considered equal
+    to each other.
+
+    **Arguments:**
+
+    - `*pytrees`: Any number of PyTrees each with any structure.
+
+    **Returns:**
+
+    A boolean.
+    """
     flat, treedef = jax.tree_flatten(pytrees[0])
-    array_types = get_args(MoreArrays)
+    array_types = (jnp.ndarray, np.ndarray)
     for pytree in pytrees[1:]:
         flat_, treedef_ = jax.tree_flatten(pytree)
         if treedef_ != treedef:
diff --git a/equinox/update.py b/equinox/update.py
index 78b4399b..ffc97e73 100644
--- a/equinox/update.py
+++ b/equinox/update.py
@@ -15,5 +15,23 @@ def _is_none(x):
 
 
 def apply_updates(model: PyTree, updates: PyTree) -> PyTree:
+    """A `jax.tree_map`-broadcasted version of
+    ```python
+    model = model if update is None else model + update
+    ```
+
+    This is often useful when updating a model's parameters via stochastic gradient
+    descent. (This function is essentially the same as `optax.apply_updates`, except
+    that it understands `None`.)
+
+    **Arguments:**
+
+    - `model`: An arbitrary PyTree.
+    - `updates`: Any PyTree that is a prefix of `model`.
+
+    **Returns:**
+
+    The updated model.
+    """
     # Assumes that updates is a prefix of model
     return jax.tree_map(_apply_update, updates, model, is_leaf=_is_none)
diff --git a/examples/build_model.py b/examples/build_model.py
deleted file mode 100644
index 88e764c9..00000000
--- a/examples/build_model.py
+++ /dev/null
@@ -1,149 +0,0 @@
-#############
-#
-# In this example, we look at Modules.
-#
-# Now I know what you're thinking! Every deep learning library out there has a class called Module. There's
-# haiku.Module, flax.linen.Module, objax.Module etc.
-#
-# And each time you have to sit down and read the documentation and understand what "Module" means for each library.
-# A lot of these also have custom notions of variables, transforms, scopes, etc. For example there's `haiku.transform`
-# or `objax.VarCollection`.
-#
-# In constrast, Equinox introduces no new abstractions. An Equinox Module is just a nice way to create a PyTree,
-# really. If you want, look up the source code for `equinox.Module` -- it's only about 100 lines long.
-#
-#############
-#
-# Now that we've finished complaining about the competition ;) let's see how it works.
-#
-# It's very simple: `Module`, and its subclasses, are PyTrees. Any attributes of the Module are also part of the
-# same PyTree. (Your whole model is just one big PyTree.) This means you can use the model in the normal way in
-# JAX, with vmap/grad/jit etc.
-#
-# Now because a `Module` is also a class, we can define methods on it. The `self` parameter -- which is a PyTree,
-# remember! -- means that this is just a function that takes PyTrees as inputs, like any other function. No method
-# is special cased. If you want you can group several related functions under different methods. If you just want
-# to define a single forward pass, then the __call__ method is a convenient choice.
-#
-#############
-#
-# In this example, we'll demonstrate how to use `equinox.Module` to create a simple MLP.
-
-import functools as ft
-from typing import Any, List
-
-import jax
-import jax.nn as jnn
-import jax.numpy as jnp
-import jax.random as jrandom
-
-import equinox as eqx
-
-
-# First pass at creating a Linear layer. `Linear1` will be a PyTree node, with `weight` and `bias` as its children.
-# A type annotation (the `Any`) is needed for these to be recognised; just use `typing.Any` if you want.
-class Linear1(eqx.Module):
-    weight: Any
-    bias: Any
-
-    def __call__(self, x):
-        return self.weight @ x + self.bias
-
-
-# We do need to initialise this. If we wanted to we could create a free function to do it:
-def linear(in_features, out_features, key):
-    wkey, bkey = jrandom.split(key)
-    weight = jrandom.normal(key, (out_features, in_features))
-    bias = jrandom.normal(key, (out_features,))
-    return Linear1(weight, bias)  # uses the default __init__ for Linear1.
-
-
-# Alternatively we can  use a custom __init__:
-class Linear2(eqx.Module):
-    weight: Any
-    bias: Any
-
-    def __init__(self, in_features, out_features, key):
-        # The super-init call isn't actually necessary here, but is good practice in Python (for
-        # co-operative multiple inheritance).
-        super().__init__()
-        wkey, bkey = jrandom.split(key)
-        weight = jrandom.normal(key, (out_features, in_features))
-        bias = jrandom.normal(key, (out_features,))
-        self.weight = weight
-        self.bias = bias
-        # An error will be thrown if you forget to set either `weight` or `bias` (or if you to try to set
-        # anything else).
-
-    def __call__(self, x):
-        return self.weight @ x + self.bias
-
-
-# And now we can compose these into a small MLP:
-class MLP(eqx.Module):
-    layers: List[Linear2]
-    activation: callable
-
-    def __init__(self, in_size, out_size, width_size, depth, key, activation=jnn.relu):
-        super().__init__()  # Once again not necessary but is good practice
-        keys = jrandom.split(key, depth + 1)
-        self.layers = []
-        self.layers.append(Linear2(in_size, width_size, keys[0]))
-        for i in range(depth - 1):
-            self.layers.append(Linear2(width_size, width_size, keys[i + 1]))
-        self.layers.append(Linear2(width_size, out_size, keys[-1]))
-        self.activation = activation
-
-    def __call__(self, x):
-        # Incidentally if you want to JIT/grad/whatever in here you can; it's completely safe to do so, unlike some
-        # other libraries.
-        for layer in self.layers[:-1]:
-            x = layer(x)
-            x = self.activation(x)
-        return self.layers[-1](x)
-
-
-# Which we can now use:
-def main():
-    key = jrandom.PRNGKey(5678)
-    model_key, data_key = jrandom.split(key, 2)
-    model = MLP(in_size=2, out_size=3, width_size=8, depth=2, key=model_key)
-    data = jrandom.normal(data_key, (2,))
-    model(data)  # Calls __call__
-
-    # Because `model` is a PyTree we can use it with normal JAX: vmap, grad, jit etc.
-    # However note that `model.activation` is some arbitrary Python function, defaulting to a ReLU.
-    # JAX has no idea how to JIT/differentiate that! We need to separate the things we want to JIT/grad from the rest.
-    params = eqx.filter(model, eqx.is_array)
-    static = eqx.filter(model, eqx.is_array, inverse=True)
-
-    @ft.partial(jax.jit, static_argnums=1)
-    def example_jit(params, static, data):
-        model = eqx.combine(params, static)
-        model(data)
-
-    @jax.grad
-    def example_grad(params, static, data):
-        model = eqx.combine(params, static)
-        return jnp.sum(model(data))  # return a scalar
-
-    @ft.partial(jax.vmap, in_axes=(None, None, 0))
-    def example_vmap(params, static, data):
-        model = eqx.combine(params, static)
-        return model(data)
-
-    example_jit(params, static, data)
-    example_grad(params, static, data)
-    example_vmap(params, static, jnp.stack([data, data]))
-
-    # Now that you've seen how to build a model, you're all set!
-    # This is literally everything you need to know to use Equinox. Like the README promises, it's a really simple way
-    # to build models, and you don't need to learn any new abstractions -- it's all just PyTrees.
-    #
-    # If you're comfortable with a little more automation, it is possible to combine JIT/grad with the filtering
-    # automatically, so you can just do fn(model) rather than fn(params, static). Have a look at
-    # filtered_transformations.py.
-
-
-if __name__ == "__main__":
-    main()
diff --git a/examples/filtered_transformations.py b/examples/filtered_transformations.py
deleted file mode 100644
index 1ed0d047..00000000
--- a/examples/filtered_transformations.py
+++ /dev/null
@@ -1,85 +0,0 @@
-##############
-#
-# This example demonstrates how to use `equinox.filter_jit` and `equinox.filter_grad`.
-#
-# Here we'll use them to facilitate training a simple MLP: to automatically take gradients and jit with respect to
-# all the jnp.arrays constituting the parameters. (But not with respect to anything else, like the choice of activation
-# function -- as that isn't something we can differentiate/JIT anyway!)
-#
-#############
-
-import jax
-import jax.numpy as jnp
-import jax.random as jrandom
-import optax
-
-import equinox as eqx
-
-
-# Toy data
-def get_data(dataset_size, *, key):
-    x = jrandom.normal(key, (dataset_size, 1))
-    y = 5 * x - 2
-    return x, y
-
-
-# Simple dataloader
-def dataloader(arrays, key, batch_size):
-    dataset_size = arrays[0].shape[0]
-    assert all(array.shape[0] == dataset_size for array in arrays)
-    indices = jnp.arange(dataset_size)
-    while True:
-        perm = jrandom.permutation(key, indices)
-        (key,) = jrandom.split(key, 1)
-        start = 0
-        end = batch_size
-        while end < dataset_size:
-            batch_perm = perm[start:end]
-            yield tuple(array[batch_perm] for array in arrays)
-            start = end
-            end = start + batch_size
-
-
-def main(
-    dataset_size=10000,
-    batch_size=256,
-    learning_rate=3e-3,
-    steps=1000,
-    width_size=8,
-    depth=1,
-    seed=5678,
-):
-    data_key, loader_key, model_key = jrandom.split(jrandom.PRNGKey(seed), 3)
-    data = get_data(dataset_size, key=data_key)
-    data = dataloader(data, batch_size=batch_size, key=loader_key)
-
-    # We happen to be using an Equinox model here, but that *is not important*.
-    # `equinox.filter_jit` and `equinox.filter_grad` will work just fine on any PyTree you like.
-    # (Here, `model` is actually a PyTree -- have a look at the `build_model.py` example for more on that.)
-    model = eqx.nn.MLP(
-        in_size=1, out_size=1, width_size=width_size, depth=depth, key=model_key
-    )
-
-    # `filter_jit` and `filter_value_and_grad` are thin wrappers around the usual `jax` functions, that automatically
-    # inspect the arguments of the function, JIT with respect to all JAX arrays, and differentiate with respect to all
-    # floating point JAX arrays (i.e. the parameters of the model).
-    @eqx.filter_jit
-    @eqx.filter_value_and_grad
-    def loss(model, x, y):
-        pred_y = jax.vmap(model)(x)
-        return jnp.mean((y - pred_y) ** 2)
-
-    optim = optax.sgd(learning_rate)
-    opt_state = optim.init(model)
-    for step, (x, y) in zip(range(steps), data):
-        value, grads = loss(model, x, y)
-        updates, opt_state = optim.update(grads, opt_state)
-        # Essentially equivalent to optax.apply_updates, it just doesn't try to update anything with a gradient
-        # of `None` (which is the gradient produced for anything we filtered out above).
-        model = eqx.apply_updates(model, updates)
-        print(step, value)
-    return value  # Final loss
-
-
-if __name__ == "__main__":
-    main()
diff --git a/examples/frozen_layer.ipynb b/examples/frozen_layer.ipynb
new file mode 100644
index 00000000..508c0e44
--- /dev/null
+++ b/examples/frozen_layer.ipynb
@@ -0,0 +1,213 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "71f37252-8e84-4521-b9c7-7fb107f599c1",
+   "metadata": {},
+   "source": [
+    "# Freezing parameters\n",
+    "\n",
+    "In this example, we demonstrate how the `filter_spec` argument of `equinox.filter_value_and_grad` can be used -- in this case, to only train some parameters and freeze the rest."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "9e127af2-03ac-47d7-b185-4b6b40d4cb7a",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import functools as ft\n",
+    "\n",
+    "import jax\n",
+    "import jax.numpy as jnp\n",
+    "import jax.random as jrandom\n",
+    "import optax  # https://github.com/deepmind/optax\n",
+    "\n",
+    "import equinox as eqx"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "a5732648-7126-465f-b8b2-8d18a325f7c8",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Toy data\n",
+    "def get_data(dataset_size, *, key):\n",
+    "    x = jrandom.normal(key, (dataset_size, 1))\n",
+    "    y = 5 * x - 2\n",
+    "    return x, y\n",
+    "\n",
+    "\n",
+    "# Toy dataloader\n",
+    "def dataloader(arrays, batch_size, *, key):\n",
+    "    dataset_size = arrays[0].shape[0]\n",
+    "    assert all(array.shape[0] == dataset_size for array in arrays)\n",
+    "    indices = jnp.arange(dataset_size)\n",
+    "    while True:\n",
+    "        perm = jrandom.permutation(key, indices)\n",
+    "        (key,) = jrandom.split(key, 1)\n",
+    "        start = 0\n",
+    "        end = batch_size\n",
+    "        while end < dataset_size:\n",
+    "            batch_perm = perm[start:end]\n",
+    "            yield tuple(array[batch_perm] for array in arrays)\n",
+    "            start = end\n",
+    "            end = start + batch_size"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "1f8d3f12-9a34-4609-984d-adfce11a3e7a",
+   "metadata": {},
+   "source": [
+    "Here, we:\n",
+    "\n",
+    "1. Set up a model. In this case, an MLP.\n",
+    "2. Set up a `filter_spec`. This will be a PyTree of the same structure as the model, with `False` on every leaf -- except for the leaves corresponding to the final layer, which we set to `True`.\n",
+    "3. Specify how to make a step. In this case we'll specify that we're still going to JIT with respect to every array, but we're only going to differentiate the ones specified by `filter_spec`."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "id": "a8dbd053-6695-469c-8734-bd4326c81add",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def main(\n",
+    "    dataset_size=10000,\n",
+    "    batch_size=256,\n",
+    "    learning_rate=3e-3,\n",
+    "    steps=1000,\n",
+    "    width_size=8,\n",
+    "    depth=1,\n",
+    "    seed=5678,\n",
+    "):\n",
+    "    data_key, loader_key, model_key = jrandom.split(jrandom.PRNGKey(seed), 3)\n",
+    "    data = get_data(dataset_size, key=data_key)\n",
+    "    data_iter = dataloader(data, batch_size, key=loader_key)\n",
+    "\n",
+    "    # Step 1\n",
+    "    model = eqx.nn.MLP(\n",
+    "        in_size=1, out_size=1, width_size=width_size, depth=depth, key=model_key\n",
+    "    )\n",
+    "\n",
+    "    # Step 2\n",
+    "    filter_spec = jax.tree_map(lambda _: False, model)\n",
+    "    filter_spec = eqx.tree_at(\n",
+    "        lambda tree: (tree.layers[-1].weight, tree.layers[-1].bias),\n",
+    "        filter_spec,\n",
+    "        replace=(True, True),\n",
+    "    )\n",
+    "\n",
+    "    # Step 3\n",
+    "    @eqx.filter_jit\n",
+    "    @ft.partial(eqx.filter_value_and_grad, filter_spec=filter_spec)\n",
+    "    def make_step(model, x, y):\n",
+    "        pred_y = jax.vmap(model)(x)\n",
+    "        return jnp.mean((y - pred_y) ** 2)\n",
+    "\n",
+    "    # And now let's train for a short while -- in exactly the usual way -- and see what\n",
+    "    # happens. We keep the original model around to compare to later.\n",
+    "    original_model = model\n",
+    "    optim = optax.sgd(learning_rate)\n",
+    "    opt_state = optim.init(model)\n",
+    "    for step, (x, y) in zip(range(steps), data_iter):\n",
+    "        value, grads = make_step(model, x, y)\n",
+    "        updates, opt_state = optim.update(grads, opt_state)\n",
+    "        model = eqx.apply_updates(model, updates)\n",
+    "    print(\n",
+    "        f\"Parameters of first layer at initialisation:\\n{jax.tree_leaves(original_model.layers[0])}\\n\"\n",
+    "    )\n",
+    "    print(\n",
+    "        f\"Parameters of first layer at end of training:\\n{jax.tree_leaves(model.layers[0])}\\n\"\n",
+    "    )\n",
+    "    print(\n",
+    "        f\"Parameters of last layer at initialisation:\\n{jax.tree_leaves(original_model.layers[-1])}\\n\"\n",
+    "    )\n",
+    "    print(\n",
+    "        f\"Parameters of last layer at end of training:\\n{jax.tree_leaves(model.layers[-1])}\\n\"\n",
+    "    )"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "6e7c6894-f98e-46cc-9869-afe30ca40828",
+   "metadata": {},
+   "source": [
+    "As we'll see, the parameters of the first layer remain unchanged throughout training. Just the parameters of the last layer are trained."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "id": "b96e16eb-20f9-4219-92b0-fe13d04db9e9",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Parameters of first layer at initialisation:\n",
+      "[DeviceArray([[-0.5500405 ],\n",
+      "             [ 0.67074966],\n",
+      "             [-0.9094155 ],\n",
+      "             [-0.5518596 ],\n",
+      "             [-0.1648488 ],\n",
+      "             [ 0.98241615],\n",
+      "             [-0.9118581 ],\n",
+      "             [ 0.32483125]], dtype=float32), DeviceArray([ 0.8876705 ,  0.4363706 , -0.878813  ,  0.26387787,\n",
+      "             -0.68248963, -0.9517925 , -0.21384668, -0.2857628 ],            dtype=float32)]\n",
+      "\n",
+      "Parameters of first layer at end of training:\n",
+      "[DeviceArray([[-0.5500405 ],\n",
+      "             [ 0.67074966],\n",
+      "             [-0.9094155 ],\n",
+      "             [-0.5518596 ],\n",
+      "             [-0.1648488 ],\n",
+      "             [ 0.98241615],\n",
+      "             [-0.9118581 ],\n",
+      "             [ 0.32483125]], dtype=float32), DeviceArray([ 0.8876705 ,  0.4363706 , -0.878813  ,  0.26387787,\n",
+      "             -0.68248963, -0.9517925 , -0.21384668, -0.2857628 ],            dtype=float32)]\n",
+      "\n",
+      "Parameters of last layer at initialisation:\n",
+      "[DeviceArray([[ 0.33031464,  0.16732198,  0.04151077, -0.01495699,\n",
+      "              -0.00766617,  0.08186949,  0.33581698,  0.13139524]],            dtype=float32), DeviceArray([-0.05869024], dtype=float32)]\n",
+      "\n",
+      "Parameters of last layer at end of training:\n",
+      "[DeviceArray([[-2.8348157 ,  3.3568215 , -0.76763505, -2.1319082 ,\n",
+      "              -0.00766617,  1.3563706 , -1.6294043 ,  0.6064026 ]],            dtype=float32), DeviceArray([-0.10058278], dtype=float32)]\n",
+      "\n"
+     ]
+    }
+   ],
+   "source": [
+    "main()"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "jax0226",
+   "language": "python",
+   "name": "jax0226"
+  },
+  "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.8.12"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}
diff --git a/examples/frozen_layer.py b/examples/frozen_layer.py
deleted file mode 100644
index 611f0253..00000000
--- a/examples/frozen_layer.py
+++ /dev/null
@@ -1,72 +0,0 @@
-import functools as ft
-
-import jax
-import jax.numpy as jnp
-import jax.random as jrandom
-import optax
-
-# Steal these utilities from the previous example
-from filtered_transformations import dataloader, get_data
-
-import equinox as eqx
-
-
-def main(
-    dataset_size=10000,
-    batch_size=256,
-    learning_rate=3e-3,
-    steps=1000,
-    width_size=8,
-    depth=1,
-    seed=5678,
-):
-    data_key, loader_key, model_key = jrandom.split(jrandom.PRNGKey(seed), 3)
-    data = get_data(dataset_size, key=data_key)
-    data = dataloader(data, batch_size=batch_size, key=loader_key)
-
-    # Let's train just the final layer of the MLP, and leave the others frozen.
-
-    model = eqx.nn.MLP(
-        in_size=1, out_size=1, width_size=width_size, depth=depth, key=model_key
-    )
-    # A PyTree (of type MLP) with False on every leaf.
-    filter_spec = jax.tree_map(lambda _: False, model)
-    # Set the leaves of the final layer to True.
-    filter_spec = eqx.tree_at(
-        lambda tree: (tree.layers[-1].weight, tree.layers[-1].bias),
-        filter_spec,
-        replace=(True, True),
-    )
-
-    # We'll still JIT with respect to every array.
-    @eqx.filter_jit
-    @ft.partial(eqx.filter_value_and_grad, filter_spec=filter_spec)
-    def loss(model, x, y):
-        pred_y = jax.vmap(model)(x)
-        return jnp.mean((y - pred_y) ** 2)
-
-    original_model = model  # Keep around to compare to later.
-    optim = optax.sgd(learning_rate)
-    opt_state = optim.init(model)
-    for step, (x, y) in zip(range(steps), data):
-        value, grads = loss(model, x, y)
-        updates, opt_state = optim.update(grads, opt_state)
-        model = eqx.apply_updates(model, updates)
-        print(step, value)
-    print(
-        f"Weights of first layer at initialisation: {jax.tree_flatten(original_model.layers[0])}"
-    )
-    print(
-        f"Weights of first layer at end of training: {jax.tree_flatten(model.layers[0])}"
-    )
-    print(
-        f"Weights of last layer at initialisation: {jax.tree_flatten(original_model.layers[-1])}"
-    )
-    print(
-        f"Weights of last layer at end of training: {jax.tree_flatten(model.layers[-1])}"
-    )
-    return original_model, model
-
-
-if __name__ == "__main__":
-    main()
diff --git a/examples/init_apply.ipynb b/examples/init_apply.ipynb
new file mode 100644
index 00000000..47e76c92
--- /dev/null
+++ b/examples/init_apply.ipynb
@@ -0,0 +1,105 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "14148da3-23fb-480b-ae31-93878cda86fa",
+   "metadata": {},
+   "source": [
+    "# Compatibility with init-apply libraries\n",
+    "\n",
+    "Existing JAX neural network libraries have sometimes followed the \"init/apply\" approach, in which the parameters of a network are initialised with a function `init()`, and then the forward pass through a model is specified with `apply()`. For example [Stax](https://jax.readthedocs.io/en/latest/jax.example_libraries.stax.html) follows this approach.\n",
+    "\n",
+    "As a result, some third-party libraries assume that your model is specified by an `init()` and an `apply()` function, and that the parameters returned from `init()` are all JIT-trace-able and grad-able.\n",
+    "\n",
+    "Equinox can be made to fit with this style very easily, like so."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "05bcc516-c80a-492b-af63-0e5dec66f438",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import equinox as eqx\n",
+    "\n",
+    "\n",
+    "def make_mlp(in_size, out_size, width_size, depth, *, key):\n",
+    "    mlp = eqx.nn.MLP(\n",
+    "        in_size, out_size, width_size, depth, key=key\n",
+    "    )  # insert your model here\n",
+    "    params, static = eqx.partition(mlp, eqx.is_inexact_array)\n",
+    "\n",
+    "    def init_fn():\n",
+    "        return params\n",
+    "\n",
+    "    def apply_fn(_params, x):\n",
+    "        model = eqx.combine(_params, static)\n",
+    "        return model(x)\n",
+    "\n",
+    "    return init_fn, apply_fn"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "5965f2e9-483a-475e-b316-a595daf6eb0f",
+   "metadata": {},
+   "source": [
+    "And that's all there is to it.\n",
+    "\n",
+    "Example usage:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "e25e63c7-3cb6-4727-9c28-0016b03acb87",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import jax\n",
+    "import jax.numpy as jnp\n",
+    "import jax.random as jrandom\n",
+    "\n",
+    "\n",
+    "def main(in_size=2, seed=5678):\n",
+    "    key = jrandom.PRNGKey(seed)\n",
+    "\n",
+    "    init_fn, apply_fn = make_mlp(\n",
+    "        in_size=in_size, out_size=1, width_size=8, depth=1, key=key\n",
+    "    )\n",
+    "\n",
+    "    x = jnp.arange(in_size)  # sample data\n",
+    "    params = init_fn()\n",
+    "    y1 = apply_fn(params, x)\n",
+    "    params = jax.tree_map(lambda p: p + 1, params)  # \"stochastic gradient descent\"\n",
+    "    y2 = apply_fn(params, x)\n",
+    "    assert y1 != y2\n",
+    "\n",
+    "\n",
+    "main()"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "jax0226",
+   "language": "python",
+   "name": "jax0226"
+  },
+  "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.8.12"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}
diff --git a/examples/modules_to_initapply.py b/examples/modules_to_initapply.py
deleted file mode 100644
index 9805774a..00000000
--- a/examples/modules_to_initapply.py
+++ /dev/null
@@ -1,57 +0,0 @@
-##############
-#
-# Existing JAX neural network libraries have sometimes followed the "init/apply"
-# approach, in which the parameters of a network are initialised with `init`, and then
-# the forward pass through a model is specified with `apply`. For example Stax follows
-# this approach.
-#
-# As a corollary, the parameters returned from `init` are sometimes assumed to all be
-# JIT-able or grad-able, e.g. by third-party libraries.
-#
-# In contrast Equinox is more general: it (a) doesn't assume that you necessarily
-# want to take gradients with respect to all your parameters, and (b) doesn't even
-# mandate that all your parameters are JAX arrays.
-#
-# If need be -- e.g. third party library compatibility -- then Equinox can be made to
-# fit this style very easily, like so.
-#
-##############
-
-import jax
-import jax.numpy as jnp
-import jax.random as jrandom
-
-import equinox as eqx
-
-
-def make_mlp(in_size, out_size, width_size, depth, *, key):
-    mlp = eqx.nn.MLP(in_size, out_size, width_size, depth, key=key)
-    params, static = eqx.partition(mlp, eqx.is_inexact_array)
-
-    def init_fn():
-        return params
-
-    def apply_fn(params, x):
-        model = eqx.combine(params, static)
-        return model(x)
-
-    return init_fn, apply_fn
-
-
-def main(in_size=2, seed=5678):
-    key = jrandom.PRNGKey(seed)
-
-    init_fn, apply_fn = make_mlp(
-        in_size=in_size, out_size=1, width_size=8, depth=1, key=key
-    )
-
-    x = jnp.arange(in_size)  # sample data
-    params = init_fn()
-    y1 = apply_fn(params, x)
-    params = jax.tree_map(lambda p: p + 1, params)  # "stochastic gradient descent"
-    y2 = apply_fn(params, x)
-    assert y1 != y2
-
-
-if __name__ == "__main__":
-    main()
diff --git a/examples/train_rnn.ipynb b/examples/train_rnn.ipynb
new file mode 100644
index 00000000..f4b7d76d
--- /dev/null
+++ b/examples/train_rnn.ipynb
@@ -0,0 +1,424 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "b4d4b318-f46e-4d2c-92cd-632027a03632",
+   "metadata": {},
+   "source": [
+    "# Train RNN\n",
+    "\n",
+    "Here we give a complete example of what using Equinox normally looks like day-to-day.\n",
+    "\n",
+    "In this example we'll train an RNN to classify clockwise vs anticlockwise spirals."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "690edba4-62d8-48ff-8ed2-2dd322399525",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import math\n",
+    "\n",
+    "import jax\n",
+    "import jax.lax as lax\n",
+    "import jax.numpy as jnp\n",
+    "import jax.random as jrandom\n",
+    "import optax  # https://github.com/deepmind/optax\n",
+    "\n",
+    "import equinox as eqx"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "2d5423ff-c18b-4394-9d1d-0bb4bfd626d6",
+   "metadata": {},
+   "source": [
+    "We begin by importing the usual libraries, setting up a very simple dataloader, and generating a toy dataset of spirals."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "14f40756-55d6-44f4-b13b-b15db35e15e9",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def dataloader(arrays, batch_size, *, key):\n",
+    "    dataset_size = arrays[0].shape[0]\n",
+    "    assert all(array.shape[0] == dataset_size for array in arrays)\n",
+    "    indices = jnp.arange(dataset_size)\n",
+    "    while True:\n",
+    "        perm = jrandom.permutation(key, indices)\n",
+    "        (key,) = jrandom.split(key, 1)\n",
+    "        start = 0\n",
+    "        end = batch_size\n",
+    "        while end < dataset_size:\n",
+    "            batch_perm = perm[start:end]\n",
+    "            yield tuple(array[batch_perm] for array in arrays)\n",
+    "            start = end\n",
+    "            end = start + batch_size\n",
+    "\n",
+    "\n",
+    "def get_data(dataset_size, *, key):\n",
+    "    t = jnp.linspace(0, 2 * math.pi, 16)\n",
+    "    offset = jrandom.uniform(key, (dataset_size, 1), minval=0, maxval=2 * math.pi)\n",
+    "    x1 = jnp.sin(t + offset) / (1 + t)\n",
+    "    x2 = jnp.cos(t + offset) / (1 + t)\n",
+    "    y = jnp.ones((dataset_size, 1))\n",
+    "\n",
+    "    half_dataset_size = dataset_size // 2\n",
+    "    x1 = x1.at[:half_dataset_size].multiply(-1)\n",
+    "    y = y.at[:half_dataset_size].set(0)\n",
+    "    x = jnp.stack([x1, x2], axis=-1)\n",
+    "\n",
+    "    return x, y"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "f4094e1f-b1b5-4e3b-9191-bf09a5f93dda",
+   "metadata": {},
+   "source": [
+    "Now for our model.\n",
+    "\n",
+    "Purely by way of example, we handle the final adding on of bias ourselves, rather than letting the `linear` layer do it. This is just so we can demonstrate how to use custom parameters in models."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "id": "533bdd6b-d660-4d94-bca9-da80ad0c0a52",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "class RNN(eqx.Module):\n",
+    "    hidden_size: int\n",
+    "    cell: eqx.Module\n",
+    "    linear: eqx.nn.Linear\n",
+    "    bias: jnp.ndarray\n",
+    "\n",
+    "    def __init__(self, in_size, out_size, hidden_size, *, key):\n",
+    "        ckey, lkey = jrandom.split(key)\n",
+    "        self.hidden_size = hidden_size\n",
+    "        self.cell = eqx.nn.GRUCell(in_size, hidden_size, key=ckey)\n",
+    "        self.linear = eqx.nn.Linear(hidden_size, out_size, use_bias=False, key=lkey)\n",
+    "        self.bias = jnp.zeros(out_size)\n",
+    "\n",
+    "    def __call__(self, input):\n",
+    "        hidden = jnp.zeros((self.hidden_size,))\n",
+    "\n",
+    "        def f(carry, inp):\n",
+    "            return self.cell(inp, carry), None\n",
+    "\n",
+    "        out, _ = lax.scan(f, hidden, input)\n",
+    "        # sigmoid because we're performing binary classification\n",
+    "        return jax.nn.sigmoid(self.linear(out) + self.bias)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "eb040df3-58c0-4546-b64d-c25147173e4b",
+   "metadata": {},
+   "source": [
+    "And finally the training loop."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "id": "edef054d-5fdb-491e-8a46-c4d80363d012",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def main(\n",
+    "    dataset_size=10000,\n",
+    "    batch_size=32,\n",
+    "    learning_rate=3e-3,\n",
+    "    steps=200,\n",
+    "    hidden_size=16,\n",
+    "    depth=1,\n",
+    "    seed=5678,\n",
+    "):\n",
+    "    data_key, loader_key, model_key = jrandom.split(jrandom.PRNGKey(seed), 3)\n",
+    "    xs, ys = get_data(dataset_size, key=data_key)\n",
+    "    iter_data = dataloader((xs, ys), batch_size, key=loader_key)\n",
+    "\n",
+    "    model = RNN(in_size=2, out_size=1, hidden_size=hidden_size, key=model_key)\n",
+    "\n",
+    "    @eqx.filter_jit\n",
+    "    @eqx.filter_value_and_grad\n",
+    "    def make_step(model, x, y):\n",
+    "        pred_y = jax.vmap(model)(x)\n",
+    "        # Trains with respect to binary cross-entropy\n",
+    "        return -jnp.mean(y * jnp.log(pred_y) + (1 - y) * jnp.log(1 - pred_y))\n",
+    "\n",
+    "    optim = optax.adam(learning_rate)\n",
+    "    opt_state = optim.init(model)\n",
+    "    for step, (x, y) in zip(range(steps), iter_data):\n",
+    "        loss, grads = make_step(model, x, y)\n",
+    "        loss = loss.item()\n",
+    "        updates, opt_state = optim.update(grads, opt_state)\n",
+    "        model = eqx.apply_updates(model, updates)\n",
+    "        print(f\"{step=}, {loss=}\")\n",
+    "\n",
+    "    pred_ys = jax.vmap(model)(xs)\n",
+    "    num_correct = jnp.sum((pred_ys > 0.5) == ys)\n",
+    "    final_accuracy = (num_correct / dataset_size).item()\n",
+    "    print(f\"{final_accuracy=}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "6498efb9-f3d5-4dfb-b3f8-68f2b42422ec",
+   "metadata": {},
+   "source": [
+    "`eqx.filter_jit` will look at all the arguments passed to `make_step`, and automatically JIT-trace every array and JIT-static everything else. For example the `model` parameters, and the data `x` and `y` will be traced. (However `model.hidden_size` is an integer, so it will be static'd instead.) This is a convenient way to JIT when your arguments are PyTrees, and if you need finer control then `eqx.filter_jit` takes a `filter_spec` argument.\n",
+    "\n",
+    "`eqx.filter_value_and_grad` will calculate the gradient with respect to the first argument (`model`). It will automatically do so for just the floating-point arrays in the `model` PyTree (i.e. its parameters) but not for anything else (e.g. `model.hidden_size`). Once again if you need finer control then `eqx.filter_value_and_grad` takes a `filter_spec` argument.\n",
+    "\n",
+    "_(Note that the above code is written for clarity, not speed. If you want maximum performance then the operations inside `dataloader`, and the `optim.update`, `eqx.apply_updates` operations, should all be brought inside the single JIT operation.)_"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "id": "4c05289e-2bd9-4f69-8a57-62d2f2f8047c",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "step=0, loss=0.695469856262207\n",
+      "step=1, loss=0.6940110921859741\n",
+      "step=2, loss=0.6931639909744263\n",
+      "step=3, loss=0.6961429119110107\n",
+      "step=4, loss=0.6957006454467773\n",
+      "step=5, loss=0.6920372247695923\n",
+      "step=6, loss=0.6938458681106567\n",
+      "step=7, loss=0.6933101415634155\n",
+      "step=8, loss=0.6897573471069336\n",
+      "step=9, loss=0.6940688490867615\n",
+      "step=10, loss=0.6896819472312927\n",
+      "step=11, loss=0.6991991996765137\n",
+      "step=12, loss=0.691841185092926\n",
+      "step=13, loss=0.6914103627204895\n",
+      "step=14, loss=0.7029447555541992\n",
+      "step=15, loss=0.7040634155273438\n",
+      "step=16, loss=0.6862525939941406\n",
+      "step=17, loss=0.700591504573822\n",
+      "step=18, loss=0.6833503246307373\n",
+      "step=19, loss=0.6996323466300964\n",
+      "step=20, loss=0.6967240571975708\n",
+      "step=21, loss=0.6946160793304443\n",
+      "step=22, loss=0.6940494775772095\n",
+      "step=23, loss=0.688583493232727\n",
+      "step=24, loss=0.694453239440918\n",
+      "step=25, loss=0.694945216178894\n",
+      "step=26, loss=0.6929985284805298\n",
+      "step=27, loss=0.6932469606399536\n",
+      "step=28, loss=0.6898558735847473\n",
+      "step=29, loss=0.6925642490386963\n",
+      "step=30, loss=0.6929633617401123\n",
+      "step=31, loss=0.6931982636451721\n",
+      "step=32, loss=0.6919726133346558\n",
+      "step=33, loss=0.6934778690338135\n",
+      "step=34, loss=0.6924077272415161\n",
+      "step=35, loss=0.6937430500984192\n",
+      "step=36, loss=0.6933866739273071\n",
+      "step=37, loss=0.6933223009109497\n",
+      "step=38, loss=0.6918285489082336\n",
+      "step=39, loss=0.6911637783050537\n",
+      "step=40, loss=0.6929216384887695\n",
+      "step=41, loss=0.695015549659729\n",
+      "step=42, loss=0.6941317319869995\n",
+      "step=43, loss=0.6889361143112183\n",
+      "step=44, loss=0.6929722428321838\n",
+      "step=45, loss=0.6959724426269531\n",
+      "step=46, loss=0.6923878788948059\n",
+      "step=47, loss=0.6954452395439148\n",
+      "step=48, loss=0.6928846836090088\n",
+      "step=49, loss=0.6919665932655334\n",
+      "step=50, loss=0.6892142295837402\n",
+      "step=51, loss=0.6910432577133179\n",
+      "step=52, loss=0.689067006111145\n",
+      "step=53, loss=0.7064827680587769\n",
+      "step=54, loss=0.6837577223777771\n",
+      "step=55, loss=0.6927756071090698\n",
+      "step=56, loss=0.6882894039154053\n",
+      "step=57, loss=0.6991082429885864\n",
+      "step=58, loss=0.6848499774932861\n",
+      "step=59, loss=0.6892214417457581\n",
+      "step=60, loss=0.6934791207313538\n",
+      "step=61, loss=0.6962975859642029\n",
+      "step=62, loss=0.7012044191360474\n",
+      "step=63, loss=0.6914117336273193\n",
+      "step=64, loss=0.6984191536903381\n",
+      "step=65, loss=0.6937364339828491\n",
+      "step=66, loss=0.7050725817680359\n",
+      "step=67, loss=0.6893508434295654\n",
+      "step=68, loss=0.6935131549835205\n",
+      "step=69, loss=0.693252444267273\n",
+      "step=70, loss=0.693372368812561\n",
+      "step=71, loss=0.6887596845626831\n",
+      "step=72, loss=0.6937415599822998\n",
+      "step=73, loss=0.6912466287612915\n",
+      "step=74, loss=0.692828893661499\n",
+      "step=75, loss=0.6928187012672424\n",
+      "step=76, loss=0.6929498314857483\n",
+      "step=77, loss=0.6911673545837402\n",
+      "step=78, loss=0.6932731866836548\n",
+      "step=79, loss=0.6969736218452454\n",
+      "step=80, loss=0.692180871963501\n",
+      "step=81, loss=0.6920971870422363\n",
+      "step=82, loss=0.6940006613731384\n",
+      "step=83, loss=0.6947032809257507\n",
+      "step=84, loss=0.6884186863899231\n",
+      "step=85, loss=0.6916593313217163\n",
+      "step=86, loss=0.6919583678245544\n",
+      "step=87, loss=0.6955068111419678\n",
+      "step=88, loss=0.6922551989555359\n",
+      "step=89, loss=0.7009901404380798\n",
+      "step=90, loss=0.6891013383865356\n",
+      "step=91, loss=0.6935603618621826\n",
+      "step=92, loss=0.69370436668396\n",
+      "step=93, loss=0.685250997543335\n",
+      "step=94, loss=0.6848466396331787\n",
+      "step=95, loss=0.6973309516906738\n",
+      "step=96, loss=0.702495813369751\n",
+      "step=97, loss=0.6917577981948853\n",
+      "step=98, loss=0.6930680871009827\n",
+      "step=99, loss=0.6879289150238037\n",
+      "step=100, loss=0.6978393197059631\n",
+      "step=101, loss=0.6950803399085999\n",
+      "step=102, loss=0.6934637427330017\n",
+      "step=103, loss=0.6898033618927002\n",
+      "step=104, loss=0.6917957067489624\n",
+      "step=105, loss=0.6918165683746338\n",
+      "step=106, loss=0.6925389766693115\n",
+      "step=107, loss=0.6907787322998047\n",
+      "step=108, loss=0.6930015087127686\n",
+      "step=109, loss=0.6905994415283203\n",
+      "step=110, loss=0.6903047561645508\n",
+      "step=111, loss=0.6895794868469238\n",
+      "step=112, loss=0.6919896602630615\n",
+      "step=113, loss=0.6929000616073608\n",
+      "step=114, loss=0.6861156225204468\n",
+      "step=115, loss=0.6916525363922119\n",
+      "step=116, loss=0.6891536712646484\n",
+      "step=117, loss=0.6949775218963623\n",
+      "step=118, loss=0.6913626790046692\n",
+      "step=119, loss=0.6878023147583008\n",
+      "step=120, loss=0.6889688968658447\n",
+      "step=121, loss=0.6912539601325989\n",
+      "step=122, loss=0.688143253326416\n",
+      "step=123, loss=0.6874911785125732\n",
+      "step=124, loss=0.6872292757034302\n",
+      "step=125, loss=0.6878771185874939\n",
+      "step=126, loss=0.6867557764053345\n",
+      "step=127, loss=0.6870120167732239\n",
+      "step=128, loss=0.684597373008728\n",
+      "step=129, loss=0.6766145825386047\n",
+      "step=130, loss=0.6783543825149536\n",
+      "step=131, loss=0.6876450777053833\n",
+      "step=132, loss=0.6898150444030762\n",
+      "step=133, loss=0.6808171272277832\n",
+      "step=134, loss=0.6852248907089233\n",
+      "step=135, loss=0.6800845265388489\n",
+      "step=136, loss=0.6718021631240845\n",
+      "step=137, loss=0.6681873798370361\n",
+      "step=138, loss=0.6753413677215576\n",
+      "step=139, loss=0.671585202217102\n",
+      "step=140, loss=0.6637449264526367\n",
+      "step=141, loss=0.6678113341331482\n",
+      "step=142, loss=0.6511460542678833\n",
+      "step=143, loss=0.650525689125061\n",
+      "step=144, loss=0.6523746252059937\n",
+      "step=145, loss=0.6346812844276428\n",
+      "step=146, loss=0.6078983545303345\n",
+      "step=147, loss=0.6069947481155396\n",
+      "step=148, loss=0.6328916549682617\n",
+      "step=149, loss=0.5921869277954102\n",
+      "step=150, loss=0.5772275328636169\n",
+      "step=151, loss=0.5446708798408508\n",
+      "step=152, loss=0.545225977897644\n",
+      "step=153, loss=0.5175529718399048\n",
+      "step=154, loss=0.49300363659858704\n",
+      "step=155, loss=0.41911324858665466\n",
+      "step=156, loss=0.390226811170578\n",
+      "step=157, loss=0.36768990755081177\n",
+      "step=158, loss=0.30429911613464355\n",
+      "step=159, loss=0.2722632884979248\n",
+      "step=160, loss=0.2577134966850281\n",
+      "step=161, loss=0.1969938725233078\n",
+      "step=162, loss=0.13784314692020416\n",
+      "step=163, loss=0.1397821009159088\n",
+      "step=164, loss=0.11436247825622559\n",
+      "step=165, loss=0.08968592435121536\n",
+      "step=166, loss=0.07026629149913788\n",
+      "step=167, loss=0.06773999333381653\n",
+      "step=168, loss=0.0786045491695404\n",
+      "step=169, loss=0.04946307837963104\n",
+      "step=170, loss=0.039137713611125946\n",
+      "step=171, loss=0.03697769343852997\n",
+      "step=172, loss=0.03442113846540451\n",
+      "step=173, loss=0.02817709371447563\n",
+      "step=174, loss=0.02260049805045128\n",
+      "step=175, loss=0.0208757221698761\n",
+      "step=176, loss=0.0238012857735157\n",
+      "step=177, loss=0.019132442772388458\n",
+      "step=178, loss=0.016685867682099342\n",
+      "step=179, loss=0.015729248523712158\n",
+      "step=180, loss=0.01557326503098011\n",
+      "step=181, loss=0.014118168503046036\n",
+      "step=182, loss=0.012221714481711388\n",
+      "step=183, loss=0.012281971052289009\n",
+      "step=184, loss=0.011437132023274899\n",
+      "step=185, loss=0.010124631226062775\n",
+      "step=186, loss=0.010237202048301697\n",
+      "step=187, loss=0.009541159495711327\n",
+      "step=188, loss=0.009026993066072464\n",
+      "step=189, loss=0.009005479514598846\n",
+      "step=190, loss=0.009186260402202606\n",
+      "step=191, loss=0.008103055879473686\n",
+      "step=192, loss=0.007458947598934174\n",
+      "step=193, loss=0.0075416406616568565\n",
+      "step=194, loss=0.007497536949813366\n",
+      "step=195, loss=0.006723443046212196\n",
+      "step=196, loss=0.00683201476931572\n",
+      "step=197, loss=0.006709215696901083\n",
+      "step=198, loss=0.006567368749529123\n",
+      "step=199, loss=0.006062147207558155\n",
+      "final_accuracy=1.0\n"
+     ]
+    }
+   ],
+   "source": [
+    "main()  # All right, let's run the code."
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "jax0226",
+   "language": "python",
+   "name": "jax0226"
+  },
+  "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.8.12"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}
diff --git a/examples/train_rnn.py b/examples/train_rnn.py
deleted file mode 100644
index 5daf372f..00000000
--- a/examples/train_rnn.py
+++ /dev/null
@@ -1,92 +0,0 @@
-##############
-#
-# This example trains an RNN to classify clockwise versus anticlockwise spirals.
-#
-#############
-
-import math
-
-import jax
-import jax.lax as lax
-import jax.numpy as jnp
-import jax.random as jrandom
-import optax
-
-# Borrow dataloader from other example
-from filtered_transformations import dataloader
-
-import equinox as eqx
-
-
-# Toy data of spirals.
-# We aim to classify clockwise versus anticlockwise.
-def get_data(dataset_size, *, key):
-    t = jnp.linspace(0, 2 * math.pi, 16)
-    offset = jrandom.uniform(key, (dataset_size, 1), minval=0, maxval=2 * math.pi)
-    x1 = jnp.sin(t + offset) / (1 + t)
-    x2 = jnp.cos(t + offset) / (1 + t)
-    y = jnp.ones((dataset_size, 1))
-
-    half_dataset_size = dataset_size // 2
-    x1 = x1.at[:half_dataset_size].multiply(-1)
-    y = y.at[:half_dataset_size].set(0)
-    x = jnp.stack([x1, x2], axis=-1)
-
-    return x, y
-
-
-class RNN(eqx.Module):
-    hidden_size: int
-    cell: eqx.Module
-    linear: eqx.nn.Linear
-
-    def __init__(self, in_size, out_size, hidden_size, *, key):
-        ckey, lkey = jrandom.split(key)
-        self.hidden_size = hidden_size
-        self.cell = eqx.nn.GRUCell(in_size, hidden_size, key=ckey)
-        self.linear = eqx.nn.Linear(hidden_size, out_size, key=lkey)
-
-    def __call__(self, input):
-        hidden = jnp.zeros((self.hidden_size,))
-
-        def f(carry, inp):
-            return self.cell(inp, carry), None
-
-        out, _ = lax.scan(f, hidden, input)
-        return jax.nn.sigmoid(self.linear(out))
-
-
-def main(
-    dataset_size=10000,
-    batch_size=32,
-    learning_rate=3e-3,
-    steps=200,
-    hidden_size=16,
-    depth=1,
-    seed=5678,
-):
-    data_key, loader_key, model_key = jrandom.split(jrandom.PRNGKey(seed), 3)
-    data = get_data(dataset_size, key=data_key)
-    data = dataloader(data, batch_size=batch_size, key=loader_key)
-
-    model = RNN(in_size=2, out_size=1, hidden_size=hidden_size, key=model_key)
-
-    @eqx.filter_jit
-    @eqx.filter_value_and_grad
-    def loss(model, x, y):
-        pred_y = jax.vmap(model)(x)
-        # Trains with respect to binary cross-entropy
-        return -jnp.mean(y * jnp.log(pred_y) + (1 - y) * jnp.log(1 - pred_y))
-
-    optim = optax.adam(learning_rate)
-    opt_state = optim.init(model)
-    for step, (x, y) in zip(range(steps), data):
-        value, grads = loss(model, x, y)
-        updates, opt_state = optim.update(grads, opt_state)
-        model = eqx.apply_updates(model, updates)
-        print(step, value)
-    return value
-
-
-if __name__ == "__main__":
-    main()
diff --git a/mkdocs.yml b/mkdocs.yml
new file mode 100644
index 00000000..c8be9ff2
--- /dev/null
+++ b/mkdocs.yml
@@ -0,0 +1,112 @@
+theme:
+    name: material
+    features:
+        - navigation.sections  # Sections are included in the navigation on the left.
+        - toc.integrate  # Table of contents is integrated on the left; does not appear separately on the right.
+        - header.autohide  # header disappears as you scroll
+    palette:
+        # Light mode / dark mode
+        # We deliberately don't automatically use `media` to check a user's preferences. We default to light mode as
+        # (a) it looks more professional, and (b) is more obvious about the fact that it offers a (dark mode) toggle.
+        - scheme: default
+          primary: white
+          accent: amber
+          toggle:
+             icon: material/weather-night
+             name: Switch to dark mode
+        - scheme: slate
+          primary: black
+          accent: amber
+          toggle:
+             icon: material/weather-sunny
+             name: Switch to light mode
+    icon:
+        repo: fontawesome/brands/github  # GitHub logo in top right
+        logo: "material/circle-opacity"  # Equinox logo in top left
+    favicon: "_static/favicon.png"
+    custom_dir: "docs/_overrides"  # Overriding part of the HTML
+
+    # These additions are my own custom ones, having overridden a partial.
+    twitter_name: "@PatrickKidger"
+    twitter_url: "https://twitter.com/PatrickKidger"
+
+site_name: Equinox
+site_description: The documentation for the Equinox software library.
+site_author: Patrick Kidger
+site_url: https://docs.kidger.site/equinox
+
+repo_url: https://github.com/patrick-kidger/equinox
+repo_name: patrick-kidger/equinox
+edit_uri: ""  # No edit button, as some of our pages are in /docs and some in /examples via symlink, so it's impossible for them all to be accurate
+
+strict: true  # Don't allow warnings during the build process
+
+# Enable if LaTeX is needed
+# extra_javascript: 
+#     # TODO: Switch to MathJax v3: https://squidfunk.github.io/mkdocs-material/reference/mathjax
+#     - https://cdnjs.cloudflare.com/ajax/libs/mathjax/2.7.7/MathJax.js?config=TeX-AMS_CHTML
+
+extra_css:
+    - _static/custom_css.css
+
+markdown_extensions:
+# Enable if LaTeX is needed
+#     - pymdownx.arithmatex  # Render LaTeX via MathJax
+    - pymdownx.superfences  # Seems to enable syntax highlighting when used with the Material theme.
+    - pymdownx.details  # Allowing hidden expandable regions denoted by ???
+    - pymdownx.snippets:  # Include one Markdown file into another
+        base_path: docs
+    - admonition
+    - toc:
+        permalink: "¤"  # Adds a clickable permalink to each section heading
+        toc_depth: 4
+
+plugins:
+    - search  # default search plugin; needs manually re-enabling when using any other plugins
+    - autorefs  # Cross-links to headings
+    - include_exclude_files:
+        include:
+            - ".htaccess"
+        exclude:
+            - "_overrides"
+    - mknotebooks  # Jupyter notebooks
+    - mkdocstrings:
+        handlers:
+            python:
+                setup_commands:
+                    - import pytkdocs_tweaks
+                    - pytkdocs_tweaks.main()
+
+                selection:
+                    inherited_members: true  # Allow looking up inherited methods
+                rendering:
+                    show_root_heading: true  # actually display anything at all...
+                    show_root_full_path: true  # display "diffrax.asdf" not just "asdf"
+                    show_if_no_docstring: true
+                    show_signature_annotations: true
+                    show_source: false  # don't include source code
+                    members_order: source  # order methods according to their order of definition in the source code, not alphabetical order
+                    heading_level: 4
+
+nav:
+    - 'index.md'
+    - 'all-of-equinox.md'
+    - Examples:
+        - Train RNN: 'examples/train_rnn.ipynb'
+        - Freezing parameters: 'examples/frozen_layer.ipynb'
+        - Compatibility with init-apply libraries: 'examples/init_apply.ipynb'
+    - Full API:
+        - 'api/module.md'
+        - Neural networks:
+            - 'api/nn/linear.md'
+            - 'api/nn/conv.md'
+            - 'api/nn/rnn.md'
+            - 'api/nn/composed.md'
+            - 'api/nn/dropout.md'
+        - Filtering:
+            - 'api/filtering/filter-operations.md'
+            - 'api/filtering/filter-functions.md'
+            - 'api/filtering/filtered-transformations.md'
+        - 'api/utilities.md'
+    - Misc:
+        - 'citation.md'
diff --git a/setup.py b/setup.py
index 5365a759..d8861e68 100644
--- a/setup.py
+++ b/setup.py
@@ -44,9 +44,9 @@
     "Topic :: Scientific/Engineering :: Mathematics",
 ]
 
-python_requires = "~=3.7"
+python_requires = "~=3.8"
 
-install_requires = ["jax>=0.2.18", "typing_extensions"]
+install_requires = ["jax>=0.2.26"]
 
 setuptools.setup(
     name=name,


Test patch:
diff --git a/.github/workflows/run_tests.yml b/.github/workflows/run_tests.yml
index 39c73e85..990b981f 100644
--- a/.github/workflows/run_tests.yml
+++ b/.github/workflows/run_tests.yml
@@ -46,7 +46,7 @@ jobs:
     needs: [ check_version ]
     strategy:
       matrix:
-        python-version: [ 3.7, 3.8, 3.9 ]
+        python-version: [ 3.8, 3.9 ]
         os: [ ubuntu-latest, macOS-latest ]
       fail-fast: false
     runs-on: ${{ matrix.os }}
diff --git a/tests/test_examples.py b/tests/test_examples.py
deleted file mode 100644
index 0df9d4b5..00000000
--- a/tests/test_examples.py
+++ /dev/null
@@ -1,62 +0,0 @@
-import pathlib
-import sys
-
-import jax.numpy as jnp
-
-
-_here = pathlib.Path(__file__).resolve().parent
-sys.path.append(str(_here / ".." / "examples"))
-
-import build_model
-import filtered_transformations
-import frozen_layer
-import modules_to_initapply
-import train_rnn
-
-
-def test_build_model():
-    build_model.main()
-
-
-def test_frozen_layer():
-    original_model, model = frozen_layer.main()
-    assert jnp.all(original_model.layers[0].weight == model.layers[0].weight)
-    assert jnp.all(original_model.layers[0].bias == model.layers[0].bias)
-    assert jnp.any(original_model.layers[-1].weight != model.layers[-1].weight)
-    assert jnp.any(original_model.layers[-1].weight != model.layers[-1].weight)
-
-
-def test_filtered_transformations():
-    loss = filtered_transformations.main()
-    assert loss < 0.1
-
-
-def test_train_rnn():
-    loss = train_rnn.main()
-    assert loss < 0.01
-
-
-def test_modules_to_initapply():
-    modules_to_initapply.main()
-
-
-def test_readme():
-    with open(_here / ".." / "README.md") as f:
-        program = []
-        maybe_program = False
-        start_program = False
-        for line in f.readlines():
-            if "```python" in line:
-                maybe_program = True
-            elif maybe_program:
-                maybe_program = False
-                if "import equinox as eqx" in line:
-                    program.append(line)
-                    start_program = True
-            elif start_program:
-                if "```" in line:
-                    exec("\n".join(program), dict())
-                    program = []
-                    start_program = False
-                else:
-                    program.append(line)