Zygote-based forward-over-reverse HesVecs

[ChrisRackauckas]

Zygote is slow:

using SparseDiffTools, BenchmarkTools
x = rand(300)
v = rand(300)
f(u) = sum(abs2,u)
du = similar(x)
c1 = similar(x); c2 = similar(x); c3 = similar(x); c4 = similar(x)
cache1 = ForwardDiff.Dual{SparseDiffTools.DeivVecTag}.(x, v)
cache2 = ForwardDiff.Dual{SparseDiffTools.DeivVecTag}.(x, v)
config = ForwardDiff.GradientConfig(f,x)
@btime num_hesvec!($du, $f, $x, $v, $c1, $c3, $c4)
@btime numauto_hesvec!($du, $f, $x, $v, $config, $c1, $c2)
@btime autonum_hesvec!($du, $f, $x, $v, $c1, $cache1, $cache2)
@btime numback_hesvec!(du, f, x, v, $c1, $c2)
@btime autoback_hesvec!(du, f, x, v, $cache1, $cache2) # fails

  50.891 μs (0 allocations: 0 bytes)
  23.646 μs (1 allocation: 16 bytes)
  198.940 μs (0 allocations: 0 bytes)
  88.932 μs (3193 allocations: 90.84 KiB) # This is Zygote :(

but if it improves this will already exist.

Note that forward diff -> Zygote currently fails. MWE:

using Zygote, ForwardDiff, SparseDiffTools

function _autoback_hesvec(f,x,v)
    g = (x) -> first(Zygote.gradient(f,x))
    ForwardDiff.partials.(g(ForwardDiff.Dual{Nothing}.(x, v)), 1)
end

x = rand(300)
v = rand(300)
f(u) = sum(abs2,u)
numback_hesvec(f, x, v) # works
_autoback_hesvec(f, x, v) # fails

fails with

MethodError: no method matching (::Zygote.Jnew{ForwardDiff.Dual{Nothing,Float64,1},Nothing,false})(::ForwardDiff.Dual{Nothing,Float64,1})
Closest candidates are:
  Jnew(!Matched::Union{Nothing, RefValue, NamedTuple}) where {T, G} at C:\Users\accou\.julia\packages\Zygote\VeaFW\src\lib\lib.jl:216
(::getfield(Zygote, Symbol("##301#back#152")){Zygote.Jnew{ForwardDiff.Dual{Nothing,Float64,1},Nothing,false}})(::ForwardDiff.Dual{Nothing,Float64,1}) at grad.jl:46
Type at dual.jl:19 [inlined]
(::typeof(∂(ForwardDiff.Dual{Nothing,Float64,1})))(::ForwardDiff.Dual{Nothing,Float64,1}) at interface2.jl:0
literal_pow at dual.jl:55 [inlined]
(::getfield(Zygote, Symbol("##1077#1084")))(::typeof(∂(Base.literal_pow)), ::ForwardDiff.Dual{Nothing,Float64,1}) at broadcast.jl:97
iterate at generator.jl:36 [inlined]
collect(::Base.Generator{Base.Iterators.Zip{Tuple{Array{typeof(∂(Base.literal_pow)),1},Zygote.FillArray{ForwardDiff.Dual{Nothing,Float64,1},1}}},getfield(Base, Symbol("##3#4")){getfield(Zygote, Symbol("##1077#1084"))}}) at array.jl:606
map at abstractarray.jl:2091 [inlined]
(::getfield(Zygote, Symbol("##1076#1083")){Tuple{Base.RefValue{typeof(^)},Array{ForwardDiff.Dual{Nothing,Float64,1},1},Base.RefValue{Val{2}}},Val{4},Array{typeof(∂(Base.literal_pow)),1}})(::Zygote.FillArray{ForwardDiff.Dual{Nothing,Float64,1},1}) at broadcast.jl:97
#2756#back at grad.jl:46 [inlined]
(::getfield(Zygote, Symbol("##129#130")){getfield(Zygote, Symbol("##2756#back#1087")){getfield(Zygote, Symbol("##1076#1083")){Tuple{Base.RefValue{typeof(^)},Array{ForwardDiff.Dual{Nothing,Float64,1},1},Base.RefValue{Val{2}}},Val{4},Array{typeof(∂(Base.literal_pow)),1}}},Tuple{NTuple{4,Nothing},Tuple{Nothing}}})(::Zygote.FillArray{ForwardDiff.Dual{Nothing,Float64,1},1}) at lib.jl:108
#254#back at grad.jl:46 [inlined]
broadcasted at broadcast.jl:1167 [inlined]
f at test_jaches_products.jl:21 [inlined]
(::typeof(∂(f)))(::ForwardDiff.Dual{Nothing,Float64,1}) at interface2.jl:0
(::getfield(Zygote, Symbol("##34#35")){typeof(∂(f))})(::ForwardDiff.Dual{Nothing,Float64,1}) at interface.jl:38
gradient(::Function, ::Array{ForwardDiff.Dual{Nothing,Float64,1},1}) at interface.jl:47
(::getfield(Main, Symbol("##17#18")){typeof(f)})(::Array{ForwardDiff.Dual{Nothing,Float64,1},1}) at test_jaches_products.jl:161
_autoback_hesvec(::Function, ::Array{Float64,1}, ::Array{Float64,1}) at test_jaches_products.jl:162
top-level scope at none:0

[ChrisRackauckas]

@MikeInnes do you know how to workaround that issue with the forward-over-reverse implementation?