一、自动微分与梯度计算：三大变换的底层逻辑

大家好，我是你们的部署工程师朋友。今天咱们聊聊JAX里最核心的东西——自动微分与梯度计算。

说实话，我第一次接触JAX时，最震撼的就是它的三大变换：grad、vmap、pmap。这三个东西，说白了就是JAX的「三驾马车」。你想想看，一个框架能把自动微分、向量化、并行化统一成几个函数变换，这设计思路确实漂亮。

1.1 grad：自动微分的「瑞士军刀」

grad是JAX最基础的变换。它做的事情很简单：给定一个函数f，返回一个新的函数，这个新函数计算f的梯度。

我习惯把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 (因为导数 2x+3 = 7)

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

1.2 vmap：自动向量化的「魔法棒」

vmap解决的是批量计算的问题。你写一个处理单个样本的函数，vmap自动把它变成处理批量的函数。

说白了，vmap就是帮你「隐式地写for循环」。但它比for循环快得多，因为它在底层做了向量化优化。

def single_predict(params, x):
    return jnp.dot(params, x)

# 批量预测
batch_predict = jax.vmap(single_predict, in_axes=(None, 0))
# in_axes=(None, 0) 表示params不批量，x按第0维批量

我曾经在部署一个NLP模型时，需要同时处理1000个句子。如果用for循环，推理时间要3秒；换成vmap后，直接降到0.2秒。这个差距，你品，你细品。

1.3 pmap：分布式并行的「加速器」

pmap是JAX的分布式并行利器。它把你的计算自动分配到多个设备（GPU/TPU）上执行。

我记得第一次用pmap时，感觉就像打开了新世界的大门。你只需要在函数上加一个pmap装饰器，JAX就会自动处理数据分片、设备间通信。

# 假设有8个GPU
parallel_grad = jax.pmap(jax.grad(loss_fn), axis_name='batch')
# 数据会自动分片到8个设备上

1.4 自定义梯度：custom_vjp的实战技巧

有时候，自动微分算出来的梯度不是我们想要的。比如某些操作不可微，或者你想手动控制梯度流。这时候就需要custom_vjp。

custom_vjp允许你自定义「前向传播」和「反向传播」的逻辑。我一般用它来处理数值不稳定的操作，或者实现一些特殊的正则化。

from jax import custom_vjp

@custom_vjp
def my_clip(x, low, high):
    return jnp.clip(x, low, high)

def my_clip_fwd(x, low, high):
    return my_clip(x, low, high), (x, low, high)

def my_clip_bwd(res, g):
    x, low, high = res
    # 只在区间内传递梯度
    grad_x = jnp.where((x >= low) & (x <= high), g, 0.0)
    return (grad_x, None, None)

my_clip.defvjp(my_clip_fwd, my_clip_bwd)

嗯，这里要注意：custom_vjp的返回值必须和原函数的参数一一对应。我刚开始写时经常忘记返回None给不需要梯度的参数，结果报错半天。

1.5 高阶梯度计算：从一阶到二阶

JAX支持任意阶的梯度计算。你可以在grad外面再套一层grad，得到二阶导数。

def f(x):
    return x**3

# 一阶导数
df = jax.grad(f)
print(df(2.0))  # 3x^2 = 12

# 二阶导数
d2f = jax.grad(jax.grad(f))
print(d2f(2.0))  # 6x = 12

# 三阶导数
d3f = jax.grad(jax.grad(jax.grad(f)))
print(d3f(2.0))  # 6

高阶梯度在优化算法中很常见。比如牛顿法就需要Hessian矩阵。但要注意，高阶梯度的计算开销是指数级增长的。我一般只在模型调试时用，生产环境很少用。

1.6 Hessian矩阵的JAX实现

Hessian矩阵是二阶偏导数的矩阵。在JAX里，实现Hessian矩阵有两种方式：

1. 直接法：用jax.hessian函数
2. 组合法：用jax.jacfwd和jax.jacrev组合

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

# 方法1：直接使用hessian
hess = jax.hessian(f)
print(hess(jnp.array([1.0, 2.0])))

# 方法2：组合jacfwd和jacrev
hess_combo = jax.jacfwd(jax.grad(f))
print(hess_combo(jnp.array([1.0, 2.0])))

方法	适用场景	计算效率
jax.hessian	小规模参数（<100维）	中等
jacfwd+grad	输出维度远小于输入维度	高
jacrev+grad	输入维度远小于输出维度	高

1.7 三大变换的关系图

下面这张图展示了三大变换如何协同工作：

从图中可以看出，三大变换可以任意组合。比如pmap(vmap(grad(f)))，就是先求导，再向量化，最后分布式并行。这种组合能力，让JAX在模型部署和推理加速方面特别灵活。

好了，这一章的内容就到这里。自动微分和三大变换是JAX的基石，理解透了，后面的部署和加速就水到渠成了。

---