Install & Compatibility
Where this runs
tested against v0.0.16 · pip install
no network on importno background threads
Install × environment matrix
Each cell = how many times install + import succeeded across repeated harness runs. Partial = flaky.
glibc = Debian/Ubuntu slim · musl = Alpine Linux
muslpy 3.10–3.920 runs
installs and imports cleanly · install 0.0s · import 0.000s · 93.7MB
glibcpy 3.10–3.920 runs
installs and imports cleanly · install 4.1s · import 0.000s · 90MB
94MB installed
● package 94MB
Code
Verified usage
Verified import paths — ran on the pinned version, not inferred.
haiku
✓ import haiku as hk
jax.numpy
✓ import jax.numpy as jnp
This quickstart demonstrates how to define a simple Multi-Layer Perceptron (MLP) using Haiku modules and then transform it into a pair of pure functions (init and apply) compatible with JAX transformations. It shows how to initialize model parameters using a JAX PRNG key and then apply the model to an input.
import haiku as hk
import jax
import jax.numpy as jnp
def forward_fn(x):
mlp = hk.nets.MLP([300, 100, 10])
return mlp(x)
transformed_forward = hk.transform(forward_fn)
rng = hk.PRNGSequence(jax.random.PRNGKey(42))
x = jnp.ones([8, 28 * 28]) # Example input
# Initialize parameters
params = transformed_forward.init(next(rng), x)
# Apply the model
logits = transformed_forward.apply(params, next(rng), x)
print("Parameters structure:", jax.tree_util.tree_map(lambda x: x.shape, params))
print("Output shape:", logits.shape)
Debug
Known issues
deprecatedAs of July 2023, Google DeepMind recommends that new projects adopt Flax instead of Haiku. Haiku is in maintenance mode, focusing on bug fixes and JAX compatibility rather than new features.fixConsider using Flax (e.g., `flax.linen`) for new projects. Haiku will continue to be supported for existing internal DeepMind usage.
affects: 0.0.15 and later
gotchaHaiku frequently updates to maintain compatibility with new JAX releases. It is highly recommended to pin your JAX and JAXlib versions to specific compatible versions to avoid unexpected breakage, especially in production environments.fixExplicitly specify `jax` and `jaxlib` versions in your `requirements.txt` or `pyproject.toml` file. Refer to Haiku's GitHub releases for notes on JAX compatibility for each version.
affects: All versions
breaking`hk.vmap(..)` now requires the `split_rng` argument to be explicitly passed.fixWhen using `hk.vmap`, ensure you provide the `split_rng` argument (e.g., `hk.vmap(func, split_rng=True)`).
affects: 0.0.7 and later
breaking`hk.jit` was removed from the public API.fixDirectly use `jax.jit` on the `apply` function of your `hk.transform`-ed model instead of `hk.jit`. For example: `apply_jit = jax.jit(transformed_forward.apply)`.
affects: 0.0.7 and later
gotchaUsing JAX transformations (like `jax.jit`, `jax.vmap`, `jax.remat`, `jax.lax.scan`) directly inside a Haiku module or within a function intended to be transformed by `hk.transform` can lead to `jax.errors.UnexpectedTracerError` or silently wrong results, because Haiku modules are side-effecting before transformation.fixAlways apply JAX transformations *after* transforming your Haiku function with `hk.transform` (or `hk.transform_with_state`). If you need to apply JAX transformations to parts of your model *within* a Haiku module, consider using `hk.lift` or carefully structuring your code to ensure pure functions are passed to JAX transformations.
affects: All versions
Upgrade
Version history
0.0.16latest on PyPI · released Dec 17, 2025
Audit
Dependencies
jaxrequiredHaiku is built on JAX; JAX must be installed separately with appropriate accelerator support (e.g., CUDA) before installing Haiku.
jaxlibrequiredRequired for JAX functionality.
absl-pyrequiredRequired for configuration and logging utilities.
jmprequiredUsed for mixed precision training.
numpyrequiredFundamental numerical computing library.
tabulaterequiredUsed for displaying data in tabular format.
flaxoptionalRecommended by DeepMind for new projects as an alternative to Haiku; Haiku can be used without Flax in Python >=3.13.