fixed dimension check, added tests to check different dimensionality

Author: racinmat

src/onehot.jl

@@ -231,8 +231,9 @@ function Base.:(*)(A::AbstractMatrix, B::OneHotMatrix{<:Any, L}) where L
   return NNlib.gather(A, B.indices)
 end
 
-function Base.:(*)(A::AbstractMatrix, B::Adjoint{Bool, <:OneHotMatrix{<:Any, L}}) where L
-  size(A, 2) == L || throw(DimensionMismatch("Matrix column must correspond with OneHot size: $(size(A, 2)) != $L"))
+function Base.:(*)(A::AbstractMatrix, B::Adjoint{Bool, <:OneHotMatrix})
+  B_dim = length(parent(B).indices)
+  size(A, 2) == B_dim || throw(DimensionMismatch("Matrix column must correspond with OneHot size: $(size(A, 2)) != $B_dim"))
   return NNlib.scatter(+, A, parent(B).indices, dstsize=(size(A,1), size(B,2)))
 end
 

test/onehot.jl

@@ -45,14 +45,16 @@ end
   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)
+  b1 = Flux.OneHotMatrix([1, 1, 2, 2], 3)
   b2 = Flux.OneHotMatrix([2, 4, 1, 3], 5)
+  b3 = Flux.OneHotMatrix([1, 1, 2], 4)
 
-  @test A*b1 == A[:,[1, 1, 2]]
+  @test A*b1 == A[:,[1, 1, 2, 2]]
   @test b1' * A == Array(b1') * A
   @test A' * b1 == A' * Array(b1)
-  @test A * b1' == A * Array(b1')
+  @test A * b3' == A * Array(b3')
   @test transpose(X) * b2 == transpose(X) * Array(b2)
+  @test_throws DimensionMismatch A*b1'
   @test_throws DimensionMismatch A*b2
   @test_throws DimensionMismatch A*b2'
 end