4. 一阶梯度计算实战

本节把 JAX 的一阶微分能力落到可运行的三类实战任务上：对标量函数求梯度、对向量函数求雅可比向量积（JVP），以及用数值微分对自动微分结果做交叉验证。三者分别对应“输出是标量”“输出是向量但只关心某个方向”“需要独立可信参照”三种典型场景。理解它们的输入输出形状与调用方式，是后续高阶导数与工程化训练循环的基础。

一、标量函数梯度：grad 的形状规则

JAX 的 jax.grad 只接受返回标量的函数。设 f(x) = x^T A x + b^T x，其中 A 为 3×3 对称矩阵，b 为长度 3 的向量，则梯度为 2Ax + b（A 对称时）。这一解析结果正好用来检验自动微分。

import jax
import jax.numpy as jnp

A = jnp.array([[2.0, 0.5, 0.0],
               [0.5, 1.0, 0.3],
               [0.0, 0.3, 4.0]])
b = jnp.array([1.0, -2.0, 0.5])

def f(x):
    return x @ A @ x + b @ x

x0 = jnp.array([0.2, -0.4, 1.1])
g = jax.grad(f)(x0)          # 形状 (3,)
g_analytic = 2 * A @ x0 + b  # 解析梯度
print(g, g_analytic, jnp.allclose(g, g_analytic, atol=1e-6))

关键规则：grad 的输出形状与输入形状一致。若输入是标量，输出也是标量；若输入是形状 (3,) 的向量，输出仍是 (3,)。当函数有多个参数时，grad 默认只对第一个位置参数求导，其余参数视为常量；若要对指定参数求导，使用 argnums：

def loss(w, x, y):
    return jnp.mean((x @ w - y) ** 2)

gw = jax.grad(loss, argnums=0)(w, x, y)   # 对 w 求导
gx = jax.grad(loss, argnums=1)(w, x, y)   # 对 x 求导

另一个常被忽略的细节是：grad 要求函数输出为标量，若返回形状 (1,) 的数组会直接报错。此时应先用 jnp.sum 或 jnp.mean 归约为标量，再交给 grad。这一约束在写损失函数时尤其重要。

二、向量-雅可比积：jvp 与 vjp 的方向语义

当函数 f: R^n → R^m 输出为向量时，完整雅可比是 m×n 矩阵，但很多场景只需要“雅可比乘一个向量”的结果，即 JVP（Jacobian-Vector Product）或 VJP（Vector-Jacobian Product）。JAX 提供 jax.jvp 与 jax.vjp 两个原语，它们不显式构造雅可比矩阵，计算与存储代价都更低。

设 f(x) = [x0*x1, sin(x0)+x2, x1**2]，在点 x0=(1.0, 2.0, 0.5) 处，取切向量 v=(0.1, -0.2, 0.3)。JVP 计算的是 J @ v，输出形状与 f 的输出一致，为 (3,)：

def fv(x):
    return jnp.array([x[0]*x[1], jnp.sin(x[0]) + x[2], x[1]**2])

x0 = jnp.array([1.0, 2.0, 0.5])
v  = jnp.array([0.1, -0.2, 0.3])
y, jvp = jax.jvp(fv, (x0,), (v,))
# jvp ≈ [0.1*2 + 1*(-0.2), cos(1)*0.1 + 0.3, 2*2*(-0.2)]
#     = [0.0, 0.354, -0.8]

VJP 则计算 u^T @ J，其中 u 是与输出同形状的余切向量，结果形状与输入一致。使用 jax.vjp 会返回一个“拉回函数”，调用它并传入 u 才得到结果：

y, vjp_fn = jax.vjp(fv, x0)
u = jnp.array([1.0, 0.0, 0.0])
(vjp,) = vjp_fn(u)   # 形状 (3,)，即 J 的第一行

方向语义要记牢：JVP 是“前向”的，输入切向量、输出变化率；VJP 是“反向”的，输入输出余切、输出输入梯度。反向模式自动微分（即 grad 背后的机制）本质就是沿计算图反向传播 VJP，因此当输入维度远大于输出维度时，VJP 更划算；反之 JVP 更合适。

三、梯度验证：用数值微分做独立参照

自动微分虽然精确，但实现中仍可能因函数不可微点、stop_gradient 误用、形状广播错误而得到“看起来对、实际错”的梯度。数值微分提供独立参照，常用中心差分：

def numeric_grad(f, x, eps=1e-5):
    g = jnp.zeros_like(x)
    for i in range(x.size):
        e = jnp.zeros_like(x).at[i].set(eps)
        g = g.at[i].set((f(x + e) - f(x - e)) / (2 * eps))
    return g

err = jnp.max(jnp.abs(g - numeric_grad(f, x0)))
print("最大绝对误差:", err)   # 示例中通常在 1e-6 量级

判断标准要结合量级：中心差分的截断误差为 O(eps²)，舍入误差约为 O(机器精度/eps)，因此 eps 取 1e-5 到 1e-6 较平衡。若最大绝对误差在 1e-5 以内，可认为自动微分结果可信；若误差达到 1e-2 以上，应优先检查函数是否在求导点附近不可微（如 abs、max、relu 的折点），或是否存在被 stop_gradient 截断的路径。

四、三种方式的对比与选用

实操建议：先用 grad 处理标量损失并核对形状；需要方向信息时优先 jvp，需要反向传播梯度时用 vjp；每次改动损失函数或引入自定义 custom_jvp/custom_vjp 后，用中心差分做一次小规模抽查。把这三类调用写成可复用的验证脚本，比在训练崩溃后回溯要高效得多。