equinox

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

Equinox provides neural network and model building on top of JAX with PyTorch-like syntax, plus PyTree manipulation, filtered transformations, and runtime error handling—all while remaining fully compatible with core JAX operations.

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

pip install equinox

uv add equinox

poetry add equinox

Installing equinox

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.

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.bias

Requires Python 3.10 or later and JAX installed.

Verify 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

Evidence: equinox-0.13.8-py3-none-any.whl

Keywords: deep-learning, equinox, jax, neural-networks

Tags

More Artificial Intelligence packages