Repository: PennyLaneAI/pennylane
Pull request #7387 (2025-05-06 19:34:30)

Allow returning diagonalized parametric MCMs
**Context:**

When we diagonalize the parametric mid-circuit measurements to get analytic results to validate our protocols, we swap them out with new, updated measurements. This is currently updated in the conditionals, but not on the measurements. That means that returning parametric mcms when using the `diagonalize_mcms` transform currently fails.

**Description of the Change:**

We also update the measurements in the transform - if any MeasurementProcess contains MCMs, we update them based on the existing dictionary that's tracking updated measurement values, and create a copy of the MP with the new MeasurementValue(s) instead of the original ones.

**Benefits:**

We can do this:

```
@diagonalize_mcms
@qml.qnode(dev)
def circ():
    qml.H(0)
    m = measure_x(0, reset=True)
    return qml.expval(m)
```

[sc-90609]


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/7387?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.65%. Comparing base [(`353fb36`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/353fb36a73bb6e9c0df10dec3a5afb20e0d85424?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) to head [(`9b720b1`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/9b720b1c09cc2554caba5ba8f0e8f6c32a0235c4?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    #7387   +/-   ##
=======================================
  Coverage   99.65%   99.65%           
=======================================
  Files         529      529           
  Lines       50603    50612    +9     
=======================================
+ Hits        50427    50436    +9     
  Misses        176      176           
```

</details>

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

<details><summary> :rocket: New features to boost your workflow: </summary>

- :snowflake: [Test Analytics](https://docs.codecov.com/docs/test-analytics): Detect flaky tests, report on failures, and find test suite problems.
</details>

Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index e25705b80..47c5aa932 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -327,6 +327,11 @@ Here's a list of deprecations made this release. For a more detailed breakdown o
 * Fixes a bug where the powers of `qml.ISWAP` and `qml.SISWAP` were decomposed incorrectly.
   [(#7361)](https://github.com/PennyLaneAI/pennylane/pull/7361)
 
+* Returning `MeasurementValue`s from the `ftqc` module's parametric mid-circuit measurements
+  (`measure_arbitrary_basis`, `measure_x` and `measure_y`) no longer raises an error in circuits 
+  using `diagonalize_mcms`.
+  [(#7387)](https://github.com/PennyLaneAI/pennylane/pull/7387)
+
 <h3>Contributors ✍️</h3>
 
 This release contains contributions from (in alphabetical order):
diff --git a/pennylane/ftqc/parametric_midmeasure.py b/pennylane/ftqc/parametric_midmeasure.py
index d5b8af297..cee5ecc99 100644
--- a/pennylane/ftqc/parametric_midmeasure.py
+++ b/pennylane/ftqc/parametric_midmeasure.py
@@ -17,6 +17,7 @@ mid-circuit measurements with a parameterized measurement axis."""
 
 import uuid
 from collections.abc import Hashable
+from copy import copy
 from functools import lru_cache
 from typing import Iterable, Optional, Union
 
@@ -784,6 +785,16 @@ def diagonalize_mcms(tape):
 
         curr_idx += 1
 
+    new_measurements = []
+    for mp in tape.measurements:
+        if mp.mv is None:
+            new_measurements.append(mp)
+        else:
+            new_mp = copy(mp)
+            mps = [mps_mapping.get(m, m) for m in mp.mv.measurements]
+            new_mp.mv.measurements = mps
+            new_measurements.append(new_mp)
+
     new_tape = tape.copy(operations=new_operations)
 
     return (new_tape,), null_postprocessing


Test patch:
diff --git a/tests/ftqc/test_parametric_mid_measure.py b/tests/ftqc/test_parametric_mid_measure.py
index 488d8a2a6..3dc222b64 100644
--- a/tests/ftqc/test_parametric_mid_measure.py
+++ b/tests/ftqc/test_parametric_mid_measure.py
@@ -916,3 +916,41 @@ class TestWorkflows:
             assert np.allclose(circ(), [np.cos(2.345), -1], atol=0.03)
         else:
             assert np.allclose(circ(), [np.cos(2.345), -1])
+
+    @pytest.mark.parametrize("mcm_method, shots", [("tree-traversal", None), ("one-shot", 10000)])
+    def test_diagonalize_mcms_returns_parametrized_mcms(self, mcm_method, shots):
+        """Test that when diagonalizing, parametrized mid-circuit measurements can be returned
+        by the QNode"""
+
+        dev = qml.device("default.qubit", shots=shots)
+
+        @diagonalize_mcms
+        @qml.qnode(dev, mcm_method=mcm_method)
+        def circ():
+            m0 = measure_x(0)
+            m1 = measure_y(1)
+            m2 = measure_arbitrary_basis(2, angle=1.23, plane="XY")
+
+            return qml.expval(m0), qml.expval(m1), qml.expval(m2)
+
+        circ()
+
+    @pytest.mark.parametrize("mcm_method, shots", [("tree-traversal", None), ("one-shot", 10000)])
+    def test_diagonalize_mcms_returns_cond_measure_result(self, mcm_method, shots):
+        """Test that when diagonalizing, the MeasurementValue output by cond_measure can be returned
+        by the QNode"""
+
+        if mcm_method == "one-shot":
+            pytest.xfail(reason="not implemented yet")  # sc-90607
+
+        dev = qml.device("default.qubit", shots=shots)
+
+        @diagonalize_mcms
+        @qml.qnode(dev, mcm_method=mcm_method)
+        def circ():
+            qml.H(0)
+            m0 = measure_x(0)
+            m1 = cond_measure(m0, measure_x, measure_y)(1)
+            return qml.expval(m1)
+
+        circ()