updated tests and outputsize gather

Author: manikyabard

src/outputsize.jl

@@ -169,4 +169,6 @@ for (fn, Dims) in ((:conv, DenseConvDims), (:depthwiseconv, DepthwiseConvDims))
   end
 end
 
-NNlib.gather!(dst::AbstractArray, ::AbstractArray, ::AbstractArray{<:Nil}) = fill(nil, size(dst)...)
+function NNlib.gather(src::AbstractArray{Tsrc, Nsrc}, idx::AbstractArray{<:Nil}) where {Tsrc, Nsrc}
+  fill(nil, (size(src)[1:Nsrc-1]..., size(idx)...))
+end

test/cuda/layers.jl

@@ -128,12 +128,12 @@ gpu_gradtest("PixelShuffle 1d", pixelshuffle, rand(Float32, 3, 18, 3), 3)
 
 embedding = [Flux.Embedding]
 gpu_gradtest("Embedding", embedding, [1,3,5], 5 => 2)
-gpu_gradtest("Embedding repeated indices", embedding, rand(1:50, 10^6), 50 => 2)
+gpu_gradtest("Embedding repeated indices", embedding, rand(1:10, 10^3), 10 => 2)
 gpu_gradtest("Embedding integer index", embedding, 1, 5 => 2)
 gpu_gradtest("Embedding 2d index", embedding, [1 2; 3 4], 5 => 2)
 gpu_gradtest("Embedding OneHotVec index", embedding, OneHotVector(1, 5), 5 => 2)
 gpu_gradtest("Embedding OneHotMatrix index", embedding,  OneHotMatrix([1,2,3], 5), 5 => 2)
-gpu_gradtest("Embedding OneHotMatrix repeated indices", embedding, OneHotMatrix(rand(1:50, 10^6), 50), 50 => 2)
+gpu_gradtest("Embedding OneHotMatrix repeated indices", embedding, OneHotMatrix(rand(1:10, 10^3), 10), 10 => 2)
 
 @testset "function layers" begin
   x = rand(Float32, 3,3)

test/outputsize.jl

@@ -155,3 +155,10 @@ end
   @test outputsize(m, (32, 32, 16, 16)) == (32, 32, 16, 16)
   @test outputsize(m, (32, 32, 16); padbatch=true) == (32, 32, 16, 1)
 end
+
+@testset "embedding" begin
+  m = Embedding(3=>5)
+  @test outputsize(m, (2,)) == (5, 2)
+  @test outputsize(m, (2, 3)) == (5, 2, 3)
+  @test outputsize(m, (2, 3, 4)) == (5, 2, 3, 4)
+end