I have a function that defines some 2D array. A toy example might be something like:

function create_array(coordinates,a)
g = zeros(typeof(a),4,4)
g[1,1] = coordinates[1]^2 + coordinates[2] + 3
g[2,2] = coordinates[2]^2 + coordinates[3] + 4
return g
end

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.

function array_term_g11(coordinates,a)
return coordinates[1]^2 + coordinates[2] + 3
end
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?

As a follow up, how would I extend this to get second derivatives i.e. the Hessian? Is this actually possible with Enzyme - I only see references to jacobians in the docs…?

On my phone, but one thing is that your inner closure function only takes one arg (x), instead of two (x,a) → …

Also any arguments that you don’t pass an activity tuple (like x on the inner autodiff), are assumed constant — which means the inner AD doesn’t do what you expect.

I’ll write a snippet tomorrow after flying out from EnzymeCon