Add projections and remove redundant conj calls

Author: Simone Carlo Surace

src/rulesets/LinearAlgebra/dense.jl

@@ -417,30 +417,36 @@ end
     end
 
     function rrule(::typeof(kron), x::AbstractMatrix{<:Number}, y::AbstractVector{<:Number})
+        project_x = ProjectTo(x)
+        project_y = ProjectTo(y)
         function kron_pullback(z̄)
             dz = reshape(unthunk(z̄), length(y), size(x)...)
-            x̄ = @thunk conj.(dot.(eachslice(dz; dims = (2, 3)), Ref(y)))
-            ȳ = @thunk conj.(dot.(eachslice(dz; dims = 1), Ref(x)))
+            x̄ = @thunk(project_x(dot.(Ref(y), eachslice(dz; dims = (2, 3)))))
+            ȳ = @thunk(project_y(dot.(Ref(x), eachslice(dz; dims = 1))))
             return NoTangent(), x̄, ȳ
         end
         return kron(x, y), kron_pullback
     end
 
     function rrule(::typeof(kron), x::AbstractVector{<:Number}, y::AbstractMatrix{<:Number})
+        project_x = ProjectTo(x)
+        project_y = ProjectTo(y)
         function kron_pullback(z̄)
             dz = reshape(unthunk(z̄), size(y, 1), length(x), size(y, 2))
-            x̄ = @thunk conj.(dot.(eachslice(dz; dims = 2), Ref(y)))
-            ȳ = @thunk conj.(dot.(eachslice(dz; dims = (1, 3)), Ref(x)))
+            x̄ = @thunk(project_x(dot.(Ref(y), eachslice(dz; dims = 2))))
+            ȳ = @thunk(project_y(dot.(Ref(x), eachslice(dz; dims = (1, 3)))))
             return NoTangent(), x̄, ȳ
         end
         return kron(x, y), kron_pullback
     end
 
     function rrule(::typeof(kron), x::AbstractMatrix{<:Number}, y::AbstractMatrix{<:Number})
+        project_x = ProjectTo(x)
+        project_y = ProjectTo(y)
         function kron_pullback(z̄)
             dz = reshape(unthunk(z̄), size(y, 1), size(x, 1), size(y, 2), size(x, 2))
-            x̄ = @thunk conj.(dot.(eachslice(dz, dims = (2, 4)), Ref(y)))
-            ȳ = @thunk conj.(dot.(eachslice(dz; dims = (1, 3)), Ref(x)))
+            x̄ = @thunk(project_x(dot.(Ref(y), eachslice(dz; dims = (2, 4)))))
+            ȳ = @thunk(project_y(dot.(Ref(x), eachslice(dz; dims = (1, 3)))))
             return NoTangent(), x̄, ȳ
         end
         return kron(x, y), kron_pullback

test/rulesets/LinearAlgebra/dense.jl

@@ -160,18 +160,24 @@
         end
     end
     @testset "kron" begin
-        @testset "AbstractVecOrMat{$T}" for T in (Float64, ComplexF64)
+        @testset "AbstractVecOrMat{$T1}, AbstractVecOrMat{$T2}" for T1 in (Float64, ComplexF64), T2 in (Float64, ComplexF64)
             @testset "frule" begin
-                test_frule(kron, randn(T, 2), randn(T, 3))
-                test_frule(kron, randn(T, 2, 3), randn(T, 5))
-                test_frule(kron, randn(T, 2), randn(T, 3, 5))
-                test_frule(kron, randn(T, 2, 3), randn(T, 5, 7))
+                test_frule(kron, randn(T1, 2), randn(T2, 3))
+                test_frule(kron, randn(T1, 2, 3), randn(T2, 5))
+                test_frule(kron, randn(T1, 2), randn(T2, 3, 5))
+                test_frule(kron, randn(T1, 2, 3), randn(T2, 5, 7))
             end
             @testset "rrule" begin
-                test_rrule(kron, randn(T, 2), randn(T, 3))
-                test_rrule(kron, randn(T, 2, 3), randn(T, 5))
-                test_rrule(kron, randn(T, 2), randn(T, 3, 5))
-                test_rrule(kron, randn(T, 2, 3), randn(T, 5, 7))
+                test_rrule(kron, randn(T1, 2), randn(T2, 3))
+
+                test_rrule(kron, Diagonal(randn(T1, 2)), randn(T2, 3))
+                test_rrule(kron, randn(T1, 2, 3), randn(T2, 5))
+
+                test_rrule(kron, randn(T1, 2), randn(T2, 3, 5))
+                test_rrule(kron, randn(T1, 2), Diagonal(randn(T2, 3)))
+
+                test_rrule(kron, randn(T1, 2, 3), randn(T2, 5, 7))
+                test_rrule(kron, Diagonal(randn(T1, 2)), Diagonal(randn(T2, 3)))
             end
         end
     end