diff --git a/Project.toml b/Project.toml index dd390100c..5105e6a06 100644 --- 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 index 37884cded..7740a80d4 100644 --- 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 index b7fd5391f..bc22a2016 100644 --- 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