Zygote cannot evaluate Hessian if the function argument is annotated

47 Views Asked by At

I am trying to calculate some Hessian matrix in Julia using Zygote package. Consider the example below:

using Zygote
function test(x::Vector{Float64})
    return x[1]*x[2]
end
x0 = Vector{Float64}([1.,2.])
hessian(x->test(x), x0)

This will produces the error below:

MethodError: no method matching test(::Matrix{ForwardDiff.Dual{Nothing, Float64, 3}})

Closest candidates are:
  test(::Vector{Float64})

The issue is I annotated the argument of the function test. If I do not annotate, it will be fine. Is there a way to overcome this issue if I choose to annotate the function's argument? Apart from code readability, I fear not making the argument's type clear might cause sub-optimal performance.

1

There are 1 best solutions below

0
mcabbott On

Zygote.hessian uses ForwardDiff over Zygote, and hence needs functions to accept arrays of Dual numbers.

Narrow type annotations aren't helpful for performance. Plain function test(x) will compile exactly the same specialisations for Vector{Float64}. As will things like test(x::AbstractVector{<:Real}) or just test(x::AbstractVector), if you want to document your intentions, and ensure an error on mistakes like test.([1,2,3]).