3. 自动微分与梯度计算：grad函数、高阶梯度、jacfwd与jacrev、vmap与pmap的自动微分

自动微分，说白了就是让框架帮我们把导数算好。你想想看，以前我们手写反向传播，一个链式法则推半天，稍微改个网络结构就得重新推导。JAX 的自动微分系统，是我用过最顺手的工具之一。它不像 PyTorch 那样维护一个动态计算图，而是把函数变换玩到了极致。

我个人习惯把 JAX 的自动微分分成三个层次：标量梯度、向量雅可比、批量并行。咱们一层层往下聊。

3.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 (2*2 + 3 = 7)

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

3.2 高阶梯度：grad 的嵌套使用

JAX 的 grad 可以嵌套使用，用来计算高阶导数。比如二阶导数、三阶导数，甚至更高阶。这在实际应用中很有用，比如物理模拟中的拉普拉斯算子，或者优化算法中的 Hessian 矩阵。

def f(x):
    return x**3

# 一阶导数
grad_f = jax.grad(f)
print(grad_f(2.0))  # 12.0 (3*2^2)

# 二阶导数
grad2_f = jax.grad(jax.grad(f))
print(grad2_f(2.0))  # 12.0 (6*2)

# 三阶导数
grad3_f = jax.grad(jax.grad(jax.grad(f)))
print(grad3_f(2.0))  # 6.0 (常数)

为什么会这样？因为 grad 返回的仍然是一个函数，你可以继续对它求导。这就像洋葱一样，一层层剥开。

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

当你的函数是向量到向量的映射时，梯度就变成了雅可比矩阵。JAX 提供了两种计算方式：jacfwd（前向模式）和 jacrev（反向模式）。

说白了，jacfwd 适合「输入小、输出大」的情况，jacrev 适合「输入大、输出小」的情况。怎么理解？

• 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)
print("jacfwd:\n", J_fwd)

# 反向模式
J_rev = jax.jacrev(f)(x)
print("jacrev:\n", J_rev)

两种方式的结果是一样的，但计算效率不同。我个人的经验是：在神经网络中，参数数量通常远大于损失函数的输出维度，所以 jacrev 更常用。但在某些生成模型中，比如 Neural ODE，输出维度可能很大，这时候 jacfwd 就派上用场了。

3.4 vmap：自动向量化

vmap 是 JAX 里我最喜欢的工具之一。它能把一个处理单个样本的函数，自动变成处理批量样本的函数。你想想看，不用手动写 for 循环，也不用操心 batch 维度，多省心。

def single_sample(x):
    return x**2 + 1

# 批量处理
batch_x = jnp.array([[1.0, 2.0], [3.0, 4.0]])
batched_fn = jax.vmap(single_sample)
print(batched_fn(batch_x))

嗯，这里要注意：vmap 默认在第一个维度上映射。如果你的 batch 维度在其他位置，可以用 in_axes 和 out_axes 指定。

3.5 pmap：分布式自动微分

pmap 是 JAX 的分布式并行工具。它把计算分布到多个设备（比如多个 GPU 或 TPU）上，每个设备处理一部分数据。和 vmap 类似，但 pmap 是跨设备的。

# 假设你有 2 个 GPU
def f(x):
    return x * 2

# 在多个设备上并行计算
parallel_f = jax.pmap(f)
data = jnp.array([[1.0, 2.0], [3.0, 4.0]])
result = parallel_f(data)

为什么会用到 pmap？说白了，当你的模型太大，单张 GPU 放不下时，或者数据量太大，单张 GPU 处理太慢时，pmap 就是你的救星。

3.6 自动微分与 vmap/pmap 的结合

JAX 最强大的地方在于，这些变换可以组合使用。比如，你可以先对函数求梯度，然后用 vmap 批量处理，最后用 pmap 分布到多个设备上。

def loss_fn(params, x, y):
    pred = jnp.dot(x, params)
    return jnp.mean((pred - y)**2)

# 对参数求梯度
grad_fn = jax.grad(loss_fn)

# 批量处理
batched_grad = jax.vmap(grad_fn, in_axes=(None, 0, 0))

# 分布式并行
parallel_grad = jax.pmap(batched_grad)

# 使用示例
params = jnp.array([1.0, 2.0])
x_batch = jnp.array([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]])
y_batch = jnp.array([3.0, 7.0, 11.0])

# 注意：pmap 需要数据在多个设备上分片
# 这里只是演示组合方式
grads = batched_grad(params, x_batch, y_batch)

这种组合能力，让 JAX 在深度学习研究中特别受欢迎。你想想看，写一次函数，就能自动获得梯度、批量处理、分布式并行，多省事。

3.7 知识体系总览

下面这张图展示了本章的知识结构，帮你理清各个概念之间的关系：

这张图把本章的核心内容串起来了。从中心的自动微分出发，四个分支分别对应 grad、jacfwd/jacrev、vmap、pmap，最后汇聚到组合使用。你想想看，掌握了这些工具，写深度学习代码就像搭积木一样简单。

---