google-deepmind/dm-haiku
JAX-based neural network library
This is a library for machine learning researchers and engineers who build neural networks using JAX. It helps convert Python functions that define neural network architectures, including their learnable parameters, into pure functions compatible with JAX's transformations. You provide a description of your network, and Haiku gives you back functions to initialize and apply it, making it easier to manage model state and integrate with JAX's automatic differentiation and parallelization tools.
3,199 stars. Actively maintained with 7 commits in the last 30 days.
Use this if you are developing neural network models with JAX and prefer an object-oriented programming style for defining your network layers, similar to Sonnet for TensorFlow.
Not ideal if you are starting a new project with JAX, as Google DeepMind now recommends Flax, which offers more features, documentation, and a larger community.
Stars
3,199
Forks
282
Language
Python
License
Apache-2.0
Category
Last pushed
Mar 12, 2026
Commits (30d)
7
Get this data via API
curl "https://pt-edge.onrender.com/api/v1/quality/ml-frameworks/google-deepmind/dm-haiku"
Open to everyone — 100 requests/day, no key needed. Get a free key for 1,000/day.
Related frameworks
explosion/thinc
🔮 A refreshing functional take on deep learning, compatible with your favorite libraries
google-deepmind/optax
Optax is a gradient processing and optimization library for JAX.
patrick-kidger/diffrax
Numerical differential equation solvers in JAX. Autodifferentiable and GPU-capable....
google/grain
Library for reading and processing ML training data.
patrick-kidger/equinox
Elegant easy-to-use neural networks + scientific computing in JAX. https://docs.kidger.site/equinox/