Repository: PennyLaneAI/pennylane
Pull request #5478 (2024-04-05 16:52:09)

Dunders combine like-op math classes together 
**Context:**

Currently, summing, matmul'ing, and multiplying are fully lazy, wrapping the two terms in operator arithmetic classes regardless of the types being combined.

```
>>> sum_op = qml.X(0) + qml.Y(0) + qml.Z(0)
>>> len(sum_op)
2
>>> sum_op.operands
(X(0) + Y(0), Z(0))
>>> 2*(3*qml.S(0))
2 * 3 * S(wires=[0])
>>> prod_op = qml.X(0) @ qml.Y(1) @ qml.Z(2)
>>> prod_op
(X(0) @ Y(1)) @ Z(2)
```

While we want to avoid full simplification (combining like terms, eliminating zero terms, reordering, rebalancing the tree) that can be expensive and do things the user does not want, basic combination of like-classes via dunders will:

* Make the output better align with intuition and expectation

For the above example, we would now have:

```
>>> len(sum_op)
3
```

* Improve readability and inspection.  Since `op2` no longer needs parentheses to denote the product of products, the repr is easier to read.

```
>>> prod_op
X(0) @ Y(1) @ Z(2)
```

* Less significant departure from pre-existing behaviour.  While we will still no longer perform the expensive, invasive, and usually unnecessary simplification that was occurring with the Legacy Hamiltonian, this better aligns with the previous dunder behavior.  By having less of a difference in the dunder behavior, less things will break in the wider ecosystem.

* Reduce the possibility of `RecursionError`s on large observables

* Reduce the memory footprint of large observables constructed via dunder.

Note that this only changes the behavior of the infix operators.  Those who want to avoid combining a sum into a larger sum can still use `qml.sum`, `qml.s_prod`, and `qml.prod`.

**What this change does not do:**

This PR does not perform any simplification beyond the very basic.  This change will not affect any of the following:

```
>>> 3 * (qml.X(0) + qml.Y(1))
3 * (X(0) + Y(1))
>>> qml.X(0) - qml.X(0)
X(0) + -1 * X(0)
>>> 0 * qml.Y(1)
0 * Y(1)
>>> qml.X(0) @ qml.Y(0)
X(0) @ Y(0)
>>> (qml.X(0) + qml.Y(1))**2
(X(0) + Y(1))**2
```

**Description of the Change:**

Set `lazy=False` on the construction of operator arithmetic classes via `Operator` dunders.

**Benefits:**

* Improve readability and inspection
* Reduced `RecursionError`
* Need for less simplification later in the pipeline
* Reduced memory footprint of large operators constructed via dunders

**Possible Drawbacks:**

Some power users may switch to direct construction with `qml.sum` if they want to create a sum of sums, or `qml.prod` if they want a product of products.

**Related GitHub Issues:**


Qottmann: Perhaps some relevant parts in https://docs.pennylane.ai/en/latest/introduction/compiling_circuits.html#simplifying-operators that need updating
trbromley: Thanks @albi3ro! I'll get back to you with a decision on whether we want to go with this.
trbromley: Hi @albi3ro! Yes we have a :+1: to proceed in this direction.

Patch:
diff --git a/doc/introduction/compiling_circuits.rst b/doc/introduction/compiling_circuits.rst
index b2beaad0b..6bcf46292 100644
--- a/doc/introduction/compiling_circuits.rst
+++ b/doc/introduction/compiling_circuits.rst
@@ -54,7 +54,7 @@ Now lets simplify a nested operator:
 
 >>> sum_op = qml.RX(1, 0) + qml.X(0)
 >>> prod1 = qml.X(0) @ sum_op
->>> nested_op = prod1 @ qml.RX(1, 0)
+>>> nested_op = qml.prod(prod1, qml.RX(1, 0))
 >>> qml.simplify(nested_op)
 (X(0) @ RX(2.0, wires=[0])) + RX(1.0, wires=[0])
 
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 38f1fdea7..5072429d7 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -263,6 +263,10 @@
 
 <h3>Breaking changes 💔</h3>
 
+* Operator dunder methods now combine like-operator arithmetic classes via `lazy=False`. This reduces the chance of `RecursionError` and makes nested
+  operators easier to work with.
+  [(#5478)](https://github.com/PennyLaneAI/pennylane/pull/5478)
+
 * The private functions `_pauli_mult`, `_binary_matrix` and `_get_pauli_map` from the `pauli` module have been removed. The same functionality can be achieved using newer features in the ``pauli`` module.
   [(#5323)](https://github.com/PennyLaneAI/pennylane/pull/5323)
 
diff --git a/pennylane/operation.py b/pennylane/operation.py
index 4a3fe4c72..a266f93eb 100644
--- a/pennylane/operation.py
+++ b/pennylane/operation.py
@@ -1510,11 +1510,15 @@ class Operator(abc.ABC):
     def __add__(self, other):
         """The addition operation of Operator-Operator objects and Operator-scalar."""
         if isinstance(other, Operator):
-            return qml.sum(self, other)
+            return qml.sum(self, other, lazy=False)
         if isinstance(other, TensorLike):
             if qml.math.allequal(other, 0):
                 return self
-            return qml.sum(self, qml.s_prod(scalar=other, operator=qml.Identity(self.wires)))
+            return qml.sum(
+                self,
+                qml.s_prod(scalar=other, operator=qml.Identity(self.wires), lazy=False),
+                lazy=False,
+            )
         return NotImplemented
 
     __radd__ = __add__
@@ -1524,7 +1528,7 @@ class Operator(abc.ABC):
         if callable(other):
             return qml.pulse.ParametrizedHamiltonian([other], [self])
         if isinstance(other, TensorLike):
-            return qml.s_prod(scalar=other, operator=self)
+            return qml.s_prod(scalar=other, operator=self, lazy=False)
         return NotImplemented
 
     def __truediv__(self, other):
@@ -1537,12 +1541,12 @@ class Operator(abc.ABC):
 
     def __matmul__(self, other):
         """The product operation between Operator objects."""
-        return qml.prod(self, other) if isinstance(other, Operator) else NotImplemented
+        return qml.prod(self, other, lazy=False) if isinstance(other, Operator) else NotImplemented
 
     def __sub__(self, other):
         """The subtraction operation of Operator-Operator objects and Operator-scalar."""
         if isinstance(other, Operator):
-            return self + qml.s_prod(-1, other)
+            return self + qml.s_prod(-1, other, lazy=False)
         if isinstance(other, TensorLike):
             return self + (qml.math.multiply(-1, other))
         return NotImplemented
@@ -1553,7 +1557,7 @@ class Operator(abc.ABC):
 
     def __neg__(self):
         """The negation operation of an Operator object."""
-        return qml.s_prod(scalar=-1, operator=self)
+        return qml.s_prod(scalar=-1, operator=self, lazy=False)
 
     def __pow__(self, other):
         r"""The power operation of an Operator object."""
diff --git a/pennylane/pulse/parametrized_hamiltonian.py b/pennylane/pulse/parametrized_hamiltonian.py
index a7ce9bda7..f241ac38e 100644
--- a/pennylane/pulse/parametrized_hamiltonian.py
+++ b/pennylane/pulse/parametrized_hamiltonian.py
@@ -253,7 +253,13 @@ class ParametrizedHamiltonian:
                 f"must be the same. Received len(params) = {len(params)} parameters but "
                 f"expected {len(self.coeffs_parametrized)} parameters."
             )
-        return self.H_fixed() + self.H_parametrized(params, t)
+        H_fixed = self.H_fixed()
+        H_param = self.H_parametrized(params, t)
+        if H_param == 0:
+            return H_fixed
+        if H_fixed == 0:
+            return H_param
+        return qml.sum(self.H_fixed(), self.H_parametrized(params, t))
 
     def __repr__(self):
         terms = []


Test patch:
diff --git a/tests/drawer/test_tape_text.py b/tests/drawer/test_tape_text.py
index 806c8594b..007381692 100644
--- a/tests/drawer/test_tape_text.py
+++ b/tests/drawer/test_tape_text.py
@@ -584,7 +584,7 @@ single_op_tests_data = [
         qml.expval(
             0.1 * qml.PauliX(0) + 0.2 * qml.PauliY(1) + 0.3 * qml.PauliZ(0) + 0.4 * qml.PauliZ(1)
         ),
-        "0: ───┤ ╭<(((0.10*X)+(0.20*Y))+(0.30*Z))+(0.40*Z)>\n1: ───┤ ╰<(((0.10*X)+(0.20*Y))+(0.30*Z))+(0.40*Z)>",
+        "0: ───┤ ╭<(0.10*X)+(0.20*Y)+(0.30*Z)+(0.40*Z)>\n1: ───┤ ╰<(0.10*X)+(0.20*Y)+(0.30*Z)+(0.40*Z)>",
     ),
     # Operations (both regular and controlled) and nested multi-valued controls
     (qml.ctrl(qml.PauliX(wires=2), control=[0, 1]), "0: ─╭●─┤  \n1: ─├●─┤  \n2: ─╰X─┤  "),
diff --git a/tests/ops/op_math/test_prod.py b/tests/ops/op_math/test_prod.py
index d1d9b861f..d35823bd7 100644
--- a/tests/ops/op_math/test_prod.py
+++ b/tests/ops/op_math/test_prod.py
@@ -1081,9 +1081,9 @@ class TestSimplify:
 
     def test_depth_property(self):
         """Test depth property."""
-        prod_op = (
-            qml.RZ(1.32, wires=0) @ qml.Identity(wires=0) @ qml.RX(1.9, wires=1) @ qml.PauliX(0)
-        )
+        o1 = qml.prod(qml.RZ(1.32, wires=0), qml.I(wires=0))
+        o2 = qml.prod(o1, qml.RX(1.9, wires=1))
+        prod_op = qml.prod(o2, qml.X(0))
         assert prod_op.arithmetic_depth == 3
 
         op_constructed = Prod(
diff --git a/tests/ops/op_math/test_sum.py b/tests/ops/op_math/test_sum.py
index 69e5e723c..44ebd22d0 100644
--- a/tests/ops/op_math/test_sum.py
+++ b/tests/ops/op_math/test_sum.py
@@ -309,7 +309,7 @@ class TestInitialization:
         ),
         (
             0.5 * (X(0) @ (0.5 * X(1))) + 0.7 * X(1) + 0.8 * (X(0) @ Y(1) @ Z(1)),
-            "(\n    0.5 * (X(0) @ (0.5 * X(1)))\n  + 0.7 * X(1)\n  + 0.8 * ((X(0) @ Y(1)) @ Z(1))\n)",
+            "(\n    0.5 * (X(0) @ (0.5 * X(1)))\n  + 0.7 * X(1)\n  + 0.8 * (X(0) @ Y(1) @ Z(1))\n)",
         ),
     )
 
@@ -870,9 +870,11 @@ class TestSimplify:
             qml.RX(1.9, wires=1),
             qml.PauliX(0),
         )
-        dunder_sum_op = sum(ops_to_sum)
+        s1 = qml.sum(ops_to_sum[0], ops_to_sum[1])
+        s2 = qml.sum(s1, ops_to_sum[2])
+        nested_sum = qml.sum(s2, ops_to_sum[3])
         class_sum_op = Sum(*ops_to_sum)
-        assert dunder_sum_op.arithmetic_depth == 3
+        assert nested_sum.arithmetic_depth == 3
         assert class_sum_op.arithmetic_depth == 1
 
     def test_simplify_method(self):
diff --git a/tests/test_operation.py b/tests/test_operation.py
index b473f4c5b..f3046cd56 100644
--- a/tests/test_operation.py
+++ b/tests/test_operation.py
@@ -2694,16 +2694,15 @@ class TestNewOpMath:
             assert op[1].scalar == 1.1
             assert qml.equal(op[1].base, qml.Identity(1))
 
-        def test_adding_many_does_not_auto_simplify(self):
+        def test_adding_many_does_auto_simplify(self):
             """Tests that adding more than two operators creates nested Sums."""
             op0, op1, op2 = qml.S(0), qml.T(0), qml.PauliZ(0)
             op = op0 + op1 + op2
             assert isinstance(op, Sum)
-            assert len(op) == 2
-            assert isinstance(op[0], Sum)
-            assert qml.equal(op[0][0], op0)
-            assert qml.equal(op[0][1], op1)
-            assert qml.equal(op[1], op2)
+            assert len(op) == 3
+            assert qml.equal(op[0], op0)
+            assert qml.equal(op[1], op1)
+            assert qml.equal(op[2], op2)
 
     class TestMul:
         """Test the __mul__/__rmul__ dunders."""
@@ -2726,13 +2725,13 @@ class TestNewOpMath:
             assert qml.math.allequal(op.scalar, 1 / scalar)
             assert qml.equal(op.base, base)
 
-        def test_mul_does_not_auto_simplify(self):
+        def test_mul_does_auto_simplify(self):
             """Tests that multiplying an SProd with a scalar creates nested SProds."""
             op = 2 * qml.PauliX(0)
             nested = 0.5 * op
             assert isinstance(nested, SProd)
-            assert nested.scalar == 0.5
-            assert qml.equal(nested.base, op)
+            assert nested.scalar == 1.0
+            assert qml.equal(nested.base, qml.X(0))
 
     class TestMatMul:
         """Test the __matmul__/__rmatmul__ dunders."""
@@ -2743,17 +2742,21 @@ class TestNewOpMath:
             op = op0 @ op1
             assert isinstance(op, Prod)
             assert qml.equal(op[0], op0)
-            assert qml.equal(op[1], op1)
+            if isinstance(op1, Prod):
+                assert qml.equal(op[1], op1[0])
+                assert qml.equal(op[2], op1[1])
+            else:
+                assert qml.equal(op[1], op1)
 
-        def test_mul_does_not_auto_simplify(self):
+        def test_mul_does_auto_simplify(self):
             """Tests that matrix-multiplying a Prod with another operator creates nested Prods."""
             op0, op1, op2 = qml.PauliX(0), qml.PauliY(1), qml.PauliZ(2)
             op = op0 @ op1 @ op2
             assert isinstance(op, Prod)
-            assert len(op) == 2
-            assert isinstance(op[0], Prod)
-            assert qml.equal(op[0], op0 @ op1)
-            assert qml.equal(op[1], op2)
+            assert len(op) == 3
+            assert qml.equal(op[0], op0)
+            assert qml.equal(op[1], op1)
+            assert qml.equal(op[2], op2)
 
     class TestHamiltonianLinearCombinationAlias:
         """Unit tests for using qml.Hamiltonian as an alias for LinearCombination"""