Haiku

Warnings

• deprecated As 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.
  Fix: Consider using Flax (e.g., `flax.linen`) for new projects. Haiku will continue to be supported for existing internal DeepMind usage.
• gotcha Haiku 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.
  Fix: Explicitly 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.
• breaking `hk.vmap(..)` now requires the `split_rng` argument to be explicitly passed.
  Fix: When using `hk.vmap`, ensure you provide the `split_rng` argument (e.g., `hk.vmap(func, split_rng=True)`).
• breaking `hk.jit` was removed from the public API.
  Fix: Directly 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)`.
• gotcha Using 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.
  Fix: Always 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.

Install

• pip install -U dm-haiku Install from PyPI

Imports

• haiku

  import haiku as hk

• jax

  import jax

• jax.numpy

  import jax.numpy as jnp

Quickstart

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)