Merge #1424

1424: multiplication of {Transpose, Adjoint} of Array and OneHotVector r=darsnack a=gxyd

Also fixes #777 . Tests have been added.

Co-authored-by: Gaurav Dhingra <gauravdhingra.gxyd@gmail.com>
Co-authored-by: Kyle Daruwalla <daruwalla@wisc.edu>

Author: 3 people

src/onehot.jl

@@ -187,3 +187,20 @@ 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
+for wrapper in [:Adjoint, :Transpose]
+  @eval begin
+    function Base.:*(A::$wrapper{<:Any, <:AbstractMatrix{T}}, b::OneHotVector{<:Any, L}) where {L, T}
+      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::$wrapper{<:Number, <:AbstractVector{T}}, b::OneHotVector{<:Any, L}) where {L, T}
+      size(A, 2) == L ||
+        throw(DimensionMismatch("Matrix column must correspond with OneHot size: $(size(A, 2)) != $L"))
+
+      return A[onecold(b)]
+    end
+  end
+end

test/onehot.jl

@@ -27,10 +27,17 @@ end
 
 @testset "abstractmatrix onehotvector 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.OneHotVector(1, 3)
   b2 = Flux.OneHotVector(3, 5)
 
   @test A*b1 == A[:,1]
+  @test b1' * A == Array(b1') * A
+  @test A' * b1 == A' * Array(b1)
+  @test v' * b2 == v' * Array(b2)
+  @test transpose(X) * b2 == transpose(X) * Array(b2)
+  @test transpose(v) * b2 == transpose(v) * Array(b2)
   @test_throws DimensionMismatch A*b2
 end
 
@@ -132,4 +139,4 @@ end
     @test map(identity, oa) == oa
     @test map(x -> 2 * x, oa) == 2 .* oa
   end
-end
+end