I’m interested in calculating second order partial derivatives w.r.t. to different arguments \partial k_j \partial x_i f.
Consider
function f(x,k)
return @. k * x^2
end
Calculation of second-order derivatives w.r.t. the same variable \partial x_j \partial x_i works well by a combination of ForwardDiff.jl and Zygote.jl.
A naive approach to get \partial k_j \partial x_i f
der(x,k) = ForwardDiff.jacobian((x)->f(x,k), x)
Zygote.jacobian(der, x, k)
fails as Zygote apparently cannot deal with closures
┌ Warning:
ForwardDiff.jacobian(f, x)
within Zygote cannot track gradients with respect tof
,
│ andf
appears to be a closure, or a struct with fields (according toissingletontype(typeof(f))
).
│ typeof(f) = var"#75#76"{Vector{Float64}}
└ @ Zygote ~/.julia/packages/Zygote/WOy6z/src/lib/forward.jl:150
Is there any workaround? Can this be achieved by (combinations) of other AD frameworks?