google-deepmind/dm-haiku
JAX-based neural network library
Provides object-oriented module abstractions (`hk.Module`) that abstract parameter management while preserving JAX's pure function transformations (`jax.jit`, `jax.grad`, `jax.pmap`). The `hk.transform` API converts functionally impure module code into pure `init` and `apply` functions, enabling deterministic RNG handling via `hk.next_rng_key()`. Designed as a library—not a framework—with Sonnet-compatible APIs, it integrates seamlessly with JAX's ecosystem and avoids reinventing optimizers or checkpointing.
3,199 stars. Actively maintained with 9 commits in the last 30 days.
Stars
3,199
Forks
282
Language
Python
License
Apache-2.0
Category
Last pushed
Mar 12, 2026
Commits (30d)
9
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
google-deepmind/optax
Optax is a gradient processing and optimization library for JAX.
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/
explosion/thinc
🔮 A refreshing functional take on deep learning, compatible with your favorite libraries
extropic-ai/thrml
Thermodynamic Hypergraphical Model Library in JAX