Second-order autodiff: which combinations should work?

[stevengj]

If you want the gradient of a scalar-valued function that depends on the gradient of another scalar-valued function, you can use forward-over-reverse combining ForwardDiff with e.g. Zygote or Enzyme or ReverseDiff. See:

Nested AD with Lux etc

For future reference, here is a manual implementation of forward-over-reverse calculation with parameter gradients (mixed second derivatives). In particular, if you have a scalar-valued h(x,p) = g(\nabla_x f) for some scalar-valued f(x,p), then one can similarly derive: \left. \nabla_p h \right|_{x,p} = \left. \frac{\partial}{\partial\alpha} \left. \nabla_p f \right|_{x + \alpha \left. \nabla g \right|_{z},p} \right|_{\alpha = 0} \, , where z = \left. \nabla_x f \right|_{x,p}. Here is an e…

(You can also use this approach for general Hessians, but it was less obvious to me that it is efficient for scalar-valued functions.)