Test evaluate subtraction with product with real

asterycs · Feb 16, 2025 · 15d5f3a · 15d5f3a
1 parent 907a857
commit 15d5f3a
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 2 deletions.
diff --git a/src/forward.jl b/src/forward.jl
@@ -728,11 +728,11 @@ function _sub_from_product(arg1::BinaryOperation{Mult}, arg2::Value)
     end
 
     if evaluate(arg1.arg1) isa Real && evaluate(arg1.arg2) == evaluate(arg2)
-        return BinaryOperation{Mult}(evaluate(arg1.arg1) - 1, evaluate(arg2))
+        return evaluate(BinaryOperation{Mult}(evaluate(arg1.arg1) - 1, evaluate(arg2)))
     end
 
     if evaluate(arg1.arg2) isa Real && evaluate(arg1.arg1) == evaluate(arg2)
-        return BinaryOperation{Mult}(evaluate(arg1.arg2) - 1, evaluate(arg2))
+        return evaluate(BinaryOperation{Mult}(evaluate(arg1.arg2) - 1, evaluate(arg2)))
     end
 
     return BinaryOperation{Sub}(evaluate(arg1), evaluate(arg2))

diff --git a/test/ForwardTest.jl b/test/ForwardTest.jl
@@ -307,6 +307,23 @@ end
     @test equivalent(evaluate(op2), Zero(Upper(2)))
 end
 
+@testset "evaluate subtraction with product with real" begin
+    A = Tensor("A", Upper(1), Lower(2))
+
+    function mul(l, r)
+        return dc.BinaryOperation{dc.Mult}(l, r)
+    end
+
+    function sub(l, r)
+        return dc.BinaryOperation{dc.Sub}(l, r)
+    end
+
+    @test dc.evaluate(sub(mul(2, A), A)) == A
+    @test dc.evaluate(sub(mul(A, 2), A)) == A
+    # @test dc.evaluate(sub(mul(A, 2), A)) == A
+    # @test dc.evaluate(sub(mul(A, 2), A)) == A
+end
+
 @testset "evaluate unary operations" begin
     A = Tensor("A", Upper(1), Lower(2))