fix: gather and scatter on empty arrays (#668)

* fix: Zero case in gather

* test: Test zero-sized array

* fix: scatter! skip empty

---------

Co-authored-by: Carlo Lucibello <carlo.lucibello@gmail.com>

Author: lenianiva

ext/NNlibCUDAExt/scatter.jl

@@ -58,6 +58,7 @@ end
 function NNlib.scatter!(op::OP, dst::AnyCuArray,
         src::AnyCuArray,
         idx::AnyCuArray) where OP
+    isempty(idx) && return dst
     dims = NNlib.scatter_dims(dst, src, idx)
     args = if dims == 0
         max_idx = length(idx)

src/gather.jl

@@ -110,6 +110,7 @@ function gather!(dst::AbstractArray, src::AbstractArray, idx::AbstractArray)
 end
 
 function gather!(dst::AnyGPUArray, src::AnyGPUArray, idx::AnyGPUArray)
+    isempty(dst) && return dst
     n_dims = scatter_dims(src, dst, idx)
     dims = size(src)[1:n_dims]
     max_dims_idx = prod(dims)

test/ext_cuda/gather.jl

@@ -103,4 +103,10 @@
         outv2 = NNlib.gather(v2, i)
         @test collect(outv2) == NNlib.gather(collect(v2), collect(i))        
     end
+
+    # Zero-sized
+    x = CT([1,2,3])
+    i = CT(Int[])
+    y = NNlib.gather(x, i)
+    @test isempty(y)
 end