当对带有 `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 组合时的静默行为异常。
