一、JAX中的梯度计算：从grad函数说起

做深度学习，梯度计算是绕不开的核心。JAX在这块做得特别优雅——你只需要调用grad，它就能自动帮你把梯度算出来。我个人觉得，这是JAX最吸引人的特性之一。

1.1 grad函数详解

grad是JAX自动微分的入口。它的用法很简单：

import jax
import jax.numpy as jnp

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

# 计算f在x=2处的梯度
grad_f = jax.grad(f)
print(grad_f(2.0))  # 输出: 7.0

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

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

# 同时对w和b求导
grad_loss = jax.grad(loss, argnums=(0, 1))

我在项目中遇到过一个问题：grad要求函数返回值是标量。如果你的函数返回向量，它会报错。解决办法？要么手动求和，要么用vmap批量处理。

1.2 value_and_grad的使用

很多时候，我们既想要损失值，又想要梯度。分开算两次？太浪费了。value_and_grad就是干这个的：

from jax import value_and_grad

def loss_fn(params):
    return jnp.sum(params ** 2)

# 一次调用，同时拿到值和梯度
loss_val, grads = value_and_grad(loss_fn)(params)
print(f"Loss: {loss_val}, Grads: {grads}")

说实话，我几乎不用单独的grad。因为训练时你总要监控loss的变化，value_and_grad一步到位，省心省力。

1.3 高阶梯度计算

JAX支持高阶梯度，说白了就是「对梯度再求导」。这在一些优化算法（比如Hessian-free优化）里很有用：

def f(x):
    return x ** 3

# 一阶导
df = jax.grad(f)
print(df(2.0))  # 12.0

# 二阶导
d2f = jax.grad(df)
print(d2f(2.0))  # 12.0

# 或者直接用高阶API
from jax import hessian
hess = hessian(f)
print(hess(2.0))  # 12.0

你想想看，为什么需要高阶梯度？我举个例子：在训练GAN时，判别器的梯度范数正则化就需要用到二阶信息。不过说实话，高阶梯度计算开销很大，非必要不用。

1.4 梯度计算的注意事项

这部分是我踩坑最多的地方，分享几个关键点：

1.4.1 随机性处理

JAX的随机数生成器（PRNG）是显式传递的。如果你在函数里用了随机操作，grad会报错：

# 错误示例
def noisy_loss(params, key):
    noise = jax.random.normal(key, shape=params.shape)
    return jnp.sum((params + noise) ** 2)

# 这样不行！grad不知道如何处理key
# grad_loss = jax.grad(noisy_loss)

# 正确做法：把随机种子作为参数，但用argnums排除
grad_loss = jax.grad(noisy_loss, argnums=0)

1.4.2 不可微操作

有些操作是不可微的，比如argmin、sort等。JAX会返回NaN或0梯度：

def bad_fn(x):
    return jnp.argmin(x)  # 不可微！

grad_bad = jax.grad(bad_fn)
print(grad_bad(jnp.array([1.0, 2.0, 3.0])))  # 可能得到0或NaN

我曾经在实现一个排序网络时踩过这个坑。解决办法？要么用软排序替代，要么用straight-through estimator。

1.4.3 数值精度

JAX默认使用float32。如果你需要更高精度：

# 设置全局精度
jax.config.update("jax_default_dtype", jnp.float64)

# 或者局部使用
x = jnp.array([1.0, 2.0], dtype=jnp.float64)

知识体系图

这张图把JAX梯度计算的核心脉络梳理清楚了。从grad函数出发，到value_and_grad的实用技巧，再到高阶梯度的进阶用法，最后是那些容易踩坑的注意事项。我个人建议：先把value_and_grad用熟，再考虑高阶梯度。毕竟大多数场景下，一阶梯度就够用了。