用 JAX 实现自动微分、vmap 与 Flax 线性回归训练

JAX 是一个可微的 NumPy：它的数组接口与 NumPy 高度相似，同一段代码可以不加修改地跑在 CPU、GPU 和 TPU 上，同时提供自动微分、自动向量化和即时编译能力。它源自 autograd，采用纯函数式编程范式——函数不产生副作用，状态由调用者显式传入传出。这篇教程面向已经用过 PyTorch、TensorFlow 或 NumPy 的开发者，目标是带你从数组、梯度一路走到手写训练循环、vmap 向量化，最后用 Optax 与 Flax 搭出一个小型编码器/解码器网络。如果你更想先确认站内有没有配套的在线运行环境，可以看 JAX 工具页。

准备工作

开始之前需要确认几件事。

• Python 环境：需要一个可用的 Python 解释器与 pip。素材中的示例在 Python 3.9 环境下运行过，但 JAX 对具体小版本的要求请以官网当前信息为准。
• 安装 JAX：至少需要 jax 与 jaxlib 两个包。CPU 版本安装最简单；要使用 GPU 或 TPU，需要按官网说明选择对应的加速器版本，安装命令与可用平台以官网当前信息为准。
• 可选依赖：本文示例还会用到 numpy、matplotlib（画散点图与拟合直线）、optax（优化器）和 flax（神经网络层）。这些都不是 JAX 核心的一部分，需要单独安装。
• 前置知识：熟悉 NumPy 的数组操作、了解梯度下降的基本形式、见过 PyTorch 的 nn.Module 写法，会让后面的内容顺畅很多。

导入方式很直接：

import jax
import jax.numpy as jnp

如果机器上没有找到 GPU 或 TPU，JAX 会打印一条警告并回退到 CPU 执行。这条警告本身不影响结果，只是提示当前后端。想看到更详细的日志，可以按提示设置环境变量 TF_CPP_MIN_LOG_LEVEL=0 后重新运行。

操作步骤

第一步：创建并理解 JAX 数组

JAX 中对应 np.array、tf.Tensor、torch.Tensor 的类型就是 JAX 数组。创建一个向量：

x = jnp.arange(10)
print(x)
# [0 1 2 3 4 5 6 7 8 9]

这段代码在 CPU、GPU、TPU 上可以逐字运行，不需要改动。与 NumPy 数组最关键的区别是：JAX 数组永远是不可变的。任何“就地修改”的想法都要换成返回新数组的写法，这一点会贯穿后面所有步骤。

第二步：测量点积与后端差异

用两个长度为一千万的向量做点积，可以直观看到后端差异：

y = jnp.arange(int(1e7))
%timeit jnp.dot(y, y).block_until_ready()

注意末尾的 block_until_ready()。JAX 的运算默认是异步派发的，如果不等待结果就计时，测到的只是“派发耗时”而不是“计算耗时”。加上它才能得到真实数字。素材中给出的参考量级是：CPU 上每次循环约 10.6 ms，GPU 上约 362 µs。具体数值取决于硬件与版本，只作量级参考。

第三步：用 jax.grad 求梯度

对一个用 Python 写的数值函数 f，把它包进 jax.grad，会返回一个新的 Python 函数，这个新函数计算原函数的梯度。数学上即 (∇f)(x)ᵢ = ∂f/∂xᵢ。

取输入数组 [1, 2, 3, 4]，定义 f(x) = Σxᵢ²，其值为 1+4+9+16 = 30，梯度为 Σ2xᵢ：

def sum_of_squares(x):
    return jnp.sum(x ** 2)

sum_of_squares_dx = jax.grad(sum_of_squares)
x = jnp.asarray([1., 2., 3., 4.])
print(sum_of_squares(x))      # 30.0
print(sum_of_squares_dx(x))   # [2. 4. 6. 8.]

这里要特别注意与 TensorFlow、PyTorch 的习惯差异：JAX 不通过损失张量本身求梯度，没有 loss.backward() 这样的调用。梯度是对函数求的，得到的是另一个函数。

第四步：处理多参数函数与 argnums

当函数有多个输入参数时，jax.grad() 默认只对第一个参数求梯度。定义 f(x, y) = Σ(xᵢ − yᵢ)²，则对 x 的梯度是 2(x − y)，对 y 的梯度是 2(y − x)：

def sum_squared_error(x, y):
    return jnp.sum((x - y) ** 2)

sum_squared_error_dx = jax.grad(sum_squared_error)
y = jnp.asarray([1.1, 2.1, 3.1, 4.1])
print(sum_squared_error_dx(x, y))
# [-0.20000005 -0.19999981 -0.19999981 -0.19999981]

可以和解析解对照：2 * (x - y) 与 2 * (y - x) 的结果与上面一致（浮点误差导致的末位差异属正常现象）。

如果要对其他参数、或同时对多个参数求梯度，用 argnums 指定参数下标，传元组即可返回多个梯度：

jax.grad(sum_squared_error, argnums=(0, 1))(x, y)
# (Array([-0.20000005, -0.19999981, -0.19999981, -0.19999981], dtype=float32),
#  Array([0.20000005, 0.19999981, 0.19999981, 0.19999981], dtype=float32))

第五步：用 .at[] 替代就地赋值

副作用指函数产生了不出现在返回值里的效果，就地修改数组就是典型例子。在 NumPy 里这样写是合法的：

import numpy as np
x = np.array([1, 2, 3])

def in_place_modify(x):
    x[0] = 123
    return None

in_place_modify(x)
x   # array([123, 2, 3])

把这个函数套到 JAX 数组上会直接抛出 TypeError，提示该对象不支持元素赋值，并建议改用 x.at[idx].set(y) 或其他 .at[] 方法。正确写法是：

x = jnp.array([5, 4, 3, 2])
# x[1] = 10  -> 抛出 TypeError
x = x.at[1].set(10)
print(x)   # [5 10 3 2]

记住这条替换规则：凡是 x[idx] = y，都改写成 x = x.at[idx].set(y)，或者使用 .at[] 下的其他方法。

第六步：手写线性回归训练循环

先造数据。真实关系为 y = w_true·x + b_true + ε，这里取 w_true = 3、b_true = −1，噪声标准差 0.1，样本 100 个：

import numpy as np
import matplotlib.pyplot as plt

xs = np.random.normal(size=(100,))
noise = np.random.normal(scale=0.1, size=(100,))
ys = xs * 3 - 1 + noise
plt.scatter(xs, ys)

模型为 ŷ(x; θ) = wx + b，把两个参数放进同一个数组 theta = [w, b]：

def model(theta, x):
    """Computes wx + b on a batch of input x."""
    w, b = theta
    return w * x + b

损失函数取均方误差 J(x, y; θ) = (ŷ − y)²：

def loss_fn(theta, x, y):
    prediction = model(theta, x)
    return jnp.mean((prediction - y) ** 2)

更新规则就是最朴素的 SGD：θ_new = θ − lr·∇θJ。因为 JAX 的梯度是对函数求的，更新函数可以写得非常短：

def update(theta, x, y, lr=0.1):
    return theta - lr * jax.grad(loss_fn)(theta, x, y)

theta = jnp.array([1., 1.])
for _ in range(1000):
    theta = update(theta, xs, ys)

plt.scatter(xs, ys)
plt.plot(xs, model(theta, xs))
w, b = theta
print(f"w: {w:.2f}, b: {b:.2f}")
# w: 3.00, b: -1.00

训练 1000 步后参数收敛到接近真实值。注意整个循环里没有任何“优化器对象”，参数就是普通数组，每次迭代由 update 返回新数组——这正是纯函数式风格在训练中的体现。

第七步：用 vmap 做自动向量化

写损失函数时经常要处理“单样本”逻辑，再套到 batch 上。手动维护 batch 维度容易出错，jax.vmap 可以把一个处理单样本的函数自动变成处理批量数据的函数。例如给二元交叉熵加 logits 的实现加上装饰器：

@jax.vmap
def binary_cross_entropy_with_logits(logits, labels):
    logits = nn.log_sigmoid(logits)
    return -jnp.sum(labels * logits + (1. - labels) * jnp.log(-jnp.expm1(logits)))

这样函数体里只写单个样本的数学表达式，批量维度由 vmap 负责，代码里不必到处携带 batch 维度。

第八步：用 Optax 替换手写更新

Optax 是配合 JAX 训练模型常用的优化库，同样以函数式方式使用：先构造优化器，初始化状态，每次迭代拿到更新量与新状态，再把更新量应用到参数上。

import optax

tx = optax.adam(learning_rate=0.03)
opt_state = tx.init(variables)
updates, opt_state = tx.update(grads, opt_state)
variables = optax.apply_updates(variables, updates)

Optax 覆盖的范围很广，从最基础的 SGD（optax.sgd(learning_rate, momentum)）到 AdamW、Lion 等较新的变体都有。具体可用的优化器列表与参数默认值以官网当前信息为准。

第九步：用 Flax Linen 定义网络

Flax 在 JAX 生态里扮演 PyTorch 中 torch.nn、TensorFlow 中 tf.keras 的角色，提供基础神经网络层。导入方式：

import flax
import flax.linen as nn

下面用编码器/解码器结构演示。与 PyTorch 不同，这里不在 __init__ 里定义参数、也不单独写 forward，而是用一个 __call__ 函数以函数式风格封装全部逻辑：

class Encoder(nn.Module):
    latents: int

    @nn.compact
    def __call__(self, x):
        x = nn.Dense(500, name='fc1')(x)
        x = nn.relu(x)
        mean_x = nn.Dense(self.latents, name='fc2_mean')(x)
        logvar_x = nn.Dense(self.latents, name='fc2_logvar')(x)
        return mean_x, logvar_x


class Decoder(nn.Module):
    @nn.compact
    def __call__(self, z):
        z = nn.Dense(500, name='fc1')(z)
        z = nn.relu(z)
        z = nn.Dense(784, name='fc2')(z)
        return z

Linen 中定义子模块和变量有两种方式：

• 显式方式（setup）：在 setup 方法里把子模块或变量赋给 self.<attr>，然后在类的任意前向方法中使用。这与 PyTorch 定义模块的方式接近。
• 内联方式（nn.compact）：在标注了 @nn.compact 的单个前向方法里直接写网络逻辑，子模块和变量就定义在使用它们的位置旁边。

内联方式的好处是：定义与使用挨在一起，不必上下翻找；遇到条件分支或循环时才定义子模块，可以减少重复代码；代码形态也更接近数学记号，比如 self.param('W', ...) @ x + self.param('b', ...) 读起来就像 y = Wx + b。

作为对照，用 setup 写一个 VAE 组合模块：

class VAE(nn.Module):
    latents: int = 20

    def setup(self):
        self.encoder = Encoder(self.latents)
        self.decoder = Decoder()

    def __call__(self, x, z_rng):
        mean, logvar = self.encoder(x)
        z = reparameterize(z_rng, mean, logvar)
        recon_x = self.decoder(z)
        return recon_x, mean, logvar

    def generate(self, z):
        return nn.sigmoid(self.decoder(z))

第十步：用 tabulate 统计参数量

Flax 提供了一个很实用的能力：为任意网络生成参数摘要，入口是 flax.linen.Module 的 tabulate() 方法。以 28×28 的 MNIST 图像（即 x ∈ ℝ⁷⁸⁴）、潜变量 z ∈ ℝ¹⁰ 为例，上面那个简单编码器的参数量在 40 万左右，属于很小的网络：

encoder = Encoder(latents=10)
print(encoder.tabulate(jax.random.PRNGKey(0), jnp.ones((128, 784))))

输出是一张表格，逐层列出 path、module、inputs、outputs 与 params。可以看到 fc1 这一层是 Dense，输入 float32[128,784]，输出 float32[128,500]，参数为 bias float32[500] 与 kernel float32[784,500]，合计 392,500 个（约 1.6 MB）；fc2_mean 与 fc2_logvar 各自是 bias float32[10] 加 kernel float32[500,10]，各 5,010 个（约 20.0 KB）。注意 tabulate 需要传入一个 PRNG key 和一份形状正确的假输入，用来推导各层形状。

一个完整示例

把前面的片段串起来，下面是一个从头到尾可运行的最小线性回归训练脚本，只依赖 jax、jax.numpy 与 numpy：

import numpy as np
import jax
import jax.numpy as jnp

# 1. 造数据：y = 3x - 1 + 噪声
rng = np.random.default_rng(0)
xs = rng.normal(size=(100,))
noise = rng.normal(scale=0.1, size=(100,))
ys = xs * 3 - 1 + noise

# 2. 模型：theta = [w, b]
def model(theta, x):
    w, b = theta
    return w * x + b

# 3. 损失：均方误差
def loss_fn(theta, x, y):
    prediction = model(theta, x)
    return jnp.mean((prediction - y) ** 2)

# 4. 单步 SGD 更新
def update(theta, x, y, lr=0.1):
    return theta - lr * jax.grad(loss_fn)(theta, x, y)

# 5. 训练
theta = jnp.array([1., 1.])
for step in range(1000):
    theta = update(theta, xs, ys)

w, b = theta
print(f"w: {w:.2f}, b: {b:.2f}")
print(f"final loss: {loss_fn(theta, xs, ys):.6f}")

运行后应看到 w 接近 3.00、b 接近 −1.00。如果想把优化器换成 Adam，只需把 update 里的手写公式替换为 Optax 的 tx.init / tx.update / optax.apply_updates 三步，参数数组本身不需要改结构。

注意事项

• 数组不可变：任何 x[idx] = y 都会抛 TypeError，必须写成 x = x.at[idx].set(y)。这是从 NumPy 迁移过来最容易踩的坑。
• 不要写副作用：函数不应修改外部状态或入参，所有状态通过参数传入、通过返回值传出。JAX 的变换（grad、vmap、jit）都建立在这个前提上。
• 梯度是对函数求的：没有 loss.backward()。多参数函数默认只对第一个参数求梯度，其余参数要用 argnums 显式指定。
• 计时前先同步：JAX 运算异步派发，测性能时务必在表达式后加 .block_until_ready()，否则数字没有意义。
• 后端回退：找不到 GPU/TPU 时会打印警告并回退到 CPU，代码仍可运行，但速度差异明显。
• 浮点误差：梯度结果与解析解对照时可能出现末位差异（如 −0.20000005 对 −0.2），属正常现象。
• tabulate 需要假输入：调用 tabulate 时必须提供 PRNG key 和形状正确的输入数组，否则无法推导各层形状。
• 版本与平台信息：安装方式、可用加速器、Optax 与 Flax 的接口细节都可能随版本变化，请以官网当前信息为准。