Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

GPU broadcast resulted in non-concrete element type Any when running several NODEs in a single training #852

Closed
jarroyoe opened this issue Aug 17, 2023 · 10 comments
Closed

GPU broadcast resulted in non-concrete element type Any when running several NODEs in a single training #852

jarroyoe opened this issue Aug 17, 2023 · 10 comments

Comments

@jarroyoe
Copy link

jarroyoe commented Aug 17, 2023 •
edited

I'm running the following MWE:

using DiffEqFlux, Lux, DifferentialEquations, ComponentArrays, CUDA
using Optimization, OptimizationOptimJL
using Random
rng = Random.default_rng()

trainingData = [collect(1:10)';rand(3,10)]
chain = Lux.Chain(Lux.Dense(4,5),Lux.Dense(5,4))

ps, st = Lux.setup(rng,chain)
st = st |> Lux.gpu
p = Float64.(Lux.gpu(ComponentArray(ps)))
trainingData = convert(CuArray,trainingData)
tsteps = convert(CuArray,trainingData[1,:])

function predict_neuralode(p)
    CUDA.@allowscalar Lux.gpu([Array(first(NeuralODE(chain,(tsteps[i],tsteps[i+5]),
         AutoTsit5(Rosenbrock23()),saveat=tsteps[i:(i+5)])(trainingData[:,i],p,st))) for i in 1:(size(trainingData,2)-5)])
end
    
function loss_function(p)
    pred = predict_neuralode(p)
    splitData = [trainingData[:,i:(i+5)] for i in 1:(size(trainingData,2)-5)]
    loss = sum(sum.(abs2,splitData .- [pred[i] for i in 1:length(pred)]))

    return loss
end

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

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

optimizedParameters = Optimization.solve(optprob,
                                    Optim.BFGS(initial_stepnorm=0.01),
                                    callback = callback)

but it leads to:

ERROR: LoadError: GPU broadcast resulted in non-concrete element type Any.
This probably means that the function you are broadcasting contains an error or type instability.
Stacktrace:
  [1] error(s::String)
    @ Base ./error.jl:35
  [2] copy
    @ ~/.julia/packages/GPUArrays/6STCb/src/host/broadcast.jl:34 [inlined]
  [3] materialize(bc::Base.Broadcast.Broadcasted{CUDA.CuArrayStyle{1}, Nothing, OrdinaryDiffEq.var"#619#620"{CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, OrdinaryDiffEq.CompositeInterpolationData{ODEFunction{false, SciMLBase.FullSpecialize, ...}, Vector{CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}}, Vector{Float64}, Vector{Vector{CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}}}, CompositeCache{Tuple{OrdinaryDiffEq.Tsit5ConstantCache, OrdinaryDiffEq.Rosenbrock23ConstantCache{Float64, SciMLBase.TimeDerivativeWrapper{ODEFunction{false, SciMLBase.FullSpecialize, ...}, 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.UDerivativeWrapper{ODEFunction{false, SciMLBase.FullSpecialize, ...}, 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())))))}}}}, CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}, LinearAlgebra.LU{Float64, CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}, CuArray{Int32, 1, CUDA.Mem.DeviceBuffer}}, Nothing}},…}}, Nothing, DataType, 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())))))}}}, Symbol, Base.RefValue{Tuple{Int64, Int64}}, Float64, CompositeCache{Tuple{OrdinaryDiffEq.Tsit5ConstantCache, OrdinaryDiffEq.Rosenbrock23ConstantCache{Float64, SciMLBase.TimeDerivativeWrapper{ODEFunction{false, SciMLBase.FullSpecialize, ...}, 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.UDerivativeWrapper{ODEFunction{false, SciMLBase.FullSpecialize, ...}, 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())))))}}}}, CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}, LinearAlgebra.LU{Float64, CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}, CuArray{Int32, 1, CUDA.Mem.DeviceBuffer}}, Nothing}},…}, ODEFunction{false, SciMLBase.FullSpecialize, ...}, Vector{Vector{CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}}}, Vector{CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}}, Vector{Float64}}, Tuple{CuArray{Int64, 1, CUDA.Mem.DeviceBuffer}}})
    @ Base.Broadcast ./broadcast.jl:860
  [4] map(::Function, ::CuArray{Int64, 1, CUDA.Mem.DeviceBuffer})
    @ GPUArrays ~/.julia/packages/GPUArrays/6STCb/src/host/broadcast.jl:84
  [5] ode_interpolation(tvals::CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, id::OrdinaryDiffEq.CompositeInterpolationData{ODEFunction{false, SciMLBase.FullSpecialize, ...}, Vector{CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}}, Vector{Float64}, Vector{Vector{CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}}}, CompositeCache{Tuple{OrdinaryDiffEq.Tsit5ConstantCache, OrdinaryDiffEq.Rosenbrock23ConstantCache{Float64, SciMLBase.TimeDerivativeWrapper{ODEFunction{false, SciMLBase.FullSpecialize, ...}, 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.UDerivativeWrapper{ODEFunction{false, SciMLBase.FullSpecialize, ...}, 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())))))}}}}, CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}, LinearAlgebra.LU{Float64, CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}, CuArray{Int32, 1, CUDA.Mem.DeviceBuffer}}, Nothing}},…}}, idxs::Nothing, deriv::Type{Val{0}}, p::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())))))}}}, continuity::Symbol)
    @ OrdinaryDiffEq ~/.julia/packages/OrdinaryDiffEq/yspeT/src/dense/generic_dense.jl:328
  [6] CompositeInterpolationData
    @ ~/.julia/packages/OrdinaryDiffEq/yspeT/src/interp_func.jl:170 [inlined]
  [7] (::ODESolution{Float64, 2, Vector{CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}}, Nothing, Nothing, Vector{Float64}, Vector{Vector{CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}}}, ODEProblem{false, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Float64, Float64}, ...}, CompositeAlgorithm{Tuple{Tsit5{Static.False,…}, Rosenbrock23{4,true,…}},…}, OrdinaryDiffEq.CompositeInterpolationData{ODEFunction{false, SciMLBase.FullSpecialize, ...}, Vector{CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}}, Vector{Float64}, Vector{Vector{CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}}}, CompositeCache{Tuple{OrdinaryDiffEq.Tsit5ConstantCache, OrdinaryDiffEq.Rosenbrock23ConstantCache{Float64, SciMLBase.TimeDerivativeWrapper{ODEFunction{false, SciMLBase.FullSpecialize, ...}, 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.UDerivativeWrapper{ODEFunction{false, SciMLBase.FullSpecialize, ...}, 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())))))}}}}, CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}, LinearAlgebra.LU{Float64, CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}, CuArray{Int32, 1, CUDA.Mem.DeviceBuffer}}, Nothing}},…}}, DiffEqBase.Stats, Vector{Int64}})(t::CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, #unused#::Type{Val{0}}, idxs::Nothing, continuity::Symbol)
    @ SciMLBase ~/.julia/packages/SciMLBase/5yyka/src/solutions/ode_solutions.jl:80
  [8] #_#397
    @ ~/.julia/packages/SciMLBase/5yyka/src/solutions/ode_solutions.jl:66 [inlined]
  [9] AbstractODESolution (repeats 2 times)
    @ ~/.julia/packages/SciMLBase/5yyka/src/solutions/ode_solutions.jl:64 [inlined]
 [10] _concrete_solve_adjoint(::ODEProblem{false, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Float64, Float64}, ...}, ::CompositeAlgorithm{Tuple{Tsit5{Static.False,…}, Rosenbrock23{0,true,…}},…}, ::InterpolatingAdjoint{0, true, Val{:central}, ZygoteVJP}, ::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; save_start::Bool, save_end::Bool, saveat::CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, save_idxs::Nothing, kwargs::Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}})
    @ SciMLSensitivity ~/.julia/packages/SciMLSensitivity/DSyJO/src/concrete_solve.jl:332
 [11] #_solve_adjoint#53
    @ ~/.julia/packages/DiffEqBase/qPmC2/src/solve.jl:1348 [inlined]
 [12] #rrule#51
    @ ~/.julia/packages/DiffEqBase/qPmC2/src/solve.jl:1301 [inlined]
 [13] chain_rrule_kw
    @ ~/.julia/packages/Zygote/oGI57/src/compiler/chainrules.jl:235 [inlined]
 [14] macro expansion
    @ ~/.julia/packages/Zygote/oGI57/src/compiler/interface2.jl:0 [inlined]
 [15] _pullback
    @ ~/.julia/packages/Zygote/oGI57/src/compiler/interface2.jl:9 [inlined]
 [16] _apply
    @ ./boot.jl:816 [inlined]
 [17] adjoint
    @ ~/.julia/packages/Zygote/oGI57/src/lib/lib.jl:203 [inlined]
 [18] _pullback
    @ ~/.julia/packages/ZygoteRules/AIbCs/src/adjoint.jl:65 [inlined]
 [19] _pullback
    @ ~/.julia/packages/DiffEqBase/qPmC2/src/solve.jl:842 [inlined]
 [20] _pullback(::Zygote.Context{false}, ::DiffEqBase.var"##solve#27", ::InterpolatingAdjoint{0, true, Val{:central}, ZygoteVJP}, ::Nothing, ::Nothing, ::Val{true}, ::Base.Pairs{Symbol, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Symbol}, NamedTuple{(:saveat,), Tuple{CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}}}}, ::typeof(solve), ::ODEProblem{false, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Float64, Float64}, ...}, ::CompositeAlgorithm{Tuple{Tsit5{Static.False,…}, Rosenbrock23{0,true,…}},…})
    @ Zygote ~/.julia/packages/Zygote/oGI57/src/compiler/interface2.jl:0
 [21] _apply(::Function, ::Vararg{Any})
    @ Core ./boot.jl:816
 [22] adjoint
    @ ~/.julia/packages/Zygote/oGI57/src/lib/lib.jl:203 [inlined]
 [23] _pullback
    @ ~/.julia/packages/ZygoteRules/AIbCs/src/adjoint.jl:65 [inlined]
 [24] _pullback
    @ ~/.julia/packages/DiffEqBase/qPmC2/src/solve.jl:832 [inlined]
 [25] _pullback(::Zygote.Context{false}, ::CommonSolve.var"#solve##kw", ::NamedTuple{(:sensealg, :saveat), Tuple{InterpolatingAdjoint{0, true, Val{:central}, ZygoteVJP}, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}}}, ::typeof(solve), ::ODEProblem{false, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Float64, Float64}, ...}, ::CompositeAlgorithm{Tuple{Tsit5{Static.False,…}, Rosenbrock23{0,true,…}},…})
    @ Zygote ~/.julia/packages/Zygote/oGI57/src/compiler/interface2.jl:0
 [26] _apply(::Function, ::Vararg{Any})
    @ Core ./boot.jl:816
 [27] adjoint
    @ ~/.julia/packages/Zygote/oGI57/src/lib/lib.jl:203 [inlined]
 [28] _pullback
    @ ~/.julia/packages/ZygoteRules/AIbCs/src/adjoint.jl:65 [inlined]
 [29] _pullback
    @ ~/.julia/packages/DiffEqFlux/Lo0SD/src/neural_de.jl:81 [inlined]
 [30] _pullback(::Zygote.Context{false}, ::NeuralODE{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)}}}}, Nothing, Nothing, Tuple{Float64, Float64}, Tuple{CompositeAlgorithm{Tuple{Tsit5{Static.False,…}, Rosenbrock23{0,true,…}},…}}, Base.Pairs{Symbol, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Symbol}, NamedTuple{(:saveat,), Tuple{CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}}}}}, ::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())))))}}}, ::NamedTuple{(:layer_1, :layer_2), Tuple{NamedTuple{(), Tuple{}}, NamedTuple{(), Tuple{}}}})
    @ Zygote ~/.julia/packages/Zygote/oGI57/src/compiler/interface2.jl:0
 [31] _pullback
    @ ./none:0 [inlined]
 [32] _pullback(ctx::Zygote.Context{false}, f::var"#2#4"{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())))))}}}}, args::Int64)
    @ Zygote ~/.julia/packages/Zygote/oGI57/src/compiler/interface2.jl:0
 [33] #469
    @ ~/.julia/packages/Zygote/oGI57/src/lib/array.jl:201 [inlined]
 [34] iterate
    @ ./generator.jl:47 [inlined]
 [35] _collect(c::UnitRange{Int64}, itr::Base.Generator{UnitRange{Int64}, Zygote.var"#469#473"{Zygote.Context{false}, var"#2#4"{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())))))}}}}}}, #unused#::Base.EltypeUnknown, isz::Base.HasShape{1})
    @ Base ./array.jl:807
 [36] collect_similar
    @ ./array.jl:716 [inlined]
 [37] map
    @ ./abstractarray.jl:2933 [inlined]
 [38] ∇map(cx::Zygote.Context{false}, f::var"#2#4"{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())))))}}}}, args::UnitRange{Int64})
    @ Zygote ~/.julia/packages/Zygote/oGI57/src/lib/array.jl:201
 [39] _pullback(cx::Zygote.Context{false}, #unused#::typeof(collect), g::Base.Generator{UnitRange{Int64}, var"#2#4"{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/lib/array.jl:244
 [40] _pullback
    @ ~/.julia/packages/GPUArraysCore/HaQcr/src/GPUArraysCore.jl:125 [inlined]
 [41] _pullback(::Zygote.Context{false}, ::var"#1#3"{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/interface2.jl:0
 [42] #rrule_via_ad#46
    @ ~/.julia/packages/Zygote/oGI57/src/compiler/chainrules.jl:260 [inlined]
 [43] rrule_via_ad
    @ ~/.julia/packages/Zygote/oGI57/src/compiler/chainrules.jl:248 [inlined]
 [44] #1255
    @ ~/.julia/packages/ChainRules/bEtjZ/src/rulesets/Base/base.jl:258 [inlined]
 [45] task_local_storage(body::ChainRules.var"#1255#1256"{Zygote.ZygoteRuleConfig{Zygote.Context{false}}, var"#1#3"{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())))))}}}}}, key::Symbol, val::GPUArraysCore.ScalarIndexing)
    @ Base ./task.jl:292
 [46] rrule
    @ ~/.julia/packages/ChainRules/bEtjZ/src/rulesets/Base/base.jl:257 [inlined]
 [47] chain_rrule
    @ ~/.julia/packages/Zygote/oGI57/src/compiler/chainrules.jl:223 [inlined]
 [48] macro expansion
    @ ~/.julia/packages/Zygote/oGI57/src/compiler/interface2.jl:0 [inlined]
 [49] _pullback(::Zygote.Context{false}, ::typeof(task_local_storage), ::var"#1#3"{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())))))}}}}, ::Symbol, ::GPUArraysCore.ScalarIndexing)
    @ Zygote ~/.julia/packages/Zygote/oGI57/src/compiler/interface2.jl:9
 [50] macro expansion
    @ ~/.julia/packages/GPUArraysCore/HaQcr/src/GPUArraysCore.jl:124 [inlined]
 [51] _pullback
    @ ~/Phytoplankton_PINN/test.jl:16 [inlined]
 [52] _pullback(ctx::Zygote.Context{false}, f::typeof(predict_neuralode), 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/interface2.jl:0
 [53] _pullback
    @ ~/Phytoplankton_PINN/test.jl:21 [inlined]
 [54] _pullback(ctx::Zygote.Context{false}, f::typeof(loss_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/interface2.jl:0
 [55] _pullback
    @ ~/Phytoplankton_PINN/test.jl:38 [inlined]
 [56] _pullback(::Zygote.Context{false}, ::var"#11#12", ::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)
    @ Zygote ~/.julia/packages/Zygote/oGI57/src/compiler/interface2.jl:0
 [57] _apply
    @ ./boot.jl:816 [inlined]
 [58] adjoint
    @ ~/.julia/packages/Zygote/oGI57/src/lib/lib.jl:203 [inlined]
 [59] _pullback
    @ ~/.julia/packages/ZygoteRules/AIbCs/src/adjoint.jl:65 [inlined]
 [60] _pullback
    @ ~/.julia/packages/SciMLBase/5yyka/src/scimlfunctions.jl:3626 [inlined]
 [61] _pullback(::Zygote.Context{false}, ::OptimizationFunction{true, Optimization.AutoZygote, ...}, ::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)
    @ Zygote ~/.julia/packages/Zygote/oGI57/src/compiler/interface2.jl:0
 [62] _apply(::Function, ::Vararg{Any})
    @ Core ./boot.jl:816
 [63] adjoint
    @ ~/.julia/packages/Zygote/oGI57/src/lib/lib.jl:203 [inlined]
 [64] _pullback
    @ ~/.julia/packages/ZygoteRules/AIbCs/src/adjoint.jl:65 [inlined]
 [65] _pullback
    @ ~/.julia/packages/Optimization/GEo8L/src/function/zygote.jl:30 [inlined]
 [66] _pullback(ctx::Zygote.Context{false}, f::Optimization.var"#156#165"{OptimizationFunction{true, Optimization.AutoZygote, ...}, SciMLBase.NullParameters}, 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/interface2.jl:0
 [67] _apply(::Function, ::Vararg{Any})
    @ Core ./boot.jl:816
 [68] adjoint
    @ ~/.julia/packages/Zygote/oGI57/src/lib/lib.jl:203 [inlined]
 [69] _pullback
    @ ~/.julia/packages/ZygoteRules/AIbCs/src/adjoint.jl:65 [inlined]
 [70] _pullback
    @ ~/.julia/packages/Optimization/GEo8L/src/function/zygote.jl:34 [inlined]
 [71] _pullback(ctx::Zygote.Context{false}, f::Optimization.var"#158#167"{Tuple{}, Optimization.var"#156#165"{OptimizationFunction{true, Optimization.AutoZygote, ...}, SciMLBase.NullParameters}}, 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/interface2.jl:0
 [72] pullback(f::Function, cx::Zygote.Context{false}, 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:44
 [73] pullback
    @ ~/.julia/packages/Zygote/oGI57/src/compiler/interface.jl:42 [inlined]
 [74] 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:96
 [75] (::Optimization.var"#157#166"{Optimization.var"#156#165"{OptimizationFunction{true, Optimization.AutoZygote, ...}, 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
 [76] (::OptimizationOptimJL.var"#5#13"{OptimizationProblem{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())))))}}}, ...}, OptimizationOptimJL.var"#4#12"{OptimizationProblem{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())))))}}}, ...}, OptimizationFunction{false, Optimization.AutoZygote, ...}}, OptimizationFunction{false, Optimization.AutoZygote, ...}})(G::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())))))}}})
    @ OptimizationOptimJL ~/.julia/packages/OptimizationOptimJL/WqQOV/src/OptimizationOptimJL.jl:119
 [77] value_gradient!!(obj::TwiceDifferentiable{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())))))}}}, ComponentMatrix{Float64, CuArray{Float64, 2, 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())))))}, 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())))))}}}}, x::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())))))}}})
    @ NLSolversBase ~/.julia/packages/NLSolversBase/kavn7/src/interface.jl:82
 [78] initial_state(method::BFGS{LineSearches.InitialStatic{Float64}, LineSearches.HagerZhang{Float64, Base.RefValue{Bool}}, Nothing, Float64, Flat}, options::Optim.Options{Float64, OptimizationOptimJL.var"#_cb#11"{var"#9#10", BFGS{LineSearches.InitialStatic{Float64}, LineSearches.HagerZhang{Float64, Base.RefValue{Bool}}, Nothing, Float64, Flat}, Base.Iterators.Cycle{Tuple{Optimization.NullData}}}}, d::TwiceDifferentiable{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())))))}}}, ComponentMatrix{Float64, CuArray{Float64, 2, 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())))))}, 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())))))}}}}, initial_x::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())))))}}})
    @ Optim ~/.julia/packages/Optim/tP8PJ/src/multivariate/solvers/first_order/bfgs.jl:94
 [79] optimize(d::TwiceDifferentiable{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())))))}}}, ComponentMatrix{Float64, CuArray{Float64, 2, 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())))))}, 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())))))}}}}, initial_x::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())))))}}}, method::BFGS{LineSearches.InitialStatic{Float64}, LineSearches.HagerZhang{Float64, Base.RefValue{Bool}}, Nothing, Float64, Flat}, options::Optim.Options{Float64, OptimizationOptimJL.var"#_cb#11"{var"#9#10", BFGS{LineSearches.InitialStatic{Float64}, LineSearches.HagerZhang{Float64, Base.RefValue{Bool}}, Nothing, Float64, Flat}, Base.Iterators.Cycle{Tuple{Optimization.NullData}}}})
    @ Optim ~/.julia/packages/Optim/tP8PJ/src/multivariate/optimize/optimize.jl:36
 [80] ___solve(prob::OptimizationProblem{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())))))}}}, ...}, opt::BFGS{LineSearches.InitialStatic{Float64}, LineSearches.HagerZhang{Float64, Base.RefValue{Bool}}, Nothing, Float64, Flat}, data::Base.Iterators.Cycle{Tuple{Optimization.NullData}}; callback::Function, maxiters::Nothing, maxtime::Nothing, abstol::Nothing, reltol::Nothing, progress::Bool, kwargs::Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}})
    @ OptimizationOptimJL ~/.julia/packages/OptimizationOptimJL/WqQOV/src/OptimizationOptimJL.jl:164
 [81] #__solve#2
    @ ~/.julia/packages/OptimizationOptimJL/WqQOV/src/OptimizationOptimJL.jl:67 [inlined]
 [82] #solve#553
    @ ~/.julia/packages/SciMLBase/5yyka/src/solve.jl:86 [inlined]
in expression starting at /home/jarroyoesquivel/Phytoplankton_PINN/test.jl:41
srun: error: vgpu-005: task 0: Exited with exit code 1

This error doesn't show up when running loss_function(p), so that means that there is something going on in the optimization process. This error doesn't show up when training on CPU either.
@jarroyoe jarroyoe changed the title GPU broadcast resulted in non-concrete element type Any when running several NODEs GPU broadcast resulted in non-concrete element type Any when running several NODEs in a single training Aug 17, 2023
@ChrisRackauckas
Copy link
Member

CUDA.@allowscalar Lux.gpu([Array(first(NeuralODE That's some wild stuff 😅 . What happens if you just simplify the code? This is definitely not a reasonable call to make.

@jarroyoe
Copy link
Author

Running on GPUs doesn't allow the use of for loops nor reduce(vcat, which would allow to simplify the code a little bit.

I'm basing myself on the case with a single NODE that's on the tutorial. I'm pretty sure there's a simpler way to predict several NODEs with different initial conditions and time steps (maybe something with multiple broadcasts?), but it escapes my head

@jarroyoe
Copy link
Author

Partitioned the predict function a little bit, still getting the same error.

using DiffEqFlux, Lux, DifferentialEquations, ComponentArrays, CUDA
using Optimization, OptimizationOptimJL
using Random
rng = Random.default_rng()

trainingData = [collect(1:10)';rand(3,10)]
chain = Lux.Chain(Lux.Dense(4,5),Lux.Dense(5,4))

ps, st = Lux.setup(rng,chain)
st = st |> Lux.gpu
p = Float64.(Lux.gpu(ComponentArray(ps)))
trainingData = convert(CuArray,trainingData)

partitionedData = CUDA.@allowscalar [trainingData[:,i:(i+5)] for i in 1:(size(trainingData,2)-5)]
function partitionedNODE(data)
        CUDA.@allowscalar NeuralODE(chain,(first(data[1,:]),last(data[1,:])),AutoTsit5(Rosenbrock23()),saveat=data[1,:])
end

function predict_neuralode(p)
        models = partitionedNODE.(partitionedData)
        CUDA.@allowscalar [Lux.gpu(first(models[i](partitionedData[i][:,1],p,st))) for i in 1:length(partitionedData)]
end    

function loss_function(p)
        pred = predict_neuralode(p)
        loss = sum(sum.(abs2, partitionedData .- pred))

        return loss
end

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

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

optimizedParameters = Optimization.solve(optprob,
                                    Optim.BFGS(initial_stepnorm=0.01),
                                    callback = callback)

ERROR: LoadError: GPU broadcast resulted in non-concrete element type Any.
This probably means that the function you are broadcasting contains an error or type instability.
Stacktrace:
  [1] error(s::String)
    @ Base ./error.jl:35
  [2] copy
    @ ~/.julia/packages/GPUArrays/6STCb/src/host/broadcast.jl:34 [inlined]
  [3] materialize(bc::Base.Broadcast.Broadcasted{CUDA.CuArrayStyle{1}, Nothing, OrdinaryDiffEq.var"#619#620"{CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, OrdinaryDiffEq.CompositeInterpolationData{ODEFunction{false, SciMLBase.FullSpecialize, ...}, Vector{CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}}, Vector{Float64}, Vector{Vector{CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}}}, CompositeCache{Tuple{OrdinaryDiffEq.Tsit5ConstantCache, OrdinaryDiffEq.Rosenbrock23ConstantCache{Float64, SciMLBase.TimeDerivativeWrapper{ODEFunction{false, SciMLBase.FullSpecialize, ...}, 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.UDerivativeWrapper{ODEFunction{false, SciMLBase.FullSpecialize, ...}, 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())))))}}}}, CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}, LinearAlgebra.LU{Float64, CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}, CuArray{Int32, 1, CUDA.Mem.DeviceBuffer}}, Nothing}},…}}, Nothing, DataType, 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())))))}}}, Symbol, Base.RefValue{Tuple{Int64, Int64}}, Float64, CompositeCache{Tuple{OrdinaryDiffEq.Tsit5ConstantCache, OrdinaryDiffEq.Rosenbrock23ConstantCache{Float64, SciMLBase.TimeDerivativeWrapper{ODEFunction{false, SciMLBase.FullSpecialize, ...}, 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.UDerivativeWrapper{ODEFunction{false, SciMLBase.FullSpecialize, ...}, 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())))))}}}}, CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}, LinearAlgebra.LU{Float64, CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}, CuArray{Int32, 1, CUDA.Mem.DeviceBuffer}}, Nothing}},…}, ODEFunction{false, SciMLBase.FullSpecialize, ...}, Vector{Vector{CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}}}, Vector{CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}}, Vector{Float64}}, Tuple{CuArray{Int64, 1, CUDA.Mem.DeviceBuffer}}})
    @ Base.Broadcast ./broadcast.jl:860
  [4] map(::Function, ::CuArray{Int64, 1, CUDA.Mem.DeviceBuffer})
    @ GPUArrays ~/.julia/packages/GPUArrays/6STCb/src/host/broadcast.jl:84
  [5] ode_interpolation(tvals::CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, id::OrdinaryDiffEq.CompositeInterpolationData{ODEFunction{false, SciMLBase.FullSpecialize, ...}, Vector{CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}}, Vector{Float64}, Vector{Vector{CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}}}, CompositeCache{Tuple{OrdinaryDiffEq.Tsit5ConstantCache, OrdinaryDiffEq.Rosenbrock23ConstantCache{Float64, SciMLBase.TimeDerivativeWrapper{ODEFunction{false, SciMLBase.FullSpecialize, ...}, 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.UDerivativeWrapper{ODEFunction{false, SciMLBase.FullSpecialize, ...}, 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())))))}}}}, CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}, LinearAlgebra.LU{Float64, CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}, CuArray{Int32, 1, CUDA.Mem.DeviceBuffer}}, Nothing}},…}}, idxs::Nothing, deriv::Type{Val{0}}, p::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())))))}}}, continuity::Symbol)
    @ OrdinaryDiffEq ~/.julia/packages/OrdinaryDiffEq/yspeT/src/dense/generic_dense.jl:328
  [6] CompositeInterpolationData
    @ ~/.julia/packages/OrdinaryDiffEq/yspeT/src/interp_func.jl:170 [inlined]
  [7] (::ODESolution{Float64, 2, Vector{CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}}, Nothing, Nothing, Vector{Float64}, Vector{Vector{CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}}}, ODEProblem{false, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Float64, Float64}, ...}, CompositeAlgorithm{Tuple{Tsit5{Static.False,…}, Rosenbrock23{4,true,…}},…}, OrdinaryDiffEq.CompositeInterpolationData{ODEFunction{false, SciMLBase.FullSpecialize, ...}, Vector{CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}}, Vector{Float64}, Vector{Vector{CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}}}, CompositeCache{Tuple{OrdinaryDiffEq.Tsit5ConstantCache, OrdinaryDiffEq.Rosenbrock23ConstantCache{Float64, SciMLBase.TimeDerivativeWrapper{ODEFunction{false, SciMLBase.FullSpecialize, ...}, 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.UDerivativeWrapper{ODEFunction{false, SciMLBase.FullSpecialize, ...}, 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())))))}}}}, CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}, LinearAlgebra.LU{Float64, CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}, CuArray{Int32, 1, CUDA.Mem.DeviceBuffer}}, Nothing}},…}}, DiffEqBase.Stats, Vector{Int64}})(t::CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, #unused#::Type{Val{0}}, idxs::Nothing, continuity::Symbol)
    @ SciMLBase ~/.julia/packages/SciMLBase/5yyka/src/solutions/ode_solutions.jl:80
  [8] #_#397
    @ ~/.julia/packages/SciMLBase/5yyka/src/solutions/ode_solutions.jl:66 [inlined]
  [9] AbstractODESolution (repeats 2 times)
    @ ~/.julia/packages/SciMLBase/5yyka/src/solutions/ode_solutions.jl:64 [inlined]
 [10] _concrete_solve_adjoint(::ODEProblem{false, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Float64, Float64}, ...}, ::CompositeAlgorithm{Tuple{Tsit5{Static.False,…}, Rosenbrock23{0,true,…}},…}, ::InterpolatingAdjoint{0, true, Val{:central}, ZygoteVJP}, ::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; save_start::Bool, save_end::Bool, saveat::CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, save_idxs::Nothing, kwargs::Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}})
    @ SciMLSensitivity ~/.julia/packages/SciMLSensitivity/DSyJO/src/concrete_solve.jl:332
 [11] #_solve_adjoint#53
    @ ~/.julia/packages/DiffEqBase/qPmC2/src/solve.jl:1348 [inlined]
 [12] #rrule#51
    @ ~/.julia/packages/DiffEqBase/qPmC2/src/solve.jl:1301 [inlined]
 [13] chain_rrule_kw
    @ ~/.julia/packages/Zygote/oGI57/src/compiler/chainrules.jl:235 [inlined]
 [14] macro expansion
    @ ~/.julia/packages/Zygote/oGI57/src/compiler/interface2.jl:0 [inlined]
 [15] _pullback
    @ ~/.julia/packages/Zygote/oGI57/src/compiler/interface2.jl:9 [inlined]
 [16] _apply
    @ ./boot.jl:816 [inlined]
 [17] adjoint
    @ ~/.julia/packages/Zygote/oGI57/src/lib/lib.jl:203 [inlined]
 [18] _pullback
    @ ~/.julia/packages/ZygoteRules/AIbCs/src/adjoint.jl:65 [inlined]
 [19] _pullback
    @ ~/.julia/packages/DiffEqBase/qPmC2/src/solve.jl:842 [inlined]
 [20] _pullback(::Zygote.Context{false}, ::DiffEqBase.var"##solve#27", ::InterpolatingAdjoint{0, true, Val{:central}, ZygoteVJP}, ::Nothing, ::Nothing, ::Val{true}, ::Base.Pairs{Symbol, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Symbol}, NamedTuple{(:saveat,), Tuple{CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}}}}, ::typeof(solve), ::ODEProblem{false, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Float64, Float64}, ...}, ::CompositeAlgorithm{Tuple{Tsit5{Static.False,…}, Rosenbrock23{0,true,…}},…})
    @ Zygote ~/.julia/packages/Zygote/oGI57/src/compiler/interface2.jl:0
 [21] _apply(::Function, ::Vararg{Any})
    @ Core ./boot.jl:816
 [22] adjoint
    @ ~/.julia/packages/Zygote/oGI57/src/lib/lib.jl:203 [inlined]
 [23] _pullback
    @ ~/.julia/packages/ZygoteRules/AIbCs/src/adjoint.jl:65 [inlined]
 [24] _pullback
    @ ~/.julia/packages/DiffEqBase/qPmC2/src/solve.jl:832 [inlined]
 [25] _pullback(::Zygote.Context{false}, ::CommonSolve.var"#solve##kw", ::NamedTuple{(:sensealg, :saveat), Tuple{InterpolatingAdjoint{0, true, Val{:central}, ZygoteVJP}, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}}}, ::typeof(solve), ::ODEProblem{false, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Float64, Float64}, ...}, ::CompositeAlgorithm{Tuple{Tsit5{Static.False,…}, Rosenbrock23{0,true,…}},…})
    @ Zygote ~/.julia/packages/Zygote/oGI57/src/compiler/interface2.jl:0
 [26] _apply(::Function, ::Vararg{Any})
    @ Core ./boot.jl:816
 [27] adjoint
    @ ~/.julia/packages/Zygote/oGI57/src/lib/lib.jl:203 [inlined]
 [28] _pullback
    @ ~/.julia/packages/ZygoteRules/AIbCs/src/adjoint.jl:65 [inlined]
 [29] _pullback
    @ ~/.julia/packages/DiffEqFlux/Lo0SD/src/neural_de.jl:81 [inlined]
 [30] _pullback(::Zygote.Context{false}, ::NeuralODE{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)}}}}, Nothing, Nothing, Tuple{Float64, Float64}, Tuple{CompositeAlgorithm{Tuple{Tsit5{Static.False,…}, Rosenbrock23{0,true,…}},…}}, Base.Pairs{Symbol, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Symbol}, NamedTuple{(:saveat,), Tuple{CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}}}}}, ::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())))))}}}, ::NamedTuple{(:layer_1, :layer_2), Tuple{NamedTuple{(), Tuple{}}, NamedTuple{(), Tuple{}}}})
    @ Zygote ~/.julia/packages/Zygote/oGI57/src/compiler/interface2.jl:0
 [31] _pullback
    @ ./none:0 [inlined]
 [32] _pullback(ctx::Zygote.Context{false}, f::var"#8#10"{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())))))}}}, Vector{NeuralODE{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)}}}}, Nothing, Nothing, Tuple{Float64, Float64}, Tuple{CompositeAlgorithm{Tuple{Tsit5{Static.False,…}, Rosenbrock23{0,true,…}},…}}, Base.Pairs{Symbol, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Symbol}, NamedTuple{(:saveat,), Tuple{CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}}}}}}}, args::Int64)
    @ Zygote ~/.julia/packages/Zygote/oGI57/src/compiler/interface2.jl:0
 [33] #469
    @ ~/.julia/packages/Zygote/oGI57/src/lib/array.jl:201 [inlined]
 [34] iterate
    @ ./generator.jl:47 [inlined]
 [35] _collect(c::UnitRange{Int64}, itr::Base.Generator{UnitRange{Int64}, Zygote.var"#469#473"{Zygote.Context{false}, var"#8#10"{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())))))}}}, Vector{NeuralODE{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)}}}}, Nothing, Nothing, Tuple{Float64, Float64}, Tuple{CompositeAlgorithm{Tuple{Tsit5{Static.False,…}, Rosenbrock23{0,true,…}},…}}, Base.Pairs{Symbol, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Symbol}, NamedTuple{(:saveat,), Tuple{CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}}}}}}}}}, #unused#::Base.EltypeUnknown, isz::Base.HasShape{1})
    @ Base ./array.jl:807
 [36] collect_similar
    @ ./array.jl:716 [inlined]
 [37] map
    @ ./abstractarray.jl:2933 [inlined]
 [38] ∇map(cx::Zygote.Context{false}, f::var"#8#10"{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())))))}}}, Vector{NeuralODE{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)}}}}, Nothing, Nothing, Tuple{Float64, Float64}, Tuple{CompositeAlgorithm{Tuple{Tsit5{Static.False,…}, Rosenbrock23{0,true,…}},…}}, Base.Pairs{Symbol, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Symbol}, NamedTuple{(:saveat,), Tuple{CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}}}}}}}, args::UnitRange{Int64})
    @ Zygote ~/.julia/packages/Zygote/oGI57/src/lib/array.jl:201
 [39] _pullback(cx::Zygote.Context{false}, #unused#::typeof(collect), g::Base.Generator{UnitRange{Int64}, var"#8#10"{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())))))}}}, Vector{NeuralODE{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)}}}}, Nothing, Nothing, Tuple{Float64, Float64}, Tuple{CompositeAlgorithm{Tuple{Tsit5{Static.False,…}, Rosenbrock23{0,true,…}},…}}, Base.Pairs{Symbol, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Symbol}, NamedTuple{(:saveat,), Tuple{CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}}}}}}}})
    @ Zygote ~/.julia/packages/Zygote/oGI57/src/lib/array.jl:244
 [40] _pullback
    @ ~/.julia/packages/GPUArraysCore/HaQcr/src/GPUArraysCore.jl:125 [inlined]
 [41] _pullback(::Zygote.Context{false}, ::var"#7#9"{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())))))}}}, Vector{NeuralODE{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)}}}}, Nothing, Nothing, Tuple{Float64, Float64}, Tuple{CompositeAlgorithm{Tuple{Tsit5{Static.False,…}, Rosenbrock23{0,true,…}},…}}, Base.Pairs{Symbol, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Symbol}, NamedTuple{(:saveat,), Tuple{CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}}}}}}})
    @ Zygote ~/.julia/packages/Zygote/oGI57/src/compiler/interface2.jl:0
 [42] #rrule_via_ad#46
    @ ~/.julia/packages/Zygote/oGI57/src/compiler/chainrules.jl:260 [inlined]
 [43] rrule_via_ad
    @ ~/.julia/packages/Zygote/oGI57/src/compiler/chainrules.jl:248 [inlined]
 [44] #1255
    @ ~/.julia/packages/ChainRules/bEtjZ/src/rulesets/Base/base.jl:258 [inlined]
 [45] task_local_storage(body::ChainRules.var"#1255#1256"{Zygote.ZygoteRuleConfig{Zygote.Context{false}}, var"#7#9"{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())))))}}}, Vector{NeuralODE{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)}}}}, Nothing, Nothing, Tuple{Float64, Float64}, Tuple{CompositeAlgorithm{Tuple{Tsit5{Static.False,…}, Rosenbrock23{0,true,…}},…}}, Base.Pairs{Symbol, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Symbol}, NamedTuple{(:saveat,), Tuple{CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}}}}}}}}, key::Symbol, val::GPUArraysCore.ScalarIndexing)
    @ Base ./task.jl:292
 [46] rrule
    @ ~/.julia/packages/ChainRules/bEtjZ/src/rulesets/Base/base.jl:257 [inlined]
 [47] chain_rrule
    @ ~/.julia/packages/Zygote/oGI57/src/compiler/chainrules.jl:223 [inlined]
 [48] macro expansion
    @ ~/.julia/packages/Zygote/oGI57/src/compiler/interface2.jl:0 [inlined]
 [49] _pullback(::Zygote.Context{false}, ::typeof(task_local_storage), ::var"#7#9"{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())))))}}}, Vector{NeuralODE{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)}}}}, Nothing, Nothing, Tuple{Float64, Float64}, Tuple{CompositeAlgorithm{Tuple{Tsit5{Static.False,…}, Rosenbrock23{0,true,…}},…}}, Base.Pairs{Symbol, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Symbol}, NamedTuple{(:saveat,), Tuple{CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}}}}}}}, ::Symbol, ::GPUArraysCore.ScalarIndexing)
    @ Zygote ~/.julia/packages/Zygote/oGI57/src/compiler/interface2.jl:9
 [50] macro expansion
    @ ~/.julia/packages/GPUArraysCore/HaQcr/src/GPUArraysCore.jl:124 [inlined]
 [51] _pullback
    @ ~/Phytoplankton_PINN/test.jl:21 [inlined]
 [52] _pullback(ctx::Zygote.Context{false}, f::typeof(predict_neuralode), 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/interface2.jl:0
 [53] _pullback
    @ ~/Phytoplankton_PINN/test.jl:25 [inlined]
 [54] _pullback(ctx::Zygote.Context{false}, f::typeof(loss_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/interface2.jl:0
 [55] _pullback
    @ ~/Phytoplankton_PINN/test.jl:41 [inlined]
 [56] _pullback(::Zygote.Context{false}, ::var"#13#14", ::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)
    @ Zygote ~/.julia/packages/Zygote/oGI57/src/compiler/interface2.jl:0
 [57] _apply
    @ ./boot.jl:816 [inlined]
 [58] adjoint
    @ ~/.julia/packages/Zygote/oGI57/src/lib/lib.jl:203 [inlined]
 [59] _pullback
    @ ~/.julia/packages/ZygoteRules/AIbCs/src/adjoint.jl:65 [inlined]
 [60] _pullback
    @ ~/.julia/packages/SciMLBase/5yyka/src/scimlfunctions.jl:3626 [inlined]
 [61] _pullback(::Zygote.Context{false}, ::OptimizationFunction{true, Optimization.AutoZygote, ...}, ::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)
    @ Zygote ~/.julia/packages/Zygote/oGI57/src/compiler/interface2.jl:0
 [62] _apply(::Function, ::Vararg{Any})
    @ Core ./boot.jl:816
 [63] adjoint
    @ ~/.julia/packages/Zygote/oGI57/src/lib/lib.jl:203 [inlined]
 [64] _pullback
    @ ~/.julia/packages/ZygoteRules/AIbCs/src/adjoint.jl:65 [inlined]
 [65] _pullback
    @ ~/.julia/packages/Optimization/GEo8L/src/function/zygote.jl:30 [inlined]
 [66] _pullback(ctx::Zygote.Context{false}, f::Optimization.var"#156#165"{OptimizationFunction{true, Optimization.AutoZygote, ...}, SciMLBase.NullParameters}, 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/interface2.jl:0
 [67] _apply(::Function, ::Vararg{Any})
    @ Core ./boot.jl:816
 [68] adjoint
    @ ~/.julia/packages/Zygote/oGI57/src/lib/lib.jl:203 [inlined]
 [69] _pullback
    @ ~/.julia/packages/ZygoteRules/AIbCs/src/adjoint.jl:65 [inlined]
 [70] _pullback
    @ ~/.julia/packages/Optimization/GEo8L/src/function/zygote.jl:34 [inlined]
 [71] _pullback(ctx::Zygote.Context{false}, f::Optimization.var"#158#167"{Tuple{}, Optimization.var"#156#165"{OptimizationFunction{true, Optimization.AutoZygote, ...}, SciMLBase.NullParameters}}, 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/interface2.jl:0
 [72] pullback(f::Function, cx::Zygote.Context{false}, 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:44
 [73] pullback
    @ ~/.julia/packages/Zygote/oGI57/src/compiler/interface.jl:42 [inlined]
 [74] 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:96
 [75] (::Optimization.var"#157#166"{Optimization.var"#156#165"{OptimizationFunction{true, Optimization.AutoZygote, ...}, 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
 [76] (::OptimizationOptimJL.var"#5#13"{OptimizationProblem{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())))))}}}, ...}, OptimizationOptimJL.var"#4#12"{OptimizationProblem{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())))))}}}, ...}, OptimizationFunction{false, Optimization.AutoZygote, ...}}, OptimizationFunction{false, Optimization.AutoZygote, ...}})(G::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())))))}}})
    @ OptimizationOptimJL ~/.julia/packages/OptimizationOptimJL/WqQOV/src/OptimizationOptimJL.jl:119
 [77] value_gradient!!(obj::TwiceDifferentiable{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())))))}}}, ComponentMatrix{Float64, CuArray{Float64, 2, 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())))))}, 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())))))}}}}, x::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())))))}}})
    @ NLSolversBase ~/.julia/packages/NLSolversBase/kavn7/src/interface.jl:82
 [78] initial_state(method::BFGS{LineSearches.InitialStatic{Float64}, LineSearches.HagerZhang{Float64, Base.RefValue{Bool}}, Nothing, Float64, Flat}, options::Optim.Options{Float64, OptimizationOptimJL.var"#_cb#11"{var"#11#12", BFGS{LineSearches.InitialStatic{Float64}, LineSearches.HagerZhang{Float64, Base.RefValue{Bool}}, Nothing, Float64, Flat}, Base.Iterators.Cycle{Tuple{Optimization.NullData}}}}, d::TwiceDifferentiable{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())))))}}}, ComponentMatrix{Float64, CuArray{Float64, 2, 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())))))}, 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())))))}}}}, initial_x::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())))))}}})
    @ Optim ~/.julia/packages/Optim/tP8PJ/src/multivariate/solvers/first_order/bfgs.jl:94
 [79] optimize(d::TwiceDifferentiable{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())))))}}}, ComponentMatrix{Float64, CuArray{Float64, 2, 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())))))}, 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())))))}}}}, initial_x::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())))))}}}, method::BFGS{LineSearches.InitialStatic{Float64}, LineSearches.HagerZhang{Float64, Base.RefValue{Bool}}, Nothing, Float64, Flat}, options::Optim.Options{Float64, OptimizationOptimJL.var"#_cb#11"{var"#11#12", BFGS{LineSearches.InitialStatic{Float64}, LineSearches.HagerZhang{Float64, Base.RefValue{Bool}}, Nothing, Float64, Flat}, Base.Iterators.Cycle{Tuple{Optimization.NullData}}}})
    @ Optim ~/.julia/packages/Optim/tP8PJ/src/multivariate/optimize/optimize.jl:36
 [80] ___solve(prob::OptimizationProblem{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())))))}}}, ...}, opt::BFGS{LineSearches.InitialStatic{Float64}, LineSearches.HagerZhang{Float64, Base.RefValue{Bool}}, Nothing, Float64, Flat}, data::Base.Iterators.Cycle{Tuple{Optimization.NullData}}; callback::Function, maxiters::Nothing, maxtime::Nothing, abstol::Nothing, reltol::Nothing, progress::Bool, kwargs::Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}})
    @ OptimizationOptimJL ~/.julia/packages/OptimizationOptimJL/WqQOV/src/OptimizationOptimJL.jl:164
 [81] #__solve#2
    @ ~/.julia/packages/OptimizationOptimJL/WqQOV/src/OptimizationOptimJL.jl:67 [inlined]
 [82] #solve#553
    @ ~/.julia/packages/SciMLBase/5yyka/src/solve.jl:86 [inlined]
in expression starting at /home/jarroyoesquivel/Phytoplankton_PINN/test.jl:44
srun: error: vgpu-001: task 0: Exited with exit code 1

@avik-pal
Copy link
Member

reduce(vcat, ...) should work on GPUs. Can you share that version of your code and the error you get?

Lux.gpu is not type stable (as mentioned in https://lux.csail.mit.edu/dev/api/Lux/utilities#Lux.gpu), and using it in this fashion is definitely not going to work.

@jarroyoe
Copy link
Author

jarroyoe commented Sep 5, 2023 •
edited

Using reduce(hcat,...) and not changing the code anywhere else changes the error stacktrace, but the error still happens in the same line:

using DiffEqFlux, Lux, DifferentialEquations, ComponentArrays, CUDA
using Optimization, OptimizationOptimJL
using Random
rng = Random.default_rng()

trainingData = [collect(1:10)';rand(3,10)]
chain = Lux.Chain(Lux.Dense(4,5),Lux.Dense(5,4))

ps, st = Lux.setup(rng,chain)
st = st |> Lux.gpu
p = Float64.(Lux.gpu(ComponentArray(ps)))
trainingData = convert(CuArray,trainingData)

partitionedData = CUDA.@allowscalar [trainingData[:,i:(i+5)] for i in 1:(size(trainingData,2)-5)]
function partitionedNODE(data)
        CUDA.@allowscalar NeuralODE(chain,(first(data[1,:]),last(data[1,:])),AutoTsit5(Rosenbrock23()),saveat=data[1,:])
end

function predict_neuralode(p)
        models = partitionedNODE.(partitionedData)
        CUDA.@allowscalar reduce(hcat,[Lux.gpu(first(models[i](partitionedData[i][:,1],p,st))) for i in 1:length(partitionedData)])
end    

function loss_function(p)
        pred = predict_neuralode(p)
        loss = sum(sum.(abs2, partitionedData .- pred))

        return loss
end

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

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

optimizedParameters = Optimization.solve(optprob,
                                    Optim.BFGS(initial_stepnorm=0.01),
                                    callback = callback)

ERROR: LoadError: GPU broadcast resulted in non-concrete element type Any.
This probably means that the function you are broadcasting contains an error or type instability.
Stacktrace:
  [1] error(s::String)
    @ Base ./error.jl:35
  [2] copy
    @ ~/.julia/packages/GPUArrays/6STCb/src/host/broadcast.jl:34 [inlined]
  [3] materialize(bc::Base.Broadcast.Broadcasted{CUDA.CuArrayStyle{1}, Nothing, OrdinaryDiffEq.var"#619#620"{CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, OrdinaryDiffEq.CompositeInterpolationData{ODEFunction{false, SciMLBase.FullSpecialize, ...}, Vector{CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}}, Vector{Float64}, Vector{Vector{CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}}}, CompositeCache{Tuple{OrdinaryDiffEq.Tsit5ConstantCache, OrdinaryDiffEq.Rosenbrock23ConstantCache{Float64, SciMLBase.TimeDerivativeWrapper{ODEFunction{false, SciMLBase.FullSpecialize, ...}, 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.UDerivativeWrapper{ODEFunction{false, SciMLBase.FullSpecialize, ...}, 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())))))}}}}, CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}, LinearAlgebra.LU{Float64, CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}, CuArray{Int32, 1, CUDA.Mem.DeviceBuffer}}, Nothing}},…}}, Nothing, DataType, 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())))))}}}, Symbol, Base.RefValue{Tuple{Int64, Int64}}, Float64, CompositeCache{Tuple{OrdinaryDiffEq.Tsit5ConstantCache, OrdinaryDiffEq.Rosenbrock23ConstantCache{Float64, SciMLBase.TimeDerivativeWrapper{ODEFunction{false, SciMLBase.FullSpecialize, ...}, 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.UDerivativeWrapper{ODEFunction{false, SciMLBase.FullSpecialize, ...}, 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())))))}}}}, CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}, LinearAlgebra.LU{Float64, CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}, CuArray{Int32, 1, CUDA.Mem.DeviceBuffer}}, Nothing}},…}, ODEFunction{false, SciMLBase.FullSpecialize, ...}, Vector{Vector{CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}}}, Vector{CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}}, Vector{Float64}}, Tuple{CuArray{Int64, 1, CUDA.Mem.DeviceBuffer}}})
    @ Base.Broadcast ./broadcast.jl:860
  [4] map(::Function, ::CuArray{Int64, 1, CUDA.Mem.DeviceBuffer})
    @ GPUArrays ~/.julia/packages/GPUArrays/6STCb/src/host/broadcast.jl:84
  [5] ode_interpolation(tvals::CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, id::OrdinaryDiffEq.CompositeInterpolationData{ODEFunction{false, SciMLBase.FullSpecialize, ...}, Vector{CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}}, Vector{Float64}, Vector{Vector{CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}}}, CompositeCache{Tuple{OrdinaryDiffEq.Tsit5ConstantCache, OrdinaryDiffEq.Rosenbrock23ConstantCache{Float64, SciMLBase.TimeDerivativeWrapper{ODEFunction{false, SciMLBase.FullSpecialize, ...}, 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.UDerivativeWrapper{ODEFunction{false, SciMLBase.FullSpecialize, ...}, 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())))))}}}}, CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}, LinearAlgebra.LU{Float64, CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}, CuArray{Int32, 1, CUDA.Mem.DeviceBuffer}}, Nothing}},…}}, idxs::Nothing, deriv::Type{Val{0}}, p::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())))))}}}, continuity::Symbol)
    @ OrdinaryDiffEq ~/.julia/packages/OrdinaryDiffEq/yspeT/src/dense/generic_dense.jl:328
  [6] CompositeInterpolationData
    @ ~/.julia/packages/OrdinaryDiffEq/yspeT/src/interp_func.jl:170 [inlined]
  [7] (::ODESolution{Float64, 2, Vector{CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}}, Nothing, Nothing, Vector{Float64}, Vector{Vector{CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}}}, ODEProblem{false, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Float64, Float64}, ...}, CompositeAlgorithm{Tuple{Tsit5{Static.False,…}, Rosenbrock23{4,true,…}},…}, OrdinaryDiffEq.CompositeInterpolationData{ODEFunction{false, SciMLBase.FullSpecialize, ...}, Vector{CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}}, Vector{Float64}, Vector{Vector{CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}}}, CompositeCache{Tuple{OrdinaryDiffEq.Tsit5ConstantCache, OrdinaryDiffEq.Rosenbrock23ConstantCache{Float64, SciMLBase.TimeDerivativeWrapper{ODEFunction{false, SciMLBase.FullSpecialize, ...}, 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.UDerivativeWrapper{ODEFunction{false, SciMLBase.FullSpecialize, ...}, 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())))))}}}}, CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}, LinearAlgebra.LU{Float64, CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}, CuArray{Int32, 1, CUDA.Mem.DeviceBuffer}}, Nothing}},…}}, DiffEqBase.Stats, Vector{Int64}})(t::CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, #unused#::Type{Val{0}}, idxs::Nothing, continuity::Symbol)
    @ SciMLBase ~/.julia/packages/SciMLBase/5yyka/src/solutions/ode_solutions.jl:80
  [8] #_#397
    @ ~/.julia/packages/SciMLBase/5yyka/src/solutions/ode_solutions.jl:66 [inlined]
  [9] AbstractODESolution (repeats 2 times)
    @ ~/.julia/packages/SciMLBase/5yyka/src/solutions/ode_solutions.jl:64 [inlined]
 [10] _concrete_solve_adjoint(::ODEProblem{false, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Float64, Float64}, ...}, ::CompositeAlgorithm{Tuple{Tsit5{Static.False,…}, Rosenbrock23{0,true,…}},…}, ::InterpolatingAdjoint{0, true, Val{:central}, ZygoteVJP}, ::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; save_start::Bool, save_end::Bool, saveat::CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, save_idxs::Nothing, kwargs::Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}})
    @ SciMLSensitivity ~/.julia/packages/SciMLSensitivity/DSyJO/src/concrete_solve.jl:332
 [11] #_solve_adjoint#53
    @ ~/.julia/packages/DiffEqBase/qPmC2/src/solve.jl:1348 [inlined]
 [12] #rrule#51
    @ ~/.julia/packages/DiffEqBase/qPmC2/src/solve.jl:1301 [inlined]
 [13] chain_rrule_kw
    @ ~/.julia/packages/Zygote/oGI57/src/compiler/chainrules.jl:235 [inlined]
 [14] macro expansion
    @ ~/.julia/packages/Zygote/oGI57/src/compiler/interface2.jl:0 [inlined]
 [15] _pullback
    @ ~/.julia/packages/Zygote/oGI57/src/compiler/interface2.jl:9 [inlined]
 [16] _apply
    @ ./boot.jl:816 [inlined]
 [17] adjoint
    @ ~/.julia/packages/Zygote/oGI57/src/lib/lib.jl:203 [inlined]
 [18] _pullback
    @ ~/.julia/packages/ZygoteRules/AIbCs/src/adjoint.jl:65 [inlined]
 [19] _pullback
    @ ~/.julia/packages/DiffEqBase/qPmC2/src/solve.jl:842 [inlined]
 [20] _pullback(::Zygote.Context{false}, ::DiffEqBase.var"##solve#27", ::InterpolatingAdjoint{0, true, Val{:central}, ZygoteVJP}, ::Nothing, ::Nothing, ::Val{true}, ::Base.Pairs{Symbol, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Symbol}, NamedTuple{(:saveat,), Tuple{CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}}}}, ::typeof(solve), ::ODEProblem{false, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Float64, Float64}, ...}, ::CompositeAlgorithm{Tuple{Tsit5{Static.False,…}, Rosenbrock23{0,true,…}},…})
    @ Zygote ~/.julia/packages/Zygote/oGI57/src/compiler/interface2.jl:0
 [21] _apply(::Function, ::Vararg{Any})
    @ Core ./boot.jl:816
 [22] adjoint
    @ ~/.julia/packages/Zygote/oGI57/src/lib/lib.jl:203 [inlined]
 [23] _pullback
    @ ~/.julia/packages/ZygoteRules/AIbCs/src/adjoint.jl:65 [inlined]
 [24] _pullback
    @ ~/.julia/packages/DiffEqBase/qPmC2/src/solve.jl:832 [inlined]
 [25] _pullback(::Zygote.Context{false}, ::CommonSolve.var"#solve##kw", ::NamedTuple{(:sensealg, :saveat), Tuple{InterpolatingAdjoint{0, true, Val{:central}, ZygoteVJP}, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}}}, ::typeof(solve), ::ODEProblem{false, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Float64, Float64}, ...}, ::CompositeAlgorithm{Tuple{Tsit5{Static.False,…}, Rosenbrock23{0,true,…}},…})
    @ Zygote ~/.julia/packages/Zygote/oGI57/src/compiler/interface2.jl:0
 [26] _apply(::Function, ::Vararg{Any})
    @ Core ./boot.jl:816
 [27] adjoint
    @ ~/.julia/packages/Zygote/oGI57/src/lib/lib.jl:203 [inlined]
 [28] _pullback
    @ ~/.julia/packages/ZygoteRules/AIbCs/src/adjoint.jl:65 [inlined]
 [29] _pullback
    @ ~/.julia/packages/DiffEqFlux/Lo0SD/src/neural_de.jl:81 [inlined]
 [30] _pullback(::Zygote.Context{false}, ::NeuralODE{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)}}}}, Nothing, Nothing, Tuple{Float64, Float64}, Tuple{CompositeAlgorithm{Tuple{Tsit5{Static.False,…}, Rosenbrock23{0,true,…}},…}}, Base.Pairs{Symbol, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Symbol}, NamedTuple{(:saveat,), Tuple{CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}}}}}, ::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())))))}}}, ::NamedTuple{(:layer_1, :layer_2), Tuple{NamedTuple{(), Tuple{}}, NamedTuple{(), Tuple{}}}})
    @ Zygote ~/.julia/packages/Zygote/oGI57/src/compiler/interface2.jl:0
 [31] _pullback
    @ ./none:0 [inlined]
 [32] _pullback(ctx::Zygote.Context{false}, f::var"#8#10"{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())))))}}}, Vector{NeuralODE{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)}}}}, Nothing, Nothing, Tuple{Float64, Float64}, Tuple{CompositeAlgorithm{Tuple{Tsit5{Static.False,…}, Rosenbrock23{0,true,…}},…}}, Base.Pairs{Symbol, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Symbol}, NamedTuple{(:saveat,), Tuple{CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}}}}}}}, args::Int64)
    @ Zygote ~/.julia/packages/Zygote/oGI57/src/compiler/interface2.jl:0
 [33] #469
    @ ~/.julia/packages/Zygote/oGI57/src/lib/array.jl:201 [inlined]
 [34] iterate
    @ ./generator.jl:47 [inlined]
 [35] _collect(c::UnitRange{Int64}, itr::Base.Generator{UnitRange{Int64}, Zygote.var"#469#473"{Zygote.Context{false}, var"#8#10"{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())))))}}}, Vector{NeuralODE{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)}}}}, Nothing, Nothing, Tuple{Float64, Float64}, Tuple{CompositeAlgorithm{Tuple{Tsit5{Static.False,…}, Rosenbrock23{0,true,…}},…}}, Base.Pairs{Symbol, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Symbol}, NamedTuple{(:saveat,), Tuple{CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}}}}}}}}}, #unused#::Base.EltypeUnknown, isz::Base.HasShape{1})
    @ Base ./array.jl:807
 [36] collect_similar
    @ ./array.jl:716 [inlined]
 [37] map
    @ ./abstractarray.jl:2933 [inlined]
 [38] ∇map(cx::Zygote.Context{false}, f::var"#8#10"{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())))))}}}, Vector{NeuralODE{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)}}}}, Nothing, Nothing, Tuple{Float64, Float64}, Tuple{CompositeAlgorithm{Tuple{Tsit5{Static.False,…}, Rosenbrock23{0,true,…}},…}}, Base.Pairs{Symbol, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Symbol}, NamedTuple{(:saveat,), Tuple{CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}}}}}}}, args::UnitRange{Int64})
    @ Zygote ~/.julia/packages/Zygote/oGI57/src/lib/array.jl:201
 [39] _pullback(cx::Zygote.Context{false}, #unused#::typeof(collect), g::Base.Generator{UnitRange{Int64}, var"#8#10"{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())))))}}}, Vector{NeuralODE{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)}}}}, Nothing, Nothing, Tuple{Float64, Float64}, Tuple{CompositeAlgorithm{Tuple{Tsit5{Static.False,…}, Rosenbrock23{0,true,…}},…}}, Base.Pairs{Symbol, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Symbol}, NamedTuple{(:saveat,), Tuple{CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}}}}}}}})
    @ Zygote ~/.julia/packages/Zygote/oGI57/src/lib/array.jl:244
 [40] _pullback
    @ ~/.julia/packages/GPUArraysCore/HaQcr/src/GPUArraysCore.jl:125 [inlined]
 [41] _pullback(::Zygote.Context{false}, ::var"#7#9"{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())))))}}}, Vector{NeuralODE{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)}}}}, Nothing, Nothing, Tuple{Float64, Float64}, Tuple{CompositeAlgorithm{Tuple{Tsit5{Static.False,…}, Rosenbrock23{0,true,…}},…}}, Base.Pairs{Symbol, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Symbol}, NamedTuple{(:saveat,), Tuple{CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}}}}}}})
    @ Zygote ~/.julia/packages/Zygote/oGI57/src/compiler/interface2.jl:0
 [42] #rrule_via_ad#46
    @ ~/.julia/packages/Zygote/oGI57/src/compiler/chainrules.jl:260 [inlined]
 [43] rrule_via_ad
    @ ~/.julia/packages/Zygote/oGI57/src/compiler/chainrules.jl:248 [inlined]
 [44] #1255
    @ ~/.julia/packages/ChainRules/bEtjZ/src/rulesets/Base/base.jl:258 [inlined]
 [45] task_local_storage(body::ChainRules.var"#1255#1256"{Zygote.ZygoteRuleConfig{Zygote.Context{false}}, var"#7#9"{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())))))}}}, Vector{NeuralODE{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)}}}}, Nothing, Nothing, Tuple{Float64, Float64}, Tuple{CompositeAlgorithm{Tuple{Tsit5{Static.False,…}, Rosenbrock23{0,true,…}},…}}, Base.Pairs{Symbol, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Symbol}, NamedTuple{(:saveat,), Tuple{CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}}}}}}}}, key::Symbol, val::GPUArraysCore.ScalarIndexing)
    @ Base ./task.jl:292
 [46] rrule
    @ ~/.julia/packages/ChainRules/bEtjZ/src/rulesets/Base/base.jl:257 [inlined]
 [47] chain_rrule
    @ ~/.julia/packages/Zygote/oGI57/src/compiler/chainrules.jl:223 [inlined]
 [48] macro expansion
    @ ~/.julia/packages/Zygote/oGI57/src/compiler/interface2.jl:0 [inlined]
 [49] _pullback(::Zygote.Context{false}, ::typeof(task_local_storage), ::var"#7#9"{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())))))}}}, Vector{NeuralODE{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)}}}}, Nothing, Nothing, Tuple{Float64, Float64}, Tuple{CompositeAlgorithm{Tuple{Tsit5{Static.False,…}, Rosenbrock23{0,true,…}},…}}, Base.Pairs{Symbol, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Symbol}, NamedTuple{(:saveat,), Tuple{CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}}}}}}}, ::Symbol, ::GPUArraysCore.ScalarIndexing)
    @ Zygote ~/.julia/packages/Zygote/oGI57/src/compiler/interface2.jl:9
 [50] macro expansion
    @ ~/.julia/packages/GPUArraysCore/HaQcr/src/GPUArraysCore.jl:124 [inlined]
 [51] _pullback
    @ ~/Phytoplankton_PINN/test.jl:21 [inlined]
 [52] _pullback(ctx::Zygote.Context{false}, f::typeof(predict_neuralode), 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/interface2.jl:0
 [53] _pullback
    @ ~/Phytoplankton_PINN/test.jl:25 [inlined]
 [54] _pullback(ctx::Zygote.Context{false}, f::typeof(loss_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/interface2.jl:0
 [55] _pullback
    @ ~/Phytoplankton_PINN/test.jl:41 [inlined]
 [56] _pullback(::Zygote.Context{false}, ::var"#13#14", ::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)
    @ Zygote ~/.julia/packages/Zygote/oGI57/src/compiler/interface2.jl:0
 [57] _apply
    @ ./boot.jl:816 [inlined]
 [58] adjoint
    @ ~/.julia/packages/Zygote/oGI57/src/lib/lib.jl:203 [inlined]
 [59] _pullback
    @ ~/.julia/packages/ZygoteRules/AIbCs/src/adjoint.jl:65 [inlined]
 [60] _pullback
    @ ~/.julia/packages/SciMLBase/5yyka/src/scimlfunctions.jl:3626 [inlined]
 [61] _pullback(::Zygote.Context{false}, ::OptimizationFunction{true, Optimization.AutoZygote, ...}, ::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)
    @ Zygote ~/.julia/packages/Zygote/oGI57/src/compiler/interface2.jl:0
 [62] _apply(::Function, ::Vararg{Any})
    @ Core ./boot.jl:816
 [63] adjoint
    @ ~/.julia/packages/Zygote/oGI57/src/lib/lib.jl:203 [inlined]
 [64] _pullback
    @ ~/.julia/packages/ZygoteRules/AIbCs/src/adjoint.jl:65 [inlined]
 [65] _pullback
    @ ~/.julia/packages/Optimization/GEo8L/src/function/zygote.jl:30 [inlined]
 [66] _pullback(ctx::Zygote.Context{false}, f::Optimization.var"#156#165"{OptimizationFunction{true, Optimization.AutoZygote, ...}, SciMLBase.NullParameters}, 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/interface2.jl:0
 [67] _apply(::Function, ::Vararg{Any})
    @ Core ./boot.jl:816
 [68] adjoint
    @ ~/.julia/packages/Zygote/oGI57/src/lib/lib.jl:203 [inlined]
 [69] _pullback
    @ ~/.julia/packages/ZygoteRules/AIbCs/src/adjoint.jl:65 [inlined]
 [70] _pullback
    @ ~/.julia/packages/Optimization/GEo8L/src/function/zygote.jl:34 [inlined]
 [71] _pullback(ctx::Zygote.Context{false}, f::Optimization.var"#158#167"{Tuple{}, Optimization.var"#156#165"{OptimizationFunction{true, Optimization.AutoZygote, ...}, SciMLBase.NullParameters}}, 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/interface2.jl:0
 [72] pullback(f::Function, cx::Zygote.Context{false}, 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:44
 [73] pullback
    @ ~/.julia/packages/Zygote/oGI57/src/compiler/interface.jl:42 [inlined]
 [74] 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:96
 [75] (::Optimization.var"#157#166"{Optimization.var"#156#165"{OptimizationFunction{true, Optimization.AutoZygote, ...}, 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
 [76] (::OptimizationOptimJL.var"#5#13"{OptimizationProblem{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())))))}}}, ...}, OptimizationOptimJL.var"#4#12"{OptimizationProblem{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())))))}}}, ...}, OptimizationFunction{false, Optimization.AutoZygote, ...}}, OptimizationFunction{false, Optimization.AutoZygote, ...}})(G::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())))))}}})
    @ OptimizationOptimJL ~/.julia/packages/OptimizationOptimJL/WqQOV/src/OptimizationOptimJL.jl:119
 [77] value_gradient!!(obj::TwiceDifferentiable{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())))))}}}, ComponentMatrix{Float64, CuArray{Float64, 2, 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())))))}, 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())))))}}}}, x::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())))))}}})
    @ NLSolversBase ~/.julia/packages/NLSolversBase/kavn7/src/interface.jl:82
 [78] initial_state(method::BFGS{LineSearches.InitialStatic{Float64}, LineSearches.HagerZhang{Float64, Base.RefValue{Bool}}, Nothing, Float64, Flat}, options::Optim.Options{Float64, OptimizationOptimJL.var"#_cb#11"{var"#11#12", BFGS{LineSearches.InitialStatic{Float64}, LineSearches.HagerZhang{Float64, Base.RefValue{Bool}}, Nothing, Float64, Flat}, Base.Iterators.Cycle{Tuple{Optimization.NullData}}}}, d::TwiceDifferentiable{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())))))}}}, ComponentMatrix{Float64, CuArray{Float64, 2, 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())))))}, 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())))))}}}}, initial_x::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())))))}}})
    @ Optim ~/.julia/packages/Optim/tP8PJ/src/multivariate/solvers/first_order/bfgs.jl:94
 [79] optimize(d::TwiceDifferentiable{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())))))}}}, ComponentMatrix{Float64, CuArray{Float64, 2, 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())))))}, 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())))))}}}}, initial_x::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())))))}}}, method::BFGS{LineSearches.InitialStatic{Float64}, LineSearches.HagerZhang{Float64, Base.RefValue{Bool}}, Nothing, Float64, Flat}, options::Optim.Options{Float64, OptimizationOptimJL.var"#_cb#11"{var"#11#12", BFGS{LineSearches.InitialStatic{Float64}, LineSearches.HagerZhang{Float64, Base.RefValue{Bool}}, Nothing, Float64, Flat}, Base.Iterators.Cycle{Tuple{Optimization.NullData}}}})
    @ Optim ~/.julia/packages/Optim/tP8PJ/src/multivariate/optimize/optimize.jl:36
 [80] ___solve(prob::OptimizationProblem{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())))))}}}, ...}, opt::BFGS{LineSearches.InitialStatic{Float64}, LineSearches.HagerZhang{Float64, Base.RefValue{Bool}}, Nothing, Float64, Flat}, data::Base.Iterators.Cycle{Tuple{Optimization.NullData}}; callback::Function, maxiters::Nothing, maxtime::Nothing, abstol::Nothing, reltol::Nothing, progress::Bool, kwargs::Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}})
    @ OptimizationOptimJL ~/.julia/packages/OptimizationOptimJL/WqQOV/src/OptimizationOptimJL.jl:164
 [81] #__solve#2
    @ ~/.julia/packages/OptimizationOptimJL/WqQOV/src/OptimizationOptimJL.jl:67 [inlined]
 [82] #solve#553
    @ ~/.julia/packages/SciMLBase/5yyka/src/solve.jl:86 [inlined]
in expression starting at /home/jarroyoesquivel/Phytoplankton_PINN/test.jl:44
srun: error: vgpu-001: task 0: Exited with exit code 1

@avik-pal
Copy link
Member

avik-pal commented Sep 5, 2023

Yes that is expected. See my previous reply

Lux.gpu is not type stable (as mentioned in https://lux.csail.mit.edu/dev/api/Lux/utilities#Lux.gpu), and using it in this fashion is definitely not going to work.

@jarroyoe
Copy link
Author

jarroyoe commented Sep 5, 2023

Makes sense. What would be an appropriate implementation? Based on the documentation of LuxDeviceUtils, I tried changing Lux.gpu to gpu_device, but it didn't work.

using DiffEqFlux, Lux, DifferentialEquations, ComponentArrays, CUDA, LuxDeviceUtils
using Optimization, OptimizationOptimJL
using Random
rng = Random.default_rng()

trainingData = [collect(1:10)';rand(3,10)]
chain = Lux.Chain(Lux.Dense(4,5),Lux.Dense(5,4))

ps, st = Lux.setup(rng,chain)
st = st |> Lux.gpu
p = Float64.(Lux.gpu(ComponentArray(ps)))
trainingData = convert(CuArray,trainingData)

partitionedData = CUDA.@allowscalar [trainingData[:,i:(i+5)] for i in 1:(size(trainingData,2)-5)]
function partitionedNODE(data)
        CUDA.@allowscalar NeuralODE(chain,(first(data[1,:]),last(data[1,:])),AutoTsit5(Rosenbrock23()),saveat=data[1,:])
end

function predict_neuralode(p)
        models = partitionedNODE.(partitionedData)
        CUDA.@allowscalar reduce(hcat,[gpu_device(first(models[i](partitionedData[i][:,1],p,st))) for i in 1:length(partitionedData)])
end    

function loss_function(p)
        pred = predict_neuralode(p)
        loss = sum(sum.(abs2, partitionedData .- pred))

        return loss
end

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

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

optimizedParameters = Optimization.solve(optprob,
                                    Optim.BFGS(initial_stepnorm=0.01),
                                    callback = callback)

and got

ERROR: LoadError: Compiling Tuple{typeof(predict_neuralode), 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())))))}}}}: 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]
  [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
    @ ~/Phytoplankton_PINN/test.jl:25 [inlined]
 [10] _pullback(ctx::Zygote.Context{false}, f::typeof(loss_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/interface2.jl:0
 [11] _pullback
    @ ~/Phytoplankton_PINN/test.jl:41 [inlined]
 [12] _pullback(::Zygote.Context{false}, ::var"#7#8", ::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)
    @ Zygote ~/.julia/packages/Zygote/oGI57/src/compiler/interface2.jl:0
 [13] _apply
    @ ./boot.jl:816 [inlined]
 [14] adjoint
    @ ~/.julia/packages/Zygote/oGI57/src/lib/lib.jl:203 [inlined]
 [15] _pullback
    @ ~/.julia/packages/ZygoteRules/OgCVT/src/adjoint.jl:66 [inlined]
 [16] _pullback
    @ ~/.julia/packages/SciMLBase/5yyka/src/scimlfunctions.jl:3626 [inlined]
 [17] _pullback(::Zygote.Context{false}, ::OptimizationFunction{true, Optimization.AutoZygote, var"#7#8", 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)
    @ Zygote ~/.julia/packages/Zygote/oGI57/src/compiler/interface2.jl:0
 [18] _apply(::Function, ::Vararg{Any})
    @ Core ./boot.jl:816
 [19] adjoint
    @ ~/.julia/packages/Zygote/oGI57/src/lib/lib.jl:203 [inlined]
 [20] _pullback
    @ ~/.julia/packages/ZygoteRules/OgCVT/src/adjoint.jl:66 [inlined]
 [21] _pullback
    @ ~/.julia/packages/Optimization/GEo8L/src/function/zygote.jl:30 [inlined]
 [22] _pullback(ctx::Zygote.Context{false}, f::Optimization.var"#156#165"{OptimizationFunction{true, Optimization.AutoZygote, var"#7#8", 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, 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/interface2.jl:0
 [23] _apply(::Function, ::Vararg{Any})
    @ Core ./boot.jl:816
 [24] adjoint
    @ ~/.julia/packages/Zygote/oGI57/src/lib/lib.jl:203 [inlined]
 [25] _pullback
    @ ~/.julia/packages/ZygoteRules/OgCVT/src/adjoint.jl:66 [inlined]
 [26] _pullback
    @ ~/.julia/packages/Optimization/GEo8L/src/function/zygote.jl:34 [inlined]
 [27] _pullback(ctx::Zygote.Context{false}, f::Optimization.var"#158#167"{Tuple{}, Optimization.var"#156#165"{OptimizationFunction{true, Optimization.AutoZygote, var"#7#8", 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, 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/interface2.jl:0
 [28] pullback(f::Function, cx::Zygote.Context{false}, 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:44
 [29] pullback
    @ ~/.julia/packages/Zygote/oGI57/src/compiler/interface.jl:42 [inlined]
 [30] 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:96
 [31] (::Optimization.var"#157#166"{Optimization.var"#156#165"{OptimizationFunction{true, Optimization.AutoZygote, var"#7#8", 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
 [32] (::OptimizationOptimJL.var"#5#13"{OptimizationProblem{true, OptimizationFunction{true, Optimization.AutoZygote, var"#7#8", 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{}}}}, OptimizationOptimJL.var"#4#12"{OptimizationProblem{true, OptimizationFunction{true, Optimization.AutoZygote, var"#7#8", 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{}}}}, OptimizationFunction{false, Optimization.AutoZygote, OptimizationFunction{true, Optimization.AutoZygote, var"#7#8", 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}, Optimization.var"#157#166"{Optimization.var"#156#165"{OptimizationFunction{true, Optimization.AutoZygote, var"#7#8", 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"#160#169"{Optimization.var"#156#165"{OptimizationFunction{true, Optimization.AutoZygote, var"#7#8", 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"#164#173", Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED_NO_TIME), Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}}, OptimizationFunction{false, Optimization.AutoZygote, OptimizationFunction{true, Optimization.AutoZygote, var"#7#8", 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}, Optimization.var"#157#166"{Optimization.var"#156#165"{OptimizationFunction{true, Optimization.AutoZygote, var"#7#8", 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"#160#169"{Optimization.var"#156#165"{OptimizationFunction{true, Optimization.AutoZygote, var"#7#8", 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"#164#173", Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED_NO_TIME), Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}})(G::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())))))}}})
    @ OptimizationOptimJL ~/.julia/packages/OptimizationOptimJL/WqQOV/src/OptimizationOptimJL.jl:119
 [33] value_gradient!!(obj::TwiceDifferentiable{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())))))}}}, ComponentMatrix{Float64, CuArray{Float64, 2, 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())))))}, 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())))))}}}}, x::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())))))}}})
    @ NLSolversBase ~/.julia/packages/NLSolversBase/kavn7/src/interface.jl:82
 [34] initial_state(method::BFGS{LineSearches.InitialStatic{Float64}, LineSearches.HagerZhang{Float64, Base.RefValue{Bool}}, Nothing, Float64, Flat}, options::Optim.Options{Float64, OptimizationOptimJL.var"#_cb#11"{var"#5#6", BFGS{LineSearches.InitialStatic{Float64}, LineSearches.HagerZhang{Float64, Base.RefValue{Bool}}, Nothing, Float64, Flat}, Base.Iterators.Cycle{Tuple{Optimization.NullData}}}}, d::TwiceDifferentiable{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())))))}}}, ComponentMatrix{Float64, CuArray{Float64, 2, 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())))))}, 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())))))}}}}, initial_x::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())))))}}})
    @ Optim ~/.julia/packages/Optim/Adqv3/src/multivariate/solvers/first_order/bfgs.jl:94
 [35] optimize(d::TwiceDifferentiable{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())))))}}}, ComponentMatrix{Float64, CuArray{Float64, 2, 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())))))}, 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())))))}}}}, initial_x::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())))))}}}, method::BFGS{LineSearches.InitialStatic{Float64}, LineSearches.HagerZhang{Float64, Base.RefValue{Bool}}, Nothing, Float64, Flat}, options::Optim.Options{Float64, OptimizationOptimJL.var"#_cb#11"{var"#5#6", BFGS{LineSearches.InitialStatic{Float64}, LineSearches.HagerZhang{Float64, Base.RefValue{Bool}}, Nothing, Float64, Flat}, Base.Iterators.Cycle{Tuple{Optimization.NullData}}}})
    @ Optim ~/.julia/packages/Optim/Adqv3/src/multivariate/optimize/optimize.jl:36
 [36] ___solve(prob::OptimizationProblem{true, OptimizationFunction{true, Optimization.AutoZygote, var"#7#8", 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::BFGS{LineSearches.InitialStatic{Float64}, LineSearches.HagerZhang{Float64, Base.RefValue{Bool}}, Nothing, Float64, Flat}, data::Base.Iterators.Cycle{Tuple{Optimization.NullData}}; callback::Function, maxiters::Nothing, maxtime::Nothing, abstol::Nothing, reltol::Nothing, progress::Bool, kwargs::Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}})
    @ OptimizationOptimJL ~/.julia/packages/OptimizationOptimJL/WqQOV/src/OptimizationOptimJL.jl:164
 [37] #__solve#2
    @ ~/.julia/packages/OptimizationOptimJL/WqQOV/src/OptimizationOptimJL.jl:67 [inlined]
 [38] #solve#553
    @ ~/.julia/packages/SciMLBase/5yyka/src/solve.jl:86 [inlined]
 [39] top-level scope
    @ ~/Phytoplankton_PINN/test.jl:44
in expression starting at /home/jarroyoesquivel/Phytoplankton_PINN/test.jl:44
srun: error: vgpu-001: task 0: Exited with exit code 1

@avik-pal
Copy link
Member

avik-pal commented Sep 5, 2023

Even gpu_device must perform runtime checks to determine which backend is functional. You don't need any data movement functionality in your code, it is already supposed to be on the correct device

@jarroyoe
Copy link
Author

In that case, would convert(CuArray,...) suffice?

@jarroyoe
Copy link
Author

Ok I got the syntax of gpu_device now. This advanced the code a little bit I think, as the error being produced is now unrelated to the type of the Arrays.

using DiffEqFlux, Lux, DifferentialEquations, ComponentArrays, CUDA, LuxDeviceUtils
using Optimization, OptimizationOptimJL
using Random
rng = Random.default_rng()
gdev = gpu_device()

trainingData = [collect(1:10)';rand(3,10)] |> gdev
chain = Lux.Chain(Lux.Dense(4,5),Lux.Dense(5,4))

ps, st = Lux.setup(rng,chain)
st = st |> gdev
p = Float64.(ComponentArray(ps)) |> gdev

partitionedData = CUDA.@allowscalar [trainingData[:,i:(i+5)] for i in 1:(size(trainingData,2)-5)]
function partitionedNODE(data)
        CUDA.@allowscalar NeuralODE(chain,(first(data[1,:]),last(data[1,:])),AutoTsit5(Rosenbrock23()),saveat=data[1,:])
end

function predict_neuralode(p)
        models = partitionedNODE.(partitionedData)
        CUDA.@allowscalar [first(models[i](partitionedData[i][:,1],p,st)) |> gdev for i in 1:length(partitionedData)]
end

function loss_function(p)
        pred = predict_neuralode(p)
        loss = sum(sum.(abs2, partitionedData .- pred))

        return loss
end

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

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

optimizedParameters = Optimization.solve(optprob,
                                    Optim.BFGS(initial_stepnorm=0.01),
                                    callback = callback)

ERROR: LoadError: Compiling Tuple{typeof(predict_neuralode), ComponentVector{Float32, CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}, Tuple{Axis{(layer_1 = ViewAxis(1:25, Axis(weight = ViewAxis(1:20, S$
Refer to the Zygote documentation for fixes.
https://fluxml.ai/Zygote.jl/latest/limitations

@avik-pal avik-pal linked a pull request Nov 16, 2023 that will close this issue
Fix Tests and Cleanup #873
Merged
6 tasks
@avik-pal avik-pal removed a link to a pull request Nov 16, 2023
Fix Tests and Cleanup #873
Merged
6 tasks
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Assignees
No one assigned
Labels
None yet
Projects
None yet
Milestone
No milestone
Development

No branches or pull requests
3 participants
@ChrisRackauckas @avik-pal @jarroyoe