本文详解如何使用 jax 的 `@custom_jvp` 为具有多个返回值的函数（如 `(out₁, out₂)`）精确实现自定义前向模式微分，涵盖 `primals` 与 `tangents` 元组的结构匹配、雅可比矩阵验证及关键注意事项。
在 JAX 中，@custom_jvp 并不限于单输出函数——它天然支持任意数量的输入与输出。核心在于：primals_out 和 tangents_out 必须严格保持与原函数返回值相同的嵌套结构与数量。例如，若函数返回二元组 (out_1, out_2)，则 tangents_out 也必须是长度为 2 的元组，分别对应每个输出沿各输入方向的方向导数（即雅可比行向量）。
以下是一个典型示例，定义函数 f(x, y) = (x * y, x / y) 并为其注册自定义前向微分规则：
import jax
import jax.numpy as jnp
@jax.custom_jvp
def f(x, y):
return x * y, x / y
@f.defjvp
def f_jvp(primals, tangents):
x, y = primals # 解包原始输入（与函数签名一致）
x_dot, y_dot = tangents # 解包切向量（顺序、数量必须与 primals 完全匹配）
# 原始输出（forward pass）
primals_out = f(x, y) # → (x*y, x/y)
# 手动推导的切向量：对每个输出分别应用链式法则
# ∂(x*y)/∂x * x_dot + ∂(x*y)/∂y * y_dot = y * x_dot + x * y_dot
# ∂(x/y)/∂x * x_dot + ∂(x/y)/∂y * y_dot = (1/y) * x_dot + (-x/y²) * y_dot
tangents_out = (
x_dot * y + x * y_dot,
x_dot / y - x * y_dot / (y ** 2)
)
return primals_out, tangents_out注册完成后，即可无缝调用高阶自动微分 API。例如，计算完整雅可比矩阵（对两个输入分别求两个输出的偏导）：
x = jnp.float32(0.5) y = jnp.float32(2.0) # 返回形状为 ((∂f₀/∂x, ∂f₀/∂y), (∂f₁/∂x, ∂f₁/∂y)) jac = jax.jacobian(f, argnums=(0, 1))(x, y) print(jac) # ((Array(2., dtype=float32), Array(0.5, dtype=float32)), # (Array(0.5, dtype=float32), Array(-0.125, dtype=float32)))
✅ 验证表明：该结果与未使用 custom_jvp 的原生 f2(x,y) 完全一致，证明自定义规则数学正确且实现无误。
⚠️ 关键注意事项：
- tangents 元组长度必须等于函数参数个数（len(primals)），且每个元素形状需与对应输入一致（标量/数组）；
- tangents_out 元组长度必须等于函数返回值个数，每个元素形状需与对应输出一致（如 out_1 是标量，则 tangents_out[0] 也应为标量）；
- 若函数返回嵌套结构（如 {'a': out1, 'b': out2}），tangents_out 也需保持相同字典结构；
- @custom_jvp 定义的是前向模式微分；若需反向模式（如 grad 或 vjp），应使用 @custom_vjp 并正确定义 fwd 与 bwd 函数（原理类似，但 bwd 接收 cotangents 并返回输入梯度）。
掌握此模式后，你可灵活为任意复杂多输出函数（如物理模拟器、多任务神经网络头、带辅助损失的训练函数）注入高效、可控的梯度逻辑。
