# AD of Neural ODEs with \`SplitODEProblem\` using Zygote or Enzyme

**URL:** <https://discourse.julialang.org/t/ad-of-neural-odes-with-splitodeproblem-using-zygote-or-enzyme/112114>\
**Category:** Modelling & Simulations\
**Tags:** question\
**Created:** [March 26, 2024, 12:53am UTC](https://discourse.julialang.org/t/ad-of-neural-odes-with-splitodeproblem-using-zygote-or-enzyme/112114 "2024-03-26T00:53:41Z")\
**Posts on this page:** 2\
**Page:** 1

<div class="post-metadata">

**Author:** ![xkykai](https://avatars.discourse-cdn.com/v4/letter/x/e9bcb4/32.png) [@xkykai](https://discourse.julialang.org/u/xkykai)\
**Post date:** [March 26, 2024, 12:53am UTC](https://discourse.julialang.org/t/ad-of-neural-odes-with-splitodeproblem-using-zygote-or-enzyme/112114/1 "2024-03-26T00:53:41Z")

</div>

Hello, I am trying to perform AD in the Neural ODE setup using a `SplitODEProblem` as it is a PDE discretization of the diffusion equation which is rather stiff when using explicit timesteppers. The neural network appears on the `rhs` part of the `SplitODEProblem` to capture missing fluxes. Here is a MWE:

```julia
using LinearAlgebra
using DiffEqBase
using OrdinaryDiffEq: SplitODEProblem, solve, IMEXEuler
import SciMLBase
using Lux, OptimizationOptimisers, ComponentArrays, Random, SciMLSensitivity
using Optimization
using Printf

n = 32
zC = collect(1.:1:32)
Δz = zC[2] - zC[1]

rng = Random.default_rng()

NN = Chain(Dense(33, 10, leakyrelu), Dense(10, 31))
ps, st = Lux.setup(rng, NN)

ps = ps |> ComponentArray .|> Float64
ps .*= 0

function Dᶜ(N, Δ)
    D = zeros(N, N+1)
    for k in 1:N
        D[k, k] = -1.0
        D[k, k+1] = 1.0
    end
    D .= 1/Δ .* D
    return D
end

function Dᶠ(N, Δ)
    D = zeros(N+1, N)
    for k in 2:N
        D[k, k-1] = -1.0
        D[k, k] = 1.0
    end
    D .= 1/Δ .* D
    return D
end

D_center = Dᶜ(n, Δz)
D_face = Dᶠ(n, Δz)

u0 = collect(range(0, 1, length=n))

function compute_diffusivity(∂u∂z)
    return ifelse(∂u∂z < 0, 0.2, 1e-5)
end

α_initial = compute_diffusivity.(D_face * u0)

D = Tridiagonal(D_center * (α_initial .* D_face))

params = (top=3e-3, bottom=0., NN=NN, st=st, f=1e-4)

function rhs(u, p, t)
    x′ = vcat(u, params.f)

    residual_flux = vcat(params.bottom, first(params.NN(x′, p.ps, params.st)), params.top)

    du = -D_center * residual_flux
    return du
end

function update_diffusivity(A, u, p, t)
    ∂u∂z = D_face * u
    α = compute_diffusivity.(∂u∂z)
    return Tridiagonal(D_center * (α .* D_face))
end

D2 = SciMLBase.MatrixOperator(D, update_func=update_diffusivity)

ps_training = ComponentArray(;ps)

times = collect(0:0.1:1)
tspan = (times[1], times[end])

prob = SplitODEProblem(D2, rhs, u0, tspan, ps_training)

alg = IMEXEuler()
println("Solving...")
sol = solve(prob, alg, dt = 1e-3, saveat = times)

truth = rand(32, 11)

function loss(p)
    prob = SplitODEProblem(D2, rhs, u0, tspan, p)
    sol = Array(solve(prob, alg, dt = 1e-3, saveat = times))
    return sum(abs2, sol - truth)
end

iter = 0
maxiter = 3

callback = function (p, l)
    @printf("loss total %6.10e\n", l,)
    return false
end

loss(ps_training)

adtype = Optimization.AutoZygote()
optf = Optimization.OptimizationFunction((x, p) -> loss(x), adtype)
optprob = Optimization.OptimizationProblem(optf, ps_training)
res = Optimization.solve(optprob, OptimizationOptimisers.Adam(1e-5), callback=callback, maxiters=maxiter)

```

When I use `adtype = Optimization.AutoForwardDiff()`, the code runs without issues. However when I use `adtype = Optimization.AutoReverseDiff()` or `adtype = Optimization.AutoZygote()`, the error below is thrown:

```julia
ERROR: type ODEFunction has no field f1
Stacktrace:

getproperty at Base.jl

initialize!(integrator::OrdinaryDiffEq.ODEIntegrator{…}, cache::OrdinaryDiffEq.SBDFCache{…}) at bdf_perform_step.jl

__init(prob::ODEProblem{…}, alg::OrdinaryDiffEq.SBDF{…}, timeseries_init::Tuple{}, ts_init::Tuple{}, ks_init::Tuple{}, recompile::Type{…}; saveat::Vector{…}, tstops::Vector{…}, d_discontinuities::Tuple{}, save_idxs::Nothing, save_everystep::Bool, save_on::Bool, save_start::Bool, save_end::Nothing, callback::CallbackSet{…}, dense::Bool, calck::Bool, dt::Float64, dtmin::Float64, dtmax::Float64, force_dtmin::Bool, adaptive::Bool, gamma::Int64, abstol::Float64, reltol::Float64, qmin::Int64, qmax::Int64, qsteady_min::Int64, qsteady_max::Int64, beta1::Nothing, beta2::Nothing, qoldinit::Int64, controller::Nothing, fullnormalize::Bool, failfactor::Int64, maxiters::Int64, internalnorm::typeof(DiffEqBase.ODE_DEFAULT_NORM), internalopnorm::typeof(opnorm), isoutofdomain::typeof(DiffEqBase.ODE_DEFAULT_ISOUTOFDOMAIN), unstable_check::typeof(DiffEqBase.ODE_DEFAULT_UNSTABLE_CHECK), verbose::Bool, timeseries_errors::Bool, dense_errors::Bool, advance_to_tstop::Bool, stop_at_next_tstop::Bool, initialize_save::Bool, progress::Bool, progress_steps::Int64, progress_name::String, progress_message::typeof(DiffEqBase.ODE_DEFAULT_PROG_MESSAGE), progress_id::Symbol, userdata::Nothing, allow_extrapolation::Bool, initialize_integrator::Bool, alias_u0::Bool, alias_du0::Bool, initializealg::OrdinaryDiffEq.DefaultInit, kwargs::Base.Pairs{…}) at solve.jl

__init at solve.jl (repeats 5 times)

#__solve#757 at solve.jl

__solve at solve.jl

solve_call(_prob::ODEProblem{…}, args::OrdinaryDiffEq.SBDF{…}; merge_callbacks::Bool, kwargshandle::Nothing, kwargs::Base.Pairs{…}) at solve.jl

kwcall(::@NamedTuple{…}, ::typeof(DiffEqBase.solve_call), _prob::ODEProblem{…}, args::OrdinaryDiffEq.SBDF{…}) at solve.jl

solve_up(prob::ODEProblem{…}, sensealg::Nothing, u0::Vector{…}, p::ComponentVector{…}, args::OrdinaryDiffEq.SBDF{…}; kwargs::Base.Pairs{…}) at solve.jl

solve_up at solve.jl

solve(prob::ODEProblem{…}, args::OrdinaryDiffEq.SBDF{…}; sensealg::Nothing, u0::Nothing, p::Nothing, wrap::Val{…}, kwargs::Base.Pairs{…}) at solve.jl

kwcall(::@NamedTuple{…}, ::typeof(solve), prob::ODEProblem{…}, args::OrdinaryDiffEq.SBDF{…}) at solve.jl

_adjoint_sensitivities(sol::ODESolution{…}, sensealg::InterpolatingAdjoint{…}, alg::OrdinaryDiffEq.SBDF{…}; t::Vector{…}, dgdu_discrete::Function, dgdp_discrete::Nothing, dgdu_continuous::Nothing, dgdp_continuous::Nothing, g::Nothing, abstol::Float64, reltol::Float64, checkpoints::Vector{…}, corfunc_analytical::Nothing, callback::Nothing, kwargs::Base.Pairs{…}) at sensitivity_interface.jl

_adjoint_sensitivities at sensitivity_interface.jl

#adjoint_sensitivities#63 at sensitivity_interface.jl

kwcall(::@NamedTuple{…}, ::typeof(adjoint_sensitivities), sol::ODESolution{…}, args::OrdinaryDiffEq.SBDF{…}) at sensitivity_interface.jl

(::SciMLSensitivity.var"#adjoint_sensitivity_backpass#307"{…})(Δ::ODESolution{…}) at concrete_solve.jl

ZBack at chainrules.jl

(::Zygote.var"#kw_zpullback#53"{…})(dy::ODESolution{…}) at chainrules.jl

#291 at lib.jl

(::Zygote.var"#2169#back#293"{…})(Δ::ODESolution{…}) at adjoint.jl

#solve#41 at solve.jl

(::Zygote.Pullback{…})(Δ::ODESolution{…}) at interface2.jl

#291 at lib.jl

(::Zygote.var"#2169#back#293"{…})(Δ::ODESolution{…}) at adjoint.jl

solve at solve.jl

(::Zygote.Pullback{…})(Δ::ODESolution{…}) at interface2.jl

loss at IMEX_SplitODE_MWE.jl

(::Zygote.Pullback{Tuple{…}, Tuple{…}})(Δ::Float64) at interface2.jl

#11 at IMEX_SplitODE_MWE.jl

#291 at lib.jl

#2169#back at adjoint.jl

OptimizationFunction at scimlfunctions.jl

#291 at lib.jl

(::Zygote.var"#2169#back#293"{Zygote.var"#291#292"{Tuple{…}, Zygote.Pullback{…}}})(Δ::Float64) at adjoint.jl

#37 at OptimizationZygoteExt.jl

(::Zygote.Pullback{Tuple{…}, Tuple{…}})(Δ::Float64) at interface2.jl

#291 at lib.jl

#2169#back at adjoint.jl

#39 at OptimizationZygoteExt.jl

(::Zygote.Pullback{Tuple{…}, Tuple{…}})(Δ::Float64) at interface2.jl

(::Zygote.var"#75#76"{Zygote.Pullback{Tuple{…}, Tuple{…}}})(Δ::Float64) at interface.jl

gradient(f::Function, args::ComponentVector{Float64, Vector{Float64}, Tuple{Axis{…}}}) at interface.jl

(::OptimizationZygoteExt.var"#38#56"{…})(::ComponentVector{…}, ::ComponentVector{…}) at OptimizationZygoteExt.jl

macro expansion at OptimizationOptimisers.jl

macro expansion at utils.jl

__solve(cache::OptimizationCache{…}) at OptimizationOptimisers.jl

solve!(cache::OptimizationCache{…}) at solve.jl

solve(::OptimizationProblem{…}, ::Adam; kwargs::Base.Pairs{…}) at solve.jl

kwcall(::@NamedTuple{…}, ::typeof(solve), ::OptimizationProblem{…}, ::Adam) at solve.jl

top-level scope at IMEX_SplitODE_MWE.jl

Some type information was truncated. Use show(err) to see complete types.

ERROR: type ODEFunction has no field f1
Stacktrace:

getproperty at Base.jl

initialize!(integrator::OrdinaryDiffEq.ODEIntegrator{…}, cache::OrdinaryDiffEq.SBDFCache{…}) at bdf_perform_step.jl

__init(prob::ODEProblem{…}, alg::OrdinaryDiffEq.SBDF{…}, timeseries_init::Tuple{}, ts_init::Tuple{}, ks_init::Tuple{}, recompile::Type{…}; saveat::Vector{…}, tstops::Vector{…}, d_discontinuities::Tuple{}, save_idxs::Nothing, save_everystep::Bool, save_on::Bool, save_start::Bool, save_end::Nothing, callback::CallbackSet{…}, dense::Bool, calck::Bool, dt::Float64, dtmin::Float64, dtmax::Float64, force_dtmin::Bool, adaptive::Bool, gamma::Int64, abstol::Float64, reltol::Float64, qmin::Int64, qmax::Int64, qsteady_min::Int64, qsteady_max::Int64, beta1::Nothing, beta2::Nothing, qoldinit::Int64, controller::Nothing, fullnormalize::Bool, failfactor::Int64, maxiters::Int64, internalnorm::typeof(DiffEqBase.ODE_DEFAULT_NORM), internalopnorm::typeof(opnorm), isoutofdomain::typeof(DiffEqBase.ODE_DEFAULT_ISOUTOFDOMAIN), unstable_check::typeof(DiffEqBase.ODE_DEFAULT_UNSTABLE_CHECK), verbose::Bool, timeseries_errors::Bool, dense_errors::Bool, advance_to_tstop::Bool, stop_at_next_tstop::Bool, initialize_save::Bool, progress::Bool, progress_steps::Int64, progress_name::String, progress_message::typeof(DiffEqBase.ODE_DEFAULT_PROG_MESSAGE), progress_id::Symbol, userdata::Nothing, allow_extrapolation::Bool, initialize_integrator::Bool, alias_u0::Bool, alias_du0::Bool, initializealg::OrdinaryDiffEq.DefaultInit, kwargs::Base.Pairs{…}) at solve.jl

__init at solve.jl (repeats 5 times)

#__solve#757 at solve.jl

__solve at solve.jl

solve_call(_prob::ODEProblem{…}, args::OrdinaryDiffEq.SBDF{…}; merge_callbacks::Bool, kwargshandle::Nothing, kwargs::Base.Pairs{…}) at solve.jl

kwcall(::@NamedTuple{…}, ::typeof(DiffEqBase.solve_call), _prob::ODEProblem{…}, args::OrdinaryDiffEq.SBDF{…}) at solve.jl

solve_up(prob::ODEProblem{…}, sensealg::Nothing, u0::Vector{…}, p::ComponentVector{…}, args::OrdinaryDiffEq.SBDF{…}; kwargs::Base.Pairs{…}) at solve.jl

solve_up at solve.jl

solve(prob::ODEProblem{…}, args::OrdinaryDiffEq.SBDF{…}; sensealg::Nothing, u0::Nothing, p::Nothing, wrap::Val{…}, kwargs::Base.Pairs{…}) at solve.jl

kwcall(::@NamedTuple{…}, ::typeof(solve), prob::ODEProblem{…}, args::OrdinaryDiffEq.SBDF{…}) at solve.jl

_adjoint_sensitivities(sol::ODESolution{…}, sensealg::InterpolatingAdjoint{…}, alg::OrdinaryDiffEq.SBDF{…}; t::Vector{…}, dgdu_discrete::Function, dgdp_discrete::Nothing, dgdu_continuous::Nothing, dgdp_continuous::Nothing, g::Nothing, abstol::Float64, reltol::Float64, checkpoints::Vector{…}, corfunc_analytical::Nothing, callback::Nothing, kwargs::Base.Pairs{…}) at sensitivity_interface.jl

_adjoint_sensitivities at sensitivity_interface.jl

#adjoint_sensitivities#63 at sensitivity_interface.jl

kwcall(::@NamedTuple{…}, ::typeof(adjoint_sensitivities), sol::ODESolution{…}, args::OrdinaryDiffEq.SBDF{…}) at sensitivity_interface.jl

(::SciMLSensitivity.var"#adjoint_sensitivity_backpass#307"{…})(Δ::ODESolution{…}) at concrete_solve.jl

ZBack at chainrules.jl

(::Zygote.var"#kw_zpullback#53"{…})(dy::ODESolution{…}) at chainrules.jl

#291 at lib.jl

(::Zygote.var"#2169#back#293"{…})(Δ::ODESolution{…}) at adjoint.jl

#solve#41 at solve.jl

(::Zygote.Pullback{…})(Δ::ODESolution{…}) at interface2.jl

#291 at lib.jl

(::Zygote.var"#2169#back#293"{…})(Δ::ODESolution{…}) at adjoint.jl

solve at solve.jl

(::Zygote.Pullback{…})(Δ::ODESolution{…}) at interface2.jl

loss at IMEX_SplitODE_MWE.jl

(::Zygote.Pullback{Tuple{…}, Tuple{…}})(Δ::Float64) at interface2.jl

#11 at IMEX_SplitODE_MWE.jl

#291 at lib.jl

#2169#back at adjoint.jl

OptimizationFunction at scimlfunctions.jl

#291 at lib.jl

(::Zygote.var"#2169#back#293"{Zygote.var"#291#292"{Tuple{…}, Zygote.Pullback{…}}})(Δ::Float64) at adjoint.jl

#37 at OptimizationZygoteExt.jl

(::Zygote.Pullback{Tuple{…}, Tuple{…}})(Δ::Float64) at interface2.jl

#291 at lib.jl

#2169#back at adjoint.jl

#39 at OptimizationZygoteExt.jl

(::Zygote.Pullback{Tuple{…}, Tuple{…}})(Δ::Float64) at interface2.jl

(::Zygote.var"#75#76"{Zygote.Pullback{Tuple{…}, Tuple{…}}})(Δ::Float64) at interface.jl

gradient(f::Function, args::ComponentVector{Float64, Vector{Float64}, Tuple{Axis{…}}}) at interface.jl

(::OptimizationZygoteExt.var"#38#56"{…})(::ComponentVector{…}, ::ComponentVector{…}) at OptimizationZygoteExt.jl

macro expansion at OptimizationOptimisers.jl

macro expansion at utils.jl

__solve(cache::OptimizationCache{…}) at OptimizationOptimisers.jl

solve!(cache::OptimizationCache{…}) at solve.jl

solve(::OptimizationProblem{…}, ::Adam; kwargs::Base.Pairs{…}) at solve.jl

kwcall(::@NamedTuple{…}, ::typeof(solve), ::OptimizationProblem{…}, ::Adam) at solve.jl

top-level scope at IMEX_SplitODE_MWE.jl

```

When I use `adtype = AutoEnzyme()`, the error below is thrown:

```julia
ERROR: Enzyme execution failed.
Mismatched activity for: store {} addrspace(10)* %.fca.0.extract3, {} addrspace(10)* addrspace(10)* %.sroa.04.0..sroa_cast, align 8, !dbg !15, !alias.scope !17, !noalias !21 const val: %.fca.0.extract3 = extractvalue { {} addrspace(10)* } %1, 0, !dbg !8
Type tree: {[-1]:Pointer, [-1,-1]:Pointer, [-1,-1,0]:Pointer, [-1,-1,0,-1]:Float@double, [-1,-1,8]:Integer, [-1,-1,9]:Integer, [-1,-1,10]:Integer, [-1,-1,11]:Integer, [-1,-1,12]:Integer, [-1,-1,13]:Integer, [-1,-1,14]:Integer, [-1,-1,15]:Integer, [-1,-1,16]:Integer, [-1,-1,17]:Integer, [-1,-1,18]:Integer, [-1,-1,19]:Integer, [-1,-1,20]:Integer, [-1,-1,21]:Integer, [-1,-1,22]:Integer, [-1,-1,23]:Integer, [-1,-1,24]:Integer, [-1,-1,25]:Integer, [-1,-1,26]:Integer, [-1,-1,27]:Integer, [-1,-1,28]:Integer, [-1,-1,29]:Integer, [-1,-1,30]:Integer, [-1,-1,31]:Integer, [-1,-1,32]:Integer, [-1,-1,33]:Integer, [-1,-1,34]:Integer, [-1,-1,35]:Integer, [-1,-1,36]:Integer, [-1,-1,37]:Integer, [-1,-1,38]:Integer, [-1,-1,39]:Integer}
You may be using a constant variable as temporary storage for active memory (https://enzyme.mit.edu/julia/stable/#Activity-of-temporary-storage). If not, please open an issue, and either rewrite this variable to not be conditionally active or use Enzyme.API.runtimeActivity!(true) as a workaround for now

Stacktrace:
 [1] SplitODEProblem
   @ C:\Users\xinle\.julia\packages\SciMLBase\Dwomw\src\problems\ode_problems.jl:444
 [2] SplitODEProblem
   @ C:\Users\xinle\.julia\packages\SciMLBase\Dwomw\src\problems\ode_problems.jl:0

Stacktrace:

throwerr(cstr::Cstring) at compiler.jl

```

Thank you so much in advance for your help! I am trying to train the Neural ODEs in a more performant way, hence I am trying the `SplitODEProblem` with the hope of using Zygote or Enzyme.

---

<div class="post-metadata">

**Author:** ![ChrisRackauckas](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/chrisrackauckas/32/77_2.png) [@ChrisRackauckas](https://discourse.julialang.org/u/ChrisRackauckas)\
**Post date:** [March 26, 2024, 2:29pm UTC](https://discourse.julialang.org/t/ad-of-neural-odes-with-splitodeproblem-using-zygote-or-enzyme/112114/2 "2024-03-26T14:29:53Z")

</div>

I think there is already an issue on this. It is indeed a missing part of the adjoints that it doesn’t support a split ODE problem with a linear operator as the first part. It is something that is possible to support but we just haven’t gotten around to it yet.