Release Equinox v0.10.3 · patrick-kidger/equinox

Features

Added equinox.nn.{State, StateIndex}. This has been one of the longest-requested features for Equinox: we now have proper stateful operations! (In a carefully-controlled way -- see the new stateful docs and the new stateful example.)
As an application of these new stateful operations: added equinox.nn.{BatchNorm, SpectralNorm}, which have graduated from experimental! Note that these have a slightly different API to their previous experimental versions.
Added equinox.Partial, which is a tidied-up version of jax.tree_util.Partial.
equinox.filter_{jit, pmap} are now compatibile with ahead-of-time compilation. (#325)
equinox.nn.LayerNorm now supports use_weight and use_bias arguments to disable each individually. This is reflecting the fact that many modern transformer architectures now use layer normalisation without bias. (#310; thanks @lockwo!)
Added equinox.internal.{AbstractVar, AbstractClassVar} to denote abstract instance attributes and abstract class attributes respectively. (Analogous to abc.abstractmethod denoting abstract methods.) The downstream scientific ecosystem is making heavy use of abstract base classes (e.g. all the ABCs in Diffrax) and these have turned out to be a really useful feature. See this docstring for more details. Right now these are an undocumented internal-only feature, but we could plausibly spin these out into their own library.

equinox.nn.Conv should now be compatible with disabled rank promotion (#308; thanks @lockwo!)
equinox.internal.loop should now be compatible with jax.experimental.xmap (patrick-kidger/diffrax#246)
Normalisation layers should now be tolerance to floating-point inaccuracies that occasionally produce negative variances. (#314; thanks @anh-tong!)
The BERT example now has fixed dropout behaviour (#316; thanks @j5b!)
Some doc fixes (#303; thanks @RaderJason!)

Everything in equinox.experimental.* has been removed. See the new stateful functionality described above.

Full Changelog: v0.10.2...v0.10.3