optimized hcat of onehot vectors and matrices

[racinmat]

PR Checklist

• Tests are added
• Entry in NEWS.md
• Documentation, if applicable
• API changes require approval from a committer (different from the author, if applicable)

Fixes #1594 , adds tests for it and also optimized reduce(hcat, xs).
No new features are added, only performance optimization.

[darsnack]

Thanks for this! I left some suggestions that make sure we hit reshaped arrays as well.

[johnnychen94]

I believe the root issue is that we don't have similar defined for OneHotArray; once it's defined, it's very likely that many of these overly-verbose optimizations will be unnecessary.

[darsnack]

Let's go back to the previously suggested changes (I have commented them again with some updates). If the constructor is causing issues, then why not just change the constructor to:

OneHotArray(indices::I, L::Integer) where {T, N, I<: AbstractArray{T, N}} = OneHotArray{T, L, N, I}(indices)

Then much more than hcat can see the performance benefit.

[darsnack]

I believe the root is is that we don't have similar defined for OneHotArray; once it's defined, it's very likely that many of these overly-verbose optimizations will be unnecessary.

Good point! Though this means we would need to pair similar with setindex! which would require checks to stop someone from breaking the one-hot behavior. I think directly operating on the index arrays will still see a performance benefit.

We might want to define similar anyways to guarantee more type stability in fallback cases (e.g. the reduce(hcat, ...) one).

[DhairyaLGandhi]

Ideally we would fix it generally rather than for special cases. That's the flux way.

[darsnack]

Here is an implementation of similar:

function Base.similar(x::OneHotArray, ::Type{Bool}, dims::Dims)
  indices = similar(_indices(x), Base.tail(dims))

  return OneHotArray(indices, first(dims))
end

This would need to be paired with Base.setindex! to address the reduce(hcat, ...) type instability (too late for me to figure out the index combos tonight 😅). Though I am guessing that since setindex! will need to perform checks to ensure the one hot property is maintained, this will not be as fast as an optimized custom reduce(::typeof(hcat), ...). I think we should define similar no matter what though.

[racinmat]

Ok, the performance is now good even for reshaped arrays. Should I now include the Base.similar in this PR, or not, because of pairing with Base.setindex!?
Is it ready to merge?
I agree that generic fix would be the best solution, but I think this is good as an intermediate thing before the proper implementation.

[darsnack]

I think the best move is to remove the reduce(hcat, ...) paths from this PR. Leave #1596 open, and we address it later with similar and setindex!.

[darsnack]

Just some readability changes in addition to the request to remove the reduce(hcat, ...) paths for now.

[racinmat]

Done, reduce and tests for it are removed.

[darsnack]

Looks great, thanks!

[DhairyaLGandhi]

bors r+

[bors]

Build succeeded:

• buildkite/flux-dot-jl