make outputsize work with Embedding

Author: mcabbott

src/layers/basic.jl

@@ -655,13 +655,14 @@ or the corresponding [`onehot encoding`](@ref OneHotArrays.onehotbatch).
 
 For indices `x`, the result is of size `(out, size(x)...)`, allowing several batch dimensions.
 For one-hot `ohx`, the result is of size `(out, size(ohx)[2:end]...)`.
+Note that [`outputsize`](@ref Flux.outputsize) expects `size(x)`, the indices not the one-hot array.
 
 # Examples
 ```jldoctest
 julia> emb = Embedding(26 => 4, init=Flux.identity_init(gain=22))
 Embedding(26 => 4)  # 104 parameters
 
-julia> emb(2)  # one column of e.weight (here not random!)
+julia> emb(2)  # one column of emb.weight (here not random!)
 4-element Vector{Float32}:
   0.0
  22.0
@@ -680,6 +681,9 @@ true
 
 julia> emb(rand(1:26, (10, 1, 12))) |> size  # three batch dimensions
 (4, 10, 1, 12)
+
+julia> Flux.outputsize(emb, (10, 1, 12))  # outputsize wants indices, not OneHotArray
+(4, 10, 1, 12)
 ```
 """
 struct Embedding{W<:AbstractMatrix}
@@ -691,8 +695,11 @@ end
 Embedding((in, out)::Pair{<:Integer, <:Integer}; init = randn32) = Embedding(init(out, in))
 
 (m::Embedding)(x::Integer) = 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::AbstractVector{<:Integer}) = NNlib.gather(m.weight, x)
+(m::Embedding)(x::AbstractArray{<:Integer}) = reshape(m(vec(x)), :, size(x)...)
+
+(m::Embedding)(x::Nil) = similar(m.weight, Nil, size(m.weight, 1))
+(m::Embedding)(x::AbstractArray{Nil}) = similar(m.weight, Nil, size(m.weight, 1), size(x)...)
 
 (m::Embedding)(x::AbstractVector{Bool}) = m.weight * x  # usually OneHotVector
 (m::Embedding)(x::AbstractMatrix{Bool}) = m.weight * x  # usually OneHotMatrix

src/outputsize.jl

@@ -87,6 +87,12 @@ DimensionMismatch("Input channels must match! (7 vs. 3)")
 julia> outputsize([Dense(10 => 4), Dense(4 => 2)], (10, 1)) # Vector of layers becomes a Chain
 (2, 1)
 ```
+
+Limitations:
+* `Embedding` accepts either integers or one-hot arrays, and `ohx = onehotbatch(x, ...)`
+  has one more dimension than `x`. Here `outputsize` uses `size(x)`.
+* At present `outputsize` does not work with recurrent layers,
+  `outputsize(RNN(2 => 3), (2, 1))` gives an error. This is a bug.
 """
 function outputsize(m, inputsizes::Tuple...; padbatch=false)
   x = nil_input(padbatch, inputsizes...)

test/outputsize.jl

@@ -5,6 +5,8 @@
   m = Dense(10, 5)
   @test_throws DimensionMismatch outputsize(m, (5, 2)) == (5, 1)
   @test outputsize(m, (10,); padbatch=true) == (5, 1)
+  @test outputsize(m, (10,)) == (5,)
+  @test outputsize(m, (10, 6, 7)) == (5, 6, 7)
 
   m = Chain(Dense(10, 8, σ), Dense(8, 5), Dense(5, 2))
   @test outputsize(m, (10,); padbatch=true) == (2, 1)
@@ -41,6 +43,19 @@
   @test outputsize(m, (10, 10, 3, 1)) == (10, 10, 19, 1)
 end
 
+@testset "embeddings" begin
+  # Here outputsize expects indices, not one-hot representation:
+  m = Embedding(3 => 4)
+  @test outputsize(m, (3, 7)) == (4, 3, 7) == size(m(rand(1:3, 3, 7)))
+  @test outputsize(m, (5, 6, 7)) == (4, 5, 6, 7) == size(m(rand(1:3, 5, 6, 7)))
+
+  m = Chain(x -> Flux.onehotbatch(x, 1:5), Embedding(5 => 7))
+  @test size(m([3,4])) == (7, 2)
+  @test outputsize(m, (2,)) == (7, 2)
+  # This works because Flux.onehotbatch([nil, nil], 1:5) makes a 5×2 OneHotMatrix
+  # But e.g. Flux.onehotbatch([nil, nil], 'a':'e') will not work.
+end
+
 @testset "multiple inputs" begin
   m = Parallel(vcat, Dense(2, 4, relu), Dense(3, 6, relu))
   @test outputsize(m, (2,), (3,)) == (10,)