Scalar indexing problem for the NeuralODE example

[gabrevaya]

Hi, firstly, thank you very much for this great package with super complete and didactical documentation! :)

While going through the documentation I realized that the NeuralODE example is not working properly on GPU. It throws the scalar indexing error and I think it is because of having the parameters as a ComponentArray, but I don't know how to fix it.

Error log

```julia 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/rSIl2/src/GPUArraysCore.jl:78 [3] getindex(xs::CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}, I::Int64) @ GPUArrays ~/.julia/packages/GPUArrays/gok9K/src/host/indexing.jl:9 [4] setindex! @ ./array.jl:979 [inlined] [5] macro expansion @ ~/.julia/packages/ComponentArrays/NEqmD/src/array_interface.jl:0 [inlined] [6] _setindex!(x::ComponentVector{Float32}, v::CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}, idx::Val{:bias}) @ ComponentArrays ~/.julia/packages/ComponentArrays/NEqmD/src/array_interface.jl:129 [7] setproperty! @ ~/.julia/packages/ComponentArrays/NEqmD/src/namedtuple_interface.jl:17 [inlined] [8] (::ComponentArrays.var"#getproperty_adjoint#88"{ComponentVector{Float32, CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}, Tuple{Axis{(weight = ViewAxis(1:200, ShapedAxis((10, 20), NamedTuple())), bias = ViewAxis(201:210, ShapedAxis((10, 1), NamedTuple())))}}}, Symbol})(Δ::CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}) @ ComponentArrays ~/.julia/packages/ComponentArrays/NEqmD/src/compat/chainrulescore.jl:4 [9] ZBack @ ~/.julia/packages/Zygote/IoW2g/src/compiler/chainrules.jl:205 [inlined] [10] Pullback @ ~/.julia/packages/Lux/lEqCI/src/layers/basic.jl:639 [inlined] [11] macro expansion @ ~/.julia/packages/Lux/lEqCI/src/layers/basic.jl:0 [inlined] [12] Pullback @ ~/.julia/packages/Lux/lEqCI/src/layers/basic.jl:507 [inlined] [13] (::typeof(∂(applychain)))(Δ::Tuple{CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}, Nothing}) @ Zygote ~/.julia/packages/Zygote/IoW2g/src/compiler/interface2.jl:0 [14] Pullback @ ~/.julia/packages/Lux/lEqCI/src/layers/basic.jl:504 [inlined] [15] (::typeof(∂(λ)))(Δ::Tuple{CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}, Nothing}) @ Zygote ~/.julia/packages/Zygote/IoW2g/src/compiler/interface2.jl:0 [16] Pullback @ /network/scratch/a/abrevayg/.julia/packages/Lux/SApdg/examples/NeuralODE/main.jl:103 [inlined] [17] Pullback @ /network/scratch/a/abrevayg/.julia/packages/Lux/SApdg/examples/NeuralODE/main.jl:134 [inlined] [18] (::typeof(∂(λ)))(Δ::Tuple{Float32, Nothing}) @ Zygote ~/.julia/packages/Zygote/IoW2g/src/compiler/interface2.jl:0 [19] (::Zygote.var"#60#61"{typeof(∂(λ))})(Δ::Tuple{Float32, Nothing}) @ Zygote ~/.julia/packages/Zygote/IoW2g/src/compiler/interface.jl:41 [20] train() @ Main /network/scratch/a/abrevayg/.julia/packages/Lux/SApdg/examples/NeuralODE/main.jl:135 [21] top-level scope @ /network/scratch/a/abrevayg/.julia/packages/Lux/SApdg/examples/NeuralODE/main.jl:155 [22] include(fname::String) @ Base.MainInclude ./client.jl:476 [23] top-level scope @ REPL[6]:1 [24] top-level scope @ ~/.julia/packages/CUDA/DfvRa/src/initialization.jl:52 in expression starting at /network/scratch/a/abrevayg/.julia/packages/Lux/SApdg/examples/NeuralODE/main.jl:155 ``` ```julia (examples) pkg> st Status `~/.julia/packages/Lux/lEqCI/examples/Project.toml` [c29ec348] AbstractDifferentiation v0.4.3 [c7e460c6] ArgParse v1.1.4 [02898b10] Augmentor v0.6.6 [052768ef] CUDA v3.12.0 ⌅ [b0b7db55] ComponentArrays v0.11.17 [2e981812] DataLoaders v0.1.3 [41bf760c] DiffEqSensitivity v6.79.0 [587475ba] Flux v0.13.4 ⌅ [acf642fa] FluxMPI v0.5.3 [59287772] Formatting v0.4.2 [f6369f11] ForwardDiff v0.10.30 ⌅ [d9f16b24] Functors v0.2.8 [6218d12a] ImageMagick v1.2.2 ⌃ [916415d5] Images v0.24.1 [b835a17e] JpegTurbo v0.1.1 [b2108857] Lux v0.4.9 [cc2ba9b6] MLDataUtils v0.5.4 [eb30cadb] MLDatasets v0.7.4 [f1d291b0] MLUtils v0.2.9 [dbeba491] Metalhead v0.7.3 [872c559c] NNlib v0.8.8 [3bd65402] Optimisers v0.2.7 [1dea7af3] OrdinaryDiffEq v6.18.2 [d7d3b36b] ParameterSchedulers v0.3.3 [91a5bcdd] Plots v1.31.3 [37e2e3b7] ReverseDiff v1.14.1 ⌅ [efcf1570] Setfield v0.8.2 [fce5fe82] Turing v0.21.9 [e88e6eb3] Zygote v0.6.41 [de0858da] Printf [9a3f8284] Random [10745b16] Statistics Info Packages marked with ⌃ and ⌅ have new versions available, but those with ⌅ cannot be upgraded. To see why use `status --outdated` ``` ```julia julia> VERSION v"1.8.0-rc3" ```

[gabrevaya]

Here is a MWE:

using Lux, Random, NNlib, Zygote, CUDA, ComponentArrays
CUDA.allowscalar(false)

model = Chain(Dense(2 => 4))
rng = Random.default_rng()
x = randn(rng, 2, 4) |> gpu

ps, st = Lux.setup(rng, model)
ps = ps |> ComponentArray |> gpu
st = st |> gpu

model(x, ps, st)
l, back = pullback(ps -> sum(first(model(x, ps, st))), ps)
grad = back(one(l))

Error log

```julia ERROR: 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/rSIl2/src/GPUArraysCore.jl:78 [3] getindex(xs::CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}, I::Int64) @ GPUArrays ~/.julia/packages/GPUArrays/gok9K/src/host/indexing.jl:9 [4] setindex! @ ./array.jl:979 [inlined] [5] macro expansion @ ~/.julia/packages/ComponentArrays/NEqmD/src/array_interface.jl:0 [inlined] [6] _setindex!(x::ComponentVector{Float32}, v::CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}, idx::Val{:bias}) @ ComponentArrays ~/.julia/packages/ComponentArrays/NEqmD/src/array_interface.jl:129 [7] setproperty! @ ~/.julia/packages/ComponentArrays/NEqmD/src/namedtuple_interface.jl:17 [inlined] [8] (::ComponentArrays.var"#getproperty_adjoint#88"{ComponentVector{Float32, CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}, Tuple{Axis{(weight = ViewAxis(1:8, ShapedAxis((4, 2), NamedTuple())), bias = ViewAxis(9:12, ShapedAxis((4, 1), NamedTuple())))}}}, Symbol})(Δ::CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}) @ ComponentArrays ~/.julia/packages/ComponentArrays/NEqmD/src/compat/chainrulescore.jl:4 [9] ZBack @ ~/.julia/packages/Zygote/IoW2g/src/compiler/chainrules.jl:205 [inlined] [10] Pullback @ ~/.julia/packages/Lux/lEqCI/src/layers/basic.jl:639 [inlined] [11] Pullback @ ./REPL[15]:1 [inlined] [12] (::typeof(∂(#5)))(Δ::Float32) @ Zygote ~/.julia/packages/Zygote/IoW2g/src/compiler/interface2.jl:0 [13] (::Zygote.var"#60#61"{typeof(∂(#5))})(Δ::Float32) @ Zygote ~/.julia/packages/Zygote/IoW2g/src/compiler/interface.jl:41 [14] top-level scope @ REPL[16]:1 [15] top-level scope @ ~/.julia/packages/CUDA/DfvRa/src/initialization.jl:52 ``` ```julia julia> VERSION v"1.8.0-rc3" ```

I noticed that it was already reported here too.

[avik-pal]

Can you try pinning ComponentArrays to a prior version and check. I think the new CRC rules there broke CuArrays support.

[gabrevaya]

With ComponentArrays v0.12.2 it works well.

[YichengDWu]

You can update ComponentArrays to v0.12.4 and it should be working

[avik-pal]

Awesome closing this issue. Reopen if it isn't resolved.