[Bug] Embedding forward pass breaks for onehotbatch with multiple batch dimensions

[reachtarunhere]

To reproduce:

using Flux: Embedding, onehotbatch, onehot
emb = Embedding(26, 2)
x1 = rand('a':'z', 10)
x1 = onehotbatch(x1, 'a':'z')
emb(x1) # works perfectly

# Breaking example below
x2 = rand('a':'z', (3, 10))
x2 = onehotbatch(x2, 'a':'z')
emb(x2) # breaks!

Error:

ArgumentError: invalid index: false of type Bool

to_index(i::Bool) at indices.jl:293
to_index(A::Matrix{Float32}, i::Bool) at indices.jl:277
to_indices at indices.jl:333 [inlined]
to_indices at multidimensional.jl:823 [inlined]
to_indices at indices.jl:324 [inlined]
view at subarray.jl:176 [inlined]
_view(X::Matrix{Float32}, colons::Tuple{Colon}, k::Bool) at scatter.jl:38
gather!(dst::Matrix{Float32}, src::Matrix{Float32}, idx::Base.ReshapedArray{Bool, 1, Flux.OneHotArray{UInt32, 26, 2, 3, Matrix{UInt32}}, Tuple{Base.MultiplicativeInverses.SignedMultiplicativeInverse{Int64}, Base.MultiplicativeInverses.SignedMultiplicativeInverse{Int64}}}) at gather.jl:27
gather at gather.jl:77 [inlined]
Embedding at basic.jl:533 [inlined]
(::Flux.Embedding{Matrix{Float32}})(x::Flux.OneHotArray{UInt32, 26, 2, 3, Matrix{UInt32}}) at basic.jl:534
top-level scope at mlp.jl:41

The issue is because the second call dispatches to the wrong method

Flux.jl/src/layers/basic.jl

Line 702 in 7ae67a3

(m::Embedding)(x::AbstractArray{Bool}) = reshape(m(reshape(x, size(x,1), :)), :, size(x)[2:end]...)

If it had dispatched to a method with similar approach of emb.weight * x2 it would work perfectly.

Flux.jl/src/layers/basic.jl

Line 701 in 7ae67a3

(m::Embedding)(x::AbstractMatrix{Bool}) = m.weight * x # usually OneHotMatrix

OneHotMatrix or tensors should be treated the same?

Thanks for reading!

[CarloLucibello]

We should make a reshape that preserves the first dimension return a OneHotArray instead of a ReshapedArray.
For instance, by adding this method (which is wrong since doesn't check the sizes) the issue is fixed

Base.reshape(x::OneHotArray, n::Int, ::Colon) = OneHotMatrix(vec(x.indices), x.nlabels)

[CarloLucibello]

Another option is to add specialized methods (m::Embedding)(x::OneHotArray)

[mcabbott]

On what version of Flux (and Juila?) do you see the above stack trace?

julia> x2 |> summary
"26×3×10 OneHotArray(::Matrix{UInt32}) with eltype Bool"

julia> emb(x2) |> summary # breaks!
"2×3×10 Array{Float32, 3}"

(@v1.10) pkg> st Flux
Status `~/.julia/environments/v1.10/Project.toml`
  [587475ba] Flux v0.13.11

[darsnack]

I am fairly certain you are using Flux <0.13.7 based on the stack trace. An up to date version would hit matmul first, then correctly call the corresponding method in OneHotArrays.jl.

julia> oh = Flux.onehotbatch(rand('a':'b', (3, 2)), 'a':'b')
2×3×2 OneHotArray(::Matrix{UInt32}) with eltype Bool:
[:, :, 1] =
 1  0  0
 0  1  1

[:, :, 2] =
 0  0  0
 1  1  1

julia> x = reshape(oh, 2, :)
2×6 reshape(OneHotArray(::Matrix{UInt32}), 2, 6) with eltype Bool:
 1  ⋅  ⋅  ⋅  ⋅  ⋅
 ⋅  1  1  1  1  1

julia> x isa AbstractMatrix{Bool}
true

julia> m = Embedding(2, 2)
Embedding(2 => 2)   # 4 parameters

julia> m(x)
2×6 Matrix{Float32}:
 -0.425469  -0.976092  -0.976092  -0.976092  -0.976092  -0.976092
  0.770382  -0.165354  -0.165354  -0.165354  -0.165354  -0.165354

[reachtarunhere]

Thanks all for the prompt responses. I tried with Flux 0.13.11 and it works perfectly. Just to point out the example posted by @darsnack is different than mine. What he has posted worked for me earlier too since you reshape the 3 dimensional tensor into a matrix and the correct method was dispatched. I did the same thing without the reshape (see x2 in my code above) and got the error with the old version.

Anyway thank you all. Everything works for the new version. For reference my previous version was 0.12.10. I am using Julia 1.7