JAX中jax.jacfwd计算多输出偏导时是否重复计算中间导数？

问题场景 #

用户实现了以下JAX函数：

def helper_1(params1):
   ...calculations...
   return z

def helper_2(z, params2):
   ...calculations...
   return y

def main(params1, params2):
   z = helper_1(params1)
   y = helper_2(z, params2)
   return z,y

使用jax.jacfwd(main, argnums=(0,1))(params1, params2)计算输出z、y对params1、params2的偏导数，疑问点：

1. 计算y对params1的导数时，JAX是否会重复计算dz/dparams1？
2. helper_1包含TensorFlow for JAX的库函数，是否会影响上述逻辑？

回答 #

• 关于中间结果复用：JAX的前向模式自动微分（jax.jacfwd）会自动做公共子表达式消除，不会重复计算dz/dparams1。在计算z对params1的导数时，已经得到了dz/dparams1的结果；计算y对params1的导数时，根据链式法则需要用到dz/dparams1，JAX会直接复用之前计算的结果，不会重新执行helper_1的梯度计算流程。
• 关于TensorFlow for JAX库函数的影响：只要这些库函数兼容JAX的计算追踪机制（即能被JAX的自动微分、JIT正确处理），就不会影响中间结果的复用逻辑。JAX的自动微分是基于计算图追踪实现的，只要函数能被正确追踪，就能识别并共享中间梯度结果，和函数内部使用的是兼容的第三方库还是原生JAX函数无关。

内容的提问来源于stack exchange，提问作者Dan Leonte