Missing `_shadow` method in Enzyme extension

[timholy]

using LinearAlgebra
using DifferentiationInterface
using Enzyme: Enzyme

function lagrangian(A, b, Bdiag, y, x, s, z, q, λ, η, xₖ)
    r1 = A*x - y
    r2 = s - x
    return sum(abs2, b - y) / 2 + λ * dot(x - xₖ, Diagonal(Bdiag), x - xₖ) / 2 - η * sum(log, s) + dot(z, r1) + dot(q, r2)
end
function lagrangian(A, b, Bdiag, top, bottom, λ, η, xₖ)
    m, n = size(A)
    y, x, s = view(top, 1:m), view(top, m+1:m+n), view(top, m+n+1:m+2n)
    z, q = view(bottom, 1:m), view(bottom, m+1:m+n)
    return lagrangian(A, b, Bdiag, y, x, s, z, q, λ, η, xₖ)
end
function lagrangian(A, b, Bdiag, topbottom, λ, η, xₖ)
    m, n = size(A)
    top, bottom = view(topbottom, 1:m+2n), view(topbottom, m+2n+1:2m+3n)
    return lagrangian(A, b, Bdiag, top, bottom, λ, η, xₖ)
end

A = [ -0.24  -0.82   1.35   0.36   0.35
        -0.53  -0.20  -0.76   0.98  -0.54
        0.22   1.25  -1.60  -1.37  -1.94
        -0.51  -0.56  -0.08   0.96   0.46
        0.48  -2.25   0.38   0.06  -1.29 ]
m, n = size(A)
b = [-1.6, 0.19, 0.17, 0.31, -1.27]
Bdiag = diag(A'*A)
η = 0.01
λ = 1e-4
top0, bottom0 = vcat(zeros(m+n), ones(n)), ones(m+n)
xₖ = top0[m+1:m+n]
tb0 = vcat(top0, bottom0)
function makeL(A, b, Bdiag, λ, η, xₖ)
    L(tb) = lagrangian(A, b, Bdiag, tb, λ, η, xₖ)
    return L
end
L = makeL(A, b, Bdiag, λ, η, xₖ)

@assert isfinite(L(tb0))

backend = AutoEnzyme(; mode=nothing, function_annotation=Enzyme.Duplicated)
gradient(L, backend, tb0)  # fine
hessian(L, backend, tb0)   # errors

The error is

julia> hessian(L, backend, tb0)   # errors
ERROR: MethodError: no method matching _shadow(::AutoEnzyme{Nothing, EnzymeCore.Duplicated}, ::EnzymeCore.ForwardMode{true, EnzymeCore.FFIABI, false, false, false}, ::Val{13}, ::DifferentiationInterface.FunctionContext{var"#L#1"{Matrix{Float64}, Vector{Float64}, Vector{Float64}, Float64, Float64, Vector{Float64}}})

Closest candidates are:
  _shadow(::AutoEnzyme{M, <:Union{Nothing, EnzymeCore.Const}}, ::EnzymeCore.Mode, ::Val{B}, ::DifferentiationInterface.FunctionContext) where {M, B}
   @ DifferentiationInterfaceEnzymeExt C:\Users\timho\.julia\dev\DifferentiationInterface\DifferentiationInterface\ext\DifferentiationInterfaceEnzymeExt\utils.jl:89
  _shadow(::AutoEnzyme, ::EnzymeCore.Mode, ::Val{B}, ::ConstantOrCache) where B
   @ DifferentiationInterfaceEnzymeExt C:\Users\timho\.julia\dev\DifferentiationInterface\DifferentiationInterface\ext\DifferentiationInterfaceEnzymeExt\utils.jl:73
  _shadow(::AutoEnzyme, ::EnzymeCore.Mode, ::Val{B}, ::Cache) where B
   @ DifferentiationInterfaceEnzymeExt C:\Users\timho\.julia\dev\DifferentiationInterface\DifferentiationInterface\ext\DifferentiationInterfaceEnzymeExt\utils.jl:64
  ...

Stacktrace:
  [1] (::DifferentiationInterfaceEnzymeExt.var"#7#8"{13, AutoEnzyme{…}, EnzymeCore.ForwardMode{…}})(c_wrapped::DifferentiationInterface.FunctionContext{var"#L#1"{…}})
    @ DifferentiationInterfaceEnzymeExt C:\Users\timho\.julia\dev\DifferentiationInterface\DifferentiationInterface\ext\DifferentiationInterfaceEnzymeExt\utils.jl:103
  [2] map
    @ .\tuple.jl:293 [inlined]
  [3] make_context_shadows(::AutoEnzyme{…}, ::EnzymeCore.ForwardMode{…}, ::Val{…}, ::DifferentiationInterface.FunctionContext{…}, ::Constant{…}, ::Constant{…})
    @ DifferentiationInterfaceEnzymeExt C:\Users\timho\.julia\dev\DifferentiationInterface\DifferentiationInterface\ext\DifferentiationInterfaceEnzymeExt\utils.jl:102
  [4] prepare_pushforward_nokwarg(::Val{…}, ::typeof(DifferentiationInterface.shuffled_gradient), ::AutoEnzyme{…}, ::Vector{…}, ::NTuple{…}, ::DifferentiationInterface.FunctionContext{…}, ::Constant{…}, ::Constant{…})
    @ DifferentiationInterfaceEnzymeExt C:\Users\timho\.julia\dev\DifferentiationInterface\DifferentiationInterface\ext\DifferentiationInterfaceEnzymeExt\forward_onearg.jl:20
  [5] _prepare_hvp_aux(::Val{…}, ::DifferentiationInterface.ForwardOverReverse, ::DifferentiationInterface.DontPrepareInner, ::var"#L#1"{…}, ::AutoEnzyme{…}, ::Vector{…}, ::NTuple{…})
    @ DifferentiationInterface C:\Users\timho\.julia\dev\DifferentiationInterface\DifferentiationInterface\src\second_order\hvp.jl:166
  [6] prepare_hvp_nokwarg(::Val{true}, ::var"#L#1"{Matrix{Float64}, Vector{Float64}, Vector{Float64}, Float64, Float64, Vector{Float64}}, ::AutoEnzyme{Nothing, EnzymeCore.Duplicated}, ::Vector{Float64}, ::NTuple{13, Vector{Float64}})
    @ DifferentiationInterface C:\Users\timho\.julia\dev\DifferentiationInterface\DifferentiationInterface\src\second_order\hvp.jl:125
  [7] _prepare_hessian_aux(::Val{true}, ::DifferentiationInterface.BatchSizeSettings{13, false, false}, ::var"#L#1"{Matrix{Float64}, Vector{Float64}, Vector{Float64}, Float64, Float64, Vector{Float64}}, ::AutoEnzyme{Nothing, EnzymeCore.Duplicated}, ::Vector{Float64})
    @ DifferentiationInterface C:\Users\timho\.julia\dev\DifferentiationInterface\DifferentiationInterface\src\second_order\hessian.jl:125
  [8] prepare_hessian_nokwarg(::Val{true}, ::var"#L#1"{Matrix{Float64}, Vector{Float64}, Vector{Float64}, Float64, Float64, Vector{Float64}}, ::AutoEnzyme{Nothing, EnzymeCore.Duplicated}, ::Vector{Float64})
    @ DifferentiationInterface C:\Users\timho\.julia\dev\DifferentiationInterface\DifferentiationInterface\src\second_order\hessian.jl:106
  [9] hessian(::var"#L#1"{Matrix{Float64}, Vector{Float64}, Vector{Float64}, Float64, Float64, Vector{Float64}}, ::AutoEnzyme{Nothing, EnzymeCore.Duplicated}, ::Vector{Float64})
    @ DifferentiationInterface C:\Users\timho\.julia\dev\DifferentiationInterface\DifferentiationInterface\src\second_order\hessian.jl:34
 [10] top-level scope
    @ REPL[3]:1

[gdalle]

Thanks for reporting this. It's an easy fix, if you want to try and contribute it yourself, otherwise I'll do it. I think I never added it because second order with Enzyme used to already be quite brittle with constant functions, so closing over active data didn't seem relevant at the time.

[gdalle]

Essentially you want to imitate whatever is done for ConstantOrCache here

[timholy]

I'll give it a whirl, it indeed looked pretty simple. Main holdup was a suitable test, but it looks like adding AutoEnzyme(; mode=nothing, function_annotation = Enzyme.Duplicated) to

DifferentiationInterface.jl/DifferentiationInterface/test/Back/Enzyme/test.jl

Lines 67 to 72 in 0dd1abf

	[
	AutoEnzyme(),
	SecondOrder(
	AutoEnzyme(; mode = Enzyme.Reverse), AutoEnzyme(; mode = Enzyme.Forward)
	),
	],

is enough to catch the bug.

[timholy]

Hmm, maybe not so simple?

Correctness: Error During Test at C:\Users\timho\.julia\dev\DifferentiationInterface\DifferentiationInterfaceTest\src\test_differentiation.jl:156
  Got exception outside of a @testtion_annotation=EnzymeCore.Duplicated) - 1/2
  Type of ghost or constant type EnzymeCore.BatchDuplicated{typeof(DifferentiationInterface.shuffled_gradient!), 2} is marked as differentiable.
  Stacktrace:ario: 6/24
    [1] error(s::String)
      @ Base .\error.jl:35
    [2] enzyme!(job::GPUCompiler.CompilerJob{Enzyme.Compiler.EnzymeTarget{GPUCompiler.NativeCompilerTarget}, Enzyme.Compiler.EnzymeCompilerParams{Enzyme.Compiler.PrimalCompilerParams}}, interp::Enzyme.Compiler.Interpreter.EnzymeInterpreter{Nothing}, mod::LLVM.Module, primalf::LLVM.Function, TT::Type, mode::Enzyme.API.CDerivativeMode, width::Int64, parallel::Bool, actualRetType::Type, wrap::Bool, modifiedBetween::Tuple{Vararg{Bool, N}} where N, returnPrimal::Bool, expectedTapeType::Type, loweredArgs::Set{Int64}, boxedArgs::Set{Int64}, removedRoots::Set{Int64})  input_size: (6,)
      @ Enzyme.Compiler C:\Users\timho\.julia\packages\Enzyme\OhLFb\src\compiler.jl:2516
    [3] compile_unhooked(output::Symbol, job::GPUCompiler.CompilerJob{Enzyme.Compiler.EnzymeTarget{GPUCompiler.NativeCompilerTarget}, Enzyme.Compiler.EnzymeCompilerParams{Enzyme.Compiler.PrimalCompilerParams}})
      @ Enzyme.Compiler C:\Users\timho\.julia\packages\Enzyme\OhLFb\src\compiler.jl:5646
    [4] compile(target::Symbol, job::GPUCompiler.CompilerJob; kwargs::@Kwargs{})
      @ GPUCompiler C:\Users\timho\.julia\packages\GPUCompiler\j4HFa\src\driver.jl:67
    [5] compile
      @ C:\Users\timho\.julia\packages\GPUCompiler\j4HFa\src\driver.jl:55 [inlined]
    [6] _thunk(job::GPUCompiler.CompilerJob{Enzyme.Compiler.EnzymeTarget{GPUCompiler.NativeCompilerTarget}, Enzyme.Compiler.EnzymeCompilerParams{Enzyme.Compiler.PrimalCompilerParams}}, postopt::Bool)
      @ Enzyme.Compiler C:\Users\timho\.julia\packages\Enzyme\OhLFb\src\compiler.jl:6558
    [7] _thunk
      @ C:\Users\timho\.julia\packages\Enzyme\OhLFb\src\compiler.jl:6556 [inlined]
    [8] cached_compilation
      @ C:\Users\timho\.julia\packages\Enzyme\OhLFb\src\compiler.jl:6613 [inlined]
    [9] thunkbase(mi::Core.MethodInstance, World::UInt64, FA::Type{<:EnzymeCore.Annotation}, A::Type{<:EnzymeCore.Annotation}, TT::Type, Mode::Enzyme.API.CDerivativeMode, width::Int64, ModifiedBetween::Tuple{Vararg{Bool, N}} where N, ReturnPrimal::Bool, ShadowInit::Bool, ABI::Type, ErrIfFuncWritten::Bool, RuntimeActivity::Bool, StrongZero::Bool, edges::Vector{Any})
      @ Enzyme.Compiler C:\Users\timho\.julia\packages\Enzyme\OhLFb\src\compiler.jl:6729
   [10] thunk_generator(world::UInt64, source::Union{LineNumberNode, Method}, FA::Type, A::Type, TT::Type, Mode::Enzyme.API.CDerivativeMode, Width::Int64, ModifiedBetween::Tuple{Vararg{Bool, N}} where N, ReturnPrimal::Bool, ShadowInit::Bool, ABI::Type, ErrIfFuncWritten::Bool, RuntimeActivity::Bool, StrongZero::Bool, self::Any, fakeworld::Any, fa::Type, a::Type, tt::Type, mode::Type, width::Type, modifiedbetween::Type, returnprimal::Type, shadowinit::Type, abi::Type, erriffuncwritten::Type, runtimeactivity::Type, strongzero::Type)
      @ Enzyme.Compiler C:\Users\timho\.julia\packages\Enzyme\OhLFb\src\compiler.jl:6873
   [11] autodiff
      @ C:\Users\timho\.julia\packages\Enzyme\OhLFb\src\Enzyme.jl:658 [inlined]
   [12] value_and_pushforward!(::typeof(DifferentiationInterface.shuffled_gradient!), ::Vector{Float64}, ::Tuple{Vector{Float64}, Vector{Float64}}, ::DifferentiationInterfaceEnzymeExt.EnzymeTwoArgPushforwardPrep{Tuple{typeof(DifferentiationInterface.shuffled_gradient!), Vector{Float64}, AutoEnzyme{Nothing, EnzymeCore.Duplicated}, Vector{Float64}, Tuple{Vector{Float64}, Vector{Float64}}, Tuple{DifferentiationInterface.FunctionContext{DifferentiationInterfaceTest.MultiplyByConstant{:out, typeof(DifferentiationInterfaceTest.arr_to_num_linalg)}}, Constant{AutoEnzyme{Nothing, EnzymeCore.Duplicated}}, Constant{DifferentiationInterface.Rewrap{1, Tuple{typeof(DifferentiationInterface.constant_maker)}}}, Constant{Float64}}}, Tuple{typeof(DifferentiationInterface.shuffled_gradient!), typeof(DifferentiationInterface.shuffled_gradient!)}, Tuple{Tuple{DifferentiationInterfaceTest.MultiplyByConstant{:out, typeof(DifferentiationInterfaceTest.arr_to_num_linalg)}, DifferentiationInterfaceTest.MultiplyByConstant{:out, typeof(DifferentiationInterfaceTest.arr_to_num_linalg)}}, Nothing, Nothing, Nothing}}, ::AutoEnzyme{Nothing, EnzymeCore.Duplicated}, ::Vector{Float64}, ::Tuple{Vector{Float64}, Vector{Float64}}, ::DifferentiationInterface.FunctionContext{DifferentiationInterfaceTest.MultiplyByConstant{:out, typeof(DifferentiationInterfaceTest.arr_to_num_linalg)}}, ::Constant{AutoEnzyme{Nothing, EnzymeCore.Duplicated}}, ::Constant{DifferentiationInterface.Rewrap{1, Tuple{typeof(DifferentiationInterface.constant_maker)}}}, ::Constant{Float64})
      @ DifferentiationInterfaceEnzymeExt C:\Users\timho\.julia\dev\DifferentiationInterface\DifferentiationInterface\ext\DifferentiationInterfaceEnzymeExt\forward_twoarg.jl:99
   [13] _gradient_and_hvp_aux!
      @ C:\Users\timho\.julia\dev\DifferentiationInterface\DifferentiationInterface\src\second_order\hvp.jl:471 [inlined]
   [14] gradient_and_hvp!
      @ C:\Users\timho\.julia\dev\DifferentiationInterface\DifferentiationInterface\src\second_order\hvp.jl:446 [inlined]
   [15] gradient_and_hvp!
      @ C:\Users\timho\.julia\dev\DifferentiationInterface\DifferentiationInterface\src\second_order\hvp.jl:117 [inlined]
   [16] (::DifferentiationInterfaceTest.var"#230#234"{AutoEnzyme{Nothing, EnzymeCore.Duplicated}, Scenario{:hvp, :in, :out, DifferentiationInterfaceTest.MultiplyByConstant{:out, typeof(DifferentiationInterfaceTest.arr_to_num_linalg)}, Vector{Float64}, Float64, Tuple{Vector{Float64}, Vector{Float64}}, Tuple{Constant{Float64}}, Vector{Float64}, Tuple{Vector{Float64}, Vector{Float64}}, @NamedTuple{x::Vector{Float64}, contexts::Tuple{Constant{Float64}}, t::Tuple{Vector{Float64}, Vector{Float64}}}}, Tuple{Constant{Float64}}, Tuple{Vector{Float64}, Vector{Float64}}, Vector{Float64}, Tuple{Vector{Float64}, Vector{Float64}}, Vector{Float64}, DifferentiationInterfaceTest.MultiplyByConstant{:out, typeof(DifferentiationInterfaceTest.arr_to_num_linalg)}, DifferentiationInterfaceTest.var"#≈#232"{typeof(isapprox), Int64, Float64}})(preptup_val::Tuple{})
      @ DifferentiationInterfaceTest C:\Users\timho\.julia\dev\DifferentiationInterface\DifferentiationInterfaceTest\src\tests\correctness_eval.jl:829
   [17] #62
      @ .\tuple.jl:627 [inlined]
   [18] BottomRF
      @ .\reduce.jl:86 [inlined]
   [19] afoldl
      @ .\operators.jl:544 [inlined]
   [20] _foldl_impl
      @ .\reduce.jl:68 [inlined]
   [21] foldl_impl
      @ .\reduce.jl:48 [inlined]
   [22] mapfoldl_impl
      @ .\reduce.jl:44 [inlined]
   [23] mapfoldl
      @ .\reduce.jl:175 [inlined]
   [24] foldl
      @ .\reduce.jl:198 [inlined]
   [25] foreach
      @ .\tuple.jl:627 [inlined]
   [26] test_correctness(ba::AutoEnzyme{Nothing, EnzymeCore.Duplicated}, scen::Scenario{:hvp, :in, :out, DifferentiationInterfaceTest.MultiplyByConstant{:out, typeof(DifferentiationInterfaceTest.arr_to_num_linalg)}, Vector{Float64}, Float64, Tuple{Vector{Float64}, Vector{Float64}}, Tuple{Constant{Float64}}, Vector{Float64}, Tuple{Vector{Float64}, Vector{Float64}}, @NamedTuple{x::Vector{Float64}, contexts::Tuple{Constant{Float64}}, t::Tuple{Vector{Float64}, Vector{Float64}}}}; isapprox::typeof(isapprox), atol::Int64, rtol::Float64, scenario_intact::Bool, sparsity::Bool, reprepare::Bool)
      @ DifferentiationInterfaceTest C:\Users\timho\.julia\dev\DifferentiationInterface\DifferentiationInterfaceTest\src\tests\correctness_eval.jl:826
   [27] macro expansion
      @ C:\Users\timho\.julia\dev\DifferentiationInterface\DifferentiationInterfaceTest\src\test_differentiation.jl:157 [inlined]
   [28] macro expansion
      @ C:\Users\timho\.julia\juliaup\julia-1.10.10+0.x64.w64.mingw32\share\julia\stdlib\v1.10\Test\src\Test.jl:1577 [inlined]
   [29] macro expansion
      @ C:\Users\timho\.julia\dev\DifferentiationInterface\DifferentiationInterfaceTest\src\test_differentiation.jl:157 [inlined]
   [30] macro expansion
      @ C:\Users\timho\.julia\juliaup\julia-1.10.10+0.x64.w64.mingw32\share\julia\stdlib\v1.10\Test\src\Test.jl:1669 [inlined]
   [31] macro expansion
      @ C:\Users\timho\.julia\dev\DifferentiationInterface\DifferentiationInterfaceTest\src\test_differentiation.jl:133 [inlined]
   [32] macro expansion
      @ C:\Users\timho\.julia\juliaup\julia-1.10.10+0.x64.w64.mingw32\share\julia\stdlib\v1.10\Test\src\Test.jl:1669 [inlined]
   [33] macro expansion
      @ C:\Users\timho\.julia\dev\DifferentiationInterface\DifferentiationInterfaceTest\src\test_differentiation.jl:131 [inlined]
   [34] macro expansion
      @ C:\Users\timho\.julia\juliaup\julia-1.10.10+0.x64.w64.mingw32\share\julia\stdlib\v1.10\Test\src\Test.jl:1669 [inlined]
   [35] macro expansion
      @ C:\Users\timho\.julia\dev\DifferentiationInterface\DifferentiationInterfaceTest\src\test_differentiation.jl:128 [inlined]
   [36] macro expansion
      @ C:\Users\timho\.julia\juliaup\julia-1.10.10+0.x64.w64.mingw32\share\julia\stdlib\v1.10\Test\src\Test.jl:1577 [inlined]
   [37] test_differentiation(backends::Vector{AbstractADType}, scenarios::Vector{Scenario}; testset_name::Nothing, correctness::Bool, type_stability::Symbol, allocations::Symbol, benchmark::Symbol, excluded::Vector{Symbol}, detailed::Bool, logging::Bool, isapprox::typeof(isapprox), atol::Int64, rtol::Float64, scenario_intact::Bool, sparsity::Bool, reprepare::Bool, ignored_modules::Nothing, function_filter::DifferentiationInterfaceTest.var"#736#743", skip_allocations::Bool, count_calls::Bool, benchmark_test::Bool, benchmark_seconds::Int64, benchmark_aggregation::typeof(minimum), adaptive_batchsize::Bool)
      @ DifferentiationInterfaceTest C:\Users\timho\.julia\dev\DifferentiationInterface\DifferentiationInterfaceTest\src\test_differentiation.jl:128
   [38] macro expansion
      @ REPL[14]:2 [inlined]
   [39] macro expansion
      @ C:\Users\timho\.julia\juliaup\julia-1.10.10+0.x64.w64.mingw32\share\julia\stdlib\v1.10\Test\src\Test.jl:1577 [inlined]
   [40] top-level scope
      @ REPL[14]:2
   [41] eval
      @ .\boot.jl:385 [inlined]
   [42] eval_user_input(ast::Any, backend::REPL.REPLBackend, mod::Module)
      @ REPL C:\Users\timho\.julia\juliaup\julia-1.10.10+0.x64.w64.mingw32\share\julia\stdlib\v1.10\REPL\src\REPL.jl:150
   [43] repl_backend_loop(backend::REPL.REPLBackend, get_module::Function)
      @ REPL C:\Users\timho\.julia\juliaup\julia-1.10.10+0.x64.w64.mingw32\share\julia\stdlib\v1.10\REPL\src\REPL.jl:246
   [44] start_repl_backend(backend::REPL.REPLBackend, consumer::Any; get_module::Function)
      @ REPL C:\Users\timho\.julia\juliaup\julia-1.10.10+0.x64.w64.mingw32\share\julia\stdlib\v1.10\REPL\src\REPL.jl:231
   [45] run_repl(repl::REPL.AbstractREPL, consumer::Any; backend_on_current_task::Bool, backend::Any)
      @ REPL C:\Users\timho\.julia\juliaup\julia-1.10.10+0.x64.w64.mingw32\share\julia\stdlib\v1.10\REPL\src\REPL.jl:389
   [46] run_repl(repl::REPL.AbstractREPL, consumer::Any)
      @ REPL C:\Users\timho\.julia\juliaup\julia-1.10.10+0.x64.w64.mingw32\share\julia\stdlib\v1.10\REPL\src\REPL.jl:375
   [47] (::Base.var"#1016#1018"{Bool, Bool, Bool})(REPL::Module)
      @ Base .\client.jl:437
   [48] #invokelatest#2
      @ .\essentials.jl:892 [inlined]
   [49] invokelatest
      @ .\essentials.jl:889 [inlined]
   [50] run_main_repl(interactive::Bool, quiet::Bool, banner::Bool, history_file::Bool, color_set::Bool)
      @ Base .\client.jl:421
   [51] exec_options(opts::Base.JLOptions)
      @ Base .\client.jl:338

This is with

diff --git a/DifferentiationInterface/ext/DifferentiationInterfaceEnzymeExt/utils.jl b/DifferentiationInterface/ext/DifferentiationInterfaceEnzymeExt/utils.jl
index 57c654cf..f1b84dc4 100644
--- a/DifferentiationInterface/ext/DifferentiationInterfaceEnzymeExt/utils.jl
+++ b/DifferentiationInterface/ext/DifferentiationInterfaceEnzymeExt/utils.jl
@@ -87,7 +87,7 @@ function _shadow(
 end

 function _shadow(
-        backend::AutoEnzyme{M, <:Union{Const, Nothing}},
+        backend::AutoEnzyme{M, <:Union{Const, AnyDuplicated, Nothing}},
         ::Mode,
         ::Val{B},
         c_wrapped::DI.FunctionContext,

which by eye seems to mimic manually implementing it based on DI.ConstantOrCache.

[timholy]

Also,

backend2 = SecondOrder(
                 AutoEnzyme(; mode = Enzyme.Reverse), AutoEnzyme(; mode = Enzyme.Forward, function_annotation=Enzyme.Duplicated)
             )

gives a SIGABRT on Julia 1.10.

[gdalle]

I suspect that DI is not the package causing the SIGABRT, so there is probably a deeper Enzyme issue there (at least a more friendly error to throw). Trouble is, Enzyme devs don't accept DI-based MWEs, so I'll have to minimize it, and second-order is hard to minimize precisely because DI does a lot of things for us.
In a nutshell, I see the issue, but it might take some sweating on my end to fix

[gdalle]

Okay I figured out some more. When you use DI in the second order, it doesn't differentiate your provided function L twice: it differentiates L once and shuffled_gradient once. However, it applies the backend's function_annotation to both indiscriminately, even though shuffled_gradient contains no differentiable data. An immediate fix for your case is to use

backend = SecondOrder(
    AutoEnzyme(; mode=nothing, function_annotation=Enzyme.Const),  # for shuffled_gradient
    AutoEnzyme(; mode=nothing, function_annotation=Enzyme.Duplicated),  # for L
)

which gives you another Enzyme error:

julia> hessian(L, backend, tb0)   # errors
┌ Warning: Using fallback BLAS replacements for (["dger_64_"]), performance may be degraded
└ @ Enzyme.Compiler ~/.julia/packages/Enzyme/OhLFb/src/compiler.jl:5154
┌ Warning: Using fallback BLAS replacements for (["dger_64_"]), performance may be degraded
└ @ Enzyme.Compiler ~/.julia/packages/Enzyme/OhLFb/src/compiler.jl:5154
ERROR: EnzymeNoDerivativeError: 
No forward mode derivative found for jl_field_isdefined_checked

I think Enzyme devs should be able to do something about that. Meanwhile, I'll try to improve activity annotation on functions inside DI to avoid this trap.

[gdalle]

The PR #943 should contain the fix on the DI side