Repository: PennyLaneAI/pennylane
Pull request #6903 (2025-01-29 21:44:59)

[Bug fix] Use qnode mcm_method information when creating transform program for construct batch
**Context:**
A bug was found in `construct_tape` (and `construct_batch`) while testing behaviour of circuits with mid-circuit measurements for another task: the QNode's `mcm_method` information is never passed to the `construct_batch` function, so it always constructs a tape as if deferred measurements were used.

**Description of the Change:**
Create an `mcm_config` from the information stored in `qnode.execution_kwargs` and use it when getting the transform program to apply.

**Related GitHub Issues:**
#6902


[sc-83393]

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/6903?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.59%. Comparing base [(`58d4f4f`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/58d4f4fb4893e1d26e4797a9ff489cf98f474510?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) to head [(`b739b12`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/b739b1264f9646a9426ee3d5be3c604900e50061?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).

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


```diff
@@           Coverage Diff           @@
##           master    #6903   +/-   ##
=======================================
  Coverage   99.59%   99.59%           
=======================================
  Files         478      478           
  Lines       45295    45296    +1     
=======================================
+ Hits        45113    45114    +1     
  Misses        182      182           
```

</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/6903?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 64eb1cb86..e9d91a9b2 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -156,6 +156,10 @@
 * `BasisState` now casts its input to integers.
   [(#6844)](https://github.com/PennyLaneAI/pennylane/pull/6844)
 
+* The `workflow.contstruct_batch` and `workflow.construct_tape` functions now correctly reflect the `mcm_method`
+  passed to the `QNode`, instead of assuming the method is always `deferred`.
+  [(#6903)](https://github.com/PennyLaneAI/pennylane/pull/6903)
+
 <h3>Contributors ✍️</h3>
 
 This release contains contributions from (in alphabetical order):
@@ -163,6 +167,7 @@ This release contains contributions from (in alphabetical order):
 Yushao Chen,
 Isaac De Vlugt,
 Diksha Dhawan,
+Lillian M.A. Frederiksen,
 Pietropaolo Frisoni,
 Marcus Gisslén,
 Christina Lee,
diff --git a/pennylane/workflow/construct_batch.py b/pennylane/workflow/construct_batch.py
index 0f677d5d2..8451b8d77 100644
--- a/pennylane/workflow/construct_batch.py
+++ b/pennylane/workflow/construct_batch.py
@@ -65,7 +65,12 @@ def _get_full_transform_program(
             **qnode.gradient_kwargs,
         )
 
-    config = _make_execution_config(qnode, gradient_fn)
+    mcm_config = qml.devices.MCMConfig(
+        postselect_mode=qnode.execute_kwargs["postselect_mode"],
+        mcm_method=qnode.execute_kwargs["mcm_method"],
+    )
+
+    config = _make_execution_config(qnode, gradient_fn, mcm_config)
     return program + qnode.device.preprocess_transforms(config)
 
 


Test patch:
diff --git a/tests/workflow/test_construct_batch.py b/tests/workflow/test_construct_batch.py
index d4288da65..bf6f4f1b9 100644
--- a/tests/workflow/test_construct_batch.py
+++ b/tests/workflow/test_construct_batch.py
@@ -448,3 +448,20 @@ class TestConstructBatch:
         )
         qml.assert_equal(batch[0], expected)
         assert fn(("a",)) == ("a",)
+
+    @pytest.mark.parametrize(
+        "mcm_method, expected_op",
+        [("deferred", qml.CNOT), ("tree-traversal", qml.measurements.MidMeasureMP)],
+    )
+    def test_mcm_method(self, mcm_method, expected_op):
+        """Test that the tape is constructed using the mcm_method specified on the QNode"""
+
+        @qml.qnode(qml.device("default.qubit"), mcm_method=mcm_method)
+        def circuit():
+            qml.measure(0)
+            return qml.expval(qml.Z(0))
+
+        (tape,), _ = qml.workflow.construct_batch(circuit, level="device")()
+
+        assert len(tape.operations) == 1
+        assert isinstance(tape.operations[0], expected_op)