support empty source array for gather

Author: yuehhua

src/gather.jl

@@ -66,6 +66,7 @@ function NNlib.gather!(dst::AnyCuArray, src::AnyCuArray, idx::AnyCuArray)
 
     # check bounds
     chk = checkbounds_src(src, dims, eltype(idx))
+    isempty(src) && return dst
     in_bnd = map(chk, collect(idx))
     if !all(in_bnd)
         j = findfirst(i -> !i, in_bnd)

test/gather.jl

@@ -91,4 +91,14 @@
     @test y isa CuArray{Float32,3}
     @test size(y) == (size(src)[1:Nsrc-M]..., size(index)...) 
     gputest(src -> NNlib.gather(src, index), src, checkgrad=true)
+
+    ## empty 2d src, 2d index of ints -> 3d output
+    src = CT(zeros(Int, 0, 3))
+    index = cu([1 2 3;
+                2 2 1;
+                3 1 3])
+    
+    y = NNlib.gather(src, index)
+    @test y isa CuArray{Float32,2}
+    @test size(y) == (0, 3, 3)
 end