Repository: PennyLaneAI/pennylane
Pull request #5987 (2024-07-11 15:21:54)

clean `unflatten` in `PrepSelPrep` 

**Context:**

LinearCombination was being cast within unflaten. It has been moved to init to avoid issues



github-actions[bot]: Hello. You may have forgotten to update the changelog!
Please edit [`doc/releases/changelog-dev.md`](/PennyLaneAI/pennylane/blob/master/doc/releases/changelog-dev.md) with:
* A one-to-two sentence description of the change. You may include a small working example for new features.
* A link back to this PR.
* Your name (or GitHub username) in the contributors section.
codecov[bot]: ## [Codecov](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5987?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.53%. Comparing base [(`6d0a528`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/6d0a528f8797cc60422d48e00ad0168ca14c5c3f?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) to head [(`bd62d2f`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/bd62d2f896c17cbc212c2082d3ae11c3264eef50?dropdown=coverage&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    #5987      +/-   ##
==========================================
- Coverage   99.67%   99.53%   -0.14%     
==========================================
  Files         426      426              
  Lines       40924    40808     -116     
==========================================
- Hits        40791    40620     -171     
- Misses        133      188      +55     
```

</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5987?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).


Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 8559de76b..ec65e17c2 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -14,6 +14,8 @@
   combination of unitaries.
   [(#5756)](https://github.com/PennyLaneAI/pennylane/pull/5756)
 
+[(#5987)](https://github.com/PennyLaneAI/pennylane/pull/5987)
+
 * The `split_to_single_terms` transform is added. This transform splits expectation values of sums
   into multiple single-term measurements on a single tape, providing better support for simulators
   that can handle non-commuting observables but don't natively support multi-term observables.
diff --git a/pennylane/templates/subroutines/prepselprep.py b/pennylane/templates/subroutines/prepselprep.py
index f6c02265a..53df7f96c 100644
--- a/pennylane/templates/subroutines/prepselprep.py
+++ b/pennylane/templates/subroutines/prepselprep.py
@@ -75,9 +75,10 @@ class PrepSelPrep(Operation):
     """
 
     def __init__(self, lcu, control=None, id=None):
+
         coeffs, ops = lcu.terms()
         control = qml.wires.Wires(control)
-        self.hyperparameters["lcu"] = lcu
+        self.hyperparameters["lcu"] = qml.ops.LinearCombination(coeffs, ops)
         self.hyperparameters["coeffs"] = coeffs
         self.hyperparameters["ops"] = ops
         self.hyperparameters["control"] = control
@@ -95,14 +96,11 @@ class PrepSelPrep(Operation):
         super().__init__(*self.data, wires=all_wires, id=id)
 
     def _flatten(self):
-        return tuple(self.lcu), (self.control)
+        return (self.lcu,), (self.control,)
 
     @classmethod
     def _unflatten(cls, data, metadata) -> "PrepSelPrep":
-        coeffs = [term.terms()[0][0] for term in data]
-        ops = [term.terms()[1][0] for term in data]
-        lcu = qml.ops.LinearCombination(coeffs, ops)
-        return cls(lcu, metadata)
+        return cls(data[0], metadata[0])
 
     def __repr__(self):
         return f"PrepSelPrep(coeffs={tuple(self.coeffs)}, ops={tuple(self.ops)}, control={self.control})"


Test patch:
diff --git a/tests/templates/test_subroutines/test_prepselprep.py b/tests/templates/test_subroutines/test_prepselprep.py
index 587bbdccb..826299738 100644
--- a/tests/templates/test_subroutines/test_prepselprep.py
+++ b/tests/templates/test_subroutines/test_prepselprep.py
@@ -27,6 +27,10 @@ import pennylane as qml
     ("lcu", "control"),
     [
         (qml.ops.LinearCombination([0.25, 0.75], [qml.Z(2), qml.X(1) @ qml.X(2)]), [0]),
+        (qml.dot([0.25, 0.75], [qml.Z(2), qml.X(1) @ qml.X(2)]), [0]),
+        (qml.Hamiltonian([0.25, 0.75], [qml.Z(2), qml.X(1) @ qml.X(2)]), [0]),
+        (0.25 * qml.Z(2) - 0.75 * qml.X(1) @ qml.X(2), [0]),
+        (qml.Z(2) + qml.X(1) @ qml.X(2), [0]),
         (qml.ops.LinearCombination([-0.25, 0.75j], [qml.Z(3), qml.X(2) @ qml.X(3)]), [0, 1]),
         (
             qml.ops.LinearCombination([-0.25 + 0.1j, 0.75j], [qml.Z(4), qml.X(4) @ qml.X(5)]),
@@ -257,25 +261,35 @@ class TestPrepSelPrep:
 
         assert qml.equal(op, op_copy)
 
-    def test_flatten_unflatten(self):
+    @pytest.mark.parametrize(
+        ("lcu"),
+        [
+            qml.ops.LinearCombination([0.25, 0.75], [qml.Z(2), qml.X(1) @ qml.X(2)]),
+            qml.dot([0.25, 0.75], [qml.Z(2), qml.X(1) @ qml.X(2)]),
+            qml.Hamiltonian([0.25, 0.75], [qml.Z(2), qml.X(1) @ qml.X(2)]),
+            0.25 * qml.Z(2) - 0.75 * qml.X(1) @ qml.X(2),
+            qml.Z(2) + qml.X(1) @ qml.X(2),
+            qml.ops.LinearCombination([-0.25, 0.75j], [qml.Z(3), qml.X(2) @ qml.X(3)]),
+            qml.ops.LinearCombination([-0.25 + 0.1j, 0.75j], [qml.Z(4), qml.X(4) @ qml.X(5)]),
+        ],
+    )
+    def test_flatten_unflatten(self, lcu):
         """Test that the class can be correctly flattened and unflattened"""
 
-        lcu = qml.ops.LinearCombination([1 / 2, 1 / 2], [qml.Identity(1), qml.PauliZ(1)])
         lcu_coeffs, lcu_ops = lcu.terms()
 
         op = qml.PrepSelPrep(lcu, control=0)
         data, metadata = op._flatten()
 
-        data_coeffs = [term.terms()[0][0] for term in data]
-        data_ops = [term.terms()[1][0] for term in data]
+        data_coeffs, data_ops = data[0].terms()
 
         assert hash(metadata)
 
-        assert len(data) == len(lcu)
+        assert len(data[0]) == len(lcu)
         assert all(coeff1 == coeff2 for coeff1, coeff2 in zip(lcu_coeffs, data_coeffs))
         assert all(op1 == op2 for op1, op2 in zip(lcu_ops, data_ops))
 
-        assert metadata == op.control
+        assert metadata[0] == op.control
 
         new_op = type(op)._unflatten(*op._flatten())
         assert op.lcu == new_op.lcu