dm-haiku

What it is and what it does

Haiku is a neural network library built on JAX that lets you write models using familiar object-oriented patterns—defining modules with parameters and methods—while automatically transforming them into pure functions compatible with JAX's transformations. It provides two core abstractions: `hk.Module` for encapsulating network state and computation, and `hk.transform` for converting module-based code into pure `init` and `apply` functions.

The library is designed to be minimal and composable, handling parameter initialization and state management without imposing custom optimizers, checkpointing formats, or replication APIs. It draws its API and programming model from Sonnet, making it familiar to users migrating from TensorFlow. As of July 2023, Google DeepMind recommends new projects adopt Flax instead; Haiku now operates in maintenance mode, receiving bug fixes and compatibility updates but no new features.

Use it for:

• Building and training image classification models with parameter management and automatic differentiation.
• Implementing reinforcement learning agents where you need state management with functional transformations.
• Migrating existing Sonnet/TensorFlow neural network code to JAX with minimal API changes.
• Prototyping generative models that benefit from functional programming patterns.
• Research projects where DeepMind has validated Haiku's reliability in large-scale experiments.

Worth the install?

AI-flagged interpretation of the facts on this page — verify before relying

Haiku is a neural network library for JAX that provides an object-oriented module abstraction and function transformation to build and train models while retaining access to JAX's pure function transformations.

Yes, if you are already committed to JAX and value a lightweight, Sonnet-like API for parameter management. No, if you are starting a new project—Google DeepMind officially recommends Flax instead, which has more active development and broader adoption. Haiku remains best-effort supported indefinitely but will not gain new features; install it only when migrating existing code or when its specific design philosophy is a deliberate fit.

Install

dm-haiku on PyPI

pip install dm-haiku

uv add dm-haiku

poetry add dm-haiku

Installing dm-haiku

Before you install

Low install friction with a pure Python wheel. Actively maintained as of 2026-08-06 with recent releases, though the project entered maintenance mode in July 2023 and now focuses on bug fixes and compatibility rather than new features.

License in practice

Licensed under Apache 2.0 (permissive), allowing free use, modification, and distribution with minimal restrictions—suitable for both open-source and commercial projects.

Quickstart

pip install dm-haiku

import haiku as hk
import numpy

def loss_fn(images, labels):
  mlp = hk.Sequential([hk.Linear(300), hk.Linear(10)])
  logits = mlp(images)
  return logits

loss_fn_t = hk.transform(loss_fn)
rng = numpy.random.PRNGKey(42)
params = loss_fn_t.init(rng, dummy_images, dummy_labels)

Requires Python 3.10 or later; depends on absl-py, jmp, numpy, and tabulate.

Verify before relying

• Whether Haiku's maintenance-mode status affects long-term viability for new projects given the official recommendation to use Flax instead.
• Performance characteristics and scalability limits compared to alternatives for production workloads.
• Community adoption trends and availability of third-party examples or integrations beyond DeepMind's internal use.

Package facts

Evidence: dm_haiku-0.0.17-py3-none-any.whl

Tags

More Libraries packages