Scalar indexing problem for the NeuralODE example #92
Labels
Comments
|
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 logERROR: 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:52julia> VERSION
v"1.8.0-rc3"I noticed that it was already reported here too. |
|
Can you try pinning ComponentArrays to a prior version and check. I think the new CRC rules there broke CuArrays support. |
|
With ComponentArrays v0.12.2 it works well. |
|
You can update ComponentArrays to v0.12.4 and it should be working |
|
Awesome closing this issue. Reopen if it isn't resolved. |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
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
The text was updated successfully, but these errors were encountered: