Dm-haiku

Haiku is a simple neural network library for [JAX] developed by some of the
authors of [Sonnet], a neural network library for [TensorFlow].

**Changelog:**

Features:

- Exposed `hk.nets.ResNet` and addeed `hk.nets.ResNet{18,34,101,152,200}`
- Added `IdentityCore`.
- Added `custom_getter` API for advanced parameter manipulation.
- Added `ConvND` and lifted `N<=3` restriction.
- Added `tree_size` and `tree_bytes` to easily compute parameter counts.
- `hk.remat` now only threads changed values (faster compilation).
- Added support for `dataclass` to define modules.
- Added support for splitting >1 key at a time `k1, k2 = hk.next_rng_keys(2)`.
- *Experimental:* Added `profiler_name_scopes` API to add Haiku names to XProf.
- *Experimental:* Added `optimize_rng_use` to improve compilation time for models with lots of RNG keys.

Examples:

- Added language model example.
- Added `VQVAE` example.

Bug fixes:

- `LayerNorm` now correctly handles bf16 inputs.
- `TruncatedNormal` initializer now respects dtype.

Usability:

- Improved error messages for `get_parameter`, `to_module` and others.
- Reimplemented core modules with "public" API (easier to read and fork).
- Added tests that ensure all public symbols are included in documentation.
- Added type annotations to more internal code.

[JAX]: https://github.com/google/jax
[Sonnet]: https://github.com/deepmind/sonnet
[Tensorflow]: https://github.com/tensorflow/tensorflow

Changes

Examples

- Added VAE example.
- Added pruning example (https://arxiv.org/abs/1710.01878).
- MNIST example uses 300-100-10 MLP.
- Updated imagenet dataset to return correctly scaled examples.

Breaking changes

- State arg to `hk.transform` dropped in favor of `transform_with_state`.
- Decay argument is now required in `BatchNorm`.

Features

- Added `hk.maybe_next_rng_key()`.
- BatchNorm and LayerNorm speed improvements.
- Added support for partition/filter/merge params.
- Haiku now allows running with `jax_numpy_rank_promotion`.

Experimental features

- `hk.experimental.to_dot` - experimental visualisation support.
- `hk.experimental.lift` - experimental purification support.

Usability

- Improved error message when RNG arg is not and RNG.
- Improved documentation.
- Improved test coverage.

Haiku is a neural network library for JAX.