From 9adb77bd30a875e806012c9dd6f15d23cd00f071 Mon Sep 17 00:00:00 2001 From: Will Tebbutt Date: Tue, 26 Sep 2023 11:22:26 +0100 Subject: [PATCH] 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