Equinox

Equinox 是你的 JAX 一站式库，提供了核心 JAX 中尚未包含的所有功能：

• 神经网络（或更一般地，任何模型），采用易于使用的类似 PyTorch 的语法；
• 针对变换的过滤 API；
• 有用的 PyTree 操作工具；
• 运行时错误等高级特性；

最重要的是，Equinox 并非一个框架：你在 Equinox 中编写的所有代码都与 JAX 或其生态系统中的其他组件完全兼容。

如果你是 JAX 的新手，请从这个 MNIST 数据集上的 CNN 示例 开始。

你来自 Flax 或 Haiku 吗？主要区别在于，Equinox (a) 提供了许多这些库中没有的高级特性，比如 PyTree 操作或运行时错误；(b) 使用更简单的方式构建模型：它们只是 PyTrees，因此可以顺畅地跨越 JIT/grad 等边界。

安装

需要 Python 3.10 或更高版本。

pip install equinox

Equinox 也可以通过 conda-forge 上社区支持的构建进行安装。

文档

可在 https://docs.kidger.site/equinox 查阅。

快速示例

模型使用类似 PyTorch 的语法定义：

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

并且与常规 JAX 操作完全兼容：

@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)

最后，背后没有任何魔法。eqx.Module 做的只是将你的类注册为一个 PyTree。从那以后，JAX 就已经知道如何处理 PyTrees 了。

引用

如果你在学术工作中发现这个库很有用，请引用以下内容：(arXiv 链接)

@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}
}

（同时也可以考虑在 GitHub 上给该项目加星。）

另请参阅：JAX 生态系统中的其他库

始终有用
jaxtyping：用于数组形状和数据类型的类型注解。

深度学习
Optax：一阶梯度优化器（SGD、Adam 等）。
Orbax：检查点保存（异步/多主机/多设备）。
Levanter：基础模型（如 LLM）的可扩展且可靠的训练。
paramax：PyTrees 的参数化和约束。

科学计算
Diffrax：数值微分方程求解器。
Optimistix：根查找、最小化、不动点和最小二乘法。
Lineax：线性方程组求解器。
BlackJAX：概率和贝叶斯采样。
sympy2jax：SymPy 到 JAX 的转换；可通过梯度下降训练符号表达式。
PySR：符号回归。（非 JAX 的荣誉提及！）

精彩的 JAX 资源
Awesome Equinox
Awesome JAX：一份更长的其他 JAX 项目列表。

Equinox 快速上手指南

Equinox 是一个基于 JAX 的一站式深度学习库，提供类似 PyTorch 的模型定义语法，同时完全兼容 JAX 生态。它不是独立的框架，而是对核心 JAX 功能的增强，支持神经网络构建、过滤变换 API、PyTree 操作及运行时错误检查等高级特性。

环境准备

• Python 版本：3.10 或更高
• 前置依赖：需先安装 JAX（Equinox 安装时通常会自动处理，但建议确保 JAX 可用）
• 系统要求：支持 Linux、macOS 和 Windows（需配置适当后端）

💡 国内用户若遇到 JAX 安装缓慢，可尝试使用清华或阿里镜像源加速 pip 安装。

安装步骤

通过 pip 安装（推荐）：

pip install equinox

如需使用 conda（社区维护版本）：

conda install -c conda-forge equinox

🚀 国内加速建议：

pip install equinox -i https://pypi.tuna.tsinghua.edu.cn/simple

基本使用

Equinox 允许你用类定义模型，语法贴近 PyTorch，且天然支持 JAX 的 jit、grad、vmap 等变换。

1. 定义一个简单的线性模型

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

2. 结合 JAX 进行训练（计算损失与梯度）

@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)

✅ 关键点：eqx.Module 本质是将你的类注册为 PyTree，因此无需特殊处理即可无缝融入 JAX 工作流。

更多示例请参考官方文档：https://docs.kidger.site/equinox