Speedup and fix of multiplication by OneHotMatrix

NEWS.md

@@ -1,5 +1,8 @@
 # Flux Release Notes
 
+## v0.12.8
+* Optimized inference and gradient calculation of OneHotMatrix[pr](https://github.com/FluxML/Flux.jl/pull/1756)
+
 ## v0.12.7
 * Added support for [`GRUv3`](https://github.com/FluxML/Flux.jl/pull/1675)
 * The layers within `Chain` and `Parallel` may now [have names](https://github.com/FluxML/Flux.jl/issues/1680).

src/onehot.jl

@@ -1,5 +1,6 @@
 import Adapt
 import .CUDA
+using LinearAlgebra, NNlib
 
 """
     OneHotArray{T,L,N,M,I} <: AbstractArray{Bool,M}
@@ -224,6 +225,13 @@ function Base.:(*)(A::AbstractMatrix, B::OneHotLike{<:Any, L}) where L
   size(A, 2) == L || throw(DimensionMismatch("Matrix column must correspond with OneHot size: $(size(A, 2)) != $L"))
   return A[:, onecold(B)]
 end
+
+function Base.:(*)(A::AbstractMatrix, B::OneHotMatrix{<:Any, L}) where L
+  _isonehot(B) || return invoke(*, Tuple{AbstractMatrix, AbstractMatrix}, A, B)
+  size(A, 2) == L || throw(DimensionMismatch("Matrix column must correspond with OneHot size: $(size(A, 2)) != $L"))
+  return NNlib.gather(A, B.indices)
+end
+
 for wrapper in [:Adjoint, :Transpose]
   @eval begin
     function Base.:*(A::$wrapper{<:Any, <:AbstractMatrix{T}}, b::OneHotVector{<:Any, L}) where {L, T}

test/cuda/cuda.jl

@@ -42,6 +42,9 @@ end
 @testset "onehot gpu" begin
   y = Flux.onehotbatch(ones(3), 1:2) |> gpu;
   @test (repr("text/plain", y); true)
+
+  gA = rand(3, 3) |> gpu;
+  @test gradient(A -> sum(A * y), gA)[1] isa CuArray
 end
 
 @testset "onecold gpu" begin

test/onehot.jl

@@ -41,6 +41,21 @@ end
   @test_throws DimensionMismatch A*b2
 end
 
+@testset "AbstractMatrix-OneHotMatrix multiplication" begin
+  A = [1 3 5; 2 4 6; 3 6 9]
+  v = [1, 2, 3, 4, 5]
+  X = reshape(v, (5, 1))
+  b1 = Flux.OneHotMatrix([1, 1, 2], 3)
+  b2 = Flux.OneHotMatrix([2, 4, 1, 3], 5)
+
+  @test A*b1 == A[:,[1, 1, 2]]
+  @test b1' * A == Array(b1') * A
+  @test A' * b1 == A' * Array(b1)
+  @test A * b1' == A * Array(b1')
+  @test transpose(X) * b2 == transpose(X) * Array(b2)
+  @test_throws DimensionMismatch A*b2
+end
+
 @testset "OneHotArray" begin
   using Flux: OneHotArray, OneHotVector, OneHotMatrix, OneHotLike