Apply suggestions from code review

Co-authored-by: Michael Abbott <[email protected]>
Co-authored-by: Kyle Daruwalla <[email protected]>

Author: 3 people

src/layers/basic.jl

@@ -442,30 +442,25 @@ or the corresponding [onehot encoding](@ref Flux.OneHotArray).
 
 # Examples
 
-```julia-repl
-julia> using Flux: Embedding
-
-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
 
@@ -479,5 +474,5 @@ Embedding(in::Integer, out::Integer; init = randn32) = Embedding(init(out, in))
 (m::Embedding)(x::AbstractArray) = 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