update Embedding constructor

Updated Embedding constructor to use `=>` and added OneHotLikeVector and OneHotLikeMatrix consts.

Author: manikyabard

src/layers/basic.jl

@@ -374,7 +374,7 @@ function Base.show(io::IO, m::Parallel)
 end
 
 """
-    Embedding(in, out; init=randn)
+    Embedding(in => out; init=randn)
 
 A lookup table that stores embeddings of dimension `out` 
 for a vocabulary of size `in`. 
@@ -409,9 +409,9 @@ end
 
 @functor Embedding
 
-function Embedding(in::Integer, out::Integer;
+function Embedding(dims::Pair{<:Integer, <:Integer};
                init = (i...) -> randn(Float32, i...))
-  return Embedding(init(out, in))
+  return Embedding(init(last(dims), first(dims)))
 end
 
 (m::Embedding)(x::Union{OneHotLikeVector, OneHotLikeMatrix}) = m.weight * x # equivalent to m.weight[:,onecold(x)]

src/onehot.jl

@@ -24,6 +24,9 @@ const OneHotLike{T, L, N, var"N+1", I} =
   Union{OneHotArray{T, L, N, var"N+1", I},
         Base.ReshapedArray{Bool, var"N+1", <:OneHotArray{T, L, <:Any, <:Any, I}}}
 
+const OneHotLikeVector{T, L} = OneHotLike{T, L, 0, 1, T}
+const OneHotLikeMatrix{T, L, I} = OneHotLike{T, L, 1, 2, I}  
+
 _isonehot(x::OneHotArray) = true
 _isonehot(x::Base.ReshapedArray{<:Any, <:Any, <:OneHotArray{<:Any, L}}) where L = (size(x, 1) == L)
 

test/cuda/layers.jl

@@ -221,7 +221,7 @@ end
 
 @testset "Embedding" begin
   vocab_size, embed_size = 10, 4
-  m = Embedding(vocab_size, embed_size)
+  m = Embedding(vocab_size => embed_size)
   x = rand(1:vocab_size, 3)
   y = m(x)
   m_g = m |> gpu

test/layers/basic.jl

@@ -156,7 +156,7 @@ import Flux: activations
 
   @testset "Embedding" begin
     vocab_size, embed_size = 10, 4
-    m = Embedding(vocab_size, embed_size)
+    m = Embedding(vocab_size => embed_size)
     @test size(m.weight) == (embed_size, vocab_size)
 
     x = rand(1:vocab_size, 3)