Add rules and tests

Author: Simone Carlo Surace

src/rulesets/LinearAlgebra/dense.jl

@@ -394,3 +394,51 @@ function rrule(
     end
     return Ω, lyap_pullback
 end
+
+#####
+##### `kron`
+#####
+
+function frule((_, Δx, Δy), ::typeof(kron), x, y)
+    return kron(x, y), kron(Δx, y) + kron(x, Δy)
+end
+
+function rrule(::typeof(kron), x::AbstractMatrix, y::AbstractVector)
+    z = kron(x, y)
+
+    function kron_pullback(z̄)
+        x̄ = zero(x)
+        ȳ = zero(y)
+        m = firstindex(z̄)
+        @inbounds for j in axes(x,2), i in axes(x,1)
+            xij = x[i,j]
+            for k in eachindex(y)
+                x̄[i, j] += y[k]' * z̄[m]
+                ȳ[k] += xij * z̄[m]
+                m += 1
+            end
+        end
+        NoTangent(), x̄, ȳ
+    end
+    z, kron_pullback
+end
+
+function rrule(::typeof(kron), x::AbstractVector, y::AbstractMatrix)
+    z = kron(x, y)
+
+    function kron_pullback(z̄)
+        x̄ = zero(x)
+        ȳ = zero(y)
+        m = firstindex(z̄)
+        @inbounds for l in axes(y,2), i in eachindex(x)
+            xi = x[i]
+            for k in axes(y,1)
+                x̄[i] += y[k, l]' * z̄[m]
+                ȳ[k, l] += xi * z̄[m]
+                m += 1
+            end
+        end
+        NoTangent(), x̄, ȳ
+    end
+    z, kron_pullback
+end

test/rulesets/LinearAlgebra/dense.jl

@@ -159,4 +159,12 @@
             test_rrule(lyap, A, C)
         end
     end
+    @testset "kron" begin
+        @testset "AbstractVecOrMat{$T}" for T in (Float64, ComplexF64)
+            test_frule(kron, randn(T, 3), randn(T, 3))
+            test_frule(kron, randn(T, 3, 2), randn(T, 3))
+            test_frule(kron, randn(T, 3), randn(T, 3, 4))
+            test_frule(kron, randn(T, 3, 4), randn(T, 2, 2))
+        end
+    end
 end