Add straightforward simplifications for products involving quotients

Author: asterycs

src/forward.jl

@@ -308,6 +308,26 @@ function evaluate(::Mult, arg1::BinaryOperation{Div}, arg2::KrD)
     return BinaryOperation{Mult}(arg1, arg2)
 end
 
+function evaluate(::Mult, arg1::BinaryOperation{Div}, arg2::Tensor)
+    if arg1.arg2 == arg2
+        return simplify(arg1.arg1)
+    end
+
+    if arg1.arg1 isa Literal
+        if arg1.arg1.value == 1 && get_free_indices(arg1.arg1) == get_free_indices(arg2)
+            return BinaryOperation{Div}(arg2, arg1.arg2)
+        end
+    end
+
+    return BinaryOperation{Mult}(arg1, arg2)
+end
+
+function evaluate(::Mult, arg1::BinaryOperation{Div}, arg2::Zero)
+    free_indices = unique(eliminate_indices([get_indices(arg1); get_indices(arg2)]))
+
+    return Zero(free_indices...)
+end
+
 function evaluate(::Mult, arg1::BinaryOperation{Mult}, arg2::KrD)
     ci = indices_in_common(arg1.arg1, arg1.arg2)
 

src/simplify.jl

@@ -111,6 +111,10 @@ function simplify(::Mult, arg1::BinaryOperation{Mult}, arg2::Literal)
     return BinaryOperation{Mult}(arg1, arg2)
 end
 
+function simplify(::Mult, arg1::Tensor, arg2::BinaryOperation{Mult})
+    return simplify(Mult(), arg2, arg1)
+end
+
 function simplify(::Mult, arg1::BinaryOperation{Mult}, arg2::Tensor)
     if is_diag(arg1) && !is_elementwise_multiplication(arg1, arg2)
         d = get_diag_delta(arg1)
@@ -165,6 +169,10 @@ function simplify(::Mult, arg1::BinaryOperation{Mult}, arg2::Tensor)
     return BinaryOperation{Mult}(arg1, arg2)
 end
 
+function simplify(::Mult, arg1::BinaryOperation{Mult}, arg2::BinaryOperation{Mult})
+    return BinaryOperation{Mult}(arg1, arg2)
+end
+
 function simplify(::Div, arg1::Value, arg2::Value)
     return evaluate(
         BinaryOperation{Div}(simplify(evaluate(arg1)), simplify(evaluate(arg2))),

test/StdStrTest.jl

@@ -21,7 +21,7 @@
     @test to_std(gradient(sin.(x)' * y * a, x)) == "a(cos(x) ⊙ y)"
     @test to_std(gradient(x' * sin.(y) * a, x)) == "asin(y)"
     @test to_std(gradient(y' * sin.(x) * a, x)) == "a(cos(x) ⊙ y)"
-    @test to_std(gradient(sin.(x .* y)' * x, x)) == "(cos(x ⊙ y) ⊙ y ⊙ x) + sin(x ⊙ y)"
+    @test to_std(gradient(sin.(x .* y)' * x, x)) == "(y ⊙ cos(x ⊙ y) ⊙ x) + sin(x ⊙ y)"
     @test to_std(gradient(sum(x), x)) == "vec(1)"
     @test to_std(gradient(2 * sum(x), x)) == "2vec(1)"
     @test to_std(gradient(sum(2 * x), x)) == "2vec(1)"
@@ -36,9 +36,8 @@
     @test to_std(gradient(sum((x + y) .^ 2), x)) == "2(x + y)"
     @test to_std(gradient(sum((x .* y) .^ 2), x)) == "2(x ⊙ y ⊙ y)"
     @test to_std(gradient(sum((A * x - y) .^ 2), x)) == "2Aᵀ(Ax - y)"
-    @test to_std(gradient(log.(x)'*x, x)) == "(vec(1) ⊘ x ⊙ x) + log(x)" # TODO: Add simplification rule for quotients
-    @test to_std(gradient(log.(x)'*log.(x), x)) ==
-          "diag(vec(1)ᵀ ⊘ xᵀ)Iᵀlog(x) + (vec(1) ⊘ x ⊙ log(x))"
+    @test to_std(gradient(log.(x)'*x, x)) == "vec(1) + log(x)"
+    @test to_std(gradient(log.(x)'*log.(x), x)) == "2(log(x) ⊘ x)"
     @test to_std(gradient((x' * A * x) ^ (-2), x)) == "(-2)(xᵀAᵀx)^(-3)(Aᵀx + Ax)"
     @test to_std(gradient((x' * A * x) ^ 2, x)) == "2xᵀAᵀx(Aᵀx + Ax)"
     @test to_std(gradient(((A .* B) * C * x)' * x, x)) == "(A ⊙ B)Cx + Cᵀ(Aᵀ ⊙ Bᵀ)x"