Dm-haiku

- Various changes to remain compatible with latest JAX release (e.g. 779, 783, 789).
- Add `shard_map` flops estimate to jaxpr_info (791).
- Correct an oversight in `TruncatedNormal`'s documentation (793).

- `hk.LayerStack` transparency map now supports custom pytree leaves 9be7c9a9c6160c089b3d32f9384a41357f47d23c.
- `hk.LayerStack` now supports layers using state 134763d60a66b9246b1aaac67e5e63b4a4f4e0f5.
- Calling `repr(self)` in module constructors is now supported (fixing 428).

- `hk.layer_stack` now allows transparent application (no prefix on module names).
- `hk.MultiHeadAttention` allows bias initializer to be configured or biases to be removed.
- `hk.DepthwiseConvND` now supports `dilation`.
- `hk.dropout` supports `broadcast_dims`.
- `hk.BatchApply` avoids an unnecessary h2d copy during tracing.
- `hk.experimental.profiler_name_scopes` has been removed, these are on by default.
- Added `hk.map` mirroring `jax.lax.map`.

- Added `hk.mixed_precision.push_policy`.
- Added `hk.experimental.{get_params,get_initial_state,get_current_state}`.
- Added `hk.experimental.{maybe_get_rng_sequence_state,maybe_replace_rng_sequence_state}`.
- `hk.switch` now supports multiple operands.
- `hk.get_parameter` now supports `init=None`.
- `hk.MethodContext` now includes `orig_class`.
- `hk.GetterContext` now includes `lifted_prefix_name`.
- `hk.layer_stack` now allows parameter reuse.
- Haiku is now compatible with `jax.enable_custom_prng`.
- `TruncatedNormal` now exports lower and upper bounds.
- Haiku init/apply functions now return `dict` rather than `Mapping`.
- `hk.dropout` now supports `broadcast_dims`.

What's Changed
* Support vmap where in_axes is a list rather than a tuple in https://github.com/deepmind/dm-haiku/commit/307cf7dbda64d637ca423cacc9978f0ca19dc8a6
* Pass pmap axis specs optionally to make_model_info in https://github.com/deepmind/dm-haiku/commit/d0ba451c96a6ac4f44fb9457e252b1d675a5416a
* Remove use of jax_experimental_name_stack flag in https://github.com/deepmind/dm-haiku/commit/dbc0b1f2ffee9b348a3cb67460f28f9cc4667f08
* Add param_axis argument to RMSNorm to allow setting scale param shape in https://github.com/deepmind/dm-haiku/commit/a4998a02bc4e8303f9897e5c32ded90cc38fa84f
* Add documentation and error messages for w_init and w_init_scale to avoid confusion in https://github.com/deepmind/dm-haiku/pull/541
* Fix hk.while_loop carrying state when reserving variable sizes of rng keys. by copybara-service in https://github.com/deepmind/dm-haiku/pull/551
* Add ensemble example to hk.lift documentation. by copybara-service in https://github.com/deepmind/dm-haiku/pull/556

**Full Changelog**: https://github.com/deepmind/dm-haiku/compare/v0.0.8...v0.0.9

* Added [`experimental.force_name`](https://dm-haiku.readthedocs.io/en/latest/api.html#haiku.experimental.force_name).
* Added ability to simulate a method name in [`experimental.name_scope`](https://dm-haiku.readthedocs.io/en/latest/api.html#haiku.experimental.name_scope).
* Added a [config option](https://dm-haiku.readthedocs.io/en/latest/api.html?highlight=rng_reserve_size#module-haiku.config) for PRNG key block size.
* Added [`unroll`](https://dm-haiku.readthedocs.io/en/latest/api.html?highlight=unroll#haiku.dynamic_unroll) parameter to `dynamic_unroll`.
* Remove use of deprecated `jax.tree_*` functions.
* Many improvements to our examples.
* Improve error messages in [`vmap`](https://dm-haiku.readthedocs.io/en/latest/api.html#haiku.vmap).
* Support `jax_experimental_name_stack` in jaxpr_info.
* [`transform_and_run`](https://dm-haiku.readthedocs.io/en/latest/api.html?highlight=map_rng#haiku.testing.transform_and_run) now supports a map on PRNG keys.
* `remat` now uses the new JAX remat implementation.
* Scale parameter is now optional in [`RMSNorm`](https://dm-haiku.readthedocs.io/en/latest/api.html?highlight=create_scale#haiku.RMSNorm).