return OneHotArray on proper reshape?

[CarloLucibello]

Right now a reshape of a OneHotArray that preserves the first dimension doesn't return a OneHotArray:

julia> oh = onehotbatch(rand('a':'b', (3, 2)), 'a':'b')
2×3×2 OneHotArray(::Matrix{UInt32}) with eltype Bool:
[:, :, 1] =
 0  0  0
 1  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 Base.ReshapedArray
true

It would be convenient to return onehotarrays in that case, for convenience and in order to avoid downstream such as
FluxML/Flux.jl#2160.

If we want to do this we have two options:

1. give up type instability, and return a onehotarray only if the first dimension is preserved, and a reshaped array otherwise
2. error out on any reshape that doesn't preserve the first dimension.

I remember we already had a discussion on reshape somewhere, but maybe worth revisiting.

[CarloLucibello]

The discussion is at FluxML/Flux.jl#1459.

Some arguments were made in favor of allowing arbitrary reshaping, so we went with the ReshapedArray thing. I wonder how problematic type instability would be instead.

[mcabbott]

The reshape there is implementing a particular multiplication, but looks like it could implement anything. A possible solution is to specify what it's doing at a higher level, and overload that operation for OneHotArray. Is this the right operation?

https://juliamath.github.io/TensorCore.jl/dev/#TensorCore.boxdot

[CarloLucibello]

Yes, \boxdot is the right operation. Dense does the same as well, again with the reshape trick
https://github.com/FluxML/Flux.jl/blob/7ae67a3d41d0cacbe616410087af3a60e50d95ec/src/layers/basic.jl#L175

[darsnack]

OneHotArrays.jl already has a solution for this: to dispatch on OneHotLike. This keeps type stability and allows any code paths to be the same for real and properly reshaped one hot arrays. I don't think we should be returning a OneHotArray. A basic array package should try to be type-stable as much as possible.

I suspect the downstream issue is caused by an outdated Flux (<0.13.7) before FluxML/Flux.jl#2084. Previously, Flux's Embedding dispatched on OneHotVector/OneHotMatrix directly which is not correct (it should use OneHotLike). After that PR, this does not even matter, because OneHotArrays.jl isn't even in the picture...just AbstractArray{Bool}. An up to date Flux would dispatch correctly as seen below.

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