Generalise kron implementation

FluxML · Sep 26, 2023 · 9adb77b · 9adb77b
1 parent 9ea933b
commit 9adb77b
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 5 deletions.
diff --git a/src/lib/array.jl b/src/lib/array.jl
@@ -368,12 +368,10 @@ function _kron(mat1::AbstractMatrix,mat2::AbstractMatrix)
     return reshape(mat1_rsh.*mat2_rsh, (m1*m2,n1*n2))
 end
 _kron(a::AbstractVector, b::AbstractVector) = vec(_kron(reshape(a, :, 1), reshape(b, :, 1)))
+_kron(a::AbstractVector, b::AbstractMatrix) = _kron(reshape(a, :, 1), b)
+_kron(a::AbstractMatrix, b::AbstractVector) = _kron(a, reshape(b, :, 1))
 
-function _pullback(cx::AContext, ::typeof(kron), a::AbstractVector, b::AbstractVector)
-  res, back = _pullback(cx, _kron, a, b)
-  return res, back ∘ unthunk_tangent
-end
-function _pullback(cx::AContext, ::typeof(kron), a::AbstractMatrix, b::AbstractMatrix)
+function _pullback(cx::AContext, ::typeof(kron), a::AbstractVecOrMat, b::AbstractVecOrMat)
   res, back = _pullback(cx, _kron, a, b)
   return res, back ∘ unthunk_tangent
 end

diff --git a/test/gradcheck.jl b/test/gradcheck.jl
@@ -275,6 +275,8 @@ end
 @test gradtest(kron, rand(5,1), rand(3,1))
 @test gradtest(kron, rand(5,1), rand(3,1), rand(8,1))
 @test gradtest(kron, rand(5,2), rand(3,2), rand(8,2))
+@test gradtest(kron, rand(5), rand(3, 2))
+@test gradtest(kron, rand(3, 2), rand(5))
 
 for mapfunc in [map,pmap]
   @testset "$mapfunc" begin