Cannot store inactive but differentiable variable into active tuple

[vleon1234]

Full discussion is here: https://discourse.julialang.org/t/lux-enzyme-and-zygote-neuralode-segmentation-fault/123645/4

MWE (with full error below the code):

println("Load libraries")
using OrdinaryDiffEq, DifferentialEquations, ComponentArrays
using Enzyme, Lux
using Zygote, SciMLSensitivity
using StableRNGs

using Optimization, OptimizationOptimisers

using Plots

rng = StableRNG(1111)

println("Setup NODE")

in_size = 4
# This script works with Zygote.gradient when layer_size=5, but julia crashes on type inference when layer_size=10 when calling Zygote.gradient()
layer_size = 10

const sc = Lux.Chain(Lux.Dense(in_size,layer_size,tanh), 
                     Lux.Dense(layer_size,layer_size,tanh), 
                     Lux.Dense(layer_size,2))
# Get the initial parameters and state variables of the model
p_nn, st = Lux.setup(rng, sc)
const _st = st

function NODE!(du,u,p,t)
    NN = sc([u;vs],p,_st)[1]
    du[1] = u[2] + NN[1]
    du[2] = NN[2]
end

println("Test NODE")
u0_test = [1.0,2.0]
vs = [3.,4.]

theta = ComponentArray(p=p_nn, u0=u0_test)
sc([theta.u0;vs], theta.p, _st)

println("Check if ODE solves")

ts = [0,.1,.2,.3,.4,.5]

prob = ODEProblem(NODE!, theta.u0, (0.,1.))
pred = solve(prob, Tsit5(), u0=theta.u0, p=theta.p, saveat=ts)[1,:]

R = rand(length(ts))

function predict_neuralode(theta)
    Array(solve(prob, Tsit5(), u0 = theta.u0, p = theta.p, saveat = ts))
end

function loss_neuralode(theta)
    pred = predict_neuralode(theta)[1,:]
    loss = sum(abs2, R .- pred)
    return loss
end

loss_neuralode(theta)

println("Check if works with Enzyme autodiff")

Enzyme.gradient(Reverse, loss_neuralode, theta)

# Testing Enzyme autodiff based on: https://docs.sciml.ai/SciMLSensitivity/dev/faq/

prob = ODEProblem(NODE!, theta.u0, (0.,1.), theta.p)
u0 = prob.u0
p = prob.p
tmp2 = Enzyme.make_zero(p)
t = prob.tspan[1]
du = zero(u0)

if DiffEqBase.isinplace(prob)
    _f = prob.f
else
    _f = (du, u, p, t) -> (du .= prob.f(u, p, t); nothing)
end

_tmp6 = Enzyme.make_zero(_f)
tmp3 = zero(u0)
tmp4 = zero(u0)
ytmp = zero(u0)
tmp1 = zero(u0)

# Error here
Enzyme.autodiff(Enzyme.Reverse, Enzyme.Duplicated(_f, _tmp6),
    Enzyme.Const, Enzyme.Duplicated(tmp3, tmp4),
    Enzyme.Duplicated(ytmp, tmp1),
    Enzyme.Duplicated(p, tmp2),
    Enzyme.Const(t))

Error message:

ERROR: LoadError: "Error cannot store inactive but differentiable variable [0.0, 0.1, 0.2, 0.3, 0.4, 0.5] into active tuple"
Stacktrace:
  [1] runtime_newstruct_augfwd(activity::Type{Val{(true, true, false)}}, width::Val{1}, ModifiedBetween::Val{(true, true, true, true)}, ::Type{@NamedTuple{u0::SubArray{Float64, 1, Vector{Float64}, Tuple{UnitRange{Int64}}, true}, p::ComponentVector{Float64, SubArray{Float64, 1, Vector{Float64}, Tuple{UnitRange{Int64}}, true}, Tuple{Axis{(layer_1 = ViewAxis(1:50, Axis(weight = ViewAxis(1:40, ShapedAxis((10, 4))), bias = 41:50)), layer_2 = ViewAxis(51:160, Axis(weight = ViewAxis(1:100, ShapedAxis((10, 10))), bias = 101:110)), layer_3 = ViewAxis(161:182, Axis(weight = ViewAxis(1:20, ShapedAxis((2, 10))), bias = 21:22)))}}}, saveat::Vector{Float64}}}, RT::Val{Any}, primal_1::SubArray{Float64, 1, Vector{Float64}, Tuple{UnitRange{Int64}}, true}, shadow_1_1::SubArray{Float64, 1, Vector{Float64}, Tuple{UnitRange{Int64}}, true}, primal_2::ComponentVector{Float64, SubArray{Float64, 1, Vector{Float64}, Tuple{UnitRange{Int64}}, true}, Tuple{Axis{(layer_1 = ViewAxis(1:50, Axis(weight = ViewAxis(1:40, ShapedAxis((10, 4))), bias = 41:50)), layer_2 = ViewAxis(51:160, Axis(weight = ViewAxis(1:100, ShapedAxis((10, 10))), bias = 101:110)), layer_3 = ViewAxis(161:182, Axis(weight = ViewAxis(1:20, ShapedAxis((2, 10))), bias = 21:22)))}}}, shadow_2_1::ComponentVector{Float64, SubArray{Float64, 1, Vector{Float64}, Tuple{UnitRange{Int64}}, true}, Tuple{Axis{(layer_1 = ViewAxis(1:50, Axis(weight = ViewAxis(1:40, ShapedAxis((10, 4))), bias = 41:50)), layer_2 = ViewAxis(51:160, Axis(weight = ViewAxis(1:100, ShapedAxis((10, 10))), bias = 101:110)), layer_3 = ViewAxis(161:182, Axis(weight = ViewAxis(1:20, ShapedAxis((2, 10))), bias = 21:22)))}}}, primal_3::Vector{Float64}, shadow_3_1::Nothing)
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/aViNX/src/rules/typeunstablerules.jl:33
  [2] NamedTuple
    @ ./boot.jl:727 [inlined]
  [3] predict_neuralode
    @ ~/DropBounce/src/MWE.jl:52
  [4] loss_neuralode
    @ ~/DropBounce/src/MWE.jl:56 [inlined]
  [5] loss_neuralode
    @ ~/DropBounce/src/MWE.jl:0 [inlined]
  [6] augmented_julia_loss_neuralode_17527_inner_1wrap
    @ ~/DropBounce/src/MWE.jl:0
  [7] macro expansion
    @ ~/.julia/packages/Enzyme/aViNX/src/compiler.jl:5201 [inlined]
  [8] enzyme_call
    @ ~/.julia/packages/Enzyme/aViNX/src/compiler.jl:4747 [inlined]
  [9] AugmentedForwardThunk
    @ ~/.julia/packages/Enzyme/aViNX/src/compiler.jl:4683 [inlined]
 [10] autodiff
    @ ~/.julia/packages/Enzyme/aViNX/src/Enzyme.jl:396 [inlined]
 [11] autodiff
    @ ~/.julia/packages/Enzyme/aViNX/src/Enzyme.jl:524 [inlined]
 [12] macro expansion
    @ ~/.julia/packages/Enzyme/aViNX/src/sugar.jl:319 [inlined]
 [13] gradient(::ReverseMode{false, false, FFIABI, false, false}, ::typeof(loss_neuralode), ::ComponentVector{Float64, Vector{Float64}, Tuple{Axis{(p = ViewAxis(1:182, Axis(layer_1 = ViewAxis(1:50, Axis(weight = ViewAxis(1:40, ShapedAxis((10, 4))), bias = 41:50)), layer_2 = ViewAxis(51:160, Axis(weight = ViewAxis(1:100, ShapedAxis((10, 10))), bias = 101:110)), layer_3 = ViewAxis(161:182, Axis(weight = ViewAxis(1:20, ShapedAxis((2, 10))), bias = 21:22)))), u0 = 183:184)}}})
    @ Enzyme ~/.julia/packages/Enzyme/aViNX/src/sugar.jl:258
 [14] top-level scope
    @ ~/DropBounce/src/MWE.jl:115
in expression starting at /home/vleon/DropBounce/src/MWE.jl:115