Fix constructor for GPU arrays

[vpuri3]

julia> c = ComponentArray(a=CUDA.ones(2,3), b=CUDA.ones(4,5))        

ERROR: CuArray only supports element types that are stored inline                               
Stacktrace:                                                                                     
  [1] error(s::String)                                                                          
    @ Base ./error.jl:35                                                                        
  [2] CuArray{Vector{Any}, 2, CUDA.Mem.DeviceBuffer}(#unused#::UndefInitializer, dims::Tuple{Int
64, Int64})                                                                                     
    @ CUDA ~/.julia/packages/CUDA/DfvRa/src/array.jl:34                                         
  [3] similar(a::CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}, #unused#::Type{Vector{Any}}, dims::
Tuple{Int64, Int64})                                                                            
    @ CUDA ~/.julia/packages/CUDA/DfvRa/src/array.jl:166                                        
  [4] _mapreduce(f::typeof(identity), op::ComponentArrays.var"#28#29", As::CuArray{Float32, 2, C
UDA.Mem.DeviceBuffer}; dims::Colon, init::Vector{Any})                                          
    @ GPUArrays ~/.julia/packages/GPUArrays/Hyss4/src/host/mapreduce.jl:64                      
  [5] mapreduce(::Function, ::Function, ::CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}; dims::Func
tion, init::Vector{Any})                                                                        
    @ GPUArrays ~/.julia/packages/GPUArrays/Hyss4/src/host/mapreduce.jl:31            
  [6] reduce(op::Function, A::CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}; kw::Base.Pairs{Symbol,
 Vector{Any}, Tuple{Symbol}, NamedTuple{(:init,), Tuple{Vector{Any}}}})                         
    @ Base ./reducedim.jl:406                                                                   
  [7] pushcat!(a::Vector{Any}, b::CuArray{Float32, 2, CUDA.Mem.DeviceBuffer})          
    @ ComponentArrays ~/.julia/packages/ComponentArrays/Eoni9/src/componentarray.jl:223         
  [8] make_idx(data::Vector{Any}, x::CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}, last_val::Int64
)                                               
    @ ComponentArrays ~/.julia/packages/ComponentArrays/Eoni9/src/componentarray.jl:178         
  [9] make_idx(data::Vector{Any}, nt::NamedTuple{(:a, :b), Tuple{CuArray{Float32, 2, CUDA.Mem.De
viceBuffer}, CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}}}, last_val::Int64)           
    @ ComponentArrays ~/.julia/packages/ComponentArrays/Eoni9/src/componentarray.jl:156         
 [10] make_carray_args(A::Type{Vector}, nt::NamedTuple{(:a, :b), Tuple{CuArray{Float32, 2, CUDA.
Mem.DeviceBuffer}, CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}}})
    @ ComponentArrays ~/.julia/packages/ComponentArrays/Eoni9/src/componentarray.jl:146
 [11] make_carray_args(nt::NamedTuple{(:a, :b), Tuple{CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}
, CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}}})
    @ Lux ~/.julia/dev/Lux/src/utils.jl:144
 [12] ComponentArray(nt::NamedTuple{(:a, :b), Tuple{CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}, 
CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}}})
    @ ComponentArrays ~/.julia/packages/ComponentArrays/Eoni9/src/componentarray.jl:63
 [13] ComponentArray(; kwargs::Base.Pairs{Symbol, CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}, Tu
ple{Symbol, Symbol}, NamedTuple{(:a, :b), Tuple{CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}, CuAr
ray{Float32, 2, CUDA.Mem.DeviceBuffer}}}})
    @ ComponentArrays ~/.julia/packages/ComponentArrays/Eoni9/src/componentarray.jl:67
 [14] top-level scope
    @ REPL[2]:1
 [15] top-level scope
    @ ~/.julia/packages/CUDA/DfvRa/src/initialization.jl:52

[vpuri3]

Currently im using this janky constructor:

using Lux, ComponentArrays

N = 128

# original ComponentArray
u0 = ComponentArray(; vx=rand(N,10), vy=rand(N,10)) |> gpu

# do stuff
vxnew = CUDA.ones(N,10)
vynew = CUDA.zeros(N,10)

# construct componentarray with GPU data
ca=ComponentArray(vcat(vec(vxnew), vec(vynew)), getaxes(u0))

julia> ca.vx
128×10 CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}:
 1.0  1.0  1.0  1.0  1.0  1.0  1.0  1.0  1.0  1.0
 ⋮                        ⋮                   
 1.0  1.0  1.0  1.0  1.0  1.0  1.0  1.0  1.0  1.0

julia> ca.vy
128×10 CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}:
 0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0
 ⋮                        ⋮                   
 0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0

[vpuri3]

@jonniedie do you want me to create a PR for this constructor? Alternatively, is it possible to write a constructor that doesn't involve copying data

[YichengDWu]

Every time you call vcat it create a copy so I doubt it would be faster

[avik-pal]

Alternatively, is it possible to write a constructor that doesn't involve copying data

I think it is fundamentally impossible. When allocating the different arrays, they would not be contiguous, while ComponentArray is a single contiguous vector.