企业官网建设流程全解析

JAX 自动微分 Cookbook：从 grad 到 JVP/VJP 的前向与反向模式全解

【免费下载链接】jaxComposable transformations of Python+NumPy programs: differentiate, vectorize, JIT to GPU/TPU, and more项目地址: https://gitcode.com/GitHub_Trending/ja/jax

本篇技术指南以 JAX 官方《The Autodiff Cookbook》为骨架，系统讲解 JAX 自动微分系统的核心用法与底层原理：从最基础的grad、value_and_grad，到高阶导数（Hessian-vector product、jacfwd/jacrev与稠密 Hessian），再到支撑这一切的两个奠基性原语——前向模式的 Jacobian-vector product（JVP）与反向模式的 vector-Jacobian product（VJP），最后深入复数域的微分规则。读完本文，你将掌握在 JAX 中构建任意阶、任意模式微分程序的完整工具箱，并理解每条 API 在仓库源码中的真实实现路径。

本文所有代码与结论均以当前仓库为准，关键 API 的实现可直接在 jax/_src/api.py 中核对；配套的可执行 Notebook 版本见 docs/notebooks/autodiff_cookbook.ipynb。

准备工作

本文示例使用如下导入，注意当前仓库使用类型化随机密钥（typed keys），因此以random.key(0)而非旧式random.PRNGKey(0)构造密钥：

import jax.numpy as jnp from jax import grad, jit, vmap from jax import random key = random.key(0)

从grad开始：梯度计算基础

grad的基本用法

可以用grad对任意函数求导：

grad_tanh = grad(jnp.tanh) print(grad_tanh(2.0))

grad的输入输出都是函数：如果 Python 函数f计算数学函数 $f$，那么grad(f)就是一个计算 $\nabla f$ 的 Python 函数，即grad(f)(x)的值为 $\nabla f(x)$。

由于grad作用于函数，可以把它反复套在自己的输出上，求任意阶导数：

print(grad(grad(jnp.tanh))(2.0)) print(grad(grad(grad(jnp.tanh)))(2.0))

在仓库源码中，grad是value_and_grad的薄封装，见 jax/_src/api.py：grad先调用value_and_grad构造同时返回函数值与梯度的包装函数，再丢弃函数值、只返回梯度（has_aux=True时则返回(gradient, aux)二元组）。

实战示例：线性逻辑回归的梯度

下面用一个完整的线性逻辑回归模型演示梯度计算。首先定义模型与损失：

def sigmoid(x): return 0.5 * (jnp.tanh(x / 2) + 1) # 输出标签为 True 的概率 def predict(W, b, inputs): return sigmoid(jnp.dot(inputs, W) + b) # 构造一个玩具数据集 inputs = jnp.array([[0.52, 1.12, 0.77], [0.88, -1.08, 0.15], [0.52, 0.06, -1.30], [0.74, -2.49, 1.39]]) targets = jnp.array([True, True, False, True]) # 训练损失为训练样本的负对数似然 def loss(W, b): preds = predict(W, b, inputs) label_probs = preds * targets + (1 - preds) * (1 - targets) return -jnp.sum(jnp.log(label_probs)) # 初始化随机模型系数 key, W_key, b_key = random.split(key, 3) W = random.normal(W_key, (3,)) b = random.normal(b_key, ())

argnums：指定对哪些位置参数求导

grad的argnums参数用来指定对函数的位置参数（positional arguments）中的哪一个求导：

# 对第一个位置参数求导： W_grad = grad(loss, argnums=0)(W, b) print('W_grad', W_grad) # 因为 argnums=0 是默认值，以下写法等价： W_grad = grad(loss)(W, b) print('W_grad', W_grad) # 也可以选择其他参数，并省略关键字写法： b_grad = grad(loss, 1)(W, b) print('b_grad', b_grad) # 支持元组形式，同时求多个参数的梯度： W_grad, b_grad = grad(loss, (0, 1))(W, b) print('W_grad', W_grad) print('b_grad', b_grad)

这个gradAPI 与 Spivak 的经典教材Calculus on Manifolds（1965）以及 Sussman 与 Wisdom 的Structure and Interpretation of Classical Mechanics（2015）、Functional Differential Geometry（2013）中的记号一一对应。本质上，argnums使grad(f, i)成为计算偏导数 $\partial_i f$ 的函数。

从源码看，argnums可以是一个整数或整数序列，jax/_src/api.py 中通过argnums_partial2将目标参数从args/kwargs中提取出来作为"动态参数"参与微分，其余参数被当作常量冻结；如果传入的argnums超出实际位置参数个数，会直接抛出TypeError。

此外，grad还支持以下参数（均定义于 jax/_src/api.py）：

• has_aux：若fun返回(主输出, 辅助数据)二元组，设为True后grad返回(梯度, 辅助数据)；
• holomorphic：承诺被微分的函数是复解析（全纯）函数，要求输入输出均为复数 dtype；
• allow_int：允许对整数输入求导（此时梯度为平凡的float0向量空间 dtype），默认False。

注意grad只接受输出为标量（包括 shape 为()的数组，但不包括 shape(1,)等）的函数；输出非标量时会得到带提示的TypeError，建议改用jax.jacobian（见 jax/_src/api.py 中的_check_scalar）。

对嵌套 list、tuple 和 dict 求导

对标准 Python 容器求导"开箱即用"，tuple、list、dict 以及任意嵌套都可以直接使用：

def loss2(params_dict): preds = predict(params_dict['W'], params_dict['b'], inputs) label_probs = preds * targets + (1 - preds) * (1 - targets) return -jnp.sum(jnp.log(label_probs)) print(grad(loss2)({'W': W, 'b': b}))

这是 JAX pytree 机制的功劳：任意容器被展平为叶子数组的树，梯度保持相同的树结构返回。你还可以注册自定义容器类型，使其不仅适用于grad，也适用于所有 JAX 变换（jit、vmap等）。

用value_and_grad同时求函数值与梯度

value_and_grad可以一次调用同时高效地返回函数值与梯度：

from jax import value_and_grad loss_value, Wb_grad = value_and_grad(loss, (0, 1))(W, b) print('loss value', loss_value) print('loss value', loss(W, b))

从源码看，value_and_grad是比grad更底层的实现：它先对部分应用后的函数调用vjp得到(ans, vjp_py)，再把单位余切向量（lax_internal._one_vjp(ans)）送入 VJP 得到梯度，见 jax/_src/api.py。这正是"grad是vjp的一个特例"这一事实的实现体现。

用数值差分校验梯度

导数的一个好处是可以用有限差分直接验证。对上述模型：

# 有限差分步长 eps = 1e-4 # 对标量 b_grad 做标量有限差分 b_grad_numerical = (loss(W, b + eps / 2.) - loss(W, b - eps / 2.)) / eps print('b_grad_numerical', b_grad_numerical) print('b_grad_autodiff', grad(loss, 1)(W, b)) # 对 W_grad 沿随机方向做有限差分 key, subkey = random.split(key) vec = random.normal(subkey, W.shape) unitvec = vec / jnp.sqrt(jnp.vdot(vec, vec)) W_grad_numerical = (loss(W + eps / 2. * unitvec, b) - loss(W - eps / 2. * unitvec, b)) / eps print('W_dirderiv_numerical', W_grad_numerical) print('W_dirderiv_autodiff', jnp.vdot(grad(loss)(W, b), unitvec))

JAX 还提供了封装好的便捷函数，可以校验到任意阶导数：

from jax.test_util import check_grads check_grads(loss, (W, b), order=2) # 校验到 2 阶导数

check_grads的实现位于 jax/_src/public_test_util.py：它支持modes=("fwd", "rev", "lin")三种校验模式，默认同时校验前向（JVP）与反向（VJP），递归地对order阶以内的导数逐一做 JVP/VJP 与数值差分的一致性检查；梯度只在一个随机选择的方向上校验，从而避免大输入/输出空间下有限差分代价过高。

用grad-of-grad构造 Hessian-vector product

利用高阶grad可以构造 Hessian-vector product（HVP）函数。HVP 在截断牛顿共轭梯度算法（truncated Newton conjugate-gradient）中用于最小化光滑凸函数，也常用于研究神经网络训练目标的曲率。

对于具有连续二阶导数的标量函数 $f : \mathbb{R}^n \to \mathbb{R}$（Hessian 矩阵对称），Hessian-vector product 函数需要能够对任意 $v \in \mathbb{R}^n$ 计算

$\qquad v \mapsto \partial^2 f(x) \cdot v$

关键在于不要物化完整的 Hessian 矩阵：当 $n$ 达到数百万甚至数十亿（神经网络场景）时，存储完整 Hessian 是不可能的。

幸运的是，利用恒等式

$\qquad \partial^2 f (x) v = \partial [x \mapsto \partial f(x) \cdot v] = \partial g(x)$

其中 $g(x) = \partial f(x) \cdot v$ 是把 $f$ 在 $x$ 处的梯度与向量 $v$ 做点积得到的新标量函数。注意我们全程只在对"向量值参数上的标量函数"求导，这正是grad最擅长的场景。JAX 代码如下：

def hvp(f, x, v): return grad(lambda x: jnp.vdot(grad(f)(x), v))(x)

这个例子也说明 JAX 对词法闭包（lexical closure）的处理非常稳健，任意嵌套都不会混淆。

用jacfwd与jacrev计算 Jacobian 与 Hessian

计算稠密 Jacobian

jacfwd和jacrev可以计算完整的 Jacobian 矩阵：

from jax import jacfwd, jacrev # 隔离出从权重矩阵到预测的函数 f = lambda W: predict(W, b, inputs) J = jacfwd(f)(W) print("jacfwd result, with shape", J.shape) print(J) J = jacrev(f)(W) print("jacrev result, with shape", J.shape) print(J)

两个函数计算相同的数值结果（差异只在机器数值精度范围内），但实现方式不同：

• jacfwd使用前向模式自动微分，对"高瘦"型 Jacobian（输出数多于输入数）更高效；
• jacrev使用反向模式，对"宽扁"型 Jacobian（输入数多于输出数）更高效；
• 接近方阵时，jacfwd通常略占优势。

jacfwd/jacrev同样支持容器类型：

def predict_dict(params, inputs): return predict(params['W'], params['b'], inputs) J_dict = jacrev(predict_dict)({'W': W, 'b': b}, inputs) for k, v in J_dict.items(): print("Jacobian from {} to logits is".format(k)) print(v)

用jacfwd(jacrev(f))计算稠密 Hessian

组合两个函数即可计算稠密 Hessian 矩阵：

def hessian(f): return jacfwd(jacrev(f)) H = hessian(f)(W) print("hessian, with shape", H.shape) print(H)

形状的直觉：对 $f : \mathbb{R}^n \to \mathbb{R}^m$，在点 $x$ 处依次有

• $f(x) \in \mathbb{R}^m$，$f$ 在 $x$ 处的值；
• $\partial f(x) \in \mathbb{R}^{m \times n}$，$x$ 处的 Jacobian 矩阵；
• $\partial^2 f(x) \in \mathbb{R}^{m \times n \times n}$，$x$ 处的 Hessian；

以此类推。

实现hessian时理论上可以用jacfwd(jacrev(f))、jacrev(jacfwd(f))或任意组合，但前向套反向（forward-over-reverse）通常最高效：内层 Jacobian 常常是在对"宽扁"函数求导（如损失函数 $f : \mathbb{R}^n \to \mathbb{R}$），反向模式占优；而外层 Jacobian 是在对 $\nabla f : \mathbb{R}^n \to \mathbb{R}^n$ 求导，其 Jacobian 是方阵，此时前向模式胜出。仓库中jax.hessian正是以此为默认实现，见 jax/_src/api.py。

基石之一：Jacobian-vector product（JVP，前向模式）

JAX 同时实现了高效且通用的前向与反向自动微分。grad建立在反向模式之上，但要理解两种模式的差异与各自适用场景，需要一点数学背景。

JVP 的数学定义

给定函数 $f : \mathbb{R}^n \to \mathbb{R}^m$，$f$ 在输入点 $x$ 处的 Jacobian $\partial f(x)$ 通常被视为 $\mathbb{R}^{m \times n}$ 矩阵，但也可以看作一个线性映射——把 $f$ 定义域在 $x$ 处的切空间（即另一份 $\mathbb{R}^n$）映射到 $f$ 值域在 $f(x)$ 处的切空间（一份 $\mathbb{R}^m$）：

$\qquad \partial f(x) : \mathbb{R}^n \to \mathbb{R}^m$

这个映射称为 $f$ 在 $x$ 处的 pushforward map（推前映射），Jacobian 矩阵就是该线性映射在标准基下的矩阵表示。

若不对具体输入点 $x$ 做承诺，可以把 $\partial f$ 视为"先取输入点、再返回该点的 Jacobian 线性映射"的函数。将输入点 $x \in \mathbb{R}^n$ 与切向量 $v \in \mathbb{R}^n$ 配对，得到输出切向量（在 $\mathbb{R}^m$ 中），这个从 $(x, v)$ 到输出切向量的映射就是Jacobian-vector product，记作

$\qquad (x, v) \mapsto \partial f(x) v$

JVP 的 JAX 代码

JAX 的jvp函数正是这一变换的建模：给定计算 $f$ 的 Python 函数，jvp给出计算 $(x, v) \mapsto (f(x), \partial f(x) v)$ 的 Python 函数。

from jax import jvp # 隔离出从权重矩阵到预测的函数 f = lambda W: predict(W, b, inputs) key, subkey = random.split(key) v = random.normal(subkey, W.shape) # 将向量 v 沿 f 在 W 处推前 y, u = jvp(f, (W,), (v,))

用类 Haskell 的类型签名可以写作：

jvp :: (a -> b) -> a -> T a -> (b, T b)

其中T a表示a的切空间类型。即jvp接收一个a -> b函数、一个a类型的值、一个T a类型的切向量，返回由b类型值与T b类型输出切向量组成的二元组。

jvp变换后的函数求值方式与原函数类似，只是每个 primal 值旁边都伴随一个T a类型的切向量。对于原函数执行的每个原始数值运算，jvp变换后的函数都会执行该原语的"JVP 规则"：既在 primal 上求值原语，又在这些 primal 值上应用该原语的 JVP。

这种求值策略直接决定了计算复杂度：由于 JVP 是边算边走，无需为后续步骤保存任何中间量，内存开销与计算的深度无关；同时 JVP 变换后函数的 FLOP 开销约为原函数的 3 倍（一份用于求值原函数如sin(x)，一份用于线性化如cos(x)，一份用于把线性化结果作用到向量上如cos_x * v）。换言之，固定 primal 点 $x$ 后，计算 $v \mapsto \partial f(x) \cdot v$ 的边际成本与求值 $f$ 相当。

为什么机器学习中不常用前向模式？

既然内存复杂度如此诱人，为什么前向模式在机器学习中很少见？

思考如何用 JVP 构建完整 Jacobian：把 JVP 作用到 one-hot 切向量上，就会揭示 Jacobian 中与该非零分量对应的一列。于是可以一列一列地构建完整 Jacobian，每列的成本约等于一次函数求值。这对"高瘦" Jacobian 高效，对"宽扁" Jacobian 则低效。

而在基于梯度的机器学习优化中，目标通常是从 $\mathbb{R}^n$ 的参数空间到标量损失 $\mathbb{R}$，其 Jacobian 是极宽的矩阵 $\partial f(x) \in \mathbb{R}^{1 \times n}$（通常被等同于梯度向量 $\nabla f(x) \in \mathbb{R}^n$）。一列一列地构建该矩阵、每次调用消耗与原始函数相近的 FLOP，显然低效。尤其对于神经网络训练损失，$n$ 可达数百万甚至数十亿，这种方式根本无法扩展。要做得更好，就需要反向模式。

基石之二：Vector-Jacobian product（VJP，反向模式）

前向模式返回计算 Jacobian-vector product 的函数，可以一列一列构建 Jacobian；反向模式则返回计算 vector-Jacobian product（等价于 Jacobian 转置乘向量）的函数，可以一行一行构建 Jacobian。

VJP 的数学定义

再次考虑 $f : \mathbb{R}^n \to \mathbb{R}^m$，VJP 的记号非常简洁：

$\qquad (x, v) \mapsto v \partial f(x)$

其中 $v$ 是 $f$ 在 $x$ 处余切空间（cotangent space，与另一份 $\mathbb{R}^m$ 同构）中的元素。严格地说，应把 $v$ 视为线性映射 $v : \mathbb{R}^m \to \mathbb{R}$，$v \partial f(x)$ 表示复合 $v \circ \partial f(x)$；但常见情形下可以把 $v$ 等同于 $\mathbb{R}^m$ 中的向量，二者几乎可互换使用（如同在"列向量"与"行向量"之间来回切换而不加说明）。

利用这一等同关系，VJP 的线性部分也可以看成 JVP 线性部分的转置（伴随）：

$\qquad (x, v) \mapsto \partial f(x)^\mathsf{T} v$

对给定点 $x$，签名写作

$\qquad \partial f(x)^\mathsf{T} : \mathbb{R}^m \to \mathbb{R}^n$

余切空间上的这一对应映射常被称为 $f$ 在 $x$ 处的 pullback（拉回映射）。关键点在于：它从"看起来像 $f$ 的输出"映射到"看起来像 $f$ 的输入"，正如对转置线性函数的预期。

VJP 的 JAX 代码

JAX 函数vjp接收计算 $f$ 的 Python 函数，返回计算 VJP $(x, v) \mapsto (f(x), v^\mathsf{T} \partial f(x))$ 的 Python 函数：

from jax import vjp # 隔离出从权重矩阵到预测的函数 f = lambda W: predict(W, b, inputs) y, vjp_fun = vjp(f, W) key, subkey = random.split(key) u = random.normal(subkey, y.shape) # 将余向量 u 沿 f 在 W 处拉回 v = vjp_fun(u)

类 Haskell 类型签名：

vjp :: (a -> b) -> a -> (b, CT b -> CT a)

其中CT a表示a的余切空间类型。即vjp接收a -> b函数与a类型点，返回由b类型值与CT b -> CT a线性映射组成的二元组。

这正是它的价值：可以一行一行构建 Jacobian，且计算 $(x, v) \mapsto (f(x), v^\mathsf{T} \partial f(x))$ 的 FLOP 开销仅为求值 $f$ 的约 3 倍。特别地，对 $f : \mathbb{R}^n \to \mathbb{R}$ 求梯度只需一次调用——这就是grad对基于梯度的优化如此高效的原因，哪怕目标函数是拥有数百万乃至数十亿参数的神经网络训练损失。

代价是：虽然 FLOP 友好，内存随计算的深度增长；且实现传统上比前向模式更复杂。在仓库源码中，vjp的完整实现见 jax/_src/api.py，其返回值vjpfun接收与primals_out形状相同的余切向量，返回与primals数量、形状相同的余切向量元组；它还支持has_aux、saveable_args、in_nzs等进阶参数。

用 VJP 实现向量值梯度

如果你需要向量值梯度（类似tf.gradients的语义）：

from jax import vjp def vgrad(f, x): y, vjp_fn = vjp(f, x) return vjp_fn(jnp.ones(y.shape))[0] print(vgrad(lambda x: 3*x**2, jnp.ones((2, 2))))

前向-反向混合的 Hessian-vector product

之前我们用纯反向模式实现了 HVP（假设二阶导连续）：

def hvp(f, x, v): return grad(lambda x: jnp.vdot(grad(f)(x), v))(x)

它已经高效，但把前向模式与反向模式结合起来还能进一步节省内存。数学上，对 $f : \mathbb{R}^n \to \mathbb{R}$、线性化点 $x$ 与向量 $v$，我们想要

$(x, v) \mapsto \partial^2 f(x) v$

考虑辅助函数 $g : \mathbb{R}^n \to \mathbb{R}^n$，即 $f$ 的导数（梯度）$g(x) = \partial f(x)$。我们只需要它的 JVP，因为

$(x, v) \mapsto \partial g(x) v = \partial^2 f(x) v$

几乎可以直接翻译成代码：

from jax import jvp, grad # forward-over-reverse（前向套反向） def hvp(f, primals, tangents): return jvp(grad(f), primals, tangents)[1]

更好的地方在于：由于不需要直接调用jnp.dot，这个hvp对任意形状的数组、任意容器类型（如存为嵌套 list/dict/tuple 的向量）都适用，甚至不依赖jax.numpy。

使用示例：

def f(X): return jnp.sum(jnp.tanh(X)**2) key, subkey1, subkey2 = random.split(key, 3) X = random.normal(subkey1, (30, 40)) V = random.normal(subkey2, (30, 40)) ans1 = hvp(f, (X,), (V,)) ans2 = jnp.tensordot(hessian(f)(X), V, 2) print(jnp.allclose(ans1, ans2, 1e-4, 1e-4))

另一种写法是反向套前向（reverse-over-forward）：

# reverse-over-forward def hvp_revfwd(f, primals, tangents): g = lambda primals: jvp(f, primals, tangents)[1] return grad(g)(primals)

它稍逊一筹，因为前向模式的开销小于反向模式，而外层的微分算子需要微分的计算规模比内层更大，所以把前向模式放在外层效果最佳：

# reverse-over-reverse，仅适用于单参数 def hvp_revrev(f, primals, tangents): x, = primals v, = tangents return grad(lambda x: jnp.vdot(grad(f)(x), v))(x) print("Forward over reverse") %timeit -n10 -r3 hvp(f, (X,), (V,)) print("Reverse over forward") %timeit -n10 -r3 hvp_revfwd(f, (X,), (V,)) print("Reverse over reverse") %timeit -n10 -r3 hvp_revrev(f, (X,), (V,)) print("Naive full Hessian materialization") %timeit -n10 -r3 jnp.tensordot(hessian(f)(X), V, 2)

组合 VJP、JVP 与vmap

Jacobian-Matrix 与 Matrix-Jacobian 乘积

有了jvp与vjp这两个"单次推前/拉回一个向量"的变换，就可以用vmap一次性推前/拉回整个基。先看 Matrix-Jacobian 乘积（拉回多个余向量）：

# 隔离出从权重矩阵到预测的函数 f = lambda W: predict(W, b, inputs) # 沿 f 在 W 处拉回余向量 m_i，对矩阵 M 的所有行 i。 # 先用列表推导式在矩阵 M 的行上循环。 def loop_mjp(f, x, M): y, vjp_fun = vjp(f, x) return jnp.vstack([jnp.asarray(vjp_fun(mi)) for mi in M]) # 再用 vmap 构建一次快速的矩阵-矩阵乘法， # 而非外层循环的向量-矩阵乘法。 def vmap_mjp(f, x, M): y, vjp_fun = vjp(f, x) outs, = vmap(vjp_fun)(M) return outs key = random.key(0) num_covecs = 128 U = random.normal(key, (num_covecs,) + y.shape) loop_vs = loop_mjp(f, W, M=U) print('Non-vmapped Matrix-Jacobian product') %timeit -n10 -r3 loop_mjp(f, W, M=U) print('\nVmapped Matrix-Jacobian product') vmap_vs = vmap_mjp(f, W, M=U) %timeit -n10 -r3 vmap_mjp(f, W, M=U) assert jnp.allclose(loop_vs, vmap_vs), 'Vmap and non-vmapped Matrix-Jacobian Products should be identical'

再看 Jacobian-Matrix 乘积（推前多个切向量）：

def loop_jmp(f, W, M): # jvp 立即返回 primal 与 tangent 组成的元组， # 因此我们在列表推导式中计算并选取 tangent 部分 return jnp.vstack([jvp(f, (W,), (mi,))[1] for mi in M]) def vmap_jmp(f, W, M): _jvp = lambda s: jvp(f, (W,), (s,))[1] return vmap(_jvp)(M) num_vecs = 128 S = random.normal(key, (num_vecs,) + W.shape) loop_vs = loop_jmp(f, W, M=S) print('Non-vmapped Jacobian-Matrix product') %timeit -n10 -r3 loop_jmp(f, W, M=S) vmap_vs = vmap_jmp(f, W, M=S) print('\nVmapped Jacobian-Matrix product') %timeit -n10 -r3 vmap_jmp(f, W, M=S) assert jnp.allclose(loop_vs, vmap_vs), 'Vmap and non-vmapped Jacobian-Matrix products should be identical'

jacfwd与jacrev的实现

有了上述快速的 Jacobian-Matrix 与 Matrix-Jacobian 乘积，jacfwd/jacrev的实现思路就非常直接了：用同样的技术一次性推前/拉回一整个标准基（与单位矩阵同构）。先看反向模式版：

from jax import jacrev as builtin_jacrev def our_jacrev(f): def jacfun(x): y, vjp_fun = vjp(f, x) # 用 vmap 做 Matrix-Jacobian 乘积。 # 这里矩阵是欧几里得基，因此一次得到 Jacobian 的所有元素。 J, = vmap(vjp_fun, in_axes=0)(jnp.eye(len(y))) return J return jacfun assert jnp.allclose(builtin_jacrev(f)(W), our_jacrev(f)(W)), 'Incorrect reverse-mode Jacobian results!'

再看前向模式版：

from jax import jacfwd as builtin_jacfwd def our_jacfwd(f): def jacfun(x): _jvp = lambda s: jvp(f, (x,), (s,))[1] Jt = vmap(_jvp, in_axes=1)(jnp.eye(len(x))) return jnp.transpose(Jt) return jacfun assert jnp.allclose(builtin_jacfwd(f)(W), our_jacfwd(f)(W)), 'Incorrect forward-mode Jacobian results!'

这与仓库源码的实现思路完全一致：jacfwd本质上是vmap(_jvp)(_std_basis(...))加上_jacfwd_unravel的重组（见 jax/_src/api.py），jacrev则是vjp之后用vmap(pullback)(_std_basis(y))一次拉回整个标准基（见 jax/_src/api.py）。值得一提的还有：jax.hessian正是jacfwd(jacrev(fun))的别名（jax/_src/api.py），其默认采用 forward-over-reverse 组合。

一个有趣的历史背景是：早期的 Autograd 库做不到这一点——它实现反向模式jacobian时必须用外层map循环一次只拉回一个向量。一次一个向量地穿过计算图，远不如用vmap把所有向量批量合并起来高效。

Autograd 做不到的另一件事是jit。无论被微分函数中使用多少 Python 动态特性，JAX 都可以对计算的线性部分使用jit。例如：

def f(x): try: if x < 3: return 2 * x ** 3 else: raise ValueError except ValueError: return jnp.pi * x y, f_vjp = vjp(f, 4.) print(jit(f_vjp)(1.))

复数与微分

JAX 对复数和微分支持良好。为同时支持全纯（holomorphic）与非全纯微分，最好以 JVP 和 VJP 的视角来思考。

考虑复到复函数 $f: \mathbb{C} \to \mathbb{C}$，并把它等同于相应的 $g: \mathbb{R}^2 \to \mathbb{R}^2$：

def f(z): x, y = jnp.real(z), jnp.imag(z) return u(x, y) + v(x, y) * 1j def g(x, y): return (u(x, y), v(x, y))

即分解 $f(z) = u(x, y) + v(x, y) i$（其中 $z = x + y i$），把 $\mathbb{C}$ 等同于 $\mathbb{R}^2$ 得到 $g$。

由于 $g$ 只涉及实数输入输出，我们已经知道如何为它写 Jacobian-vector product：给定切向量 $(c, d) \in \mathbb{R}^2$，即

$\begin{bmatrix} \partial_0 u(x, y) & \partial_1 u(x, y) \ \partial_0 v(x, y) & \partial_1 v(x, y) \end{bmatrix} \begin{bmatrix} c \ d \end{bmatrix}$

要对原函数 $f$ 作用于切向量 $c + di \in \mathbb{C}$ 得到 JVP，只需沿用同一套定义并把结果等同为另一个复数：

$\partial f(x + y i)(c + d i) = \begin{bmatrix} 1 & i \end{bmatrix} \begin{bmatrix} \partial_0 u(x, y) & \partial_1 u(x, y) \ \partial_0 v(x, y) & \partial_1 v(x, y) \end{bmatrix} \begin{bmatrix} c \ d \end{bmatrix}$

这就是 $\mathbb{C} \to \mathbb{C}$ 函数 JVP 的定义！注意 $f$ 是否全纯无关紧要：JVP 是无歧义的。

下面做一次验证：

def check(seed): key = random.key(seed) # 为 u 和 v 生成随机系数 key, subkey = random.split(key) a, b, c, d = random.uniform(subkey, (4,)) def fun(z): x, y = jnp.real(z), jnp.imag(z) return u(x, y) + v(x, y) * 1j def u(x, y): return a * x + b * y def v(x, y): return c * x + d * y # primal 点 key, subkey = random.split(key) x, y = random.uniform(subkey, (2,)) z = x + y * 1j # 切向量 key, subkey = random.split(key) c, d = random.uniform(subkey, (2,)) z_dot = c + d * 1j # 检查 jvp _, ans = jvp(fun, (z,), (z_dot,)) expected = (grad(u, 0)(x, y) * c + grad(u, 1)(x, y) * d + grad(v, 0)(x, y) * c * 1j+ grad(v, 1)(x, y) * d * 1j) print(jnp.allclose(ans, expected))

check(0) check(1) check(2)

那么 VJP 呢？做法类似：对余切向量 $c + di \in \mathbb{C}$，把 $f$ 的 VJP 定义为

$(c + di)^* ; \partial f(x + y i) = \begin{bmatrix} c & -d \end{bmatrix} \begin{bmatrix} \partial_0 u(x, y) & \partial_1 u(x, y) \ \partial_0 v(x, y) & \partial_1 v(x, y) \end{bmatrix} \begin{bmatrix} 1 \ -i \end{bmatrix}$

为什么有负号？只是为了处理复共轭，以及我们是在与余向量（covector）打交道。

VJP 规则的验证：

def check(seed): key = random.key(seed) # 为 u 和 v 生成随机系数 key, subkey = random.split(key) a, b, c, d = random.uniform(subkey, (4,)) def fun(z): x, y = jnp.real(z), jnp.imag(z) return u(x, y) + v(x, y) * 1j def u(x, y): return a * x + b * y def v(x, y): return c * x + d * y # primal 点 key, subkey = random.split(key) x, y = random.uniform(subkey, (2,)) z = x + y * 1j # 余切向量 key, subkey = random.split(key) c, d = random.uniform(subkey, (2,)) z_bar = jnp.array(c + d * 1j) # 用于控制 dtype # 检查 vjp _, fun_vjp = vjp(fun, z) ans, = fun_vjp(z_bar) expected = (grad(u, 0)(x, y) * c + grad(v, 0)(x, y) * (-d) + grad(u, 1)(x, y) * c * (-1j) + grad(v, 1)(x, y) * (-d) * (-1j)) assert jnp.allclose(ans, expected, atol=1e-5, rtol=1e-5)

check(0) check(1) check(2)

grad、jacfwd、jacrev的复数行为

回忆对 $\mathbb{R} \to \mathbb{R}$ 函数，grad(f)(x)被定义为vjp(f, x)1——把 VJP 作用到值1.0上即可揭示梯度（即 Jacobian，即导数）。对 $\mathbb{C} \to \mathbb{R}$ 函数可以做同样的事：仍然使用1.0作为余切向量，得到的复数结果概括了完整 Jacobian：

def f(z): x, y = jnp.real(z), jnp.imag(z) return x**2 + y**2 z = 3. + 4j grad(f)(z)

对一般的 $\mathbb{C} \to \mathbb{C}$ 函数，Jacobian 有 4 个实数自由度（如上述 2×2 Jacobian 矩阵），无法全部塞进一个复数。但对全纯函数可以：全纯函数正是导数能用单个复数表示的 $\mathbb{C} \to \mathbb{C}$ 函数（Cauchy-Riemann 方程保证上述 2×2 Jacobian 具有复平面中"缩放-旋转"矩阵的特殊形式，即单个复数乘法的作用）。而这个复数可以用一次vjp（余向量取1.0）揭示出来。

由于这一技巧只对全纯函数成立，使用前需要向 JAX 承诺函数是全纯的；否则 JAX 会在对复数输出函数使用grad时报错：

def f(z): return jnp.sin(z) z = 3. + 4j grad(f, holomorphic=True)(z)

holomorphic=True承诺的全部作用就是关闭"输出为复数时"的错误。也可以对并非全纯的函数写holomorphic=True，但得到的答案并不代表完整 Jacobian，而是"丢弃输出虚部后的函数"的 Jacobian：

def f(z): return jnp.conjugate(z) z = 3. + 4j grad(f, holomorphic=True)(z) # f 实际上并不是全纯的！

由此grad在复数域有几个实用的推论：

1. 可以对全纯的 $\mathbb{C} \to \mathbb{C}$ 函数使用grad。
2. 可以用grad优化 $f : \mathbb{C} \to \mathbb{R}$ 函数（如以复数参数 $x$ 的实值损失函数），方法是沿grad(f)(x)的共轭方向步进。
3. 如果一个 $\mathbb{R} \to \mathbb{R}$ 函数内部恰好用了某些复数运算（其中一些必然非全纯，例如卷积中使用的 FFT），grad依然可用，且结果与纯实数实现给出的结果一致。

无论如何，JVP 与 VJP 总是无歧义的；若要计算非全纯 $\mathbb{C} \to \mathbb{C}$ 函数的完整 Jacobian 矩阵，用 JVP 或 VJP 即可。

可以预期复数在 JAX 中处处可用。下面是对复矩阵做 Cholesky 分解并求导的示例：

A = jnp.array([[5., 2.+3j, 5j], [2.-3j, 7., 1.+7j], [-5j, 1.-7j, 12.]]) def f(X): L = jnp.linalg.cholesky(X) return jnp.sum((L - jnp.sin(L))**2) grad(f, holomorphic=True)(A)

进阶自动微分展望

本文从易到难、循序渐进地演示了 JAX 中自动微分的各种应用。若想进一步深入，仓库内的 docs/advanced_autodiff.md 提供了"高级自动微分"的完整进阶指南。

自动微分的世界里还有大量其他技巧与功能，本文未覆盖、未来可能在"Advanced Autodiff Cookbook"中展开的主题包括：

• Gauss-Newton Vector Products：只线性化一次的 Gauss-Newton 向量积；
• Custom VJPs and JVPs：自定义 VJP 与 JVP 规则（仓库中jax.custom_jvp/jax.custom_vjp的实现位于 jax/_src/custom_derivatives.py，对应教程见 docs/notebooks/Custom_derivative_rules_for_Python_code.ipynb）；
• Efficient derivatives at fixed-points：不动点处的高效求导；
• Estimating the trace of a Hessian：用随机 Hessian-vector product 估计 Hessian 的迹；
• Forward-mode autodiff using only reverse-mode：只用反向模式实现前向模式自动微分；
• Taking derivatives with respect to custom data types：对自定义数据类型求导；
• Checkpointing：二项式 checkpointing（为高效反向模式，而非模型快照），仓库中jax.checkpoint/jax.remat的实现在 jax/_src/ad_checkpoint.py；
• Optimizing VJPs with Jacobian pre-accumulation：通过 Jacobian 预累加优化 VJP。

关键 API 速查与源码索引

整体而言，JAX 的自动微分体系可以用一条简洁的调用链概括：grad是value_and_grad的薄封装，value_and_grad内部调用vjp并把单位余切向量送入 VJP；jacrev用vmap把vjp的 pullback 一次性作用到整个标准基；jacfwd用vmap把jvp一次性推前整个标准基；hessian则是二者的组合。理解这条调用链与 JVP/VJP 两种模式的数学定义，就掌握了在 JAX 中编写任意自动微分程序的核心方法论。

【免费下载链接】jaxComposable transformations of Python+NumPy programs: differentiate, vectorize, JIT to GPU/TPU, and more项目地址: https://gitcode.com/GitHub_Trending/ja/jax