cannot broadcast on gpu(x)

[YichengDWu]

julia> using Lux

julia> ps = Lux.ComponentArray(a=[2,42],b=[3,4,5])
ComponentVector{Int64}(a = [2, 42], b = [3, 4, 5])

julia> ps .* 0
ComponentVector{Int64}(a = [0, 0], b = [0, 0, 0])

julia> gpu(ps) .* 0
ERROR: This object is not a GPU array
Stacktrace:
 [1] error(s::String)
   @ Base .\error.jl:33
 [2] backend(#unused#::Type)
   @ GPUArrays C:\Users\Luffy\.julia\packages\GPUArrays\gok9K\src\device\execution.jl:15
 [3] backend(x::ComponentVector{Int64, CuArray{Int64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Axis{(a = 1:2, b = 3:5)}}})
   @ GPUArrays C:\Users\Luffy\.julia\packages\GPUArrays\gok9K\src\device\execution.jl:16
 [4] _copyto!
   @ C:\Users\Luffy\.julia\packages\GPUArrays\gok9K\src\host\broadcast.jl:73 [inlined]
 [5] copyto!
   @ C:\Users\Luffy\.julia\packages\GPUArrays\gok9K\src\host\broadcast.jl:56 [inlined]
 [6] copy
   @ C:\Users\Luffy\.julia\packages\GPUArrays\gok9K\src\host\broadcast.jl:47 [inlined]
 [7] materialize(bc::Base.Broadcast.Broadcasted{CUDA.CuArrayStyle{1}, Nothing, typeof(*), Tuple{ComponentVector{Int64, CuArray{Int64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Axis{(a = 1:2, b = 3:5)}}}, Int64}})
   @ Base.Broadcast .\broadcast.jl:860
 [8] top-level scope
   @ REPL[72]:1
 [9] top-level scope
   @ C:\Users\Luffy\.julia\packages\CUDA\DfvRa\src\initialization.jl:52

[YichengDWu]

A similar but different error

julia> map(identity, gpu(ps))
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:33
  [2] assertscalar(op::String)
    @ GPUArraysCore C:\Users\Luffy\.julia\packages\GPUArraysCore\rSIl2\src\GPUArraysCore.jl:78
  [3] getindex
    @ C:\Users\Luffy\.julia\packages\GPUArrays\gok9K\src\host\indexing.jl:9 [inlined]
  [4] getindex
    @ C:\Users\Luffy\.julia\packages\ComponentArrays\IQNQW\src\array_interface.jl:96 [inlined]
  [5] iterate
    @ .\abstractarray.jl:1144 [inlined]
  [6] iterate
    @ .\abstractarray.jl:1142 [inlined]
  [7] iterate
    @ .\generator.jl:44 [inlined]
  [8] _collect(c::ComponentVector{Int64, CuArray{Int64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Axis{(a = 1:2, b = 3:5)}}}, itr::Base.Generator{ComponentVector{Int64, CuArray{Int64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Axis{(a = 1:2, b = 3:5)}}}, typeof(identity)}, #unused#::Base.EltypeUnknown, isz::Base.HasShape{1})
    @ Base .\array.jl:744
  [9] collect_similar(cont::ComponentVector{Int64, CuArray{Int64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Axis{(a = 1:2, b = 3:5)}}}, itr::Base.Generator{ComponentVector{Int64, CuArray{Int64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Axis{(a = 1:2, b = 3:5)}}}, typeof(identity)})
    @ Base .\array.jl:653
 [10] map(f::Function, A::ComponentVector{Int64, CuArray{Int64, 1, CUDA.Mem.DeviceBuffer}, Tuple{Axis{(a = 1:2, b = 3:5)}}})
    @ Base .\abstractarray.jl:2849
 [11] top-level scope
    @ REPL[160]:1
 [12] top-level scope
    @ C:\Users\Luffy\.julia\packages\CUDA\DfvRa\src\initialization.jl:52

[jonniedie]

Interesting. ComponentArrays take the broadcast style of their inner arrays rather than having their own broadcast style now. The axes of the ComponentArray that get passed in to broadcast copy allow it to recreate a ComponentArray with the correct inner array type. I used to have it the other way where they had their own broadcast style and would explicitly try to recreate the inner array. This was basically impossible to do in a generic way because the interface for similar in different array packages is inconsistently followed.

I’ll see what I can do on my end, but if you have any ideas, I’d be open to a PR.

[YichengDWu]

Closing this since it's a Lux's issue