JAX 入门

PyTorch 原地修改张量。TensorFlow 构建图。JAX 编译纯函数。最后这一条改变了你思考深度学习的方式。

类型： Build 语言： Python 前置要求： 阶段 03 第 01-10 课、基础 NumPy 预计时间： ~90 分钟

学习目标

• 用 JAX 的函数式 API（jax.numpy、jax.grad、jax.jit、jax.vmap）写纯函数的神经网络代码
• 解释 PyTorch 的即时原地修改和 JAX 的函数式编译模型之间的关键设计差异
• 用 jit 编译和 vmap 向量化加速训练循环，对比朴素的 Python
• 在 JAX 里训练一个简单网络，把显式的状态管理和 PyTorch 的面向对象做法做对比

问题背景

你知道怎么在 PyTorch 里搭神经网络。你定义一个 nn.Module，调 .backward()，让优化器走一步。它管用。几百万人在用它。

但 PyTorch 的 DNA 里刻着一个约束：它即时地、一次一个、在 Python 里追踪操作。每个 tensor + tensor 都是一次单独的内核启动。每个训练步都重新解释同样的 Python 代码。这在你需要跨 2048 张 TPU 训练一个 5400 亿参数模型之前都还行。到那时，开销就要了你的命。

Google DeepMind 在 JAX 上训练 Gemini。Anthropic 在 JAX 上训练了 Claude。这些不是小活——它们是地球上最大的神经网络训练。它们选 JAX，是因为它把你的训练循环当成一个可编译的程序，而不是一连串 Python 调用。

JAX 是带三种超能力的 NumPy：自动微分、JIT 编译到 XLA、自动向量化。你写一个处理单个样本的函数。JAX 给你一个能处理一个批次、算梯度、编译成机器码、跨多个设备运行的函数。全程都不用改原来那个函数。

核心概念

JAX 哲学

JAX 是个函数式框架。没有类，没有可变状态，没有 .backward() 方法。取而代之：

PyTorch	JAX
带状态的 nn.Module 类	纯函数：f(params, x) -> y
loss.backward()	jax.grad(loss_fn)(params, x, y)
即时执行	通过 XLA 做 JIT 编译
for x in batch: 手动循环	jax.vmap(f) 自动向量化
DataParallel / FSDP	jax.pmap(f) 自动并行
可变的 model.parameters()	不可变的数组 pytree

这不是风格偏好。这是编译器约束。JIT 编译要求纯函数——同样的输入永远产出同样的输出，没有副作用。正是那条限制让 100 倍的提速成为可能。

jax.numpy：熟悉的表层

JAX 在加速器上重新实现了 NumPy API：

import jax.numpy as jnp

a = jnp.array([1.0, 2.0, 3.0])
b = jnp.array([4.0, 5.0, 6.0])
c = jnp.dot(a, b)

同样的函数名。同样的广播规则。同样的切片语义。但数组活在 GPU/TPU 上，每个操作都能被编译器追踪。

一个关键区别：JAX 数组是不可变的。没有 a[0] = 5。而是：a = a.at[0].set(5)。这会别扭一个礼拜，然后就开窍了——正是不可变性让 grad、jit、vmap 这些变换可以组合。

jax.grad：函数式自动微分

PyTorch 把梯度挂在张量上（.grad）。JAX 把梯度挂在函数上。

import jax

def f(x):
    return x ** 2

df = jax.grad(f)
df(3.0)

jax.grad 接收一个函数，返回一个计算梯度的新函数。没有 .backward() 调用。没有存在张量上的计算图。梯度就是另一个你可以调用、组合或 JIT 编译的函数。

它能任意组合：

d2f = jax.grad(jax.grad(f))
d2f(3.0)

二阶导。三阶导。雅可比。海森。全靠组合 grad。PyTorch 也能做（torch.autograd.functional.hessian），但那是硬贴上去的。在 JAX 里，这是根基。

约束是：grad 只对纯函数管用。里面不能有 print 语句（它们在追踪时跑，而不是执行时）。不能修改外部状态。不能在没有显式 key 管理的情况下生成随机数。

jit：编译到 XLA

@jax.jit
def train_step(params, x, y):
    loss = loss_fn(params, x, y)
    return loss

fast_step = jax.jit(train_step)

第一次调用时，JAX 追踪这个函数——它记录发生了哪些操作，但不执行它们。然后把这份追踪交给 XLA（Accelerated Linear Algebra，Google 为 TPU 和 GPU 写的编译器）。XLA 融合操作、消除冗余的内存拷贝、生成优化过的机器码。

后续调用完全跳过 Python。编译后的代码以 C++ 速度在加速器上跑。

JIT 何时有帮助：

• 训练步（同样的计算重复几千次）
• 推理（同一个模型，不同输入）
• 任何被以相似形状的输入调用一次以上的函数

JIT 何时反而拖累：

• 含有依赖于值的 Python 控制流的函数（if x > 0，而 x 是个被追踪的数组）
• 一次性计算（编译开销超过运行时间）
• 调试（追踪藏起了实际执行）

控制流限制是真实的。jax.lax.cond 替换 if/else。jax.lax.scan 替换 for 循环。这些不是可选的——它们是编译的代价。

vmap：自动向量化

你写一个处理单个样本的函数：

def predict(params, x):
    return jnp.dot(params['w'], x) + params['b']

vmap 把它提升为处理一个批次：

batch_predict = jax.vmap(predict, in_axes=(None, 0))

in_axes=(None, 0) 意思是：不在 params 上分批（共享），在 x 的第 0 轴上分批。没有手动 for 循环。没有 reshape。没有手动穿引批维度。JAX 自己搞清楚批维度并把整个计算向量化。

这不是语法糖。vmap 生成融合后的向量化代码，比 Python 循环快 10-100 倍。而且它和 jit、grad 能组合：

per_example_grads = jax.vmap(jax.grad(loss_fn), in_axes=(None, 0, 0))

逐样本梯度。一行。这在 PyTorch 里不靠 hack 几乎做不到。

pmap：跨设备的数据并行

parallel_step = jax.pmap(train_step, axis_name='devices')

pmap 把函数复制到所有可用设备（GPU/TPU）上、并切分批次。函数内部，jax.lax.pmean 和 jax.lax.psum 跨设备同步梯度。

Google 用 pmap（及其继任者 shard_map）跨数千张 TPU v5e 芯片训练 Gemini。编程模型是：写单设备版本，用 pmap 一包，搞定。

Pytree：通用数据结构

JAX 操作的是"pytree"——列表、元组、字典和数组的嵌套组合。你的模型参数就是一个 pytree：

params = {
    'layer1': {'w': jnp.zeros((784, 256)), 'b': jnp.zeros(256)},
    'layer2': {'w': jnp.zeros((256, 128)), 'b': jnp.zeros(128)},
    'layer3': {'w': jnp.zeros((128, 10)),  'b': jnp.zeros(10)},
}

每个 JAX 变换——grad、jit、vmap——都知道怎么遍历 pytree。jax.tree.map(f, tree) 把 f 施加到每个叶子上。优化器就是这么一次更新所有参数的：

params = jax.tree.map(lambda p, g: p - lr * g, params, grads)

没有 .parameters() 方法。没有参数注册。树结构就是模型。

函数式 vs 面向对象

PyTorch 把状态存在对象里：

class Model(nn.Module):
    def __init__(self):
        self.linear = nn.Linear(784, 10)

    def forward(self, x):
        return self.linear(x)

JAX 用带显式状态的纯函数：

def predict(params, x):
    return jnp.dot(x, params['w']) + params['b']

params 是传进来的。什么都不存。什么都不改。这让每个函数都可测试、可组合、可编译。这也意味着你自己管理 params——或者用 Flax、Equinox 这类库。

JAX 生态

JAX 给你原语。库给你人体工学：

库	角色	风格
Flax（Google）	神经网络层	带显式状态的 nn.Module
Equinox（Patrick Kidger）	神经网络层	基于 pytree、Python 风
Optax（DeepMind）	优化器 + LR 调度	可组合的梯度变换
Orbax（Google）	检查点	保存/恢复 pytree
CLU（Google）	指标 + 日志	训练循环工具

Optax 是标准的优化器库。它把梯度变换（Adam、SGD、裁剪）和参数更新分开，让组合变得轻而易举：

optimizer = optax.chain(
    optax.clip_by_global_norm(1.0),
    optax.adam(learning_rate=1e-3),
)

何时用 JAX vs PyTorch

因素	JAX	PyTorch
TPU 支持	一等公民（Google 两个都造）	社区维护（torch_xla）
GPU 支持	好（通过 XLA 用 CUDA）	业界最佳（原生 CUDA）
调试	难（追踪 + 编译）	简单（即时、逐行）
生态	偏研究（Flax、Equinox）	庞大（HuggingFace、torchvision 等）
招聘	小众（Google/DeepMind/Anthropic）	主流（到处都是）
大规模训练	更优（XLA、pmap、mesh）	好（FSDP、DeepSpeed）
原型速度	更慢（函数式开销）	更快（改完就跑）
生产推理	TensorFlow Serving、Vertex AI	TorchServe、Triton、ONNX
谁在用	DeepMind（Gemini）、Anthropic（Claude）	Meta（Llama）、OpenAI（GPT）、Stability AI

老实说：除非你有具体理由用 JAX，否则用 PyTorch。那些理由是——能用 TPU、需要逐样本梯度、超大规模的多设备训练，或者你在 Google/DeepMind/Anthropic 上班。

JAX 里的随机数

JAX 没有全局随机状态。每个随机操作都要求一个显式的 PRNG key：

key = jax.random.PRNGKey(42)
key1, key2 = jax.random.split(key)
w = jax.random.normal(key1, shape=(784, 256))

一开始很烦。但它保证了跨设备和跨编译的可复现性——这是 PyTorch 的 torch.manual_seed 在多 GPU 场景下保证不了的属性。

batchnorm-effect

动手构建

我们要用 JAX 和 Optax 在 MNIST 上训练一个 3 层 MLP。784 个输入，两个隐藏层分别 256 和 128 个神经元，10 个输出类。

第 1 步：环境与数据

import jax
import jax.numpy as jnp
from jax import random
import optax

def get_mnist_data():
    from sklearn.datasets import fetch_openml
    mnist = fetch_openml('mnist_784', version=1, as_frame=False, parser='auto')
    X = mnist.data.astype('float32') / 255.0
    y = mnist.target.astype('int')
    X_train, X_test = X[:60000], X[60000:]
    y_train, y_test = y[:60000], y[60000:]
    return X_train, y_train, X_test, y_test

第 2 步：初始化参数

没有类。就一个返回 pytree 的函数：

def init_params(key):
    k1, k2, k3 = random.split(key, 3)
    scale1 = jnp.sqrt(2.0 / 784)
    scale2 = jnp.sqrt(2.0 / 256)
    scale3 = jnp.sqrt(2.0 / 128)
    params = {
        'layer1': {
            'w': scale1 * random.normal(k1, (784, 256)),
            'b': jnp.zeros(256),
        },
        'layer2': {
            'w': scale2 * random.normal(k2, (256, 128)),
            'b': jnp.zeros(128),
        },
        'layer3': {
            'w': scale3 * random.normal(k3, (128, 10)),
            'b': jnp.zeros(10),
        },
    }
    return params

He 初始化，手工做的。从一个种子分出三个 PRNG key。每个权重都是嵌套字典里的一个不可变数组。

第 3 步：前向传播

def forward(params, x):
    x = jnp.dot(x, params['layer1']['w']) + params['layer1']['b']
    x = jax.nn.relu(x)
    x = jnp.dot(x, params['layer2']['w']) + params['layer2']['b']
    x = jax.nn.relu(x)
    x = jnp.dot(x, params['layer3']['w']) + params['layer3']['b']
    return x

def loss_fn(params, x, y):
    logits = forward(params, x)
    one_hot = jax.nn.one_hot(y, 10)
    return -jnp.mean(jnp.sum(jax.nn.log_softmax(logits) * one_hot, axis=-1))

纯函数。params 进，预测出。没有 self，没有存下来的状态。loss_fn 从零算交叉熵——softmax、log、负均值。

第 4 步：JIT 编译的训练步

@jax.jit
def train_step(params, opt_state, x, y):
    loss, grads = jax.value_and_grad(loss_fn)(params, x, y)
    updates, opt_state = optimizer.update(grads, opt_state, params)
    params = optax.apply_updates(params, updates)
    return params, opt_state, loss

@jax.jit
def accuracy(params, x, y):
    logits = forward(params, x)
    preds = jnp.argmax(logits, axis=-1)
    return jnp.mean(preds == y)

jax.value_and_grad 一遍就返回损失值和梯度。@jax.jit 装饰器把两个函数都编译到 XLA。第一次调用之后，每个训练步运行起来都不碰 Python。

第 5 步：训练循环

optimizer = optax.adam(learning_rate=1e-3)

X_train, y_train, X_test, y_test = get_mnist_data()
X_train, X_test = jnp.array(X_train), jnp.array(X_test)
y_train, y_test = jnp.array(y_train), jnp.array(y_test)

key = random.PRNGKey(0)
params = init_params(key)
opt_state = optimizer.init(params)

batch_size = 128
n_epochs = 10

for epoch in range(n_epochs):
    key, subkey = random.split(key)
    perm = random.permutation(subkey, len(X_train))
    X_shuffled = X_train[perm]
    y_shuffled = y_train[perm]

    epoch_loss = 0.0
    n_batches = len(X_train) // batch_size
    for i in range(n_batches):
        start = i * batch_size
        xb = X_shuffled[start:start + batch_size]
        yb = y_shuffled[start:start + batch_size]
        params, opt_state, loss = train_step(params, opt_state, xb, yb)
        epoch_loss += loss

    train_acc = accuracy(params, X_train[:5000], y_train[:5000])
    test_acc = accuracy(params, X_test, y_test)
    print(f"Epoch {epoch + 1:2d} | Loss: {epoch_loss / n_batches:.4f} | "
          f"Train Acc: {train_acc:.4f} | Test Acc: {test_acc:.4f}")

10 个 epoch。约 97% 测试准确率。第一个 epoch 慢（JIT 编译）。第 2-10 个 epoch 快。

注意缺了什么：没有 .zero_grad()，没有 .backward()，没有 .step()。整个更新就是一个组合起来的函数调用。梯度被算出来、被 Adam 变换、被施加到参数上——全在 train_step 里面。

实际使用

Flax：Google 的标准

Flax 是最常见的 JAX 神经网络库。它把 nn.Module 加了回来，但带显式状态管理：

import flax.linen as nn

class MLP(nn.Module):
    @nn.compact
    def __call__(self, x):
        x = nn.Dense(256)(x)
        x = nn.relu(x)
        x = nn.Dense(128)(x)
        x = nn.relu(x)
        x = nn.Dense(10)(x)
        return x

model = MLP()
params = model.init(jax.random.PRNGKey(0), jnp.ones((1, 784)))
logits = model.apply(params, x_batch)

结构和 PyTorch 一样，但 params 和模型是分开的。model.init() 创建 params。model.apply(params, x) 跑前向传播。模型对象没有状态。

Equinox：Python 风的替代方案

Equinox（Patrick Kidger 写的）把模型表示成 pytree：

import equinox as eqx

model = eqx.nn.MLP(
    in_size=784, out_size=10, width_size=256, depth=2,
    activation=jax.nn.relu, key=jax.random.PRNGKey(0)
)
logits = model(x)

模型本身就是个 pytree。不用 .apply()。参数就是模型的叶子。这更贴近 JAX 思考问题的方式。

Optax：可组合的优化器

Optax 把梯度变换和更新解耦：

schedule = optax.warmup_cosine_decay_schedule(
    init_value=0.0, peak_value=1e-3,
    warmup_steps=1000, decay_steps=50000
)

optimizer = optax.chain(
    optax.clip_by_global_norm(1.0),
    optax.adamw(learning_rate=schedule, weight_decay=0.01),
)

梯度裁剪、学习率 warmup、权重衰减——全都组合成一串变换。每个变换看到梯度、改它、传给下一个。没有铁板一块的优化器类。

拿去用

安装：

pip install jax jaxlib optax flax

GPU 支持：

pip install jax[cuda12]

TPU（Google Cloud）：

pip install jax[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html

性能坑：

• 第一次 JIT 调用很慢（编译）。基准测试前先热身。
• 别在 JIT 内部对 JAX 数组做 Python 循环。用 jax.lax.scan 或 jax.lax.fori_loop。
• jax.debug.print() 在 JIT 内部能用。普通 print() 不能。
• 用 jax.profiler 或 TensorBoard 做剖析。XLA 编译可能藏起瓶颈。
• JAX 默认预分配 75% 的 GPU 内存。设 XLA_PYTHON_CLIENT_PREALLOCATE=false 来关掉。

检查点：

import orbax.checkpoint as ocp
checkpointer = ocp.PyTreeCheckpointer()
checkpointer.save('/tmp/model', params)
restored = checkpointer.restore('/tmp/model')

本课产出：

• outputs/prompt-jax-optimizer.md —— 一个挑选合适 JAX 优化器配置的提示词
• outputs/skill-jax-patterns.md —— 一个讲 JAX 函数式套路的 skill

练习

1. 给 MLP 加 dropout。在 JAX 里，dropout 需要一个 PRNG key——把一个 key 穿引过前向传播，并为每个 dropout 层分出 key。对比加和不加的测试准确率。

2. 用 jax.vmap 为一批 32 张 MNIST 图像算逐样本梯度。算每个样本的梯度范数。哪些样本梯度最大，为什么？

3. 把手写的前向函数换成一个对任意层数都适用的通用 mlp_forward(params, x)。用 jax.tree.leaves 自动判断深度。

4. 对带和不带 @jax.jit 的训练步做基准测试。各计时 100 步。在你的硬件上提速有多大？第一次调用的编译开销是多少？

5. 通过组合 optax.chain(optax.clip_by_global_norm(1.0), optax.adam(1e-3)) 实现梯度裁剪。带和不带裁剪各训练一次。把训练过程中的梯度范数画出来看效果。

关键术语

术语	大家怎么说	实际是什么
XLA	"让 JAX 快的那个东西"	Accelerated Linear Algebra——一个编译器，融合操作并从计算图生成优化过的 GPU/TPU 内核
JIT	"即时编译"	JAX 在第一次调用时追踪函数、编译到 XLA，后续调用就跑编译后的版本
纯函数（Pure function）	"没有副作用"	输出只依赖输入的函数——没有全局状态、没有原地修改、没有显式 key 之外的随机性
vmap	"自动分批"	把一个处理单样本的函数变换成处理一个批次的函数，不用重写
pmap	"自动并行"	把一个函数复制到多个设备上、并切分输入批次
Pytree	"嵌套的数组字典"	任何 JAX 能遍历和变换的列表、元组、字典、数组的嵌套结构
追踪（Tracing）	"记录计算"	JAX 用抽象值执行函数来构建计算图，但不计算真实结果
函数式自动微分	"函数的 grad"	通过变换函数来计算导数，而不是给张量挂上梯度存储
Optax	"JAX 的优化器库"	一个可组合的梯度变换库——Adam、SGD、裁剪、调度——能串在一起
Flax	"JAX 的 nn.Module"	Google 为 JAX 写的神经网络库，加上层抽象的同时让状态保持显式

延伸阅读

• JAX 文档：https://jax.readthedocs.io/ —— 官方文档，对 grad、jit、vmap 有很棒的教程
• 《JAX: composable transformations of Python+NumPy programs》（Bradbury 等人，2018）—— 解释设计哲学的原始论文
• Flax 文档：https://flax.readthedocs.io/ —— Google 为 JAX 写的神经网络库
• Patrick Kidger，《Equinox: neural networks in JAX via callable PyTrees and filtered transformations》（2021）—— Flax 的 Python 风替代方案
• DeepMind，《Optax: composable gradient transformation and optimisation》—— 标准的优化器库
• 《You Don't Know JAX》（Colin Raffel，2020）—— 一份实用的 JAX 坑与套路指南，出自 T5 作者之一