Apply suggestions from code review

Co-authored-by: Michael Abbott <32575566+mcabbott@users.noreply.github.com>
Co-authored-by: Kyle Daruwalla <daruwalla.k.public@icloud.com>

Author: 3 people

src/layers/basic.jl

@@ -385,28 +385,25 @@ or the corresponding [onehot encoding](@ref Flux.OneHotArray).
 
 # Examples
 
-```julia-repl
-julia> vocab_size, embed_size = 1000, 4;
-
-julia> model = Embedding(vocab_size, embed_size)
-Embedding(1000, 4)
-
-julia> vocab_idxs = [1, 722, 53, 220, 3]
+```jldoctest
+julia> m = Embedding(reshape(-6:45, 2, 26) .+ 0.01f0)
+Embedding(26 => 2)
 
-julia> x = OneHotMatrix(vocab_idxs, vocab_size);
+julia> m(5)  # embedding vector for 5th element
+2-element Vector{Float32}:
+ 2.01
+ 3.01
 
-julia> model(x)
-4×5 Matrix{Float32}:
-  0.91139    0.670462    0.463217   0.670462    0.110932
-  0.247225  -0.0823874   0.698694  -0.0823874   0.945958
- -0.393626  -0.590136   -0.545422  -0.590136    0.77743
- -0.497621   0.87595    -0.870251   0.87595    -0.772696
-```
+julia> m([6, 15, 15])  # applied to a batch
+2×3 Matrix{Float32}:
+ 4.01  22.01  22.01
+ 5.01  23.01  23.01
 
-julia> model(vocab_idxs) == model(x)
+julia> ans == m(Flux.OneHotMatrix([6, 15, 15], 26))
 true
+```
 """
-struct Embedding{W}
+struct Embedding{W <: AbstractMatrix}
   weight::W
 end
 
@@ -417,11 +414,11 @@ function Embedding(in::Integer, out::Integer;
   return Embedding(init(out, in))
 end
 
-(m::Embedding)(x::Union{OneHotVector, OneHotMatrix}) = m.weight * x # equivalent to m.weight[:,onecold(x)]
-(m::Embedding)(x::Int) = m.weight[:, x]
-(m::Embedding)(x::AbstractVector) = NNlib.gather(m.weight, x)
-(m::Embedding)(x::AbstractArray) = reshape(m(vec(x)), :, size(x)...)
+(m::Embedding)(x::Union{OneHotLikeVector, OneHotLikeMatrix}) = m.weight * x # equivalent to m.weight[:,onecold(x)]
+(m::Embedding)(x::Integer) = m.weight[:, x]
+(m::Embedding)(x::AbstractVector{<:Integer}) = NNlib.gather(m.weight, x)
+(m::Embedding)(x::AbstractArray{<:Integer}) = reshape(m(vec(x)), :, size(x)...)
 
 function Base.show(io::IO, m::Embedding)
-  print(io, "Embedding($(size(m.weight, 2)), $(size(m.weight, 1)))")
+  print(io, "Embedding($(size(m.weight, 2)) => $(size(m.weight, 1)))")
 end