Circuit cutting: Add expansion to cut_circuit transform (#2340)

* Add expansion

Co-authored-by: Josh Izaac <josh146@gmail.com>

* Reword

* Reword raise

* Add tests

* Add to changelog

* Update changelog

* Fix CI

* Update changelog

* Update max_depth default

* Fix

* Fix tests

Co-authored-by: Josh Izaac <josh146@gmail.com>

doc/releases/changelog-dev.md

@@ -20,7 +20,10 @@
   `ControlledOperation`. Control values of `0` are implemented by `qml.PauliX` applied
   before and after the controlled operation
   [(#2288)](https://github.com/PennyLaneAI/pennylane/pull/2288)
-
+
+* Circuit cutting now performs expansion to search for wire cuts in contained operations or tapes.
+  [(#2340)](https://github.com/PennyLaneAI/pennylane/pull/2340)
+
 <h3>Deprecations</h3>
 
 <h3>Breaking changes</h3>
@@ -48,4 +51,4 @@
 
 This release contains contributions from (in alphabetical order):
 
-Karim Alaa El-Din, Anthony Hayes, Josh Izaac, Christina Lee.
+Karim Alaa El-Din, Thomas Bromley, Anthony Hayes, Josh Izaac, Christina Lee.

pennylane/transforms/qcut.py

@@ -940,7 +940,10 @@ def qcut_processing_fn(
 
 @batch_transform
 def cut_circuit(
-    tape: QuantumTape, use_opt_einsum: bool = False, device_wires: Optional[Wires] = None
+    tape: QuantumTape,
+    use_opt_einsum: bool = False,
+    device_wires: Optional[Wires] = None,
+    max_depth: int = 1,
 ) -> Tuple[Tuple[QuantumTape], Callable]:
     """
     Cut up a quantum circuit into smaller circuit fragments.
@@ -965,6 +968,7 @@ def cut_circuit(
             e.g., ``pip install opt_einsum``. Both settings for ``use_opt_einsum`` result in a
             differentiable contraction.
         device_wires (Wires): wires of the device that the cut circuits are to be run on
+        max_depth (int): the maximum depth used to expand the circuit while searching for wire cuts
 
     Returns:
         Tuple[Tuple[QuantumTape], Callable]: the tapes corresponding to the circuit fragments as a
@@ -1155,6 +1159,7 @@ def circuit(x):
         ... )
         0.47165198882111165
     """
+    # pylint: disable=unused-argument
     if len(tape.measurements) != 1:
         raise ValueError(
             "The circuit cutting workflow only supports circuits with a single output "
@@ -1177,10 +1182,6 @@ def circuit(x):
                 "installed using:\npip install opt_einsum"
             ) from e
 
-    num_cut = len([op for op in tape.operations if isinstance(op, WireCut)])
-    if num_cut == 0:
-        raise ValueError("Cannot apply the circuit cutting workflow to a circuit without any cuts")
-
     g = tape_to_graph(tape)
     replace_wire_cut_nodes(g)
     fragments, communication_graph = fragment_graph(g)
@@ -1216,6 +1217,33 @@ def qnode_execution_wrapper(self, qnode, targs, tkwargs):
     return self.default_qnode_wrapper(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)
+
+    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 remap_tape_wires(tape: QuantumTape, wires: Sequence) -> QuantumTape:
     """Map the wires of a tape to a new set of wires.
 

tests/transforms/test_qcut.py

@@ -2691,6 +2691,7 @@ def test_multiple_measurements_raises(self):
         to be cut"""
 
         with qml.tape.QuantumTape() as tape:
+            qml.WireCut(wires=0)
             qml.expval(qml.PauliZ(0))
             qml.expval(qml.PauliZ(1))
 
@@ -2700,14 +2701,18 @@ def test_multiple_measurements_raises(self):
     def test_no_measurements_raises(self):
         """Tests if a ValueError is raised when a tape with no measurement is requested
         to be cut"""
+        with qml.tape.QuantumTape() as tape:
+            qml.WireCut(wires=0)
+
         with pytest.raises(ValueError, match="The circuit cutting workflow only supports circuits"):
-            qcut.cut_circuit(qml.tape.QuantumTape())
+            qcut.cut_circuit(tape)
 
     def test_non_expectation_raises(self):
         """Tests if a ValueError is raised when a tape with measurements that are not expectation
         values is requested to be cut"""
 
         with qml.tape.QuantumTape() as tape:
+            qml.WireCut(wires=0)
             qml.var(qml.PauliZ(0))
 
         with pytest.raises(ValueError, match="workflow only supports circuits with expectation"):
@@ -2716,6 +2721,7 @@ def test_non_expectation_raises(self):
     def test_fail_import(self, monkeypatch):
         """Test if an ImportError is raised when opt_einsum is requested but not installed"""
         with qml.tape.QuantumTape() as tape:
+            qml.WireCut(wires=0)
             qml.expval(qml.PauliZ(0))
 
         with monkeypatch.context() as m:
@@ -2730,10 +2736,85 @@ def test_no_cuts_raises(self):
         with qml.tape.QuantumTape() as tape:
             qml.expval(qml.PauliZ(0))
 
-        with pytest.raises(ValueError, match="to a circuit without any cuts"):
+        with pytest.raises(ValueError, match="No WireCut operations found in the circuit."):
             qcut.cut_circuit(tape)
 
 
+class TestCutCircuitExpansion:
+    """Test of expansion in the cut_circuit function"""
+
+    def test_no_expansion(self, mocker):
+        """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))
+
+        spy = mocker.spy(qcut.cut_circuit, "expand_fn")
+        qcut.cut_circuit(tape, device_wires=[0])
+        spy.assert_called_once()
+
+    def test_expansion(self, mocker):
+        """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))
+
+        spy = mocker.spy(qcut.cut_circuit, "expand_fn")
+        qcut.cut_circuit(tape, device_wires=[0])
+
+        assert spy.call_count == 2
+
+    def test_expansion_error(self):
+        """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))
+
+        with pytest.raises(ValueError, match="No WireCut operations found in the circuit."):
+            qcut.cut_circuit(tape, device_wires=[0], max_depth=1)
+
+    def test_expansion_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)
+        dev_big = qml.device("default.qubit", wires=4)
+
+        def circuit(template_weights):
+            qml.TTN(range(n_wires), n_block_wires, block, n_params_block, template_weights)
+            return qml.expval(qml.PauliZ(wires=n_wires - 1))
+
+        qnode = qml.QNode(circuit, dev_big)
+        qnode_cut = qcut.cut_circuit(qml.QNode(circuit, dev_cut))
+
+        spy = mocker.spy(qcut.cut_circuit, "expand_fn")
+        res = qnode_cut(template_weights)
+        assert spy.call_count == 2
+
+        assert np.isclose(res, qnode(template_weights))
+
+
 class TestCutStrategy:
     """Tests for class CutStrategy"""