Equinox

Equinox is your one-stop JAX library, for everything you need that isn't already in core JAX:

neural networks (or more generally any model), with easy-to-use PyTorch-like syntax;
filtered APIs for transformations;
useful PyTree manipulation routines;
advanced features like runtime errors;

and best of all, Equinox isn't a framework: everything you write in Equinox is compatible with anything else in JAX or the ecosystem.

If you're completely new to JAX, then start with this CNN on MNIST example.

Coming from Flax or Haiku? The main difference is that Equinox (a) offers a lot of advanced features not found in these libraries, like PyTree manipulation or runtime errors; (b) has a simpler way of building models: they're just PyTrees, so they can pass across JIT/grad/etc. boundaries smoothly.

Installation

Requires Python 3.10+.

pip install equinox

Equinox is also available through a community-supported build on conda-forge.

Documentation

Available at https://docs.kidger.site/equinox.

Quick example

Models are defined using PyTorch-like syntax:

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

and are fully compatible with normal JAX operations:

@jax.jit
@jax.grad
def loss_fn(model, x, y):
    pred_y = jax.vmap(model)(x)
    return jax.numpy.mean((y - pred_y) ** 2)

batch_size, in_size, out_size = 32, 2, 3
model = Linear(in_size, out_size, key=jax.random.PRNGKey(0))
x = jax.numpy.zeros((batch_size, in_size))
y = jax.numpy.zeros((batch_size, out_size))
grads = loss_fn(model, x, y)

Finally, there's no magic behind the scenes. All eqx.Module does is register your class as a PyTree. From that point onwards, JAX already knows how to work with PyTrees.

Citation

If you found this library to be useful in academic work, then please cite: (arXiv link)

@article{kidger2021equinox,
    author={Patrick Kidger and Cristian Garcia},
    title={{E}quinox: neural networks in {JAX} via callable {P}y{T}rees and filtered transformations},
    year={2021},
    journal={Differentiable Programming workshop at Neural Information Processing Systems 2021}
}

(Also consider starring the project on GitHub.)

Metadata

Release files for equinox 0.13.8

For a detailed explanation of source distributions (sdists) and built distributions (wheels), please see the package formats documentation.

Source distribution (sdist)

Source distribution for equinox 0.13.8
File	Size	Uploaded
equinox-0.13.8.tar.gz	145.3 kB	May 5, 2026	Details

Built distribution (wheel)

Table of built distributions (wheels) for equinox 0.13.8
File	Interpreter	ABI	Platform	Reset
equinox-0.13.8-py3-none-any.whl 185.8 kB May 5, 2026	Python 3	none	any	Details

Total release size: 331.1 kB

Release files / equinox-0.13.8.tar.gz

Download URL	equinox-0.13.8.tar.gz
Size	145.3 kB
Tags	Source
SHA-256 checksum How to use checksums	`dd075050018e2dd02e252e9d29d3060f7e67f085622d8d27a8e89e24bb8523db`
BLAKE2b-256 checksum How to use checksums	`19ff522336d2f8264f2ad97119710b76e2cddf66145d03a1e89899175d26b192`
Upload date	May 5, 2026
Uploaded using Trusted Publishing? What is trusted publishing?	No
Uploaded via	`twine/6.2.0 CPython/3.11.15`

Release files / equinox-0.13.8-py3-none-any.whl

Download URL	equinox-0.13.8-py3-none-any.whl
Size	185.8 kB
Tags	Python 3
SHA-256 checksum How to use checksums	`ca004348533cc30a63ebe8823d7dd4bb626dce17743d40bbddb89b402ef2a240`
BLAKE2b-256 checksum How to use checksums	`aed669a76c8ccdef14af687c497040292a46e59fc7a0ab24724b60e50ca61030`
Upload date	May 5, 2026
Uploaded using Trusted Publishing? What is trusted publishing?	No
Uploaded via	`twine/6.2.0 CPython/3.11.15`