4. 自动微分：grad函数、高阶导数、jacfwd与jacrev、vjp与jvp

自动微分，说白了就是让框架帮你算导数。你写个函数，JAX 自动给你把梯度算出来。这事在深度学习里太常见了——反向传播嘛。但 JAX 的自动微分比 PyTorch 或 TensorFlow 要「纯粹」得多，它不维护计算图，而是用函数式变换的方式来做。

我个人习惯把自动微分分成两类：前向模式和反向模式。JAX 两者都支持，而且用起来非常灵活。咱们一个一个来看。

4.1 grad：最常用的梯度函数

grad 是 JAX 里最基础的自动微分工具。你给它一个函数，它返回一个新函数，这个新函数计算原函数的梯度。

import jax
import jax.numpy as jnp

def f(x):
    return x ** 2 + 3 * x + 1

grad_f = jax.grad(f)
print(grad_f(2.0))  # 输出 7.0，因为 f'(x) = 2x + 3

嗯，这里要注意：grad 默认只对第一个参数求导。如果你的函数有多个参数，需要指定 argnums。

def g(x, y):
    return x ** 2 + y ** 3

grad_g_x = jax.grad(g, argnums=0)
grad_g_y = jax.grad(g, argnums=1)

print(grad_g_x(2.0, 3.0))  # 4.0
print(grad_g_y(2.0, 3.0))  # 27.0

4.2 高阶导数：grad 套 grad

JAX 的自动微分是可组合的。你可以对 grad 的结果再求 grad，得到二阶导数。这就是高阶导数的本质——函数变换的嵌套。

def f(x):
    return x ** 3

first_deriv = jax.grad(f)
second_deriv = jax.grad(first_deriv)
third_deriv = jax.grad(second_deriv)

print(first_deriv(2.0))   # 12.0
print(second_deriv(2.0))  # 12.0
print(third_deriv(2.0))   # 6.0

为什么会这样？因为 grad 返回的仍然是一个 JAX 可追踪的函数。你可以无限嵌套下去。不过实际项目中，我很少用到三阶以上——计算量太大，而且数值稳定性会变差。

4.3 jacfwd 与 jacrev：雅可比矩阵的两种计算方式

当你的函数是 R^n → R^m 时，梯度就不够用了。你需要雅可比矩阵——一个 m×n 的矩阵，包含所有偏导数。

JAX 提供了两种方式：jacfwd（前向模式）和 jacrev（反向模式）。

def f(x):
    return jnp.array([x[0] ** 2, x[1] ** 3, x[0] * x[1]])

x = jnp.array([2.0, 3.0])

J_fwd = jax.jacfwd(f)(x)
J_rev = jax.jacrev(f)(x)

print(J_fwd)
# [[4. 0.]
#  [0. 27.]
#  [3. 2.]]

两者结果一样，但计算方式不同。你想想看：jacfwd 适合「输入小、输出大」的场景（比如 n=1, m=1000），而 jacrev 适合「输入大、输出小」的场景（比如 n=1000, m=1）。

方法	适用场景	计算复杂度
jacfwd	输入维度小，输出维度大	O(n) 次前向传播
jacrev	输入维度大，输出维度小	O(m) 次反向传播

4.4 vjp 与 jvp：向量-雅可比积与雅可比-向量积

很多时候你不需要完整的雅可比矩阵，只需要它和一个向量的乘积。这就是 vjp（向量-雅可比积）和 jvp（雅可比-向量积）的用武之地。

jvp 是前向模式：给定一个输入方向向量，计算函数在该方向上的变化率。

def f(x):
    return x ** 2

x = 3.0
tangent = 1.0  # 方向向量

val, jvp_out = jax.jvp(f, (x,), (tangent,))
print(val)      # 9.0
print(jvp_out)  # 6.0

vjp 是反向模式：给定输出侧的向量，反向传播到输入侧。

val, vjp_func = jax.vjp(f, x)
cotangent = 1.0
vjp_out = vjp_func(cotangent)
print(val)      # 9.0
print(vjp_out)  # (6.0,)

说白了，jvp 和 vjp 是构建更复杂自动微分系统的基础。比如在神经网络的 Hessian 向量积计算中，我就经常用 jvp 来避免显式构建 Hessian 矩阵。

4.5 知识体系总览

下面这张图展示了 JAX 自动微分的核心逻辑。你可以看到，所有工具都围绕「函数变换」这个核心理念展开。

4.6 实战建议

最后，给你几个我在实际项目中总结的经验：

• 优先用 grad：如果你的函数输出是标量，直接用 grad，简单高效。
• 雅可比矩阵要谨慎：显式构建雅可比矩阵很贵。能用 vjp/jvp 就别用 jacfwd/jacrev。
• 高阶导数要测试：写个简单函数验证一下数值正确性，别直接上大模型。
• 组合使用：比如 jvp(grad(f)) 可以高效计算 Hessian 向量积，这在二阶优化器中非常有用。

自动微分是 JAX 最强大的特性之一。它让你可以像写数学公式一样写代码，然后自动得到导数。嗯，掌握了这些，你就能在科学计算和深度学习里游刃有余了。