JAXvjp报错原因与vmap使用教程

大家好，今天本人给大家带来文章《JAX vjp 失败原因及 vmap 使用指南》，文中内容主要涉及到，如果你对文章方面的知识点感兴趣，那就请各位朋友继续看下去吧~希望能真正帮到你们，谢谢！

当对带有 `custom_vjp` 的函数先 `vmap` 再调用 `vjp` 时，若在定义 `vmap` 版本后覆盖了原始函数名，会导致前向传播中递归调用错误的 vmapped 版本，从而引发 cotangent 形状不匹配的错误。

在 JAX 中，custom_vjp 的前向函数（fwd）必须严格调用原始未变换的函数，以确保其输入/输出形状与 vjp 约定一致：即前向传播返回的 primal_out 形状应与后续 vjp 接收的 cotangent 形状完全匹配（即 cotangent.shape == primal_out.shape）。

问题代码中，关键错误在于：

test_func = vmap(test_func, in_axes=(None, 0))  # ❌ 覆盖了原始 test_func

这导致 test_func_fwd 内部调用的 test_func(jnp.dot(R, R)) 实际执行的是 已 vmapped 的版本，而 jnp.dot(R, R) 的输入 R 是标量（因 R 是 jnp.dot 的结果，shape 为 ()），但 vmapped test_func 期望 R 具有 batch 维度（如 (10, 3)），于是内部逻辑错乱，最终使 primal_out 的隐式形状与 vjp 期望不符——vjp 认为输出是 (10,)，但 bwd 接收到的 residual 和 cotangent 却因前向误调而维度失配，触发报错：

ValueError: Shape of cotangent input to vjp pullback function (10,) must be the same as the shape of corresponding primal input (10, 3).

该错误信息虽表述为“cotangent 应与 primal input 同形”，实则是 JAX 在反向传播校验阶段，因前向路径被污染，无法正确推导出梯度传播所需的张量结构所致。

✅ 正确做法是：保留原始 test_func 不变，仅将 vmap 结果赋给新变量名：

# ✅ 保持原始 test_func 不被覆盖
test_func_mapped = vmap(test_func, in_axes=(None, 0))

# 在 vjp 中使用映射后的版本
primal, f_vjp = vjp(partial(test_func_mapped, f), jnp.ones((10, 3)))
cotangent = jnp.ones(10)  # shape matches primal_out: (10,)
cotangent_out = f_vjp(cotangent)

print(cotangent_out[0].shape)  # → (10, 3)

? 补充注意事项：

• custom_vjp 的 fwd 函数中禁止调用任何高阶变换（如 vmap, jit, grad）后的函数，除非明确设计为支持嵌套；
• 若需批量处理并保留 vjp 可用性，推荐使用 vmap 包裹整个 vjp 调用（即 vmap(vjp(...))），而非先 vmap 函数再 vjp；
• 对于复杂控制流或状态依赖场景，建议通过 jax.custom_vjp + non-differentiable arguments 显式隔离不变参数，并始终在 fwd/bwd 中使用原始函数引用。

遵循“函数变换不覆盖原名”这一原则，可避免绝大多数 vmap 与 custom_vjp 组合时的静默行为异常。

本篇关于《JAXvjp报错原因与vmap使用教程》的介绍就到此结束啦，但是学无止境，想要了解学习更多关于文章的相关知识，请关注golang学习网公众号！