Hessian of a buffer array in Zygote

General Usage
,
1

I have a function which takes some coordinates, a parameter, and creates a 2D array.

I want to get the Jacobian and the Hessian of this array w.r.t the input coordinates using Zygote.

The Jacobian works fine, but the Hessian does not:



using Zygote


function some_array(coords,a) # `coords` is a vector length 4, `a` is a scalar 

    xs = zeros(typeof(a),4,4)
    g = Zygote.bufferfrom(xs) #a 4x4 buffered Zygote object

    t,r,θ,ϕ =  coords[1],coords[2],coords[3],coords[4] #unpack the coordinates

    g[1,1] = 2*r #assign array element

    return copy(g)


end 

#Coordinates [t,r,θ,ϕ] and parameter `a` 
t = 0.0
r = 10.0
θ = π/2.0 
ϕ = 0.0 
a = 0.1
coords = [t,r,θ,ϕ]

#Create an array. Works ok.
some_array(coords,a)

#Get jacobian of array. Works ok.
array_jacobian = jacobian(x -> some_array(x,a), coords)[1]
array_jacobian = reshape(array_jacobian,(4,4,4)) #array_jacobian[:, :, i] is the derivative of some_array w.r.t. coords[i]

#Get hessian of array. Does not work?
array_hessian = hessian(x -> some_array(x,a), coords)[1]

Whilst array_jacobian returns what I expect, array_hessian throws an error:

LoadError: MethodError: no method matching Float64(::ForwardDiff.Dual{Nothing, Float64, 4})

Can someone guide me as to where I am going wrong?