From 9ea933b8fb7171ae42916a0a430e12261f3a57e2 Mon Sep 17 00:00:00 2001 From: Will Tebbutt Date: Tue, 26 Sep 2023 11:22:01 +0100 Subject: [PATCH 1/2] Bump patch --- Project.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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" From 9adb77bd30a875e806012c9dd6f15d23cd00f071 Mon Sep 17 00:00:00 2001 From: Will Tebbutt Date: Tue, 26 Sep 2023 11:22:26 +0100 Subject: [PATCH 2/2] Generalise kron implementation --- src/lib/array.jl | 8 +++----- test/gradcheck.jl | 2 ++ 2 files changed, 5 insertions(+), 5 deletions(-) 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