3. 函数变换基础

JAX 与 NumPy 最本质的区别，不在于它能算矩阵，而在于它把“函数”本身当作可以加工的对象。JAX 提供了一组函数变换（function transformations）：输入一个普通 Python 函数，输出一个具有新能力的新函数。本节要建立的，就是这四个最常用变换的定位感——它们各自解决什么问题、为什么必须组合、以及组合时对函数本身提出了什么要求。

四个变换各自的“职责边界”

先给出一张定位表，避免把四个变换混为一谈：

变换作用对象解决的问题典型输入输出
jax.jit整个函数把 Python 执行编译成 XLA 计算图，减少解释开销函数 → 编译后的函数
jax.grad标量输出函数自动求梯度（反向模式）函数 → 梯度函数
jax.vmap含批维的函数把单样本函数自动向量化为批处理函数 → 批处理函数
jax.pmap含批维的函数把批维切分到多设备并行函数 → 多设备函数

关键认识是：它们不是四个并列的工具，而是四个正交的“维度”。jit 管的是“怎么执行”，grad 管的是“对谁求导”，vmap 管的是“沿哪个轴并行”，pmap 管的是“在哪些设备上并行”。正因为正交，它们才能任意嵌套组合，且组合顺序往往可以互换（在满足纯函数前提下）。

为什么必须要求纯函数

JAX 的变换在“追踪”（tracing）阶段执行你的 Python 代码，此时参数是抽象的 tracer 而非真实数值。这意味着函数体内不能依赖真实数值做控制流，也不能有副作用。具体来说，以下写法会出问题：

  • Python 副作用：在函数里 print、修改全局列表、写文件。追踪时只执行一次，实际运行时不会按预期重复。
  • 依赖具体值的分支：if x > 0: 在 jit 下会因 tracer 无法比较而报错，应改用 jax.lax.cond 或 jnp.where。
  • 原地修改数组：arr[0] = 1 不被支持，应使用 arr.at[0].set(1) 的函数式写法。
  • 随机数：不能用全局状态，必须显式传入 PRNGKey。
纯函数要求不是“风格建议”，而是变换能正确工作的前提。一旦函数有副作用，jit 的缓存会复用第一次追踪的结果，导致行为与直觉不符，且极难排查。

组合方式：从单样本到批量训练

考虑一个最小但完整的例子：线性回归的单样本损失。先定义纯函数形式的损失：

import jax
import jax.numpy as jnp

def loss(params, x, y):
    w, b = params
    pred = w * x + b
    return (pred - y) ** 2

现在逐层叠加变换。第一步，对参数求梯度：

grad_fn = jax.grad(loss)          # 对标量输出求导，返回 (dw, db)

第二步，把“单样本”变成“一批样本”。注意 in_axes 指定每个参数沿哪个轴批处理，params 不批处理所以是 None：

batch_grad = jax.vmap(grad_fn, in_axes=(None, 0, 0))
# 输入: params, x[batch], y[batch]
# 输出: 每个样本的梯度, 形状 (batch, 2)

第三步，把批梯度聚合成一次参数更新，并整体编译：

def step(params, x, y, lr=0.1):
    grads = batch_grad(params, x, y)      # (batch, 2)
    g = jax.tree.map(lambda t: t.mean(axis=0), grads)
    return jax.tree.map(lambda p, gi: p - lr * gi, params, g)

step_jit = jax.jit(step)

这里 jax.tree.map 对参数 pytree（元组 (w, b)）逐叶子操作，是 JAX 处理结构化参数的惯用手法。整个 step 被 jit 编译后，vmap 与 grad 都已被融合进同一张 XLA 图，运行时只有一次编译调用。

变换链的嵌套结构

下图把上面的组合关系画成一条“变换链”，重点在于看清每个变换作用在哪一层、输入输出的形状如何变化：

变换链：loss → grad → vmap → jit loss(params,x,y) 标量输出 jax.grad → (dw, db) jax.vmap in_axes=(None,0,0) jax.jit 编译为 XLA 图 单样本梯度 (batch, 2) 一次编译调用 组合语义（由内向外） 1. grad 把 loss 变成梯度函数 2. vmap 把梯度函数沿样本轴并行 3. jit 把整条链编译成一张图 等价写法： jax.jit(jax.vmap(jax.grad(loss), in_axes=(None,0,0))) 纯函数约束（追踪期） ✗ print / 全局状态 / 写文件 ✗ if x > 0（依赖 tracer 值） ✗ arr[0] = 1（原地修改） ✗ 全局随机数状态 ✓ lax.cond / jnp.where / .at[].set ✓ 显式传入 PRNGKey

组合顺序与常见陷阱

上面的链是 jit(vmap(grad(loss)))。由于四个变换正交，很多顺序在数学上等价，但工程上仍有讲究：

  • jit 放最外层：让 vmap 和 grad 的展开都发生在编译图内部，避免多次编译。若把 jit 放内层，外层 vmap 会触发多次编译，反而变慢。
  • grad 与 vmap 可互换：vmap(grad(f)) 与 grad(vmap(f)) 在 f 为逐样本标量时结果一致，但前者更直观，且能避免对批维求和的歧义。
  • pmap 替代 vmap 的位置：当需要多设备时，把 vmap 换成 pmap，但 pmap 要求批维大小等于设备数，且函数内不能有跨设备通信以外的隐式依赖。
判断组合是否正确的实用方法：先用小批量（如 batch=4）在 CPU 上跑通 jit(vmap(grad(f)))，确认输出形状与数值合理，再切换到 pmap 或更大批量。形状错误几乎总是 in_axes 写错导致的。

一个可复现的验证片段

下面这段代码可以直接运行，用来验证变换链的输出形状与梯度正确性（数值为示例，非实测）：

params = (jnp.array(0.5), jnp.array(0.1))
x = jnp.array([1.0, 2.0, 3.0, 4.0])
y = 2.0 * x + 1.0

g = batch_grad(params, x, y)
print(g.shape)          # (4, 2)，每个样本一组 (dw, db)
print(step_jit(params, x, y))  # 更新后的 (w, b)

预期观察：g.shape 为 (4, 2)，说明 vmap 正确地把 4 个样本各自求导；step_jit 返回的 w 应向 2.0 靠近、b 应向 1.0 靠近。若形状是 (2,)，说明 vmap 的 in_axes 漏了样本轴；若运行报 tracer 比较错误，说明函数体内混入了依赖具体值的 Python 分支。

把四个变换理解为“执行方式、求导、批并行、设备并行”四个正交维度，再牢记纯函数这一前提，就能在后续的优化器、训练循环与多设备扩展中，稳定地把它们像积木一样组合起来，而不至于在形状或追踪错误上反复踩坑。

← 上一章 📖 返回目录 下一章 →