equinox
Elegant easy-to-use neural networks in JAX.
Decision gist · record as of 2026-08-14
Yes. Equinox is actively maintained, has no known vulnerabilities, installs with low friction, and offers a permissive Apache 2.0 license. It fills a genuine gap for developers who want neural network abstractions in JAX without framework overhead. The main caveat is the Alpha development status and the requirement for Python 3.10+; if you need production stability or support for older Python versions, verify API stability first.AI-flagged interpretation of the facts on this page — verify before relying
Before you install
- Requires Python 3.10 or later and JAX installed.
- Low install friction with a pure Python wheel.
- Actively maintained with a recent release; last commit on 2026-08-10 and 2948 GitHub stars indicate ongoing development and community use.
License · maintenance · safety
permissive license (permissive) — Apache License 2.0 is permissive; you may use, modify, and distribute Equinox and derivative works freely, provided you include the license and state any changes.
last release 2026-05-05 (101 days) · last repo commit 2026-08-10 · 2,948 stars
0 known vulnerabilities (OSV.dev, 2026-08-14) · 1,232,737 downloads/mo, #4,183 on PyPI
Alternatives
Verify before relying
pip install equinox
import equinox as eqx
import jax
class Linear(eqx.Module):
weight: jax.Array
bias: jax.Array
def __init__(self, in_size, out_size, key):
wkey, bkey = jax.random.split(key)
self.weight = jax.random.normal(wkey, (out_size, in_size))
self.bias = jax.random.normal(bkey, (out_size,))
def __call__(self, x):
return self.weight @ x + self.bias- Whether Equinox's advanced features (runtime errors, PyTree manipulation) are documented with examples beyond the MNIST tutorial.
- Performance characteristics compared to Flax or Haiku in typical training scenarios.
- Maturity of the API given the 'Alpha' development status classifier.
What it is and what it does
Equinox is a JAX library that fills gaps in core JAX for machine learning by providing neural network and model abstractions with familiar PyTorch-like syntax. Models are defined as PyTrees—JAX's native data structure—so they integrate seamlessly with JAX transformations like jit, grad, and vmap without special handling. The library includes utilities for PyTree manipulation, filtered APIs for transformations, and runtime error support.
Unlike frameworks, Equinox does not enforce a specific training loop or impose constraints on how you use JAX. Everything you write remains compatible with the broader JAX ecosystem. It depends on jax, jaxtyping, typing-extensions, and wadler-lindig, all of which are lightweight. The package is actively maintained and positioned for developers who want neural network convenience without sacrificing JAX's composability and functional programming model.
Use it for
- Building and training neural networks in JAX with PyTorch-familiar class syntax while preserving JAX's functional composition.
- Defining custom models as PyTrees that can be passed directly through jit-compiled and grad-transformed functions.
- Manipulating model parameters and structure using Equinox's PyTree utilities without manual pytree registration.
- Prototyping deep learning research where you need advanced JAX features like vmap and grad but want simpler model definition syntax.
- Migrating from Flax or Haiku to JAX while retaining model-building ergonomics and gaining access to lower-level JAX control.
Worth the install?
AI-flagged interpretation of the facts on this page. Verify before relying on it.
Yes.
Equinox is actively maintained, has no known vulnerabilities, installs with low friction, and offers a permissive Apache 2.0 license. It fills a genuine gap for developers who want neural network abstractions in JAX without framework overhead. The main caveat is the Alpha development status and the requirement for Python 3.10+; if you need production stability or support for older Python versions, verify API stability first.
Install
equinox on PyPI
Before you install
Low install friction with a pure Python wheel. Actively maintained with a recent release; last commit on 2026-08-10 and 2948 GitHub stars indicate ongoing development and community use.
Requires Python 3.10 or later and JAX installed.
License in practice
Apache License 2.0 is permissive; you may use, modify, and distribute Equinox and derivative works freely, provided you include the license and state any changes.
Quickstart
pip install equinox
import equinox as eqx
import jax
class Linear(eqx.Module):
weight: jax.Array
bias: jax.Array
def __init__(self, in_size, out_size, key):
wkey, bkey = jax.random.split(key)
self.weight = jax.random.normal(wkey, (out_size, in_size))
self.bias = jax.random.normal(bkey, (out_size,))
def __call__(self, x):
return self.weight @ x + self.biasVerify before relying
- Whether Equinox's advanced features (runtime errors, PyTree manipulation) are documented with examples beyond the MNIST tutorial.
- Performance characteristics compared to Flax or Haiku in typical training scenarios.
- Maturity of the API given the 'Alpha' development status classifier.
Package facts
| License | permissive license permissive |
| Python support | Supports the current Python release >=3.10 |
| Install friction | Low. Pure-Python wheel |
| Runtime dependencies | 4 packagesjaxjaxtypingtyping-extensionswadler-lindig |
| Maintenance | Actively maintained 101 days since the last release |
| Last repo commit | |
| First released | |
| Downloads | 1,232,737 / month, #4,183 on PyPI 30-day window, as of 2026-08-14 |
| Known vulnerabilities | None known OSV.dev, checked 2026-08-14 |
| Classifiers | Development Status :: 3 - AlphaIntended Audience :: DevelopersIntended Audience :: Financial and Insurance IndustryIntended Audience :: Information TechnologyIntended Audience :: Science/ResearchLicense :: OSI Approved :: Apache Software LicenseNatural Language :: EnglishProgramming Language :: Python :: 3Topic :: Scientific/Engineering :: Artificial IntelligenceTopic :: Scientific/Engineering :: Information AnalysisTopic :: Scientific/Engineering :: Mathematics |
Evidence: equinox-0.13.8-py3-none-any.whl
Tags
Let your AI agent find packages like this
Example. Real query, live index.
You found this page by searching. An agent finds it by wishing: SkillFed indexes 14,416 PyPI packages by what they can do, searchable in plain language.
wish › “jax neural networks”
- equinoxEquinox provides neural network and model building on top of JAX with…
- e3nn-jaxImplements equivariant convolutional neural networks for 3D…
- jraphJraph provides data structures and utilities for building and working…
Give your agent the search over MCP, or paste the wish link into any chat.
More Artificial Intelligence packages
LiteLLM provides a unified Python interface to call 100+ LLM providers (OpenAI, Anthropic, Gemini, Bedrock, Azure, and others) using OpenAI-compatible API format, available as both a Python SDK and a self-hosted AI Gateway proxy server.
Install it if you need to work with multiple LLM providers or want to centralize LLM routing in your organization.
Client library and CLI tool for downloading, uploading, and managing models, datasets, and repositories on the Hugging Face Hub platform.
Install it if you work with Hugging Face Hub models or datasets.
LangChain provides a framework for building agents and LLM-powered applications by composing language models, tools, and memory through a unified API that abstracts over multiple model providers.
hf-xet provides chunk-based deduplication and efficient file transfer for the Hugging Face Hub, enabling faster uploads and downloads of large files with local disk caching.
Tokenizers converts raw text into token sequences for NLP models, with support for training custom vocabularies and using pre-built tokenizers (BPE, WordPiece) optimized for speed via Rust.
Transformers provides a unified framework for loading, fine-tuning, and running state-of-the-art pretrained models across text, vision, audio, video, and multimodal tasks using PyTorch, JAX, or TensorFlow.
Install it if you need to run or train any transformer-based model for NLP, vision, audio, or multimodal tasks.
See also flax · jax · lineax · dm-haiku · jaxtyping · jax-dataclasses · optax · rax · jraph · e3nn-jax