OpenViSTAI · Agents · Blockchain · Crypto
Submit a project
Menu
AI Developer Tools & SDKs

Haiku

Haiku is an open-source Python neural-network library that adds a module-based model interface to JAX. Its `hk.Module` abstraction organizes model components and parameters, while `hk.transform` converts module-using functions into pure initialization and application functions that can work with JAX transformations.

Explore the practical guide ↓

Independent OpenViST project profile, based on linked sources.

Haiku — OpenViST project profile in AI Developer Tools & SDKs.
OpenViST project illustration. See the linked sources for product details.
Understand the project

Practical overview

When it fits

Use Haiku when a JAX project needs its module-based model and parameter-management abstractions, particularly when working with existing Haiku code. Account for its maintenance-mode status when choosing a dependency for a new project.

Who it is for

  • Python developers building neural-network models with JAX.
  • Researchers and engineering teams maintaining Haiku-based experiments or projects.
  • Developers familiar with Sonnet’s module-oriented programming model who want a similar interface in JAX.

Documented uses

  • Define reusable neural-network components and modules.
  • Manage model parameters and initialize them from example inputs.
  • Transform model functions for use with JAX tools such as automatic differentiation and JIT compilation.
Before you choose

Requirements & limits

The project is available as a public repository under the Apache-2.0 license. Installation instructions offer pip installation from PyPI or the repository.

  • JAX must be installed separately, with the installation choice depending on accelerator requirements.
  • Haiku is in maintenance mode; the project recommends Flax for new projects.
  • Haiku does not define custom optimizers, checkpointing formats or replication APIs.
  • The library is written in Python and depends on C++ code through JAX.

Check current requirements in the official documentation. Source-based guidance does not mean OpenViST has run or independently tested this project.

Source-based practical guide

Install Haiku after setting up JAX

Prepare a Python environment with JAX and Haiku installed.

Before you start

  • A Python environment.
  • JAX installed with accelerator support appropriate to the environment.

Steps

  1. Follow the JAX installation instructions for the relevant accelerator setup.

  2. Install Haiku with pip, either from PyPI using the command below or from the repository using the command below.

    Documented command
    pip install -U dm-haiku
    Documented command
    pip install git+https://github.com/deepmind/dm-haiku

Expected outcome

The environment has JAX and Haiku installed for use in Python projects.

Guide limitations

  • Haiku does not list JAX as a dependency because JAX installation depends on CUDA and accelerator configuration.
  • Examples may require additional packages.
Documentation 1 ↗
Source-based practical guide

Transform a model function for JAX

Turn a Haiku module-based model function into initialization and application functions.

Before you start

  • Python with JAX and Haiku installed.
  • A model function that defines or applies Haiku modules.

Steps

  1. Define the model function using Haiku modules such as hk.Linear.

  2. Pass the function to hk.transform to obtain pure init and apply functions.

  3. Initialize parameters by calling init with a random key and example inputs.

Expected outcome

The transformed model has an initialization function that returns its parameters and an application function suitable for JAX transformations.

Guide limitations

  • The repository’s example training loop also uses JAX and an input dataset; Haiku does not provide an optimizer, checkpointing format or replication API.
Documentation 1 ↗
Evidence and transparency

Sources & review

Official repository ↗
Source record details
Source availability
Active
Sources checked
10 October 2026
Original website title
google-deepmind/dm-haiku · Official homepage title · retained in its source language
How we review sourcesSuggest a correction

Related options

Flax — The Haiku repository recommends Flax for new projects and describes it as having a broader feature set.