dm-haiku: Sonnet-style modules for JAX, now in maintenance mode

What dm-haiku solves for JAX users

JAX is a numerical computing library that combines NumPy, automatic differentiation, and first-class GPU/TPU support. It is functional: parameters are values you pass around, and every transformation expects pure functions. Writing a network that way means threading parameter dictionaries through every call by hand. Haiku exists to remove that boilerplate.

The README describes Haiku as "a simple neural network library for JAX developed by some of the authors of Sonnet, a neural network library for TensorFlow." The intended audience is people who liked Sonnet's Module programming model in TensorFlow and want the same experience in JAX. The README says modules, methods, argument names, defaults, and initialization schemes aim to match Sonnet 2, so a TensorFlow and Sonnet user can move over without relearning naming conventions.

The README is also explicit about scope. Haiku is "a library, not a framework": it manages model parameters and other model state, and deliberately does not define custom optimizers, checkpointing formats, or replication APIs. That is a real boundary. If you want an all-in-one training stack, Haiku is not it, and the README says so.

How hk.Module and hk.transform actually work

Two core tools carry the design: `hk.Module` and `hk.transform`. Modules are Python objects that hold references to their own parameters, other modules, and methods that apply functions on user inputs. That is the object-oriented surface. Underneath, `hk.transform` turns functions that use these object-oriented, functionally "impure" modules into pure functions usable with `jax.jit`, `jax.grad`, `jax.pmap` and the rest.

The mechanism is a pair of functions. The README states that `transform` converts your function into pure `init` and `apply` functions. `init`, with signature `params = init(rng, ...)`, collects initial parameter values by running your function and tracking anything requested through `hk.get_parameter`. `apply` then runs the same function against supplied parameters. The README's quickstart shows the split directly: `params = loss_fn_t.init(rng, dummy_images, dummy_labels)` followed by `grads = jax.grad(loss_fn_t.apply)(params, images, labels)`.

Randomness gets the same treatment. Within a transformed function, `hk.next_rng_key()` returns a unique rng key. The README states these keys are deterministically derived from an initial random key passed into the top-level transformed function, which is what makes them safe under JAX transformations. The quickstart also calls `hk.without_apply_rng(loss_fn_t)` before training, which removes the rng argument from `apply` when the function does not need one.

The repository layout reflects the library-not-framework split. There is a `haiku/` package, an `examples/` directory with `mnist.py`, `vae.py`, `impala_lite.py`, `mnist_pruning.py`, transformer and rnn subdirectories, and a `docs/` tree, but no optimizer or checkpointing package. Those pieces are left to the rest of JAX.

Installing dm-haiku and running a first module

The package name on PyPI is `dm-haiku`, and the README's installation section points at pip. The repository's `requirements.txt` lists the runtime dependencies: `absl-py>=0.7.1`, `jmp>=0.0.2`, `numpy>=1.18.0`, and `tabulate>=0.8.9`. JAX itself is not in that file; `requirements-jax.txt` is a separate file in the repository root.

pip install dm-haiku

After installation, the quickstart pattern is to define a function that builds modules, then transform it. This snippet follows the README's example shape: an `hk.Sequential` of `hk.Linear` layers, wrapped in `hk.transform`.

import haiku as hk
import jax
import jax.numpy as jnp

def net_fn(x):
  mlp = hk.Sequential([
      hk.Linear(300), jax.nn.relu,
      hk.Linear(100), jax.nn.relu,
      hk.Linear(10),
  ])
  return mlp(x)

net = hk.transform(net_fn)
net = hk.without_apply_rng(net)
rng = jax.random.PRNGKey(42)
params = net.init(rng, jnp.zeros([1, 784]))
out = net.apply(params, jnp.zeros([1, 784]))

What you should see: `params` is a nested dictionary of arrays keyed by module name, and `out` has shape `[1, 10]`. The README's own quickstart goes one step further and wires the transformed function into a training loop with `jax.grad(loss_fn_t.apply)` and a hand-written update rule, which is the honest picture of what Haiku leaves to you.

For a longer reference, the README points at the examples directory and singles out `examples/mnist.py` as a good place to start.

Maintenance mode: what the README commits to and what it does not

The most important paragraph in the README is the notice at the top. As of July 2023, Google DeepMind recommends that new projects adopt Flax instead of Haiku. The notice states Haiku will remain best-effort supported, that the project has entered maintenance mode, and that development efforts will focus on bug fixes and compatibility with new releases of JAX. New releases will keep Haiku working with newer versions of Python and JAX, but the team will not add or accept PRs for new features. The README also states there is significant internal usage at Google DeepMind and that the plan is to support Haiku in this mode indefinitely.

The release history is consistent with that description rather than contradicting it. v0.0.15 landed on 2025-09-18, v0.0.16 on 2025-12-17, and v0.0.17 on 2026-07-27. The last push to the default branch was on 2026-09-09. Version numbers in the 0.0.x range after years of use tell you the project never treated API stability as a promise, and the README's own Sonnet-parity language means the API surface is inherited rather than designed fresh.

One practical consequence: the README does not document a deprecation or removal schedule for any API. If you depend on a specific module, there is no published timeline telling you when it might change. Compatibility work is the stated priority, not feature evolution.

Where Haiku is the wrong choice

The clearest failure case is a new project. The README's own recommendation is to adopt Flax instead, and it gives reasons: at the time of writing, Flax has a superset of the features available in Haiku, a larger and more active development team, and more adoption with users outside of Alphabet. It also points to more extensive documentation, examples in the Hugging Face transformers repository, and an active community producing end-to-end examples. Starting fresh on Haiku means starting on the library its own maintainers steer people away from.

The second case is anyone who wants training infrastructure included. Haiku does not define custom optimizers, checkpointing formats, or replication APIs. You assemble those yourself from JAX and other libraries. If your team's bottleneck is getting a standard training loop running, that assembly work is real cost with no Haiku answer for it.

A third case is users who need something the API simply does not cover. Because new features are not being added, an unmet need is not a roadmap item. The README frames the future of the project as bug fixes and JAX compatibility, so a gap you hit today is likely to stay a gap.

Finally, note the naming collision the README itself calls out: if you are looking for Haiku the operating system, that is a different project at haiku-os.org.

Haiku vs Flax: the actual difference in approach

Both libraries sit on JAX and both manage parameters, so the comparison is about programming model rather than capability. Haiku keeps Sonnet's object-oriented style: you instantiate `hk.Module` subclasses and call them like functions, and `hk.transform` converts that impure code into pure `init` and `apply` functions. The README describes this as preserving Sonnet's Module-based programming model for state management while retaining access to JAX's function transformations, and says the APIs and abstractions are as close as reasonable to Sonnet.

Flax, which the README names as the recommended alternative, was originally developed by Google Brain and is now developed by Google DeepMind. The README states that at the time of writing Flax has a superset of Haiku's features, a larger and more active development team, and more adoption outside Alphabet, plus more extensive documentation and community examples.

The practical difference for a reader deciding today is not which transform is nicer. It is that one library is adding features and the other is not. Haiku's advantage is continuity: if your codebase or your colleagues' habits come from Sonnet, the argument names, defaults, and initialization schemes should match, and the migration is mostly mechanical. Flax's advantage is that it is where new work is going. If you have no Sonnet history, the continuity argument does not apply to you.

Editorial conclusion

Adopt Haiku if you maintain existing Sonnet-style JAX code or need hk.transform and hk.next_rng_key for research reproduction, and you accept that new releases only track Python and JAX compatibility. Do not start a greenfield project here: the README states DeepMind recommends Flax instead. Before committing, verify that your JAX version is covered by the release you pin, check the requirements.txt pins (absl-py, jmp, numpy, tabulate), and run the MNIST example in examples/mnist.py against your installed versions.

Frequently asked questions

Official sources

1. google-deepmind/dm-haiku on GitHub
2. License: Apache-2.0
3. Project website
4. README
5. Releases