Dm-haiku

- **Bug fix:** modules with leading zeros (e.g. `linear_007`) are now correctly handled. 7632affafad
- **Breaking change:** [`hk.vmap(..)`](https://dm-haiku.readthedocs.io/en/latest/api.html#haiku.vmap) now requires `split_rng` to be passed.
- **Breaking change:** `hk.jit` was removed from the public API.
- Feature: we always add profiler name scopes to Haiku modules with the latest version of JAX.
- Added a [tutorial on parameter sharing](https://github.com/deepmind/dm-haiku/blob/main/docs/notebooks/parameter_sharing.ipynb).
- Added [`hk.ModuleProtocol`](https://dm-haiku.readthedocs.io/en/latest/api.html#haiku.ModuleProtocol) and [`hk.SupportsCall`](https://dm-haiku.readthedocs.io/en/latest/api.html#haiku.SupportsCall).
- Added `cross_replica_axis` to [`VectorQuantiser`](https://dm-haiku.readthedocs.io/en/latest/api.html#haiku.nets.VectorQuantizer).
- Added `allow_reuse` argument to [`hk.lift`](https://dm-haiku.readthedocs.io/en/latest/api.html#lift).
- Added `fan_in_axes` to [`VarianceScaling`](https://dm-haiku.readthedocs.io/en/latest/api.html#variancescaling) initialiser.
- Added [`hk.custom_setter(..)`](https://dm-haiku.readthedocs.io/en/latest/api.html#custom-setter) to intercept `hk.set_state(..)`.
- Added [`hk.Deferred`](https://dm-haiku.readthedocs.io/en/latest/api.html#deferred).
- Added [`hk.experimental.transparent_lift(..)`](https://dm-haiku.readthedocs.io/en/latest/api.html#transparent-lift) and [`hk.experimental.transparent_lift_with_state(..)`](https://dm-haiku.readthedocs.io/en/latest/api.html#transparent-lift-with-state).
- Added [`hk.experimental.fast_eval_shape(..)`](https://dm-haiku.readthedocs.io/en/latest/api.html#fast-eval-shape).
- Added [`hk.experimental.current_name()`](https://dm-haiku.readthedocs.io/en/latest/api.html#current-name).
- Added `hk.experimental.DO_NOT_STORE`. 2a6c03406f
- Added [config APIs](https://dm-haiku.readthedocs.io/en/latest/api.html#module-haiku.config).
- Added support for new `jax.named_call` implementation.
- The `HAIKU_FLATMAPPING` env var is no longer used.
- [`hk.dropout(..)`](https://dm-haiku.readthedocs.io/en/latest/api.html#haiku.dropout) now supports dynamic `rate`.
- [`hk.without_apply_rng(..)`](https://dm-haiku.readthedocs.io/en/latest/api.html#without-apply-rng) now supports multi transformed functions.

- Haiku now returns plain nested `dict`s rather than `FlatMap` from all APIs.
- [`hk.vmap(..)`](https://dm-haiku.readthedocs.io/en/latest/api.html#vmap) now optionally takes `split_rng`, this argument will be required in the next version of Haiku.
- [`hk.LayerNorm`](https://dm-haiku.readthedocs.io/en/latest/api.html#layernorm) now accepts `param_axis` in the constructor, this argument will be required in the next version of Haiku.
- [`hk.get_channel_index(..)`](https://dm-haiku.readthedocs.io/en/latest/api.html#get-channel-index) was added.
- [`hk.experimental.lift_with_state(..)`](https://dm-haiku.readthedocs.io/en/latest/api.html#lift-with-state) was added.
- [`hk.experimental.layer_stack(..)`](https://dm-haiku.readthedocs.io/en/latest/api.html#layer-stack) was added.
- [`hk.DepthwiseConv{1,3}D`](https://dm-haiku.readthedocs.io/en/latest/api.html?highlight=DepthwiseConv#convolutional) were added.
- [`hk.BatchNorm`](https://dm-haiku.readthedocs.io/en/latest/api.html#batchnorm) now supports sequences in `cross_replica_axis`.
- [`hk.experimental.check_jax_usage()`](https://dm-haiku.readthedocs.io/en/latest/api.html#check-jax-usage) makes Haiku check that JAX control flow/transforms are used correctly.
- [`hk.dynamic_unroll(..)`](https://dm-haiku.readthedocs.io/en/latest/api.html#dynamic-unroll) now supports `return_all_states`.
- [`hk.cond(..)`](https://dm-haiku.readthedocs.io/en/latest/api.html#cond) supports N operands.
- [`hk.experimental.module_auto_repr(False)`](https://dm-haiku.readthedocs.io/en/latest/api.html#module-auto-repr) can be used to speed up init.
- [`hk.data_structures.merge(..)`](https://dm-haiku.readthedocs.io/en/latest/api.html#merge) now supports `check_duplicates`.
- [`TruncatedNormal`](https://dm-haiku.readthedocs.io/en/latest/api.html#truncatednormal) initialiser now supports complex dtypes.
- `transform(jit(f))` now provides a nice error message.
- [`hk.multinomial(..)`](https://dm-haiku.readthedocs.io/en/latest/api.html#multinomial) now uses `jax.random.categorical`.
- Added [`hk.mixed_precision.{current,get}_policy(..)`](https://dm-haiku.readthedocs.io/en/latest/api.html#automatic-mixed-precision) for introspection.
- Mixed precision policies now support reloaded modules.

- Added support for mixed precision training (dba1fd9) via [jmp](https://github.com/deepmind/jmp)
- Added `hk.with_empty_state(..)`.
- Added `hk.multi_transform(..)` (137), supporting transforming multiple functions that share parameters.
- Added `hk.data_structures.is_subset(..)` to test whether parameters are a subset of another.
- Minimum Python version is now 3.7.
- Multiple changes in preparation for a future version of Haiku changing to plain `dict`s.
- `hk.next_rng_keys(..)` now returns a stacked array rather than a collection.
- `hk.MultiHeadAttention` now supports distinct sequence lengths in query and key/value.
- `hk.LayerNorm` now optionally supports faster (but less stable) variance computation.
- `hk.nets.MLP` now has an output_shape property.
- `hk.nets.ResNet` now supports changing strides.
- `UnexpectedTracerError` inside a Haiku transform now has a more useful error message.
- `hk.{lift,custom_creator,custom_getter}` are no longer experimental.
- Haiku now supports JAX's pluggable RNGs.
- We have made multiple improvements to our docs an error messages.

Any many other small fixes and improvements.

**Changelog:**

- **(Important Fix) Fixed strides in basic block (300e6a40be3).**
- Added map, partition_n and traverse to data_structures.
- Added "build your own Haiku" to the docs.
- Added summarise utility to Haiku.
- Added visualisation section to docs.
- Added precision arg to Linear, Conv and ConvTranspose.
- Added RMSNorm.
- Added module_name and name to GetterContext.
- Added hk.eval_shape.
- Improved performance of non cross-replica BN variance.
- Haiku branch functions are only traced once (mirroring JAX).
- Attention logits are rescaled before the softmax now.
- ModuleMetaclass now inherits from Protocol.
- Removed "dot access" to FlatMapping.
- Removed query_size from MultiHeadAttention constructor.

Any many other small fixes and improvements.

**Changelog:**

- Added `hk.experimental.intercept_methods `.
- Added `hk.running_init`.
- Added `hk.experimental.name_scope`.
- Added optional support for state in `custom_creator` and `custom_getter`.
- Added index groups to `BatchNorm`.
- Added interactive notebooks to documentation, including [basics guide](https://dm-haiku.readthedocs.io/en/latest/notebooks/basics.html).
- Added support for batch major unrolls in `static_unroll` and `dynamic_unroll`.
- Added `hk.experimental.abstract_to_dot`.
- Added step markers in imagenet example.
- Added `hk.MultiHeadAttention`.
- Added option to remove double bias from `VanillaRNN`.
- Added support for `feature_group_count` in `ConvND`.
- Added logits config to resnet models.
- Added various control flow primitives (`fori_loop`, `switch`, `while_loop`).
- Added `cross_replica_axis` to `VectorQuantizerEMA`.
- Added `original_shape` to `ParamContext`.
- Added `hk.SeparableDepthwiseConv2D `.
- Added support for `unroll` kwarg to `hk.scan`.
- Added `output_shape` argument to `ConvTranspose` modules.
- Replaced `frozendict` with `FlatMapping`, significantly reduces overheads calling jitted computations.
- Misc changes to ensure parameter dtype follows input dtype.
- Multiple changes to support JAX omnistaging.
- `ExponentialMovingAverage.initialize` now takes shape/dtype not value.
- Replaced optix with optax in examples.
- `hk.Embed` embeddings now created lazily.
- Re-indexed documentation for easier navigation.

**Changelog:**

- Changed the default value of `apply_rng` to `True` in `hk.transform` to simplify the `apply_fn` signature.
- Made `ConvND`, `ConvNDTranspose`, `ResetCore` and pooling modules optionally batched.
- Added `hk.GroupNorm`.
- Added `hk.scan`.
- Changed `hk.BatchNorm` to always create state for moving averages.
- Changed `use_projection` in `hk.nets.ResNet` to take a sequence of bools.
- Exposed `hk.net.ResNet.{BlockGroup, BlockV1, BlockV2}`.
- Added `original_dtype` to `ParamContext` to expose the original parameter dtype to custom_getters.
- Added `GAN` example notebook.