No worries, I didn’t try to actually execute it at first either. This seems to work, and the reason I thought it didn’t above seems to be that I’d defined the GradientConfig in the global scope. In the example below g1 is not type-stable, but g2 is.
using ForwardDiff, LinearAlgebra
x = randn(100)
f(x) = dot(x, x)
cfg = ForwardDiff.GradientConfig(f, x, ForwardDiff.Chunk{2}())
g1(x) = ForwardDiff.gradient(f, x, cfg)
g2 = let cfg = cfg
x -> ForwardDiff.gradient(f, x, cfg)
end
@code_warntype g1(x)
@code_warntype g2(x)
I’ll need to check my original, non-MWE code, but suspect the issue there may have also been scoping-related.