dm-haiku

Why we included this project

Building neural networks directly on JAX often gets awkward the moment you want to define and reuse layers, because JAX's pure-function model doesn't naturally keep parameters attached to the objects that use them. Haiku solves that with an object-oriented module API that still produces the pure init and apply functions JAX's jit, grad, and pmap transforms expect. It comes from the same authors as Sonnet, so the API is immediately recognizable to anyone with TensorFlow experience, and it's been used to reproduce a wide range of DeepMind experiments in vision, language, generative models, and reinforcement learning. Given that history, it's a dependable, lightweight layer for researchers who want composability without a full framework's optimizer and checkpointing machinery. One thing to keep in mind: Google DeepMind now recommends Flax for new projects, and Haiku is in maintenance mode, so it's mainly useful for existing Haiku codebases or as a clean reference implementation.