Hamiltonian support for Circuit Cutting (#4642)

**Context:** `qml.cut_circuit` does not work for quantum circuits which
computes the expectation values of Hamiltonians with more than one term.

**Description of the Change:** Use `qml.transforms.hamiltonian_expand`
to extend support for Hamiltonians.

**Benefits:** Circuit cutting could be used on more realistic circuits. 

**Possible Drawbacks:** The Number of tapes to be computed increases and
might give rise to possible overhead in the new scenario.

**Related GitHub Issues:**

---------

Co-authored-by: Romain Moyard <[email protected]>

doc/releases/changelog-dev.md

@@ -183,6 +183,10 @@
   interface-specific scalar data, eg `[(tf.Variable(1.1), tf.Variable(2.2))]`.
   [(#4603)](https://github.com/PennyLaneAI/pennylane/pull/4603)
 
+* `qml.cut_circuit` is now compatible with circuits that compute the expectation values of Hamiltonians 
+  with two or more terms.
+  [(#4642)](https://github.com/PennyLaneAI/pennylane/pull/4642)
+
 * `_qfunc_output` has been removed from `QuantumScript`, as it is no longer necessary. There is
   still a `_qfunc_output` property on `QNode` instances.
   [(#4651)](https://github.com/PennyLaneAI/pennylane/pull/4651)

pennylane/transforms/qcut/cutcircuit.py

@@ -46,7 +46,23 @@ def _cut_circuit_expand(
     def processing_fn(res):
         return res[0]
 
-    return [_qcut_expand_fn(tape, max_depth, auto_cutter)], processing_fn
+    tapes, tapes_fn = [tape], processing_fn
+
+    # Expand the tapes for handling Hamiltonian with two or more terms
+    tape_meas_ops = tape.measurements
+    if isinstance(tape_meas_ops[0].obs, qml.Hamiltonian):
+        if len(tape_meas_ops) > 1:
+            raise NotImplementedError(
+                "Hamiltonian expansion is supported only with a single Hamiltonian"
+            )
+
+        new_meas_op = type(tape_meas_ops[0])(obs=qml.Hamiltonian(*tape_meas_ops[0].obs.terms()))
+        new_tape = qml.tape.QuantumScript(tape.operations, [new_meas_op], shots=tape.shots)
+        new_tape.trainable_params = tape.trainable_params
+
+        tapes, tapes_fn = qml.transforms.hamiltonian_expand(new_tape, group=False)
+
+    return [_qcut_expand_fn(tape, max_depth, auto_cutter) for tape in tapes], tapes_fn
 
 
 @partial(transform, expand_transform=_cut_circuit_expand)
@@ -365,8 +381,10 @@ def circuit(x):
                 "installed using:\npip install opt_einsum"
             ) from e
 
+    # convert the quantum tape to a DAG structure
     g = tape_to_graph(tape)
 
+    # place WireCut(s) nodes in the DAG automatically if intended
     if auto_cutter is True or callable(auto_cutter):
         cut_strategy = kwargs.pop("cut_strategy", None) or CutStrategy(
             max_free_wires=len(device_wires)
@@ -379,12 +397,19 @@ def circuit(x):
             **kwargs,
         )
 
+    # replace the WireCut nodes in the DAG with Measure and Perpare nodes.
     replace_wire_cut_nodes(g)
+
+    # decompose the DAG into subgraphs based on the replaced WireCut(s)
+    # along with a quotient graph to store connections between them
     fragments, communication_graph = fragment_graph(g)
+
+    # convert decomposed DAGs into tapes, remap their wires for device and expand them
     fragment_tapes = [graph_to_tape(f) for f in fragments]
     fragment_tapes = [qml.map_wires(t, dict(zip(t.wires, device_wires))) for t in fragment_tapes]
     expanded = [expand_fragment_tape(t) for t in fragment_tapes]
 
+    # store the data necessary for classical post processing of results
     configurations = []
     prepare_nodes = []
     measure_nodes = []
@@ -393,6 +418,7 @@ def circuit(x):
         prepare_nodes.append(p)
         measure_nodes.append(m)
 
+    # flatten out the tapes to be returned
     tapes = tuple(tape for c in configurations for tape in c)
 
     return tapes, partial(

tests/transforms/test_qcut.py

@@ -5420,3 +5420,202 @@ class TestRedirect:
 
     def test_qcut_redirects_to_qcut_qcut(self):
         assert qml.transforms.qcut._prep_one_state == qml.transforms.qcut.qcut._prep_one_state
+
+
+class TestCutCircuitWithHamiltonians:
+    """Integration tests for `cut_circuit` transform with Hamiltonians."""
+
+    def test_circuit_with_hamiltonian(self, mocker):
+        """
+        Tests that the full automatic circuit cutting pipeline returns the correct value and
+        gradient for a complex circuit with multiple wire cut scenarios. The circuit is the
+        uncut version of the circuit in ``TestCutCircuitTransform.test_complicated_circuit``.
+
+        0: ──BasisState(M0)─╭C───RX─╭C──╭C────────────────────┤
+        1: ─────────────────╰X──────╰X──╰Z────────────────╭RX─┤ ╭<H>
+        2: ──H──────────────╭C─────────────╭RY────────╭RY─│───┤ ├<H>
+        3: ─────────────────╰RY──H──╭C───H─╰C──╭RY──H─╰C──│───┤ ╰<H>
+        4: ─────────────────────────╰RY──H─────╰C─────────╰C──┤
+        """
+
+        dev_original = qml.device("default.qubit", wires=5)
+        dev_cut = qml.device("default.qubit", wires=4)
+
+        hamiltonian = qml.Hamiltonian(
+            [1.0, 1.0],
+            [qml.PauliZ(1) @ qml.PauliZ(2) @ qml.PauliZ(3), qml.PauliY(0) @ qml.PauliX(1)],
+        )
+
+        def two_qubit_unitary(param, wires):
+            qml.Hadamard(wires=[wires[0]])
+            qml.CRY(param, wires=[wires[0], wires[1]])
+
+        def f(params):
+            qml.BasisState(np.array([1]), wires=[0])
+            qml.WireCut(wires=0)
+
+            qml.CNOT(wires=[0, 1])
+            qml.WireCut(wires=0)
+            qml.RX(params[0], wires=0)
+            qml.CNOT(wires=[0, 1])
+
+            qml.WireCut(wires=0)
+            qml.WireCut(wires=1)
+
+            qml.CZ(wires=[0, 1])
+            qml.WireCut(wires=[0, 1])
+
+            two_qubit_unitary(params[1], wires=[2, 3])
+            qml.WireCut(wires=3)
+            two_qubit_unitary(params[2] ** 2, wires=[3, 4])
+            qml.WireCut(wires=3)
+            two_qubit_unitary(np.sin(params[3]), wires=[3, 2])
+            qml.WireCut(wires=3)
+            two_qubit_unitary(np.sqrt(params[4]), wires=[4, 3])
+            qml.WireCut(wires=3)
+            two_qubit_unitary(np.cos(params[1]), wires=[3, 2])
+            qml.CRX(params[2], wires=[4, 1])
+
+            return qml.expval(hamiltonian)
+
+        params = np.array([0.4, 0.5, 0.6, 0.7, 0.8], requires_grad=True)
+
+        circuit = qml.QNode(f, dev_original)
+        cut_circuit = qcut.cut_circuit(qml.QNode(f, dev_cut))
+
+        res_expected = circuit(params)
+        grad_expected = qml.grad(circuit)(params)
+
+        spy = mocker.spy(qcut.cutcircuit, "qcut_processing_fn")
+        res = cut_circuit(params)
+        assert spy.call_count == len(hamiltonian.ops)
+
+        grad = qml.grad(cut_circuit)(params)
+
+        assert np.isclose(res, res_expected)
+        assert np.allclose(grad, grad_expected)
+
+    def test_autoscale_and_grouped_with_hamiltonian(self, mocker):
+        """
+        Tests that the full circuit cutting pipeline returns the correct value for a typical
+        scenario with auto-scaling. The circuit is the uncut version of the circuit in
+        ``TestCutCircuitTransform.test_standard_circuit``:
+
+        0: ─╭U(M1)───────────────────╭U(M4)─┤ ╭<Z@X>
+        1: ─╰U(M1)─────╭U(M2)────────╰U(M4)─┤ │
+        2: ─╭U(M0)─────╰U(M2)─╭U(M3)────────┤ │
+        3: ─╰U(M0)────────────╰U(M3)────────┤ ╰<Z@X>
+        """
+        pytest.importorskip("kahypar")
+
+        dev_original = qml.device("default.qubit")
+
+        hamiltonian = qml.Hamiltonian(
+            [1.0, 1.0, 1.0],
+            [
+                qml.PauliZ(1) @ qml.PauliZ(2) @ qml.PauliZ(3),
+                qml.PauliZ(1) @ qml.PauliZ(2),
+                qml.PauliZ(0) @ qml.PauliZ(1),
+            ],
+            grouping_type="qwc",
+        )
+
+        # We need a 3-qubit device
+        dev_cut = qml.device("default.qubit")
+        us = [unitary_group.rvs(2**2, random_state=i) for i in range(5)]
+
+        def f():
+            qml.QubitUnitary(us[0], wires=[0, 1])
+            qml.QubitUnitary(us[1], wires=[2, 3])
+            qml.QubitUnitary(us[2], wires=[1, 2])
+            qml.QubitUnitary(us[3], wires=[0, 1])
+            qml.QubitUnitary(us[4], wires=[2, 3])
+            return qml.expval(hamiltonian)
+
+        circuit = qml.QNode(f, dev_original)
+        cut_circuit = qcut.cut_circuit(
+            qml.QNode(f, dev_cut), auto_cutter=True, device_wires=qml.wires.Wires(range(3))
+        )
+
+        res_expected = circuit()
+
+        spy = mocker.spy(qcut.cutcircuit, "qcut_processing_fn")
+        res = cut_circuit()
+        assert spy.call_count == len(hamiltonian.ops)
+        assert np.isclose(res, res_expected, atol=1e-8)
+        assert cut_circuit.tape.measurements[0].obs.grouping_indices == hamiltonian.grouping_indices
+
+    def test_template_with_hamiltonian(self):
+        """Test cut with MPS Template"""
+
+        pytest.importorskip("kahypar")
+
+        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])
+
+        n_wires = 8
+        n_block_wires = 2
+        n_params_block = 2
+        n_blocks = qml.MPS.get_n_blocks(range(n_wires), n_block_wires)
+        template_weights = [[0.1, -0.3]] * n_blocks
+
+        device_size = 2
+        cut_strategy = qml.transforms.qcut.CutStrategy(max_free_wires=device_size)
+
+        hamiltonian = qml.Hamiltonian(
+            [1.0, 1.0],
+            [qml.PauliZ(1) @ qml.PauliZ(8) @ qml.PauliZ(3), qml.PauliY(5) @ qml.PauliX(4)],
+        )
+
+        with qml.queuing.AnnotatedQueue() as q0:
+            qml.MPS(range(n_wires), n_block_wires, block, n_params_block, template_weights)
+            qml.expval(hamiltonian)
+
+        tape0 = qml.tape.QuantumScript.from_queue(q0)
+        tape = tape0.expand()
+        tapes, _ = qml.transforms.hamiltonian_expand(tape, group=False)
+
+        frag_lens = [5, 7]
+        frag_ords = [[1, 6], [3, 6]]
+        for idx, tape in enumerate(tapes):
+            graph = qcut.tape_to_graph(tape)
+            cut_graph = qcut.find_and_place_cuts(
+                graph=graph,
+                cut_strategy=cut_strategy,
+                replace_wire_cuts=True,
+            )
+            frags, _ = qcut.fragment_graph(cut_graph)
+
+            assert len(frags) <= frag_lens[idx]
+
+            assert all(frag_ords[idx][0] <= f.order() <= frag_ords[idx][1] for f in frags)
+
+            # each frag should have the device size constraint satisfied.
+            assert all(len(set(e[2] for e in f.edges.data("wire"))) <= device_size for f in frags)
+
+    def test_raise_with_hamiltonian(self):
+        """Test that exception is correctly raise when caclulating expectation values of multiple Hamiltonians"""
+
+        dev_cut = qml.device("default.qubit", wires=4)
+
+        hamiltonian = qml.Hamiltonian(
+            [1.0, 1.0],
+            [qml.PauliZ(1) @ qml.PauliZ(2) @ qml.PauliZ(3), qml.PauliY(0) @ qml.PauliX(1)],
+        )
+
+        def f():
+            qml.CNOT(wires=[0, 1])
+            qml.WireCut(wires=0)
+            qml.RX(1.0, wires=0)
+            qml.CNOT(wires=[0, 1])
+
+            return [qml.expval(hamiltonian), qml.expval(hamiltonian)]
+
+        with pytest.raises(
+            NotImplementedError,
+            match="Hamiltonian expansion is supported only with a single Hamiltonian",
+        ):
+            cut_circuit = qcut.cut_circuit(qml.QNode(f, dev_cut))
+            cut_circuit()