inner second order and outer optimization AD?

[zahachtah]

Hi, I have a case where I want to use DiffEqFlux.jl collocation example:

https://docs.sciml.ai/DiffEqFlux/stable/examples/collocation/

but instead of just training a neural net, I need the first and second derivative with respect to one input.

I am trying to solve it but I keep getting errors in zygote or forwarddiff backend compabilities which are a bit over my head.

I was going to ask if there are any things I should consider when doing this that might help me. I can also give more detailed code but here are some relevant parts (note that I need to use the value_derivative_and_second_derivative! which I think is not available for AutoZygote?):

so in the inner function I use

backend = AutoForwardDiff()

u0 = MVector{3,Float64}(rand(3))
prep = prepare_second_derivative(g, backend, u0[2], Constant(u0), Constant(0.0))

# Reusable buffers (scalars)
der = zeros(Float64, 1)   # df/du2
der2 = zeros(Float64, 1)  # d2f/du2^2

function rhs!(du, u, p, t) # the differential equation to run the simulation with
    # DI computes value + 1st + 2nd derivative wrt x in one call
    f, _, _ = value_derivative_and_second_derivative!(
        g, der, der2, prep, backend, u[2], Constant(u), Constant(t)
    )
    f1 = der[1]
    f2 = der2[1]
    # Moment-closure dynamics (M0..M2) with M3/M4 from drivers
    m0, m1, m2 = u
    drv_vec = driver_vec(t, eltype(u))
    m3 = drv_vec[end-1]
    m4 = drv_vec[end]
    d0 = f * m0 + 0.5 * f2 * m2
    d1 = f1 * m2 + 0.5 * f2 * m3
    d2 = f1 * m3 + 0.5 * f2 * (m4 - m2^2)
    # d3/d4 are provided by drivers (closure)
    #println((d0, d1, d2))
    du[1] = d0
    du[2] = d1
    du[3] = d2
    return nothing
end

and then I need to train the parameters to minimize the loss of the simulation with the actual data I have.

adtype = Optimization.AutoZygote()
optf = Optimization.OptimizationFunction((x, p) -> loss_neuralode(x), adtype)
optprob = Optimization.OptimizationProblem(optf, ComponentArray(pinit))

[gdalle]

Hi! Can you give a complete runnable code with all imports and error messages?

[zahachtah]

I have two full-code examples using mock data (since the full dataload takes ages). First one is the pretraining from https://docs.sciml.ai/DiffEqFlux/stable/examples/collocation/. Here it errors as the inner AD is called directly and zygote faults.

#!/usr/bin/env julia

using DiffEqFlux
using DifferentialEquations
using Lux
using StaticArrays
using ComponentArrays
using Optimization
using OptimizationOptimisers
using ADTypes
using Random
using DifferentiationInterface: Constant, prepare_second_derivative,
    value_derivative_and_second_derivative, value_derivative_and_second_derivative!

# --- Synthetic driver, avoiding huge codebae for realdata management ---
m = 10
freqs = range(0.5, 2.5; length=m)
phases = range(0.0, 2pi; length=m)
driver_vec(t, T) = SVector{m,T}(sin.(T.(freqs) .* T(t) .+ T.(phases)))

# --- NN dudt model (same rhs path) ---
rng = Random.default_rng()
nn = Chain(Dense(3 + m, 32, tanh), Dense(32, 1))
ps, st = Lux.setup(rng, nn)
p0 = Float64.(ComponentArray(ps))

# --- f(M0,M1,M2,t) and g(x, u, t, p) ---
function f_scalar(u012, t, p)
    T = eltype(u012)
    x = vcat(u012, driver_vec(t, T))
    return nn(x, p, st)[1][1]
end
g(x, u, t, p) = f_scalar(SVector{3,typeof(x)}(u[1], x, u[3]), t, p)

# --- Inner AD ---
inner_ad = AutoForwardDiff()
u0 = [0.5, 0.1, 0.2]
prep = prepare_second_derivative(g, inner_ad, u0[2], Constant(u0), Constant(0.0), Constant(p0); strict=Val(false))
der = zeros(Float64, 1)
der2 = zeros(Float64, 1)

function rhs_core(u, p, t; use_bang=true)
    if use_bang
        f, _, _ = value_derivative_and_second_derivative!(
            g, der, der2, prep, inner_ad, u[2], Constant(u), Constant(t), Constant(p)
        )
        f1 = der[1]
        f2 = der2[1]
    else
        f, f1, f2 = value_derivative_and_second_derivative(
            g, prep, inner_ad, u[2], Constant(u), Constant(t), Constant(p)
        )
    end
    m0, m1, m2 = u
    drv = driver_vec(t, eltype(u))
    m3 = drv[end-1]
    m4 = drv[end]
    d0 = f * m0 + 0.5 * f2 * m2
    d1 = f1 * m2 + 0.5 * f2 * m3
    d2 = f1 * m3 + 0.5 * f2 * (m4 - m2^2)
    return SVector(d0, d1, d2)
end

function rhs!(du, u, p, t)
    d = rhs_core(u, p, t; use_bang=true)
    du[1] = d[1]
    du[2] = d[2]
    du[3] = d[3]
    return nothing
end

rhs_vec(u, p, t) = rhs_core(u, p, t; use_bang=false)

# --- Synthetic data from rhs ---
tspan = (0.0, 5.0)
tsteps = range(tspan[1], tspan[2]; length=200)
p_true = p0 .+ 0.05 .* randn(length(p0))
prob = ODEProblem(rhs!, copy(u0), tspan, p0)
data = Array(solve(remake(prob; p=p_true), Tsit5(); saveat=tsteps))

# --- Collocation (u, du) ---
du_colloc, u_colloc = collocate_data(data, tsteps, EpanechnikovKernel())

function loss(p)
    cost = zero(eltype(p))
    @inbounds for i in 1:size(u_colloc, 2)
        du_tmp = rhs_vec(@view(u_colloc[:, i]), p, tsteps[i])
        du_true = @view du_colloc[:, i]
        cost += sum(abs2, du_true .- du_tmp)
    end
    return sqrt(cost / size(u_colloc, 2))
end


# Outer AD 
optf = Optimization.OptimizationFunction((x, p) -> loss(x), Optimization.AutoZygote())
optprob = Optimization.OptimizationProblem(optf, p0)
res = Optimization.solve(optprob, Adam(0.05); maxiters=200)

println("final loss = ", loss(res.u))

then I have this one with is the second ODE trianing of the DiffEqFlux example I linked above. Here the innerAD is wrapped in a loss function and now I get a warning that zygote defaults to a numerical AD instead.

#!/usr/bin/env julia

using DifferentialEquations
using Lux
using StaticArrays
using ComponentArrays
using Optimization
using OptimizationOptimisers
using ADTypes
using Random
using DifferentiationInterface: Constant, prepare_second_derivative, value_derivative_and_second_derivative!
using Zygote
using SciMLSensitivity

struct MomentsCtx
    nn
    st
    driver_vec
    inner_ad
    prep
    der
    der2
    u0
    tsteps
    tspan
    prob
    use_euler
    euler_dt
end

function build_ctx(; m=10, dim=64, tlen=200, use_euler=false)
    freqs = range(0.5, 2.5; length=m)
    phases = range(0.0, 2pi; length=m)
    driver_vec(t, T) = SVector{m,T}(sin.(T.(freqs) .* T(t) .+ T.(phases)))

    rng = Random.default_rng()
    nn = Chain(
        Dense(3 + m, dim),
        Dense(dim, dim, swish),
        Dense(dim, 1))
    ps, st = Lux.setup(rng, nn)
    p0 = ComponentArray(ps)

    function f_scalar(u012, t, p)
        T = eltype(u012)
        x = vcat(u012, driver_vec(t, T))
        return nn(x, p, st)[1][1]
    end
    g(x, u, t, p) = f_scalar(SVector{3,typeof(x)}(u[1], x, u[3]), t, p)

    inner_ad = AutoForwardDiff()
    u0 = [0.5, 0.1, 0.2]
    prep = prepare_second_derivative(g, inner_ad, u0[2], Constant(u0), Constant(0.0), Constant(p0); strict=Val(false))
    der = zeros(Float64, 1)
    der2 = zeros(Float64, 1)

    function rhs!(du, u, p, t)
        f, _, _ = value_derivative_and_second_derivative!(
            g, der, der2, prep, inner_ad, u[2], Constant(u), Constant(t), Constant(p)
        )
        f1 = der[1]
        f2 = der2[1]
        m0, m1, m2 = u
        drv = driver_vec(t, eltype(u))
        m3 = drv[end-1]
        m4 = drv[end]
        d0 = f * m0 + 0.5 * f2 * m2
        d1 = f1 * m2 + 0.5 * f2 * m3
        d2 = f1 * m3 + 0.5 * f2 * (m4 - m2^2)
        du[1] = d0
        du[2] = d1
        du[3] = d2
        return nothing
    end

    tsteps = range(0.0, 5.0; length=tlen)
    tspan = (first(tsteps), last(tsteps))
    prob = ODEProblem(rhs!, copy(u0), tspan, p0)
    euler_dt = step(tsteps)

    ctx = MomentsCtx(nn, st, driver_vec, inner_ad, prep, der, der2, u0, tsteps, tspan, prob, use_euler, euler_dt)
    return ctx, p0
end

function predict_ode(ctx::MomentsCtx, p)
    if ctx.use_euler
        return solve(remake(ctx.prob; p=p), Euler(); dt=ctx.euler_dt, adaptive=false, save_everystep=true)
    end
    return solve(remake(ctx.prob; p=p), Tsit5(), saveat=ctx.tsteps)
end

function loss(ctx::MomentsCtx, p, target)
    sol = predict_ode(ctx, p)
    pred = reduce(hcat, sol.u)
    return sum(abs2, pred .- target)
end

function run_train(; tlen=200, maxiters=50, use_euler=false)
    ctx, p0 = build_ctx(tlen=tlen, use_euler=use_euler)
    p_true = p0 .+ 0.05f0 .* randn(Float32, length(p0))
    target = reduce(hcat, predict_ode(ctx, p_true).u)

    optf = Optimization.OptimizationFunction((x, p) -> loss(ctx, x, target), Optimization.AutoZygote())
    optprob = Optimization.OptimizationProblem(optf, p0)
    iter = Ref(0)
    callback = function (p, l)
        iter[] += 1
        println("iter=$(iter[]) loss=$(l)")
        return false
    end
    res = Optimization.solve(optprob, Adam(0.01); maxiters=maxiters, callback=callback)
    return res, ctx
end

res, ctx = run_train()

Any suggestions on how to deal with this inner/outer AD greatly appreciated and any performance tips always welcome of course :-D

cheers!

[zahachtah]

Errormessage for the first script:

julia> include("scripts/tools/test_synth_collocation_train.jl")
ERROR: LoadError: MethodError: no method matching ChainRules.OneElement(::ChainRulesCore.Tangent{Any, @NamedTuple{…}}, ::Tuple{Int64}, ::Tuple{Base.OneTo{…}})
The type `ChainRules.OneElement` exists, but no method is defined for this combination of argument types when trying to construct it.

Closest candidates are:
  ChainRules.OneElement(::T, ::I, ::A) where {N, T<:Number, I<:NTuple{N, Int64}, A<:NTuple{N, AbstractUnitRange}}
   @ ChainRules ~/.julia/packages/ChainRules/14CDN/src/rulesets/Base/indexing.jl:104

Stacktrace:
  [1] ∇getindex(x::Vector{ForwardDiff.Dual{…}}, dy::ChainRulesCore.Tangent{Any, @NamedTuple{…}}, inds::Int64)
    @ ChainRules ~/.julia/packages/ChainRules/14CDN/src/rulesets/Base/indexing.jl:86
  [2] (::ChainRules.var"#thunked_∇getindex##2#thunked_∇getindex##3"{Vector{…}, ChainRulesCore.Tangent{…}, Tuple{…}})()
    @ ChainRules ~/.julia/packages/ChainRules/14CDN/src/rulesets/Base/indexing.jl:69
  [3] unthunk
    @ ~/.julia/packages/ChainRulesCore/Vsbj9/src/tangent_types/thunks.jl:213 [inlined]
  [4] unthunk(x::ChainRulesCore.InplaceableThunk{…})
    @ ChainRulesCore ~/.julia/packages/ChainRulesCore/Vsbj9/src/tangent_types/thunks.jl:252
  [5] (::Functors.ExcludeWalk{…})(::Function, ::ChainRulesCore.InplaceableThunk{…})
    @ Functors ~/.julia/packages/Functors/LbNAu/src/walks.jl:126
  [6] (::Functors.CachedWalk{…})(::Function, ::ChainRulesCore.InplaceableThunk{…})
    @ Functors ~/.julia/packages/Functors/LbNAu/src/walks.jl:177
  [7] execute(::Functors.CachedWalk{…}, ::ChainRulesCore.InplaceableThunk{…})
    @ Functors ~/.julia/packages/Functors/LbNAu/src/walks.jl:55
  [8] fmap(::Function, ::ChainRulesCore.InplaceableThunk{…}; exclude::Function, walk::Functors.DefaultWalk, cache::IdDict{…}, prune::Functors.NoKeyword)
    @ Functors ~/.julia/packages/Functors/LbNAu/src/maps.jl:11
  [9] recursive_unthunk(x::ChainRulesCore.InplaceableThunk{…})
    @ Lux.Utils ~/.julia/packages/Lux/u3j3H/src/utils.jl:237
 [10] (::Lux.Utils.var"#6#7"{ChainRulesCore.InplaceableThunk{…}, Matrix{…}})()
    @ Lux.Utils ~/.julia/packages/Lux/u3j3H/src/utils.jl:122
 [11] unthunk
    @ ~/.julia/packages/ChainRulesCore/Vsbj9/src/tangent_types/thunks.jl:213 [inlined]
 [12] (::Functors.ExcludeWalk{…})(::Function, ::ChainRulesCore.Thunk{…})
    @ Functors ~/.julia/packages/Functors/LbNAu/src/walks.jl:126
 [13] (::Functors.CachedWalk{…})(::Function, ::ChainRulesCore.Thunk{…})
    @ Functors ~/.julia/packages/Functors/LbNAu/src/walks.jl:177
 [14] execute(::Functors.CachedWalk{…}, ::ChainRulesCore.Thunk{…})
    @ Functors ~/.julia/packages/Functors/LbNAu/src/walks.jl:55
 [15] fmap(::Function, ::ChainRulesCore.Thunk{…}; exclude::Function, walk::Functors.DefaultWalk, cache::IdDict{…}, prune::Functors.NoKeyword)
    @ Functors ~/.julia/packages/Functors/LbNAu/src/maps.jl:11
 [16] recursive_unthunk(x::ChainRulesCore.Thunk{Lux.Utils.var"#6#7"{ChainRulesCore.InplaceableThunk{…}, Matrix{…}}})
    @ LuxLib.Utils ~/.julia/packages/LuxLib/BTt3T/src/utils.jl:381
 [17] (::LuxLib.Impl.var"#246#247"{ChainRulesCore.Thunk{…}, ChainRulesCore.ProjectTo{…}, Base.ReshapedArray{…}})()
    @ LuxLib.Impl ~/.julia/packages/LuxLib/BTt3T/src/impl/matmul.jl:240
 [18] unthunk
    @ ~/.julia/packages/ChainRulesCore/Vsbj9/src/tangent_types/thunks.jl:213 [inlined]
 [19] unthunk_tangent
    @ ~/.julia/packages/Zygote/55SqB/src/compiler/chainrules.jl:3 [inlined]
 [20] (::Zygote.var"#820#821"{Zygote.var"#816#817"{…}})(Δ::ChainRulesCore.Thunk{LuxLib.Impl.var"#246#247"{…}})
    @ Zygote ~/.julia/packages/ZygoteRules/CkVIK/src/adjoint.jl:72
 [21] (::Zygote.Pullback{Tuple{…}, Tuple{…}})(Δ::ChainRulesCore.Thunk{LuxLib.Impl.var"#246#247"{…}})
    @ Zygote ~/.julia/packages/Zygote/55SqB/src/compiler/interface2.jl:100
 [22] ##d::Denseinternal#357
    @ ~/.julia/packages/Lux/u3j3H/src/layers/basic.jl:369 [inlined]
 [23] (::Zygote.Pullback{Tuple{…}, Tuple{…}})(Δ::Tuple{ChainRulesCore.InplaceableThunk{…}, Nothing})
    @ Zygote ~/.julia/packages/Zygote/55SqB/src/compiler/interface2.jl:0
 [24] Dense
    @ ./none:-1 [inlined]
 [25] (::Zygote.Pullback{Tuple{…}, Any})(Δ::Tuple{ChainRulesCore.InplaceableThunk{…}, Nothing})
    @ Zygote ~/.julia/packages/Zygote/55SqB/src/compiler/interface2.jl:0
 [26] back
    @ ~/.julia/packages/Zygote/55SqB/src/lib/lib.jl:119 [inlined]
 [27] #259
    @ ~/.julia/packages/ZygoteRules/CkVIK/src/adjoint.jl:72 [inlined]
 [28] back
    @ ~/.julia/packages/Zygote/55SqB/src/lib/lib.jl:155 [inlined]
 [29] applychain
    @ ~/.julia/packages/Lux/u3j3H/src/layers/containers.jl:-1 [inlined]
 [30] (::Zygote.Pullback{Tuple{…}, Tuple{…}})(Δ::Tuple{ChainRulesCore.InplaceableThunk{…}, Nothing})
    @ Zygote ~/.julia/packages/Zygote/55SqB/src/compiler/interface2.jl:0
 [31] (::Zygote.Pullback{Tuple{…}, Tuple{…}})(Δ::Tuple{ChainRulesCore.InplaceableThunk{…}, Nothing})
    @ Zygote ~/.julia/packages/Zygote/55SqB/src/compiler/interface2.jl:100
 [32] back
    @ ~/.julia/packages/Zygote/55SqB/src/lib/lib.jl:119 [inlined]
 [33] #259
    @ ~/.julia/packages/ZygoteRules/CkVIK/src/adjoint.jl:72 [inlined]
 [34] f_scalar
    @ ~/Coding/attractors/scripts/tools/test_synth_collocation_train.jl:31 [inlined]
 [35] (::Zygote.Pullback{…})(Δ::@NamedTuple{…})
    @ Zygote ~/.julia/packages/Zygote/55SqB/src/compiler/interface2.jl:0
 [36] g
    @ ~/Coding/attractors/scripts/tools/test_synth_collocation_train.jl:33 [inlined]
 [37] (::Zygote.Pullback{…})(Δ::@NamedTuple{…})
    @ Zygote ~/.julia/packages/Zygote/55SqB/src/compiler/interface2.jl:0
 [38] #375
    @ ~/.julia/packages/Zygote/55SqB/src/lib/lib.jl:214 [inlined]
 [39] #377
    @ ~/.julia/packages/ZygoteRules/CkVIK/src/adjoint.jl:72 [inlined]
 [40] value_derivative_and_second_derivative
    @ ~/.julia/packages/DifferentiationInterface/MgcE4/ext/DifferentiationInterfaceForwardDiffExt/onearg.jl:685 [inlined]
 [41] (::Zygote.Pullback{Tuple{…}, Tuple{…}})(Δ::Tuple{Float64, Float64, Float64})
    @ Zygote ~/.julia/packages/Zygote/55SqB/src/compiler/interface2.jl:0
 [42] #rhs_core#1024
    @ ~/Coding/attractors/scripts/tools/test_synth_collocation_train.jl:60 [inlined]
 [43] (::Zygote.Pullback{Tuple{…}, Any})(Δ::SVector{3, Float64})
    @ Zygote ~/.julia/packages/Zygote/55SqB/src/compiler/interface2.jl:0
 [44] (::Zygote.Pullback{Tuple{…}, Any})(Δ::SVector{3, Float64})
    @ Zygote ~/.julia/packages/Zygote/55SqB/src/compiler/interface2.jl:100
--- the above 1 lines are repeated 2 more times ---
 [47] (::Zygote.var"#pullback##0#pullback##1"{Zygote.Pullback{Tuple{…}, Tuple{…}}})(Δ::Float64)
    @ Zygote ~/.julia/packages/Zygote/55SqB/src/compiler/interface.jl:97
 [48] withgradient(::Function, ::ComponentVector{Float64, Vector{Float64}, Tuple{Axis{…}}}, ::Vararg{Any})
    @ Zygote ~/.julia/packages/Zygote/55SqB/src/compiler/interface.jl:0
 [49] value_and_gradient
    @ ~/.julia/packages/DifferentiationInterface/MgcE4/ext/DifferentiationInterfaceZygoteExt/DifferentiationInterfaceZygoteExt.jl:115 [inlined]
 [50] value_and_gradient!(f::Function, grad::ComponentVector{…}, prep::DifferentiationInterface.NoGradientPrep{…}, backend::AutoZygote, x::ComponentVector{…}, contexts::Constant{…})
    @ DifferentiationInterfaceZygoteExt ~/.julia/packages/DifferentiationInterface/MgcE4/ext/DifferentiationInterfaceZygoteExt/DifferentiationInterfaceZygoteExt.jl:131
 [51] (::OptimizationZygoteExt.var"#fg!#19"{…})(res::ComponentVector{…}, θ::ComponentVector{…})
    @ OptimizationZygoteExt ~/.julia/packages/OptimizationBase/CWaVk/ext/OptimizationZygoteExt.jl:53
 [52] __solve(cache::OptimizationCache{…})
    @ OptimizationOptimisers ~/.julia/packages/OptimizationOptimisers/atr4L/src/OptimizationOptimisers.jl:83
 [53] solve!(cache::OptimizationCache{…})
    @ OptimizationBase ~/.julia/packages/OptimizationBase/CWaVk/src/solve.jl:205
 [54] solve(::OptimizationProblem{…}, ::Adam{…}; kwargs::@Kwargs{…})
    @ OptimizationBase ~/.julia/packages/OptimizationBase/CWaVk/src/solve.jl:96
 [55] top-level scope
    @ ~/Coding/attractors/scripts/tools/test_synth_collocation_train.jl:98
 [56] include(mapexpr::Function, mod::Module, _path::String)
    @ Base ./Base.jl:307
 [57] top-level scope
    @ REPL[8]:1
in expression starting at /Users/jonnorberg/Coding/attractors/scripts/tools/test_synth_collocation_train.jl:98
Some type information was truncated. Use `show(err)` to see complete types.

[zahachtah]

warnings from the second script:

iter=5 loss=18.63536396991454
┌ Warning: Potential performance improvement omitted. ReverseDiffVJP tried and failed in the automated AD choice algorithm. To show the stack trace, set SciMLSensitivity.STACKTRACE_WITH_VJPWARN[] = true. To turn off this printing, add `verbose = false` to the `solve` call.
└ @ SciMLSensitivity ~/.julia/packages/SciMLSensitivity/iaNlp/src/concrete_solve.jl:68

┌ Warning: Reverse-Mode AD VJP choices all failed. Falling back to numerical VJPs
└ @ SciMLSensitivity ~/.julia/packages/SciMLSensitivity/iaNlp/src/concrete_solve.jl:208

[gdalle]

Thanks for the examples. At first glance, I think this is a SciML issue more than a DifferentiationInterface issue. There might be limitations of individual AD backends involved, but I don't see a clue telling me that DI itself is responsible. I think you may want to reopen this issue on whichever SciML repo seems most appropriate (feel free to tag me in case DI ends up being the culprit after all).

[zahachtah]

ok, thank you, will do!