Is there a way to compute the AD Hessian of a function using purely Zygote functions and not ForwardDiff?

Specific Domains Optimization (Mathematical)
1

To solve the issue raised in this post, I’ve been trying to write a Jacobian function using Zygote functions only so I can compute the AD Hessian of my function without using Dual type from ForwardDiff, which doesn’t work with FFTW. Here is the Jacobian function inspired by this.

function jacobian(f,x)
    y,back  = Zygote.pullback(f,x)
    k  = length(y)
    n  = length(x)
    J  = Matrix{eltype(y)}(undef,k,n)
    e_mat = Matrix(I,k,k)
    @inbounds for i = 1:k
        J[i,:] = back(e_mat[:,i])[1]
    end
    (J,)
end

hessian(f, x) = jacobian(x -> gradient(f, x)[1], x)

f(x) = x[1]*x[2]^2
x0=rand(Float64,(2,1))
hessian(f,x0)

However, this gives the following error:

Mutating arrays is not supported

    error(::String)@error.jl:33
    (::Zygote.var"#364#365")(::Nothing)@array.jl:58
    (::Zygote.var"#2245#back#366"{Zygote.var"#364#365"})(::Nothing)@adjoint.jl:59
    (::Zygote.var"#150#151"{Zygote.var"#2245#back#366"{Zygote.var"#364#365"},Tuple{Tuple{Nothing,Nothing},Tuple{Nothing}}})(::Nothing)@lib.jl:191
    (::Zygote.var"#1693#back#152"{Zygote.var"#150#151"{Zygote.var"#2245#back#366"{Zygote.var"#364#365"},Tuple{Tuple{Nothing,Nothing},Tuple{Nothing}}}})(::Nothing)@adjoint.jl:59
    #356@array.jl:38[inlined]
    (::typeof(∂(λ)))(::Tuple{Array{Bool,2},Nothing})@interface2.jl:0
    #2209#back@adjoint.jl:59[inlined]
    (::typeof(∂(λ)))(::Tuple{Nothing,Array{Bool,2},Nothing})@interface2.jl:0
    literal_getindex@builtins.jl:15[inlined]
    (::typeof(∂(λ)))(::Tuple{Nothing,Array{Bool,2},Nothing})@interface2.jl:0
    f@Other: 1[inlined]
    (::typeof(∂(λ)))(::Tuple{Nothing,Array{Bool,1}})@interface2.jl:0
    #41@interface.jl:40[inlined]
    (::typeof(∂(λ)))(::Tuple{Array{Bool,1}})@interface2.jl:0
    gradient@interface.jl:49[inlined]
    (::typeof(∂(gradient)))(::Tuple{Array{Bool,1}})@interface2.jl:0
    #1@Other: 1[inlined]
    (::typeof(∂(#1)))(::Array{Bool,1})@interface2.jl:0
    (::Zygote.var"#41#42"{typeof(∂(#1))})(::Array{Bool,1})@interface.jl:40
    jacobian(::Function, ::Array{Float64,2})@Other: 8
    hessian@Other: 1[inlined]
    top-level scope@Local: 1[inlined]

I cannot think of a way to avoid mutating array, since we need it to construct the Jacobian matrix. I’m just trying to figure out if trying to write a Jacobian function using only Zygote is doomed? f is a very simple function (and another example here) yet it’s giving me errors every time. I can’t imagine how hard it would be to make it compatible to a function involving complex-numbers and FFT. Thank you so much for any guidance, and happy new year!

2

As mentioned in Gradient of Gradient in Zygote - #3 by ChrisRackauckas, Zygote has known issues with nested which are set to be fixed with Diffractor.jl. I would just avoid doing this because you’ll run into other issues. ForwardDiff already works with complex numbers and making it work with FFTs is pretty easy, so that’s a very straightforward path for your actual problem.