Equinox v0.10.2

This release has lots of examples and bugfixes from several new contributors!

Features

• eqx.nn.{Linear, MLP} now support the string "scalar" for their input and output sizes, to produce an array of shape () rather than an array of shape (1,).
• Added equinox.internal.scan for a checkpointed scan implementation. (It'd be interesting to see this used for an optimally-checkpointed scan-over-layers in an LLM?)

Documentation

• Much nicer examples! Big thanks to:
  • @Artur-Galstyan for contributing a CNN-on-MNIST example;
  • to @j5b for contributing a BERT example;
  • to @Benjamin-Walker for contributing a U-Net example.

Bugfixes

• eqx.filter_closure_convert and eqx.internal.while_loop now work with tree-math.
• Improved numerical stability of MultiheadAttention, and fixed it producing NaNs in fully-masked case. (Thanks @j5b!)
• Fixed (the deprecated, but still) deterministic=True being ignored in MultiheadAttention. (Thanks @mk-0!)
• __new__ can now be overridden in subclasses of eqx.Module. (Thanks @ASEM000!)

Misc

• Now using ruff and pyright. (No longer using flake8 or isort.)
• Modules are now private-by-default, e.g. equinox._jit instead of equinox.jit. If you're broken by this change then you should make sure to import from the public interface: e.g. equinox.filter_jit instead of equinox._jit.filter_jit.
• equinox.internal.while_loop(..., kind="checkpointed") now supports readable buffers.
• eqx.filter_vmap now supports all-Nones in in_axes. (Thanks @RaderJason!)

Full Changelog: v0.10.1...v0.10.2