I’ve run into a type-stability issue when trying to compute third-order derivatives with ForwardDiff.
I have a potential function that takes a vector as input. Computing the Hessian with ForwardDiff.hessian works well and is entirely type-stable for my use case. However, I also need the Jacobian of the Hessian (i.e. a third-order derivative tensor). Ideally, I would like to construct both a HessianConfig and a JacobianConfig for the Hessian, so that I can reuse them and keep the cost per evaluation low — this part of the code runs very frequently.
Unfortunately, it seems that nested ForwardDiff calls of this sort are not supported, or at least I haven’t been able to make them type-stable.
Right now my workaround is to cache Hessian configs keyed by input dual types, which avoids repeated allocations, but the resulting Jacobian-of-Hessian function is still type-unstable.
Has anyone found an alternative approach for computing third derivatives efficiently, or a way to make this pattern type-stable? Or is there something fundamental I’m misunderstanding about how ForwardDiff handles nested AD?
Any suggestions or pointers would be appreciated!