Specialize getindex for indexing VectorOfArray with boolean mask and colon

Author: ma-sadeghi

src/vector_of_array.jl

@@ -73,6 +73,14 @@ Base.@propagate_inbounds function Base.getindex(A::AbstractVectorOfArray{T, N},
   return Adapt.adapt(__parameterless_type(T),reshape(reduce(hcat,vecs),size(A.u[1])...,length(A.u)))
 end
 
+Base.@propagate_inbounds function Base.getindex(A::AbstractVectorOfArray{T, N},
+  I::AbstractArray{Bool},J::Colon...) where {T, N}
+  @assert length(J) == ndims(A.u[1])+1-ndims(I)
+  @assert size(I) == size(A)[1:ndims(A)-length(J)]
+  vecs = vec.(A.u)
+  return Base.getindex(Adapt.adapt(__parameterless_type(T),reshape(reduce(hcat,vecs),size(A.u[1])...,length(A.u))), I, J...)
+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/vectorofarray_gpu.jl

@@ -1,10 +1,26 @@
-using RecursiveArrayTools, CUDA
+using RecursiveArrayTools, CUDA, Test
 CUDA.allowscalar(false)
 
+# Test indexing with colon
 x = zeros(5)
 y = VectorOfArray([x,x,x])
 y[:,:]
 
 x = CUDA.zeros(5)
 y = VectorOfArray([x,x,x])
 y[:,:]
+
+# Test indexing with boolean masks and colon
+nx, ny, nt = 3, 4, 5
+x = CUDA.rand(nx, ny, nt)
+m = CUDA.rand(nx, ny) .> 0.5
+x[m, :]
+
+va = VectorOfArray([slice for slice in eachslice(x, dims=3)])
+@test va[m, :] ≈ x[m, :]
+va[m, :]
+
+xc = Array(x)
+mc = Array(m)
+xc[mc, :]
+