Implement 'LinearAlgebra.diag' and one necessary simplification

Author: asterycs

src/ricci.jl

@@ -845,6 +845,20 @@ function LinearAlgebra.diagm(v::Tensor)
     return BinaryOperation{Mult}(id, v)
 end
 
+function LinearAlgebra.diag(m::Tensor, k::Integer = 0)
+    indices = get_free_indices(m)
+
+    if length(indices) != 2
+        throw(DomainError(indices, "Input is not a matrix, cannot extract the diagonal"))
+    end
+
+    if k != 0
+        throw(DomainError(k, "Cannot extract the k-th diagonal when k!=0"))
+    end
+
+    return (m .* LinearAlgebra.I) * vector(1)
+end
+
 function script(index::Lower)
     @assert index.letter >= 0
     text = []

src/simplify.jl

@@ -125,8 +125,45 @@ function simplify(::Mult, arg1::Tensor, arg2::BinaryOperation{Mult})
     return simplify(Mult(), arg2, arg1)
 end
 
+function can_apply(l::BinaryOperation{Mult}, r::KrD)
+    return can_contract(l.arg1, r) || can_contract(l.arg2, r)
+end
+
+function can_apply(l::KrD, r::BinaryOperation{Mult})
+    return can_contract(l, r.arg1) || can_contract(l, r.arg2)
+end
+
+function can_apply(l::BinaryOperation{Mult}, r::BinaryOperation{Mult})
+    return can_apply(l, r.arg1) || can_apply(l, r.arg2)
+end
+
+function can_apply(l::BinaryOperation{Mult}, r::Tensor)
+    return can_contract(l.arg1, r) || can_contract(l.arg2, r)
+end
+
+function can_apply(l::Tensor, r::BinaryOperation{Mult})
+    return can_contract(l, r.arg1) || can_contract(l, r.arg2)
+end
+
+function can_apply(l::Tensor, r::KrD)
+    return can_contract(l, r)
+end
+
+function can_apply(l::KrD, r::Tensor)
+    return can_contract(l, r)
+end
+
+function can_apply(l::KrD, r::KrD)
+    return can_contract(l, r)
+end
+
+function can_apply(l, r)
+    return false
+end
+
 function simplify(::Mult, arg1::BinaryOperation{Mult}, arg2::Tensor)
     op = BinaryOperation{Mult}(arg1, arg2)
+    target_indices = unique(get_free_indices(op))
 
     if is_diagm(arg1) &&
        !is_elementwise_multiplication(arg1, arg2) &&
@@ -135,7 +172,6 @@ function simplify(::Mult, arg1::BinaryOperation{Mult}, arg2::Tensor)
 
         @assert !isnothing(d)
 
-        target_indices = eliminate_indices(vcat(get_free_indices(arg1), get_indices(arg2)))
         factors = collect_factors(arg1)
         vector_factors = filter(f -> f != d, factors)
         reshaped = []
@@ -189,6 +225,32 @@ function simplify(::Mult, arg1::BinaryOperation{Mult}, arg2::Tensor)
         return to_binary_operation(Mult(), reshaped)
     end
 
+    if length(get_free_indices(arg1)) > 2 && length(get_free_indices(arg2)) == 1
+        if can_apply(arg1.arg1, arg2) &&
+           can_apply(arg1.arg2, arg2) &&
+           arg1.arg1 isa KrD &&
+           arg1.arg2 isa KrD
+            r_free_indices = get_free_indices(arg2)
+
+            new_r = update_index(
+                arg2,
+                only(r_free_indices),
+                first(target_indices);
+                allow_shape_change = true,
+            )
+
+            eliminated = eliminated_indices([get_free_indices(arg1); r_free_indices])
+            new_l = update_index(
+                arg1.arg2,
+                first(eliminated),
+                first(target_indices);
+                allow_shape_change = true,
+            )
+
+            return BinaryOperation{Mult}(new_l, new_r)
+        end
+    end
+
     return op
 end
 

test/RicciTest.jl

@@ -9,7 +9,7 @@ using Test
 using DiffMatic: Variable, Literal, KrD, Zero
 using DiffMatic: Upper, Lower
 
-using LinearAlgebra: norm, tr, I, diagm
+using LinearAlgebra: norm, tr, I, diagm, diag
 
 dc = DiffMatic
 

test/StdStrTest.jl

@@ -32,6 +32,7 @@
     @test to_std(gradient(sum(2 * x), x)) == "2vec(1)"
     @test to_std(gradient(2 * sum(sin.(x)), x)) == "2cos(x)"
     @test to_std(gradient(sum(2 * sin.(x)), x)) == "2cos(x)"
+    @test to_std(gradient(diag(A)'*x, x)) == "(A ⊙ I)vec(1)" # TODO: Output 'diag(A)'
     @test to_std(gradient(2 * sum(cos.(A * x + y)), x)) == "(-2)Aᵀsin(Ax + y)"
     @test to_std(gradient(sum(2 * cos.(A * x + y)), x)) == "(-2)Aᵀsin(Ax + y)"
     @test to_std(gradient(sum(x .^ 2), x)) == "2x"
@@ -73,6 +74,7 @@ end
     @matrix A B C X
     @vector x y z
 
+    @test to_std(derivative(diag(A)'*x, A)) == "Iᵀdiagm(x)"
     @test to_std(derivative(sum(-y .* (X*z)), X)) == "(-1)zyᵀ"
     @test to_std(derivative(sum((A .* B) * C * x), x)) == "vec(1)ᵀ(A ⊙ B)C"
 end