3. 自动微分实战：grad函数、梯度计算、与PyTorch自动微分的对比
自动微分，说白了就是框架帮我们算梯度。
我刚开始做NLP的时候，手写梯度简直是噩梦。一个Transformer模型，反向传播的链式法则能绕晕你。后来用了JAX的grad，嗯，真香。
这一节，我们就来聊聊JAX怎么算梯度，以及它和PyTorch到底有啥不一样。
3.1 grad函数：JAX的梯度计算核心
JAX的grad函数，用法极其简单。你给它一个函数，它返回一个新函数。这个新函数能计算原函数的梯度。
看个例子：
import jax
import jax.numpy as jnp
def f(x):
return x ** 2 + 3 * x + 1
# 计算 f 在 x=2.0 处的梯度
grad_f = jax.grad(f)
print(grad_f(2.0)) # 输出: 7.0
为什么是7.0？因为f'(x) = 2x + 3，代入x=2.0，就是7.0。
我个人习惯，在写复杂模型前，先用这种小函数测试一下梯度对不对。这能省下大量调试时间。
argnums 指定。
3.2 高阶梯度与梯度惩罚
JAX的grad可以嵌套使用。这意味着你可以轻松计算二阶导数、三阶导数……
我在做GAN的梯度惩罚（Gradient Penalty）时，就用到了这个特性。看代码：
def loss_fn(params, x):
return jnp.sum(params * x)
# 一阶梯度
grad_loss = jax.grad(loss_fn)
# 二阶梯度（Hessian向量积）
hessian_loss = jax.grad(jax.grad(loss_fn))
# 实际使用：梯度惩罚项
def gradient_penalty(discriminator, real_data, fake_data):
# 插值
t = jax.random.uniform(jax.random.PRNGKey(0), (real_data.shape[0], 1))
interpolates = t * real_data + (1 - t) * fake_data
# 计算判别器输出对输入的梯度
grad_fn = jax.grad(lambda x: jnp.sum(discriminator(x)))
gradients = jax.vmap(grad_fn)(interpolates)
# 计算梯度惩罚
penalty = jnp.mean((jnp.linalg.norm(gradients, axis=1) - 1) ** 2)
return penalty
我曾经在WGAN-GP的训练中，用这个技巧稳定了训练过程。说实话，如果用手动实现，代码量至少翻三倍。
3.3 与PyTorch自动微分的对比
PyTorch和JAX的自动微分，思路完全不同。我整理了一个表格，方便你对比：
| 特性 | JAX | PyTorch |
|---|---|---|
| 计算方式 | 函数式（纯函数） | 动态计算图（状态式） |
| 梯度存储 | 不存储，每次重新计算 | 存储在 .grad 属性中 |
| 控制流 | 原生Python控制流自动支持 | 需要特殊处理（如 torch.where） |
| 副作用 | 不允许（纯函数要求） | 允许（如原地修改） |
| 高阶梯度 | 原生支持，嵌套grad即可 | 需要 create_graph=True |
| 性能 | JIT编译后极快 | 动态图灵活但略慢 |
你想想看，PyTorch的计算图是动态构建的。每次前向传播，都会生成一个新的图。而JAX呢？它把你的函数编译成一个计算图，然后直接求导。
举个例子，在PyTorch中：
import torch
x = torch.tensor(2.0, requires_grad=True)
y = x ** 2 + 3 * x + 1
y.backward()
print(x.grad) # 输出: tensor(7.)
在JAX中：
import jax.numpy as jnp
from jax import grad
def f(x):
return x ** 2 + 3 * x + 1
grad_f = grad(f)
print(grad_f(2.0)) # 输出: 7.0
表面上看结果一样，但内部机制完全不同。PyTorch需要你手动调用backward()，然后从.grad中取结果。JAX则是直接返回一个函数，你调用它就能得到梯度。
3.4 实战：用grad实现一个简单的线性回归
光说不练假把式。我们来实现一个线性回归的梯度下降：
import jax
import jax.numpy as jnp
from jax import grad, jit
# 生成数据
key = jax.random.PRNGKey(42)
x = jax.random.normal(key, (100, 1))
true_w, true_b = 2.0, -1.0
y = true_w * x + true_b + 0.1 * jax.random.normal(key, (100, 1))
# 定义损失函数
def mse_loss(params, x, y):
w, b = params
pred = w * x + b
return jnp.mean((pred - y) ** 2)
# 计算梯度
def update(params, x, y, lr=0.01):
grads = grad(mse_loss)(params, x, y)
return [p - lr * g for p, g in zip(params, grads)]
# 训练
params = [jnp.zeros((1, 1)), jnp.zeros(())]
for i in range(1000):
params = update(params, x, y)
if i % 200 == 0:
loss = mse_loss(params, x, y)
print(f"Step {i}, Loss: {loss:.4f}")
print(f"Learned w: {params[0].item():.4f}, b: {params[1].item():.4f}")
这段代码里，grad(mse_loss)直接返回了损失函数对params的梯度。然后我们手动更新参数。整个过程没有.backward()，没有.grad，干净利落。
我曾经用这个模式写过完整的BERT微调代码。说实话，一旦习惯了JAX的函数式风格，你会觉得PyTorch的nn.Module有点啰嗦。
3.5 避坑指南
最后，分享几个我踩过的坑：
- 不要修改全局变量： JAX的纯函数要求，如果你在函数内部修改了全局变量，
grad会忽略这些修改。我曾经因为这个bug调试了一整天。 - 注意数据类型： JAX默认使用float32。如果你需要更高精度，记得显式转换。
- grad vs value_and_grad： 如果你既需要损失值又需要梯度，用
jax.value_and_grad。它同时返回两者，避免重复计算。