DifferentiationInterface not generating correct shadow

Specific Domains Optimization (Mathematical)
, ,
1

Hello!
I am trying to play with the Sophia optimization algorthim. I am storing my parameters as a ComponentArray since some parameters are scalar and others are vectors. When I use a first order method (Adam) I have no issue, but when I use a SecondOrder adtype I get an error. Below is a MWE that maintains the structure of my real code. The addition u_buffer is for hacking together TBPTT.

using Enzyme
using SciMLSensitivity
using Optimisers
using DiffEqFlux
using ComponentArrays
using OrdinaryDiffEq
using Optimization
using DifferentiationInterface

function test_dynamics!(du, u, p, t)
    a = p.a; b = p.b
    du[1] = a[1] * u[1] + b * u[2]
    du[2] = a[2] * u[2] + b * u[1]
end
function dist_loss!(p, u_buffer, u_initial, prob,dt)
    x0 = 1f0; y0 = 2f0
    sol = solve(prob, Tsit5(),u0=u_initial, p=p, saveat=dt, sensealg = GaussAdjoint(autojacvec=EnzymeVJP()))#, sensealg=SensitivityADPassThrough())
    dists=zeros(Float32, length(sol.u))
    for k in 1:length(sol.u)
        dist_sq = (sol.u[k][1]-x0)^2 + (sol.u[k][2]-y0)^2
        dists[k] = dist_sq
    end
    total_loss = sum(dists) * Float32(1e-6)
    u_buffer .= sol.u[end]
    return Float32(total_loss)
end

params = ComponentArray(a=[1f0, 2f0], b=3f0)
u0 = [1f0, 1f0]; u0_buff = similar(u0)
prob = ODEProblem(test_dynamics!, u0, (0f0, 1f0), params)
loss_fn = (p, hp) -> dist_loss!(p, u0_buff, u0, prob, 0.1f0)

adtype = Optimization.AutoEnzyme(; mode=Enzyme.set_runtime_activity(Enzyme.Reverse))
second_ad = DifferentiationInterface.SecondOrder(adtype, adtype)  # inner & outer

optf = Optimization.OptimizationFunction(loss_fn, second_ad)
optprob = Optimization.OptimizationProblem(optf, params)

result = Optimization.solve( optprob, Optimization.Sophia(), maxiters = 10)

and this is the error

MethodError: no method matching Duplicated(::ComponentVector{Float32, Vector{Float32}, Tuple{Axis{(a = ViewAxis(1:2, Shaped1DAxis((2,))), b = 3)}}}, ::Vector{Float32})
The type Duplicated exists, but no method is defined for this combination of argument types when trying to construct it.
Closest candidates are:
Duplicated(::T1, !Matched::T1) where T1
@ EnzymeCore ~/.julia/packages/EnzymeCore/lmG5F/src/EnzymeCore.jl:68
Duplicated(::T1, !Matched::T1, !Matched::Bool) where T1
@ EnzymeCore ~/.julia/packages/EnzymeCore/lmG5F/src/EnzymeCore.jl:68
Duplicated(!Matched::AbstractLuxLayer, !Matched::AbstractLuxLayer)
@ LuxCoreEnzymeCoreExt ~/.julia/packages/LuxCore/SN4dl/ext/LuxCoreEnzymeCoreExt.jl:23
Stacktrace:
[1] value_and_pullback!(::typeof(DifferentiationInterface.shuffled_gradient!), ::ComponentVector{Float32, Vector{Float32}, Tuple{Axis{(a = ViewAxis(1:2, Shaped1DAxis((2,))), b = 3)}}}, ::Tuple{ComponentVector{Float32, Vector{Float32}, Tuple{Axis{(a = ViewAxis(1:2, Shaped1DAxis((2,))), b = 3)}}}}, ::DifferentiationInterfaceEnzymeExt.EnzymeReverseTwoArgPullbackPrep{Tuple{typeof(DifferentiationInterface.shuffled_gradient!), ComponentVector{Float32, Vector{Float32}, Tuple{Axis{(a = ViewAxis(1:2, Shaped1DAxis((2,))), b = 3)}}}, AutoEnzyme{ReverseMode{false, true, false, FFIABI, false, false}, Nothing}, ComponentVector{Float32, Vector{Float32}, Tuple{Axis{(a = ViewAxis(1:2, Shaped1DAxis((2,))), b = 3)}}}, Tuple{Vector{Float32}}, Tuple{DifferentiationInterface.FunctionContext{var"#981#982"}, DifferentiationInterface.Constant{AutoEnzyme{ReverseMode{false, true, false, FFIABI, false, false}, Nothing}}, DifferentiationInterface.Constant{DifferentiationInterface.Rewrap{1, Tuple{typeof(DifferentiationInterface.constant_maker)}}}, DifferentiationInterface.Constant{SciMLBase.NullParameters}}}, Nothing, NTuple{4, Nothing}, Tuple{Vector{Float32}}}, ::AutoEnzyme{ReverseMode{false, true, false, FFIABI, false, false}, Nothing}, ::ComponentVector{Float32, Vector{Float32}, Tuple{Axis{(a = ViewAxis(1:2, Shaped1DAxis((2,))), b = 3)}}}, ::Tuple{Vector{Float32}}, ::DifferentiationInterface.FunctionContext{var"#981#982"}, ::DifferentiationInterface.Constant{AutoEnzyme{ReverseMode{false, true, false, FFIABI, false, false}, Nothing}}, ::DifferentiationInterface.Constant{DifferentiationInterface.Rewrap{1, Tuple{typeof(DifferentiationInterface.constant_maker)}}}, ::DifferentiationInterface.Constant{SciMLBase.NullParameters})
@ DifferentiationInterfaceEnzymeExt ~/.julia/packages/DifferentiationInterface/zJHX8/ext/DifferentiationInterfaceEnzymeExt/reverse_twoarg.jl:140
[2] pullback!(::typeof(DifferentiationInterface.shuffled_gradient!), ::ComponentVector{Float32, Vector{Float32}, Tuple{Axis{(a = ViewAxis(1:2, Shaped1DAxis((2,))), b = 3)}}}, ::Tuple{ComponentVector{Float32, Vector{Float32}, Tuple{Axis{(a = ViewAxis(1:2, Shaped1DAxis((2,))), b = 3)}}}}, ::DifferentiationInterfaceEnzymeExt.EnzymeReverseTwoArgPullbackPrep{Tuple{typeof(DifferentiationInterface.shuffled_gradient!), ComponentVector{Float32, Vector{Float32}, Tuple{Axis{(a = ViewAxis(1:2, Shaped1DAxis((2,))), b = 3)}}}, AutoEnzyme{ReverseMode{false, true, false, FFIABI, false, false}, Nothing}, ComponentVector{Float32, Vector{Float32}, Tuple{Axis{(a = ViewAxis(1:2, Shaped1DAxis((2,))), b = 3)}}}, Tuple{Vector{Float32}}, Tuple{DifferentiationInterface.FunctionContext{var"#981#982"}, DifferentiationInterface.Constant{AutoEnzyme{ReverseMode{false, true, false, FFIABI, false, false}, Nothing}}, DifferentiationInterface.Constant{DifferentiationInterface.Rewrap{1, Tuple{typeof(DifferentiationInterface.constant_maker)}}}, DifferentiationInterface.Constant{SciMLBase.NullParameters}}}, Nothing, NTuple{4, Nothing}, Tuple{Vector{Float32}}}, ::AutoEnzyme{ReverseMode{false, true, false, FFIABI, false, false}, Nothing}, ::ComponentVector{Float32, Vector{Float32}, Tuple{Axis{(a = ViewAxis(1:2, Shaped1DAxis((2,))), b = 3)}}}, ::Tuple{Vector{Float32}}, ::DifferentiationInterface.FunctionContext{var"#981#982"}, ::DifferentiationInterface.Constant{AutoEnzyme{ReverseMode{false, true, false, FFIABI, false, false}, Nothing}}, ::DifferentiationInterface.Constant{DifferentiationInterface.Rewrap{1, Tuple{typeof(DifferentiationInterface.constant_maker)}}}, ::DifferentiationInterface.Constant{SciMLBase.NullParameters})
@ DifferentiationInterface ~/.julia/packages/DifferentiationInterface/zJHX8/src/first_order/pullback.jl:557
[3] _hvp_aux!(::DifferentiationInterface.InPlaceSupported, f::var"#981#982", tg::Tuple{ComponentVector{Float32, Vector{Float32}, Tuple{Axis{(a = ViewAxis(1:2, Shaped1DAxis((2,))), b = 3)}}}}, prep::DifferentiationInterface.ReverseOverReverseHVPPrep{Tuple{var"#981#982", SecondOrder{AutoEnzyme{ReverseMode{false, true, false, FFIABI, false, false}, Nothing}, AutoEnzyme{ReverseMode{false, true, false, FFIABI, false, false}, Nothing}}, ComponentVector{Float32, Vector{Float32}, Tuple{Axis{(a = ViewAxis(1:2, Shaped1DAxis((2,))), b = 3)}}}, Tuple{Vector{Float32}}, Tuple{DifferentiationInterface.Constant{SciMLBase.NullParameters}}}, ComponentVector{Float32, Vector{Float32}, Tuple{Axis{(a = ViewAxis(1:2, Shaped1DAxis((2,))), b = 3)}}}, DifferentiationInterfaceEnzymeExt.EnzymeReverseOneArgPullbackPrep{Tuple{typeof(DifferentiationInterface.shuffled_gradient), AutoEnzyme{ReverseMode{false, true, false, FFIABI, false, false}, Nothing}, ComponentVector{Float32, Vector{Float32}, Tuple{Axis{(a = ViewAxis(1:2, Shaped1DAxis((2,))), b = 3)}}}, Tuple{Vector{Float32}}, Tuple{DifferentiationInterface.FunctionContext{var"#981#982"}, DifferentiationInterface.Constant{AutoEnzyme{ReverseMode{false, true, false, FFIABI, false, false}, Nothing}}, DifferentiationInterface.Constant{DifferentiationInterface.Rewrap{1, Tuple{typeof(DifferentiationInterface.constant_maker)}}}, DifferentiationInterface.Constant{SciMLBase.NullParameters}}}, Nothing, NTuple{4, Nothing}, ComponentVector{Float32, Vector{Float32}, Tuple{Axis{(a = ViewAxis(1:2, Shaped1DAxis((2,))), b = 3)}}}}, DifferentiationInterfaceEnzymeExt.EnzymeReverseTwoArgPullbackPrep{Tuple{typeof(DifferentiationInterface.shuffled_gradient!), ComponentVector{Float32, Vector{Float32}, Tuple{Axis{(a = ViewAxis(1:2, Shaped1DAxis((2,))), b = 3)}}}, AutoEnzyme{ReverseMode{false, true, false, FFIABI, false, false}, Nothing}, ComponentVector{Float32, Vector{Float32}, Tuple{Axis{(a = ViewAxis(1:2, Shaped1DAxis((2,))), b = 3)}}}, Tuple{Vector{Float32}}, Tuple{DifferentiationInterface.FunctionContext{var"#981#982"}, DifferentiationInterface.Constant{AutoEnzyme{ReverseMode{false, true, false, FFIABI, false, false}, Nothing}}, DifferentiationInterface.Constant{DifferentiationInterface.Rewrap{1, Tuple{typeof(DifferentiationInterface.constant_maker)}}}, DifferentiationInterface.Constant{SciMLBase.NullParameters}}}, Nothing, NTuple{4, Nothing}, Tuple{Vector{Float32}}}}, backend::SecondOrder{AutoEnzyme{ReverseMode{false, true, false, FFIABI, false, false}, Nothing}, AutoEnzyme{ReverseMode{false, true, false, FFIABI, false, false}, Nothing}}, x::ComponentVector{Float32, Vector{Float32}, Tuple{Axis{(a = ViewAxis(1:2, Shaped1DAxis((2,))), b = 3)}}}, tx::Tuple{Vector{Float32}}, contexts::DifferentiationInterface.Constant{SciMLBase.NullParameters})
@ DifferentiationInterface ~/.julia/packages/DifferentiationInterface/zJHX8/src/second_order/hvp.jl:733
[4] hvp!(f::var"#981#982", tg::Tuple{ComponentVector{Float32, Vector{Float32}, Tuple{Axis{(a = ViewAxis(1:2, Shaped1DAxis((2,))), b = 3)}}}}, prep::DifferentiationInterface.ReverseOverReverseHVPPrep{Tuple{var"#981#982", SecondOrder{AutoEnzyme{ReverseMode{false, true, false, FFIABI, false, false}, Nothing}, AutoEnzyme{ReverseMode{false, true, false, FFIABI, false, false}, Nothing}}, ComponentVector{Float32, Vector{Float32}, Tuple{Axis{(a = ViewAxis(1:2, Shaped1DAxis((2,))), b = 3)}}}, Tuple{Vector{Float32}}, Tuple{DifferentiationInterface.Constant{SciMLBase.NullParameters}}}, ComponentVector{Float32, Vector{Float32}, Tuple{Axis{(a = ViewAxis(1:2, Shaped1DAxis((2,))), b = 3)}}}, DifferentiationInterfaceEnzymeExt.EnzymeReverseOneArgPullbackPrep{Tuple{typeof(DifferentiationInterface.shuffled_gradient), AutoEnzyme{ReverseMode{false, true, false, FFIABI, false, false}, Nothing}, ComponentVector{Float32, Vector{Float32}, Tuple{Axis{(a = ViewAxis(1:2, Shaped1DAxis((2,))), b = 3)}}}, Tuple{Vector{Float32}}, Tuple{DifferentiationInterface.FunctionContext{var"#981#982"}, DifferentiationInterface.Constant{AutoEnzyme{ReverseMode{false, true, false, FFIABI, false, false}, Nothing}}, DifferentiationInterface.Constant{DifferentiationInterface.Rewrap{1, Tuple{typeof(DifferentiationInterface.constant_maker)}}}, DifferentiationInterface.Constant{SciMLBase.NullParameters}}}, Nothing, NTuple{4, Nothing}, ComponentVector{Float32, Vector{Float32}, Tuple{Axis{(a = ViewAxis(1:2, Shaped1DAxis((2,))), b = 3)}}}}, DifferentiationInterfaceEnzymeExt.EnzymeReverseTwoArgPullbackPrep{Tuple{typeof(DifferentiationInterface.shuffled_gradient!), ComponentVector{Float32, Vector{Float32}, Tuple{Axis{(a = ViewAxis(1:2, Shaped1DAxis((2,))), b = 3)}}}, AutoEnzyme{ReverseMode{false, true, false, FFIABI, false, false}, Nothing}, ComponentVector{Float32, Vector{Float32}, Tuple{Axis{(a = ViewAxis(1:2, Shaped1DAxis((2,))), b = 3)}}}, Tuple{Vector{Float32}}, Tuple{DifferentiationInterface.FunctionContext{var"#981#982"}, DifferentiationInterface.Constant{AutoEnzyme{ReverseMode{false, true, false, FFIABI, false, false}, Nothing}}, DifferentiationInterface.Constant{DifferentiationInterface.Rewrap{1, Tuple{typeof(DifferentiationInterface.constant_maker)}}}, DifferentiationInterface.Constant{SciMLBase.NullParameters}}}, Nothing, NTuple{4, Nothing}, Tuple{Vector{Float32}}}}, backend::SecondOrder{AutoEnzyme{ReverseMode{false, true, false, FFIABI, false, false}, Nothing}, AutoEnzyme{ReverseMode{false, true, false, FFIABI, false, false}, Nothing}}, x::ComponentVector{Float32, Vector{Float32}, Tuple{Axis{(a = ViewAxis(1:2, Shaped1DAxis((2,))), b = 3)}}}, tx::Tuple{Vector{Float32}}, contexts::DifferentiationInterface.Constant{SciMLBase.NullParameters})
@ DifferentiationInterface ~/.julia/packages/DifferentiationInterface/zJHX8/src/second_order/hvp.jl:713
[5] hvp!(f::var"#981#982", tg::Tuple{ComponentVector{Float32, Vector{Float32}, Tuple{Axis{(a = ViewAxis(1:2, Shaped1DAxis((2,))), b = 3)}}}}, backend::SecondOrder{AutoEnzyme{ReverseMode{false, true, false, FFIABI, false, false}, Nothing}, AutoEnzyme{ReverseMode{false, true, false, FFIABI, false, false}, Nothing}}, x::ComponentVector{Float32, Vector{Float32}, Tuple{Axis{(a = ViewAxis(1:2, Shaped1DAxis((2,))), b = 3)}}}, tx::Tuple{Vector{Float32}}, contexts::DifferentiationInterface.Constant{SciMLBase.NullParameters})
@ DifferentiationInterface ~/.julia/packages/DifferentiationInterface/zJHX8/src/second_order/hvp.jl:89
[6] (::OptimizationBase.var"#hv!#24"{SciMLBase.NullParameters, DifferentiationInterface.ReverseOverReverseHVPPrep{Nothing, ComponentVector{Float32, Vector{Float32}, Tuple{Axis{(a = ViewAxis(1:2, Shaped1DAxis((2,))), b = 3)}}}, DifferentiationInterfaceEnzymeExt.EnzymeReverseOneArgPullbackPrep{Nothing, Nothing, NTuple{4, Nothing}, ComponentVector{Float32, Vector{Float32}, Tuple{Axis{(a = ViewAxis(1:2, Shaped1DAxis((2,))), b = 3)}}}}, DifferentiationInterfaceEnzymeExt.EnzymeReverseTwoArgPullbackPrep{Nothing, Nothing, NTuple{4, Nothing}, Tuple{Vector{Float32}}}}, OptimizationFunction{true, SecondOrder{AutoEnzyme{ReverseMode{false, true, false, FFIABI, false, false}, Nothing}, AutoEnzyme{ReverseMode{false, true, false, FFIABI, false, false}, Nothing}}, var"#981#982", Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED_NO_TIME), Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, SecondOrder{AutoEnzyme{ReverseMode{false, true, false, FFIABI, false, false}, Nothing}, AutoEnzyme{ReverseMode{false, true, false, FFIABI, false, false}, Nothing}}})(H::ComponentVector{Float32, Vector{Float32}, Tuple{Axis{(a = ViewAxis(1:2, Shaped1DAxis((2,))), b = 3)}}}, θ::ComponentVector{Float32, Vector{Float32}, Tuple{Axis{(a = ViewAxis(1:2, Shaped1DAxis((2,))), b = 3)}}}, v::Vector{Float32}, p::SciMLBase.NullParameters)
@ OptimizationBase ~/.julia/packages/OptimizationBase/1tTb9/src/OptimizationDIExt.jl:105
[7] __solve(cache::OptimizationCache{OptimizationFunction{true, AutoEnzyme{ReverseMode{false, true, false, FFIABI, false, false}, Nothing}, var"#981#982", OptimizationBase.var"#grad#16"{SciMLBase.NullParameters, OptimizationFunction{true, SecondOrder{AutoEnzyme{ReverseMode{false, true, false, FFIABI, false, false}, Nothing}, AutoEnzyme{ReverseMode{false, true, false, FFIABI, false, false}, Nothing}}, var"#981#982", Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED_NO_TIME), Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, AutoEnzyme{ReverseMode{false, true, false, FFIABI, false, false}, Nothing}}, OptimizationBase.var"#fg!#18"{SciMLBase.NullParameters, OptimizationFunction{true, SecondOrder{AutoEnzyme{ReverseMode{false, true, false, FFIABI, false, false}, Nothing}, AutoEnzyme{ReverseMode{false, true, false, FFIABI, false, false}, Nothing}}, var"#981#982", Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED_NO_TIME), Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, AutoEnzyme{ReverseMode{false, true, false, FFIABI, false, false}, Nothing}}, OptimizationBase.var"#hess#20"{SciMLBase.NullParameters, OptimizationFunction{true, SecondOrder{AutoEnzyme{ReverseMode{false, true, false, FFIABI, false, false}, Nothing}, AutoEnzyme{ReverseMode{false, true, false, FFIABI, false, false}, Nothing}}, var"#981#982", Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED_NO_TIME), Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, SecondOrder{AutoEnzyme{ReverseMode{false, true, false, FFIABI, false, false}, Nothing}, AutoEnzyme{ReverseMode{false, true, false, FFIABI, false, false}, Nothing}}}, Nothing, OptimizationBase.var"#hv!#24"{SciMLBase.NullParameters, DifferentiationInterface.ReverseOverReverseHVPPrep{Nothing, ComponentVector{Float32, Vector{Float32}, Tuple{Axis{(a = ViewAxis(1:2, Shaped1DAxis((2,))), b = 3)}}}, DifferentiationInterfaceEnzymeExt.EnzymeReverseOneArgPullbackPrep{Nothing, Nothing, NTuple{4, Nothing}, ComponentVector{Float32, Vector{Float32}, Tuple{Axis{(a = ViewAxis(1:2, Shaped1DAxis((2,))), b = 3)}}}}, DifferentiationInterfaceEnzymeExt.EnzymeReverseTwoArgPullbackPrep{Nothing, Nothing, NTuple{4, Nothing}, Tuple{Vector{Float32}}}}, OptimizationFunction{true, SecondOrder{AutoEnzyme{ReverseMode{false, true, false, FFIABI, false, false}, Nothing}, AutoEnzyme{ReverseMode{false, true, false, FFIABI, false, false}, Nothing}}, var"#981#982", Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED_NO_TIME), Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, SecondOrder{AutoEnzyme{ReverseMode{false, true, false, FFIABI, false, false}, Nothing}, AutoEnzyme{ReverseMode{false, true, false, FFIABI, false, false}, Nothing}}}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED_NO_TIME), Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, OptimizationBase.ReInitCache{ComponentVector{Float32, Vector{Float32}, Tuple{Axis{(a = ViewAxis(1:2, Shaped1DAxis((2,))), b = 3)}}}, SciMLBase.NullParameters}, Nothing, Nothing, Nothing, Nothing, Nothing, Optimization.Sophia, Bool, var"#973#974", Nothing})
@ Optimization ~/.julia/packages/Optimization/hfEs4/src/sophia.jl:108
[8] solve!(cache::OptimizationCache{OptimizationFunction{true, AutoEnzyme{ReverseMode{false, true, false, FFIABI, false, false}, Nothing}, var"#981#982", OptimizationBase.var"#grad#16"{SciMLBase.NullParameters, OptimizationFunction{true, SecondOrder{AutoEnzyme{ReverseMode{false, true, false, FFIABI, false, false}, Nothing}, AutoEnzyme{ReverseMode{false, true, false, FFIABI, false, false}, Nothing}}, var"#981#982", Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED_NO_TIME), Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, AutoEnzyme{ReverseMode{false, true, false, FFIABI, false, false}, Nothing}}, OptimizationBase.var"#fg!#18"{SciMLBase.NullParameters, OptimizationFunction{true, SecondOrder{AutoEnzyme{ReverseMode{false, true, false, FFIABI, false, false}, Nothing}, AutoEnzyme{ReverseMode{false, true, false, FFIABI, false, false}, Nothing}}, var"#981#982", Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED_NO_TIME), Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, AutoEnzyme{ReverseMode{false, true, false, FFIABI, false, false}, Nothing}}, OptimizationBase.var"#hess#20"{SciMLBase.NullParameters, OptimizationFunction{true, SecondOrder{AutoEnzyme{ReverseMode{false, true, false, FFIABI, false, false}, Nothing}, AutoEnzyme{ReverseMode{false, true, false, FFIABI, false, false}, Nothing}}, var"#981#982", Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED_NO_TIME), Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, SecondOrder{AutoEnzyme{ReverseMode{false, true, false, FFIABI, false, false}, Nothing}, AutoEnzyme{ReverseMode{false, true, false, FFIABI, false, false}, Nothing}}}, Nothing, OptimizationBase.var"#hv!#24"{SciMLBase.NullParameters, DifferentiationInterface.ReverseOverReverseHVPPrep{Nothing, ComponentVector{Float32, Vector{Float32}, Tuple{Axis{(a = ViewAxis(1:2, Shaped1DAxis((2,))), b = 3)}}}, DifferentiationInterfaceEnzymeExt.EnzymeReverseOneArgPullbackPrep{Nothing, Nothing, NTuple{4, Nothing}, ComponentVector{Float32, Vector{Float32}, Tuple{Axis{(a = ViewAxis(1:2, Shaped1DAxis((2,))), b = 3)}}}}, DifferentiationInterfaceEnzymeExt.EnzymeReverseTwoArgPullbackPrep{Nothing, Nothing, NTuple{4, Nothing}, Tuple{Vector{Float32}}}}, OptimizationFunction{true, SecondOrder{AutoEnzyme{ReverseMode{false, true, false, FFIABI, false, false}, Nothing}, AutoEnzyme{ReverseMode{false, true, false, FFIABI, false, false}, Nothing}}, var"#981#982", Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED_NO_TIME), Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, SecondOrder{AutoEnzyme{ReverseMode{false, true, false, FFIABI, false, false}, Nothing}, AutoEnzyme{ReverseMode{false, true, false, FFIABI, false, false}, Nothing}}}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED_NO_TIME), Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, OptimizationBase.ReInitCache{ComponentVector{Float32, Vector{Float32}, Tuple{Axis{(a = ViewAxis(1:2, Shaped1DAxis((2,))), b = 3)}}}, SciMLBase.NullParameters}, Nothing, Nothing, Nothing, Nothing, Nothing, Optimization.Sophia, Bool, var"#973#974", Nothing})
@ SciMLBase ~/.julia/packages/SciMLBase/rvXrA/src/solve.jl:226
[9] solve(::OptimizationProblem{true, OptimizationFunction{true, SecondOrder{AutoEnzyme{ReverseMode{false, true, false, FFIABI, false, false}, Nothing}, AutoEnzyme{ReverseMode{false, true, false, FFIABI, false, false}, Nothing}}, var"#981#982", Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED_NO_TIME), Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, ComponentVector{Float32, Vector{Float32}, Tuple{Axis{(a = ViewAxis(1:2, Shaped1DAxis((2,))), b = 3)}}}, SciMLBase.NullParameters, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, @Kwargs{}}, ::Optimization.Sophia; kwargs::@Kwargs{maxiters::Int64, callback::var"#973#974"})
@ SciMLBase ~/.julia/packages/SciMLBase/rvXrA/src/solve.jl:128

I would’ve thought that the higher level functions would’ve created a ComponentArray for the shadow.

Am I doing something wrong or is this a bug?

Thanks for your time and effort

2

@ChrisRackauckas this seems like a bug in OptimizationBase.jl, Optimization shouldn’t use DI for Enzyme?

3

Judging by this frame of the stack trace, DI.hvp! is called with x::ComponentVector but tx::Tuple{Vector}, which suggests that somewhere in the Optimization stack the type of the input is not respected?

If I understand correctly, OptimizationBase does have Enzyme-specific implementations for backend::AutoEnzyme, but here the user supplied backend::DI.SecondOrder, which is why those dispatches were not touched.

As a side remark, @alexl123 do you have any reason to choose reverse-over-reverse for your second-order backend? You’d probably be better off with

forward_adtype = Optimization.AutoEnzyme(; mode=Enzyme.set_runtime_activity(Enzyme.Forward))
reverse_adtype = Optimization.AutoEnzyme(; mode=Enzyme.set_runtime_activity(Enzyme.Reverse))
second_ad = DifferentiationInterface.SecondOrder(forward_adtype, reverse_adtype)  # inner & outer
4

After taking a closer look at Sophia, maybe the problem is here?

It seems to me that this generates a Vector for u regardless of the type of θ (in this case ComponentVector). Other autodiff backends may not raise an eyebrow but Enzyme’s Duplicated system enforces type consistency between primal and shadow, so that’s where you see the error but that’s not where it originates.

A simple way to preserve type information would be to use randn! and initialize u properly at the beginning:

julia> using ComponentArrays, Random

julia> x = ComponentVector(a=[1.0, 2.0])
ComponentVector{Float64}(a = [1.0, 2.0])

julia> randn!(x)
ComponentVector{Float64}(a = [0.07972696817241108, -0.935076728780384])
5

Thank you!

I am still very new to AD, so no particular reason for reverse over reverse. I just want to use Sophia, I only explicitly created the SecondOrder AD type because I got a warning suggesting I do that.

Should I modify sophia.jl locally to implement that change?

6

Ah I see, equivalently here I would recommend just passing in

adtype = Optimization.AutoEnzyme(; mode=Enzyme.set_runtime_activity(Enzyme.Reverse))

directly instead of using the DI.secondorder (which will force it to use the official native Enzyme implementation directly).

7

At the moment this will yield the same behavior for second-order optimization algorithms inside Optimization.jl.

julia> result = Optimization.solve(optprob, Optimization.Sophia(); maxiters=10)
┌ Warning: The selected optimization algorithm requires second order derivatives, but `SecondOrder` ADtype was not provided. 
│         So a `SecondOrder` with AutoEnzyme(mode=ReverseMode{false, true, false, FFIABI, false, false}()) for both inner and outer will be created, this can be suboptimal and not work in some cases so 
│         an explicit `SecondOrder` ADtype is recommended.
â”” @ OptimizationBase ~/.julia/packages/OptimizationBase/Lc8sB/src/cache.jl:49
ERROR: MethodError: no method matching Duplicated(::ComponentVector{Float32, Vector{Float32}, Tuple{Axis{…}}}, ::Vector{Float32})
The type `Duplicated` exists, but no method is defined for this combination of argument types when trying to construct it.

See also:

8

Oh hm, that definitely seems like an Optimization.jl bug then, passing in an incorrect type

1 Like
9

Even better, you could try to contribute a fix!
While you’re at it, it would also be nice to be able to pass an RNG to Sophia for reproducibility

10

Should be handled on latest release now.