Failing to compute Hessian via Jacobian of gradient

指定された Julia と Diffractor のバージョンで、最小限のネストされた Jacobian の例を再現します。src/codegen/forward.jl:70 から開始し、stack trace に示されている AbstractDifferentiation.jl の jacobian パスを追跡します。Diffractor の gradient の Jacobian が失敗しなくなれば完了です。直接的な Hessian のサポートは別途検討します。

I am on 1.10.0-beta2 with environment:

Status ~/1.10/Project.toml
[c29ec348] AbstractDifferentiation v0.5.2
[9f5e2b26] Diffractor v0.2.1
[31c24e10] Distributions v0.25.100
[f6369f11] ForwardDiff v0.10.36
[276daf66] SpecialFunctions v2.3.1
[4c63d2b9] StatsFuns v1.3.0

Trying to compute Hessian with Diffractor.jl.
First, it would be nice if hessian from AbstractDifferentiation.jl would be supported for Diffractor.jl, like gradient and jacobian, but as far as I can see it is not yet.
Trying to get around this, I am computing the jacobian of the gradient to get the Hessian. But this fails even on a trivial function:

using Diffractor: DiffractorForwardBackend
using AbstractDifferentiation: gradient, jacobian, hessian

θ = [1.0,2.0]
∇(θ) = Diffractor.gradient(θ -> sum(θ.^2), θ)
∇(θ) # gradient works

# Hessian directly
hessian(DiffractorForwardBackend(), θ -> sum(θ.^2), θ) # not supported 

# Hessian as jacobian of gradient
jacobian(DiffractorForwardBackend(), θ -> ∇(θ)[1], θ) # errors with error a bottom 

f(θ) = 2*θ # coding up the gradient directly
jacobian(DiffractorForwardBackend(), θ -> f(θ), θ) # works, as expected

g(θ) = ∇(θ)[1]
jacobian(DiffractorForwardBackend(), θ -> g(θ), θ) # errors with error a bottom

ERROR: Unexprected statement encountered. This is a bug in Diffractor. stmt=$(Expr(:new_opaque_closure, Tuple{Any}, Union{}, Any, ∂⃖¹₁#3(...) @ Diffractor none:0, nothing, :(%6), :(%9), :(%12), :(%15), :(%18), nothing))
Stacktrace:
[1] error(s::String)
@ Base ./error.jl:35
[2] (::Diffractor.var"#mapstmt!#220"{Core.MethodInstance, Int64, Diffractor.var"#emit!#219"{Vector{Any}, Vector{Any}}, Vector{Int64}})(stmt::Any)
@ Diffractor ~/.julia/packages/Diffractor/vzpvE/src/codegen/forward.jl:70
[3] fwd_transform!(ci::Core.CodeInfo, mi::Core.MethodInstance, nargs::Int64, N::Int64)
@ Diffractor ~/.julia/packages/Diffractor/vzpvE/src/codegen/forward.jl:91
[4] fwd_transform(::Core.CodeInfo, ::Core.MethodInstance, ::Vararg{Any})
@ Diffractor ~/.julia/packages/Diffractor/vzpvE/src/codegen/forward.jl:3
[5] perform_fwd_transform(world::UInt64, source::LineNumberNode, ff::Type{Diffractor.∂☆recurse{1}}, args::Any)
@ Diffractor ~/.julia/packages/Diffractor/vzpvE/src/stage1/recurse_fwd.jl:62
[6] ∂☆internal
@ Diffractor ~/.julia/packages/Diffractor/vzpvE/src/stage1/forward.jl:121 [inlined]
[7] ∂☆
@ Diffractor ~/.julia/packages/Diffractor/vzpvE/src/stage1/forward.jl:160 [inlined]
[8] ∂☆
@ Diffractor ~/.julia/packages/Diffractor/vzpvE/src/stage1/forward.jl:254 [inlined]
[9] ∂⃖
@ Diffractor ~/.julia/packages/Diffractor/vzpvE/src/stage1/generated.jl:215 [inlined]
[10] ∂☆internal
@ Diffractor ~/.julia/packages/Diffractor/vzpvE/src/stage1/forward.jl:121 [inlined]
[11] ∂☆
@ Diffractor ~/.julia/packages/Diffractor/vzpvE/src/stage1/forward.jl:160 [inlined]
[12] ∂☆
@ Diffractor ~/.julia/packages/Diffractor/vzpvE/src/stage1/forward.jl:254 [inlined]
[13] ∂⃖
@ Diffractor ~/.julia/packages/Diffractor/vzpvE/src/interface.jl:25 [inlined]
[14] ∂☆internal
@ Diffractor ~/.julia/packages/Diffractor/vzpvE/src/stage1/forward.jl:121 [inlined]
[15] ∂☆
@ Diffractor ~/.julia/packages/Diffractor/vzpvE/src/stage1/forward.jl:160 [inlined]
[16] ∂☆
@ Diffractor ~/.julia/packages/Diffractor/vzpvE/src/stage1/forward.jl:254 [inlined]
[17] ∇
@ Diffractor ~/.julia/packages/Diffractor/vzpvE/src/interface.jl:120 [inlined]
[18] (::Diffractor.∂☆recurse{1})(::ZeroBundle{1, Diffractor.∇{var"#3#4"}}, ::Diffractor.TangentBundle{1, Vector{Float64}, Diffractor.TaylorTangent{Tuple{Vector{Float64}}}})
@ Diffractor ~/.julia/packages/Diffractor/vzpvE/src/stage1/recurse_fwd.jl:0
[19] (::Diffractor.∂☆internal{1})(::ZeroBundle{1, Diffractor.∇{var"#3#4"}}, ::Vararg{Diffractor.AbstractTangentBundle{1}})
@ Diffractor ~/.julia/packages/Diffractor/vzpvE/src/stage1/forward.jl:121
[20] (::Diffractor.∂☆{1})(::ZeroBundle{1, Diffractor.∇{var"#3#4"}}, ::Vararg{Diffractor.AbstractTangentBundle{1}})
@ Diffractor ~/.julia/packages/Diffractor/vzpvE/src/stage1/forward.jl:160
[21] ∂☆
@ Diffractor ~/.julia/packages/Diffractor/vzpvE/src/stage1/forward.jl:254 [inlined]
[22] ∇
@ Diffractor ~/.julia/packages/Diffractor/vzpvE/src/interface.jl:127 [inlined]
[23] (::Diffractor.∂☆recurse{1})(::ZeroBundle{1, Type{Diffractor.∇}}, ::ZeroBundle{1, var"#3#4"}, ::Diffractor.TangentBundle{1, Vector{Float64}, Diffractor.TaylorTangent{Tuple{Vector{Float64}}}})
@ Diffractor ~/.julia/packages/Diffractor/vzpvE/src/stage1/recurse_fwd.jl:0
[24] (::Diffractor.∂☆internal{1})(::ZeroBundle{1, Type{Diffractor.∇}}, ::Vararg{Diffractor.AbstractTangentBundle{1}})
@ Diffractor ~/.julia/packages/Diffractor/vzpvE/src/stage1/forward.jl:121
[25] ∂☆
@ Diffractor ~/.julia/packages/Diffractor/vzpvE/src/stage1/forward.jl:160 [inlined]
[26] ∇
@ Diffractor ./REPL[8]:1 [inlined]
[27] ∂☆internal
@ Diffractor ~/.julia/packages/Diffractor/vzpvE/src/stage1/forward.jl:121 [inlined]
[28] ∂☆
@ Diffractor ~/.julia/packages/Diffractor/vzpvE/src/stage1/forward.jl:160 [inlined]
[29] #13
@ Diffractor ./REPL[16]:1 [inlined]
[30] (::Diffractor.∂☆recurse{1})(::ZeroBundle{1, var"#13#14"}, ::Diffractor.TangentBundle{1, Vector{Float64}, Diffractor.TaylorTangent{Tuple{Vector{Float64}}}})
@ Diffractor ~/.julia/packages/Diffractor/vzpvE/src/stage1/recurse_fwd.jl:0
[31] (::Diffractor.∂☆internal{1})(::ZeroBundle{1, var"#13#14"}, ::Vararg{Diffractor.AbstractTangentBundle{1}})
@ Diffractor ~/.julia/packages/Diffractor/vzpvE/src/stage1/forward.jl:121
[32] (::Diffractor.∂☆{1})(::ZeroBundle{1, var"#13#14"}, ::Vararg{Diffractor.AbstractTangentBundle{1}})
@ Diffractor ~/.julia/packages/Diffractor/vzpvE/src/stage1/forward.jl:160
[33] (::Diffractor.var"#pushforward#359"{var"#13#14", Tuple{Vector{Float64}}})(vs::Tuple{Vector{Float64}})
@ Diffractor ~/.julia/packages/Diffractor/vzpvE/src/AbstractDifferentiation.jl:13
[34] (::Diffractor.var"#358#362")(cols::Any)
@ Diffractor ~/.julia/packages/AbstractDifferentiation/eEkWP/src/AbstractDifferentiation.jl:526 [inlined]
[35] mapslices(f::Diffractor.var"#358#362"{Diffractor.var"#pushforward#359"{var"#13#14", Tuple{Vector{Float64}}}}, A::Matrix{Float64}; dims::Int64)
@ Base ./abstractarray.jl:3179
[36] jacobian(b::DiffractorForwardBackend, f::Function, args::Vector{Float64})
@ Diffractor ~/.julia/packages/AbstractDifferentiation/eEkWP/src/AbstractDifferentiation.jl:524