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,19 @@ 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::OneHotLike{<:Any, L, 1}) 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, _indices(B))
+end
+
+function Base.:(*)(A::AbstractMatrix, B::Adjoint{Bool, <:OneHotMatrix})
+  B_dim = length(_indices(parent(B)))
+  size(A, 2) == B_dim || throw(DimensionMismatch("Matrix column must correspond with OneHot size: $(size(A, 2)) != $B_dim"))
+  return NNlib.scatter(+, A, _indices(parent(B)), dstsize=(size(A,1), size(B,2)))
+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, 2) |> gpu;
+  @test gradient(A -> sum(A * y), gA)[1] isa CuArray
 end
 
 @testset "onecold gpu" begin

test/onehot.jl

@@ -32,7 +32,7 @@ end
   b1 = Flux.OneHotVector(1, 3)
   b2 = Flux.OneHotVector(3, 5)
 
-  @test A*b1 == A[:,1]
+  @test A * b1 == A[:,1]
   @test b1' * A == Array(b1') * A
   @test A' * b1 == A' * Array(b1)
   @test v' * b2 == v' * Array(b2)
@@ -41,6 +41,35 @@ 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, 2], 3)
+  b2 = Flux.OneHotMatrix([2, 4, 1, 3], 5)
+  b3 = Flux.OneHotMatrix([1, 1, 2], 4)
+  b4 = reshape(Flux.OneHotMatrix([1 2 3; 2 2 1], 3), 3, :)
+  b5 = reshape(b4, 6, :)
+  b6 = reshape(Flux.OneHotMatrix([1 2 2; 2 2 1], 2), 3, :)
+  b7 = reshape(Flux.OneHotMatrix([1 2 3; 1 2 3], 3), 6, :)
+
+  @test A * b1 == A[:,[1, 1, 2, 2]]
+  @test b1' * A == Array(b1') * A
+  @test A' * b1 == A' * Array(b1)
+  @test A * b3' == A * Array(b3')
+  @test transpose(X) * b2 == transpose(X) * Array(b2)
+  @test A * b4 == A[:,[1, 2, 2, 2, 3, 1]]
+  @test A * b5' == hcat(A[:,[1, 2, 3, 3]], A[:,1]+A[:,2], zeros(Int64, 3))
+  @test A * b6 == hcat(A[:,1], 2*A[:,2], A[:,2], A[:,1]+A[:,2])
+  @test A * b7' == A[:,[1, 2, 3, 1, 2, 3]]
+
+  @test_throws DimensionMismatch A*b1'
+  @test_throws DimensionMismatch A*b2
+  @test_throws DimensionMismatch A*b2'
+  @test_throws DimensionMismatch A*b6'
+  @test_throws DimensionMismatch A*b7
+end
+
 @testset "OneHotArray" begin
   using Flux: OneHotArray, OneHotVector, OneHotMatrix, OneHotLike