updated Embedding tests

Author: manikyabard

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, [1,3,5,3], 5 => 2)
+gpu_gradtest("Embedding repeated indices", embedding, rand(1:50, 10^6), 50 => 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([1,2,2], 5), 5 => 2)
+gpu_gradtest("Embedding OneHotMatrix repeated indices", embedding, OneHotMatrix(rand(1:50, 10^6), 50), 50 => 2)
 
 @testset "function layers" begin
   x = rand(Float32, 3,3)

test/layers/basic.jl

@@ -284,9 +284,9 @@ import Flux: activations
     @test y isa Matrix{Float32}
     @test y ≈ m.weight[:,x]
     x2 = OneHotMatrix(x, vocab_size)
-    y2 = m(x2)
-    @test y2 isa Matrix{Float32}
-    @test y2 ≈ y
+    @test m(x2) isa Matrix{Float32}
+    @test m(x2) ≈ y 
+    @test m(collect(x2)) ≈ y
     @test_throws DimensionMismatch m(OneHotMatrix(x, 1000))
 
     x = rand(1:vocab_size, 3, 4)
@@ -297,6 +297,9 @@ import Flux: activations
     @test m(2) ≈ m.weight[:,2]
     @test m(OneHotVector(3, vocab_size)) ≈ m.weight[:,3]
     @test_throws DimensionMismatch m(OneHotVector(3, 1000))
+
+    x = onehotbatch(rand(1:vocab_size, 4, 3, 4, 5), 1:vocab_size)
+    @test m(x) ≈ m(onecold(x))
   end
 end