ComponentArray with scalars on GPU

[vpuri3]

julia> using ComponentArrays, Lux                                                               
                                                
julia> c = ComponentArray(a=ones(2,3), b=0.0) |> gpu                                    
ComponentVector{Float32, CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}, Tuple{Axis{(a = ViewAxis(1:
6, ShapedAxis((2, 3), NamedTuple())), b = 7)}}}(a = Float32[1.0 1.0 1.0; 1.0 1.0 1.0], b = Error
 showing value of type ComponentVector{Float32, CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}, Tupl
e{Axis{(a = ViewAxis(1:6, ShapedAxis((2, 3), NamedTuple())), b = 7)}}}:                         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/ZBmfM/src/GPUArraysCore.jl:87
  [3] getindex                                                                                  
    @ ~/.julia/packages/GPUArrays/Hyss4/src/host/indexing.jl:9 [inlined]    
  [4] macro expansion                                                                           
    @ ~/.julia/packages/ComponentArrays/Eoni9/src/array_interface.jl:0 [inlined]
  [5] _getindex(index_fun::typeof(getindex), x::ComponentVector{Float32, CuArray{Float32, 1, CUD
A.Mem.DeviceBuffer}, Tuple{Axis{(a = ViewAxis(1:6, ShapedAxis((2, 3), NamedTuple())), b = 7)}}}, idx::Val{:b})                                  

julia> c.a
2×3 CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}:
 1.0  1.0  1.0
 1.0  1.0  1.0

julia> c.b
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/ZBmfM/src/GPUArraysCore.jl:87
 [3] getindex
   @ ~/.julia/packages/GPUArrays/Hyss4/src/host/indexing.jl:9 [inlined]
 [4] maybeview
   @ ./views.jl:147 [inlined]
 [5] macro expansion
   @ ~/.julia/packages/ComponentArrays/Eoni9/src/array_interface.jl:0 [inlined]
 [6] _getindex(index_fun::typeof(Base.maybeview), x::ComponentVector{Float32, CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}, Tuple{Axis{(a = ViewAxis(1:6, ShapedAxis((2, 3), NamedTuple())), b = 7)}}}, idx::Val{:b})
   @ ComponentArrays ~/.julia/packages/ComponentArrays/Eoni9/src/array_interface.jl:121
 [7] getproperty(x::ComponentVector{Float32, CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}, Tuple{Axis{(a = ViewAxis(1:6, ShapedAxis((2, 3), NamedTuple())), b = 7)}}}, s::Symbol)
   @ ComponentArrays ~/.julia/packages/ComponentArrays/Eoni9/src/namedtuple_interface.jl:14
 [8] top-level scope
   @ REPL[17]:1
 [9] top-level scope
   @ ~/.julia/packages/CUDA/DfvRa/src/initialization.jl:52

[jonniedie]

Yeah, that’s an interesting one. If scalar indexing is not allowed, I’m not sure there’s any use in having a component of the ComponentArray be a scalar, since you’ll never be able to get it. I don’t think there’s anything we can do there, it’s just something that GPUArrays just fundamentally doesn’t want you to do. You could, however just do this:

julia> c = ComponentArray(a=ones(2,3), b=[0.0]) |> gpu

julia> c.b[1]
 0.0

[vpuri3]

^yes that seems to be the only way forward. thanks