Equinox v0.5.2

Autogenerated release notes as follows:

What's Changed

• Fixed GroupNorm raising a spurious runtime error by @patrick-kidger in #106

Full Changelog: v0.5.1...v0.5.2

Equinox v0.5.1

This release:

• Adds equinox.nn.GroupNorm.
• Adds support for grouped convolutions and transposed convolutions, e.g. equinox.nn.Conv2d(..., groups=...). (Thanks @jatentaki!)
• Fixes exceptions raised by tree_deserialise_leaves having no message.
• Fixes a few documentation issues. (Thanks @jvmncs!)

---

Autogenerated release notes as follows:

What's Changed

• Minor doc tweaks for filter_{vmap,pmap} by @patrick-kidger in #84
• Updated examples to v0.5.0 by @patrick-kidger in #86
• Doc fix for nn.Pool by @patrick-kidger in #89
• fix two small typos in documentation by @jvmncs in #91
• Fixed uninformative errors when deserialising by @patrick-kidger in #97
• Implement num_feature_groups for Conv by @jatentaki in #100
• Tweak docs by @patrick-kidger in #101
• Added GroupNorm by @patrick-kidger in #104
• Bump version by @patrick-kidger in #103

New Contributors

• @jvmncs made their first contribution in #91
• @jatentaki made their first contribution in #100

Full Changelog: v0.5.0...v0.5.1

Equinox v0.5.0

This is a big update.

Exciting new features!

• Added filter_vmap.

  • This can be used to create ensembles of models.
  • (Closes #65.)

• Added filter_pmap.

  • (Closes #65.)

• Added pooling layers:

  • eqx.nn.Pool
  • eqx.nn.AvgPool1d
  • eqx.nn.AvgPool2d
  • eqx.nn.AvgPool3d
  • eqx.nn.MaxPool1d
  • eqx.nn.MaxPool2d
  • eqx.nn.MaxPool3d
  • (Closes #59.)
  • (Thanks to @Benjamin-Walker for implementing this.)

• Added tree_serialise_leaves and tree_deserialise_leaves.

  • This can be used to save and load models to file.
  • (Closes #46.)
  • (Thanks to @jaschau for helpful discussions on this.)

• Added tree_inference, as a convenience for toggling all inference flags through a model.

Refactoring for nicer APIs

• filter_{jit,grad,value_and_grad} now have an easier-to-use API for specifying which arguments have what behaviour.

  • Instead of having to specify (args, kwargs) as a single PyTree, then you can specify a default, args, kwargs separately. In particular this avoids doing messy stuff like filter_spec=((...), {}) when you had no kwargs.
  • You no longer have to match up the filter specification for args and kwargs against their runtime values. Both the runtime values, and the filter specification, are matched up against the function signature.
    e.g. you can do filter_jit(lambda x: x, kwargs=dict(x=True))(1), using a keyword argument at JIT-time and a positional argument at call time.
  • Currying is available: both filter_jit(fun) and filter_jit(default=...)(fun) will work.
  • The old API is still available for backward compatibility, of course.
  • (Closes #48.)

• tree_at can now replace subtrees, and not just leaves.

  • (Closes #47.)

• filter, partition now support an is_leaf argument.

  • (Closes #68.)

Miscellaneous

• Calling filter_jit(filter_grad(fun)) twice will no longer lead to unnecessary recompilation: the second filter_grad(fun) instance will be a PyTree that looks like the first filter_grad(fun) instance, and thus we won't get any recompilation.
  • This is actually an improvement over standard JAX! See https://github.com/google/jax/discussions/10284.

Full Changelog: v0.4.0...v0.5.0

Equinox v0.4.0

A new minor release as there's a few minor breaking changes:

• Some of the projections in MultiheadAttention no longer have a bias by default (#60)
• equinox.experimental.{get_state,set_state,BatchNorm} now raises RuntimeErrors for many things; this is to match a change in how jax.experimental.host_callback raises errors in jaxlib>=0.3.5. (#63)

Besides this, a couple of more exciting (?) things:

• equinox.tree_pformat (which is used when printing equinox.Modules) now pretty-prints results much more neatly. (#62)
• equinox.experimental.{get_state,BatchNorm,SpectralNorm} are now substantially faster when run in inference mode. (#61)

Both of which sound pretty minor but both of which were technically really interesting to implement ;)

---

The pull requests in this release were:

• Linear doc tweaks by @patrick-kidger in #58
• Remove default bias in MultiheadAttention by @patrick-kidger in #60
• Improvements to stateful by @patrick-kidger in #61
• Improved pretty-printing by @patrick-kidger in #62
• Bump version by @patrick-kidger in #63

Full Changelog: v0.3.2...v0.4.0

Equinox v0.3.2

Autogenerated release notes as follows:

What's Changed

• Version 0.3.2: SpectralNorm and more! by @patrick-kidger in #55

Full Changelog: v0.3.1...v0.3.2

Equinox v0.3.1

Autogenerated release notes as follows:

What's Changed

• Fix function names when logging compilation progress by @marcelroed in #52

New Contributors

• @marcelroed made their first contribution in #52

Full Changelog: v0.3.0...v0.3.1

Equinox v0.3.0

Three main things in this release.

1. equinox.experimental.BatchNorm. Hurrah, that's nice to have.
2. Very interesting from a technical point of view: stateful operations. In this case, equinox.experimental.{get_state, set_state, StateIndex}. These are the technology used to update the statistics of BatchNorm without requiring the user to faff around outputting the model themselves. They work by wrapping jax.experimental.host_callback.call to save and load external state on demand. Which is pretty magic, so these should really be used sparingly...
3. Removed several old pieces of deprecated functionality: equinox.jitf and so on.

---

Autogenerated release notes as follows:

What's Changed

• Score based example by @patrick-kidger in #42
• minor doc fixes by @patrick-kidger in #43
• Version 0.3.0 -- BatchNorm and stateful by @patrick-kidger in #50

Full Changelog: v0.2.2...v0.3.0

Equinox v0.2.2

Added several new layers:

• LayerNorm
• MultiheadAttention
• ConvTranspose, ConvTranspose1d, ConvTranspose2d, ConvTranspose3d
• Embedding

(Thanks to @andyehrenberg for implementing much of this, and to @lucidrains for reviewing the implementation of attention.)

Autogenerated release notes as follows:

What's Changed

• ConvTranspose layers, MultiheadAttention, lookup embeddings, LayerNorm by @andyehrenberg in #34
• Tidied; simplified; generalised ConvTranspose implementation. by @patrick-kidger in #40
• Attention, Transposed Convolutions, Embeddings, LayerNorm by @patrick-kidger in #38

Full Changelog: v0.2.1...v0.2.2

Equinox v0.2.1

Autogenerated release notes as follows:

What's Changed

• Added automated releases. by @patrick-kidger in #39
• Added pretty-printing for Modules by @patrick-kidger in #41

Full Changelog: v0.2.0...v0.2.1

Equinox v0.2.0

First release using GitHub releases. We'll be using this to serve as a changelog.

This bumps the minor version 0.1.6 -> 0.2.0 so this is a breaking release. (Admittedly for something pretty minor. See the autogenerated changelog below.)

What's Changed

• Fixed filter_grad(has_aux=True) returning arguments in the wrong order. by @patrick-kidger in #36

Full Changelog: v0.1.6...v0.2.0