empty array input will throw error

FluxML · Jun 2, 2022 · 0e3df4c · 0e3df4c
1 parent b564601
commit 0e3df4c
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 1 deletion.
diff --git a/src/gather.jl b/src/gather.jl
@@ -21,6 +21,8 @@ or multiple `dst` columns.
 See [`gather`](@ref) for an allocating version.
 """
 function gather!(dst::AbstractArray, src::AbstractArray, idx::AbstractArray)
+    isempty(dst) && throw(ArgumentError("dst does not support empty array of size $(size(dst))."))
+    isempty(src) && throw(ArgumentError("src does not support empty array of size $(size(src))."))
     dims = scatter_dims(src, dst, idx)
     colons = ntuple(i -> Colon(), dims)
     for k in CartesianIndices(idx)
@@ -70,7 +72,7 @@ julia> NNlib.gather([1 2 3; 4 5 6], [1,3,1,3,1])
 function gather(src::AbstractArray{Tsrc, Nsrc}, 
                 idx::AbstractArray{Tidx, Nidx}) where 
                     {Tsrc, Nsrc, Nidx, Tidx}
-
+    isempty(src) && throw(ArgumentError("src does not support empty array of size $(size(src))."))
     M = typelength(Tidx) 
     dstsize = (size(src)[1:Nsrc-M]..., size(idx)...)
     dst = similar(src, Tsrc, dstsize)

diff --git a/test/gather.jl b/test/gather.jl
@@ -18,6 +18,8 @@ using NNlib: gather, gather!
     @test y == output
     @test gather!(T.(zero(index)), src, index) == output
     @test_throws ArgumentError gather!(zeros(T, 3, 5), src, index)
+    @test_throws ArgumentError gather!(zeros(T, 0, 5), src, index)
+    @test_throws ArgumentError gather!(T.(zero(index)), zeros(T, 0, 5), index)
 
     index2 = [1 2 3 4;
               4 2 1 3;
@@ -37,6 +39,7 @@ using NNlib: gather, gather!
     @test y isa Array{T,3}
     @test size(y) == size(index)
     @test y == output
+    @test_throws ArgumentError gather(zeros(T, 0, 5), index)
 
 
     ## 2d src, 2d index of ints -> 3d output