SciMLBase.jl
SciMLBase.jl copied to clipboard
SciML
try/catch is not supported when attempting to use `remake` with Zygote
I have a code that used to run 6 months ago that is not running anymore. I mentioned this error in https://github.com/FluxML/Zygote.jl/issues/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
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?
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).
I'm not entirely sure but I think @allowscalar inserts a try block? That might help narrow things down.
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
This issue is stuck on https://github.com/SciML/LinearSolve.jl/pull/434 to avoid differentiating the linear solver kernel.