Repository: PennyLaneAI/pennylane

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.

Discussion:
codecov[bot]: # [Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3483?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
> Merging [#3483](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3483?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (66aa60c) into [master](https://codecov.io/gh/PennyLaneAI/pennylane/commit/00ff856645cdf7525c822bc8b73be01aacd10491?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (00ff856) will **increase** coverage by `99.71%`.
> The diff coverage is `100.00%`.

```diff
@@             Coverage Diff             @@
##           master    #3483       +/-   ##
===========================================
+ Coverage        0   99.71%   +99.71%     
===========================================
  Files           0      308      +308     
  Lines           0    27179    +27179     
===========================================
+ Hits            0    27102    +27102     
- Misses          0       77       +77     
```


| [Impacted Files](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3483?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | Coverage Δ | |
|---|---|---|
| [pennylane/ops/op\_math/prod.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3483/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL29wcy9vcF9tYXRoL3Byb2QucHk=) | `100.00% <100.00%> (ø)` | |
| [pennylane/ops/op\_math/sprod.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3483/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL29wcy9vcF9tYXRoL3Nwcm9kLnB5) | `100.00% <100.00%> (ø)` | |
| [pennylane/ops/op\_math/sum.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3483/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL29wcy9vcF9tYXRoL3N1bS5weQ==) | `100.00% <100.00%> (ø)` | |
| [pennylane/ops/functions/simplify.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3483/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL29wcy9mdW5jdGlvbnMvc2ltcGxpZnkucHk=) | `100.00% <0.00%> (ø)` | |
| [pennylane/pauli/\_\_init\_\_.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3483/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3BhdWxpL19faW5pdF9fLnB5) | `100.00% <0.00%> (ø)` | |
| [pennylane/transforms/unitary\_to\_rot.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3483/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3RyYW5zZm9ybXMvdW5pdGFyeV90b19yb3QucHk=) | `100.00% <0.00%> (ø)` | |
| [pennylane/ops/qubit/qchem\_ops.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3483/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL29wcy9xdWJpdC9xY2hlbV9vcHMucHk=) | `100.00% <0.00%> (ø)` | |
| [pennylane/fourier/circuit\_spectrum.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3483/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2ZvdXJpZXIvY2lyY3VpdF9zcGVjdHJ1bS5weQ==) | `100.00% <0.00%> (ø)` | |
| [pennylane/templates/tensornetworks/mps.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3483/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3RlbXBsYXRlcy90ZW5zb3JuZXR3b3Jrcy9tcHMucHk=) | `100.00% <0.00%> (ø)` | |
| [pennylane/grouping/group\_observables.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3483/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2dyb3VwaW5nL2dyb3VwX29ic2VydmFibGVzLnB5) | `100.00% <0.00%> (ø)` | |
| ... and [301 more](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3483/diff?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | |

Help us with your feedback. Take ten seconds to tell us [how you rate us](https://about.codecov.io/nps?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI). Have a feature suggestion? [Share it here.](https://app.codecov.io/gh/feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI)

kevinkayyy: Hi all, 

It seems that all checks have passed. However, please still check if the unit tests added are good enough.
kevinkayyy: Hi Christina, queuing is now added to qml.prod. I looked deeper into the corresponding classes and also their parent classes and I hope my changes are written in a coherent style. 

I also added queuing for qml.op_sum and qml.s_prod but I'm not really sure about those, because I saw in the "_queue_category" functions, there are comments saying "# don't queue scalar prods (Sum instances) as they might not be Unitary!" Please let me know if I should remove queuing for these wrappers, or write something for the classes to make all operators unitary.
kevinkayyy: > Exactly! Looks great 👍
> 
> Not too much to comment on. Exactly what I wanted, and I'm glad to finally have this feature.
> 
> Nice solution using `itertools` and adjustment to update the queuing behaviour. Nice clean tests.

Thanks. I'm glad that you like it.

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."""