OpenViSTAI · Agents · Blockchain · Crypto
Submit a project
Menu
AI Developer Tools & SDKs

Equinox

Equinox is a Python library for defining neural networks and other models as JAX PyTrees, with filtered transformation APIs and utilities for manipulating PyTrees.

Explore the practical guide ↓

Independent OpenViST project profile, based on linked sources.

Equinox — OpenViST project profile in AI Developer Tools & SDKs.
OpenViST project illustration. See the linked sources for product details.
Understand the project

Practical overview

When it fits

Use Equinox when defining models or other structured values for a JAX workflow, especially when those structures need to pass through JAX transformations.

Who it is for

  • Python developers building neural-network models with JAX.
  • Scientific-computing developers who need structured models or PyTree manipulation alongside JAX operations.

Documented uses

  • Define callable model classes whose fields include JAX arrays.
  • Apply JAX transformations such as compilation, differentiation and vectorization to models and loss functions.
  • Manipulate model parameters and other PyTrees.
Before you choose

Requirements & limits

The repository lists Equinox as open source under the Apache-2.0 license. Installation requires Python 3.10 or later; pip installation is documented, alongside a community-supported conda-forge build.

  • Equinox models are registered as PyTrees, which lets JAX work with them through its existing transformations.
  • Equinox supplies model structures and related utilities rather than a standalone training framework.
  • The documented installation requirement is Python 3.10 or later.

Check current requirements in the official documentation. Source-based guidance does not mean OpenViST has run or independently tested this project.

Source-based practical guide

Define a linear model and evaluate a JAX loss gradient

Create an Equinox model as a PyTree and compute gradients of a vectorized loss using JAX.

Before you start

  • Python 3.10 or later
  • Equinox and JAX installed

Steps

  1. Install Equinox with the command below.

    Documented command
    pip install equinox

  2. Define a class inheriting from eqx.Module, with JAX-array weight and bias fields and a __call__ method that computes the model output.

  3. Define a loss function using jax.vmap to apply the model across inputs, then use jax.grad and jax.jit as shown in the repository example.

  4. Create the model with a JAX random key, prepare input and target arrays, and call the loss function to obtain gradients.

Expected outcome

A linear Equinox model represented as a PyTree, with gradients computed for a vectorized loss through JAX transformations.

Guide limitations

  • The example uses zero-valued inputs and targets and computes gradients; it does not document an optimizer or a training loop.
  • The installation instructions specify Python 3.10+.
Documentation 1 ↗
Evidence and transparency

Sources & review

Official repository ↗
Source record details
Source availability
Active
Sources checked
14 October 2026
Original website title
patrick-kidger/equinox · Official homepage title · retained in its source language
How we review sourcesSuggest a correction