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

Scalar indexing when running UDE on GPU #832

Closed
jarroyoe opened this issue Jun 9, 2023 · 4 comments
Closed

Scalar indexing when running UDE on GPU #832

jarroyoe opened this issue Jun 9, 2023 · 4 comments

Comments

@jarroyoe
Copy link

jarroyoe commented Jun 9, 2023

I'm trying to run the following MWE of a UDE with CUDA:

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

training_data = rand(10,20)
neuralNetwork = Lux.Chain(Lux.Dense(10,20),Lux.Dense(20,10))
knownDynamics(x,p,q) = -x
pinit, st = Lux.setup(rng,neuralNetwork)
st = st |> Lux.gpu
p64 = Float64.(Lux.gpu(ComponentArray(pinit)))
training_data = Lux.gpu(training_data)
x0 = training_data[:,1] |> Lux.gpu


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

        du = knownPred .+ nnPred
end

nn_dynamics!(du,u,p,t) = ude!(du,u,p,t,nothing)
prob_nn = ODEProblem(nn_dynamics!,x0, (Float64(1),Float64(size(training_data,2))), p64)
function predict(p, X = x0)
        _prob = remake(prob_nn, u0 = X, tspan = (Float64(1),Float64(size(training_data,2))), p = p)
        CUDA.@allowscalar Lux.gpu(solve(_prob, AutoTsit5(Rosenbrock23()), saveat = 1.,
                abstol=1e-6, reltol=1e-6
                ))
end

function loss_function(p)
        pred = predict(p)
        loss = sum(abs2,training_data .- pred)/size(training_data,2)
        return loss
end

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

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

I get the following error, despite not seeing where I'm doing any scalar indexing when solving the UDE.

ERROR: LoadError: Scalar indexing is disallowed.
Invocation of getindex resulted in scalar indexing of a GPU array.
This is typically caused by calling an iterating implementation of a method.
Such implementations *do not* execute on the GPU, but very slowly on the CPU,
and therefore are only permitted from the REPL for prototyping purposes.
If you did intend to index this array, annotate the caller with @allowscalar.
Stacktrace:
  [1] error(s::String)
    @ Base ./error.jl:35
  [2] assertscalar(op::String)
    @ GPUArraysCore ~/.julia/packages/GPUArraysCore/HaQcr/src/GPUArraysCore.jl:103
  [3] getindex
    @ ~/.julia/packages/GPUArrays/6STCb/src/host/indexing.jl:9 [inlined]
  [4] _generic_matmatmul!(C::Adjoint{Float64, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}}, tA::Char, tB::Char, A::CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, B::CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}, _add::LinearAlgebra.MulAddMul{true, true, Bool, Bool})
    @ LinearAlgebra /carnegie/binaries/centos7/julia/1.8.5/share/julia/stdlib/v1.8/LinearAlgebra/src/matmul.jl:990
  [5] generic_matmatmul!
    @ /carnegie/binaries/centos7/julia/1.8.5/share/julia/stdlib/v1.8/LinearAlgebra/src/matmul.jl:847 [inlined]
  [6] mul!
    @ /carnegie/binaries/centos7/julia/1.8.5/share/julia/stdlib/v1.8/LinearAlgebra/src/matmul.jl:437 [inlined]
  [7] mul!
    @ /carnegie/binaries/centos7/julia/1.8.5/share/julia/stdlib/v1.8/LinearAlgebra/src/matmul.jl:276 [inlined]
  [8] _vecjacobian!
    @ ~/.julia/packages/SciMLSensitivity/DSyJO/src/derivative_wrappers.jl:275 [inlined]
  [9] #vecjacobian!#28
    @ ~/.julia/packages/SciMLSensitivity/DSyJO/src/derivative_wrappers.jl:224 [inlined]
 [10] Long CUDA Description @ SciMLSensitivity ~/.julia/packages/SciMLSensitivity/DSyJO/src/interpolating_adjoint.jl:133
 [11] ODEFunction
    @ ~/.julia/packages/SciMLBase/5yyka/src/scimlfunctions.jl:2126 [inlined]
 [12] initialize!(integrator::ODEIntegrator{true,CompositeAlgorithm{Tuple{Tsit5{Static.False,…}, Rosenbrock23{12,true,…}},…},CuArray{Float64, 1, CUDA.Mem.DeviceBuffer},Float64,…}, cache::Tsit5Cache{CuArray{Float64, 1, CUDA.Mem.DeviceBuffer},…})
    @ OrdinaryDiffEq ~/.julia/packages/OrdinaryDiffEq/yspeT/src/perform_step/low_order_rk_perform_step.jl:766
 [13] initialize!(integrator::ODEIntegrator{true,CompositeAlgorithm{Tuple{Tsit5{Static.False,…}, Rosenbrock23{12,true,…}},…},CuArray{Float64, 1, CUDA.Mem.DeviceBuffer},Float64,…}, cache::CompositeCache{Tuple{Tsit5Cache{CuArray{Float64, 1, CUDA.Mem.DeviceBuffer},…}, Rosenbrock23Cache{CuArray{Float64, 1, CUDA.Mem.DeviceBuffer},…}},…})
    @ OrdinaryDiffEq ~/.julia/packages/OrdinaryDiffEq/yspeT/src/perform_step/composite_perform_step.jl:51
 [14] __init(prob::ODEProblem{true, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Float64, Float64}, ...}, alg::CompositeAlgorithm{Tuple{Tsit5{Static.False,…}, Rosenbrock23{12,true,…}},…}, timeseries_init::Tuple{}, ts_init::Tuple{}, ks_init::Tuple{}, recompile::Type{Val{true}}; saveat::Vector{Float32}, tstops::Vector{Float64}, d_discontinuities::Tuple{}, save_idxs::Nothing, save_everystep::Bool, save_on::Bool, save_start::Bool, save_end::Nothing, callback::CallbackSet{Tuple{}, Tuple{DiscreteCallback{...}}}, dense::Bool, calck::Bool, dt::Float64, dtmin::Nothing, dtmax::Float64, force_dtmin::Bool, adaptive::Bool, gamma::Rational{Int64}, abstol::Float64, reltol::Float64, qmin::Rational{Int64}, qmax::Int64, qsteady_min::Int64, qsteady_max::Int64, beta1::Nothing, beta2::Nothing, qoldinit::Rational{Int64}, controller::Nothing, fullnormalize::Bool, failfactor::Int64, maxiters::Int64, internalnorm::typeof(DiffEqBase.ODE_DEFAULT_NORM), internalopnorm::typeof(opnorm), isoutofdomain::typeof(DiffEqBase.ODE_DEFAULT_ISOUTOFDOMAIN), unstable_check::typeof(DiffEqBase.ODE_DEFAULT_UNSTABLE_CHECK), verbose::Bool, timeseries_errors::Bool, dense_errors::Bool, advance_to_tstop::Bool, stop_at_next_tstop::Bool, initialize_save::Bool, progress::Bool, progress_steps::Int64, progress_name::String, progress_message::typeof(DiffEqBase.ODE_DEFAULT_PROG_MESSAGE), userdata::Nothing, allow_extrapolation::Bool, initialize_integrator::Bool, alias_u0::Bool, alias_du0::Bool, initializealg::OrdinaryDiffEq.DefaultInit, kwargs::Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}})
    @ OrdinaryDiffEq ~/.julia/packages/OrdinaryDiffEq/yspeT/src/solve.jl:499
 [15] __solve(::ODEProblem{true, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Float64, Float64}, ...}, ::CompositeAlgorithm{Tuple{Tsit5{Static.False,…}, Rosenbrock23{12,true,…}},…}; kwargs::Base.Pairs{Symbol, Any, NTuple{8, Symbol}, NamedTuple{(:callback, :save_everystep, :save_start, :saveat, :tstops, :abstol, :reltol, :verbose), Tuple{CallbackSet{Tuple{}, Tuple{DiscreteCallback{...}}}, Bool, Bool, Vector{Float32}, Vector{Float64}, Float64, Float64, Bool}}})
    @ OrdinaryDiffEq ~/.julia/packages/OrdinaryDiffEq/yspeT/src/solve.jl:5
 [16] #solve_call#22
    @ ~/.julia/packages/DiffEqBase/qPmC2/src/solve.jl:509 [inlined]
 [17] solve_up(prob::ODEProblem{true, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Float64, Float64}, ...}, sensealg::Nothing, u0::CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, p::ComponentVector{Float64, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Axis{(layer_1 = ViewAxis(1:220, Axis(weight = ViewAxis(1:200, ShapedAxis((20, 10), NamedTuple())), bias = ViewAxis(201:220, ShapedAxis((20, 1), NamedTuple())))), layer_2 = ViewAxis(221:430, Axis(weight = ViewAxis(1:200, ShapedAxis((10, 20), NamedTuple())), bias = ViewAxis(201:210, ShapedAxis((10, 1), NamedTuple())))))}}}, args::CompositeAlgorithm{Tuple{Tsit5{Static.False,…}, Rosenbrock23{0,true,…}},…}; kwargs::Base.Pairs{Symbol, Any, NTuple{7, Symbol}, NamedTuple{(:save_everystep, :save_start, :saveat, :tstops, :abstol, :reltol, :verbose), Tuple{Bool, Bool, Vector{Float32}, Vector{Float64}, Float64, Float64, Bool}}})
    @ DiffEqBase ~/.julia/packages/DiffEqBase/qPmC2/src/solve.jl:932
 [18] solve(prob::ODEProblem{true, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Float64, Float64}, ...}, args::CompositeAlgorithm{Tuple{Tsit5{Static.False,…}, Rosenbrock23{0,true,…}},…}; sensealg::Nothing, u0::Nothing, p::Nothing, wrap::Val{true}, kwargs::Base.Pairs{Symbol, Any, NTuple{7, Symbol}, NamedTuple{(:save_everystep, :save_start, :saveat, :tstops, :abstol, :reltol, :verbose), Tuple{Bool, Bool, Vector{Float32}, Vector{Float64}, Float64, Float64, Bool}}})
    @ DiffEqBase ~/.julia/packages/DiffEqBase/qPmC2/src/solve.jl:842
 [19] _adjoint_sensitivities(sol::ODESolution{Float32, 2, Vector{CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}}, Nothing, Nothing, Vector{Float64}, Vector{Vector{CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}}}, ODEProblem{true, CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}, Tuple{Float64, Float64}, ...}, CompositeAlgorithm{Tuple{Tsit5{Static.False,…}, Rosenbrock23{10,true,…}},…}, OrdinaryDiffEq.CompositeInterpolationData{ODEFunction{true, SciMLBase.AutoSpecialize, ...}, Vector{CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}}, Vector{Float64}, Vector{Vector{CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}}}, CompositeCache{Tuple{Tsit5Cache{CuArray{Float32, 1, CUDA.Mem.DeviceBuffer},…}, Rosenbrock23Cache{CuArray{Float32, 1, CUDA.Mem.DeviceBuffer},…}},…}}, DiffEqBase.Stats, Vector{Int64}}, sensealg::InterpolatingAdjoint{0, false, Val{:central}, Bool}, alg::CompositeAlgorithm{Tuple{Tsit5{Static.False,…}, Rosenbrock23{0,true,…}},…}; t::StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}, dgdu_discrete::Function, dgdp_discrete::Nothing, dgdu_continuous::Nothing, dgdp_continuous::Nothing, g::Nothing, abstol::Float64, reltol::Float64, checkpoints::Vector{Float64}, corfunc_analytical::Nothing, callback::Nothing, kwargs::Base.Pairs{Symbol, Bool, Tuple{Symbol}, NamedTuple{(:verbose,), Tuple{Bool}}})
    @ SciMLSensitivity ~/.julia/packages/SciMLSensitivity/DSyJO/src/sensitivity_interface.jl:433
 [20] #adjoint_sensitivities#68
    @ ~/.julia/packages/SciMLSensitivity/DSyJO/src/sensitivity_interface.jl:386 [inlined]
 [21] (::SciMLSensitivity.var"#adjoint_sensitivity_backpass#274"{Base.Pairs{Symbol, Real, Tuple{Symbol, Symbol, Symbol}, NamedTuple{(:verbose, :abstol, :reltol), Tuple{Bool, Float64, Float64}}}, CompositeAlgorithm{Tuple{Tsit5{Static.False,…}, Rosenbrock23{0,true,…}},…}, InterpolatingAdjoint{0, false, Val{:central}, Bool}, CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}, ComponentVector{Float64, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Axis{(layer_1 = ViewAxis(1:220, Axis(weight = ViewAxis(1:200, ShapedAxis((20, 10), NamedTuple())), bias = ViewAxis(201:220, ShapedAxis((20, 1), NamedTuple())))), layer_2 = ViewAxis(221:430, Axis(weight = ViewAxis(1:200, ShapedAxis((10, 20), NamedTuple())), bias = ViewAxis(201:210, ShapedAxis((10, 1), NamedTuple())))))}}}, SciMLBase.ChainRulesOriginator, Tuple{}, Colon, NamedTuple{(:verbose, :abstol, :reltol), Tuple{Bool, Float64, Float64}}})(Δ::CuArray{Float32, 2, CUDA.Mem.DeviceBuffer})
    @ SciMLSensitivity ~/.julia/packages/SciMLSensitivity/DSyJO/src/concrete_solve.jl:476
 [22] ZBack
    @ ~/.julia/packages/Zygote/oGI57/src/compiler/chainrules.jl:211 [inlined]
 [23] (::Zygote.var"#kw_zpullback#45"{SciMLSensitivity.var"#adjoint_sensitivity_backpass#274"{Base.Pairs{Symbol, Real, Tuple{Symbol, Symbol, Symbol}, NamedTuple{(:verbose, :abstol, :reltol), Tuple{Bool, Float64, Float64}}}, CompositeAlgorithm{Tuple{Tsit5{Static.False,…}, Rosenbrock23{0,true,…}},…}, InterpolatingAdjoint{0, false, Val{:central}, Bool}, CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}, ComponentVector{Float64, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Axis{(layer_1 = ViewAxis(1:220, Axis(weight = ViewAxis(1:200, ShapedAxis((20, 10), NamedTuple())), bias = ViewAxis(201:220, ShapedAxis((20, 1), NamedTuple())))), layer_2 = ViewAxis(221:430, Axis(weight = ViewAxis(1:200, ShapedAxis((10, 20), NamedTuple())), bias = ViewAxis(201:210, ShapedAxis((10, 1), NamedTuple())))))}}}, SciMLBase.ChainRulesOriginator, Tuple{}, Colon, NamedTuple{(:verbose, :abstol, :reltol), Tuple{Bool, Float64, Float64}}}})(dy::CuArray{Float32, 2, CUDA.Mem.DeviceBuffer})
    @ Zygote ~/.julia/packages/Zygote/oGI57/src/compiler/chainrules.jl:237
 [24] #208
    @ ~/.julia/packages/Zygote/oGI57/src/lib/lib.jl:206 [inlined]
 [25] (::Zygote.var"#2066#back#210"{Zygote.var"#208#209"{Tuple{NTuple{6, Nothing}, Tuple{Nothing}}, Zygote.var"#kw_zpullback#45"{SciMLSensitivity.var"#adjoint_sensitivity_backpass#274"{Base.Pairs{Symbol, Real, Tuple{Symbol, Symbol, Symbol}, NamedTuple{(:verbose, :abstol, :reltol), Tuple{Bool, Float64, Float64}}}, CompositeAlgorithm{Tuple{Tsit5{Static.False,…}, Rosenbrock23{0,true,…}},…}, InterpolatingAdjoint{0, false, Val{:central}, Bool}, CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}, ComponentVector{Float64, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Axis{(layer_1 = ViewAxis(1:220, Axis(weight = ViewAxis(1:200, ShapedAxis((20, 10), NamedTuple())), bias = ViewAxis(201:220, ShapedAxis((20, 1), NamedTuple())))), layer_2 = ViewAxis(221:430, Axis(weight = ViewAxis(1:200, ShapedAxis((10, 20), NamedTuple())), bias = ViewAxis(201:210, ShapedAxis((10, 1), NamedTuple())))))}}}, SciMLBase.ChainRulesOriginator, Tuple{}, Colon, NamedTuple{(:verbose, :abstol, :reltol), Tuple{Bool, Float64, Float64}}}}}})(Δ::CuArray{Float32, 2, CUDA.Mem.DeviceBuffer})
    @ Zygote ~/.julia/packages/ZygoteRules/AIbCs/src/adjoint.jl:67
 [26] Pullback
    @ ~/.julia/packages/DiffEqBase/qPmC2/src/solve.jl:842 [inlined]
 [27] (::Zygote.Pullback{Tuple{DiffEqBase.var"##solve#27", Nothing, Nothing, Nothing, Val{true}, Base.Pairs{Symbol, Float64, Tuple{Symbol, Symbol, Symbol}, NamedTuple{(:saveat, :abstol, :reltol), Tuple{Float64, Float64, Float64}}}, typeof(solve), ODEProblem{true, CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}, Tuple{Float64, Float64}, ...}, CompositeAlgorithm{Tuple{Tsit5{Static.False,…}, Rosenbrock23{0,true,…}},…}}, Any})(Δ::CuArray{Float32, 2, CUDA.Mem.DeviceBuffer})
    @ Zygote ~/.julia/packages/Zygote/oGI57/src/compiler/interface2.jl:0
 [28] (::Zygote.var"#208#209"{Tuple{NTuple{7, Nothing}, Tuple{Nothing}}, Zygote.Pullback{Tuple{DiffEqBase.var"##solve#27", Nothing, Nothing, Nothing, Val{true}, Base.Pairs{Symbol, Float64, Tuple{Symbol, Symbol, Symbol}, NamedTuple{(:saveat, :abstol, :reltol), Tuple{Float64, Float64, Float64}}}, typeof(solve), ODEProblem{true, CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}, Tuple{Float64, Float64}, ...}, CompositeAlgorithm{Tuple{Tsit5{Static.False,…}, Rosenbrock23{0,true,…}},…}}, Any}})(Δ::CuArray{Float32, 2, CUDA.Mem.DeviceBuffer})
    @ Zygote ~/.julia/packages/Zygote/oGI57/src/lib/lib.jl:206
 [29] (::Zygote.var"#2066#back#210"{Zygote.var"#208#209"{Tuple{NTuple{7, Nothing}, Tuple{Nothing}}, Zygote.Pullback{Tuple{DiffEqBase.var"##solve#27", Nothing, Nothing, Nothing, Val{true}, Base.Pairs{Symbol, Float64, Tuple{Symbol, Symbol, Symbol}, NamedTuple{(:saveat, :abstol, :reltol), Tuple{Float64, Float64, Float64}}}, typeof(solve), ODEProblem{true, CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}, Tuple{Float64, Float64}, ...}, CompositeAlgorithm{Tuple{Tsit5{Static.False,…}, Rosenbrock23{0,true,…}},…}}, Any}}})(Δ::CuArray{Float32, 2, CUDA.Mem.DeviceBuffer})
    @ Zygote ~/.julia/packages/ZygoteRules/AIbCs/src/adjoint.jl:67
 [30] Pullback
    @ ~/.julia/packages/DiffEqBase/qPmC2/src/solve.jl:832 [inlined]
 [31] (::Zygote.Pullback{Tuple{CommonSolve.var"#solve##kw", NamedTuple{(:saveat, :abstol, :reltol), Tuple{Float64, Float64, Float64}}, typeof(solve), ODEProblem{true, CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}, Tuple{Float64, Float64}, ...}, CompositeAlgorithm{Tuple{Tsit5{Static.False,…}, Rosenbrock23{0,true,…}},…}}, Any})(Δ::CuArray{Float32, 2, CUDA.Mem.DeviceBuffer})
    @ Zygote ~/.julia/packages/Zygote/oGI57/src/compiler/interface2.jl:0
 [32] Pullback
    @ ~/.julia/packages/GPUArraysCore/HaQcr/src/GPUArraysCore.jl:125 [inlined]
 [33] Long CUDA Description
    @ Zygote ~/.julia/packages/Zygote/oGI57/src/compiler/interface2.jl:0
 [34] Long CUDA Description
    @ Zygote ~/.julia/packages/Zygote/oGI57/src/compiler/chainrules.jl:263
 [35] task_local_storage_pullback
    @ ~/.julia/packages/ChainRules/bEtjZ/src/rulesets/Base/base.jl:261 [inlined]
 [36] ZBack
    @ ~/.julia/packages/Zygote/oGI57/src/compiler/chainrules.jl:211 [inlined]
 [37] macro expansion
    @ ~/.julia/packages/GPUArraysCore/HaQcr/src/GPUArraysCore.jl:124 [inlined]
 [38] Pullback
    @ ~/NODE_Community_Forecast/test.jl:28 [inlined]
 [39]Long CUDA Description
    @ Zygote ~/.julia/packages/Zygote/oGI57/src/compiler/interface2.jl:0
 [40] Pullback
    @ ~/NODE_Community_Forecast/test.jl:27 [inlined]
 [41] Long CUDA Description
    @ Zygote ~/.julia/packages/Zygote/oGI57/src/compiler/interface2.jl:0
 [42] Pullback
    @ ~/NODE_Community_Forecast/test.jl:34 [inlined]
 [43] Pullback
    @ ~/NODE_Community_Forecast/test.jl:40 [inlined]
 [44] #208
    @ ~/.julia/packages/Zygote/oGI57/src/lib/lib.jl:206 [inlined]
 [45] #2066#back
    @ ~/.julia/packages/ZygoteRules/AIbCs/src/adjoint.jl:67 [inlined]
 [46] Pullback
    @ ~/.julia/packages/SciMLBase/5yyka/src/scimlfunctions.jl:3626 [inlined]
 [47] #208
    @ ~/.julia/packages/Zygote/oGI57/src/lib/lib.jl:206 [inlined]
 [48] #2066#back
    @ ~/.julia/packages/ZygoteRules/AIbCs/src/adjoint.jl:67 [inlined]
 [49] Pullback
    @ ~/.julia/packages/Optimization/GEo8L/src/function/zygote.jl:30 [inlined]
 [50] Long CUDA Description
    @ Zygote ~/.julia/packages/Zygote/oGI57/src/compiler/interface2.jl:0
 [51] #208
    @ ~/.julia/packages/Zygote/oGI57/src/lib/lib.jl:206 [inlined]
 [52] #2066#back
    @ ~/.julia/packages/ZygoteRules/AIbCs/src/adjoint.jl:67 [inlined]
 [53] Pullback
    @ ~/.julia/packages/Optimization/GEo8L/src/function/zygote.jl:34 [inlined]
 [54] Long CUDA Description
    @ Zygote ~/.julia/packages/Zygote/oGI57/src/compiler/interface2.jl:0
 [55] Long CUDA Description
    @ Zygote ~/.julia/packages/Zygote/oGI57/src/compiler/interface.jl:45
 [56] gradient(f::Function, args::ComponentVector{Float64, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Axis{(layer_1 = ViewAxis(1:220, Axis(weight = ViewAxis(1:200, ShapedAxis((20, 10), NamedTuple())), bias = ViewAxis(201:220, ShapedAxis((20, 1), NamedTuple())))), layer_2 = ViewAxis(221:430, Axis(weight = ViewAxis(1:200, ShapedAxis((10, 20), NamedTuple())), bias = ViewAxis(201:210, ShapedAxis((10, 1), NamedTuple())))))}}})
    @ Zygote ~/.julia/packages/Zygote/oGI57/src/compiler/interface.jl:97
 [57] (::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:220, Axis(weight = ViewAxis(1:200, ShapedAxis((20, 10), NamedTuple())), bias = ViewAxis(201:220, ShapedAxis((20, 1), NamedTuple())))), layer_2 = ViewAxis(221:430, Axis(weight = ViewAxis(1:200, ShapedAxis((10, 20), NamedTuple())), bias = ViewAxis(201:210, ShapedAxis((10, 1), NamedTuple())))))}}}, ::ComponentVector{Float64, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Axis{(layer_1 = ViewAxis(1:220, Axis(weight = ViewAxis(1:200, ShapedAxis((20, 10), NamedTuple())), bias = ViewAxis(201:220, ShapedAxis((20, 1), NamedTuple())))), layer_2 = ViewAxis(221:430, Axis(weight = ViewAxis(1:200, ShapedAxis((10, 20), NamedTuple())), bias = ViewAxis(201:210, ShapedAxis((10, 1), NamedTuple())))))}}})
    @ Optimization ~/.julia/packages/Optimization/GEo8L/src/function/zygote.jl:32
 [58] macro expansion
    @ ~/.julia/packages/OptimizationOptimisers/FWIuf/src/OptimizationOptimisers.jl:31 [inlined]
 [59] macro expansion
    @ ~/.julia/packages/Optimization/GEo8L/src/utils.jl:37 [inlined]
 [60] __solve(prob::OptimizationProblem{true, ComponentVector{Float64, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Axis{(layer_1 = ViewAxis(1:220, Axis(weight = ViewAxis(1:200, ShapedAxis((20, 10), NamedTuple())), bias = ViewAxis(201:220, ShapedAxis((20, 1), NamedTuple())))), layer_2 = ViewAxis(221:430, Axis(weight = ViewAxis(1:200, ShapedAxis((10, 20), NamedTuple())), bias = ViewAxis(201:210, ShapedAxis((10, 1), NamedTuple())))))}}}, ...}, opt::Optimisers.Adam{Float32}, data::Base.Iterators.Cycle{Tuple{Optimization.NullData}}; maxiters::Int64, callback::Function, progress::Bool, save_best::Bool, kwargs::Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}})
    @ OptimizationOptimisers ~/.julia/packages/OptimizationOptimisers/FWIuf/src/OptimizationOptimisers.jl:30
 [61] #solve#553
    @ ~/.julia/packages/SciMLBase/5yyka/src/solve.jl:86 [inlined]
in expression starting at /home/jarroyoesquivel/NODE_Community_Forecast/test.jl:43
@ChrisRackauckas
Copy link
Member

It's in the adjoint handling of the mutation. Do:

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

        knownPred .+ nnPred
end

You weren't even mutating du anyways, so this is just better.

@jarroyoe
Copy link
Author

jarroyoe commented Jun 9, 2023

Thanks for your prompt response! There seems to be a hidden Float32 when doing that. I know from running NODEs that the neural network gives a Float64, so what else could it be?

ERROR: LoadError: Detected non-constant types in an out-of-place ODE solve, i.e. for
`du = f(u,p,t)` we see `typeof(du) !== typeof(u/t)`. This is not
supported by OrdinaryDiffEq.jl's solvers. Please either make `f`
type-constant (i.e. typeof(du) === typeof(u/t)) or use the mutating
in-place form `f(du,u,p,t)` (which is type-constant by construction).

Note that one common case for this is when computing with GPUs, using
`Float32` for `u0` and `Float64` for `tspan`. To correct this, ensure
that the element type of `tspan` matches the preferred compute type,
for example `ODEProblem(f,0f0,(0f0,1f0))` for `Float32`-based time.

typeof(u/t) = CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}
typeof(du) = CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}

Stacktrace:
  [1] ode_determine_initdt(u0::CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}, t::Float64, tdir::Float64, dtmax::Float64, abstol::Float64, reltol::Float64, internalnorm::typeof(DiffEqBase.ODE_DEFAULT_NORM), prob::ODEProblem{false, CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}, Tuple{Float64, Float64}, ...}, integrator::ODEIntegrator{false,CompositeAlgorithm{Tuple{Tsit5{Static.False,…}, Rosenbrock23{10,true,…}},…},CuArray{Float32, 1, CUDA.Mem.DeviceBuffer},Float64,…})
    @ OrdinaryDiffEq ~/.julia/packages/OrdinaryDiffEq/yspeT/src/initdt.jl:253
  [2] auto_dt_reset!(integrator::ODEIntegrator{false,CompositeAlgorithm{Tuple{Tsit5{Static.False,…}, Rosenbrock23{10,true,…}},…},CuArray{Float32, 1, CUDA.Mem.DeviceBuffer},Float64,…})
    @ OrdinaryDiffEq ~/.julia/packages/OrdinaryDiffEq/yspeT/src/integrators/integrator_interface.jl:442
  [3] handle_dt!(integrator::ODEIntegrator{false,CompositeAlgorithm{Tuple{Tsit5{Static.False,…}, Rosenbrock23{10,true,…}},…},CuArray{Float32, 1, CUDA.Mem.DeviceBuffer},Float64,…})
    @ OrdinaryDiffEq ~/.julia/packages/OrdinaryDiffEq/yspeT/src/solve.jl:547
  [4] __init(prob::ODEProblem{false, CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}, Tuple{Float64, Float64}, ...}, alg::CompositeAlgorithm{Tuple{Tsit5{Static.False,…}, Rosenbrock23{10,true,…}},…}, timeseries_init::Tuple{}, ts_init::Tuple{}, ks_init::Tuple{}, recompile::Type{Val{true}}; saveat::Tuple{}, tstops::Tuple{}, d_discontinuities::Tuple{}, save_idxs::Nothing, save_everystep::Bool, save_on::Bool, save_start::Bool, save_end::Bool, callback::Nothing, dense::Bool, calck::Bool, dt::Float64, dtmin::Nothing, dtmax::Float64, force_dtmin::Bool, adaptive::Bool, gamma::Rational{Int64}, abstol::Float64, reltol::Float64, qmin::Rational{Int64}, qmax::Int64, qsteady_min::Int64, qsteady_max::Int64, beta1::Nothing, beta2::Nothing, qoldinit::Rational{Int64}, controller::Nothing, fullnormalize::Bool, failfactor::Int64, maxiters::Int64, internalnorm::typeof(DiffEqBase.ODE_DEFAULT_NORM), internalopnorm::typeof(opnorm), isoutofdomain::typeof(DiffEqBase.ODE_DEFAULT_ISOUTOFDOMAIN), unstable_check::typeof(DiffEqBase.ODE_DEFAULT_UNSTABLE_CHECK), verbose::Bool, timeseries_errors::Bool, dense_errors::Bool, advance_to_tstop::Bool, stop_at_next_tstop::Bool, initialize_save::Bool, progress::Bool, progress_steps::Int64, progress_name::String, progress_message::typeof(DiffEqBase.ODE_DEFAULT_PROG_MESSAGE), userdata::Nothing, allow_extrapolation::Bool, initialize_integrator::Bool, alias_u0::Bool, alias_du0::Bool, initializealg::OrdinaryDiffEq.DefaultInit, kwargs::Base.Pairs{Symbol, Bool, Tuple{Symbol}, NamedTuple{(:save_noise,), Tuple{Bool}}})
    @ OrdinaryDiffEq ~/.julia/packages/OrdinaryDiffEq/yspeT/src/solve.jl:509
  [5] __solve(::ODEProblem{false, CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}, Tuple{Float64, Float64}, ...}, ::CompositeAlgorithm{Tuple{Tsit5{Static.False,…}, Rosenbrock23{10,true,…}},…}; kwargs::Base.Pairs{Symbol, Real, NTuple{6, Symbol}, NamedTuple{(:save_noise, :save_start, :save_end, :verbose, :abstol, :reltol), Tuple{Bool, Bool, Bool, Bool, Float64, Float64}}})
    @ OrdinaryDiffEq ~/.julia/packages/OrdinaryDiffEq/yspeT/src/solve.jl:5
  [6] #solve_call#22
    @ ~/.julia/packages/DiffEqBase/qPmC2/src/solve.jl:509 [inlined]
  [7] solve_up(prob::ODEProblem{false, CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}, Tuple{Float64, Float64}, ...}, sensealg::Nothing, u0::CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}, p::ComponentVector{Float64, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Axis{(layer_1 = ViewAxis(1:220, Axis(weight = ViewAxis(1:200, ShapedAxis((20, 10), NamedTuple())), bias = ViewAxis(201:220, ShapedAxis((20, 1), NamedTuple())))), layer_2 = ViewAxis(221:430, Axis(weight = ViewAxis(1:200, ShapedAxis((10, 20), NamedTuple())), bias = ViewAxis(201:210, ShapedAxis((10, 1), NamedTuple())))))}}}, args::CompositeAlgorithm{Tuple{Tsit5{Static.False,…}, Rosenbrock23{0,true,…}},…}; kwargs::Base.Pairs{Symbol, Real, NTuple{6, Symbol}, NamedTuple{(:save_noise, :save_start, :save_end, :verbose, :abstol, :reltol), Tuple{Bool, Bool, Bool, Bool, Float64, Float64}}})
    @ DiffEqBase ~/.julia/packages/DiffEqBase/qPmC2/src/solve.jl:932
  [8] solve(prob::ODEProblem{false, CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}, Tuple{Float64, Float64}, ...}, args::CompositeAlgorithm{Tuple{Tsit5{Static.False,…}, Rosenbrock23{0,true,…}},…}; sensealg::Nothing, u0::Nothing, p::Nothing, wrap::Val{true}, kwargs::Base.Pairs{Symbol, Real, NTuple{6, Symbol}, NamedTuple{(:save_noise, :save_start, :save_end, :verbose, :abstol, :reltol), Tuple{Bool, Bool, Bool, Bool, Float64, Float64}}})
    @ DiffEqBase ~/.julia/packages/DiffEqBase/qPmC2/src/solve.jl:842
  [9] _concrete_solve_adjoint(::ODEProblem{false, CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}, Tuple{Float64, Float64}, ...}, ::CompositeAlgorithm{Tuple{Tsit5{Static.False,…}, Rosenbrock23{0,true,…}},…}, ::InterpolatingAdjoint{0, false, Val{:central}, Bool}, ::CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}, ::ComponentVector{Float64, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Axis{(layer_1 = ViewAxis(1:220, Axis(weight = ViewAxis(1:200, ShapedAxis((20, 10), NamedTuple())), bias = ViewAxis(201:220, ShapedAxis((20, 1), NamedTuple())))), layer_2 = ViewAxis(221:430, Axis(weight = ViewAxis(1:200, ShapedAxis((10, 20), NamedTuple())), bias = ViewAxis(201:210, ShapedAxis((10, 1), NamedTuple())))))}}}, ::SciMLBase.ChainRulesOriginator; save_start::Bool, save_end::Bool, saveat::Float64, save_idxs::Nothing, kwargs::Base.Pairs{Symbol, Real, Tuple{Symbol, Symbol, Symbol}, NamedTuple{(:verbose, :abstol, :reltol), Tuple{Bool, Float64, Float64}}})
    @ SciMLSensitivity ~/.julia/packages/SciMLSensitivity/DSyJO/src/concrete_solve.jl:279
 [10] _concrete_solve_adjoint(::ODEProblem{false, CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}, Tuple{Float64, Float64}, ...}, ::CompositeAlgorithm{Tuple{Tsit5{Static.False,…}, Rosenbrock23{0,true,…}},…}, ::Nothing, ::CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}, ::ComponentVector{Float64, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Axis{(layer_1 = ViewAxis(1:220, Axis(weight = ViewAxis(1:200, ShapedAxis((20, 10), NamedTuple())), bias = ViewAxis(201:220, ShapedAxis((20, 1), NamedTuple())))), layer_2 = ViewAxis(221:430, Axis(weight = ViewAxis(1:200, ShapedAxis((10, 20), NamedTuple())), bias = ViewAxis(201:210, ShapedAxis((10, 1), NamedTuple())))))}}}, ::SciMLBase.ChainRulesOriginator; verbose::Bool, kwargs::Base.Pairs{Symbol, Float64, Tuple{Symbol, Symbol, Symbol}, NamedTuple{(:saveat, :abstol, :reltol), Tuple{Float64, Float64, Float64}}})
    @ SciMLSensitivity ~/.julia/packages/SciMLSensitivity/DSyJO/src/concrete_solve.jl:163
 [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", ::Nothing, ::Nothing, ::Nothing, ::Val{true}, ::Base.Pairs{Symbol, Float64, Tuple{Symbol, Symbol, Symbol}, NamedTuple{(:saveat, :abstol, :reltol), Tuple{Float64, Float64, Float64}}}, ::typeof(solve), ::ODEProblem{false, CuArray{Float32, 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{(:saveat, :abstol, :reltol), Tuple{Float64, Float64, Float64}}, ::typeof(solve), ::ODEProblem{false, CuArray{Float32, 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] _pullback
    @ ~/.julia/packages/GPUArraysCore/HaQcr/src/GPUArraysCore.jl:125 [inlined]
 [27] _pullback(::Zygote.Context{false}, ::var"#1#2"{ODEProblem{false, CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}, Tuple{Float64, Float64}, ...}})
    @ Zygote ~/.julia/packages/Zygote/oGI57/src/compiler/interface2.jl:0
 [28] #rrule_via_ad#46
    @ ~/.julia/packages/Zygote/oGI57/src/compiler/chainrules.jl:260 [inlined]
 [29] rrule_via_ad
    @ ~/.julia/packages/Zygote/oGI57/src/compiler/chainrules.jl:248 [inlined]
 [30] #1255
    @ ~/.julia/packages/ChainRules/bEtjZ/src/rulesets/Base/base.jl:258 [inlined]
 [31] task_local_storage(body::ChainRules.var"#1255#1256"{Zygote.ZygoteRuleConfig{Zygote.Context{false}}, var"#1#2"{ODEProblem{false, CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}, Tuple{Float64, Float64}, ...}}}, key::Symbol, val::GPUArraysCore.ScalarIndexing)
    @ Base ./task.jl:292
 [32] rrule(config::Zygote.ZygoteRuleConfig{Zygote.Context{false}}, #unused#::typeof(task_local_storage), body::Function, key::Symbol, value::GPUArraysCore.ScalarIndexing)
    @ ChainRules ~/.julia/packages/ChainRules/bEtjZ/src/rulesets/Base/base.jl:257
 [33] chain_rrule(::Zygote.ZygoteRuleConfig{Zygote.Context{false}}, ::Function, ::Function, ::Vararg{Any})
    @ Zygote ~/.julia/packages/Zygote/oGI57/src/compiler/chainrules.jl:223
 [34] macro expansion
    @ ~/.julia/packages/Zygote/oGI57/src/compiler/interface2.jl:0 [inlined]
 [35] _pullback(::Zygote.Context{false}, ::typeof(task_local_storage), ::var"#1#2"{ODEProblem{false, CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}, Tuple{Float64, Float64}, ...}}, ::Symbol, ::GPUArraysCore.ScalarIndexing)
    @ Zygote ~/.julia/packages/Zygote/oGI57/src/compiler/interface2.jl:9
 [36] macro expansion
    @ ~/.julia/packages/GPUArraysCore/HaQcr/src/GPUArraysCore.jl:124 [inlined]
 [37] _pullback
    @ ~/NODE_Community_Forecast/test.jl:28 [inlined]
 [38] _pullback(::Zygote.Context{false}, ::typeof(predict), ::ComponentVector{Float64, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Axis{(layer_1 = ViewAxis(1:220, Axis(weight = ViewAxis(1:200, ShapedAxis((20, 10), NamedTuple())), bias = ViewAxis(201:220, ShapedAxis((20, 1), NamedTuple())))), layer_2 = ViewAxis(221:430, Axis(weight = ViewAxis(1:200, ShapedAxis((10, 20), NamedTuple())), bias = ViewAxis(201:210, ShapedAxis((10, 1), NamedTuple())))))}}}, ::CuArray{Float32, 1, CUDA.Mem.DeviceBuffer})
    @ Zygote ~/.julia/packages/Zygote/oGI57/src/compiler/interface2.jl:0
 [39] _pullback
    @ ~/NODE_Community_Forecast/test.jl:27 [inlined]
 [40] _pullback(ctx::Zygote.Context{false}, f::typeof(predict), args::ComponentVector{Float64, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Axis{(layer_1 = ViewAxis(1:220, Axis(weight = ViewAxis(1:200, ShapedAxis((20, 10), NamedTuple())), bias = ViewAxis(201:220, ShapedAxis((20, 1), NamedTuple())))), layer_2 = ViewAxis(221:430, Axis(weight = ViewAxis(1:200, ShapedAxis((10, 20), NamedTuple())), bias = ViewAxis(201:210, ShapedAxis((10, 1), NamedTuple())))))}}})
    @ Zygote ~/.julia/packages/Zygote/oGI57/src/compiler/interface2.jl:0
 [41] _pullback
    @ ~/NODE_Community_Forecast/test.jl:34 [inlined]
 [42] _pullback(ctx::Zygote.Context{false}, f::typeof(loss_function), args::ComponentVector{Float64, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Axis{(layer_1 = ViewAxis(1:220, Axis(weight = ViewAxis(1:200, ShapedAxis((20, 10), NamedTuple())), bias = ViewAxis(201:220, ShapedAxis((20, 1), NamedTuple())))), layer_2 = ViewAxis(221:430, Axis(weight = ViewAxis(1:200, ShapedAxis((10, 20), NamedTuple())), bias = ViewAxis(201:210, ShapedAxis((10, 1), NamedTuple())))))}}})
    @ Zygote ~/.julia/packages/Zygote/oGI57/src/compiler/interface2.jl:0
 [43] _pullback
    @ ~/NODE_Community_Forecast/test.jl:40 [inlined]
 [44] _pullback(::Zygote.Context{false}, ::var"#3#4", ::ComponentVector{Float64, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Axis{(layer_1 = ViewAxis(1:220, Axis(weight = ViewAxis(1:200, ShapedAxis((20, 10), NamedTuple())), bias = ViewAxis(201:220, ShapedAxis((20, 1), NamedTuple())))), layer_2 = ViewAxis(221:430, Axis(weight = ViewAxis(1:200, ShapedAxis((10, 20), NamedTuple())), bias = ViewAxis(201:210, ShapedAxis((10, 1), NamedTuple())))))}}}, ::SciMLBase.NullParameters)
    @ Zygote ~/.julia/packages/Zygote/oGI57/src/compiler/interface2.jl:0
 [45] _apply
    @ ./boot.jl:816 [inlined]
 [46] adjoint
    @ ~/.julia/packages/Zygote/oGI57/src/lib/lib.jl:203 [inlined]
 [47] _pullback
    @ ~/.julia/packages/ZygoteRules/AIbCs/src/adjoint.jl:65 [inlined]
 [48] _pullback
    @ ~/.julia/packages/SciMLBase/5yyka/src/scimlfunctions.jl:3626 [inlined]
 [49] _pullback(::Zygote.Context{false}, ::OptimizationFunction{true, Optimization.AutoZygote, ...}, ::ComponentVector{Float64, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Axis{(layer_1 = ViewAxis(1:220, Axis(weight = ViewAxis(1:200, ShapedAxis((20, 10), NamedTuple())), bias = ViewAxis(201:220, ShapedAxis((20, 1), NamedTuple())))), layer_2 = ViewAxis(221:430, Axis(weight = ViewAxis(1:200, ShapedAxis((10, 20), NamedTuple())), bias = ViewAxis(201:210, ShapedAxis((10, 1), NamedTuple())))))}}}, ::SciMLBase.NullParameters)
    @ Zygote ~/.julia/packages/Zygote/oGI57/src/compiler/interface2.jl:0
 [50] _apply(::Function, ::Vararg{Any})
    @ Core ./boot.jl:816
 [51] adjoint
    @ ~/.julia/packages/Zygote/oGI57/src/lib/lib.jl:203 [inlined]
 [52] _pullback
    @ ~/.julia/packages/ZygoteRules/AIbCs/src/adjoint.jl:65 [inlined]
 [53] _pullback
    @ ~/.julia/packages/Optimization/GEo8L/src/function/zygote.jl:30 [inlined]
 [54] _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:220, Axis(weight = ViewAxis(1:200, ShapedAxis((20, 10), NamedTuple())), bias = ViewAxis(201:220, ShapedAxis((20, 1), NamedTuple())))), layer_2 = ViewAxis(221:430, Axis(weight = ViewAxis(1:200, ShapedAxis((10, 20), NamedTuple())), bias = ViewAxis(201:210, ShapedAxis((10, 1), NamedTuple())))))}}})
    @ Zygote ~/.julia/packages/Zygote/oGI57/src/compiler/interface2.jl:0
 [55] _apply(::Function, ::Vararg{Any})
    @ Core ./boot.jl:816
 [56] adjoint
    @ ~/.julia/packages/Zygote/oGI57/src/lib/lib.jl:203 [inlined]
 [57] _pullback
    @ ~/.julia/packages/ZygoteRules/AIbCs/src/adjoint.jl:65 [inlined]
 [58] _pullback
    @ ~/.julia/packages/Optimization/GEo8L/src/function/zygote.jl:34 [inlined]
 [59] _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:220, Axis(weight = ViewAxis(1:200, ShapedAxis((20, 10), NamedTuple())), bias = ViewAxis(201:220, ShapedAxis((20, 1), NamedTuple())))), layer_2 = ViewAxis(221:430, Axis(weight = ViewAxis(1:200, ShapedAxis((10, 20), NamedTuple())), bias = ViewAxis(201:210, ShapedAxis((10, 1), NamedTuple())))))}}})
    @ Zygote ~/.julia/packages/Zygote/oGI57/src/compiler/interface2.jl:0
 [60] pullback(f::Function, cx::Zygote.Context{false}, args::ComponentVector{Float64, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Axis{(layer_1 = ViewAxis(1:220, Axis(weight = ViewAxis(1:200, ShapedAxis((20, 10), NamedTuple())), bias = ViewAxis(201:220, ShapedAxis((20, 1), NamedTuple())))), layer_2 = ViewAxis(221:430, Axis(weight = ViewAxis(1:200, ShapedAxis((10, 20), NamedTuple())), bias = ViewAxis(201:210, ShapedAxis((10, 1), NamedTuple())))))}}})
    @ Zygote ~/.julia/packages/Zygote/oGI57/src/compiler/interface.jl:44
 [61] pullback
    @ ~/.julia/packages/Zygote/oGI57/src/compiler/interface.jl:42 [inlined]
 [62] gradient(f::Function, args::ComponentVector{Float64, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Axis{(layer_1 = ViewAxis(1:220, Axis(weight = ViewAxis(1:200, ShapedAxis((20, 10), NamedTuple())), bias = ViewAxis(201:220, ShapedAxis((20, 1), NamedTuple())))), layer_2 = ViewAxis(221:430, Axis(weight = ViewAxis(1:200, ShapedAxis((10, 20), NamedTuple())), bias = ViewAxis(201:210, ShapedAxis((10, 1), NamedTuple())))))}}})
    @ Zygote ~/.julia/packages/Zygote/oGI57/src/compiler/interface.jl:96
 [63] (::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:220, Axis(weight = ViewAxis(1:200, ShapedAxis((20, 10), NamedTuple())), bias = ViewAxis(201:220, ShapedAxis((20, 1), NamedTuple())))), layer_2 = ViewAxis(221:430, Axis(weight = ViewAxis(1:200, ShapedAxis((10, 20), NamedTuple())), bias = ViewAxis(201:210, ShapedAxis((10, 1), NamedTuple())))))}}}, ::ComponentVector{Float64, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Axis{(layer_1 = ViewAxis(1:220, Axis(weight = ViewAxis(1:200, ShapedAxis((20, 10), NamedTuple())), bias = ViewAxis(201:220, ShapedAxis((20, 1), NamedTuple())))), layer_2 = ViewAxis(221:430, Axis(weight = ViewAxis(1:200, ShapedAxis((10, 20), NamedTuple())), bias = ViewAxis(201:210, ShapedAxis((10, 1), NamedTuple())))))}}})
    @ Optimization ~/.julia/packages/Optimization/GEo8L/src/function/zygote.jl:32
 [64] macro expansion
    @ ~/.julia/packages/OptimizationOptimisers/FWIuf/src/OptimizationOptimisers.jl:31 [inlined]
 [65] macro expansion
    @ ~/.julia/packages/Optimization/GEo8L/src/utils.jl:37 [inlined]
 [66] __solve(prob::OptimizationProblem{true, ComponentVector{Float64, CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Axis{(layer_1 = ViewAxis(1:220, Axis(weight = ViewAxis(1:200, ShapedAxis((20, 10), NamedTuple())), bias = ViewAxis(201:220, ShapedAxis((20, 1), NamedTuple())))), layer_2 = ViewAxis(221:430, Axis(weight = ViewAxis(1:200, ShapedAxis((10, 20), NamedTuple())), bias = ViewAxis(201:210, ShapedAxis((10, 1), NamedTuple())))))}}}, ...}, opt::Optimisers.Adam{Float32}, data::Base.Iterators.Cycle{Tuple{Optimization.NullData}}; maxiters::Int64, callback::Function, progress::Bool, save_best::Bool, kwargs::Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}})
    @ OptimizationOptimisers ~/.julia/packages/OptimizationOptimisers/FWIuf/src/OptimizationOptimisers.jl:30
 [67] #solve#553
    @ ~/.julia/packages/SciMLBase/5yyka/src/solve.jl:86 [inlined]
in expression starting at /home/jarroyoesquivel/NODE_Community_Forecast/test.jl:43

@prbzrg
Copy link
Member

prbzrg commented Jun 9, 2023

Output of Lux.gpu is Float32, for arrays instead of

Lux.gpu(...)

use

convert(CuArray, ...)

@jarroyoe
Copy link
Author

Thank you, converting everything to CuArray ended up solving some other issues I've had when running on GPUs as well!

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 @jarroyoe @prbzrg