I have a function that defines some 2D array. A toy example might be something like:
g = zeros(typeof(a),4,4)
g[1,1] = coordinates^2 + coordinates + 3
g[2,2] = coordinates^2 + coordinates + 4
I want to be able to compute the derivative of the array with respect to the coordinates variables, but not the argument a, which is a constant. There are 4 coordinate variables.
What is the best way to do this using an AD library?
The approach I have used initially is to define a pure function for each of the elements of the array, and to differentiate that e.g.
return coordinates^2 + coordinates + 3
Zygote.jacobian(x -> array_term_g11(x,a), [t r θ ϕ])
However this is obviously verbose and clunky, and it seems like Enzyme is now preferred over Zygote. However it is not obvious to me from reading the docs how to implement this derivative of an array w.r.t a vector input in Enzyme?