try/catch is not supported when attempting to use `remake` with Zygote

[jarroyoe]

I have a code that used to run 6 months ago that is not running anymore. I mentioned this error in FluxML/Zygote.jl#1479 but remake seems to be the main issue.
The code

using OrdinaryDiffEq, DiffEqFlux
using Optimization, OptimizationOptimisers, OptimizationOptimJL
using ComponentArrays, Lux, Zygote, Random, CUDA, LinearAlgebra
rng = Random.default_rng()

function trainUDEModel(neuralNetwork,knownDynamics,training_data;needed_ps = Float64[],p_true = Float64[])
    pinit, st = Lux.setup(rng,neuralNetwork)
    st = st |> Lux.gpu
    p64 = Float64.(Lux.gpu(ComponentArray(pinit)))
    training_data = Float64.(Lux.gpu(training_data))
    x0 = Float64.(Lux.gpu(training_data[:,1]))

   function ude(du,u,p,t,q)
        knownPred = convert(CuArray,knownDynamics(u,nothing,q))
        nnPred = convert(CuArray,first(neuralNetwork(u,p,st)))

        du .= knownPred .+ nnPred
    end

    # Closure with the known parameter
    nn_dynamics(du,u,p,t) = ude(du,u,p,t,p_true)
    # Define the problem
    prob_nn = ODEProblem(nn_dynamics,x0, (Float64(1),Float64(size(training_data,2))), p64)
    ## Function to train the network
    # Define a predictor
    function predict(p, X = x0)
        _prob = remake(prob_nn, u0 = X, tspan = (Float64(1),Float64(size(training_data,2))), p = p)
        CUDA.@allowscalar convert(CuArray,solve(_prob, AutoTsit5(Rosenbrock23()), saveat = 1.,
                abstol=1e-6, reltol=1e-6
                ))
    end

    lipschitz_regularizer = 0.5
    function loss_function(p)
	    W1 = p.layer_1.weight
	    W2 = p.layer_2.weight
        lipschitz_constant = spectralRadius(W1)*spectralRadius(W2)

        pred = predict(p)
        loss = sum(abs2,training_data .- pred)/size(training_data,2) + lipschitz_regularizer*lipschitz_constant
        return loss
    end

    losses = Float64[]

    callback = function (p, l)
      push!(losses, l)
    if length(losses)%50==0
          println("Current loss after $(length(losses)) iterations: $(losses[end])")
    end
    return false
    end

    ## Training
    #callback(pinit, loss_function(pinit)...; doplot=true)

    adtype = Optimization.AutoZygote()
    optf = Optimization.OptimizationFunction((x, p) -> loss_function(x), adtype)
    optprob = Optimization.OptimizationProblem(optf, p64)

    result_neuralode = Optimization.solve(optprob,
                                           ADAM(),
                                           #callback = callback,
                                           maxiters = 300)

    optprob2 = remake(optprob,u0 = result_neuralode.u)
    result_neuralode2 = Optimization.solve(optprob2,
                                            Optim.BFGS(initial_stepnorm=0.01),
                                            #callback=callback,
                                            allow_f_increases = false)


    return result_neuralode2.u
end

function spectralRadius(X,niters=10)
    y = randn!(similar(X, size(X, 2)))
    tmp = X * y
    for i in 1:niters
        tmp = X*y
        tmp = tmp / norm(tmp)
        y = X' * tmp
        y = y / norm(y)
    end
    return norm(X*y)
end

training_data = rand(4,10)
neuralnetwork = Lux.Chain(Lux.Dense(4,5),Lux.Dense(5,4))
knownDynamics(x,p,q)=-q
trainUDEModel(neuralnetwork,knownDynamics,training_data;p_true=1)
println("Done!")

yields

┌ Info: The GPU function is being called but the GPU is not accessible.
│ Defaulting back to the CPU. (No action is required if you want
└ to run on the CPU).
┌ Info: The GPU function is being called but the GPU is not accessible.
│ Defaulting back to the CPU. (No action is required if you want
└ to run on the CPU).
ERROR: LoadError: Compiling Tuple{var"#predict#6"{Vector{Float64}, Matrix{Float64}, ODEProblem{Vector{Float64}, Tuple{Float64, Float64}, true, ComponentVector{Float64, Vector{Float64}, Tuple{Axis{(layer_1 = ViewAxis(1:25, Axis(weight = ViewAxis(1:20, ShapedAxis((5, 4), NamedTuple())), bias = ViewAxis(21:25, ShapedAxis((5, 1), NamedTuple())))), layer_2 = ViewAxis(26:49, Axis(weight = ViewAxis(1:20, ShapedAxis((4, 5), NamedTuple())), bias = ViewAxis(21:24, ShapedAxis((4, 1), NamedTuple())))))}}}, ODEFunction{true, SciMLBase.AutoSpecialize, var"#nn_dynamics#5"{Int64, var"#ude#4"{Lux.Chain{NamedTuple{(:layer_1, :layer_2), Tuple{Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}, Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}}}}, typeof(knownDynamics)}}, UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing}, Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}, SciMLBase.StandardODEProblem}}, ComponentVector{Float64, Vector{Float64}, Tuple{Axis{(layer_1 = ViewAxis(1:25, Axis(weight = ViewAxis(1:20, ShapedAxis((5, 4), NamedTuple())), bias = ViewAxis(21:25, ShapedAxis((5, 1), NamedTuple())))), layer_2 = ViewAxis(26:49, Axis(weight = ViewAxis(1:20, ShapedAxis((4, 5), NamedTuple())), bias = ViewAxis(21:24, ShapedAxis((4, 1), NamedTuple())))))}}}, Vector{Float64}}: try/catch is not supported.
Refer to the Zygote documentation for fixes.
https://fluxml.ai/Zygote.jl/latest/limitations

Stacktrace:
ERROR: LoadError: Compiling Tuple{var"#predict#6"{Vector{Float64}, Matrix{Float64}, ODEProblem{Vector{Float64}, Tuple{Float64, Float64}, true, ComponentVector{Float64, Vector{Float64}, Tuple{Axis{(layer_1 = ViewAxis(1:25, Axis(weight = ViewAxis(1:20, ShapedAxis((5, 4), NamedTuple())), bias = ViewAxis(21:25, ShapedAxis((5, 1), NamedTuple())))), layer_2 = ViewAxis(26:49, Axis(weight = ViewAxis(1:20, ShapedAxis((4, 5), NamedTuple())), bias = ViewAxis(21:24, ShapedAxis((4, 1), NamedTuple())))))}}}, ODEFunction{true, SciMLBase.AutoSpecialize, var"#nn_dynamics#5"{Int64, var"#ude#4"{Lux.Chain{NamedTuple{(:layer_1, :layer_2), Tuple{Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}, Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}}}}, typeof(knownDynamics)}}, UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing}, Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}, SciMLBase.StandardODEProblem}}, ComponentVector{Float64, Vector{Float64}, Tuple{Axis{(layer_1 = ViewAxis(1:25, Axis(weight = ViewAxis(1:20, ShapedAxis((5, 4), NamedTuple())), bias = ViewAxis(21:25, ShapedAxis((5, 1), NamedTuple())))), layer_2 = ViewAxis(26:49, Axis(weight = ViewAxis(1:20, ShapedAxis((4, 5), NamedTuple())), bias = ViewAxis(21:24, ShapedAxis((4, 1), NamedTuple())))))}}}, Vector{Float64}}: try/catch is not supported.
Refer to the Zygote documentation for fixes.
https://fluxml.ai/Zygote.jl/latest/limitations

Stacktrace:
  [1] error(s::String)
    @ Base ./error.jl:35
  [2] instrument(ir::IRTools.Inner.IR)
    @ Zygote ~/.julia/packages/Zygote/oGI57/src/compiler/reverse.jl:128
  [3] #Primal#31
    @ ~/.julia/packages/Zygote/oGI57/src/compiler/reverse.jl:227 [inlined]
  [1] error(s::String)
    @ Base ./error.jl:35
  [2] instrument(ir::IRTools.Inner.IR)
    @ Zygote ~/.julia/packages/Zygote/oGI57/src/compiler/reverse.jl:128
  [3] #Primal#31
    @ ~/.julia/packages/Zygote/oGI57/src/compiler/reverse.jl:227 [inlined]
  [4] Zygote.Adjoint(ir::IRTools.Inner.IR; varargs::Nothing, normalise::Bool)
    @ Zygote ~/.julia/packages/Zygote/oGI57/src/compiler/reverse.jl:352
  [5] _generate_pullback_via_decomposition(T::Type)
    @ Zygote ~/.julia/packages/Zygote/oGI57/src/compiler/emit.jl:101
  [6] #s3181#1581
    @ ~/.julia/packages/Zygote/oGI57/src/compiler/interface2.jl:28 [inlined]
  [7] var"#s3181#1581"(::Any, ctx::Any, f::Any, args::Any)
    @ Zygote ./none:0
  [4] Zygote.Adjoint(ir::IRTools.Inner.IR; varargs::Nothing, normalise::Bool)
    @ Zygote ~/.julia/packages/Zygote/oGI57/src/compiler/reverse.jl:352
  [5] _generate_pullback_via_decomposition(T::Type)
    @ Zygote ~/.julia/packages/Zygote/oGI57/src/compiler/emit.jl:101
  [6] #s3181#1581
    @ ~/.julia/packages/Zygote/oGI57/src/compiler/interface2.jl:28 [inlined]
  [7] var"#s3181#1581"(::Any, ctx::Any, f::Any, args::Any)
    @ Zygote ./none:0
  [8] (::Core.GeneratedFunctionStub)(::Any, ::Vararg{Any})
    @ Core ./boot.jl:582
  [9] _pullback
    @ ~/NODE_Community_Forecast/test.jl:27 [inlined]
  [8] (::Core.GeneratedFunctionStub)(::Any, ::Vararg{Any})
    @ Core ./boot.jl:582
  [9] _pullback
    @ ~/NODE_Community_Forecast/test.jl:27 [inlined]
 [10] _pullback(ctx::Zygote.Context{false}, f::var"#predict#6"{Vector{Float64}, Matrix{Float64}, ODEProblem{Vector{Float64}, Tuple{Float64, Float64}, true, ComponentVector{Float64, Vector{Float64}, Tuple{Axis{(layer_1 = ViewAxis(1:25, Axis(weight = ViewAxis(1:20, ShapedAxis((5, 4), NamedTuple())), bias = ViewAxis(21:25, ShapedAxis((5, 1), NamedTuple())))), layer_2 = ViewAxis(26:49, Axis(weight = ViewAxis(1:20, ShapedAxis((4, 5), NamedTuple())), bias = ViewAxis(21:24, ShapedAxis((4, 1), NamedTuple())))))}}}, ODEFunction{true, SciMLBase.AutoSpecialize, var"#nn_dynamics#5"{Int64, var"#ude#4"{Lux.Chain{NamedTuple{(:layer_1, :layer_2), Tuple{Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}, Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}}}}, typeof(knownDynamics)}}, UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing}, Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}, SciMLBase.StandardODEProblem}}, args::ComponentVector{Float64, Vector{Float64}, Tuple{Axis{(layer_1 = ViewAxis(1:25, Axis(weight = ViewAxis(1:20, ShapedAxis((5, 4), NamedTuple())), bias = ViewAxis(21:25, ShapedAxis((5, 1), NamedTuple())))), layer_2 = ViewAxis(26:49, Axis(weight = ViewAxis(1:20, ShapedAxis((4, 5), NamedTuple())), bias = ViewAxis(21:24, ShapedAxis((4, 1), NamedTuple())))))}}})
    @ Zygote ~/.julia/packages/Zygote/oGI57/src/compiler/interface2.jl:0
 [11] _pullback
    @ ~/NODE_Community_Forecast/test.jl:39 [inlined]
 [12] _pullback(ctx::Zygote.Context{false}, f::var"#loss_function#7"{Matrix{Float64}, Float64}, args::ComponentVector{Float64, Vector{Float64}, Tuple{Axis{(layer_1 = ViewAxis(1:25, Axis(weight = ViewAxis(1:20, ShapedAxis((5, 4), NamedTuple())), bias = ViewAxis(21:25, ShapedAxis((5, 1), NamedTuple())))), layer_2 = ViewAxis(26:49, Axis(weight = ViewAxis(1:20, ShapedAxis((4, 5), NamedTuple())), bias = ViewAxis(21:24, ShapedAxis((4, 1), NamedTuple())))))}}})
    @ Zygote ~/.julia/packages/Zygote/oGI57/src/compiler/interface2.jl:0
 [13] _pullback
    @ ~/NODE_Community_Forecast/test.jl:58 [inlined]
 [14] _apply
    @ ./boot.jl:816 [inlined]
 [15] adjoint
    @ ~/.julia/packages/Zygote/oGI57/src/lib/lib.jl:203 [inlined]
 [16] _pullback
    @ ~/.julia/packages/ZygoteRules/4nXuu/src/adjoint.jl:66 [inlined]
 [17] _pullback
    @ ~/.julia/packages/SciMLBase/l4PVV/src/scimlfunctions.jl:3772 [inlined]
 [18] _pullback(::Zygote.Context{false}, ::OptimizationFunction{true, Optimization.AutoZygote, var"#3#9"{var"#loss_function#7"{Matrix{Float64}, Float64}}, 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}, ::ComponentVector{Float64, Vector{Float64}, Tuple{Axis{(layer_1 = ViewAxis(1:25, Axis(weight = ViewAxis(1:20, ShapedAxis((5, 4), NamedTuple())), bias = ViewAxis(21:25, ShapedAxis((5, 1), NamedTuple())))), layer_2 = ViewAxis(26:49, Axis(weight = ViewAxis(1:20, ShapedAxis((4, 5), NamedTuple())), bias = ViewAxis(21:24, ShapedAxis((4, 1), NamedTuple())))))}}}, ::SciMLBase.NullParameters)
    @ Zygote ~/.julia/packages/Zygote/oGI57/src/compiler/interface2.jl:0
 [19] _apply(::Function, ::Vararg{Any})
    @ Core ./boot.jl:816
 [20] adjoint
    @ ~/.julia/packages/Zygote/oGI57/src/lib/lib.jl:203 [inlined]
 [21] _pullback
    @ ~/.julia/packages/ZygoteRules/4nXuu/src/adjoint.jl:66 [inlined]
 [22] _pullback
    @ ~/.julia/packages/Optimization/GEo8L/src/function/zygote.jl:30 [inlined]
 [23] _pullback(ctx::Zygote.Context{false}, f::Optimization.var"#156#165"{OptimizationFunction{true, Optimization.AutoZygote, var"#3#9"{var"#loss_function#7"{Matrix{Float64}, Float64}}, 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}, SciMLBase.NullParameters}, args::ComponentVector{Float64, Vector{Float64}, Tuple{Axis{(layer_1 = ViewAxis(1:25, Axis(weight = ViewAxis(1:20, ShapedAxis((5, 4), NamedTuple())), bias = ViewAxis(21:25, ShapedAxis((5, 1), NamedTuple())))), layer_2 = ViewAxis(26:49, Axis(weight = ViewAxis(1:20, ShapedAxis((4, 5), NamedTuple())), bias = ViewAxis(21:24, ShapedAxis((4, 1), NamedTuple())))))}}})
    @ Zygote ~/.julia/packages/Zygote/oGI57/src/compiler/interface2.jl:0
 [24] _apply(::Function, ::Vararg{Any})
    @ Core ./boot.jl:816
 [25] adjoint
    @ ~/.julia/packages/Zygote/oGI57/src/lib/lib.jl:203 [inlined]
 [26] _pullback
    @ ~/.julia/packages/ZygoteRules/4nXuu/src/adjoint.jl:66 [inlined]
 [27] _pullback
    @ ~/.julia/packages/Optimization/GEo8L/src/function/zygote.jl:34 [inlined]
 [28] _pullback(ctx::Zygote.Context{false}, f::Optimization.var"#158#167"{Tuple{}, Optimization.var"#156#165"{OptimizationFunction{true, Optimization.AutoZygote, var"#3#9"{var"#loss_function#7"{Matrix{Float64}, Float64}}, 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}, SciMLBase.NullParameters}}, args::ComponentVector{Float64, Vector{Float64}, Tuple{Axis{(layer_1 = ViewAxis(1:25, Axis(weight = ViewAxis(1:20, ShapedAxis((5, 4), NamedTuple())), bias = ViewAxis(21:25, ShapedAxis((5, 1), NamedTuple())))), layer_2 = ViewAxis(26:49, Axis(weight = ViewAxis(1:20, ShapedAxis((4, 5), NamedTuple())), bias = ViewAxis(21:24, ShapedAxis((4, 1), NamedTuple())))))}}})
    @ Zygote ~/.julia/packages/Zygote/oGI57/src/compiler/interface2.jl:0
 [29] pullback(f::Function, cx::Zygote.Context{false}, args::ComponentVector{Float64, Vector{Float64}, Tuple{Axis{(layer_1 = ViewAxis(1:25, Axis(weight = ViewAxis(1:20, ShapedAxis((5, 4), NamedTuple())), bias = ViewAxis(21:25, ShapedAxis((5, 1), NamedTuple())))), layer_2 = ViewAxis(26:49, Axis(weight = ViewAxis(1:20, ShapedAxis((4, 5), NamedTuple())), bias = ViewAxis(21:24, ShapedAxis((4, 1), NamedTuple())))))}}})
    @ Zygote ~/.julia/packages/Zygote/oGI57/src/compiler/interface.jl:44
 [30] pullback
    @ ~/.julia/packages/Zygote/oGI57/src/compiler/interface.jl:42 [inlined]
 [10] _pullback(ctx::Zygote.Context{false}, f::var"#predict#6"{Vector{Float64}, Matrix{Float64}, ODEProblem{Vector{Float64}, Tuple{Float64, Float64}, true, ComponentVector{Float64, Vector{Float64}, Tuple{Axis{(layer_1 = ViewAxis(1:25, Axis(weight = ViewAxis(1:20, ShapedAxis((5, 4), NamedTuple())), bias = ViewAxis(21:25, ShapedAxis((5, 1), NamedTuple())))), layer_2 = ViewAxis(26:49, Axis(weight = ViewAxis(1:20, ShapedAxis((4, 5), NamedTuple())), bias = ViewAxis(21:24, ShapedAxis((4, 1), NamedTuple())))))}}}, ODEFunction{true, SciMLBase.AutoSpecialize, var"#nn_dynamics#5"{Int64, var"#ude#4"{Lux.Chain{NamedTuple{(:layer_1, :layer_2), Tuple{Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}, Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}}}}, typeof(knownDynamics)}}, UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing}, Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}, SciMLBase.StandardODEProblem}}, args::ComponentVector{Float64, Vector{Float64}, Tuple{Axis{(layer_1 = ViewAxis(1:25, Axis(weight = ViewAxis(1:20, ShapedAxis((5, 4), NamedTuple())), bias = ViewAxis(21:25, ShapedAxis((5, 1), NamedTuple())))), layer_2 = ViewAxis(26:49, Axis(weight = ViewAxis(1:20, ShapedAxis((4, 5), NamedTuple())), bias = ViewAxis(21:24, ShapedAxis((4, 1), NamedTuple())))))}}})
 [31] gradient(f::Function, args::ComponentVector{Float64, Vector{Float64}, Tuple{Axis{(layer_1 = ViewAxis(1:25, Axis(weight = ViewAxis(1:20, ShapedAxis((5, 4), NamedTuple())), bias = ViewAxis(21:25, ShapedAxis((5, 1), NamedTuple())))), layer_2 = ViewAxis(26:49, Axis(weight = ViewAxis(1:20, ShapedAxis((4, 5), NamedTuple())), bias = ViewAxis(21:24, ShapedAxis((4, 1), NamedTuple())))))}}})
    @ Zygote ~/.julia/packages/Zygote/oGI57/src/compiler/interface.jl:96
    @ Zygote ~/.julia/packages/Zygote/oGI57/src/compiler/interface2.jl:0
 [11] _pullback
    @ ~/NODE_Community_Forecast/test.jl:39 [inlined]
 [12] _pullback(ctx::Zygote.Context{false}, f::var"#loss_function#7"{Matrix{Float64}, Float64}, args::ComponentVector{Float64, Vector{Float64}, Tuple{Axis{(layer_1 = ViewAxis(1:25, Axis(weight = ViewAxis(1:20, ShapedAxis((5, 4), NamedTuple())), bias = ViewAxis(21:25, ShapedAxis((5, 1), NamedTuple())))), layer_2 = ViewAxis(26:49, Axis(weight = ViewAxis(1:20, ShapedAxis((4, 5), NamedTuple())), bias = ViewAxis(21:24, ShapedAxis((4, 1), NamedTuple())))))}}})
    @ Zygote ~/.julia/packages/Zygote/oGI57/src/compiler/interface2.jl:0
 [32] (::Optimization.var"#157#166"{Optimization.var"#156#165"{OptimizationFunction{true, Optimization.AutoZygote, var"#3#9"{var"#loss_function#7"{Matrix{Float64}, Float64}}, 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}, SciMLBase.NullParameters}})(::ComponentVector{Float64, Vector{Float64}, Tuple{Axis{(layer_1 = ViewAxis(1:25, Axis(weight = ViewAxis(1:20, ShapedAxis((5, 4), NamedTuple())), bias = ViewAxis(21:25, ShapedAxis((5, 1), NamedTuple())))), layer_2 = ViewAxis(26:49, Axis(weight = ViewAxis(1:20, ShapedAxis((4, 5), NamedTuple())), bias = ViewAxis(21:24, ShapedAxis((4, 1), NamedTuple())))))}}}, ::ComponentVector{Float64, Vector{Float64}, Tuple{Axis{(layer_1 = ViewAxis(1:25, Axis(weight = ViewAxis(1:20, ShapedAxis((5, 4), NamedTuple())), bias = ViewAxis(21:25, ShapedAxis((5, 1), NamedTuple())))), layer_2 = ViewAxis(26:49, Axis(weight = ViewAxis(1:20, ShapedAxis((4, 5), NamedTuple())), bias = ViewAxis(21:24, ShapedAxis((4, 1), NamedTuple())))))}}})
    @ Optimization ~/.julia/packages/Optimization/GEo8L/src/function/zygote.jl:32
 [33] macro expansion
    @ ~/.julia/packages/OptimizationOptimisers/FWIuf/src/OptimizationOptimisers.jl:31 [inlined]
 [34] macro expansion
    @ ~/.julia/packages/Optimization/GEo8L/src/utils.jl:37 [inlined]
 [13] _pullback
    @ ~/NODE_Community_Forecast/test.jl:58 [inlined]
 [14] _apply
    @ ./boot.jl:816 [inlined]
 [15] adjoint
    @ ~/.julia/packages/Zygote/oGI57/src/lib/lib.jl:203 [inlined]
 [16] _pullback
    @ ~/.julia/packages/ZygoteRules/4nXuu/src/adjoint.jl:66 [inlined]
 [17] _pullback
    @ ~/.julia/packages/SciMLBase/l4PVV/src/scimlfunctions.jl:3772 [inlined]
 [18] _pullback(::Zygote.Context{false}, ::OptimizationFunction{true, Optimization.AutoZygote, var"#3#9"{var"#loss_function#7"{Matrix{Float64}, Float64}}, 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}, ::ComponentVector{Float64, Vector{Float64}, Tuple{Axis{(layer_1 = ViewAxis(1:25, Axis(weight = ViewAxis(1:20, ShapedAxis((5, 4), NamedTuple())), bias = ViewAxis(21:25, ShapedAxis((5, 1), NamedTuple())))), layer_2 = ViewAxis(26:49, Axis(weight = ViewAxis(1:20, ShapedAxis((4, 5), NamedTuple())), bias = ViewAxis(21:24, ShapedAxis((4, 1), NamedTuple())))))}}}, ::SciMLBase.NullParameters)
    @ Zygote ~/.julia/packages/Zygote/oGI57/src/compiler/interface2.jl:0
 [35] __solve(prob::OptimizationProblem{true, OptimizationFunction{true, Optimization.AutoZygote, var"#3#9"{var"#loss_function#7"{Matrix{Float64}, Float64}}, 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}, ComponentVector{Float64, Vector{Float64}, Tuple{Axis{(layer_1 = ViewAxis(1:25, Axis(weight = ViewAxis(1:20, ShapedAxis((5, 4), NamedTuple())), bias = ViewAxis(21:25, ShapedAxis((5, 1), NamedTuple())))), layer_2 = ViewAxis(26:49, Axis(weight = ViewAxis(1:20, ShapedAxis((4, 5), NamedTuple())), bias = ViewAxis(21:24, ShapedAxis((4, 1), NamedTuple())))))}}}, SciMLBase.NullParameters, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}}, opt::Optimisers.Adam{Float32}, data::Base.Iterators.Cycle{Tuple{Optimization.NullData}}; maxiters::Int64, callback::Function, progress::Bool, save_best::Bool, kwargs::Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}})
    @ OptimizationOptimisers ~/.julia/packages/OptimizationOptimisers/FWIuf/src/OptimizationOptimisers.jl:30
 [36] #solve#595
    @ ~/.julia/packages/SciMLBase/l4PVV/src/solve.jl:86 [inlined]
 [19] _apply(::Function, ::Vararg{Any})
    @ Core ./boot.jl:816
 [20] adjoint
    @ ~/.julia/packages/Zygote/oGI57/src/lib/lib.jl:203 [inlined]
 [21] _pullback
    @ ~/.julia/packages/ZygoteRules/4nXuu/src/adjoint.jl:66 [inlined]
 [22] _pullback
    @ ~/.julia/packages/Optimization/GEo8L/src/function/zygote.jl:30 [inlined]
 [37] trainUDEModel(neuralNetwork::Lux.Chain{NamedTuple{(:layer_1, :layer_2), Tuple{Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}, Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}}}}, knownDynamics::typeof(knownDynamics), training_data::Matrix{Float64}; needed_ps::Vector{Float64}, p_true::Int64)
    @ Main ~/NODE_Community_Forecast/test.jl:61
 [38] top-level scope
    @ ~/NODE_Community_Forecast/test.jl:91
in expression starting at /home/jarroyoesquivel/NODE_Community_Forecast/test.jl:91
 [23] _pullback(ctx::Zygote.Context{false}, f::Optimization.var"#156#165"{OptimizationFunction{true, Optimization.AutoZygote, var"#3#9"{var"#loss_function#7"{Matrix{Float64}, Float64}}, 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}, SciMLBase.NullParameters}, args::ComponentVector{Float64, Vector{Float64}, Tuple{Axis{(layer_1 = ViewAxis(1:25, Axis(weight = ViewAxis(1:20, ShapedAxis((5, 4), NamedTuple())), bias = ViewAxis(21:25, ShapedAxis((5, 1), NamedTuple())))), layer_2 = ViewAxis(26:49, Axis(weight = ViewAxis(1:20, ShapedAxis((4, 5), NamedTuple())), bias = ViewAxis(21:24, ShapedAxis((4, 1), NamedTuple())))))}}})
    @ Zygote ~/.julia/packages/Zygote/oGI57/src/compiler/interface2.jl:0
 [24] _apply(::Function, ::Vararg{Any})
    @ Core ./boot.jl:816
 [25] adjoint
    @ ~/.julia/packages/Zygote/oGI57/src/lib/lib.jl:203 [inlined]
 [26] _pullback
    @ ~/.julia/packages/ZygoteRules/4nXuu/src/adjoint.jl:66 [inlined]
 [27] _pullback
    @ ~/.julia/packages/Optimization/GEo8L/src/function/zygote.jl:34 [inlined]
 [28] _pullback(ctx::Zygote.Context{false}, f::Optimization.var"#158#167"{Tuple{}, Optimization.var"#156#165"{OptimizationFunction{true, Optimization.AutoZygote, var"#3#9"{var"#loss_function#7"{Matrix{Float64}, Float64}}, 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}, SciMLBase.NullParameters}}, args::ComponentVector{Float64, Vector{Float64}, Tuple{Axis{(layer_1 = ViewAxis(1:25, Axis(weight = ViewAxis(1:20, ShapedAxis((5, 4), NamedTuple())), bias = ViewAxis(21:25, ShapedAxis((5, 1), NamedTuple())))), layer_2 = ViewAxis(26:49, Axis(weight = ViewAxis(1:20, ShapedAxis((4, 5), NamedTuple())), bias = ViewAxis(21:24, ShapedAxis((4, 1), NamedTuple())))))}}})
    @ Zygote ~/.julia/packages/Zygote/oGI57/src/compiler/interface2.jl:0
 [29] pullback(f::Function, cx::Zygote.Context{false}, args::ComponentVector{Float64, Vector{Float64}, Tuple{Axis{(layer_1 = ViewAxis(1:25, Axis(weight = ViewAxis(1:20, ShapedAxis((5, 4), NamedTuple())), bias = ViewAxis(21:25, ShapedAxis((5, 1), NamedTuple())))), layer_2 = ViewAxis(26:49, Axis(weight = ViewAxis(1:20, ShapedAxis((4, 5), NamedTuple())), bias = ViewAxis(21:24, ShapedAxis((4, 1), NamedTuple())))))}}})
    @ Zygote ~/.julia/packages/Zygote/oGI57/src/compiler/interface.jl:44
 [30] pullback
    @ ~/.julia/packages/Zygote/oGI57/src/compiler/interface.jl:42 [inlined]
 [31] gradient(f::Function, args::ComponentVector{Float64, Vector{Float64}, Tuple{Axis{(layer_1 = ViewAxis(1:25, Axis(weight = ViewAxis(1:20, ShapedAxis((5, 4), NamedTuple())), bias = ViewAxis(21:25, ShapedAxis((5, 1), NamedTuple())))), layer_2 = ViewAxis(26:49, Axis(weight = ViewAxis(1:20, ShapedAxis((4, 5), NamedTuple())), bias = ViewAxis(21:24, ShapedAxis((4, 1), NamedTuple())))))}}})
    @ Zygote ~/.julia/packages/Zygote/oGI57/src/compiler/interface.jl:96
 [32] (::Optimization.var"#157#166"{Optimization.var"#156#165"{OptimizationFunction{true, Optimization.AutoZygote, var"#3#9"{var"#loss_function#7"{Matrix{Float64}, Float64}}, 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}, SciMLBase.NullParameters}})(::ComponentVector{Float64, Vector{Float64}, Tuple{Axis{(layer_1 = ViewAxis(1:25, Axis(weight = ViewAxis(1:20, ShapedAxis((5, 4), NamedTuple())), bias = ViewAxis(21:25, ShapedAxis((5, 1), NamedTuple())))), layer_2 = ViewAxis(26:49, Axis(weight = ViewAxis(1:20, ShapedAxis((4, 5), NamedTuple())), bias = ViewAxis(21:24, ShapedAxis((4, 1), NamedTuple())))))}}}, ::ComponentVector{Float64, Vector{Float64}, Tuple{Axis{(layer_1 = ViewAxis(1:25, Axis(weight = ViewAxis(1:20, ShapedAxis((5, 4), NamedTuple())), bias = ViewAxis(21:25, ShapedAxis((5, 1), NamedTuple())))), layer_2 = ViewAxis(26:49, Axis(weight = ViewAxis(1:20, ShapedAxis((4, 5), NamedTuple())), bias = ViewAxis(21:24, ShapedAxis((4, 1), NamedTuple())))))}}})
    @ Optimization ~/.julia/packages/Optimization/GEo8L/src/function/zygote.jl:32
 [33] macro expansion
    @ ~/.julia/packages/OptimizationOptimisers/FWIuf/src/OptimizationOptimisers.jl:31 [inlined]
 [34] macro expansion
    @ ~/.julia/packages/Optimization/GEo8L/src/utils.jl:37 [inlined]
 [35] __solve(prob::OptimizationProblem{true, OptimizationFunction{true, Optimization.AutoZygote, var"#3#9"{var"#loss_function#7"{Matrix{Float64}, Float64}}, 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}, ComponentVector{Float64, Vector{Float64}, Tuple{Axis{(layer_1 = ViewAxis(1:25, Axis(weight = ViewAxis(1:20, ShapedAxis((5, 4), NamedTuple())), bias = ViewAxis(21:25, ShapedAxis((5, 1), NamedTuple())))), layer_2 = ViewAxis(26:49, Axis(weight = ViewAxis(1:20, ShapedAxis((4, 5), NamedTuple())), bias = ViewAxis(21:24, ShapedAxis((4, 1), NamedTuple())))))}}}, SciMLBase.NullParameters, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}}, opt::Optimisers.Adam{Float32}, data::Base.Iterators.Cycle{Tuple{Optimization.NullData}}; maxiters::Int64, callback::Function, progress::Bool, save_best::Bool, kwargs::Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}})
    @ OptimizationOptimisers ~/.julia/packages/OptimizationOptimisers/FWIuf/src/OptimizationOptimisers.jl:30
 [36] #solve#595
    @ ~/.julia/packages/SciMLBase/l4PVV/src/solve.jl:86 [inlined]
 [37] trainUDEModel(neuralNetwork::Lux.Chain{NamedTuple{(:layer_1, :layer_2), Tuple{Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}, Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}}}}, knownDynamics::typeof(knownDynamics), training_data::Matrix{Float64}; needed_ps::Vector{Float64}, p_true::Int64)
    @ Main ~/NODE_Community_Forecast/test.jl:61
 [38] top-level scope
    @ ~/NODE_Community_Forecast/test.jl:91
in expression starting at /home/jarroyoesquivel/NODE_Community_Forecast/test.jl:91

My current Pkg.status():

⌅ [052768ef] CUDA v3.12.2
⌅ [b0b7db55] ComponentArrays v0.13.8
⌅ [aae7a2af] DiffEqFlux v2.4.0
⌃ [b2108857] Lux v0.4.37
⌃ [7f7a1694] Optimization v3.12.1
⌃ [36348300] OptimizationOptimJL v0.1.5
⌃ [42dfb2eb] OptimizationOptimisers v0.1.2
⌅ [1dea7af3] OrdinaryDiffEq v6.49.1
⌃ [e88e6eb3] Zygote v0.6.58
  [37e2e46d] LinearAlgebra
  [9a3f8284] Random

[ChrisRackauckas]

Can you help narrow this down? Does the MWE need GPUs to make the issue happen, etc? I don't think there are any try/catch statements in the remake code at all, so making a minimal example would be extremely helpful for narrowing this down. Right now there are just so many things in there, ComponentArrays, CUDA, Lux, etc. and I presume you don't need all of those combined to get the issue. So which is it?

[jarroyoe]

That's the tricky part, I couldn't pinpoint what exactly would cause the error to generate a MWE. Without CUDA, the following code

using OrdinaryDiffEq, DiffEqFlux
using Optimization, OptimizationOptimisers, OptimizationOptimJL
using ComponentArrays, Lux, Zygote, Random, LinearAlgebra
rng = Random.default_rng()

function trainUDEModel(neuralNetwork,knownDynamics,training_data;needed_ps = Float64[],p_true = Float64[])
    pinit, st = Lux.setup(rng,neuralNetwork)
    p64 = Float64.(ComponentArray(pinit))
    x0 = Float64.(training_data[:,1])

   function ude(du,u,p,t,q)
        knownPred = knownDynamics(u,nothing,q)
        nnPred = first(neuralNetwork(u,p,st))

        du .= knownPred .+ nnPred
    end

    # Closure with the known parameter
    nn_dynamics(du,u,p,t) = ude(du,u,p,t,p_true)
    # Define the problem
    prob_nn = ODEProblem(nn_dynamics,x0, (Float64(1),Float64(size(training_data,2))), p64)
    ## Function to train the network
    # Define a predictor
    function predict(p, X = x0)
        _prob = remake(prob_nn, u0 = X, tspan = (Float64(1),Float64(size(training_data,2))), p = p)
        solve(_prob, AutoTsit5(Rosenbrock23()), saveat = 1.,
                abstol=1e-6, reltol=1e-6
                )
    end


    lipschitz_regularizer = 0.5
    function loss_function(p)
	   W1 = p.layer_1.weight
	    W2 = p.layer_2.weight
          lipschitz_constant = spectralRadius(W1)*spectralRadius(W2)


        pred = predict(p)
        loss = sum(abs2,training_data .- pred)/size(training_data,2) + lipschitz_regularizer*lipschitz_constant
        return loss
    end


    ## Training

    adtype = Optimization.AutoZygote()
    optf = Optimization.OptimizationFunction((x, p) -> loss_function(x), adtype)
    optprob = Optimization.OptimizationProblem(optf, p64)

    result_neuralode = Optimization.solve(optprob,
                                           ADAM(),
                                           maxiters = 300)

    optprob2 = remake(optprob,u0 = result_neuralode.u)
    result_neuralode2 = Optimization.solve(optprob2,
                                            Optim.BFGS(initial_stepnorm=0.01),
                                            allow_f_increases = false)


    return result_neuralode2.u
end

function spectralRadius(X,niters=10)
    y = randn!(similar(X, size(X, 2)))
    tmp = X * y
    for i in 1:niters
        tmp = X*y
        tmp = tmp / norm(tmp)
        y = X' * tmp
        y = y / norm(y)
    end
    return norm(X*y)
end

training_data = rand(4,10)
neuralnetwork = Lux.Chain(Lux.Dense(4,5),Lux.Dense(5,4))
knownDynamics(x,p,q)=-q
trainUDEModel(neuralnetwork,knownDynamics,training_data;p_true=1)
println("Done!")

runs no problem, although it is the same as running the other code without CUDA (which should be the same as running the one with CUDA without GPU).

[darsnack]

I'm not entirely sure but I think @allowscalar inserts a try block? That might help narrow things down.

[jarroyoe]

Thanks, that actually progressed the code a little bit. Now I got into a weird error regarding Rosenbrock23 compatibility with CUDA. It is worth mentioning the code now runs without GPUs.

using OrdinaryDiffEq, DiffEqFlux
using Optimization, OptimizationOptimisers, OptimizationOptimJL
using ComponentArrays, Lux, Zygote, Random, CUDA, LinearAlgebra
rng = Random.default_rng()

function trainUDEModel(neuralNetwork,knownDynamics,training_data;needed_ps = Float64[],p_true = Float64[])
    pinit, st = Lux.setup(rng,neuralNetwork)
    st = st |> Lux.gpu
    p64 = Float64.(Lux.gpu(ComponentArray(pinit)))
    training_data = Float64.(Lux.gpu(training_data))
    x0 = Float64.(Lux.gpu(training_data[:,1]))

   function ude(du,u,p,t,q)
        knownPred = convert(CuArray,knownDynamics(u,nothing,q))
        nnPred = convert(CuArray,first(neuralNetwork(u,p,st)))

        du .= knownPred .+ nnPred
    end

    # Closure with the known parameter
    nn_dynamics(du,u,p,t) = ude(du,u,p,t,p_true)
    # Define the problem
    prob_nn = ODEProblem(nn_dynamics,x0, (Float64(1),Float64(size(training_data,2))), p64)
    ## Function to train the network
    # Define a predictor
    function predict(p, X = x0)
        _prob = remake(prob_nn, u0 = X, tspan = (Float64(1),Float64(size(training_data,2))), p = p)
	convert(CuArray,solve(_prob, AutoTsit5(Rosenbrock23()), saveat = 1.,
                abstol=1e-6, reltol=1e-6))
    end

    lipschitz_regularizer = 0.5
    function loss_function(p)
	    W1 = p.layer_1.weight
	    W2 = p.layer_2.weight
        lipschitz_constant = spectralRadius(W1)*spectralRadius(W2)

        pred = predict(p)
        loss = sum(abs2,training_data .- pred)/size(training_data,2) + lipschitz_regularizer*lipschitz_constant
        return loss
    end

    adtype = Optimization.AutoZygote()
    optf = Optimization.OptimizationFunction((x, p) -> loss_function(x), adtype)
    optprob = Optimization.OptimizationProblem(optf, p64)

    result_neuralode = Optimization.solve(optprob,
                                           ADAM(),
                                           #callback = callback,
                                           maxiters = 300)

    optprob2 = remake(optprob,u0 = result_neuralode.u)
    result_neuralode2 = Optimization.solve(optprob2,
                                            Optim.BFGS(initial_stepnorm=0.01),
                                            #callback=callback,
                                            allow_f_increases = false)


    return result_neuralode2.u
end

function spectralRadius(X,niters=10)
    y = randn!(similar(X, size(X, 2)))
    tmp = X * y
    for i in 1:niters
        tmp = X*y
        tmp = tmp / norm(tmp)
        y = X' * tmp
        y = y / norm(y)
    end
    return norm(X*y)
end

training_data = rand(4,10)
neuralnetwork = Lux.Chain(Lux.Dense(4,5),Lux.Dense(5,4))
knownDynamics(x,p,q)=[-q;-q;-q;-q]
trainUDEModel(neuralnetwork,knownDynamics,training_data;p_true=1)
println("Done!")

ERROR: LoadError: MethodError: no method matching geqrf!(::CuArray{ForwardDiff.Dual{ForwardDiff.Tag{ODEFunction{true, SciMLBase.AutoSpecialize, var"#nn_dynamics#4"{Int64, var"#ude#3"{Lux.Chain{NamedTuple{(:layer_1, :layer_2), Tuple{Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}, Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}}}}, typeof(knownDynamics)}}, UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing}, Float64}, Float64, 4}, 2, CUDA.Mem.DeviceBuffer})
Closest candidates are:
  geqrf!(!Matched::StridedCuMatrix{Float32}) at ~/.julia/packages/CUDA/OYQsb/lib/cusolver/dense.jl:151
  geqrf!(!Matched::StridedCuMatrix{Float64}) at ~/.julia/packages/CUDA/OYQsb/lib/cusolver/dense.jl:151
  geqrf!(!Matched::StridedCuMatrix{ComplexF32}) at ~/.julia/packages/CUDA/OYQsb/lib/cusolver/dense.jl:151
  ...
Stacktrace:
  [1] qr!(A::CuArray{ForwardDiff.Dual{ForwardDiff.Tag{ODEFunction{true, SciMLBase.AutoSpecialize, var"#nn_dynamics#4"{Int64, var"#ude#3"{Lux.Chain{NamedTuple{(:layer_1, :layer_2), Tuple{Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}, Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}}}}, typeof(knownDynamics)}}, UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing}, Float64}, Float64, 4}, 2, CUDA.Mem.DeviceBuffer})
    @ CUDA.CUSOLVER ~/.julia/packages/CUDA/OYQsb/lib/cusolver/linalg.jl:80
  [2] qr(::CuArray{ForwardDiff.Dual{ForwardDiff.Tag{ODEFunction{true, SciMLBase.AutoSpecialize, var"#nn_dynamics#4"{Int64, var"#ude#3"{Lux.Chain{NamedTuple{(:layer_1, :layer_2), Tuple{Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}, Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}}}}, typeof(knownDynamics)}}, UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing}, Float64}, Float64, 4}, 2, CUDA.Mem.DeviceBuffer}; kwargs::Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}})
    @ LinearAlgebra /carnegie/binaries/centos7/julia/1.8.5/share/julia/stdlib/v1.8/LinearAlgebra/src/qr.jl:428
  [3] qr(::CuArray{ForwardDiff.Dual{ForwardDiff.Tag{ODEFunction{true, SciMLBase.AutoSpecialize, var"#nn_dynamics#4"{Int64, var"#ude#3"{Lux.Chain{NamedTuple{(:layer_1, :layer_2), Tuple{Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}, Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}}}}, typeof(knownDynamics)}}, UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing}, Float64}, Float64, 4}, 2, CUDA.Mem.DeviceBuffer})
    @ LinearAlgebra /carnegie/binaries/centos7/julia/1.8.5/share/julia/stdlib/v1.8/LinearAlgebra/src/qr.jl:425
  [4] qr_instance(a::CuArray{ForwardDiff.Dual{ForwardDiff.Tag{ODEFunction{true, SciMLBase.AutoSpecialize, var"#nn_dynamics#4"{Int64, var"#ude#3"{Lux.Chain{NamedTuple{(:layer_1, :layer_2), Tuple{Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}, Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}}}}, typeof(knownDynamics)}}, UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing}, Float64}, Float64, 4}, 2, CUDA.Mem.DeviceBuffer}, pivot::NoPivot) (repeats 2 times)
    @ ArrayInterface ~/.julia/packages/ArrayInterface/KwW6K/src/ArrayInterface.jl:664
  [5] init_cacheval(alg::LinearSolve.QRFactorization{NoPivot}, A::CuArray{ForwardDiff.Dual{ForwardDiff.Tag{ODEFunction{true, SciMLBase.AutoSpecialize, var"#nn_dynamics#4"{Int64, var"#ude#3"{Lux.Chain{NamedTuple{(:layer_1, :layer_2), Tuple{Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}, Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}}}}, typeof(knownDynamics)}}, UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing}, Float64}, Float64, 4}, 2, CUDA.Mem.DeviceBuffer}, b::CuArray{ForwardDiff.Dual{ForwardDiff.Tag{ODEFunction{true, SciMLBase.AutoSpecialize, var"#nn_dynamics#4"{Int64, var"#ude#3"{Lux.Chain{NamedTuple{(:layer_1, :layer_2), Tuple{Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}, Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}}}}, typeof(knownDynamics)}}, UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing}, Float64}, Float64, 4}, 1, CUDA.Mem.DeviceBuffer}, u::CuArray{ForwardDiff.Dual{ForwardDiff.Tag{ODEFunction{true, SciMLBase.AutoSpecialize, var"#nn_dynamics#4"{Int64, var"#ude#3"{Lux.Chain{NamedTuple{(:layer_1, :layer_2), Tuple{Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}, Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}}}}, typeof(knownDynamics)}}, UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing}, Float64}, Float64, 4}, 1, CUDA.Mem.DeviceBuffer}, Pl::LinearSolve.InvPreconditioner{Diagonal{ForwardDiff.Dual{ForwardDiff.Tag{ODEFunction{true, SciMLBase.AutoSpecialize, var"#nn_dynamics#4"{Int64, var"#ude#3"{Lux.Chain{NamedTuple{(:layer_1, :layer_2), Tuple{Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}, Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}}}}, typeof(knownDynamics)}}, UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing}, Float64}, Float64, 4}, CuArray{ForwardDiff.Dual{ForwardDiff.Tag{ODEFunction{true, SciMLBase.AutoSpecialize, var"#nn_dynamics#4"{Int64, var"#ude#3"{Lux.Chain{NamedTuple{(:layer_1, :layer_2), Tuple{Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}, Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}}}}, typeof(knownDynamics)}}, UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing}, Float64}, Float64, 4}, 1, CUDA.Mem.DeviceBuffer}}}, Pr::Diagonal{ForwardDiff.Dual{ForwardDiff.Tag{ODEFunction{true, SciMLBase.AutoSpecialize, var"#nn_dynamics#4"{Int64, var"#ude#3"{Lux.Chain{NamedTuple{(:layer_1, :layer_2), Tuple{Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}, Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}}}}, typeof(knownDynamics)}}, UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing}, Float64}, Float64, 4}, CuArray{ForwardDiff.Dual{ForwardDiff.Tag{ODEFunction{true, SciMLBase.AutoSpecialize, var"#nn_dynamics#4"{Int64, var"#ude#3"{Lux.Chain{NamedTuple{(:layer_1, :layer_2), Tuple{Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}, Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}}}}, typeof(knownDynamics)}}, UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing}, Float64}, Float64, 4}, 1, CUDA.Mem.DeviceBuffer}}, maxiters::Int64, abstol::Float64, reltol::Float64, verbose::Bool, assumptions::LinearSolve.OperatorAssumptions{true, LinearSolve.OperatorCondition.IllConditioned})
    @ LinearSolve ~/.julia/packages/LinearSolve/LD2dF/src/factorization.jl:105
  [6] init(::LinearProblem{CuArray{ForwardDiff.Dual{ForwardDiff.Tag{ODEFunction{true, SciMLBase.AutoSpecialize, var"#nn_dynamics#4"{Int64, var"#ude#3"{Lux.Chain{NamedTuple{(:layer_1, :layer_2), Tuple{Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}, Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}}}}, typeof(knownDynamics)}}, UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing}, Float64}, Float64, 4}, 1, CUDA.Mem.DeviceBuffer}, true, CuArray{ForwardDiff.Dual{ForwardDiff.Tag{ODEFunction{true, SciMLBase.AutoSpecialize, var"#nn_dynamics#4"{Int64, var"#ude#3"{Lux.Chain{NamedTuple{(:layer_1, :layer_2), Tuple{Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}, Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}}}}, typeof(knownDynamics)}}, UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing}, Float64}, Float64, 4}, 2, CUDA.Mem.DeviceBuffer}, CuArray{ForwardDiff.Dual{ForwardDiff.Tag{ODEFunction{true, SciMLBase.AutoSpecialize, var"#nn_dynamics#4"{Int64, var"#ude#3"{Lux.Chain{NamedTuple{(:layer_1, :layer_2), Tuple{Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}, Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}}}}, typeof(knownDynamics)}}, UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing}, Float64}, Float64, 4}, 1, CUDA.Mem.DeviceBuffer}, SciMLBase.NullParameters, Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}}, ::LinearSolve.QRFactorization{NoPivot}; alias_A::Bool, alias_b::Bool, abstol::Float64, reltol::Float64, maxiters::Int64, verbose::Bool, Pl::LinearSolve.InvPreconditioner{Diagonal{ForwardDiff.Dual{ForwardDiff.Tag{ODEFunction{true, SciMLBase.AutoSpecialize, var"#nn_dynamics#4"{Int64, var"#ude#3"{Lux.Chain{NamedTuple{(:layer_1, :layer_2), Tuple{Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}, Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}}}}, typeof(knownDynamics)}}, UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing}, Float64}, Float64, 4}, CuArray{ForwardDiff.Dual{ForwardDiff.Tag{ODEFunction{true, SciMLBase.AutoSpecialize, var"#nn_dynamics#4"{Int64, var"#ude#3"{Lux.Chain{NamedTuple{(:layer_1, :layer_2), Tuple{Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}, Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}}}}, typeof(knownDynamics)}}, UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing}, Float64}, Float64, 4}, 1, CUDA.Mem.DeviceBuffer}}}, Pr::Diagonal{ForwardDiff.Dual{ForwardDiff.Tag{ODEFunction{true, SciMLBase.AutoSpecialize, var"#nn_dynamics#4"{Int64, var"#ude#3"{Lux.Chain{NamedTuple{(:layer_1, :layer_2), Tuple{Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}, Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}}}}, typeof(knownDynamics)}}, UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing}, Float64}, Float64, 4}, CuArray{ForwardDiff.Dual{ForwardDiff.Tag{ODEFunction{true, SciMLBase.AutoSpecialize, var"#nn_dynamics#4"{Int64, var"#ude#3"{Lux.Chain{NamedTuple{(:layer_1, :layer_2), Tuple{Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}, Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}}}}, typeof(knownDynamics)}}, UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing}, Float64}, Float64, 4}, 1, CUDA.Mem.DeviceBuffer}}, assumptions::LinearSolve.OperatorAssumptions{true, LinearSolve.OperatorCondition.IllConditioned}, kwargs::Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}})
    @ LinearSolve ~/.julia/packages/LinearSolve/LD2dF/src/common.jl:179
  [7] alg_cache(alg::Rosenbrock23{1,true,LinearSolve.QRFactorization{NoPivot},typeof(OrdinaryDiffEq.DEFAULT_PRECS),Val{:forward},true,nothing}, u::CuArray{ForwardDiff.Dual{ForwardDiff.Tag{ODEFunction{true, SciMLBase.AutoSpecialize, var"#nn_dynamics#4"{Int64, var"#ude#3"{Lux.Chain{NamedTuple{(:layer_1, :layer_2), Tuple{Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}, Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}}}}, typeof(knownDynamics)}}, UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing}, Float64}, Float64, 4}, 1, CUDA.Mem.DeviceBuffer}, rate_prototype::CuArray{ForwardDiff.Dual{ForwardDiff.Tag{ODEFunction{true, SciMLBase.AutoSpecialize, var"#nn_dynamics#4"{Int64, var"#ude#3"{Lux.Chain{NamedTuple{(:layer_1, :layer_2), Tuple{Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}, Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}}}}, typeof(knownDynamics)}}, UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing}, Float64}, Float64, 4}, 1, CUDA.Mem.DeviceBuffer}, #unused#::Type{ForwardDiff.Dual{ForwardDiff.Tag{ODEFunction{true, SciMLBase.AutoSpecialize, var"#nn_dynamics#4"{Int64, var"#ude#3"{Lux.Chain{NamedTuple{(:layer_1, :layer_2), Tuple{Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}, Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}}}}, typeof(knownDynamics)}}, UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing}, Float64}, Float64, 4}}, #unused#::Type{ForwardDiff.Dual{ForwardDiff.Tag{ODEFunction{true, SciMLBase.AutoSpecialize, var"#nn_dynamics#4"{Int64, var"#ude#3"{Lux.Chain{NamedTuple{(:layer_1, :layer_2), Tuple{Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}, Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}}}}, typeof(knownDynamics)}}, UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing}, Float64}, Float64, 4}}, #unused#::Type{Float64}, uprev::CuArray{ForwardDiff.Dual{ForwardDiff.Tag{ODEFunction{true, SciMLBase.AutoSpecialize, var"#nn_dynamics#4"{Int64, var"#ude#3"{Lux.Chain{NamedTuple{(:layer_1, :layer_2), Tuple{Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}, Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}}}}, typeof(knownDynamics)}}, UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing}, Float64}, Float64, 4}, 1, CUDA.Mem.DeviceBuffer}, uprev2::CuArray{ForwardDiff.Dual{ForwardDiff.Tag{ODEFunction{true, SciMLBase.AutoSpecialize, var"#nn_dynamics#4"{Int64, var"#ude#3"{Lux.Chain{NamedTuple{(:layer_1, :layer_2), Tuple{Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}, Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}}}}, typeof(knownDynamics)}}, UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing}, Float64}, Float64, 4}, 1, CUDA.Mem.DeviceBuffer}, f::ODEFunction{true, SciMLBase.FullSpecialize, ODEFunction{true, SciMLBase.AutoSpecialize, var"#nn_dynamics#4"{Int64, var"#ude#3"{Lux.Chain{NamedTuple{(:layer_1, :layer_2), Tuple{Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}, Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}}}}, typeof(knownDynamics)}}, UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing}, UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing}, t::Float64, dt::Float64, reltol::Float64, p::ComponentVector{ForwardDiff.Dual{ForwardDiff.Tag{ODEFunction{true, SciMLBase.AutoSpecialize, var"#nn_dynamics#4"{Int64, var"#ude#3"{Lux.Chain{NamedTuple{(:layer_1, :layer_2), Tuple{Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}, Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}}}}, typeof(knownDynamics)}}, UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing}, Float64}, Float64, 4}, CuArray{ForwardDiff.Dual{ForwardDiff.Tag{ODEFunction{true, SciMLBase.AutoSpecialize, var"#nn_dynamics#4"{Int64, var"#ude#3"{Lux.Chain{NamedTuple{(:layer_1, :layer_2), Tuple{Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}, Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}}}}, typeof(knownDynamics)}}, UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing}, Float64}, Float64, 4}, 1, CUDA.Mem.DeviceBuffer}, Tuple{Axis{(layer_1 = ViewAxis(1:25, Axis(weight = ViewAxis(1:20, ShapedAxis((5, 4), NamedTuple())), bias = ViewAxis(21:25, ShapedAxis((5, 1), NamedTuple())))), layer_2 = ViewAxis(26:49, Axis(weight = ViewAxis(1:20, ShapedAxis((4, 5), NamedTuple())), bias = ViewAxis(21:24, ShapedAxis((4, 1), NamedTuple())))))}}}, calck::Bool, #unused#::Val{true})
    @ OrdinaryDiffEq ~/.julia/packages/OrdinaryDiffEq/yspeT/src/caches/rosenbrock_caches.jl:114
  [8] macro expansion
    @ ~/.julia/packages/OrdinaryDiffEq/yspeT/src/caches/basic_caches.jl:0 [inlined]
  [9] __alg_cache(algs::Tuple{Tsit5{typeof(OrdinaryDiffEq.trivial_limiter!),typeof(OrdinaryDiffEq.trivial_limiter!),Static.False}, Rosenbrock23{1,true,LinearSolve.QRFactorization{NoPivot},typeof(OrdinaryDiffEq.DEFAULT_PRECS),Val{:forward},true,nothing}}, u::CuArray{ForwardDiff.Dual{ForwardDiff.Tag{ODEFunction{true, SciMLBase.AutoSpecialize, var"#nn_dynamics#4"{Int64, var"#ude#3"{Lux.Chain{NamedTuple{(:layer_1, :layer_2), Tuple{Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}, Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}}}}, typeof(knownDynamics)}}, UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing}, Float64}, Float64, 4}, 1, CUDA.Mem.DeviceBuffer}, rate_prototype::CuArray{ForwardDiff.Dual{ForwardDiff.Tag{ODEFunction{true, SciMLBase.AutoSpecialize, var"#nn_dynamics#4"{Int64, var"#ude#3"{Lux.Chain{NamedTuple{(:layer_1, :layer_2), Tuple{Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}, Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}}}}, typeof(knownDynamics)}}, UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing}, Float64}, Float64, 4}, 1, CUDA.Mem.DeviceBuffer}, #unused#::Type{ForwardDiff.Dual{ForwardDiff.Tag{ODEFunction{true, SciMLBase.AutoSpecialize, var"#nn_dynamics#4"{Int64, var"#ude#3"{Lux.Chain{NamedTuple{(:layer_1, :layer_2), Tuple{Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}, Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}}}}, typeof(knownDynamics)}}, UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing}, Float64}, Float64, 4}}, #unused#::Type{ForwardDiff.Dual{ForwardDiff.Tag{ODEFunction{true, SciMLBase.AutoSpecialize, var"#nn_dynamics#4"{Int64, var"#ude#3"{Lux.Chain{NamedTuple{(:layer_1, :layer_2), Tuple{Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}, Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}}}}, typeof(knownDynamics)}}, UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing}, Float64}, Float64, 4}}, #unused#::Type{Float64}, uprev::CuArray{ForwardDiff.Dual{ForwardDiff.Tag{ODEFunction{true, SciMLBase.AutoSpecialize, var"#nn_dynamics#4"{Int64, var"#ude#3"{Lux.Chain{NamedTuple{(:layer_1, :layer_2), Tuple{Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}, Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}}}}, typeof(knownDynamics)}}, UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing}, Float64}, Float64, 4}, 1, CUDA.Mem.DeviceBuffer}, uprev2::CuArray{ForwardDiff.Dual{ForwardDiff.Tag{ODEFunction{true, SciMLBase.AutoSpecialize, var"#nn_dynamics#4"{Int64, var"#ude#3"{Lux.Chain{NamedTuple{(:layer_1, :layer_2), Tuple{Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}, Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}}}}, typeof(knownDynamics)}}, UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing}, Float64}, Float64, 4}, 1, CUDA.Mem.DeviceBuffer}, f::ODEFunction{true, SciMLBase.FullSpecialize, ODEFunction{true, SciMLBase.AutoSpecialize, var"#nn_dynamics#4"{Int64, var"#ude#3"{Lux.Chain{NamedTuple{(:layer_1, :layer_2), Tuple{Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}, Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}}}}, typeof(knownDynamics)}}, UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing}, UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing}, t::Float64, dt::Float64, reltol::Float64, p::ComponentVector{ForwardDiff.Dual{ForwardDiff.Tag{ODEFunction{true, SciMLBase.AutoSpecialize, var"#nn_dynamics#4"{Int64, var"#ude#3"{Lux.Chain{NamedTuple{(:layer_1, :layer_2), Tuple{Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}, Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}}}}, typeof(knownDynamics)}}, UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing}, Float64}, Float64, 4}, CuArray{ForwardDiff.Dual{ForwardDiff.Tag{ODEFunction{true, SciMLBase.AutoSpecialize, var"#nn_dynamics#4"{Int64, var"#ude#3"{Lux.Chain{NamedTuple{(:layer_1, :layer_2), Tuple{Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}, Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}}}}, typeof(knownDynamics)}}, UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing}, Float64}, Float64, 4}, 1, CUDA.Mem.DeviceBuffer}, Tuple{Axis{(layer_1 = ViewAxis(1:25, Axis(weight = ViewAxis(1:20, ShapedAxis((5, 4), NamedTuple())), bias = ViewAxis(21:25, ShapedAxis((5, 1), NamedTuple())))), layer_2 = ViewAxis(26:49, Axis(weight = ViewAxis(1:20, ShapedAxis((4, 5), NamedTuple())), bias = ViewAxis(21:24, ShapedAxis((4, 1), NamedTuple())))))}}}, calck::Bool, #unused#::Val{true})
    @ OrdinaryDiffEq ~/.julia/packages/OrdinaryDiffEq/yspeT/src/caches/basic_caches.jl:50
 [10] alg_cache(alg::CompositeAlgorithm{Tuple{Tsit5{typeof(OrdinaryDiffEq.trivial_limiter!),typeof(OrdinaryDiffEq.trivial_limiter!),Static.False}, Rosenbrock23{1,true,LinearSolve.QRFactorization{NoPivot},typeof(OrdinaryDiffEq.DEFAULT_PRECS),Val{:forward},true,nothing}},OrdinaryDiffEq.AutoSwitchCache{Tsit5{typeof(OrdinaryDiffEq.trivial_limiter!),typeof(OrdinaryDiffEq.trivial_limiter!),Static.False}, Rosenbrock23{0,true,Nothing,typeof(OrdinaryDiffEq.DEFAULT_PRECS),Val{:forward},true,nothing}, Rational{Int64}, Int64}}, u::CuArray{ForwardDiff.Dual{ForwardDiff.Tag{ODEFunction{true, SciMLBase.AutoSpecialize, var"#nn_dynamics#4"{Int64, var"#ude#3"{Lux.Chain{NamedTuple{(:layer_1, :layer_2), Tuple{Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}, Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}}}}, typeof(knownDynamics)}}, UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing}, Float64}, Float64, 4}, 1, CUDA.Mem.DeviceBuffer}, rate_prototype::CuArray{ForwardDiff.Dual{ForwardDiff.Tag{ODEFunction{true, SciMLBase.AutoSpecialize, var"#nn_dynamics#4"{Int64, var"#ude#3"{Lux.Chain{NamedTuple{(:layer_1, :layer_2), Tuple{Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}, Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}}}}, typeof(knownDynamics)}}, UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing}, Float64}, Float64, 4}, 1, CUDA.Mem.DeviceBuffer}, #unused#::Type{ForwardDiff.Dual{ForwardDiff.Tag{ODEFunction{true, SciMLBase.AutoSpecialize, var"#nn_dynamics#4"{Int64, var"#ude#3"{Lux.Chain{NamedTuple{(:layer_1, :layer_2), Tuple{Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}, Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}}}}, typeof(knownDynamics)}}, UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing}, Float64}, Float64, 4}}, #unused#::Type{ForwardDiff.Dual{ForwardDiff.Tag{ODEFunction{true, SciMLBase.AutoSpecialize, var"#nn_dynamics#4"{Int64, var"#ude#3"{Lux.Chain{NamedTuple{(:layer_1, :layer_2), Tuple{Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}, Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}}}}, typeof(knownDynamics)}}, UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing}, Float64}, Float64, 4}}, #unused#::Type{Float64}, uprev::CuArray{ForwardDiff.Dual{ForwardDiff.Tag{ODEFunction{true, SciMLBase.AutoSpecialize, var"#nn_dynamics#4"{Int64, var"#ude#3"{Lux.Chain{NamedTuple{(:layer_1, :layer_2), Tuple{Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}, Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}}}}, typeof(knownDynamics)}}, UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing}, Float64}, Float64, 4}, 1, CUDA.Mem.DeviceBuffer}, uprev2::CuArray{ForwardDiff.Dual{ForwardDiff.Tag{ODEFunction{true, SciMLBase.AutoSpecialize, var"#nn_dynamics#4"{Int64, var"#ude#3"{Lux.Chain{NamedTuple{(:layer_1, :layer_2), Tuple{Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}, Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}}}}, typeof(knownDynamics)}}, UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing}, Float64}, Float64, 4}, 1, CUDA.Mem.DeviceBuffer}, f::ODEFunction{true, SciMLBase.FullSpecialize, ODEFunction{true, SciMLBase.AutoSpecialize, var"#nn_dynamics#4"{Int64, var"#ude#3"{Lux.Chain{NamedTuple{(:layer_1, :layer_2), Tuple{Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}, Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}}}}, typeof(knownDynamics)}}, UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing}, UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing}, t::Float64, dt::Float64, reltol::Float64, p::ComponentVector{ForwardDiff.Dual{ForwardDiff.Tag{ODEFunction{true, SciMLBase.AutoSpecialize, var"#nn_dynamics#4"{Int64, var"#ude#3"{Lux.Chain{NamedTuple{(:layer_1, :layer_2), Tuple{Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}, Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}}}}, typeof(knownDynamics)}}, UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing}, Float64}, Float64, 4}, CuArray{ForwardDiff.Dual{ForwardDiff.Tag{ODEFunction{true, SciMLBase.AutoSpecialize, var"#nn_dynamics#4"{Int64, var"#ude#3"{Lux.Chain{NamedTuple{(:layer_1, :layer_2), Tuple{Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}, Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}}}}, typeof(knownDynamics)}}, UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing}, Float64}, Float64, 4}, 1, CUDA.Mem.DeviceBuffer}, Tuple{Axis{(layer_1 = ViewAxis(1:25, Axis(weight = ViewAxis(1:20, ShapedAxis((5, 4), NamedTuple())), bias = ViewAxis(21:25, ShapedAxis((5, 1), NamedTuple())))), layer_2 = ViewAxis(26:49, Axis(weight = ViewAxis(1:20, ShapedAxis((4, 5), NamedTuple())), bias = ViewAxis(21:24, ShapedAxis((4, 1), NamedTuple())))))}}}, calck::Bool, #unused#::Val{true})
    @ OrdinaryDiffEq ~/.julia/packages/OrdinaryDiffEq/yspeT/src/caches/basic_caches.jl:31
 [11] __init(prob::ODEProblem{CuArray{ForwardDiff.Dual{ForwardDiff.Tag{ODEFunction{true, SciMLBase.AutoSpecialize, var"#nn_dynamics#4"{Int64, var"#ude#3"{Lux.Chain{NamedTuple{(:layer_1, :layer_2), Tuple{Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}, Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}}}}, typeof(knownDynamics)}}, UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing}, Float64}, Float64, 4}, 1, CUDA.Mem.DeviceBuffer}, Tuple{Float64, Float64}, true, ComponentVector{ForwardDiff.Dual{ForwardDiff.Tag{ODEFunction{true, SciMLBase.AutoSpecialize, var"#nn_dynamics#4"{Int64, var"#ude#3"{Lux.Chain{NamedTuple{(:layer_1, :layer_2), Tuple{Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}, Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}}}}, typeof(knownDynamics)}}, UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing}, Float64}, Float64, 4}, CuArray{ForwardDiff.Dual{ForwardDiff.Tag{ODEFunction{true, SciMLBase.AutoSpecialize, var"#nn_dynamics#4"{Int64, var"#ude#3"{Lux.Chain{NamedTuple{(:layer_1, :layer_2), Tuple{Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}, Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}}}}, typeof(knownDynamics)}}, UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing}, Float64}, Float64, 4}, 1, CUDA.Mem.DeviceBuffer}, Tuple{Axis{(layer_1 = ViewAxis(1:25, Axis(weight = ViewAxis(1:20, ShapedAxis((5, 4), NamedTuple())), bias = ViewAxis(21:25, ShapedAxis((5, 1), NamedTuple())))), layer_2 = ViewAxis(26:49, Axis(weight = ViewAxis(1:20, ShapedAxis((4, 5), NamedTuple())), bias = ViewAxis(21:24, ShapedAxis((4, 1), NamedTuple())))))}}}, ODEFunction{true, SciMLBase.FullSpecialize, ODEFunction{true, SciMLBase.AutoSpecialize, var"#nn_dynamics#4"{Int64, var"#ude#3"{Lux.Chain{NamedTuple{(:layer_1, :layer_2), Tuple{Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}, Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}}}}, typeof(knownDynamics)}}, UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing}, UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing}, Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}, SciMLBase.StandardODEProblem}, alg::CompositeAlgorithm{Tuple{Tsit5{typeof(OrdinaryDiffEq.trivial_limiter!),typeof(OrdinaryDiffEq.trivial_limiter!),Static.False}, Rosenbrock23{1,true,LinearSolve.QRFactorization{NoPivot},typeof(OrdinaryDiffEq.DEFAULT_PRECS),Val{:forward},true,nothing}},AutoSwitch{Tsit5{typeof(OrdinaryDiffEq.trivial_limiter!),typeof(OrdinaryDiffEq.trivial_limiter!),Static.False}, Rosenbrock23{0,true,Nothing,typeof(OrdinaryDiffEq.DEFAULT_PRECS),Val{:forward},true,nothing}, Rational{Int64}, Int64}}, timeseries_init::Tuple{}, ts_init::Tuple{}, ks_init::Tuple{}, recompile::Type{Val{true}}; saveat::Vector{Float64}, tstops::Tuple{}, d_discontinuities::Tuple{}, save_idxs::Nothing, save_everystep::Bool, save_on::Bool, save_start::Bool, save_end::Nothing, callback::Nothing, dense::Bool, calck::Bool, dt::Float64, dtmin::Nothing, dtmax::Float64, force_dtmin::Bool, adaptive::Bool, gamma::Rational{Int64}, abstol::Float64, reltol::Float64, qmin::Rational{Int64}, qmax::Int64, qsteady_min::Int64, qsteady_max::Int64, beta1::Nothing, beta2::Nothing, qoldinit::Rational{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), userdata::Nothing, allow_extrapolation::Bool, initialize_integrator::Bool, alias_u0::Bool, alias_du0::Bool, initializealg::OrdinaryDiffEq.DefaultInit, kwargs::Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}})
    @ OrdinaryDiffEq ~/.julia/packages/OrdinaryDiffEq/yspeT/src/solve.jl:322
 [12] __solve(::ODEProblem{CuArray{ForwardDiff.Dual{ForwardDiff.Tag{ODEFunction{true, SciMLBase.AutoSpecialize, var"#nn_dynamics#4"{Int64, var"#ude#3"{Lux.Chain{NamedTuple{(:layer_1, :layer_2), Tuple{Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}, Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}}}}, typeof(knownDynamics)}}, UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing}, Float64}, Float64, 4}, 1, CUDA.Mem.DeviceBuffer}, Tuple{Float64, Float64}, true, ComponentVector{ForwardDiff.Dual{ForwardDiff.Tag{ODEFunction{true, SciMLBase.AutoSpecialize, var"#nn_dynamics#4"{Int64, var"#ude#3"{Lux.Chain{NamedTuple{(:layer_1, :layer_2), Tuple{Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}, Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}}}}, typeof(knownDynamics)}}, UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing}, Float64}, Float64, 4}, CuArray{ForwardDiff.Dual{ForwardDiff.Tag{ODEFunction{true, SciMLBase.AutoSpecialize, var"#nn_dynamics#4"{Int64, var"#ude#3"{Lux.Chain{NamedTuple{(:layer_1, :layer_2), Tuple{Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}, Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}}}}, typeof(knownDynamics)}}, UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing}, Float64}, Float64, 4}, 1, CUDA.Mem.DeviceBuffer}, Tuple{Axis{(layer_1 = ViewAxis(1:25, Axis(weight = ViewAxis(1:20, ShapedAxis((5, 4), NamedTuple())), bias = ViewAxis(21:25, ShapedAxis((5, 1), NamedTuple())))), layer_2 = ViewAxis(26:49, Axis(weight = ViewAxis(1:20, ShapedAxis((4, 5), NamedTuple())), bias = ViewAxis(21:24, ShapedAxis((4, 1), NamedTuple())))))}}}, ODEFunction{true, SciMLBase.FullSpecialize, ODEFunction{true, SciMLBase.AutoSpecialize, var"#nn_dynamics#4"{Int64, var"#ude#3"{Lux.Chain{NamedTuple{(:layer_1, :layer_2), Tuple{Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}, Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}}}}, typeof(knownDynamics)}}, UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing}, UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing}, Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}, SciMLBase.StandardODEProblem}, ::CompositeAlgorithm{Tuple{Tsit5{typeof(OrdinaryDiffEq.trivial_limiter!),typeof(OrdinaryDiffEq.trivial_limiter!),Static.False}, Rosenbrock23{1,true,LinearSolve.QRFactorization{NoPivot},typeof(OrdinaryDiffEq.DEFAULT_PRECS),Val{:forward},true,nothing}},AutoSwitch{Tsit5{typeof(OrdinaryDiffEq.trivial_limiter!),typeof(OrdinaryDiffEq.trivial_limiter!),Static.False}, Rosenbrock23{0,true,Nothing,typeof(OrdinaryDiffEq.DEFAULT_PRECS),Val{:forward},true,nothing}, Rational{Int64}, Int64}}; kwargs::Base.Pairs{Symbol, Any, NTuple{4, Symbol}, NamedTuple{(:saveat, :verbose, :abstol, :reltol), Tuple{Vector{Float64}, Bool, Float64, Float64}}})
    @ OrdinaryDiffEq ~/.julia/packages/OrdinaryDiffEq/yspeT/src/solve.jl:5
 [13] solve_call(_prob::ODEProblem{CuArray{ForwardDiff.Dual{ForwardDiff.Tag{ODEFunction{true, SciMLBase.AutoSpecialize, var"#nn_dynamics#4"{Int64, var"#ude#3"{Lux.Chain{NamedTuple{(:layer_1, :layer_2), Tuple{Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}, Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}}}}, typeof(knownDynamics)}}, UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing}, Float64}, Float64, 4}, 1, CUDA.Mem.DeviceBuffer}, Tuple{Float64, Float64}, true, ComponentVector{ForwardDiff.Dual{ForwardDiff.Tag{ODEFunction{true, SciMLBase.AutoSpecialize, var"#nn_dynamics#4"{Int64, var"#ude#3"{Lux.Chain{NamedTuple{(:layer_1, :layer_2), Tuple{Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}, Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}}}}, typeof(knownDynamics)}}, UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing}, Float64}, Float64, 4}, CuArray{ForwardDiff.Dual{ForwardDiff.Tag{ODEFunction{true, SciMLBase.AutoSpecialize, var"#nn_dynamics#4"{Int64, var"#ude#3"{Lux.Chain{NamedTuple{(:layer_1, :layer_2), Tuple{Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}, Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}}}}, typeof(knownDynamics)}}, UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing}, Float64}, Float64, 4}, 1, CUDA.Mem.DeviceBuffer}, Tuple{Axis{(layer_1 = ViewAxis(1:25, Axis(weight = ViewAxis(1:20, ShapedAxis((5, 4), NamedTuple())), bias = ViewAxis(21:25, ShapedAxis((5, 1), NamedTuple())))), layer_2 = ViewAxis(26:49, Axis(weight = ViewAxis(1:20, ShapedAxis((4, 5), NamedTuple())), bias = ViewAxis(21:24, ShapedAxis((4, 1), NamedTuple())))))}}}, ODEFunction{true, SciMLBase.FullSpecialize, ODEFunction{true, SciMLBase.AutoSpecialize, var"#nn_dynamics#4"{Int64, var"#ude#3"{Lux.Chain{NamedTuple{(:layer_1, :layer_2), Tuple{Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}, Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}}}}, typeof(knownDynamics)}}, UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing}, UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing}, Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}, SciMLBase.StandardODEProblem}, args::CompositeAlgorithm{Tuple{Tsit5{typeof(OrdinaryDiffEq.trivial_limiter!),typeof(OrdinaryDiffEq.trivial_limiter!),Static.False}, Rosenbrock23{1,true,LinearSolve.QRFactorization{NoPivot},typeof(OrdinaryDiffEq.DEFAULT_PRECS),Val{:forward},true,nothing}},AutoSwitch{Tsit5{typeof(OrdinaryDiffEq.trivial_limiter!),typeof(OrdinaryDiffEq.trivial_limiter!),Static.False}, Rosenbrock23{0,true,Nothing,typeof(OrdinaryDiffEq.DEFAULT_PRECS),Val{:forward},true,nothing}, Rational{Int64}, Int64}}; merge_callbacks::Bool, kwargshandle::KeywordArgError, kwargs::Base.Pairs{Symbol, Any, NTuple{4, Symbol}, NamedTuple{(:saveat, :verbose, :abstol, :reltol), Tuple{Vector{Float64}, Bool, Float64, Float64}}})
    @ DiffEqBase ~/.julia/packages/DiffEqBase/s433k/src/solve.jl:559
 [14] solve_up(prob::ODEProblem{CuArray{ForwardDiff.Dual{ForwardDiff.Tag{ODEFunction{true, SciMLBase.AutoSpecialize, var"#nn_dynamics#4"{Int64, var"#ude#3"{Lux.Chain{NamedTuple{(:layer_1, :layer_2), Tuple{Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}, Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}}}}, typeof(knownDynamics)}}, UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing}, Float64}, Float64, 4}, 1, CUDA.Mem.DeviceBuffer}, Tuple{Float64, Float64}, true, ComponentVector{ForwardDiff.Dual{ForwardDiff.Tag{ODEFunction{true, SciMLBase.AutoSpecialize, var"#nn_dynamics#4"{Int64, var"#ude#3"{Lux.Chain{NamedTuple{(:layer_1, :layer_2), Tuple{Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}, Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}}}}, typeof(knownDynamics)}}, UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing}, Float64}, Float64, 4}, CuArray{ForwardDiff.Dual{ForwardDiff.Tag{ODEFunction{true, SciMLBase.AutoSpecialize, var"#nn_dynamics#4"{Int64, var"#ude#3"{Lux.Chain{NamedTuple{(:layer_1, :layer_2), Tuple{Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}, Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}}}}, typeof(knownDynamics)}}, UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing}, Float64}, Float64, 4}, 1, CUDA.Mem.DeviceBuffer}, Tuple{Axis{(layer_1 = ViewAxis(1:25, Axis(weight = ViewAxis(1:20, ShapedAxis((5, 4), NamedTuple())), bias = ViewAxis(21:25, ShapedAxis((5, 1), NamedTuple())))), layer_2 = ViewAxis(26:49, Axis(weight = ViewAxis(1:20, ShapedAxis((4, 5), NamedTuple())), bias = ViewAxis(21:24, ShapedAxis((4, 1), NamedTuple())))))}}}, ODEFunction{true, SciMLBase.FullSpecialize, ODEFunction{true, SciMLBase.AutoSpecialize, var"#nn_dynamics#4"{Int64, var"#ude#3"{Lux.Chain{NamedTuple{(:layer_1, :layer_2), Tuple{Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}, Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}}}}, typeof(knownDynamics)}}, UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing}, UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing}, Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}, SciMLBase.StandardODEProblem}, sensealg::Nothing, u0::CuArray{ForwardDiff.Dual{ForwardDiff.Tag{ODEFunction{true, SciMLBase.AutoSpecialize, var"#nn_dynamics#4"{Int64, var"#ude#3"{Lux.Chain{NamedTuple{(:layer_1, :layer_2), Tuple{Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}, Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}}}}, typeof(knownDynamics)}}, UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing}, Float64}, Float64, 4}, 1, CUDA.Mem.DeviceBuffer}, p::ComponentVector{ForwardDiff.Dual{ForwardDiff.Tag{ODEFunction{true, SciMLBase.AutoSpecialize, var"#nn_dynamics#4"{Int64, var"#ude#3"{Lux.Chain{NamedTuple{(:layer_1, :layer_2), Tuple{Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}, Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}}}}, typeof(knownDynamics)}}, UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing}, Float64}, Float64, 4}, CuArray{ForwardDiff.Dual{ForwardDiff.Tag{ODEFunction{true, SciMLBase.AutoSpecialize, var"#nn_dynamics#4"{Int64, var"#ude#3"{Lux.Chain{NamedTuple{(:layer_1, :layer_2), Tuple{Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}, Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}}}}, typeof(knownDynamics)}}, UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing}, Float64}, Float64, 4}, 1, CUDA.Mem.DeviceBuffer}, Tuple{Axis{(layer_1 = ViewAxis(1:25, Axis(weight = ViewAxis(1:20, ShapedAxis((5, 4), NamedTuple())), bias = ViewAxis(21:25, ShapedAxis((5, 1), NamedTuple())))), layer_2 = ViewAxis(26:49, Axis(weight = ViewAxis(1:20, ShapedAxis((4, 5), NamedTuple())), bias = ViewAxis(21:24, ShapedAxis((4, 1), NamedTuple())))))}}}, args::CompositeAlgorithm{Tuple{Tsit5{typeof(OrdinaryDiffEq.trivial_limiter!),typeof(OrdinaryDiffEq.trivial_limiter!),Static.False}, Rosenbrock23{0,true,Nothing,typeof(OrdinaryDiffEq.DEFAULT_PRECS),Val{:forward},true,nothing}},AutoSwitch{Tsit5{typeof(OrdinaryDiffEq.trivial_limiter!),typeof(OrdinaryDiffEq.trivial_limiter!),Static.False}, Rosenbrock23{0,true,Nothing,typeof(OrdinaryDiffEq.DEFAULT_PRECS),Val{:forward},true,nothing}, Rational{Int64}, Int64}}; kwargs::Base.Pairs{Symbol, Any, NTuple{4, Symbol}, NamedTuple{(:saveat, :verbose, :abstol, :reltol), Tuple{Vector{Float64}, Bool, Float64, Float64}}})
    @ DiffEqBase ~/.julia/packages/DiffEqBase/s433k/src/solve.jl:1020
 [15] solve(prob::ODEProblem{CuArray{ForwardDiff.Dual{ForwardDiff.Tag{ODEFunction{true, SciMLBase.AutoSpecialize, var"#nn_dynamics#4"{Int64, var"#ude#3"{Lux.Chain{NamedTuple{(:layer_1, :layer_2), Tuple{Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}, Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}}}}, typeof(knownDynamics)}}, UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing}, Float64}, Float64, 4}, 1, CUDA.Mem.DeviceBuffer}, Tuple{Float64, Float64}, true, ComponentVector{ForwardDiff.Dual{ForwardDiff.Tag{ODEFunction{true, SciMLBase.AutoSpecialize, var"#nn_dynamics#4"{Int64, var"#ude#3"{Lux.Chain{NamedTuple{(:layer_1, :layer_2), Tuple{Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}, Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}}}}, typeof(knownDynamics)}}, UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing}, Float64}, Float64, 4}, CuArray{ForwardDiff.Dual{ForwardDiff.Tag{ODEFunction{true, SciMLBase.AutoSpecialize, var"#nn_dynamics#4"{Int64, var"#ude#3"{Lux.Chain{NamedTuple{(:layer_1, :layer_2), Tuple{Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}, Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}}}}, typeof(knownDynamics)}}, UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing}, Float64}, Float64, 4}, 1, CUDA.Mem.DeviceBuffer}, Tuple{Axis{(layer_1 = ViewAxis(1:25, Axis(weight = ViewAxis(1:20, ShapedAxis((5, 4), NamedTuple())), bias = ViewAxis(21:25, ShapedAxis((5, 1), NamedTuple())))), layer_2 = ViewAxis(26:49, Axis(weight = ViewAxis(1:20, ShapedAxis((4, 5), NamedTuple())), bias = ViewAxis(21:24, ShapedAxis((4, 1), NamedTuple())))))}}}, ODEFunction{true, SciMLBase.FullSpecialize, ODEFunction{true, SciMLBase.AutoSpecialize, var"#nn_dynamics#4"{Int64, var"#ude#3"{Lux.Chain{NamedTuple{(:layer_1, :layer_2), Tuple{Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}, Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}}}}, typeof(knownDynamics)}}, UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing}, UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing}, Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}, SciMLBase.StandardODEProblem}, args::CompositeAlgorithm{Tuple{Tsit5{typeof(OrdinaryDiffEq.trivial_limiter!),typeof(OrdinaryDiffEq.trivial_limiter!),Static.False}, Rosenbrock23{0,true,Nothing,typeof(OrdinaryDiffEq.DEFAULT_PRECS),Val{:forward},true,nothing}},AutoSwitch{Tsit5{typeof(OrdinaryDiffEq.trivial_limiter!),typeof(OrdinaryDiffEq.trivial_limiter!),Static.False}, Rosenbrock23{0,true,Nothing,typeof(OrdinaryDiffEq.DEFAULT_PRECS),Val{:forward},true,nothing}, Rational{Int64}, Int64}}; sensealg::Nothing, u0::Nothing, p::Nothing, wrap::Val{true}, kwargs::Base.Pairs{Symbol, Any, NTuple{4, Symbol}, NamedTuple{(:saveat, :verbose, :abstol, :reltol), Tuple{Vector{Float64}, Bool, Float64, Float64}}})
    @ DiffEqBase ~/.julia/packages/DiffEqBase/s433k/src/solve.jl:930
 [16] (::SciMLSensitivity.var"#300#309"{0, CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}, Float64, Base.Pairs{Symbol, Real, Tuple{Symbol, Symbol, Symbol}, NamedTuple{(:verbose, :abstol, :reltol), Tuple{Bool, Float64, Float64}}}, ODEProblem{CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Float64, Float64}, true, ComponentVector{Float64, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Axis{(layer_1 = ViewAxis(1:25, Axis(weight = ViewAxis(1:20, ShapedAxis((5, 4), NamedTuple())), bias = ViewAxis(21:25, ShapedAxis((5, 1), NamedTuple())))), layer_2 = ViewAxis(26:49, Axis(weight = ViewAxis(1:20, ShapedAxis((4, 5), NamedTuple())), bias = ViewAxis(21:24, ShapedAxis((4, 1), NamedTuple())))))}}}, ODEFunction{true, SciMLBase.AutoSpecialize, var"#nn_dynamics#4"{Int64, var"#ude#3"{Lux.Chain{NamedTuple{(:layer_1, :layer_2), Tuple{Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}, Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}}}}, typeof(knownDynamics)}}, UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing}, Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}, SciMLBase.StandardODEProblem}, CompositeAlgorithm{Tuple{Tsit5{typeof(OrdinaryDiffEq.trivial_limiter!),typeof(OrdinaryDiffEq.trivial_limiter!),Static.False}, Rosenbrock23{0,true,Nothing,typeof(OrdinaryDiffEq.DEFAULT_PRECS),Val{:forward},true,nothing}},AutoSwitch{Tsit5{typeof(OrdinaryDiffEq.trivial_limiter!),typeof(OrdinaryDiffEq.trivial_limiter!),Static.False}, Rosenbrock23{0,true,Nothing,typeof(OrdinaryDiffEq.DEFAULT_PRECS),Val{:forward},true,nothing}, Rational{Int64}, Int64}}, ForwardDiffSensitivity{0, nothing}, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, ComponentVector{Float64, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Axis{(layer_1 = ViewAxis(1:25, Axis(weight = ViewAxis(1:20, ShapedAxis((5, 4), NamedTuple())), bias = ViewAxis(21:25, ShapedAxis((5, 1), NamedTuple())))), layer_2 = ViewAxis(26:49, Axis(weight = ViewAxis(1:20, ShapedAxis((4, 5), NamedTuple())), bias = ViewAxis(21:24, ShapedAxis((4, 1), NamedTuple())))))}}}, Tuple{}, Vector{Float64}})()
    @ SciMLSensitivity ~/.julia/packages/SciMLSensitivity/DSyJO/src/concrete_solve.jl:873
 [17] unthunk
    @ ~/.julia/packages/ChainRulesCore/7MWx2/src/tangent_types/thunks.jl:204 [inlined]
 [18] wrap_chainrules_output
    @ ~/.julia/packages/Zygote/oGI57/src/compiler/chainrules.jl:110 [inlined]
 [19] map
    @ ./tuple.jl:223 [inlined]
 [20] map (repeats 3 times)
    @ ./tuple.jl:224 [inlined]
 [21] wrap_chainrules_output
    @ ~/.julia/packages/Zygote/oGI57/src/compiler/chainrules.jl:111 [inlined]
 [22] ZBack
    @ ~/.julia/packages/Zygote/oGI57/src/compiler/chainrules.jl:211 [inlined]
 [23] kw_zpullback
    @ ~/.julia/packages/Zygote/oGI57/src/compiler/chainrules.jl:237 [inlined]
 [24] #280
    @ ~/.julia/packages/Zygote/oGI57/src/lib/lib.jl:206 [inlined]
 [25] (::Zygote.var"#2155#back#282"{Zygote.var"#280#281"{Tuple{NTuple{6, Nothing}, Tuple{Nothing}}, Zygote.var"#kw_zpullback#53"{SciMLSensitivity.var"#forward_sensitivity_backpass#305"{0, Float64, Base.Pairs{Symbol, Real, Tuple{Symbol, Symbol, Symbol}, NamedTuple{(:verbose, :abstol, :reltol), Tuple{Bool, Float64, Float64}}}, ODEProblem{CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Float64, Float64}, true, ComponentVector{Float64, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Axis{(layer_1 = ViewAxis(1:25, Axis(weight = ViewAxis(1:20, ShapedAxis((5, 4), NamedTuple())), bias = ViewAxis(21:25, ShapedAxis((5, 1), NamedTuple())))), layer_2 = ViewAxis(26:49, Axis(weight = ViewAxis(1:20, ShapedAxis((4, 5), NamedTuple())), bias = ViewAxis(21:24, ShapedAxis((4, 1), NamedTuple())))))}}}, ODEFunction{true, SciMLBase.AutoSpecialize, var"#nn_dynamics#4"{Int64, var"#ude#3"{Lux.Chain{NamedTuple{(:layer_1, :layer_2), Tuple{Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}, Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}}}}, typeof(knownDynamics)}}, UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing}, Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}, SciMLBase.StandardODEProblem}, CompositeAlgorithm{Tuple{Tsit5{typeof(OrdinaryDiffEq.trivial_limiter!),typeof(OrdinaryDiffEq.trivial_limiter!),Static.False}, Rosenbrock23{0,true,Nothing,typeof(OrdinaryDiffEq.DEFAULT_PRECS),Val{:forward},true,nothing}},AutoSwitch{Tsit5{typeof(OrdinaryDiffEq.trivial_limiter!),typeof(OrdinaryDiffEq.trivial_limiter!),Static.False}, Rosenbrock23{0,true,Nothing,typeof(OrdinaryDiffEq.DEFAULT_PRECS),Val{:forward},true,nothing}, Rational{Int64}, Int64}}, ForwardDiffSensitivity{0, nothing}, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, ComponentVector{Float64, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Axis{(layer_1 = ViewAxis(1:25, Axis(weight = ViewAxis(1:20, ShapedAxis((5, 4), NamedTuple())), bias = ViewAxis(21:25, ShapedAxis((5, 1), NamedTuple())))), layer_2 = ViewAxis(26:49, Axis(weight = ViewAxis(1:20, ShapedAxis((4, 5), NamedTuple())), bias = ViewAxis(21:24, ShapedAxis((4, 1), NamedTuple())))))}}}, SciMLBase.ChainRulesOriginator, Tuple{}, Vector{Float64}}}}})(Δ::CuArray{Float64, 2, CUDA.Mem.DeviceBuffer})
    @ Zygote ~/.julia/packages/ZygoteRules/4nXuu/src/adjoint.jl:71
 [26] Pullback
    @ ~/.julia/packages/DiffEqBase/s433k/src/solve.jl:930 [inlined]
 [27] (::Zygote.Pullback{Tuple{DiffEqBase.var"##solve#40", Nothing, Nothing, Nothing, Val{true}, Base.Pairs{Symbol, Float64, Tuple{Symbol, Symbol, Symbol}, NamedTuple{(:saveat, :abstol, :reltol), Tuple{Float64, Float64, Float64}}}, typeof(solve), ODEProblem{CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Float64, Float64}, true, ComponentVector{Float64, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Axis{(layer_1 = ViewAxis(1:25, Axis(weight = ViewAxis(1:20, ShapedAxis((5, 4), NamedTuple())), bias = ViewAxis(21:25, ShapedAxis((5, 1), NamedTuple())))), layer_2 = ViewAxis(26:49, Axis(weight = ViewAxis(1:20, ShapedAxis((4, 5), NamedTuple())), bias = ViewAxis(21:24, ShapedAxis((4, 1), NamedTuple())))))}}}, ODEFunction{true, SciMLBase.AutoSpecialize, var"#nn_dynamics#4"{Int64, var"#ude#3"{Lux.Chain{NamedTuple{(:layer_1, :layer_2), Tuple{Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}, Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}}}}, typeof(knownDynamics)}}, UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing}, Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}, SciMLBase.StandardODEProblem}, CompositeAlgorithm{Tuple{Tsit5{typeof(OrdinaryDiffEq.trivial_limiter!),typeof(OrdinaryDiffEq.trivial_limiter!),Static.False}, Rosenbrock23{0,true,Nothing,typeof(OrdinaryDiffEq.DEFAULT_PRECS),Val{:forward},true,nothing}},AutoSwitch{Tsit5{typeof(OrdinaryDiffEq.trivial_limiter!),typeof(OrdinaryDiffEq.trivial_limiter!),Static.False}, Rosenbrock23{0,true,Nothing,typeof(OrdinaryDiffEq.DEFAULT_PRECS),Val{:forward},true,nothing}, Rational{Int64}, Int64}}}, Any})(Δ::CuArray{Float64, 2, CUDA.Mem.DeviceBuffer})
    @ Zygote ~/.julia/packages/Zygote/oGI57/src/compiler/interface2.jl:0
 [28] (::Zygote.var"#280#281"{Tuple{NTuple{7, Nothing}, Tuple{Nothing}}, Zygote.Pullback{Tuple{DiffEqBase.var"##solve#40", Nothing, Nothing, Nothing, Val{true}, Base.Pairs{Symbol, Float64, Tuple{Symbol, Symbol, Symbol}, NamedTuple{(:saveat, :abstol, :reltol), Tuple{Float64, Float64, Float64}}}, typeof(solve), ODEProblem{CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Float64, Float64}, true, ComponentVector{Float64, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Axis{(layer_1 = ViewAxis(1:25, Axis(weight = ViewAxis(1:20, ShapedAxis((5, 4), NamedTuple())), bias = ViewAxis(21:25, ShapedAxis((5, 1), NamedTuple())))), layer_2 = ViewAxis(26:49, Axis(weight = ViewAxis(1:20, ShapedAxis((4, 5), NamedTuple())), bias = ViewAxis(21:24, ShapedAxis((4, 1), NamedTuple())))))}}}, ODEFunction{true, SciMLBase.AutoSpecialize, var"#nn_dynamics#4"{Int64, var"#ude#3"{Lux.Chain{NamedTuple{(:layer_1, :layer_2), Tuple{Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}, Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}}}}, typeof(knownDynamics)}}, UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing}, Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}, SciMLBase.StandardODEProblem}, CompositeAlgorithm{Tuple{Tsit5{typeof(OrdinaryDiffEq.trivial_limiter!),typeof(OrdinaryDiffEq.trivial_limiter!),Static.False}, Rosenbrock23{0,true,Nothing,typeof(OrdinaryDiffEq.DEFAULT_PRECS),Val{:forward},true,nothing}},AutoSwitch{Tsit5{typeof(OrdinaryDiffEq.trivial_limiter!),typeof(OrdinaryDiffEq.trivial_limiter!),Static.False}, Rosenbrock23{0,true,Nothing,typeof(OrdinaryDiffEq.DEFAULT_PRECS),Val{:forward},true,nothing}, Rational{Int64}, Int64}}}, Any}})(Δ::CuArray{Float64, 2, CUDA.Mem.DeviceBuffer})
    @ Zygote ~/.julia/packages/Zygote/oGI57/src/lib/lib.jl:206
 [29] (::Zygote.var"#2155#back#282"{Zygote.var"#280#281"{Tuple{NTuple{7, Nothing}, Tuple{Nothing}}, Zygote.Pullback{Tuple{DiffEqBase.var"##solve#40", Nothing, Nothing, Nothing, Val{true}, Base.Pairs{Symbol, Float64, Tuple{Symbol, Symbol, Symbol}, NamedTuple{(:saveat, :abstol, :reltol), Tuple{Float64, Float64, Float64}}}, typeof(solve), ODEProblem{CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Float64, Float64}, true, ComponentVector{Float64, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Axis{(layer_1 = ViewAxis(1:25, Axis(weight = ViewAxis(1:20, ShapedAxis((5, 4), NamedTuple())), bias = ViewAxis(21:25, ShapedAxis((5, 1), NamedTuple())))), layer_2 = ViewAxis(26:49, Axis(weight = ViewAxis(1:20, ShapedAxis((4, 5), NamedTuple())), bias = ViewAxis(21:24, ShapedAxis((4, 1), NamedTuple())))))}}}, ODEFunction{true, SciMLBase.AutoSpecialize, var"#nn_dynamics#4"{Int64, var"#ude#3"{Lux.Chain{NamedTuple{(:layer_1, :layer_2), Tuple{Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}, Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}}}}, typeof(knownDynamics)}}, UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing}, Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}, SciMLBase.StandardODEProblem}, CompositeAlgorithm{Tuple{Tsit5{typeof(OrdinaryDiffEq.trivial_limiter!),typeof(OrdinaryDiffEq.trivial_limiter!),Static.False}, Rosenbrock23{0,true,Nothing,typeof(OrdinaryDiffEq.DEFAULT_PRECS),Val{:forward},true,nothing}},AutoSwitch{Tsit5{typeof(OrdinaryDiffEq.trivial_limiter!),typeof(OrdinaryDiffEq.trivial_limiter!),Static.False}, Rosenbrock23{0,true,Nothing,typeof(OrdinaryDiffEq.DEFAULT_PRECS),Val{:forward},true,nothing}, Rational{Int64}, Int64}}}, Any}}})(Δ::CuArray{Float64, 2, CUDA.Mem.DeviceBuffer})
    @ Zygote ~/.julia/packages/ZygoteRules/4nXuu/src/adjoint.jl:71
 [30] Pullback
    @ ~/.julia/packages/DiffEqBase/s433k/src/solve.jl:920 [inlined]
 [31] (::Zygote.Pullback{Tuple{CommonSolve.var"#solve##kw", NamedTuple{(:saveat, :abstol, :reltol), Tuple{Float64, Float64, Float64}}, typeof(solve), ODEProblem{CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Float64, Float64}, true, ComponentVector{Float64, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Axis{(layer_1 = ViewAxis(1:25, Axis(weight = ViewAxis(1:20, ShapedAxis((5, 4), NamedTuple())), bias = ViewAxis(21:25, ShapedAxis((5, 1), NamedTuple())))), layer_2 = ViewAxis(26:49, Axis(weight = ViewAxis(1:20, ShapedAxis((4, 5), NamedTuple())), bias = ViewAxis(21:24, ShapedAxis((4, 1), NamedTuple())))))}}}, ODEFunction{true, SciMLBase.AutoSpecialize, var"#nn_dynamics#4"{Int64, var"#ude#3"{Lux.Chain{NamedTuple{(:layer_1, :layer_2), Tuple{Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}, Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}}}}, typeof(knownDynamics)}}, UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing}, Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}, SciMLBase.StandardODEProblem}, CompositeAlgorithm{Tuple{Tsit5{typeof(OrdinaryDiffEq.trivial_limiter!),typeof(OrdinaryDiffEq.trivial_limiter!),Static.False}, Rosenbrock23{0,true,Nothing,typeof(OrdinaryDiffEq.DEFAULT_PRECS),Val{:forward},true,nothing}},AutoSwitch{Tsit5{typeof(OrdinaryDiffEq.trivial_limiter!),typeof(OrdinaryDiffEq.trivial_limiter!),Static.False}, Rosenbrock23{0,true,Nothing,typeof(OrdinaryDiffEq.DEFAULT_PRECS),Val{:forward},true,nothing}, Rational{Int64}, Int64}}}, Any})(Δ::CuArray{Float64, 2, CUDA.Mem.DeviceBuffer})
    @ Zygote ~/.julia/packages/Zygote/oGI57/src/compiler/interface2.jl:0
 [32] Pullback
    @ ~/NODE_Community_Forecast/test.jl:28 [inlined]
 [33] (::Zygote.Pullback{Tuple{var"#predict#5"{CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}, ODEProblem{CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Float64, Float64}, true, ComponentVector{Float64, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Axis{(layer_1 = ViewAxis(1:25, Axis(weight = ViewAxis(1:20, ShapedAxis((5, 4), NamedTuple())), bias = ViewAxis(21:25, ShapedAxis((5, 1), NamedTuple())))), layer_2 = ViewAxis(26:49, Axis(weight = ViewAxis(1:20, ShapedAxis((4, 5), NamedTuple())), bias = ViewAxis(21:24, ShapedAxis((4, 1), NamedTuple())))))}}}, ODEFunction{true, SciMLBase.AutoSpecialize, var"#nn_dynamics#4"{Int64, var"#ude#3"{Lux.Chain{NamedTuple{(:layer_1, :layer_2), Tuple{Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}, Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}}}}, typeof(knownDynamics)}}, UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing}, Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}, SciMLBase.StandardODEProblem}}, ComponentVector{Float64, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Axis{(layer_1 = ViewAxis(1:25, Axis(weight = ViewAxis(1:20, ShapedAxis((5, 4), NamedTuple())), bias = ViewAxis(21:25, ShapedAxis((5, 1), NamedTuple())))), layer_2 = ViewAxis(26:49, Axis(weight = ViewAxis(1:20, ShapedAxis((4, 5), NamedTuple())), bias = ViewAxis(21:24, ShapedAxis((4, 1), NamedTuple())))))}}}, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}}, Tuple{Zygote.Pullback{Tuple{Type{Rosenbrock23}}, Tuple{}}, Zygote.Pullback{Tuple{SciMLBase.var"#remake##kw", NamedTuple{(:u0, :tspan, :p), Tuple{CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Float64, Float64}, ComponentVector{Float64, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Axis{(layer_1 = ViewAxis(1:25, Axis(weight = ViewAxis(1:20, ShapedAxis((5, 4), NamedTuple())), bias = ViewAxis(21:25, ShapedAxis((5, 1), NamedTuple())))), layer_2 = ViewAxis(26:49, Axis(weight = ViewAxis(1:20, ShapedAxis((4, 5), NamedTuple())), bias = ViewAxis(21:24, ShapedAxis((4, 1), NamedTuple())))))}}}}}, typeof(remake), ODEProblem{CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Float64, Float64}, true, ComponentVector{Float64, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Axis{(layer_1 = ViewAxis(1:25, Axis(weight = ViewAxis(1:20, ShapedAxis((5, 4), NamedTuple())), bias = ViewAxis(21:25, ShapedAxis((5, 1), NamedTuple())))), layer_2 = ViewAxis(26:49, Axis(weight = ViewAxis(1:20, ShapedAxis((4, 5), NamedTuple())), bias = ViewAxis(21:24, ShapedAxis((4, 1), NamedTuple())))))}}}, ODEFunction{true, SciMLBase.AutoSpecialize, var"#nn_dynamics#4"{Int64, var"#ude#3"{Lux.Chain{NamedTuple{(:layer_1, :layer_2), Tuple{Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}, Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}}}}, typeof(knownDynamics)}}, UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing}, Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}, SciMLBase.StandardODEProblem}}, Any}, Zygote.ZBack{ChainRules.var"#size_pullback#919"}, Zygote.Pullback{Tuple{typeof(convert), Type{CuArray}, ODESolution{Float64, 2, Vector{CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}}, Nothing, Nothing, Vector{Float64}, Vector{Vector{CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}}}, ODEProblem{CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Float64, Float64}, true, ComponentVector{Float64, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Axis{(layer_1 = ViewAxis(1:25, Axis(weight = ViewAxis(1:20, ShapedAxis((5, 4), NamedTuple())), bias = ViewAxis(21:25, ShapedAxis((5, 1), NamedTuple())))), layer_2 = ViewAxis(26:49, Axis(weight = ViewAxis(1:20, ShapedAxis((4, 5), NamedTuple())), bias = ViewAxis(21:24, ShapedAxis((4, 1), NamedTuple())))))}}}, ODEFunction{true, SciMLBase.AutoSpecialize, var"#nn_dynamics#4"{Int64, var"#ude#3"{Lux.Chain{NamedTuple{(:layer_1, :layer_2), Tuple{Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}, Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}}}}, typeof(knownDynamics)}}, UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing}, Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}, SciMLBase.StandardODEProblem}, CompositeAlgorithm{Tuple{Tsit5{typeof(OrdinaryDiffEq.trivial_limiter!),typeof(OrdinaryDiffEq.trivial_limiter!),Static.False}, Rosenbrock23{4,true,LinearSolve.QRFactorization{NoPivot},typeof(OrdinaryDiffEq.DEFAULT_PRECS),Val{:forward},true,nothing}},OrdinaryDiffEq.AutoSwitchCache{Tsit5{typeof(OrdinaryDiffEq.trivial_limiter!),typeof(OrdinaryDiffEq.trivial_limiter!),Static.False}, Rosenbrock23{0,true,Nothing,typeof(OrdinaryDiffEq.DEFAULT_PRECS),Val{:forward},true,nothing}, Rational{Int64}, Int64}}, OrdinaryDiffEq.CompositeInterpolationData{ODEFunction{true, SciMLBase.AutoSpecialize, var"#nn_dynamics#4"{Int64, var"#ude#3"{Lux.Chain{NamedTuple{(:layer_1, :layer_2), Tuple{Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}, Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}}}}, typeof(knownDynamics)}}, UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing}, Vector{CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}}, Vector{Float64}, Vector{Vector{CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}}}, CompositeCache{Tuple{Tsit5Cache{CuArray{Float64, 1, CUDA.Mem.DeviceBuffer},CuArray{Float64, 1, CUDA.Mem.DeviceBuffer},CuArray{Float64, 1, CUDA.Mem.DeviceBuffer},typeof(OrdinaryDiffEq.trivial_limiter!),typeof(OrdinaryDiffEq.trivial_limiter!),Static.False}, Rosenbrock23Cache{CuArray{Float64, 1, CUDA.Mem.DeviceBuffer},CuArray{Float64, 1, CUDA.Mem.DeviceBuffer},CuArray{Float64, 1, CUDA.Mem.DeviceBuffer},CuArray{Float64, 2, CUDA.Mem.DeviceBuffer},CuArray{Float64, 2, CUDA.Mem.DeviceBuffer},OrdinaryDiffEq.Rosenbrock23Tableau{Float64},SciMLBase.TimeGradientWrapper{ODEFunction{true, SciMLBase.AutoSpecialize, var"#nn_dynamics#4"{Int64, var"#ude#3"{Lux.Chain{NamedTuple{(:layer_1, :layer_2), Tuple{Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}, Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}}}}, typeof(knownDynamics)}}, UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing}, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, ComponentVector{Float64, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Axis{(layer_1 = ViewAxis(1:25, Axis(weight = ViewAxis(1:20, ShapedAxis((5, 4), NamedTuple())), bias = ViewAxis(21:25, ShapedAxis((5, 1), NamedTuple())))), layer_2 = ViewAxis(26:49, Axis(weight = ViewAxis(1:20, ShapedAxis((4, 5), NamedTuple())), bias = ViewAxis(21:24, ShapedAxis((4, 1), NamedTuple())))))}}}},SciMLBase.UJacobianWrapper{ODEFunction{true, SciMLBase.AutoSpecialize, var"#nn_dynamics#4"{Int64, var"#ude#3"{Lux.Chain{NamedTuple{(:layer_1, :layer_2), Tuple{Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}, Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}}}}, typeof(knownDynamics)}}, UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing}, Float64, ComponentVector{Float64, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Axis{(layer_1 = ViewAxis(1:25, Axis(weight = ViewAxis(1:20, ShapedAxis((5, 4), NamedTuple())), bias = ViewAxis(21:25, ShapedAxis((5, 1), NamedTuple())))), layer_2 = ViewAxis(26:49, Axis(weight = ViewAxis(1:20, ShapedAxis((4, 5), NamedTuple())), bias = ViewAxis(21:24, ShapedAxis((4, 1), NamedTuple())))))}}}},LinearSolve.LinearCache{CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, SciMLBase.NullParameters, LinearSolve.QRFactorization{NoPivot}, QR{Float64, CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}}, LinearSolve.InvPreconditioner{Diagonal{Float64, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}}}, Diagonal{Float64, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}}, Float64, true, LinearSolve.OperatorCondition.IllConditioned},SparseDiffTools.ForwardColorJacCache{CuArray{ForwardDiff.Dual{ForwardDiff.Tag{DiffEqBase.OrdinaryDiffEqTag, Float64}, Float64, 4}, 1, CUDA.Mem.DeviceBuffer}, CuArray{ForwardDiff.Dual{ForwardDiff.Tag{DiffEqBase.OrdinaryDiffEqTag, Float64}, Float64, 4}, 1, CUDA.Mem.DeviceBuffer}, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Vector{CuArray{NTuple{4, Float64}, 1, CUDA.Mem.DeviceBuffer}}, UnitRange{Int64}, Nothing},CuArray{ForwardDiff.Dual{ForwardDiff.Tag{DiffEqBase.OrdinaryDiffEqTag, Float64}, Float64, 1}, 1, CUDA.Mem.DeviceBuffer},Float64,Rosenbrock23{4,true,LinearSolve.QRFactorization{NoPivot},typeof(OrdinaryDiffEq.DEFAULT_PRECS),Val{:forward},true,nothing},Nothing}},OrdinaryDiffEq.AutoSwitchCache{Tsit5{typeof(OrdinaryDiffEq.trivial_limiter!),typeof(OrdinaryDiffEq.trivial_limiter!),Static.False}, Rosenbrock23{0,true,Nothing,typeof(OrdinaryDiffEq.DEFAULT_PRECS),Val{:forward},true,nothing}, Rational{Int64}, Int64}}}, DiffEqBase.Stats, Vector{Int64}}}, Tuple{Zygote.ZBack{RecursiveArrayTools.RecursiveArrayToolsZygoteExt.var"#1#2"}}}, Zygote.Pullback{Tuple{typeof(AutoTsit5), Rosenbrock23{0,true,Nothing,typeof(OrdinaryDiffEq.DEFAULT_PRECS),Val{:forward},true,nothing}}, Tuple{Zygote.Pullback{Tuple{Type{NamedTuple}}, Tuple{}}, Zygote.var"#2348#back#407"{Zygote.var"#pairs_namedtuple_pullback#406"{(), NamedTuple{(), Tuple{}}}}, Zygote.Pullback{Tuple{OrdinaryDiffEq.var"##AutoTsit5#645", Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}, typeof(AutoTsit5), Rosenbrock23{0,true,Nothing,typeof(OrdinaryDiffEq.DEFAULT_PRECS),Val{:forward},true,nothing}}, Any}}}, Zygote.Pullback{Tuple{Type{NamedTuple{(:saveat, :abstol, :reltol)}}, Tuple{Float64, Float64, Float64}}, Tuple{Zygote.Pullback{Tuple{Type{NamedTuple{(:saveat, :abstol, :reltol), Tuple{Float64, Float64, Float64}}}, Tuple{Float64, Float64, Float64}}, Tuple{Zygote.var"#2203#back#304"{Zygote.Jnew{NamedTuple{(:saveat, :abstol, :reltol), Tuple{Float64, Float64, Float64}}, Nothing, true}}}}}}, Zygote.var"#1999#back#193"{typeof(identity)}, Zygote.ZBack{Zygote.var"#Real_pullback#320"}, Zygote.var"#2166#back#292"{Zygote.var"#back#291"{:training_data, Zygote.Context{false}, var"#predict#5"{CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}, ODEProblem{CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Float64, Float64}, true, ComponentVector{Float64, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Axis{(layer_1 = ViewAxis(1:25, Axis(weight = ViewAxis(1:20, ShapedAxis((5, 4), NamedTuple())), bias = ViewAxis(21:25, ShapedAxis((5, 1), NamedTuple())))), layer_2 = ViewAxis(26:49, Axis(weight = ViewAxis(1:20, ShapedAxis((4, 5), NamedTuple())), bias = ViewAxis(21:24, ShapedAxis((4, 1), NamedTuple())))))}}}, ODEFunction{true, SciMLBase.AutoSpecialize, var"#nn_dynamics#4"{Int64, var"#ude#3"{Lux.Chain{NamedTuple{(:layer_1, :layer_2), Tuple{Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}, Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}}}}, typeof(knownDynamics)}}, UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing}, Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}, SciMLBase.StandardODEProblem}}, CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}}}, Zygote.Pullback{Tuple{Type{NamedTuple{(:u0, :tspan, :p)}}, Tuple{CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Float64, Float64}, ComponentVector{Float64, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Axis{(layer_1 = ViewAxis(1:25, Axis(weight = ViewAxis(1:20, ShapedAxis((5, 4), NamedTuple())), bias = ViewAxis(21:25, ShapedAxis((5, 1), NamedTuple())))), layer_2 = ViewAxis(26:49, Axis(weight = ViewAxis(1:20, ShapedAxis((4, 5), NamedTuple())), bias = ViewAxis(21:24, ShapedAxis((4, 1), NamedTuple())))))}}}}}, Tuple{Zygote.Pullback{Tuple{Type{NamedTuple{(:u0, :tspan, :p), Tuple{CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Float64, Float64}, ComponentVector{Float64, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Axis{(layer_1 = ViewAxis(1:25, Axis(weight = ViewAxis(1:20, ShapedAxis((5, 4), NamedTuple())), bias = ViewAxis(21:25, ShapedAxis((5, 1), NamedTuple())))), layer_2 = ViewAxis(26:49, Axis(weight = ViewAxis(1:20, ShapedAxis((4, 5), NamedTuple())), bias = ViewAxis(21:24, ShapedAxis((4, 1), NamedTuple())))))}}}}}}, Tuple{CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Float64, Float64}, ComponentVector{Float64, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Axis{(layer_1 = ViewAxis(1:25, Axis(weight = ViewAxis(1:20, ShapedAxis((5, 4), NamedTuple())), bias = ViewAxis(21:25, ShapedAxis((5, 1), NamedTuple())))), layer_2 = ViewAxis(26:49, Axis(weight = ViewAxis(1:20, ShapedAxis((4, 5), NamedTuple())), bias = ViewAxis(21:24, ShapedAxis((4, 1), NamedTuple())))))}}}}}, Tuple{Zygote.var"#2203#back#304"{Zygote.Jnew{NamedTuple{(:u0, :tspan, :p), Tuple{CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Float64, Float64}, ComponentVector{Float64, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Axis{(layer_1 = ViewAxis(1:25, Axis(weight = ViewAxis(1:20, ShapedAxis((5, 4), NamedTuple())), bias = ViewAxis(21:25, ShapedAxis((5, 1), NamedTuple())))), layer_2 = ViewAxis(26:49, Axis(weight = ViewAxis(1:20, ShapedAxis((4, 5), NamedTuple())), bias = ViewAxis(21:24, ShapedAxis((4, 1), NamedTuple())))))}}}}}, Nothing, true}}}}}}, Zygote.var"#1999#back#193"{typeof(identity)}, Zygote.ZBack{Zygote.var"#Real_pullback#320"}, Zygote.Pullback{Tuple{CommonSolve.var"#solve##kw", NamedTuple{(:saveat, :abstol, :reltol), Tuple{Float64, Float64, Float64}}, typeof(solve), ODEProblem{CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Float64, Float64}, true, ComponentVector{Float64, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Axis{(layer_1 = ViewAxis(1:25, Axis(weight = ViewAxis(1:20, ShapedAxis((5, 4), NamedTuple())), bias = ViewAxis(21:25, ShapedAxis((5, 1), NamedTuple())))), layer_2 = ViewAxis(26:49, Axis(weight = ViewAxis(1:20, ShapedAxis((4, 5), NamedTuple())), bias = ViewAxis(21:24, ShapedAxis((4, 1), NamedTuple())))))}}}, ODEFunction{true, SciMLBase.AutoSpecialize, var"#nn_dynamics#4"{Int64, var"#ude#3"{Lux.Chain{NamedTuple{(:layer_1, :layer_2), Tuple{Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}, Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}}}}, typeof(knownDynamics)}}, UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing}, Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}, SciMLBase.StandardODEProblem}, CompositeAlgorithm{Tuple{Tsit5{typeof(OrdinaryDiffEq.trivial_limiter!),typeof(OrdinaryDiffEq.trivial_limiter!),Static.False}, Rosenbrock23{0,true,Nothing,typeof(OrdinaryDiffEq.DEFAULT_PRECS),Val{:forward},true,nothing}},AutoSwitch{Tsit5{typeof(OrdinaryDiffEq.trivial_limiter!),typeof(OrdinaryDiffEq.trivial_limiter!),Static.False}, Rosenbrock23{0,true,Nothing,typeof(OrdinaryDiffEq.DEFAULT_PRECS),Val{:forward},true,nothing}, Rational{Int64}, Int64}}}, Any}, Zygote.var"#2166#back#292"{Zygote.var"#back#291"{:prob_nn, Zygote.Context{false}, var"#predict#5"{CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}, ODEProblem{CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Float64, Float64}, true, ComponentVector{Float64, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Axis{(layer_1 = ViewAxis(1:25, Axis(weight = ViewAxis(1:20, ShapedAxis((5, 4), NamedTuple())), bias = ViewAxis(21:25, ShapedAxis((5, 1), NamedTuple())))), layer_2 = ViewAxis(26:49, Axis(weight = ViewAxis(1:20, ShapedAxis((4, 5), NamedTuple())), bias = ViewAxis(21:24, ShapedAxis((4, 1), NamedTuple())))))}}}, ODEFunction{true, SciMLBase.AutoSpecialize, var"#nn_dynamics#4"{Int64, var"#ude#3"{Lux.Chain{NamedTuple{(:layer_1, :layer_2), Tuple{Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}, Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}}}}, typeof(knownDynamics)}}, UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing}, Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}, SciMLBase.StandardODEProblem}}, ODEProblem{CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Float64, Float64}, true, ComponentVector{Float64, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Axis{(layer_1 = ViewAxis(1:25, Axis(weight = ViewAxis(1:20, ShapedAxis((5, 4), NamedTuple())), bias = ViewAxis(21:25, ShapedAxis((5, 1), NamedTuple())))), layer_2 = ViewAxis(26:49, Axis(weight = ViewAxis(1:20, ShapedAxis((4, 5), NamedTuple())), bias = ViewAxis(21:24, ShapedAxis((4, 1), NamedTuple())))))}}}, ODEFunction{true, SciMLBase.AutoSpecialize, var"#nn_dynamics#4"{Int64, var"#ude#3"{Lux.Chain{NamedTuple{(:layer_1, :layer_2), Tuple{Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}, Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}}}}, typeof(knownDynamics)}}, UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing}, Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}, SciMLBase.StandardODEProblem}}}, Zygote.var"#1999#back#193"{typeof(identity)}}})(Δ::CuArray{Float64, 2, CUDA.Mem.DeviceBuffer})
    @ Zygote ~/.julia/packages/Zygote/oGI57/src/compiler/interface2.jl:0
 [34] Pullback
    @ ~/NODE_Community_Forecast/test.jl:27 [inlined]
 [35] (::Zygote.Pullback{Tuple{var"#predict#5"{CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}, ODEProblem{CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Float64, Float64}, true, ComponentVector{Float64, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Axis{(layer_1 = ViewAxis(1:25, Axis(weight = ViewAxis(1:20, ShapedAxis((5, 4), NamedTuple())), bias = ViewAxis(21:25, ShapedAxis((5, 1), NamedTuple())))), layer_2 = ViewAxis(26:49, Axis(weight = ViewAxis(1:20, ShapedAxis((4, 5), NamedTuple())), bias = ViewAxis(21:24, ShapedAxis((4, 1), NamedTuple())))))}}}, ODEFunction{true, SciMLBase.AutoSpecialize, var"#nn_dynamics#4"{Int64, var"#ude#3"{Lux.Chain{NamedTuple{(:layer_1, :layer_2), Tuple{Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}, Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}}}}, typeof(knownDynamics)}}, UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing}, Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}, SciMLBase.StandardODEProblem}}, ComponentVector{Float64, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Axis{(layer_1 = ViewAxis(1:25, Axis(weight = ViewAxis(1:20, ShapedAxis((5, 4), NamedTuple())), bias = ViewAxis(21:25, ShapedAxis((5, 1), NamedTuple())))), layer_2 = ViewAxis(26:49, Axis(weight = ViewAxis(1:20, ShapedAxis((4, 5), NamedTuple())), bias = ViewAxis(21:24, ShapedAxis((4, 1), NamedTuple())))))}}}}, Tuple{Zygote.var"#2166#back#292"{Zygote.var"#back#291"{:x0, Zygote.Context{false}, var"#predict#5"{CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}, ODEProblem{CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Float64, Float64}, true, ComponentVector{Float64, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Axis{(layer_1 = ViewAxis(1:25, Axis(weight = ViewAxis(1:20, ShapedAxis((5, 4), NamedTuple())), bias = ViewAxis(21:25, ShapedAxis((5, 1), NamedTuple())))), layer_2 = ViewAxis(26:49, Axis(weight = ViewAxis(1:20, ShapedAxis((4, 5), NamedTuple())), bias = ViewAxis(21:24, ShapedAxis((4, 1), NamedTuple())))))}}}, ODEFunction{true, SciMLBase.AutoSpecialize, var"#nn_dynamics#4"{Int64, var"#ude#3"{Lux.Chain{NamedTuple{(:layer_1, :layer_2), Tuple{Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}, Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}}}}, typeof(knownDynamics)}}, UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing}, Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}, SciMLBase.StandardODEProblem}}, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}}}, Zygote.Pullback{Tuple{var"#predict#5"{CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}, ODEProblem{CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Float64, Float64}, true, ComponentVector{Float64, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Axis{(layer_1 = ViewAxis(1:25, Axis(weight = ViewAxis(1:20, ShapedAxis((5, 4), NamedTuple())), bias = ViewAxis(21:25, ShapedAxis((5, 1), NamedTuple())))), layer_2 = ViewAxis(26:49, Axis(weight = ViewAxis(1:20, ShapedAxis((4, 5), NamedTuple())), bias = ViewAxis(21:24, ShapedAxis((4, 1), NamedTuple())))))}}}, ODEFunction{true, SciMLBase.AutoSpecialize, var"#nn_dynamics#4"{Int64, var"#ude#3"{Lux.Chain{NamedTuple{(:layer_1, :layer_2), Tuple{Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}, Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}}}}, typeof(knownDynamics)}}, UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing}, Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}, SciMLBase.StandardODEProblem}}, ComponentVector{Float64, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Axis{(layer_1 = ViewAxis(1:25, Axis(weight = ViewAxis(1:20, ShapedAxis((5, 4), NamedTuple())), bias = ViewAxis(21:25, ShapedAxis((5, 1), NamedTuple())))), layer_2 = ViewAxis(26:49, Axis(weight = ViewAxis(1:20, ShapedAxis((4, 5), NamedTuple())), bias = ViewAxis(21:24, ShapedAxis((4, 1), NamedTuple())))))}}}, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}}, Tuple{Zygote.Pullback{Tuple{Type{Rosenbrock23}}, Tuple{}}, Zygote.Pullback{Tuple{SciMLBase.var"#remake##kw", NamedTuple{(:u0, :tspan, :p), Tuple{CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Float64, Float64}, ComponentVector{Float64, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Axis{(layer_1 = ViewAxis(1:25, Axis(weight = ViewAxis(1:20, ShapedAxis((5, 4), NamedTuple())), bias = ViewAxis(21:25, ShapedAxis((5, 1), NamedTuple())))), layer_2 = ViewAxis(26:49, Axis(weight = ViewAxis(1:20, ShapedAxis((4, 5), NamedTuple())), bias = ViewAxis(21:24, ShapedAxis((4, 1), NamedTuple())))))}}}}}, typeof(remake), ODEProblem{CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Float64, Float64}, true, ComponentVector{Float64, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Axis{(layer_1 = ViewAxis(1:25, Axis(weight = ViewAxis(1:20, ShapedAxis((5, 4), NamedTuple())), bias = ViewAxis(21:25, ShapedAxis((5, 1), NamedTuple())))), layer_2 = ViewAxis(26:49, Axis(weight = ViewAxis(1:20, ShapedAxis((4, 5), NamedTuple())), bias = ViewAxis(21:24, ShapedAxis((4, 1), NamedTuple())))))}}}, ODEFunction{true, SciMLBase.AutoSpecialize, var"#nn_dynamics#4"{Int64, var"#ude#3"{Lux.Chain{NamedTuple{(:layer_1, :layer_2), Tuple{Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}, Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}}}}, typeof(knownDynamics)}}, UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing}, Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}, SciMLBase.StandardODEProblem}}, Any}, Zygote.ZBack{ChainRules.var"#size_pullback#919"}, Zygote.Pullback{Tuple{typeof(convert), Type{CuArray}, ODESolution{Float64, 2, Vector{CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}}, Nothing, Nothing, Vector{Float64}, Vector{Vector{CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}}}, ODEProblem{CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Float64, Float64}, true, ComponentVector{Float64, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Axis{(layer_1 = ViewAxis(1:25, Axis(weight = ViewAxis(1:20, ShapedAxis((5, 4), NamedTuple())), bias = ViewAxis(21:25, ShapedAxis((5, 1), NamedTuple())))), layer_2 = ViewAxis(26:49, Axis(weight = ViewAxis(1:20, ShapedAxis((4, 5), NamedTuple())), bias = ViewAxis(21:24, ShapedAxis((4, 1), NamedTuple())))))}}}, ODEFunction{true, SciMLBase.AutoSpecialize, var"#nn_dynamics#4"{Int64, var"#ude#3"{Lux.Chain{NamedTuple{(:layer_1, :layer_2), Tuple{Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}, Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}}}}, typeof(knownDynamics)}}, UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing}, Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}, SciMLBase.StandardODEProblem}, CompositeAlgorithm{Tuple{Tsit5{typeof(OrdinaryDiffEq.trivial_limiter!),typeof(OrdinaryDiffEq.trivial_limiter!),Static.False}, Rosenbrock23{4,true,LinearSolve.QRFactorization{NoPivot},typeof(OrdinaryDiffEq.DEFAULT_PRECS),Val{:forward},true,nothing}},OrdinaryDiffEq.AutoSwitchCache{Tsit5{typeof(OrdinaryDiffEq.trivial_limiter!),typeof(OrdinaryDiffEq.trivial_limiter!),Static.False}, Rosenbrock23{0,true,Nothing,typeof(OrdinaryDiffEq.DEFAULT_PRECS),Val{:forward},true,nothing}, Rational{Int64}, Int64}}, OrdinaryDiffEq.CompositeInterpolationData{ODEFunction{true, SciMLBase.AutoSpecialize, var"#nn_dynamics#4"{Int64, var"#ude#3"{Lux.Chain{NamedTuple{(:layer_1, :layer_2), Tuple{Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}, Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}}}}, typeof(knownDynamics)}}, UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing}, Vector{CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}}, Vector{Float64}, Vector{Vector{CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}}}, CompositeCache{Tuple{Tsit5Cache{CuArray{Float64, 1, CUDA.Mem.DeviceBuffer},CuArray{Float64, 1, CUDA.Mem.DeviceBuffer},CuArray{Float64, 1, CUDA.Mem.DeviceBuffer},typeof(OrdinaryDiffEq.trivial_limiter!),typeof(OrdinaryDiffEq.trivial_limiter!),Static.False}, Rosenbrock23Cache{CuArray{Float64, 1, CUDA.Mem.DeviceBuffer},CuArray{Float64, 1, CUDA.Mem.DeviceBuffer},CuArray{Float64, 1, CUDA.Mem.DeviceBuffer},CuArray{Float64, 2, CUDA.Mem.DeviceBuffer},CuArray{Float64, 2, CUDA.Mem.DeviceBuffer},OrdinaryDiffEq.Rosenbrock23Tableau{Float64},SciMLBase.TimeGradientWrapper{ODEFunction{true, SciMLBase.AutoSpecialize, var"#nn_dynamics#4"{Int64, var"#ude#3"{Lux.Chain{NamedTuple{(:layer_1, :layer_2), Tuple{Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}, Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}}}}, typeof(knownDynamics)}}, UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing}, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, ComponentVector{Float64, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Axis{(layer_1 = ViewAxis(1:25, Axis(weight = ViewAxis(1:20, ShapedAxis((5, 4), NamedTuple())), bias = ViewAxis(21:25, ShapedAxis((5, 1), NamedTuple())))), layer_2 = ViewAxis(26:49, Axis(weight = ViewAxis(1:20, ShapedAxis((4, 5), NamedTuple())), bias = ViewAxis(21:24, ShapedAxis((4, 1), NamedTuple())))))}}}},SciMLBase.UJacobianWrapper{ODEFunction{true, SciMLBase.AutoSpecialize, var"#nn_dynamics#4"{Int64, var"#ude#3"{Lux.Chain{NamedTuple{(:layer_1, :layer_2), Tuple{Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}, Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}}}}, typeof(knownDynamics)}}, UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing}, Float64, ComponentVector{Float64, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Axis{(layer_1 = ViewAxis(1:25, Axis(weight = ViewAxis(1:20, ShapedAxis((5, 4), NamedTuple())), bias = ViewAxis(21:25, ShapedAxis((5, 1), NamedTuple())))), layer_2 = ViewAxis(26:49, Axis(weight = ViewAxis(1:20, ShapedAxis((4, 5), NamedTuple())), bias = ViewAxis(21:24, ShapedAxis((4, 1), NamedTuple())))))}}}},LinearSolve.LinearCache{CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, SciMLBase.NullParameters, LinearSolve.QRFactorization{NoPivot}, QR{Float64, CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}}, LinearSolve.InvPreconditioner{Diagonal{Float64, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}}}, Diagonal{Float64, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}}, Float64, true, LinearSolve.OperatorCondition.IllConditioned},SparseDiffTools.ForwardColorJacCache{CuArray{ForwardDiff.Dual{ForwardDiff.Tag{DiffEqBase.OrdinaryDiffEqTag, Float64}, Float64, 4}, 1, CUDA.Mem.DeviceBuffer}, CuArray{ForwardDiff.Dual{ForwardDiff.Tag{DiffEqBase.OrdinaryDiffEqTag, Float64}, Float64, 4}, 1, CUDA.Mem.DeviceBuffer}, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Vector{CuArray{NTuple{4, Float64}, 1, CUDA.Mem.DeviceBuffer}}, UnitRange{Int64}, Nothing},CuArray{ForwardDiff.Dual{ForwardDiff.Tag{DiffEqBase.OrdinaryDiffEqTag, Float64}, Float64, 1}, 1, CUDA.Mem.DeviceBuffer},Float64,Rosenbrock23{4,true,LinearSolve.QRFactorization{NoPivot},typeof(OrdinaryDiffEq.DEFAULT_PRECS),Val{:forward},true,nothing},Nothing}},OrdinaryDiffEq.AutoSwitchCache{Tsit5{typeof(OrdinaryDiffEq.trivial_limiter!),typeof(OrdinaryDiffEq.trivial_limiter!),Static.False}, Rosenbrock23{0,true,Nothing,typeof(OrdinaryDiffEq.DEFAULT_PRECS),Val{:forward},true,nothing}, Rational{Int64}, Int64}}}, DiffEqBase.Stats, Vector{Int64}}}, Tuple{Zygote.ZBack{RecursiveArrayTools.RecursiveArrayToolsZygoteExt.var"#1#2"}}}, Zygote.Pullback{Tuple{typeof(AutoTsit5), Rosenbrock23{0,true,Nothing,typeof(OrdinaryDiffEq.DEFAULT_PRECS),Val{:forward},true,nothing}}, Tuple{Zygote.Pullback{Tuple{Type{NamedTuple}}, Tuple{}}, Zygote.var"#2348#back#407"{Zygote.var"#pairs_namedtuple_pullback#406"{(), NamedTuple{(), Tuple{}}}}, Zygote.Pullback{Tuple{OrdinaryDiffEq.var"##AutoTsit5#645", Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}, typeof(AutoTsit5), Rosenbrock23{0,true,Nothing,typeof(OrdinaryDiffEq.DEFAULT_PRECS),Val{:forward},true,nothing}}, Any}}}, Zygote.Pullback{Tuple{Type{NamedTuple{(:saveat, :abstol, :reltol)}}, Tuple{Float64, Float64, Float64}}, Tuple{Zygote.Pullback{Tuple{Type{NamedTuple{(:saveat, :abstol, :reltol), Tuple{Float64, Float64, Float64}}}, Tuple{Float64, Float64, Float64}}, Tuple{Zygote.var"#2203#back#304"{Zygote.Jnew{NamedTuple{(:saveat, :abstol, :reltol), Tuple{Float64, Float64, Float64}}, Nothing, true}}}}}}, Zygote.var"#1999#back#193"{typeof(identity)}, Zygote.ZBack{Zygote.var"#Real_pullback#320"}, Zygote.var"#2166#back#292"{Zygote.var"#back#291"{:training_data, Zygote.Context{false}, var"#predict#5"{CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}, ODEProblem{CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Float64, Float64}, true, ComponentVector{Float64, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Axis{(layer_1 = ViewAxis(1:25, Axis(weight = ViewAxis(1:20, ShapedAxis((5, 4), NamedTuple())), bias = ViewAxis(21:25, ShapedAxis((5, 1), NamedTuple())))), layer_2 = ViewAxis(26:49, Axis(weight = ViewAxis(1:20, ShapedAxis((4, 5), NamedTuple())), bias = ViewAxis(21:24, ShapedAxis((4, 1), NamedTuple())))))}}}, ODEFunction{true, SciMLBase.AutoSpecialize, var"#nn_dynamics#4"{Int64, var"#ude#3"{Lux.Chain{NamedTuple{(:layer_1, :layer_2), Tuple{Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}, Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}}}}, typeof(knownDynamics)}}, UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing}, Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}, SciMLBase.StandardODEProblem}}, CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}}}, Zygote.Pullback{Tuple{Type{NamedTuple{(:u0, :tspan, :p)}}, Tuple{CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Float64, Float64}, ComponentVector{Float64, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Axis{(layer_1 = ViewAxis(1:25, Axis(weight = ViewAxis(1:20, ShapedAxis((5, 4), NamedTuple())), bias = ViewAxis(21:25, ShapedAxis((5, 1), NamedTuple())))), layer_2 = ViewAxis(26:49, Axis(weight = ViewAxis(1:20, ShapedAxis((4, 5), NamedTuple())), bias = ViewAxis(21:24, ShapedAxis((4, 1), NamedTuple())))))}}}}}, Tuple{Zygote.Pullback{Tuple{Type{NamedTuple{(:u0, :tspan, :p), Tuple{CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Float64, Float64}, ComponentVector{Float64, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Axis{(layer_1 = ViewAxis(1:25, Axis(weight = ViewAxis(1:20, ShapedAxis((5, 4), NamedTuple())), bias = ViewAxis(21:25, ShapedAxis((5, 1), NamedTuple())))), layer_2 = ViewAxis(26:49, Axis(weight = ViewAxis(1:20, ShapedAxis((4, 5), NamedTuple())), bias = ViewAxis(21:24, ShapedAxis((4, 1), NamedTuple())))))}}}}}}, Tuple{CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Float64, Float64}, ComponentVector{Float64, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Axis{(layer_1 = ViewAxis(1:25, Axis(weight = ViewAxis(1:20, ShapedAxis((5, 4), NamedTuple())), bias = ViewAxis(21:25, ShapedAxis((5, 1), NamedTuple())))), layer_2 = ViewAxis(26:49, Axis(weight = ViewAxis(1:20, ShapedAxis((4, 5), NamedTuple())), bias = ViewAxis(21:24, ShapedAxis((4, 1), NamedTuple())))))}}}}}, Tuple{Zygote.var"#2203#back#304"{Zygote.Jnew{NamedTuple{(:u0, :tspan, :p), Tuple{CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Float64, Float64}, ComponentVector{Float64, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Axis{(layer_1 = ViewAxis(1:25, Axis(weight = ViewAxis(1:20, ShapedAxis((5, 4), NamedTuple())), bias = ViewAxis(21:25, ShapedAxis((5, 1), NamedTuple())))), layer_2 = ViewAxis(26:49, Axis(weight = ViewAxis(1:20, ShapedAxis((4, 5), NamedTuple())), bias = ViewAxis(21:24, ShapedAxis((4, 1), NamedTuple())))))}}}}}, Nothing, true}}}}}}, Zygote.var"#1999#back#193"{typeof(identity)}, Zygote.ZBack{Zygote.var"#Real_pullback#320"}, Zygote.Pullback{Tuple{CommonSolve.var"#solve##kw", NamedTuple{(:saveat, :abstol, :reltol), Tuple{Float64, Float64, Float64}}, typeof(solve), ODEProblem{CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Float64, Float64}, true, ComponentVector{Float64, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Axis{(layer_1 = ViewAxis(1:25, Axis(weight = ViewAxis(1:20, ShapedAxis((5, 4), NamedTuple())), bias = ViewAxis(21:25, ShapedAxis((5, 1), NamedTuple())))), layer_2 = ViewAxis(26:49, Axis(weight = ViewAxis(1:20, ShapedAxis((4, 5), NamedTuple())), bias = ViewAxis(21:24, ShapedAxis((4, 1), NamedTuple())))))}}}, ODEFunction{true, SciMLBase.AutoSpecialize, var"#nn_dynamics#4"{Int64, var"#ude#3"{Lux.Chain{NamedTuple{(:layer_1, :layer_2), Tuple{Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}, Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}}}}, typeof(knownDynamics)}}, UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing}, Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}, SciMLBase.StandardODEProblem}, CompositeAlgorithm{Tuple{Tsit5{typeof(OrdinaryDiffEq.trivial_limiter!),typeof(OrdinaryDiffEq.trivial_limiter!),Static.False}, Rosenbrock23{0,true,Nothing,typeof(OrdinaryDiffEq.DEFAULT_PRECS),Val{:forward},true,nothing}},AutoSwitch{Tsit5{typeof(OrdinaryDiffEq.trivial_limiter!),typeof(OrdinaryDiffEq.trivial_limiter!),Static.False}, Rosenbrock23{0,true,Nothing,typeof(OrdinaryDiffEq.DEFAULT_PRECS),Val{:forward},true,nothing}, Rational{Int64}, Int64}}}, Any}, Zygote.var"#2166#back#292"{Zygote.var"#back#291"{:prob_nn, Zygote.Context{false}, var"#predict#5"{CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}, ODEProblem{CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Float64, Float64}, true, ComponentVector{Float64, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Axis{(layer_1 = ViewAxis(1:25, Axis(weight = ViewAxis(1:20, ShapedAxis((5, 4), NamedTuple())), bias = ViewAxis(21:25, ShapedAxis((5, 1), NamedTuple())))), layer_2 = ViewAxis(26:49, Axis(weight = ViewAxis(1:20, ShapedAxis((4, 5), NamedTuple())), bias = ViewAxis(21:24, ShapedAxis((4, 1), NamedTuple())))))}}}, ODEFunction{true, SciMLBase.AutoSpecialize, var"#nn_dynamics#4"{Int64, var"#ude#3"{Lux.Chain{NamedTuple{(:layer_1, :layer_2), Tuple{Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}, Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}}}}, typeof(knownDynamics)}}, UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing}, Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}, SciMLBase.StandardODEProblem}}, ODEProblem{CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Float64, Float64}, true, ComponentVector{Float64, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Axis{(layer_1 = ViewAxis(1:25, Axis(weight = ViewAxis(1:20, ShapedAxis((5, 4), NamedTuple())), bias = ViewAxis(21:25, ShapedAxis((5, 1), NamedTuple())))), layer_2 = ViewAxis(26:49, Axis(weight = ViewAxis(1:20, ShapedAxis((4, 5), NamedTuple())), bias = ViewAxis(21:24, ShapedAxis((4, 1), NamedTuple())))))}}}, ODEFunction{true, SciMLBase.AutoSpecialize, var"#nn_dynamics#4"{Int64, var"#ude#3"{Lux.Chain{NamedTuple{(:layer_1, :layer_2), Tuple{Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}, Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}}}}, typeof(knownDynamics)}}, UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing}, Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}, SciMLBase.StandardODEProblem}}}, Zygote.var"#1999#back#193"{typeof(identity)}}}}})(Δ::CuArray{Float64, 2, CUDA.Mem.DeviceBuffer})
    @ Zygote ~/.julia/packages/Zygote/oGI57/src/compiler/interface2.jl:0
 [36] Pullback
    @ ~/NODE_Community_Forecast/test.jl:38 [inlined]
 [37] (::Zygote.Pullback{Tuple{var"#loss_function#6"{CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}, Float64}, ComponentVector{Float64, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Axis{(layer_1 = ViewAxis(1:25, Axis(weight = ViewAxis(1:20, ShapedAxis((5, 4), NamedTuple())), bias = ViewAxis(21:25, ShapedAxis((5, 1), NamedTuple())))), layer_2 = ViewAxis(26:49, Axis(weight = ViewAxis(1:20, ShapedAxis((4, 5), NamedTuple())), bias = ViewAxis(21:24, ShapedAxis((4, 1), NamedTuple())))))}}}}, Any})(Δ::Float64)
    @ Zygote ~/.julia/packages/Zygote/oGI57/src/compiler/interface2.jl:0
 [38] Pullback
    @ ~/NODE_Community_Forecast/test.jl:44 [inlined]
 [39] (::Zygote.Pullback{Tuple{var"#2#7"{var"#loss_function#6"{CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}, Float64}}, ComponentVector{Float64, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Axis{(layer_1 = ViewAxis(1:25, Axis(weight = ViewAxis(1:20, ShapedAxis((5, 4), NamedTuple())), bias = ViewAxis(21:25, ShapedAxis((5, 1), NamedTuple())))), layer_2 = ViewAxis(26:49, Axis(weight = ViewAxis(1:20, ShapedAxis((4, 5), NamedTuple())), bias = ViewAxis(21:24, ShapedAxis((4, 1), NamedTuple())))))}}}, SciMLBase.NullParameters}, Tuple{Zygote.var"#2166#back#292"{Zygote.var"#back#291"{:loss_function, Zygote.Context{false}, var"#2#7"{var"#loss_function#6"{CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}, Float64}}, var"#loss_function#6"{CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}, Float64}}}, Zygote.Pullback{Tuple{var"#loss_function#6"{CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}, Float64}, ComponentVector{Float64, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Axis{(layer_1 = ViewAxis(1:25, Axis(weight = ViewAxis(1:20, ShapedAxis((5, 4), NamedTuple())), bias = ViewAxis(21:25, ShapedAxis((5, 1), NamedTuple())))), layer_2 = ViewAxis(26:49, Axis(weight = ViewAxis(1:20, ShapedAxis((4, 5), NamedTuple())), bias = ViewAxis(21:24, ShapedAxis((4, 1), NamedTuple())))))}}}}, Any}}})(Δ::Float64)
    @ Zygote ~/.julia/packages/Zygote/oGI57/src/compiler/interface2.jl:0
 [40] #280
    @ ~/.julia/packages/Zygote/oGI57/src/lib/lib.jl:206 [inlined]
 [41] #2155#back
    @ ~/.julia/packages/ZygoteRules/4nXuu/src/adjoint.jl:71 [inlined]
 [42] Pullback
    @ ~/.julia/packages/SciMLBase/l4PVV/src/scimlfunctions.jl:3772 [inlined]
 [43] (::Zygote.Pullback{Tuple{OptimizationFunction{true, Optimization.AutoZygote, var"#2#7"{var"#loss_function#6"{CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}, Float64}}, 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}, ComponentVector{Float64, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Axis{(layer_1 = ViewAxis(1:25, Axis(weight = ViewAxis(1:20, ShapedAxis((5, 4), NamedTuple())), bias = ViewAxis(21:25, ShapedAxis((5, 1), NamedTuple())))), layer_2 = ViewAxis(26:49, Axis(weight = ViewAxis(1:20, ShapedAxis((4, 5), NamedTuple())), bias = ViewAxis(21:24, ShapedAxis((4, 1), NamedTuple())))))}}}, SciMLBase.NullParameters}, Tuple{Zygote.var"#2166#back#292"{Zygote.var"#back#291"{:f, Zygote.Context{false}, OptimizationFunction{true, Optimization.AutoZygote, var"#2#7"{var"#loss_function#6"{CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}, Float64}}, 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}, var"#2#7"{var"#loss_function#6"{CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}, Float64}}}}, Zygote.var"#2155#back#282"{Zygote.var"#280#281"{Tuple{Tuple{Nothing, Nothing}}, Zygote.Pullback{Tuple{var"#2#7"{var"#loss_function#6"{CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}, Float64}}, ComponentVector{Float64, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Axis{(layer_1 = ViewAxis(1:25, Axis(weight = ViewAxis(1:20, ShapedAxis((5, 4), NamedTuple())), bias = ViewAxis(21:25, ShapedAxis((5, 1), NamedTuple())))), layer_2 = ViewAxis(26:49, Axis(weight = ViewAxis(1:20, ShapedAxis((4, 5), NamedTuple())), bias = ViewAxis(21:24, ShapedAxis((4, 1), NamedTuple())))))}}}, SciMLBase.NullParameters}, Tuple{Zygote.var"#2166#back#292"{Zygote.var"#back#291"{:loss_function, Zygote.Context{false}, var"#2#7"{var"#loss_function#6"{CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}, Float64}}, var"#loss_function#6"{CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}, Float64}}}, Zygote.Pullback{Tuple{var"#loss_function#6"{CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}, Float64}, ComponentVector{Float64, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Axis{(layer_1 = ViewAxis(1:25, Axis(weight = ViewAxis(1:20, ShapedAxis((5, 4), NamedTuple())), bias = ViewAxis(21:25, ShapedAxis((5, 1), NamedTuple())))), layer_2 = ViewAxis(26:49, Axis(weight = ViewAxis(1:20, ShapedAxis((4, 5), NamedTuple())), bias = ViewAxis(21:24, ShapedAxis((4, 1), NamedTuple())))))}}}}, Any}}}}}}})(Δ::Float64)
    @ Zygote ~/.julia/packages/Zygote/oGI57/src/compiler/interface2.jl:0
 [44] #280
    @ ~/.julia/packages/Zygote/oGI57/src/lib/lib.jl:206 [inlined]
 [45] #2155#back
    @ ~/.julia/packages/ZygoteRules/4nXuu/src/adjoint.jl:71 [inlined]
 [46] Pullback
    @ ~/.julia/packages/Optimization/GEo8L/src/function/zygote.jl:30 [inlined]
 [47] (::Zygote.Pullback{Tuple{Optimization.var"#156#165"{OptimizationFunction{true, Optimization.AutoZygote, var"#2#7"{var"#loss_function#6"{CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}, Float64}}, 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}, SciMLBase.NullParameters}, ComponentVector{Float64, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Axis{(layer_1 = ViewAxis(1:25, Axis(weight = ViewAxis(1:20, ShapedAxis((5, 4), NamedTuple())), bias = ViewAxis(21:25, ShapedAxis((5, 1), NamedTuple())))), layer_2 = ViewAxis(26:49, Axis(weight = ViewAxis(1:20, ShapedAxis((4, 5), NamedTuple())), bias = ViewAxis(21:24, ShapedAxis((4, 1), NamedTuple())))))}}}}, Tuple{Zygote.var"#1999#back#193"{typeof(identity)}, Zygote.var"#2166#back#292"{Zygote.var"#back#291"{:f, Zygote.Context{false}, Optimization.var"#156#165"{OptimizationFunction{true, Optimization.AutoZygote, var"#2#7"{var"#loss_function#6"{CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}, Float64}}, 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}, SciMLBase.NullParameters}, OptimizationFunction{true, Optimization.AutoZygote, var"#2#7"{var"#loss_function#6"{CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}, Float64}}, 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}}}, Zygote.var"#2155#back#282"{Zygote.var"#280#281"{Tuple{Tuple{Nothing, Nothing}, Tuple{}}, Zygote.Pullback{Tuple{OptimizationFunction{true, Optimization.AutoZygote, var"#2#7"{var"#loss_function#6"{CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}, Float64}}, 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}, ComponentVector{Float64, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Axis{(layer_1 = ViewAxis(1:25, Axis(weight = ViewAxis(1:20, ShapedAxis((5, 4), NamedTuple())), bias = ViewAxis(21:25, ShapedAxis((5, 1), NamedTuple())))), layer_2 = ViewAxis(26:49, Axis(weight = ViewAxis(1:20, ShapedAxis((4, 5), NamedTuple())), bias = ViewAxis(21:24, ShapedAxis((4, 1), NamedTuple())))))}}}, SciMLBase.NullParameters}, Tuple{Zygote.var"#2166#back#292"{Zygote.var"#back#291"{:f, Zygote.Context{false}, OptimizationFunction{true, Optimization.AutoZygote, var"#2#7"{var"#loss_function#6"{CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}, Float64}}, 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}, var"#2#7"{var"#loss_function#6"{CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}, Float64}}}}, Zygote.var"#2155#back#282"{Zygote.var"#280#281"{Tuple{Tuple{Nothing, Nothing}}, Zygote.Pullback{Tuple{var"#2#7"{var"#loss_function#6"{CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}, Float64}}, ComponentVector{Float64, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Axis{(layer_1 = ViewAxis(1:25, Axis(weight = ViewAxis(1:20, ShapedAxis((5, 4), NamedTuple())), bias = ViewAxis(21:25, ShapedAxis((5, 1), NamedTuple())))), layer_2 = ViewAxis(26:49, Axis(weight = ViewAxis(1:20, ShapedAxis((4, 5), NamedTuple())), bias = ViewAxis(21:24, ShapedAxis((4, 1), NamedTuple())))))}}}, SciMLBase.NullParameters}, Tuple{Zygote.var"#2166#back#292"{Zygote.var"#back#291"{:loss_function, Zygote.Context{false}, var"#2#7"{var"#loss_function#6"{CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}, Float64}}, var"#loss_function#6"{CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}, Float64}}}, Zygote.Pullback{Tuple{var"#loss_function#6"{CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}, Float64}, ComponentVector{Float64, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Axis{(layer_1 = ViewAxis(1:25, Axis(weight = ViewAxis(1:20, ShapedAxis((5, 4), NamedTuple())), bias = ViewAxis(21:25, ShapedAxis((5, 1), NamedTuple())))), layer_2 = ViewAxis(26:49, Axis(weight = ViewAxis(1:20, ShapedAxis((4, 5), NamedTuple())), bias = ViewAxis(21:24, ShapedAxis((4, 1), NamedTuple())))))}}}}, Any}}}}}}}}}, Zygote.var"#2166#back#292"{Zygote.var"#back#291"{:p, Zygote.Context{false}, Optimization.var"#156#165"{OptimizationFunction{true, Optimization.AutoZygote, var"#2#7"{var"#loss_function#6"{CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}, Float64}}, 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}, SciMLBase.NullParameters}, SciMLBase.NullParameters}}, Zygote.Pullback{Tuple{typeof(Zygote.literal_getindex), Float64, Val{1}}, Tuple{Zygote.Pullback{Tuple{typeof(getindex), Float64, Int64}, Any}}}}})(Δ::Float64)
    @ Zygote ~/.julia/packages/Zygote/oGI57/src/compiler/interface2.jl:0
 [48] #280
    @ ~/.julia/packages/Zygote/oGI57/src/lib/lib.jl:206 [inlined]
 [49] #2155#back
    @ ~/.julia/packages/ZygoteRules/4nXuu/src/adjoint.jl:71 [inlined]
 [50] Pullback
    @ ~/.julia/packages/Optimization/GEo8L/src/function/zygote.jl:34 [inlined]
 [51] (::Zygote.Pullback{Tuple{Optimization.var"#158#167"{Tuple{}, Optimization.var"#156#165"{OptimizationFunction{true, Optimization.AutoZygote, var"#2#7"{var"#loss_function#6"{CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}, Float64}}, 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}, SciMLBase.NullParameters}}, ComponentVector{Float64, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Axis{(layer_1 = ViewAxis(1:25, Axis(weight = ViewAxis(1:20, ShapedAxis((5, 4), NamedTuple())), bias = ViewAxis(21:25, ShapedAxis((5, 1), NamedTuple())))), layer_2 = ViewAxis(26:49, Axis(weight = ViewAxis(1:20, ShapedAxis((4, 5), NamedTuple())), bias = ViewAxis(21:24, ShapedAxis((4, 1), NamedTuple())))))}}}}, Tuple{Zygote.var"#2166#back#292"{Zygote.var"#back#291"{:_f, Zygote.Context{false}, Optimization.var"#158#167"{Tuple{}, Optimization.var"#156#165"{OptimizationFunction{true, Optimization.AutoZygote, var"#2#7"{var"#loss_function#6"{CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}, Float64}}, 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}, SciMLBase.NullParameters}}, Optimization.var"#156#165"{OptimizationFunction{true, Optimization.AutoZygote, var"#2#7"{var"#loss_function#6"{CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}, Float64}}, 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}, SciMLBase.NullParameters}}}, Zygote.var"#1999#back#193"{typeof(identity)}, Zygote.var"#2155#back#282"{Zygote.var"#280#281"{Tuple{Tuple{Nothing}, Tuple{}}, Zygote.Pullback{Tuple{Optimization.var"#156#165"{OptimizationFunction{true, Optimization.AutoZygote, var"#2#7"{var"#loss_function#6"{CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}, Float64}}, 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}, SciMLBase.NullParameters}, ComponentVector{Float64, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Axis{(layer_1 = ViewAxis(1:25, Axis(weight = ViewAxis(1:20, ShapedAxis((5, 4), NamedTuple())), bias = ViewAxis(21:25, ShapedAxis((5, 1), NamedTuple())))), layer_2 = ViewAxis(26:49, Axis(weight = ViewAxis(1:20, ShapedAxis((4, 5), NamedTuple())), bias = ViewAxis(21:24, ShapedAxis((4, 1), NamedTuple())))))}}}}, Tuple{Zygote.var"#1999#back#193"{typeof(identity)}, Zygote.var"#2166#back#292"{Zygote.var"#back#291"{:f, Zygote.Context{false}, Optimization.var"#156#165"{OptimizationFunction{true, Optimization.AutoZygote, var"#2#7"{var"#loss_function#6"{CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}, Float64}}, 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}, SciMLBase.NullParameters}, OptimizationFunction{true, Optimization.AutoZygote, var"#2#7"{var"#loss_function#6"{CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}, Float64}}, 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}}}, Zygote.var"#2155#back#282"{Zygote.var"#280#281"{Tuple{Tuple{Nothing, Nothing}, Tuple{}}, Zygote.Pullback{Tuple{OptimizationFunction{true, Optimization.AutoZygote, var"#2#7"{var"#loss_function#6"{CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}, Float64}}, 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}, ComponentVector{Float64, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Axis{(layer_1 = ViewAxis(1:25, Axis(weight = ViewAxis(1:20, ShapedAxis((5, 4), NamedTuple())), bias = ViewAxis(21:25, ShapedAxis((5, 1), NamedTuple())))), layer_2 = ViewAxis(26:49, Axis(weight = ViewAxis(1:20, ShapedAxis((4, 5), NamedTuple())), bias = ViewAxis(21:24, ShapedAxis((4, 1), NamedTuple())))))}}}, SciMLBase.NullParameters}, Tuple{Zygote.var"#2166#back#292"{Zygote.var"#back#291"{:f, Zygote.Context{false}, OptimizationFunction{true, Optimization.AutoZygote, var"#2#7"{var"#loss_function#6"{CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}, Float64}}, 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}, var"#2#7"{var"#loss_function#6"{CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}, Float64}}}}, Zygote.var"#2155#back#282"{Zygote.var"#280#281"{Tuple{Tuple{Nothing, Nothing}}, Zygote.Pullback{Tuple{var"#2#7"{var"#loss_function#6"{CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}, Float64}}, ComponentVector{Float64, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Axis{(layer_1 = ViewAxis(1:25, Axis(weight = ViewAxis(1:20, ShapedAxis((5, 4), NamedTuple())), bias = ViewAxis(21:25, ShapedAxis((5, 1), NamedTuple())))), layer_2 = ViewAxis(26:49, Axis(weight = ViewAxis(1:20, ShapedAxis((4, 5), NamedTuple())), bias = ViewAxis(21:24, ShapedAxis((4, 1), NamedTuple())))))}}}, SciMLBase.NullParameters}, Tuple{Zygote.var"#2166#back#292"{Zygote.var"#back#291"{:loss_function, Zygote.Context{false}, var"#2#7"{var"#loss_function#6"{CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}, Float64}}, var"#loss_function#6"{CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}, Float64}}}, Zygote.Pullback{Tuple{var"#loss_function#6"{CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}, Float64}, ComponentVector{Float64, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Axis{(layer_1 = ViewAxis(1:25, Axis(weight = ViewAxis(1:20, ShapedAxis((5, 4), NamedTuple())), bias = ViewAxis(21:25, ShapedAxis((5, 1), NamedTuple())))), layer_2 = ViewAxis(26:49, Axis(weight = ViewAxis(1:20, ShapedAxis((4, 5), NamedTuple())), bias = ViewAxis(21:24, ShapedAxis((4, 1), NamedTuple())))))}}}}, Any}}}}}}}}}, Zygote.var"#2166#back#292"{Zygote.var"#back#291"{:p, Zygote.Context{false}, Optimization.var"#156#165"{OptimizationFunction{true, Optimization.AutoZygote, var"#2#7"{var"#loss_function#6"{CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}, Float64}}, 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}, SciMLBase.NullParameters}, SciMLBase.NullParameters}}, Zygote.Pullback{Tuple{typeof(Zygote.literal_getindex), Float64, Val{1}}, Tuple{Zygote.Pullback{Tuple{typeof(getindex), Float64, Int64}, Any}}}}}}}, Zygote.var"#2166#back#292"{Zygote.var"#back#291"{:args, Zygote.Context{false}, Optimization.var"#158#167"{Tuple{}, Optimization.var"#156#165"{OptimizationFunction{true, Optimization.AutoZygote, var"#2#7"{var"#loss_function#6"{CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}, Float64}}, 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}, SciMLBase.NullParameters}}, Tuple{}}}}})(Δ::Float64)
    @ Zygote ~/.julia/packages/Zygote/oGI57/src/compiler/interface2.jl:0
 [52] (::Zygote.var"#68#69"{Zygote.Pullback{Tuple{Optimization.var"#158#167"{Tuple{}, Optimization.var"#156#165"{OptimizationFunction{true, Optimization.AutoZygote, var"#2#7"{var"#loss_function#6"{CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}, Float64}}, 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}, SciMLBase.NullParameters}}, ComponentVector{Float64, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Axis{(layer_1 = ViewAxis(1:25, Axis(weight = ViewAxis(1:20, ShapedAxis((5, 4), NamedTuple())), bias = ViewAxis(21:25, ShapedAxis((5, 1), NamedTuple())))), layer_2 = ViewAxis(26:49, Axis(weight = ViewAxis(1:20, ShapedAxis((4, 5), NamedTuple())), bias = ViewAxis(21:24, ShapedAxis((4, 1), NamedTuple())))))}}}}, Tuple{Zygote.var"#2166#back#292"{Zygote.var"#back#291"{:_f, Zygote.Context{false}, Optimization.var"#158#167"{Tuple{}, Optimization.var"#156#165"{OptimizationFunction{true, Optimization.AutoZygote, var"#2#7"{var"#loss_function#6"{CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}, Float64}}, 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}, SciMLBase.NullParameters}}, Optimization.var"#156#165"{OptimizationFunction{true, Optimization.AutoZygote, var"#2#7"{var"#loss_function#6"{CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}, Float64}}, 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}, SciMLBase.NullParameters}}}, Zygote.var"#1999#back#193"{typeof(identity)}, Zygote.var"#2155#back#282"{Zygote.var"#280#281"{Tuple{Tuple{Nothing}, Tuple{}}, Zygote.Pullback{Tuple{Optimization.var"#156#165"{OptimizationFunction{true, Optimization.AutoZygote, var"#2#7"{var"#loss_function#6"{CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}, Float64}}, 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}, SciMLBase.NullParameters}, ComponentVector{Float64, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Axis{(layer_1 = ViewAxis(1:25, Axis(weight = ViewAxis(1:20, ShapedAxis((5, 4), NamedTuple())), bias = ViewAxis(21:25, ShapedAxis((5, 1), NamedTuple())))), layer_2 = ViewAxis(26:49, Axis(weight = ViewAxis(1:20, ShapedAxis((4, 5), NamedTuple())), bias = ViewAxis(21:24, ShapedAxis((4, 1), NamedTuple())))))}}}}, Tuple{Zygote.var"#1999#back#193"{typeof(identity)}, Zygote.var"#2166#back#292"{Zygote.var"#back#291"{:f, Zygote.Context{false}, Optimization.var"#156#165"{OptimizationFunction{true, Optimization.AutoZygote, var"#2#7"{var"#loss_function#6"{CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}, Float64}}, 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}, SciMLBase.NullParameters}, OptimizationFunction{true, Optimization.AutoZygote, var"#2#7"{var"#loss_function#6"{CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}, Float64}}, 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}}}, Zygote.var"#2155#back#282"{Zygote.var"#280#281"{Tuple{Tuple{Nothing, Nothing}, Tuple{}}, Zygote.Pullback{Tuple{OptimizationFunction{true, Optimization.AutoZygote, var"#2#7"{var"#loss_function#6"{CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}, Float64}}, 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}, ComponentVector{Float64, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Axis{(layer_1 = ViewAxis(1:25, Axis(weight = ViewAxis(1:20, ShapedAxis((5, 4), NamedTuple())), bias = ViewAxis(21:25, ShapedAxis((5, 1), NamedTuple())))), layer_2 = ViewAxis(26:49, Axis(weight = ViewAxis(1:20, ShapedAxis((4, 5), NamedTuple())), bias = ViewAxis(21:24, ShapedAxis((4, 1), NamedTuple())))))}}}, SciMLBase.NullParameters}, Tuple{Zygote.var"#2166#back#292"{Zygote.var"#back#291"{:f, Zygote.Context{false}, OptimizationFunction{true, Optimization.AutoZygote, var"#2#7"{var"#loss_function#6"{CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}, Float64}}, 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}, var"#2#7"{var"#loss_function#6"{CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}, Float64}}}}, Zygote.var"#2155#back#282"{Zygote.var"#280#281"{Tuple{Tuple{Nothing, Nothing}}, Zygote.Pullback{Tuple{var"#2#7"{var"#loss_function#6"{CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}, Float64}}, ComponentVector{Float64, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Axis{(layer_1 = ViewAxis(1:25, Axis(weight = ViewAxis(1:20, ShapedAxis((5, 4), NamedTuple())), bias = ViewAxis(21:25, ShapedAxis((5, 1), NamedTuple())))), layer_2 = ViewAxis(26:49, Axis(weight = ViewAxis(1:20, ShapedAxis((4, 5), NamedTuple())), bias = ViewAxis(21:24, ShapedAxis((4, 1), NamedTuple())))))}}}, SciMLBase.NullParameters}, Tuple{Zygote.var"#2166#back#292"{Zygote.var"#back#291"{:loss_function, Zygote.Context{false}, var"#2#7"{var"#loss_function#6"{CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}, Float64}}, var"#loss_function#6"{CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}, Float64}}}, Zygote.Pullback{Tuple{var"#loss_function#6"{CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}, Float64}, ComponentVector{Float64, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Axis{(layer_1 = ViewAxis(1:25, Axis(weight = ViewAxis(1:20, ShapedAxis((5, 4), NamedTuple())), bias = ViewAxis(21:25, ShapedAxis((5, 1), NamedTuple())))), layer_2 = ViewAxis(26:49, Axis(weight = ViewAxis(1:20, ShapedAxis((4, 5), NamedTuple())), bias = ViewAxis(21:24, ShapedAxis((4, 1), NamedTuple())))))}}}}, Any}}}}}}}}}, Zygote.var"#2166#back#292"{Zygote.var"#back#291"{:p, Zygote.Context{false}, Optimization.var"#156#165"{OptimizationFunction{true, Optimization.AutoZygote, var"#2#7"{var"#loss_function#6"{CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}, Float64}}, 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}, SciMLBase.NullParameters}, SciMLBase.NullParameters}}, Zygote.Pullback{Tuple{typeof(Zygote.literal_getindex), Float64, Val{1}}, Tuple{Zygote.Pullback{Tuple{typeof(getindex), Float64, Int64}, Any}}}}}}}, Zygote.var"#2166#back#292"{Zygote.var"#back#291"{:args, Zygote.Context{false}, Optimization.var"#158#167"{Tuple{}, Optimization.var"#156#165"{OptimizationFunction{true, Optimization.AutoZygote, var"#2#7"{var"#loss_function#6"{CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}, Float64}}, 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}, SciMLBase.NullParameters}}, Tuple{}}}}}})(Δ::Float64)
    @ Zygote ~/.julia/packages/Zygote/oGI57/src/compiler/interface.jl:45
 [53] gradient(f::Function, args::ComponentVector{Float64, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Axis{(layer_1 = ViewAxis(1:25, Axis(weight = ViewAxis(1:20, ShapedAxis((5, 4), NamedTuple())), bias = ViewAxis(21:25, ShapedAxis((5, 1), NamedTuple())))), layer_2 = ViewAxis(26:49, Axis(weight = ViewAxis(1:20, ShapedAxis((4, 5), NamedTuple())), bias = ViewAxis(21:24, ShapedAxis((4, 1), NamedTuple())))))}}})
    @ Zygote ~/.julia/packages/Zygote/oGI57/src/compiler/interface.jl:97
 [54] (::Optimization.var"#157#166"{Optimization.var"#156#165"{OptimizationFunction{true, Optimization.AutoZygote, var"#2#7"{var"#loss_function#6"{CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}, Float64}}, 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}, SciMLBase.NullParameters}})(::ComponentVector{Float64, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Axis{(layer_1 = ViewAxis(1:25, Axis(weight = ViewAxis(1:20, ShapedAxis((5, 4), NamedTuple())), bias = ViewAxis(21:25, ShapedAxis((5, 1), NamedTuple())))), layer_2 = ViewAxis(26:49, Axis(weight = ViewAxis(1:20, ShapedAxis((4, 5), NamedTuple())), bias = ViewAxis(21:24, ShapedAxis((4, 1), NamedTuple())))))}}}, ::ComponentVector{Float64, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Axis{(layer_1 = ViewAxis(1:25, Axis(weight = ViewAxis(1:20, ShapedAxis((5, 4), NamedTuple())), bias = ViewAxis(21:25, ShapedAxis((5, 1), NamedTuple())))), layer_2 = ViewAxis(26:49, Axis(weight = ViewAxis(1:20, ShapedAxis((4, 5), NamedTuple())), bias = ViewAxis(21:24, ShapedAxis((4, 1), NamedTuple())))))}}})
    @ Optimization ~/.julia/packages/Optimization/GEo8L/src/function/zygote.jl:32
 [55] macro expansion
    @ ~/.julia/packages/OptimizationOptimisers/FWIuf/src/OptimizationOptimisers.jl:31 [inlined]
 [56] macro expansion
    @ ~/.julia/packages/Optimization/GEo8L/src/utils.jl:37 [inlined]
 [57] __solve(prob::OptimizationProblem{true, OptimizationFunction{true, Optimization.AutoZygote, var"#2#7"{var"#loss_function#6"{CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}, Float64}}, 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}, ComponentVector{Float64, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Axis{(layer_1 = ViewAxis(1:25, Axis(weight = ViewAxis(1:20, ShapedAxis((5, 4), NamedTuple())), bias = ViewAxis(21:25, ShapedAxis((5, 1), NamedTuple())))), layer_2 = ViewAxis(26:49, Axis(weight = ViewAxis(1:20, ShapedAxis((4, 5), NamedTuple())), bias = ViewAxis(21:24, ShapedAxis((4, 1), NamedTuple())))))}}}, SciMLBase.NullParameters, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}}, opt::Optimisers.Adam{Float32}, data::Base.Iterators.Cycle{Tuple{Optimization.NullData}}; maxiters::Int64, callback::Function, progress::Bool, save_best::Bool, kwargs::Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}})
    @ OptimizationOptimisers ~/.julia/packages/OptimizationOptimisers/FWIuf/src/OptimizationOptimisers.jl:30
 [58] #solve#595
    @ ~/.julia/packages/SciMLBase/l4PVV/src/solve.jl:86 [inlined]
 [59] trainUDEModel(neuralNetwork::Lux.Chain{NamedTuple{(:layer_1, :layer_2), Tuple{Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}, Lux.Dense{true, typeof(identity), typeof(Lux.glorot_uniform), typeof(Lux.zeros32)}}}}, knownDynamics::typeof(knownDynamics), training_data::Matrix{Float64}; needed_ps::Vector{Float64}, p_true::Int64)
    @ Main ~/NODE_Community_Forecast/test.jl:47
in expression starting at /home/jarroyoesquivel/NODE_Community_Forecast/test.jl:77
srun: error: vgpu-002: task 0: Exited with exit code 1

[ChrisRackauckas]

This issue is stuck on SciML/LinearSolve.jl#434 to avoid differentiating the linear solver kernel.