Better errors
Equinox now includes several additional checks to guard against various bugs. If you have a new error, then this is probably an indication that your code always had a silent bug, and should be updated.
- `eqx.nn.LayerNorm` now correctly validates that the shape of its input. This was a common cause of silent bugs. (Thanks dlwh for pointing this one out!)
- Equinox now prints out a warning if you supply both `__init__` and `__post_init__` -- the former actually overwrites the latter. (This is normal Python dataclass behaviour, but probably unexpected.)
- Equinox now prevents you from assigning Module attributes with a bound method of your current instance, e.g.
python
class Model(eqx.Module):
foo: Callable
def __init__(self):
self.foo = self.bar
def bar(self):
...
Otherwise, you end up with two different copies of your model! One at `self`, the other at `self.foo.__self__`. (The latter being in the bound method.)
- `eqx.tree_at` now gives a better error message if you use it try to and update something that isn't a PyTree leaf. (Thanks LouisDesdoigts!)
API changes
These should all be very minor.
- **Breaking change:** `eqx.nn.StateIndex` now takes the initial value, rather than a function that returns the initial value.
- **Breaking change:** If using `eqx.field(converter=...)`, then conversion now happens before `__post_init__`, rather than after it.
- Prefer `eqx.nn.make_with_state` over `eqx.nn.State`. The latter will continue to work, but the former is more memory-efficient. (It deletes the original copy of the initial state.)
- Prefer `eqx.nn.inference_mode` over `eqx.tree_inference`. The latter will continue to exist for backward compatibility. These are the same function, this is really just a matter of moving it into the `eqx.nn` namespace where it always belonged.
Sharing layers
Equinox now supports sharing a layer between multiple parts of your model! This has probably been our longest-requested feature -- in large part because of how intractable it seemed. Equinox models are Py*Trees*, not Py*DAGs*, so how exactly are we supposed to have two different parts of our model point at the same layer?
The answer turned out to be the following -- in this example, we're reusing the embedding weight matrix between the initial embedding layer, and the final readout layer, of a language model.
python
class LanguageModel(eqx.Module):
shared: eqx.nn.Shared
def __init__(self):
embedding = eqx.nn.Embedding(...)
linear = eqx.nn.Linear(...)
These two weights will now be tied together.
where = lambda embed_and_lin: embed_and_lin[1].weight
get = lambda embed_and_lin: embed_and_lin[0].weight
self.shared = eqx.nn.Shared((embedding, linear), where, get)
def __call__(self, tokens):
Expand back out so we can evaluate these layers.
embedding, linear = self.shared()
assert embedding.weight is linear.weight same parameter!
Now go ahead and evaluate your language model.
...
here, `eqx.nn.Shared(...)` simply removes all of the nodes at `where`, so that we don't have two separate copies. Then when it is called at `self.shared()`, it puts them back again. Note that this isn't a copy and doesn't incur any additional memory overhead; this all happens at the Python level, not the XLA level.
(The curious may like to take a look at the implementation in `equinox/nn/_shared.py`, which turned out to be very simple.)
_On a meta level, I'd like to comment that I'm quite proud of having gotten this one in! It means that Equinox now supports both stateful layers and shared layers, which have always been the two pieces that seemed out of reach when using something as simple as PyTrees to represent models. But it turns out that PyTrees really are all you need. :D_
Other changes
Documentation
- Many documentation fixes courtesy of colehaus and Artur-Galstyan!
- Added two new examples to the documentation. Thank you to ahmed-alllam for both of them!
- Deep convolutional GAN
- Vision Transformer
- Added an FAQ entry on comparisons between Equinox and PyTorch/Keras/Julia/Flax. It's a common enough question that should probably have had an answer before now.
- Added an FAQ entry on debugging recompilation.
Features
- Added `eqx.filter_checkpoint`, which as you might expect is a filtered version of `jax.checkpoint`. (Thanks dlwh!)
- Added `eqx.Module.__check_init__`. This is run in a similar fashion to `__post_init__`; see the documentation. This can be used to check that invariants of your module hold after initialisation.
- Added support for vmap'ing stateful layers, by adding `eqx.nn.State.{substate, update}`. This offers a way to subset or update a `State` object, that so only the parts of it that need to be vmap'd are passed in. See the stateful documentation for an example of how to do this.
- Runtime error should now produce much more readable results, without any of the terrifying `INTERNAL: Generated function failed: CpuCallback error` stuff! This clean-up of the runtime error message is done by `eqx.filter_jit`, so that will need to be your top-level way of JIT'ing your computation.
- Added `eqx.nn.StatefulLayer` -- this is (only!) with `eqx.nn.Sequential`, to indicate that the layer should be called with `x, state`, and not just `x`. If you would like a custom stateful layer to be compatible with `Sequential` then go ahead and subclass this, and potentially implement the `is_stateful` method. (Thanks paganpasta!)
- The forward pass of each `eqx.nn.*` layer is now wrapped in a `jax.named_scope`, for better debugging experience. (Thanks ahmed-alllam!)
- `eqx.module_update_wrapper` no longer requires a second argument; it will look at the `__wrapped__` attribute of its first argument.
- Added `eqx.internal.closure_to_pytree`, for... you guessed it, turning function closures into PyTrees. The closed-over variables are treated as the subnodes in the PyTree. This will operate recursively so that closed-over closures will themselves become PyTrees, etc. Note that closed-over global variables are not included.
Bugfixes
- `eqx.tree_{serialise,deserialise}_leaves` now correctly handle unusual NumPy scalars, like `bfloat16`. (Thanks colehaus!)
- `eqx.field(metadata=...)` arguments no longer results in the `static`/`converter` arguments being ignored. (Thanks mjo22!)
- `eqx.filter_custom_vjp` now supports residuals that are not arrays. (The residuals are the pytree that is passed between the forward and backward pass.)
- `eqx.{AbstractVar,AbstractClassVar}` should now support overriden generics in subclasses. That is, something like this:
python
class Foo(eqx.Module):
x: eqx.AbstractVar[list[str]]
class Bar(Foo):
x: list[str]
should no longer raise spurious errors under certain conditions.
- `eqx.internal.while_loop` now supports using custom (non-Equinox) pytrees in the state.
- `eqx.tree_check` no longer raises some false positives.
- Equinox modules now support `__init_subclass__` with additional class creation kwargs. (Thanks ASEM000, Roger-luo!)
New Contributors
* homerjed made their first contribution in https://github.com/patrick-kidger/equinox/pull/445
* LouisDesdoigts made their first contribution in https://github.com/patrick-kidger/equinox/pull/460
* knyazer made their first contribution in https://github.com/patrick-kidger/equinox/pull/474
**Full Changelog**: https://github.com/patrick-kidger/equinox/compare/v0.10.11...v0.11.0