Qcut sampling template expansion (#2399)

* extend graph_to_tape for sampling, add unit test

* add PR to changelog

* add fragment sample circuit unit test

* Apply suggestions from code review

Co-authored-by: Tom Bromley <49409390+trbromley@users.noreply.github.com>

* add sampling tensor obs error, add test

* add test for sample obs

* Apply suggestions from code review

Co-authored-by: Tom Bromley <49409390+trbromley@users.noreply.github.com>

* extend graph to tape logic to sample measurements, add unit test

* remove unused logic, add mid circuit meas test

* update changelog

* fix variables

* Apply suggestions from code review

Co-authored-by: Tom Bromley <49409390+trbromley@users.noreply.github.com>

* add suggestions from code review

* add upsupported meas exception, add unit test

* fix codefactor issue

* add expand_fragment_tapes_mc function and unit test

* update docs

* update changelog

* update unit test

* add unit tests for MC measurements and state preps

* update expansion logic, add unit test

* fix wire bug

* Apply suggestions from code review

Co-authored-by: Tom Bromley <49409390+trbromley@users.noreply.github.com>

* add suggestions from code review

* begin mc postprocessing function, add unit test

* attempt MC postprocessing

* update sample postprocessing, add mc, add unit test

* add mc unit test, update docs and changlog

* Apply suggestions from code review

Co-authored-by: Tom Bromley <49409390+trbromley@users.noreply.github.com>

* add suggestions from code review

* add remaining suggestions

* add cut_circuit_mc transform, add unit test

* update docs, changelog, tests

* Apply suggestions from code review

Co-authored-by: Tom Bromley <49409390+trbromley@users.noreply.github.com>

* add suggestions from code review

* update shots logic, add unit tests

* draft usage details

* add validation, add unit tests

* update usage details wording

* update _reshape_results helper function and unit test

* update usage details with working example

* Apply suggestions from code review

Co-authored-by: Tom Bromley <49409390+trbromley@users.noreply.github.com>

* add review suggestions

* add unit tests for validation

* fix test

* add _qcut_expand_fn_mc and unit tests

* update changelog

* retrigger checks

* revert merge conflict

* restructure logic

* add docstrings

* Minor changes

* Compact parametrize

* Correct typo

* Use kwargs

* Fix spies

* Reword

* Fix CI

* retrigger checks

Co-authored-by: Tom Bromley <49409390+trbromley@users.noreply.github.com>
Co-authored-by: trbromley <brotho02@gmail.com>

doc/releases/changelog-dev.md

@@ -138,6 +138,10 @@
     can be decorated with `@qml.cut_circuit_mc()` to perform this type of cutting.
     [(#2382)](https://github.com/PennyLaneAI/pennylane/pull/2382)
 
+  - Add expansion to `qcut.cut_circuit_mc()` to search for wire cuts in
+    contained operations or tapes.
+    [(#2399)](https://github.com/PennyLaneAI/pennylane/pull/2399)
+
 <h3>Improvements</h3>
 
 * The parameter-shift Hessian can now be computed for arbitrary

pennylane/transforms/qcut.py

@@ -15,12 +15,12 @@
 Functions for performing quantum circuit cutting.
 """
 
-from collections.abc import Sequence as SequenceType
 import copy
 import inspect
 import string
 import uuid
 import warnings
+from collections.abc import Sequence as SequenceType
 from dataclasses import InitVar, dataclass
 from functools import partial
 from itertools import compress, product
@@ -1897,29 +1897,52 @@ def qnode_execution_wrapper(self, qnode, targs, tkwargs):
 
 def _qcut_expand_fn(
     tape: QuantumTape,
-    use_opt_einsum: bool = False,
-    device_wires: Optional[Wires] = None,
     max_depth: int = 1,
 ):
     """Expansion function for circuit cutting.
 
     Expands operations until reaching a depth that includes :class:`~.WireCut` operations.
     """
-    # pylint: disable=unused-argument
     for op in tape.operations:
         if isinstance(op, WireCut):
             return tape
 
     if max_depth > 0:
-        return cut_circuit.expand_fn(tape.expand(), max_depth=max_depth - 1)
+        return _qcut_expand_fn(tape.expand(), max_depth=max_depth - 1)
 
     raise ValueError(
         "No WireCut operations found in the circuit. Consider increasing the max_depth value if "
         "operations or nested tapes contain WireCut operations."
     )
 
 
-cut_circuit.expand_fn = _qcut_expand_fn
+def _cut_circuit_expand(
+    tape: QuantumTape,
+    use_opt_einsum: bool = False,
+    device_wires: Optional[Wires] = None,
+    max_depth: int = 1,
+):
+    """Main entry point for expanding operations until reaching a depth that
+    includes :class:`~.WireCut` operations."""
+    # pylint: disable=unused-argument
+    return _qcut_expand_fn(tape, max_depth)
+
+
+def _cut_circuit_mc_expand(
+    tape: QuantumTape,
+    shots: Optional[int] = None,
+    device_wires: Optional[Wires] = None,
+    classical_processing_fn: Optional[callable] = None,
+    max_depth: int = 1,
+):
+    """Main entry point for expanding operations in sample-based tapes until
+    reaching a depth that includes :class:`~.WireCut` operations."""
+    # pylint: disable=unused-argument
+    return _qcut_expand_fn(tape, max_depth)
+
+
+cut_circuit.expand_fn = _cut_circuit_expand
+cut_circuit_mc.expand_fn = _cut_circuit_mc_expand
 
 
 def remap_tape_wires(tape: QuantumTape, wires: Sequence) -> QuantumTape:

tests/transforms/test_qcut.py

@@ -3628,47 +3628,61 @@ def test_no_cuts_raises(self):
 
 
 class TestCutCircuitExpansion:
-    """Test of expansion in the cut_circuit function"""
+    """Test of expansion in the cut_circuit and cut_circuit_mc functions"""
 
-    def test_no_expansion(self, mocker):
+    transform_measurement_pairs = [
+        (qcut.cut_circuit, qml.expval(qml.PauliZ(0))),
+        (qcut.cut_circuit_mc, qml.sample(wires=[0])),
+    ]
+
+    @pytest.mark.parametrize("cut_transform, measurement", transform_measurement_pairs)
+    def test_no_expansion(self, mocker, cut_transform, measurement):
         """Test if no/trivial expansion occurs if WireCut operations are already present in the
         tape"""
         with qml.tape.QuantumTape() as tape:
             qml.RX(0.3, wires=0)
             qml.WireCut(wires=0)
             qml.RY(0.4, wires=0)
-            qml.expval(qml.PauliZ(0))
+            qml.apply(measurement)
+
+        spy = mocker.spy(qcut, "_qcut_expand_fn")
+
+        kwargs = {"shots": 10} if measurement.return_type is qml.measurements.Sample else {}
+        cut_transform(tape, device_wires=[0], **kwargs)
 
-        spy = mocker.spy(qcut.cut_circuit, "expand_fn")
-        qcut.cut_circuit(tape, device_wires=[0])
         spy.assert_called_once()
 
-    def test_expansion(self, mocker):
+    @pytest.mark.parametrize("cut_transform, measurement", transform_measurement_pairs)
+    def test_expansion(self, mocker, cut_transform, measurement):
         """Test if expansion occurs if WireCut operations are present in a nested tape"""
         with qml.tape.QuantumTape() as tape:
             qml.RX(0.3, wires=0)
             with qml.tape.QuantumTape() as _:
                 qml.WireCut(wires=0)
             qml.RY(0.4, wires=0)
-            qml.expval(qml.PauliZ(0))
+            qml.apply(measurement)
 
-        spy = mocker.spy(qcut.cut_circuit, "expand_fn")
-        qcut.cut_circuit(tape, device_wires=[0])
+        spy = mocker.spy(qcut, "_qcut_expand_fn")
+
+        kwargs = {"shots": 10} if measurement.return_type is qml.measurements.Sample else {}
+        cut_transform(tape, device_wires=[0], **kwargs)
 
         assert spy.call_count == 2
 
-    def test_expansion_error(self):
+    @pytest.mark.parametrize("cut_transform, measurement", transform_measurement_pairs)
+    def test_expansion_error(self, cut_transform, measurement):
         """Test if a ValueError is raised if expansion continues beyond the maximum depth"""
         with qml.tape.QuantumTape() as tape:
             qml.RX(0.3, wires=0)
             with qml.tape.QuantumTape() as _:
                 with qml.tape.QuantumTape() as __:
                     qml.WireCut(wires=0)
             qml.RY(0.4, wires=0)
-            qml.expval(qml.PauliZ(0))
+            qml.apply(measurement)
 
         with pytest.raises(ValueError, match="No WireCut operations found in the circuit."):
-            qcut.cut_circuit(tape, device_wires=[0], max_depth=1)
+            kwargs = {"shots": 10} if measurement.return_type is qml.measurements.Sample else {}
+            cut_transform(tape, device_wires=[0], **kwargs)
 
     def test_expansion_ttn(self, mocker):
         """Test if wire cutting is compatible with the tree tensor network operation"""
@@ -3695,12 +3709,39 @@ def circuit(template_weights):
         qnode = qml.QNode(circuit, dev_big)
         qnode_cut = qcut.cut_circuit(qml.QNode(circuit, dev_cut))
 
-        spy = mocker.spy(qcut.cut_circuit, "expand_fn")
+        spy = mocker.spy(qcut, "_qcut_expand_fn")
         res = qnode_cut(template_weights)
         assert spy.call_count == 2
 
         assert np.isclose(res, qnode(template_weights))
 
+    def test_expansion_mc_ttn(self, mocker):
+        """Test if wire cutting is compatible with the tree tensor network operation"""
+
+        def block(weights, wires):
+            qml.CNOT(wires=[wires[0], wires[1]])
+            qml.RY(weights[0], wires=wires[0])
+            qml.RY(weights[1], wires=wires[1])
+            qml.WireCut(wires=wires[1])
+
+        n_wires = 4
+        n_block_wires = 2
+        n_params_block = 2
+        n_blocks = qml.TTN.get_n_blocks(range(n_wires), n_block_wires)
+        template_weights = [[0.1, -0.3]] * n_blocks
+
+        dev_cut = qml.device("default.qubit", wires=2, shots=10)
+
+        def circuit(template_weights):
+            qml.TTN(range(n_wires), n_block_wires, block, n_params_block, template_weights)
+            return qml.sample(wires=[n_wires - 1])
+
+        qnode_cut = qcut.cut_circuit_mc(qml.QNode(circuit, dev_cut))
+
+        spy = mocker.spy(qcut, "_qcut_expand_fn")
+        qnode_cut(template_weights)
+        assert spy.call_count == 2
+
 
 class TestCutStrategy:
     """Tests for class CutStrategy"""