Extend kron support (FluxML#1458)

* Bump patch * Generalise kron implementation
Pangoraw · Dec 4, 2023 · 9bd69a9 · 9bd69a9
1 parent 29d4ecf
commit 9bd69a9
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 6 deletions.
diff --git a/Project.toml b/Project.toml
@@ -1,6 +1,6 @@
 name = "Zygote"
 uuid = "e88e6eb3-aa80-5325-afca-941959d7151f"
-version = "0.6.64"
+version = "0.6.65"
 
 [deps]
 AbstractFFTs = "621f4979-c628-5d54-868e-fcf4e3e8185c"

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