more embedding tests; keep Embedding unexported

Author: CarloLucibello

Project.toml

@@ -1,12 +1,14 @@
 name = "Flux"
 uuid = "587475ba-b771-5e3f-ad9e-33799f191a9c"
-version = "0.12.4"
+version = "0.12.5"
 
 [deps]
 AbstractTrees = "1520ce14-60c1-5f80-bbc7-55ef81b5835c"
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
+ChainRules = "082447d4-558c-5d27-93f4-14fc19e9eca2"
+ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 CodecZlib = "944b1d66-785c-5afd-91f1-9de20f533193"
 Colors = "5ae59095-9a9b-59fe-a467-6f913c188581"
 DelimitedFiles = "8bb1440f-4735-579b-a4ab-409b98df4dab"
@@ -37,7 +39,7 @@ Colors = "0.12"
 Functors = "0.2.1"
 Juno = "0.8"
 MacroTools = "0.5"
-NNlib = "0.7.14"
+NNlib = "0.7.24"
 NNlibCUDA = "0.1"
 Reexport = "0.2, 1.0"
 StatsBase = "0.33"

src/Flux.jl

@@ -11,7 +11,7 @@ using Zygote: Params, @adjoint, gradient, pullback, @nograd
 export gradient
 
 export Chain, Dense, Maxout, SkipConnection, Parallel, flatten,
-       RNN, LSTM, GRU, Embedding,
+       RNN, LSTM, GRU,
        SamePad, Conv, CrossCor, ConvTranspose, DepthwiseConv,
        AdaptiveMaxPool, AdaptiveMeanPool, GlobalMaxPool, GlobalMeanPool, MaxPool, MeanPool,
        Dropout, AlphaDropout, LayerNorm, BatchNorm, InstanceNorm, GroupNorm,

test/cuda/layers.jl

@@ -261,16 +261,37 @@ end
 end
 
 @testset "Embedding" begin
-  vocab_size, embed_size = 10, 4
+  vocab_size, embed_size = 5, 2
   m = Embedding(vocab_size, embed_size)
-  x = rand(1:vocab_size, 3)
+
+  x = [1, 3, 5]
   y = m(x)
   m_g = m |> gpu
   x_g = x |> gpu
   y_g = m_g(x_g)
   @test collect(y_g) == y
+  
+  gs = gradient(() -> sum(m(x)), params(m))
+  gs_g = gradient(() -> sum(m_g(x_g)), params(m_g))
+  @test collect(gs_g[m_g.weight]) ≈ gs[m.weight]
+
   gs = gradient(() -> sum(tanh.(m(x))), params(m))
   gs_g = gradient(() -> sum(tanh.(m_g(x_g))), params(m_g))
   @test collect(gs_g[m_g.weight]) ≈ gs[m.weight]
+
+  @testset "repeated indexes" begin
+    vocab_size, embed_size = 5, 2
+    m = Embedding(vocab_size, embed_size)
+
+    x = [1, 3, 5, 3] # repeated indexes
+    y = m(x)
+    m_g = m |> gpu
+    x_g = x |> gpu
+    y_g = m_g(x_g)
+    @test collect(y_g) == y
+    gs = gradient(() -> sum(m(x)), params(m))
+    gs_g = gradient(() -> sum(m_g(x_g)), params(m_g))
+    @test collect(gs_g[m_g.weight]) ≈ gs[m.weight]
+  end
 end