Failing to compute Hessian via Jacobian of gradient

Tái hiện các ví dụ Jacobian lồng nhau tối thiểu với các phiên bản Julia và Diffractor được liệt kê. Bắt đầu tại src/codegen/forward.jl:70 và lần theo đường dẫn jacobian của AbstractDifferentiation.jl được hiển thị trong stack trace. Hoàn thành khi Jacobian của gradient Diffractor không còn thất bại; hỗ trợ Hessian trực tiếp sẽ được xem xét riêng.

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