Equinox: callable PyTrees as the model abstraction in JAX
Elegant easy-to-use neural networks + scientific computing in JAX. https://docs.kidger.site/equinox/
At a glance
- What is it?
- Equinox is a JAX library that registers ordinary Python classes as PyTrees, so a model is just a class with arrays in it. The README frames it as a non-framework: nothing you write in it is locked out of the rest of JAX.
- Who is it for?
- Adopt Equinox if you already work in JAX and want model code that survives being passed through jax.jit, jax.grad and jax.vmap without a separate framework state object. Do not adopt it if you want a batteries-included training stack: the README points at Optax, Orbax and Levanter for optimisers, checkpointing and large-scale training rather than providing them.
- Can I use it commercially?
- Yes. Apache-2.0 is a permissive licence: you can use, modify and sell software built on it, as long as you keep its copyright and licence notices.
- Is it still maintained?
- Yes. The repository last received commits 10 days ago.
- What is it written in?
- Mainly Python, according to GitHub's language statistics.
Answers come from the project's GitHub data, last synced on September 15, 2026, and from our analysis. They are not legal advice.
DEEP OPEN-SOURCE ANALYSIS
The problem Equinox solves is JAX's PyTree boundary, not neural networks
JAX already knows how to differentiate, vectorise and compile functions over PyTrees. What it does not know is how to treat your model as one. The usual workaround is a container of parameter arrays that gets threaded through every function by hand, or a framework that keeps parameters and architecture in separate structures and gives you a different API at each boundary. Equinox takes the direct route: eqx.Module registers your class as a PyTree, and from that point onwards, in the README's words, JAX already knows how to work with PyTrees. The audience is therefore narrower than the tagline suggests. It is for people who have already accepted JAX and want their model code to behave like any other value in it. If you have not written JAX before, the README sends you to a CNN on MNIST example rather than to a tutorial on the library itself.
A model is a class with array attributes, and nothing else is registered
The quick example in the README defines Linear as an eqx.Module subclass with two annotated attributes, weight and bias, both typed as jax.Array. The constructor splits a PRNG key with jax.random.split and assigns the two random draws. The forward pass is __call__, a single matrix multiply and bias add. There is no build step, no collection dictionary, no separate apply function. That is the whole mechanism. The README is explicit that there is no magic behind the scenes and that registering the class as a PyTree is all eqx.Module does. The consequence is that the same object can be passed into a jitted, differentiated loss function directly, as the second README example does with a stacked jax.jit and jax.grad decorator over a function that calls jax.vmap(model). Because the model is a PyTree, vmap maps over it the same way it maps over an array.
Filtered transformations and PyTree manipulation are the part Flax and Haiku do not have
The README's own comparison to Flax and Haiku is worth taking at face value because it is unusually specific. It claims two differences: advanced features such as PyTree manipulation and runtime errors, and a simpler construction path because models are PyTrees and can cross JIT, grad and other boundaries smoothly. The filtered APIs are the mechanism behind the first claim. A filtered transformation lets you apply a JAX transformation to some leaves of a PyTree and not others, which is the standard way to freeze part of a model, exclude non-array fields from a gradient, or apply a different update rule per parameter group. In a framework where parameters live in a separate collection, that selection is done by name against the collection. Here it is done against the tree itself. The README does not spell out the filter API surface, so if your workflow depends on it, read the documentation site rather than the repository front page.
Installing it and the version constraints that come with it
Installation is one command, pip install equinox, and the README states the requirement plainly: Python 3.10 or newer. There is also a community-supported conda-forge build, which the README links to as a feedstock rather than describing. No configuration keys, environment variables or registration steps appear in the supplied material, which is consistent with the library's design: there is no global state to configure because a model is just a class. The releases listed run v0.13.6 through v0.13.8, with v0.13.8 dated 2026-05-05, and the repository's last push is later than the most recent release. The 0.x version number is the practical signal here. Pin the version in your requirements file and read the release notes before moving between minor versions, because a library whose entire public surface is a class decorator and a set of transformation wrappers has little room to deprecate things gradually.
Where Equinox is the wrong tool
The README lists what Equinox is not by omission. Optimisers are Optax, checkpointing is Orbax, large-scale foundation model training is Levanter. If you want a single package that takes you from model definition to a trained checkpoint, Equinox is not that package and does not claim to be. The second limitation is more structural. Because a model is an ordinary Python class, anything you put in it that is not a PyTree leaf has to be handled deliberately, and anything you want JAX to see has to be registered. That is the price of the simplicity. A framework with an explicit parameter collection makes the boundary visible by construction; Equinox makes it invisible until you hit a field that should not be differentiated. The README also notes the library is not a framework, which cuts both ways: nothing is locked in, and nothing is provided for free either.
How it differs from Flax in approach
Flax and Haiku, the two libraries the README names, both separate the model definition from the parameter state. You define a module, then initialise it to obtain parameters, then apply it with those parameters passed back in. Equinox collapses those into one object: the module instance holds the arrays and is itself callable. The practical difference shows up at transformation boundaries. With a separate state object, every jitted or differentiated function signature has to carry the parameters alongside the inputs, and any code that wants to inspect or transform the model has to know the collection layout. With a callable PyTree, the model goes in as one argument and comes out as one argument, and tree operations apply to it directly. The trade is that you no longer have a framework telling you which arrays are parameters and which are not; you decide, using filtered transformations.
Maintenance cost and the Apache-2.0 terms
The repository is not archived, the last push is 2026-09-06, and releases have arrived on a roughly monthly cadence across the three most recent versions. That is an actively developed project, but it is also a 0.x project, so the upgrade cost is real: budget for reading release notes and re-running your model tests at each minor bump rather than assuming compatibility. The licence is Apache-2.0, which permits commercial and closed-source use and includes an explicit patent grant, with the usual requirements around preserving notices and stating changes. That is a summary of the identifier, not legal advice; if you are shipping the library inside a product, have counsel read the LICENSE file. There is no stated dual-licensing or commercial tier in the supplied material, and the README asks for a citation in academic work rather than any payment.
Who should take the dependency
The deciding question is whether you have JAX code that already treats models as data and keeps running into the framework boundary. If you are writing research code where a model needs to be passed through jit, grad and vmap in the same function, or where you want to manipulate the model as a tree, Equinox removes a layer rather than adding one. If you are starting from PyTorch and want a familiar training loop, the README's own pointer to the MNIST example is the honest starting point, and you should expect to assemble the optimiser and checkpointing halves yourself from Optax and Orbax. Before committing, check two things against your own code: that your Python is 3.10 or newer, and that you are not relying on Flax or Haiku collection semantics anywhere in the call path, because that is the one boundary Equinox does not smooth over.
Editorial conclusion
Adopt Equinox if you already work in JAX and want model code that survives being passed through jax.jit, jax.grad and jax.vmap without a separate framework state object. Do not adopt it if you want a batteries-included training stack: the README points at Optax, Orbax and Levanter for optimisers, checkpointing and large-scale training rather than providing them. Verify first that your Python is 3.10 or newer and that your existing model code does not depend on Flax or Haiku collection semantics, because that is the boundary Equinox does not cross.
Community notes