Repository: PennyLaneAI/pennylane (python)

Issue:
Add `lazy=True` keyword argument to `qml.s_prod`, `qml.prod`, and `qml.op_sum`
Currently, `qml.pow` and `qml.adjoint` both have a keyword argument `lazy=True` that indicates whether or not to perform some basic simplification steps:

```
>>> qml.adjoint(qml.PauliX(0), lazy=True)
Adjoint(PauliX(wires=[0]))
>>> qml.adjoint(qml.PauliX(0), lazy=False)
PauliX(wires=[0])
>>> qml.pow(qml.RX(1.2, 0), 2.0, lazy=False)
RX(2.4, wires=[0])
```

For `s_prod`, `prod`, and `op_sum`, if `lazy=False`, we can do something similar.

For `qml.s_prod(scalar, op, lazy=False)`, we should check if `op` is already a scalar product `SProd`.  If so, we combine the coefficients.

```pycon
>>> qml.s_prod(3, qml.s_prod(4, qml.PauliX(0)))
3*(4*(PauliX(wires=[0])))
>>> qml.s_prod(3, qml.s_prod(4, qml.PauliX(0)), lazy=False) #target behavior
12*(PauliZ(wires=[0]))
```

For `qml.prod(*ops, lazy=False)`, we should check if any of the provided operators are already of type `Prod`.  If so, we use its `operands` instead.

```pycon
>>> op = qml.prod(qml.S(0), qml.prod(qml.S(1), qml.T(1)))
>>> op
S(wires=[0]) @ (S(wires=[1]) @ T(wires=[1]))
>>> op.operands
(S(wires=[0]), S(wires=[1]) @ T(wires=[1]))
>>> qml.prod(qml.S(0), qml.prod(qml.S(1), qml.T(1)), lazy=False).operands # target behavior
(S(wires=[0]), S(wires=[1]), T(wires=[1]))
```

Same for `qml.op_sum`.  If any operators are of type `Sum`, we would use its constituents instead of itself.

Pull request description:
Add lazy=True keyword argument to qml.s_prod, qml.prod, and qml.op_sum
**Context:** Currently, qml.pow and qml.adjoint both have a keyword argument lazy=True that indicates whether or not to perform some basic simplification steps. We are now doing similar things for qml.s_prod, qml.prod, and qml.op_sum, if lazy=False.

**Description of the Change:** 

For qml.s_prod(scalar, op, lazy=False), we check if op is already a scalar product SProd. If so, we combine the coefficients.
For qml.prod(*ops, lazy=False), we check if any of the provided operators are already of type Prod. If so, we use its operands instead.
For qml.op_sum(*summands, lazy=False), we check if any of the summands are of type Sum. If so, we use its constituents instead. 

**Related GitHub Issues:**

Closes #3094 


Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index aed8a4a2c..0570aa7a9 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -409,6 +409,9 @@
   the optional `use_grouping` attribute.
   [(#3456)](https://github.com/PennyLaneAI/pennylane/pull/3456)
 
+* Add the optional argument lazy=True to functions `qml.s_prod`, `qml.prod` and `qml.op_sum` to allow simplification.
+  [(#3483)](https://github.com/PennyLaneAI/pennylane/pull/3483)
+
 * Updated `zyz_decomposition` function such that it now supports broadcast operators. This
   means that single-qubit `QubitUnitary` operators, instantiated from a batch of unitaries,
   can now be decomposed.
@@ -847,6 +850,7 @@ Edward Jiang
 Christina Lee
 Albert Mitjans Coma
 Romain Moyard
+Kevin Shen
 Mudit Pandey
 Matthew Silverman
 Jay Soni
diff --git a/pennylane/ops/op_math/prod.py b/pennylane/ops/op_math/prod.py
index c75904a21..e87e07393 100644
--- a/pennylane/ops/op_math/prod.py
+++ b/pennylane/ops/op_math/prod.py
@@ -30,6 +30,7 @@ from pennylane.ops.op_math.pow import Pow
 from pennylane.ops.op_math.sprod import SProd
 from pennylane.ops.op_math.sum import Sum
 from pennylane.ops.qubit.non_parametric_ops import PauliX, PauliY, PauliZ
+from pennylane.queuing import QueuingManager
 from pennylane.wires import Wires
 
 from .composite import CompositeOp
@@ -39,7 +40,7 @@ MAX_NUM_WIRES_KRON_PRODUCT = 9
 computing the sparse matrix representation."""
 
 
-def prod(*ops, do_queue=True, id=None):
+def prod(*ops, do_queue=True, id=None, lazy=True):
     """Construct an operator which represents the generalized product of the
     operators provided.
 
@@ -53,6 +54,7 @@ def prod(*ops, do_queue=True, id=None):
     Keyword Args:
         do_queue (bool): determines if the product operator will be queued. Default is True.
         id (str or None): id for the product operator. Default is None.
+        lazy=True (bool): If ``lazy=False``, a simplification will be peformed such that when any of the operators is already a product operator, its operands will be used instead.
 
     Returns:
         ~ops.op_math.Prod: the operator representing the product.
@@ -68,7 +70,21 @@ def prod(*ops, do_queue=True, id=None):
     array([[ 0, -1],
            [ 1,  0]])
     """
-    return Prod(*ops, do_queue=do_queue, id=id)
+    if lazy:
+        return Prod(*ops, do_queue=do_queue, id=id)
+
+    ops_simp = Prod(
+        *itertools.chain.from_iterable([op if isinstance(op, Prod) else [op] for op in ops]),
+        do_queue=do_queue,
+        id=id,
+    )
+
+    if do_queue:
+        for op in ops:
+            QueuingManager.update_info(op, owner=ops_simp)
+        QueuingManager.update_info(ops_simp, owns=ops)
+
+    return ops_simp
 
 
 class Prod(CompositeOp):
diff --git a/pennylane/ops/op_math/sprod.py b/pennylane/ops/op_math/sprod.py
index cdcb22c2d..169740170 100644
--- a/pennylane/ops/op_math/sprod.py
+++ b/pennylane/ops/op_math/sprod.py
@@ -26,11 +26,12 @@ from pennylane.interfaces import SUPPORTED_INTERFACES
 from pennylane.operation import Operator
 from pennylane.ops.op_math.pow import Pow
 from pennylane.ops.op_math.sum import Sum
+from pennylane.queuing import QueuingManager
 
 from .symbolicop import SymbolicOp
 
 
-def s_prod(scalar, operator, do_queue=True, id=None):
+def s_prod(scalar, operator, lazy=True, do_queue=True, id=None):
     r"""Construct an operator which is the scalar product of the
     given scalar and operator provided.
 
@@ -39,9 +40,9 @@ def s_prod(scalar, operator, do_queue=True, id=None):
         operator (~.operation.Operator): the operator which will get scaled.
 
     Keyword Args:
+        lazy=True (bool): If ``lazy=False`` and the operator is already a scalar product operator, the scalar provided will simply be combined with the existing scaling factor.
         do_queue (bool): determines if the scalar product operator will be queued. Default is True.
         id (str or None): id for the scalar product operator. Default is None.
-
     Returns:
         ~ops.op_math.SProd: The operator representing the scalar product.
 
@@ -56,7 +57,16 @@ def s_prod(scalar, operator, do_queue=True, id=None):
     array([[ 0., 2.],
            [ 2., 0.]])
     """
-    return SProd(scalar, operator, do_queue=do_queue, id=id)
+    if lazy or not isinstance(operator, SProd):
+        return SProd(scalar, operator, do_queue=do_queue, id=id)
+
+    sprod_op = SProd(scalar=scalar * operator.scalar, base=operator.base, do_queue=do_queue, id=id)
+
+    if do_queue:
+        QueuingManager.update_info(operator, owner=sprod_op)
+        QueuingManager.update_info(sprod_op, owns=operator)
+
+    return sprod_op
 
 
 class SProd(SymbolicOp):
diff --git a/pennylane/ops/op_math/sum.py b/pennylane/ops/op_math/sum.py
index 6a2a0d701..f3e596b37 100644
--- a/pennylane/ops/op_math/sum.py
+++ b/pennylane/ops/op_math/sum.py
@@ -15,6 +15,7 @@
 This file contains the implementation of the Sum class which contains logic for
 computing the sum of operations.
 """
+import itertools
 from copy import copy
 from typing import List
 from functools import reduce
@@ -24,11 +25,12 @@ import numpy as np
 import pennylane as qml
 from pennylane import math
 from pennylane.operation import Operator
+from pennylane.queuing import QueuingManager
 
 from .composite import CompositeOp
 
 
-def op_sum(*summands, do_queue=True, id=None):
+def op_sum(*summands, do_queue=True, id=None, lazy=True):
     r"""Construct an operator which is the sum of the given operators.
 
     Args:
@@ -38,6 +40,7 @@ def op_sum(*summands, do_queue=True, id=None):
         do_queue (bool): determines if the sum operator will be queued (currently not supported).
             Default is True.
         id (str or None): id for the Sum operator. Default is None.
+        lazy=True (bool): If ``lazy=False``, a simplification will be peformed such that when any of the operators is already a sum operator, its operands (summands) will be used instead.
 
     Returns:
         ~ops.op_math.Sum: The operator representing the sum of summands.
@@ -53,7 +56,21 @@ def op_sum(*summands, do_queue=True, id=None):
     array([[ 1,  1],
            [ 1, -1]])
     """
-    return Sum(*summands, do_queue=do_queue, id=id)
+    if lazy:
+        return Sum(*summands, do_queue=do_queue, id=id)
+
+    summands_simp = Sum(
+        *itertools.chain.from_iterable([op if isinstance(op, Sum) else [op] for op in summands]),
+        do_queue=do_queue,
+        id=id,
+    )
+
+    if do_queue:
+        for op in summands:
+            QueuingManager.update_info(op, owner=summands_simp)
+        QueuingManager.update_info(summands_simp, owns=summands)
+
+    return summands_simp
 
 
 class Sum(CompositeOp):


Test patch:
diff --git a/tests/ops/op_math/test_prod.py b/tests/ops/op_math/test_prod.py
index 0fbae773c..cba54421e 100644
--- a/tests/ops/op_math/test_prod.py
+++ b/tests/ops/op_math/test_prod.py
@@ -1026,6 +1026,30 @@ class TestWrapperFunc:
         assert prod_class_op.wires == prod_func_op.wires
         assert prod_class_op.parameters == prod_func_op.parameters
 
+    def test_lazy_mode(self):
+        """Test that by default, the operator is simply wrapped in `Prod`, even if a simplification exists."""
+        op = prod(qml.S(0), Prod(qml.S(1), qml.T(1)))
+
+        assert isinstance(op, Prod)
+        assert len(op) == 2
+
+    def test_non_lazy_mode(self):
+        """Test the lazy=False keyword."""
+        op = prod(qml.S(0), Prod(qml.S(1), qml.T(1)), lazy=False)
+
+        assert isinstance(op, Prod)
+        assert len(op) == 3
+
+    def test_nonlazy_mode_queueing(self):
+        """Test that if a simpification is accomplished, the metadata for the original op
+        and the new simplified op is updated."""
+        with qml.queuing.AnnotatedQueue() as q:
+            prod1 = prod(qml.S(1), qml.T(1))
+            prod2 = prod(qml.S(0), prod1, lazy=False)
+
+        assert q[prod1]["owner"] == prod2
+        assert prod1 in q[prod2]["owns"]
+
 
 class TestIntegration:
     """Integration tests for the Prod class."""
diff --git a/tests/ops/op_math/test_sprod.py b/tests/ops/op_math/test_sprod.py
index 4e750c93b..9107fb0de 100644
--- a/tests/ops/op_math/test_sprod.py
+++ b/tests/ops/op_math/test_sprod.py
@@ -645,6 +645,32 @@ class TestWrapperFunc:
         assert sprod_class_op.wires == sprod_func_op.wires
         assert sprod_class_op.parameters == sprod_func_op.parameters
 
+    def test_lazy_mode(self):
+        """Test that by default, the operator is simply wrapped in `SProd`, even if a simplification exists."""
+        op = s_prod(3, s_prod(4, qml.PauliX(0)))
+
+        assert isinstance(op, SProd)
+        assert op.scalar == 3
+        assert isinstance(op.base, SProd)
+
+    def test_non_lazy_mode(self):
+        """Test the lazy=False keyword."""
+        op = s_prod(3, s_prod(4, qml.PauliX(0)), lazy=False)
+
+        assert isinstance(op, SProd)
+        assert op.scalar == 12
+        assert qml.equal(op.base, qml.PauliX(0))
+
+    def test_non_lazy_mode_queueing(self):
+        """Test that if a simpification is accomplished, the metadata for the original op
+        and the new simplified op is updated."""
+        with qml.queuing.AnnotatedQueue() as q:
+            sprod1 = s_prod(4, qml.PauliX(0))
+            sprod2 = s_prod(3, sprod1, lazy=False)
+
+        assert q[sprod1]["owner"] is sprod2
+        assert q[sprod2]["owns"] is sprod1
+
 
 class TestIntegration:
     """Integration tests for SProd with a QNode."""
diff --git a/tests/ops/op_math/test_sum.py b/tests/ops/op_math/test_sum.py
index 51be9193f..c032bf8bc 100644
--- a/tests/ops/op_math/test_sum.py
+++ b/tests/ops/op_math/test_sum.py
@@ -835,6 +835,30 @@ class TestWrapperFunc:
         assert sum_class_op.wires == sum_func_op.wires
         assert sum_class_op.parameters == sum_func_op.parameters
 
+    def test_lazy_mode(self):
+        """Test that by default, the operator is simply wrapped in `Sum`, even if a simplification exists."""
+        op = op_sum(qml.S(0), Sum(qml.S(1), qml.T(1)))
+
+        assert isinstance(op, Sum)
+        assert len(op) == 2
+
+    def test_non_lazy_mode(self):
+        """Test the lazy=False keyword."""
+        op = op_sum(qml.S(0), Sum(qml.S(1), qml.T(1)), lazy=False)
+
+        assert isinstance(op, Sum)
+        assert len(op) == 3
+
+    def test_non_lazy_mode_queueing(self):
+        """Test that if a simpification is accomplished, the metadata for the original op
+        and the new simplified op is updated."""
+        with qml.queuing.AnnotatedQueue() as q:
+            sum1 = op_sum(qml.S(1), qml.T(1))
+            sum2 = op_sum(qml.S(0), sum1, lazy=False)
+
+        assert q[sum1]["owner"] is sum2
+        assert sum1 in q[sum2]["owns"]
+
 
 class TestIntegration:
     """Integration tests for the Sum class."""