Skip to content

Haiku 0.0.1
Compare
Choose a tag to compare
@tomhennigan tomhennigan released this 04 Jun 11:18
v0.0.1
064ddf5

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.
Assets 2
Loading