dm-haiku: object-oriented modules on top of JAX, now in maintenance mode
JAX-based neural network library
At a glance
- What is it?
- Haiku wraps JAX's pure function transformations in Sonnet-style hk.Module objects, and its own README tells new projects to use Flax instead. Here is what the library still does, how init/apply work, and who should keep it.
- Who is it for?
- Adopt Haiku if you already have Sonnet 2 or TensorFlow module code to port, or you need its explicit parameter handling and hk.next_rng_key inside an existing Haiku codebase. Do not adopt it for a new project: the README states that as of July 2023 Google DeepMind recommends Flax for new work.
- Can I use it commercially?
- Yes. Apache-2.0 is a permissive licence: you can use, modify and sell software built on it, as long as you keep its copyright and licence notices.
- Is it still maintained?
- Yes. The repository last received commits 1 day ago.
- What is it written in?
- Mainly Python, according to GitHub's language statistics.
Answers come from the project's GitHub data, last synced on September 15, 2026, and from our analysis. They are not legal advice.
DEEP OPEN-SOURCE ANALYSIS
The gap Haiku fills between raw JAX and a full framework
JAX gives you NumPy-style arrays, automatic differentiation and GPU/TPU support as pure functions. It does not give you a place to put weights. You can carry a nested dictionary of arrays through your training loop by hand, but every layer then needs its own shape logic, its own initialiser and its own key handling, and that boilerplate grows with the model. Haiku's answer is to let you write layers as Python objects with attributes and methods, then convert the whole function into something JAX can differentiate and compile. The README frames the library as a tool for building neural networks, and describes it as Sonnet for JAX, written by some of the authors of Sonnet for TensorFlow. The intended audience is therefore specific: people who already think in terms of modules that own their parameters, and who want that mental model to survive the move to JAX. If you are happy writing plain functions over pytrees, Haiku adds an abstraction you may not want. The README is explicit that Haiku is a library rather than a framework: it does not define custom optimizers, checkpointing formats or replication APIs. That narrow scope is the design, not an omission.
hk.Module plus hk.transform: how parameters actually get collected
Two pieces do the work. hk.Module is a Python object that holds references to its own parameters, to other modules, and to the methods that apply functions to inputs. hk.transform takes a function built from those modules and returns a pair of pure functions, init and apply, which is what JAX requires. The mechanism the README describes is a collection pass: init has the signature params = init(rng, ...) where the remaining arguments are the arguments to your untransformed function, and it collects initial values by running the function and tracking every parameter requested through hk.get_parameter. apply then runs the same function against supplied parameters. The quickstart shows the shape of it: a loss_fn builds hk.Sequential of hk.Linear layers, loss_fn_t = hk.transform(loss_fn) produces the transformed pair, and params = loss_fn_t.init(rng, dummy_images, dummy_labels) needs a real batch of dummy inputs because shapes are only known at trace time. Randomness gets its own treatment. Inside a transformed function, hk.next_rng_key() returns a unique key derived deterministically from the initial key passed to the top-level transformed function, which is what makes the result safe under JAX transformations. The README also shows hk.without_apply_rng, used in the quickstart to strip the rng argument from apply when the loss does not need it at call time.
Installation and the quickstart in practice
The README's installation section points at the documentation rather than reproducing a command, and the PyPI badge indicates the package is published as dm-haiku, so the install path is a pip install of that distribution name against a JAX version of your choosing. The quickstart is the part worth reading closely, because it exposes the workflow cost. You write an ordinary Python function using hk.Sequential and hk.Linear, transform it, and then split your program into two phases that never merge: an init call that needs a representative batch and an rng, and an apply call used inside jax.grad. In the example, grads = jax.grad(loss_fn_t.apply)(params, images, labels) and the update rule is a plain jax.tree.map over parameters, because Haiku deliberately ships no optimizer. The dummy inputs are not a formality. Any conditional that depends on input shape, or any layer whose parameter shape depends on the batch, will be resolved during init and fixed thereafter, which is standard JAX tracing behaviour that Haiku inherits rather than hides. The README points to the examples directory and singles out the MNIST example as a good starting point; that is where the pattern of init, apply and a hand-written update loop is laid out end to end.
The maintenance-mode notice is the first thing to read
The README carries an important banner stating that as of July 2023 Google DeepMind recommends new projects adopt Flax instead of Haiku. It gives reasons: Flax is described as having a superset of Haiku's features, a larger and more active development team, more adoption outside Alphabet, more extensive documentation and examples, and an active community producing end-to-end examples. Haiku will remain best-effort supported and the project is in maintenance mode, meaning development effort is focused on bug fixes and compatibility with new JAX releases. New releases will keep Haiku working with newer versions of Python and JAX, and the README states that new features will not be added and pull requests for them will not be accepted. It also states that Google DeepMind has significant internal usage of Haiku and currently plans to support it in this mode indefinitely. Those two facts sit together: the library is not abandoned, and it is not growing. For a dependency you already run, indefinite bug-fix support is a reasonable position. For a greenfield project, the same paragraph is a direct instruction to look elsewhere.
Where Haiku is the wrong tool
Three cases stand out. First, new projects where the team has no Sonnet background: the README itself routes you to Flax, and choosing Haiku means accepting a feature set its own maintainers describe as a subset, with no path to new features. Second, teams that want an ecosystem rather than a parameter-management layer. Haiku does not define optimizers, checkpoint formats or replication APIs, so you assemble those yourself, and the README presents that as a deliberate boundary. If you want a library that also covers training utilities and a large body of example code, Haiku is the wrong layer. Third, anyone whose model depends on shape-dependent control flow that cannot be captured by a single representative init batch. Because init traces your function to discover parameters, conditional architectures need care at that boundary, and the README does not present a workaround. There is also a quieter cost: the README notes that outside of new features such as hk.transform, Haiku aims to match the API of Sonnet 2, including modules, methods, argument names, defaults and initialisation schemes. That fidelity is a benefit for porting and a constraint for anyone expecting Haiku to diverge and modernise.
Flax: the same problem, a different stance on modules
The alternative named in the README is Flax, originally developed by Google Brain and now by Google DeepMind. The difference in approach is not cosmetic. Haiku keeps state in Python objects during tracing and then extracts it: hk.Module instances hold parameters, and hk.transform converts the impure function into pure init and apply. Flax's model treats modules as dataclasses and makes state explicit in the function signature from the start, which is why the README can describe Flax as a superset of Haiku's features without contradiction. In practice the choice shows up in how you read a model file. Haiku code looks like a Sonnet model with hk.Linear and hk.Sequential calls and a separate transform step; Flax code threads parameters and collections through the call. If your team's existing code is Sonnet 2, Haiku's stated aim of matching module names, argument names and defaults makes the port mechanical. If your team is starting fresh, Flax's larger contributor base and documentation set, as described in the README, are the practical reason to start there. Neither library gives you an optimizer, so that part of the decision is neutral.
Versioning, licence and what to verify before you depend on it
Haiku is licensed under Apache-2.0, which permits commercial and closed-source use and requires preserving notices and stating changes; that is a summary of the identifier, not legal advice, and your counsel should review it if the dependency is load-bearing. Releases are infrequent and versioned in the 0.0.x series, with v0.0.17 dated 2026-07-27, v0.0.16 dated 2025-12-17 and v0.0.15 dated 2025-09-18. The gap pattern is consistent with the maintenance-mode statement: releases exist to keep the library working with newer Python and JAX, not to add capability. Two things follow. Pin your JAX version and test upgrades rather than tracking latest, because the compatibility work is the maintenance activity. And treat the 0.0.x numbering as a signal that no API stability guarantee is being advertised. Before adopting, check three things against the current repository: that the release you install declares support for your Python and JAX versions, that the maintenance-mode banner still reads as it does now, and whether the modules you need exist in Haiku or only in Flax. The README's quickstart, the examples directory and the MNIST example are enough to judge the programming model in an afternoon.
Editorial conclusion
Adopt Haiku if you already have Sonnet 2 or TensorFlow module code to port, or you need its explicit parameter handling and hk.next_rng_key inside an existing Haiku codebase. Do not adopt it for a new project: the README states that as of July 2023 Google DeepMind recommends Flax for new work. Before committing, verify that the pinned JAX and Python versions in your environment are covered by the current release, and read the maintenance-mode notice in the repository README in full.
Community notes