Haiku

Haiku is a simple neural network library for JAX that enables users to use familiar object-oriented programming models while allowing full access to JAX's pure function transformations. It provides a module abstraction (`hk.Module`) and a function transformation (`hk.transform`) to manage model parameters and state. As of July 2023, Google DeepMind recommends Flax for new projects, with Haiku having entered maintenance mode, focusing on bug fixes and compatibility with new JAX releases. The current version is 0.0.16.