Inox

- Feature: added `equinox.filter_eval_shape`.
- Feature: `equinox.filter_pmap(out=...)` now supports callable arguments.
- Upgrade: `equinox.filter_vmap(out=...)` now properly supports callable arguments. (Previously they were experimental. No part of the API has changed.)
- Bugfix for passing PyTrees to `equinox.filter_{vmap,pmap}(out=...)` not working. (Thanks jatentaki!)

**Full Changelog**: https://github.com/patrick-kidger/equinox/compare/v0.5.3...v0.5.4

Autogenerated release notes as follows:

What's Changed
* Fix pmap named axes by jatentaki in https://github.com/patrick-kidger/equinox/pull/113


**Full Changelog**: https://github.com/patrick-kidger/equinox/compare/v0.5.2...v0.5.3

Autogenerated release notes as follows:

What's Changed
* Fixed GroupNorm raising a spurious runtime error by patrick-kidger in https://github.com/patrick-kidger/equinox/pull/106


**Full Changelog**: https://github.com/patrick-kidger/equinox/compare/v0.5.1...v0.5.2

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 https://github.com/patrick-kidger/equinox/pull/84
* Updated examples to v0.5.0 by patrick-kidger in https://github.com/patrick-kidger/equinox/pull/86
* Doc fix for `nn.Pool` by patrick-kidger in https://github.com/patrick-kidger/equinox/pull/89
* fix two small typos in documentation by jvmncs in https://github.com/patrick-kidger/equinox/pull/91
* Fixed uninformative errors when deserialising by patrick-kidger in https://github.com/patrick-kidger/equinox/pull/97
* Implement num_feature_groups for Conv by jatentaki in https://github.com/patrick-kidger/equinox/pull/100
* Tweak docs by patrick-kidger in https://github.com/patrick-kidger/equinox/pull/101
* Added GroupNorm by patrick-kidger in https://github.com/patrick-kidger/equinox/pull/104
* Bump version by patrick-kidger in https://github.com/patrick-kidger/equinox/pull/103

New Contributors
* jvmncs made their first contribution in https://github.com/patrick-kidger/equinox/pull/91
* jatentaki made their first contribution in https://github.com/patrick-kidger/equinox/pull/100

**Full Changelog**: https://github.com/patrick-kidger/equinox/compare/v0.5.0...v0.5.1

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**: https://github.com/patrick-kidger/equinox/compare/v0.4.0...v0.5.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 `RuntimeError`s 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.Module`s) 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 https://github.com/patrick-kidger/equinox/pull/58
* Remove default bias in `MultiheadAttention` by patrick-kidger in https://github.com/patrick-kidger/equinox/pull/60
* Improvements to stateful by patrick-kidger in https://github.com/patrick-kidger/equinox/pull/61
* Improved pretty-printing by patrick-kidger in https://github.com/patrick-kidger/equinox/pull/62
* Bump version by patrick-kidger in https://github.com/patrick-kidger/equinox/pull/63


**Full Changelog**: https://github.com/patrick-kidger/equinox/compare/v0.3.2...v0.4.0