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.

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.

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.