Repository: PennyLaneAI/pennylane
Pull request #5524 (2024-04-16 14:40:15)

Minor updates for PR #5511
**Context:**

When testing for PR #5511, I found a couple necessary fixes to `Sum` and `LinearCombination`, and a bug in the queuing behavior of `TrotterProduct`. I'm extracting this fixes out into a separate PR for ease of review.

With program capture, everything we do needs to be purely functional, so we cannot set grouping indices after construction the operation.  So I am just adding the option to set them during initialization.

**Description of the Change:**

* Adds `grouping_indices` to the call signatures of `Sum` and `LinearCombination`. I'm ignoring `Hamiltonian` because it's now deprecated.

* Updates `LinearCombination._flatten` so that we aren't duplicating leaves. This could have led to potential issues when repacking.

* Updates the queuing behavior of `TrotterProduct` so it de-queues the input observable.


**Benefits:**

**Possible Drawbacks:**

**Related GitHub Issues:**


codecov[bot]: ## [Codecov](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5524?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Project coverage is 99.66%. Comparing base [(`0747c0a`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/0747c0ab22a277b1a3356656af33c9b6c2ce2609?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) to head [(`5dcbda5`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5524?dropdown=coverage&src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).
> Report is 1 commits behind head on master.


<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##           master    #5524      +/-   ##
==========================================
- Coverage   99.67%   99.66%   -0.01%     
==========================================
  Files         404      406       +2     
  Lines       37852    37582     -270     
==========================================
- Hits        37728    37457     -271     
- Misses        124      125       +1     
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5524?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).

dwierichs: @albi3ro While writing #5523 I noticed that `ApproxTimeEvolution`, `CommutingEvolution`, `QDrift` also do not queue correctly.

trbromley: @albi3ro do you think this should be part of the v0.36 milestone?
albi3ro: > @albi3ro do you think this should be part of the v0.36 milestone?

Don't think it necessarily matters too much, but couldn't hurt.

Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index aeeeedf42..01be34366 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -298,6 +298,9 @@
 * `qml.transforms.split_non_commuting` will now work with single-term operator arithmetic.
   [(#5314)](https://github.com/PennyLaneAI/pennylane/pull/5314)
 
+* `LinearCombination` and `Sum` now accept `_grouping_indices` on initialization.
+  [(#5524)](https://github.com/PennyLaneAI/pennylane/pull/5524)
+
 <h4>Mid-circuit measurements and dynamic circuits</h4>
 
 * The `QubitDevice` class and children classes support the `dynamic_one_shot` transform provided that they support `MidMeasureMP` operations natively.
@@ -465,6 +468,10 @@
 
 <h3>Bug fixes 🐛</h3>
 
+* `ApproxTimeEvolution`, `CommutingEvolution`, `QDrift`, and `TrotterProduct` 
+  now de-queue their input observable.
+  [(#5524)](https://github.com/PennyLaneAI/pennylane/pull/5524)
+
 * (In)equality of `qml.HilbertSchmidt` instances is now reported correctly by `qml.equal`.
   [(#5538)](https://github.com/PennyLaneAI/pennylane/pull/5538)
 
diff --git a/pennylane/ops/op_math/linear_combination.py b/pennylane/ops/op_math/linear_combination.py
index 60c604b44..e4d943f42 100644
--- a/pennylane/ops/op_math/linear_combination.py
+++ b/pennylane/ops/op_math/linear_combination.py
@@ -100,14 +100,11 @@ class LinearCombination(Sum):
 
     def _flatten(self):
         # note that we are unable to restore grouping type or method without creating new properties
-        return (self._coeffs, self._ops, self.data), (self.grouping_indices,)
+        return self.terms(), (self.grouping_indices,)
 
     @classmethod
     def _unflatten(cls, data, metadata):
-        new_op = cls(data[0], data[1])
-        new_op._grouping_indices = metadata[0]  # pylint: disable=protected-access
-        new_op.data = data[2]
-        return new_op
+        return cls(data[0], data[1], _grouping_indices=metadata[0])
 
     def __init__(
         self,
@@ -116,6 +113,7 @@ class LinearCombination(Sum):
         simplify=False,
         grouping_type=None,
         method="rlf",
+        _grouping_indices=None,
         _pauli_rep=None,
         id=None,
     ):
@@ -147,7 +145,12 @@ class LinearCombination(Sum):
             operands = tuple(qml.s_prod(c, op) for c, op in zip(coeffs, observables))
 
         super().__init__(
-            *operands, grouping_type=grouping_type, method=method, id=id, _pauli_rep=_pauli_rep
+            *operands,
+            grouping_type=grouping_type,
+            method=method,
+            id=id,
+            _grouping_indices=_grouping_indices,
+            _pauli_rep=_pauli_rep,
         )
 
     @staticmethod
diff --git a/pennylane/ops/op_math/sum.py b/pennylane/ops/op_math/sum.py
index cb9a580f9..725540bf6 100644
--- a/pennylane/ops/op_math/sum.py
+++ b/pennylane/ops/op_math/sum.py
@@ -216,17 +216,24 @@ class Sum(CompositeOp):
 
     @classmethod
     def _unflatten(cls, data, metadata):
-        # pylint: disable=protected-access
-        new_op = cls(*data)
-        new_op._grouping_indices = metadata[0]
-        return new_op
+        return cls(*data, _grouping_indices=metadata[0])
 
     def __init__(
-        self, *operands: Operator, grouping_type=None, method="rlf", id=None, _pauli_rep=None
+        self,
+        *operands: Operator,
+        grouping_type=None,
+        method="rlf",
+        id=None,
+        _grouping_indices=None,
+        _pauli_rep=None,
     ):
         super().__init__(*operands, id=id, _pauli_rep=_pauli_rep)
 
-        self._grouping_indices = None
+        self._grouping_indices = _grouping_indices
+        if _grouping_indices is not None and grouping_type is not None:
+            raise ValueError(
+                "_grouping_indices and grouping_type cannot be specified at the same time."
+            )
         if grouping_type is not None:
             self.compute_grouping(grouping_type=grouping_type, method=method)
 
diff --git a/pennylane/templates/subroutines/approx_time_evolution.py b/pennylane/templates/subroutines/approx_time_evolution.py
index 694821865..c6d0a06c7 100644
--- a/pennylane/templates/subroutines/approx_time_evolution.py
+++ b/pennylane/templates/subroutines/approx_time_evolution.py
@@ -139,6 +139,11 @@ class ApproxTimeEvolution(Operation):
         # trainable parameters are passed to the base init method
         super().__init__(*hamiltonian.data, time, wires=wires, id=id)
 
+    def queue(self, context=qml.QueuingManager):
+        context.remove(self.hyperparameters["hamiltonian"])
+        context.append(self)
+        return self
+
     @staticmethod
     def compute_decomposition(
         *coeffs_and_time, wires, hamiltonian, n
diff --git a/pennylane/templates/subroutines/commuting_evolution.py b/pennylane/templates/subroutines/commuting_evolution.py
index e10252e91..547dc2158 100644
--- a/pennylane/templates/subroutines/commuting_evolution.py
+++ b/pennylane/templates/subroutines/commuting_evolution.py
@@ -141,6 +141,11 @@ class CommutingEvolution(Operation):
 
         super().__init__(time, *hamiltonian.parameters, wires=hamiltonian.wires, id=id)
 
+    def queue(self, context=qml.QueuingManager):
+        context.remove(self.hyperparameters["hamiltonian"])
+        context.append(self)
+        return self
+
     @staticmethod
     def compute_decomposition(
         time, *_, wires, hamiltonian, **__
diff --git a/pennylane/templates/subroutines/qdrift.py b/pennylane/templates/subroutines/qdrift.py
index 35299e3a7..989e7a110 100644
--- a/pennylane/templates/subroutines/qdrift.py
+++ b/pennylane/templates/subroutines/qdrift.py
@@ -190,6 +190,11 @@ class QDrift(Operation):
         }
         super().__init__(time, wires=hamiltonian.wires, id=id)
 
+    def queue(self, context=qml.QueuingManager):
+        context.remove(self.hyperparameters["base"])
+        context.append(self)
+        return self
+
     @classmethod
     def _unflatten(cls, data, metadata):
         """Recreate an operation from its serialized format.
diff --git a/pennylane/templates/subroutines/trotter.py b/pennylane/templates/subroutines/trotter.py
index 52a56a0c0..68946909e 100644
--- a/pennylane/templates/subroutines/trotter.py
+++ b/pennylane/templates/subroutines/trotter.py
@@ -218,6 +218,11 @@ class TrotterProduct(ErrorOperation):
         }
         super().__init__(time, wires=hamiltonian.wires, id=id)
 
+    def queue(self, context=qml.QueuingManager):
+        context.remove(self.hyperparameters["base"])
+        context.append(self)
+        return self
+
     def error(
         self, method: str = "commutator", fast: bool = True
     ):  # pylint: disable=arguments-differ


Test patch:
diff --git a/tests/ops/op_math/test_linear_combination.py b/tests/ops/op_math/test_linear_combination.py
index 52970279b..4436a7121 100644
--- a/tests/ops/op_math/test_linear_combination.py
+++ b/tests/ops/op_math/test_linear_combination.py
@@ -618,12 +618,11 @@ class TestLinearCombination:
         data, metadata = H._flatten()
         assert metadata[0] == H.grouping_indices
         assert hash(metadata)
-        assert len(data) == 3
+        assert len(data) == 2
         assert qml.math.allequal(
             data[0], H._coeffs
         )  # Previously checking "is" instead of "==", problem?
         assert data[1] == H._ops
-        assert data[2] == H.data
 
         new_H = LinearCombination._unflatten(*H._flatten())
         assert qml.equal(H, new_H)
@@ -1457,6 +1456,12 @@ class TestLinearCombinationArithmeticJax:
 class TestGrouping:
     """Tests for the grouping functionality"""
 
+    def test_set_on_initialization(self):
+        """Test that grouping indices can be set on initialization."""
+
+        op = qml.ops.LinearCombination([1, 1], [qml.X(0), qml.Y(1)], _grouping_indices=[[0, 1]])
+        assert op.grouping_indices == [[0, 1]]
+
     def test_indentities_preserved(self):
         """Tests that the grouping indices do not drop identity terms when the wire order is nonstandard."""
 
diff --git a/tests/ops/op_math/test_sum.py b/tests/ops/op_math/test_sum.py
index be266eab9..3707dc544 100644
--- a/tests/ops/op_math/test_sum.py
+++ b/tests/ops/op_math/test_sum.py
@@ -1326,6 +1326,18 @@ class TestArithmetic:
 class TestGrouping:
     """Test grouping functionality of Sum"""
 
+    def test_set_on_initialization(self):
+        """Test that grouping indices can be set on initialization."""
+
+        op = qml.ops.Sum(qml.X(0), qml.Y(1), _grouping_indices=[[0, 1]])
+        assert op.grouping_indices == [[0, 1]]
+        op_ac = qml.ops.Sum(qml.X(0), qml.Y(1), grouping_type="anticommuting")
+        assert op_ac.grouping_indices == ((0,), (1,))
+        with pytest.raises(ValueError, match=r"cannot be specified at the same time."):
+            qml.ops.Sum(
+                qml.X(0), qml.Y(1), grouping_type="anticommuting", _grouping_indices=[[0, 1]]
+            )
+
     def test_non_pauli_error(self):
         """Test that grouping non-Pauli observables is not supported."""
         op = Sum(qml.PauliX(0), Prod(qml.PauliZ(0), qml.PauliX(1)), qml.Hadamard(2))
diff --git a/tests/templates/test_subroutines/test_approx_time_evolution.py b/tests/templates/test_subroutines/test_approx_time_evolution.py
index 5924d3e75..d61b91b4b 100644
--- a/tests/templates/test_subroutines/test_approx_time_evolution.py
+++ b/tests/templates/test_subroutines/test_approx_time_evolution.py
@@ -47,6 +47,17 @@ def test_flatten_unflatten():
     assert new_op is not op
 
 
+def test_queuing():
+    """Test that ApproxTimeEvolution de-queues the input hamiltonian."""
+
+    with qml.queuing.AnnotatedQueue() as q:
+        H = qml.X(0) + qml.Y(1)
+        op = qml.ApproxTimeEvolution(H, 0.1, n=20)
+
+    assert len(q.queue) == 1
+    assert q.queue[0] is op
+
+
 class TestDecomposition:
     """Tests that the template defines the correct decomposition."""
 
diff --git a/tests/templates/test_subroutines/test_commuting_evolution.py b/tests/templates/test_subroutines/test_commuting_evolution.py
index 60124cf55..5de57ea4c 100644
--- a/tests/templates/test_subroutines/test_commuting_evolution.py
+++ b/tests/templates/test_subroutines/test_commuting_evolution.py
@@ -85,6 +85,17 @@ def test_adjoint():
     assert all(np.isclose(state1, state2))
 
 
+def test_queuing():
+    """Test that CommutingEvolution de-queues the input hamiltonian."""
+
+    with qml.queuing.AnnotatedQueue() as q:
+        H = qml.X(0) + qml.Y(1)
+        op = qml.CommutingEvolution(H, 0.1, (2,))
+
+    assert len(q.queue) == 1
+    assert q.queue[0] is op
+
+
 def test_decomposition_expand():
     """Test that the decomposition of CommutingEvolution is an ApproxTimeEvolution with one step."""
 
diff --git a/tests/templates/test_subroutines/test_qdrift.py b/tests/templates/test_subroutines/test_qdrift.py
index 4e7b5a808..ff47bc46f 100644
--- a/tests/templates/test_subroutines/test_qdrift.py
+++ b/tests/templates/test_subroutines/test_qdrift.py
@@ -44,6 +44,16 @@ test_hamiltonians = (
 class TestInitialization:
     """Test that the class is intialized correctly."""
 
+    def test_queuing(self):
+        """Test that QDrift de-queues the input hamiltonian."""
+
+        with qml.queuing.AnnotatedQueue() as q:
+            H = qml.X(0) + qml.Y(1)
+            op = qml.QDrift(H, 0.1, n=20)
+
+        assert len(q.queue) == 1
+        assert q.queue[0] is op
+
     @pytest.mark.parametrize("n", (1, 2, 3))
     @pytest.mark.parametrize("time", (0.5, 1, 2))
     @pytest.mark.parametrize("seed", (None, 1234, 42))
diff --git a/tests/templates/test_subroutines/test_trotter.py b/tests/templates/test_subroutines/test_trotter.py
index a6dea66db..88f87801a 100644
--- a/tests/templates/test_subroutines/test_trotter.py
+++ b/tests/templates/test_subroutines/test_trotter.py
@@ -399,6 +399,16 @@ class TestInitialization:
             qml.matrix(op2, wire_order=hamiltonian.wires),
         )
 
+    def test_queuing(self):
+        """Test that the target operator is removed from the queue."""
+
+        with qml.queuing.AnnotatedQueue() as q:
+            H = qml.X(0) + qml.Y(1)
+            op = qml.TrotterProduct(H, time=2)
+
+        assert len(q.queue) == 1
+        assert q.queue[0] is op
+
 
 class TestPrivateFunctions:
     """Test the private helper functions."""