dm-haiku
Haiku is a library for building neural networks in JAX.
Decision gist · record as of 2026-08-14
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.AI-flagged interpretation of the facts on this page — verify before relying
Before you install
- Requires Python 3.10 or later; depends on absl-py, jmp, numpy, and tabulate.
- 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 · maintenance · safety
Apache 2.0 (permissive) — Licensed under Apache 2.0 (permissive), allowing free use, modification, and distribution with minimal restrictions—suitable for both open-source and commercial projects.
last release 2026-07-27 (18 days) · last repo commit 2026-08-06 · 3,270 stars
0 known vulnerabilities (OSV.dev, 2026-08-14) · 262,617 downloads/mo, #8,370 on PyPI
Alternatives
Verify before relying
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)- 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.
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 on it.
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
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.
Requires Python 3.10 or later; depends on absl-py, jmp, numpy, and tabulate.
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)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
| License | Apache 2.0 permissive |
| Python support | Supports the current Python release >=3.10 |
| Install friction | Low. Pure-Python wheel |
| Runtime dependencies | 4 packagesabsl-pyjmpnumpytabulate |
| Maintenance | Actively maintained 18 days since the last release |
| Last repo commit | |
| First released | |
| Downloads | 262,617 / month, #8,370 on PyPI 30-day window, as of 2026-08-14 |
| Known vulnerabilities | None known OSV.dev, checked 2026-08-14 |
| Classifiers | Development Status :: 4 - BetaIntended Audience :: DevelopersIntended Audience :: EducationIntended Audience :: Science/ResearchLicense :: OSI Approved :: Apache Software LicenseProgramming Language :: Python :: 3Programming Language :: Python :: 3.10Programming Language :: Python :: 3.11Topic :: Scientific/Engineering :: MathematicsTopic :: Software Development :: LibrariesTopic :: Software Development :: Libraries :: Python Modules |
Evidence: dm_haiku-0.0.17-py3-none-any.whl
Tags
Let your AI agent find packages like this
Example. Real query, live index.
You found this page by searching. An agent finds it by wishing: SkillFed indexes 14,416 PyPI packages by what they can do, searchable in plain language.
wish › “object-oriented neural networks jax”
- dm-haikuHaiku is a neural network library for JAX that provides an…
- equinoxEquinox provides neural network and model building on top of JAX with…
- e3nn-jaxImplements equivariant convolutional neural networks for 3D…
Give your agent the search over MCP, or paste the wish link into any chat.
More Libraries packages
urllib3 is an HTTP client library that provides thread-safe connection pooling, SSL/TLS verification, multipart file uploads, request retries, compression support, and proxy handling for Python applications.
Requests is a Python HTTP library that simplifies sending HTTP/1.1 requests with automatic handling of headers, authentication, cookies, and response parsing.
Pluggy provides a plugin system that lets you define hook specifications and register implementations to be called in sequence, enabling extensible Python applications without tight coupling.
Install it if you're building an extensible application or framework.
Provides parsing, arithmetic, and recurrence rule computation for dates and times, with timezone support and iCalendar RFC compliance.
Install it if you need to parse flexible date strings, compute relative dates, handle timezones, or work with recurrence rules—it's the de facto choice for these tasks.
Six provides utility functions to write Python code that runs on both Python 2.7 and Python 3.3+, smoothing over language differences between the two versions.
pytest is a testing framework that lets you write test functions using plain assert statements and automatically discovers and runs them, with detailed failure reporting.
See also flax · jraph · jmp · equinox · clu · rax · qwix · jax · keras-nightly · google-tunix