4. 自动微分（Grad）：grad函数的使用，高阶导数计算，处理复杂函数（如分支、循环）

自动微分，说白了就是让框架帮我们把导数算好。你想想看，以前手推梯度的时候，一个链式法则写满一页纸，还容易漏掉某个中间变量。JAX 的 grad 函数，就是专门干这个的。

我个人习惯，拿到一个新模型，第一件事就是用 grad 验证一下梯度是否正确。这步省了，后面调 bug 会调到怀疑人生。

4.1 grad 函数的基本用法

grad 是 JAX 自动微分的核心入口。它接受一个函数，返回这个函数的梯度函数。看个最简单的例子：

import jax
import jax.numpy as jnp

def f(x):
    return x ** 2

grad_f = jax.grad(f)
print(grad_f(3.0))  # 输出 6.0

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

def g(x, y):
    return x * y + jnp.sin(y)

# 对 x 求导
grad_g_x = jax.grad(g, argnums=0)
# 对 y 求导
grad_g_y = jax.grad(g, argnums=1)
# 同时对 x 和 y 求导
grad_g_xy = jax.grad(g, argnums=(0, 1))

print(grad_g_x(2.0, 3.0))   # 输出 3.0
print(grad_g_y(2.0, 3.0))   # 输出 2.0 + cos(3.0)
小技巧： 如果你需要同时返回函数值和梯度，用 value_and_grad。这比分别调用两次高效得多，我在项目中经常这么用。 
val_grad_fn = jax.value_and_grad(f)
val, grad = val_grad_fn(3.0)
print(f"值: {val}, 梯度: {grad}")

4.2 高阶导数计算

JAX 的 grad 可以嵌套使用，这就意味着你可以轻松计算高阶导数。说白了，就是对梯度函数再求一次梯度。

def f(x):
    return x ** 3

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

# 二阶导数: 6x
grad2_f = jax.grad(grad_f)
print(grad2_f(2.0))  # 12.0

# 三阶导数: 6
grad3_f = jax.grad(grad2_f)
print(grad3_f(2.0))  # 6.0

我记得有一次做物理信息神经网络（PINN），需要计算损失函数对输入的二阶导数。用 JAX 的嵌套 grad，几行代码就搞定了。要是用 PyTorch，还得手动处理 create_graph=True 和 retain_graph，麻烦得很。

注意： 高阶导数计算会显著增加计算图和内存开销。我曾经在训练一个 10 层网络时，试图计算三阶 Hessian 向量积，结果显存直接爆了。实际应用中，二阶导数基本够用，三阶以上很少用到。

4.3 处理复杂函数：分支与循环

很多框架在处理控制流（if-else、for 循环）时，自动微分会出问题。但 JAX 不一样，它通过 jax.lax.cond 和 jax.lax.scan 等函数，让控制流变得可微分。

4.3.1 分支结构

先看一个带 if-else 的函数：

def relu(x):
    if x > 0:
        return x
    else:
        return 0.0

# 直接使用 grad 会报错
# grad_relu = jax.grad(relu)  # 错误！

为什么会这样？因为 Python 的 if-else 是动态的，JAX 在追踪计算图时无法确定走哪条分支。解决办法是用 jax.lax.cond：

def relu_jax(x):
    return jax.lax.cond(
        x > 0,
        lambda x: x,
        lambda x: 0.0,
        x
    )

grad_relu = jax.grad(relu_jax)
print(grad_relu(2.0))   # 输出 1.0
print(grad_relu(-1.0))  # 输出 0.0

嗯，这里要注意：jax.lax.cond 的三个参数分别是条件、真分支函数、假分支函数、输入。两个分支函数的签名必须一致。

4.3.2 循环结构

循环也是类似的问题。Python 的 for 循环在 JAX 中无法直接求导，需要用 jax.lax.scan 或 jax.lax.fori_loop。

看一个累加的例子：

def sum_squares_scan(xs):
    def body_fn(carry, x):
        return carry + x ** 2, None
    total, _ = jax.lax.scan(body_fn, 0.0, xs)
    return total

xs = jnp.array([1.0, 2.0, 3.0])
grad_fn = jax.grad(sum_squares_scan)
print(grad_fn(xs))  # 输出 [2.0, 4.0, 6.0]

scan 的语义是：body_fn(carry, x) -> (new_carry, output)。它会把 xs 中的每个元素依次传入，并更新 carry。最终返回的是最后一个 carry 和所有 output 的堆叠。

核心要点： JAX 的自动微分之所以能处理控制流，是因为它把分支和循环都转换成了纯函数式的操作。每个分支、每次迭代都变成了计算图中的一个节点，梯度自然可以沿着这些节点反向传播。

4.4 实战：一个带分支和循环的复杂函数

把前面学的串起来，看一个稍微复杂点的例子：

def complex_function(x, n):
    # 分支：根据 x 的正负选择不同路径
    def pos_path(x):
        # 循环：累加 x 的幂次
        def body_fn(carry, i):
            return carry + x ** (i + 1), None
        total, _ = jax.lax.scan(body_fn, 0.0, jnp.arange(n))
        return total

    def neg_path(x):
        return jnp.sin(x) * n

    return jax.lax.cond(x > 0, pos_path, neg_path, x)

# 测试
x_val = 2.0
n_val = 3
result = complex_function(x_val, n_val)
print(f"函数值: {result}")  # 2 + 4 + 8 = 14

grad_fn = jax.grad(complex_function, argnums=0)
print(f"梯度: {grad_fn(x_val, n_val)}")  # 1 + 2*2 + 3*4 = 17

这个例子中，complex_function 既有分支（cond），又有循环（scan）。JAX 的自动微分依然能正确计算梯度。我在做强化学习策略网络时，经常需要这种带控制流的函数，JAX 处理起来非常顺手。

4.5 知识体系总览

下面这张图总结了本章的核心内容：

JAX 自动微分知识体系 jax.grad 基本用法 grad(f) → 梯度函数 argnums 指定参数 高阶导数 grad(grad(f)) → 二阶导 嵌套使用，注意内存 复杂函数 分支: lax.cond 循环: lax.scan value_and_grad 多参数求导 Hessian 矩阵 内存优化技巧 fori_loop while_loop 核心思想：一切皆可微分 将控制流转换为纯函数式操作，构建可微计算图

4.6 避坑指南

我曾经在项目里踩过一个坑：用 grad 对一个包含 Python 原生循环的函数求导，结果跑了一整天都没出结果。后来才发现，JAX 的 grad 要求所有操作都是 JAX 可追踪的。原生 Python 循环会打断计算图的构建。

总结几个常见问题：

  • 不要用 Python 原生 if-else：改用 jax.lax.cond 或 jnp.where
  • 不要用 Python 原生 for/while：改用 jax.lax.scan、fori_loop 或 while_loop
  • 注意数据类型：grad 要求输入输出都是浮点数，整数类型会报错
  • 高阶导数慎用：二阶以上时，计算图和内存开销会指数级增长
调试技巧： 如果你不确定某个函数能否被 grad 正确处理，先用 jax.make_jaxpr(f)(x) 查看它的 JAX 表达式。如果输出里出现了 lambda 或 cond 等原始操作，说明 JAX 已经正确追踪了你的函数。

好了，自动微分的核心内容就这些。说白了，记住两件事：一是用 grad 替代手推梯度，二是把控制流换成 JAX 的函数式版本。做到这两点，你就能在 JAX 里自由地求导了。

← 上一章 📖 返回目录 下一章 →