Fix all colon dispatch for GPUs

.buildkite/pipeline.yml

@@ -0,0 +1,17 @@
+steps:
+  - label: "GPU"
+    plugins:
+      - JuliaCI/julia#v1:
+          version: "1"
+      - JuliaCI/julia-test#v1:
+           coverage: false # 1000x slowdown
+    agents:
+      queue: "juliagpu"
+      cuda: "*"
+    env:
+      GROUP: 'GPU'
+      JULIA_PKG_SERVER: "" # it often struggles with our large artifacts
+      # SECRET_CODECOV_TOKEN: "..."
+    timeout_in_minutes: 30
+    # Don't run Buildkite if the commit message includes the text [skip tests]
+    if: build.message !~ /\[skip tests\]/

.github/workflows/CI.yml

@@ -13,6 +13,7 @@ jobs:
       matrix:
         group:
           - Core
+          - Downstream
         version:
           - '1'
           - '1.6'

Project.toml

@@ -4,6 +4,7 @@ authors = ["Chris Rackauckas <[email protected]>"]
 version = "2.21.1"
 
 [deps]
+Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"

src/RecursiveArrayTools.jl

@@ -10,7 +10,7 @@ using Requires, RecipesBase, StaticArrays, Statistics,
 
 import ChainRulesCore
 import ChainRulesCore: NoTangent
-import ZygoteRules
+import ZygoteRules, Adapt
 
 using FillArrays
 

src/vector_of_array.jl

@@ -65,6 +65,14 @@ Base.@propagate_inbounds function Base.getindex(A::AbstractDiffEqArray{T, N},
   I::Union{Int,AbstractArray{Int},CartesianIndex,Colon,BitArray,AbstractArray{Bool}}...) where {T, N}
   RecursiveArrayTools.VectorOfArray(A.u)[I...]
 end
+
+Base.@propagate_inbounds function Base.getindex(A::AbstractVectorOfArray{T, N},
+  I::Colon...) where {T, N}
+  @assert length(I) == ndims(A.u[1])+1
+  vecs = vec.(A.u)
+  return Adapt.adapt(__parameterless_type(T),reshape(reduce(hcat,vecs),size(A.u[1])...,length(A.u)))
+end
+
 Base.@propagate_inbounds Base.getindex(A::AbstractDiffEqArray{T, N}, i::Int,::Colon) where {T, N} = [A.u[j][i] for j in 1:length(A)]
 Base.@propagate_inbounds Base.getindex(A::AbstractDiffEqArray{T, N}, ::Colon,i::Int) where {T, N} = A.u[i]
 Base.@propagate_inbounds Base.getindex(A::AbstractDiffEqArray{T, N}, i::Int,II::AbstractArray{Int}) where {T, N} = [A.u[j][i] for j in II]

test/gpu/Project.toml

@@ -0,0 +1,2 @@
+[deps]
+CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"

test/gpu/vectorofarray_gpu.jl

@@ -0,0 +1,10 @@
+using RecursiveArrayTools, CUDA
+CUDA.allowscalar(false)
+
+x = zeros(5)
+y = VectorOfArray([x,x,x])
+y[:,:]
+
+x = CUDA.zeros(5)
+y = VectorOfArray([x,x,x])
+y[:,:]

test/runtests.jl

@@ -11,7 +11,15 @@ function activate_downstream_env()
   Pkg.instantiate()
 end
 
+function activate_gpu_env()
+  Pkg.activate("gpu")
+  Pkg.develop(PackageSpec(path=dirname(@__DIR__)))
+  Pkg.instantiate()
+end
+
 @time begin
+
+if !is_APPVEYOR && GROUP == "Core"
   @time @testset "Utils Tests" begin include("utils_test.jl") end
   @time @testset "Partitions Tests" begin include("partitions_test.jl") end
   @time @testset "VecOfArr Indexing Tests" begin include("basic_indexing.jl") end
@@ -20,9 +28,15 @@ end
   @time @testset "Linear Algebra Tests" begin include("linalg.jl") end
   @time @testset "Upstream Tests" begin include("upstream.jl") end
   @time @testset "Adjoint Tests" begin include("adjoints.jl") end
+end
+
+if !is_APPVEYOR && GROUP == "Downstream"
+  activate_downstream_env()
+  @time @testset "DiffEqArray Indexing Tests" begin include("downstream/symbol_indexing.jl") end
+end
 
-  if !is_APPVEYOR && GROUP == "Downstream"
-    activate_downstream_env()
-    @time @testset "DiffEqArray Indexing Tests" begin include("downstream/symbol_indexing.jl") end
-  end
+if !is_APPVEYOR && GROUP == "GPU"
+  activate_gpu_env()
+  @time @testset "VectorOfArray GPU" begin include("gpu/vectorofarray_gpu.jl") end
+end
 end