Repository: PennyLaneAI/pennylane
Pull request #2234 (2022-02-25 22:25:02)

Circuit cutting: add mid-circuit measurement integration tests
**Context:**

Adds tests for the `cut_circuit` transform that ensures circuits resulting in fragments that contain mid-circuit measurements are supported.

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://codecov.io/gh/PennyLaneAI/pennylane/pull/2234?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
> Merging [#2234](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2234?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (01e507c) into [qcut_integ_tests](https://codecov.io/gh/PennyLaneAI/pennylane/commit/7192f3dc21a78aba49d6500585d116d7c9f2ef27?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (7192f3d) will **decrease** coverage by `2.71%`.
> The diff coverage is `100.00%`.

[![Impacted file tree graph](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2234/graphs/tree.svg?width=650&height=150&src=pr&token=9S3YzlUgGk&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI)](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2234?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI)

```diff
@@                 Coverage Diff                  @@
##           qcut_integ_tests    #2234      +/-   ##
====================================================
- Coverage             99.26%   96.55%   -2.72%     
====================================================
  Files                   231      231              
  Lines                 18349    18359      +10     
====================================================
- Hits                  18215    17726     -489     
- Misses                  134      633     +499     
```


| [Impacted Files](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2234?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | Coverage Δ | |
|---|---|---|
| [pennylane/transforms/qcut.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2234/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3RyYW5zZm9ybXMvcWN1dC5weQ==) | `100.00% <100.00%> (ø)` | |
| [pennylane/interfaces/batch/jax.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2234/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2ludGVyZmFjZXMvYmF0Y2gvamF4LnB5) | `0.00% <0.00%> (-100.00%)` | :arrow_down: |
| [pennylane/interfaces/batch/jax\_jit.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2234/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2ludGVyZmFjZXMvYmF0Y2gvamF4X2ppdC5weQ==) | `0.00% <0.00%> (-100.00%)` | :arrow_down: |
| [pennylane/interfaces/jax.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2234/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2ludGVyZmFjZXMvamF4LnB5) | `0.00% <0.00%> (-97.44%)` | :arrow_down: |
| [pennylane/devices/default\_qubit\_jax.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2234/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2RldmljZXMvZGVmYXVsdF9xdWJpdF9qYXgucHk=) | `0.00% <0.00%> (-96.63%)` | :arrow_down: |
| [pennylane/qnn/cost.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2234/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3Fubi9jb3N0LnB5) | `42.85% <0.00%> (-57.15%)` | :arrow_down: |
| [pennylane/transforms/classical\_jacobian.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2234/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3RyYW5zZm9ybXMvY2xhc3NpY2FsX2phY29iaWFuLnB5) | `77.08% <0.00%> (-22.92%)` | :arrow_down: |
| [pennylane/math/utils.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2234/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL21hdGgvdXRpbHMucHk=) | `84.41% <0.00%> (-15.59%)` | :arrow_down: |
| [pennylane/math/multi\_dispatch.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2234/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL21hdGgvbXVsdGlfZGlzcGF0Y2gucHk=) | `87.73% <0.00%> (-12.27%)` | :arrow_down: |
| [pennylane/math/is\_independent.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2234/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL21hdGgvaXNfaW5kZXBlbmRlbnQucHk=) | `88.50% <0.00%> (-11.50%)` | :arrow_down: |
| ... and [14 more](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2234/diff?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | |

------

[Continue to review full report at Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2234?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2234?src=pr&el=footer&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI). Last update [7192f3d...01e507c](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2234?src=pr&el=lastupdated&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).

trbromley: [sc-14724]

Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 7395cda0c..78a419eac 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -91,6 +91,7 @@
 
   A suite of integration tests has been added.
   [(#2231)](https://github.com/PennyLaneAI/pennylane/pull/2231)
+  [(#2234)](https://github.com/PennyLaneAI/pennylane/pull/2234)
 
 <h3>Improvements</h3>
 
diff --git a/pennylane/transforms/qcut.py b/pennylane/transforms/qcut.py
index d800ca8e1..701f59bdf 100644
--- a/pennylane/transforms/qcut.py
+++ b/pennylane/transforms/qcut.py
@@ -361,12 +361,16 @@ def graph_to_tape(graph: MultiDiGraph) -> QuantumTape:
     wire_map = {w: w for w in wires}
     reverse_wire_map = {v: k for k, v in wire_map.items()}
 
-    copy_ops = [copy.copy(op) for _, op in ordered_ops]
+    copy_ops = [copy.copy(op) for _, op in ordered_ops if not isinstance(op, MeasurementProcess)]
+    copy_meas = [copy.copy(op) for _, op in ordered_ops if isinstance(op, MeasurementProcess)]
+    observables = []
 
     with QuantumTape() as tape:
         for op in copy_ops:
-            new_wires = [wire_map[w] for w in op.wires]
-            op._wires = Wires(new_wires)  # TODO: find a better way to update operation wires
+            new_wires = Wires([wire_map[w] for w in op.wires])
+
+            # TODO: find a better way to update operation wires
+            op._wires = new_wires
             apply(op)
 
             if isinstance(op, MeasureNode):
@@ -380,6 +384,19 @@ def graph_to_tape(graph: MultiDiGraph) -> QuantumTape:
                 wire_map[original_wire] = new_wire
                 reverse_wire_map[new_wire] = original_wire
 
+        for meas in copy_meas:
+            obs = meas.obs
+            obs._wires = Wires([wire_map[w] for w in obs.wires])
+            observables.append(obs)
+
+        # We assume that each MeasurementProcess node in the graph contributes to a single
+        # expectation value of an observable, given by the tensor product over the observables of
+        # each MeasurementProcess.
+        if len(observables) > 1:
+            qml.expval(Tensor(*observables))
+        elif len(observables) == 1:
+            qml.expval(obs)
+
     return tape
 
 
@@ -424,7 +441,7 @@ def _get_measurements(
 
     obs = measurement.obs
 
-    return [expval(obs @ g) for g in group]
+    return [expval(copy.copy(obs) @ g) for g in group]
 
 
 def _prep_zero_state(wire):


Test patch:
diff --git a/tests/transforms/test_qcut.py b/tests/transforms/test_qcut.py
index f9af242dc..08e179f10 100644
--- a/tests/transforms/test_qcut.py
+++ b/tests/transforms/test_qcut.py
@@ -987,6 +987,41 @@ class TestGraphToTape:
         for tape1, tape2 in zip(tapes1, tapes2):
             compare_tapes(tape1, tape2)
 
+    def test_identity(self):
+        """Tests that the graph_to_tape function correctly performs the inverse of the tape_to_graph
+        function, including converting a tensor product expectation value into separate nodes in the
+        graph returned by tape_to_graph, and then combining those nodes again into a single tensor
+        product in the circuit returned by graph_to_tape"""
+
+        with qml.tape.QuantumTape() as tape:
+            qml.CNOT(wires=[0, 1])
+            qml.expval(qml.PauliZ(0) @ qml.PauliZ(1))
+
+        graph = qcut.tape_to_graph(tape)
+        tape_out = qcut.graph_to_tape(graph)
+
+        compare_tapes(tape, tape_out)
+        assert len(tape_out.measurements) == 1
+
+    def test_change_obs_wires(self):
+        """Tests that the graph_to_tape function correctly swaps the wires of observables when
+        the tape contains mid-circuit measurements"""
+
+        with qml.tape.QuantumTape() as tape:
+            qml.CNOT(wires=[0, 1])
+            qcut.MeasureNode(wires=1)
+            qcut.PrepareNode(wires=1)
+            qml.CNOT(wires=[0, 1])
+            qml.expval(qml.PauliZ(1))
+
+        graph = qcut.tape_to_graph(tape)
+        tape_out = qcut.graph_to_tape(graph)
+
+        m = tape_out.measurements
+        assert len(m) == 1
+        assert m[0].wires == Wires([2])
+        assert m[0].obs.name == "PauliZ"
+
 
 class TestGetMeasurements:
     """Tests for the _get_measurements function"""
@@ -1940,6 +1975,136 @@ class TestCutCircuitTransform:
         assert np.isclose(res, res_expected)
         assert np.isclose(grad, grad_expected)
 
+    def test_with_mid_circuit_measurement(self, mocker, use_opt_einsum):
+        """Tests the full circuit cutting pipeline returns the correct value and gradient for a
+        circuit that contains mid-circuit measurements, using the `cut_circuit` transform."""
+        dev = qml.device("default.qubit", wires=3)
+
+        @qml.qnode(dev)
+        def circuit(x):
+            qml.RX(x, wires=0)
+            qml.CNOT(wires=[0, 1])
+            qml.WireCut(wires=1)
+            qml.RX(np.sin(x) ** 2, wires=1)
+            qml.CNOT(wires=[1, 2])
+            qml.WireCut(wires=1)
+            qml.CNOT(wires=[0, 1])
+            return qml.expval(qml.PauliZ(0) @ qml.PauliZ(1))
+
+        spy = mocker.spy(qcut, "qcut_processing_fn")
+        x = np.array(0.531, requires_grad=True)
+        cut_circuit = qcut.cut_circuit(circuit, use_opt_einsum=use_opt_einsum)
+
+        assert np.isclose(cut_circuit(x), float(circuit(x)))
+        spy.assert_called_once()
+
+        gradient = qml.grad(circuit)(x)
+        cut_gradient = qml.grad(cut_circuit)(x)
+
+        assert np.isclose(gradient, cut_gradient)
+
+    def test_simple_cut_circuit_torch(self, use_opt_einsum):
+        """
+        Tests the full circuit cutting pipeline returns the correct value and
+        gradient for a simple circuit using the `cut_circuit` transform with the torch interface.
+        """
+        torch = pytest.importorskip("torch")
+
+        dev = qml.device("default.qubit", wires=2)
+
+        @qml.qnode(dev, interface="torch")
+        def circuit(x):
+            qml.RX(x, wires=0)
+            qml.RY(0.543, wires=1)
+            qml.WireCut(wires=0)
+            qml.CNOT(wires=[0, 1])
+            qml.RZ(0.240, wires=0)
+            qml.RZ(0.133, wires=1)
+            return qml.expval(qml.PauliZ(wires=[0]))
+
+        x = torch.tensor(0.531, requires_grad=True)
+        cut_circuit = qcut.cut_circuit(circuit, use_opt_einsum=use_opt_einsum)
+
+        res = cut_circuit(x)
+        res_expected = circuit(x)
+        assert np.isclose(res.detach().numpy(), res_expected.detach().numpy())
+
+        res.backward()
+        grad = x.grad.detach().numpy()
+
+        x.grad = None
+        res_expected.backward()
+        grad_expected = x.grad.detach().numpy()
+
+        assert np.isclose(grad, grad_expected)
+
+    def test_simple_cut_circuit_tf(self, use_opt_einsum):
+        """
+        Tests the full circuit cutting pipeline returns the correct value and
+        gradient for a simple circuit using the `cut_circuit` transform with the TF interface.
+        """
+        tf = pytest.importorskip("tensorflow")
+
+        dev = qml.device("default.qubit", wires=2)
+
+        @qml.qnode(dev, interface="tf")
+        def circuit(x):
+            qml.RX(x, wires=0)
+            qml.RY(0.543, wires=1)
+            qml.WireCut(wires=0)
+            qml.CNOT(wires=[0, 1])
+            qml.RZ(0.240, wires=0)
+            qml.RZ(0.133, wires=1)
+            return qml.expval(qml.PauliZ(wires=[0]))
+
+        x = tf.Variable(0.531)
+        cut_circuit = qcut.cut_circuit(circuit, use_opt_einsum=use_opt_einsum)
+
+        with tf.GradientTape() as tape:
+            res = cut_circuit(x)
+
+        grad = tape.gradient(res, x)
+
+        with tf.GradientTape() as tape:
+            res_expected = circuit(x)
+
+        grad_expected = tape.gradient(res_expected, x)
+
+        assert np.isclose(res, res_expected)
+        assert np.isclose(grad, grad_expected)
+
+    def test_simple_cut_circuit_jax(self, use_opt_einsum):
+        """
+        Tests the full circuit cutting pipeline returns the correct value and
+        gradient for a simple circuit using the `cut_circuit` transform with the Jax interface.
+        """
+        jax = pytest.importorskip("jax")
+        import jax.numpy as jnp
+
+        dev = qml.device("default.qubit", wires=2)
+
+        @qml.qnode(dev, interface="jax")
+        def circuit(x):
+            qml.RX(x, wires=0)
+            qml.RY(0.543, wires=1)
+            qml.WireCut(wires=0)
+            qml.CNOT(wires=[0, 1])
+            qml.RZ(0.240, wires=0)
+            qml.RZ(0.133, wires=1)
+            return qml.expval(qml.PauliZ(wires=[0]))
+
+        x = jnp.array(0.531)
+        cut_circuit = qcut.cut_circuit(circuit, use_opt_einsum=use_opt_einsum)
+
+        res = cut_circuit(x)
+        res_expected = circuit(x)
+
+        grad = jax.grad(cut_circuit)(x)
+        grad_expected = jax.grad(circuit)(x)
+
+        assert np.isclose(res, res_expected)
+        assert np.isclose(grad, grad_expected)
+
 
 class TestCutStrategy:
     """Tests for class CutStrategy"""