Update quantum monte carlo (#4708)

**Description of the Change:**

- Update the quantum monte carlo to the new transform system
- Upadte applied_controlle_Q to the new transform system
- Update the documentation

**Benefits:**

The transform can now be applied directly on QNodes.

doc/releases/changelog-dev.md

@@ -383,6 +383,9 @@
 
 <h4>Other improvements</h4>
 
+* `qml.quantum_monte_carlo` now uses the new transform system.
+  [(#4708)](https://github.com/PennyLaneAI/pennylane/pull/4708/)
+
 * `qml.simplify` now uses the new transforms API.
   [(#4949)](https://github.com/PennyLaneAI/pennylane/pull/4949)
 
@@ -648,10 +651,9 @@ Juan Giraldo,
 Emiliano Godinez Ramirez,
 Ankit Khandelwal,
 Christina Lee,
-Romain Moyard,
 Vincent Michaud-Rioux,
-Romain Moyard,
 Anurav Modak,
+Romain Moyard,
 Mudit Pandey,
 Matthew Silverman,
 Jay Soni,

pennylane/transforms/qmc.py

@@ -14,10 +14,14 @@
 """
 Contains the quantum_monte_carlo transform.
 """
-from functools import wraps
+from copy import copy
+from typing import Sequence, Callable
+
+import pennylane as qml
 from pennylane import PauliX, Hadamard, MultiControlledX, CZ, adjoint
 from pennylane.wires import Wires
 from pennylane.templates import QFT
+from pennylane.transforms.core import transform
 
 
 def _apply_controlled_z(wires, control_wire, work_wires):
@@ -27,7 +31,7 @@ def _apply_controlled_z(wires, control_wire, work_wires):
     The multi-qubit gate :math:`Z = I - 2|0\rangle \langle 0|` can be performed using the
     conventional multi-controlled-Z gate with an additional bit flip on each qubit before and after.
 
-    This function performs the multi-controlled-Z gate via a multi-controlled-X gate by picking an
+    This function returns the multi-controlled-Z gate via a multi-controlled-X gate by picking an
     arbitrary target wire to perform the X and adding a Hadamard on that wire either side of the
     transformation.
 
@@ -39,19 +43,24 @@ def _apply_controlled_z(wires, control_wire, work_wires):
         work_wires (Wires): the work wires used in the decomposition
     """
     target_wire = wires[0]
-    PauliX(target_wire)
-    Hadamard(target_wire)
+    updated_operations = []
+    updated_operations.append(PauliX(target_wire))
+    updated_operations.append(Hadamard(target_wire))
 
     control_values = "0" * (len(wires) - 1) + "1"
     control_wires = wires[1:] + control_wire
-    MultiControlledX(
-        wires=[*control_wires, target_wire],
-        control_values=control_values,
-        work_wires=work_wires,
+    updated_operations.append(
+        MultiControlledX(
+            wires=[*control_wires, target_wire],
+            control_values=control_values,
+            work_wires=work_wires,
+        )
     )
 
-    Hadamard(target_wire)
-    PauliX(target_wire)
+    updated_operations.append(Hadamard(target_wire))
+    updated_operations.append(PauliX(target_wire))
+
+    return updated_operations
 
 
 def _apply_controlled_v(target_wire, control_wire):
@@ -67,23 +76,26 @@ def _apply_controlled_v(target_wire, control_wire):
         target_wire (Wires): the ancilla wire in which the expectation value is encoded
         control_wire (Wires): the control wire from the register of phase estimation qubits
     """
-    CZ(wires=[control_wire[0], target_wire[0]])
+    return [CZ(wires=[control_wire[0], target_wire[0]])]
 
 
-def apply_controlled_Q(fn, wires, target_wire, control_wire, work_wires):
-    r"""Provides the circuit to apply a controlled version of the :math:`\mathcal{Q}` unitary
+@transform
+def apply_controlled_Q(
+    tape: qml.tape.QuantumTape, wires, target_wire, control_wire, work_wires
+) -> (Sequence[qml.tape.QuantumTape], Callable):
+    r"""Applies the transform that performs a controlled version of the :math:`\mathcal{Q}` unitary
     defined in `this <https://arxiv.org/abs/1805.00109>`__ paper.
 
-    The input ``fn`` should be the quantum circuit corresponding to the :math:`\mathcal{F}` unitary
+    The input ``tape`` should be the quantum circuit corresponding to the :math:`\mathcal{F}` unitary
     in the paper above. This function transforms this circuit into a controlled version of the
     :math:`\mathcal{Q}` unitary, which forms part of the quantum Monte Carlo algorithm. The
     :math:`\mathcal{Q}` unitary encodes the target expectation value as a phase in one of its
     eigenvalues. This phase can be estimated using quantum phase estimation (see
     :class:`~.QuantumPhaseEstimation` for more details).
 
     Args:
-        fn (Callable): a quantum function that applies quantum operations according to the
-            :math:`\mathcal{F}` unitary used as part of quantum Monte Carlo estimation
+        tape (QNode or QuantumTape or Callable): the quantum circuit that applies quantum operations
+            according to the :math:`\mathcal{F}` unitary used as part of quantum Monte Carlo estimation
         wires (Union[Wires or Sequence[int]]): the wires acted upon by the ``fn`` circuit
         target_wire (Union[Wires, int]): The wire in which the expectation value is encoded. Must be
             contained within ``wires``.
@@ -93,41 +105,57 @@ def apply_controlled_Q(fn, wires, target_wire, control_wire, work_wires):
             decomposing :math:`\mathcal{Q}`
 
     Returns:
-        function: The input function transformed to the :math:`\mathcal{Q}` unitary
+        qnode (QNode) or quantum function (Callable) or tuple[List[QuantumTape], function]:
+
+        The transformed circuit as described in :func:`qml.transform <pennylane.transform>`. Executing this circuit
+        will perform control on :math:`\mathcal{Q}` unitary.
 
     Raises:
         ValueError: if ``target_wire`` is not in ``wires``
     """
-    fn_inv = adjoint(fn)
-
-    wires = Wires(wires)
-    target_wire = Wires(target_wire)
-    control_wire = Wires(control_wire)
-    work_wires = Wires(work_wires)
-
-    if not wires.contains_wires(target_wire):
-        raise ValueError("The target wire must be contained within wires")
-
-    @wraps(fn)
-    def wrapper(*args, **kwargs):
-        _apply_controlled_v(target_wire=target_wire, control_wire=control_wire)
-        fn_inv(*args, **kwargs)
-        _apply_controlled_z(wires=wires, control_wire=control_wire, work_wires=work_wires)
-        fn(*args, **kwargs)
-
-        _apply_controlled_v(target_wire=target_wire, control_wire=control_wire)
-        fn_inv(*args, **kwargs)
-        _apply_controlled_z(wires=wires, control_wire=control_wire, work_wires=work_wires)
-        fn(*args, **kwargs)
-
-    return wrapper
-
-
-def quantum_monte_carlo(fn, wires, target_wire, estimation_wires):
-    r"""Provides the circuit to perform the
+    operations = tape.operations.copy()
+    updated_operations = []
+
+    with qml.queuing.QueuingManager.stop_recording():
+        op_inv = [adjoint(copy(op)) for op in reversed(operations)]
+
+        wires = Wires(wires)
+        target_wire = Wires(target_wire)
+        control_wire = Wires(control_wire)
+        work_wires = Wires(work_wires)
+
+        if not wires.contains_wires(target_wire):
+            raise ValueError("The target wire must be contained within wires")
+
+        updated_operations.extend(
+            _apply_controlled_v(target_wire=target_wire, control_wire=control_wire)
+        )
+        updated_operations.extend(op_inv)
+        updated_operations.extend(
+            _apply_controlled_z(wires=wires, control_wire=control_wire, work_wires=work_wires)
+        )
+        updated_operations.extend(operations)
+        updated_operations.extend(
+            _apply_controlled_v(target_wire=target_wire, control_wire=control_wire)
+        )
+        updated_operations.extend(op_inv)
+        updated_operations.extend(
+            _apply_controlled_z(wires=wires, control_wire=control_wire, work_wires=work_wires)
+        )
+        updated_operations.extend(operations)
+
+    tape = type(tape)(updated_operations, tape.measurements, shots=tape.shots)
+    return [tape], lambda x: x[0]
+
+
+@transform
+def quantum_monte_carlo(
+    tape: qml.tape.QuantumTape, wires, target_wire, estimation_wires
+) -> (Sequence[qml.tape.QuantumTape], Callable):
+    r"""Applies the transform
     `quantum Monte Carlo estimation <https://arxiv.org/abs/1805.00109>`__ algorithm.
 
-    The input ``fn`` should be the quantum circuit corresponding to the :math:`\mathcal{F}` unitary
+    The input `tape`` should be the quantum circuit corresponding to the :math:`\mathcal{F}` unitary
     in the paper above. This unitary encodes the probability distribution and random variable onto
     ``wires`` so that measurement of the ``target_wire`` provides the expectation value to be
     estimated. The quantum Monte Carlo algorithm then estimates the expectation value using quantum
@@ -146,15 +174,19 @@ def quantum_monte_carlo(fn, wires, target_wire, estimation_wires):
         simulators, but may perform faster and is suited to quick prototyping.
 
     Args:
-        fn (Callable): a quantum function that applies quantum operations according to the
+        tape (QNode or QuantumTape or Callable): the quantum circuit that applies quantum operations according to the
             :math:`\mathcal{F}` unitary used as part of quantum Monte Carlo estimation
         wires (Union[Wires or Sequence[int]]): the wires acted upon by the ``fn`` circuit
         target_wire (Union[Wires, int]): The wire in which the expectation value is encoded. Must be
             contained within ``wires``.
         estimation_wires (Union[Wires, Sequence[int], or int]): the wires used for phase estimation
 
     Returns:
-        function: The circuit for quantum Monte Carlo estimation
+        qnode (QNode) or quantum function (Callable) or tuple[List[QuantumTape], function]:
+
+        The transformed circuit as described in :func:`qml.transform <pennylane.transform>`. Executing this circuit
+        will perform the quantum Monte Carlo estimation.
+
 
     Raises:
         ValueError: if ``wires`` and ``estimation_wires`` share a common wire
@@ -322,34 +354,35 @@ def qmc():
          'diff_method': 'best',
          'gradient_fn': 'backprop'}
     """
+    operations = tape.operations.copy()
     wires = Wires(wires)
     target_wire = Wires(target_wire)
     estimation_wires = Wires(estimation_wires)
 
     if Wires.shared_wires([wires, estimation_wires]):
         raise ValueError("No wires can be shared between the wires and estimation_wires registers")
 
-    @wraps(fn)
-    def wrapper(*args, **kwargs):
-        fn(*args, **kwargs)
+    updated_operations = []
+    with qml.queuing.QueuingManager.stop_recording():
+        updated_operations.extend(operations)
         for i, control_wire in enumerate(estimation_wires):
-            Hadamard(control_wire)
+            updated_operations.append(Hadamard(control_wire))
 
             # Find wires eligible to be used as helper wires
             work_wires = estimation_wires.toset() - {control_wire}
             n_reps = 2 ** (len(estimation_wires) - (i + 1))
 
-            q = apply_controlled_Q(
-                fn,
+            tapes_q, _ = apply_controlled_Q(
+                tape,
                 wires=wires,
                 target_wire=target_wire,
                 control_wire=control_wire,
                 work_wires=work_wires,
             )
-
+            tape_q = tapes_q[0]
             for _ in range(n_reps):
-                q(*args, **kwargs)
-
-        adjoint(QFT(wires=estimation_wires))
+                updated_operations.extend(tape_q.operations)
 
-    return wrapper
+        updated_operations.append(adjoint(QFT(wires=estimation_wires), lazy=False))
+    updated_tape = type(tape)(updated_operations, tape.measurements, shots=tape.shots)
+    return [updated_tape], lambda x: x[0]

tests/transforms/test_qmc_transform.py

@@ -176,7 +176,7 @@ def test_raises(self):
         with pytest.raises(ValueError, match="The target wire must be contained within wires"):
             apply_controlled_Q(
                 lambda: ..., wires=range(3), target_wire=4, control_wire=5, work_wires=None
-            )
+            )()
 
 
 class TestQuantumMonteCarlo:
@@ -225,7 +225,7 @@ def test_shared_wires(self):
         with pytest.raises(ValueError, match="No wires can be shared between the wires"):
             quantum_monte_carlo(
                 lambda: None, wires=wires, target_wire=0, estimation_wires=estimation_wires
-            )
+            )()
 
     @pytest.mark.slow
     def test_integration(self):
@@ -258,27 +258,21 @@ def fn():
             fn, wires=wires, target_wire=target_wire, estimation_wires=estimation_wires
         )
 
-        with qml.queuing.AnnotatedQueue() as q:
-            qmc_circuit()
-            qml.probs(estimation_wires)
-
-        tape = qml.tape.QuantumScript.from_queue(q)
-        tape = tape.expand(depth=2)
-
-        assert all(
-            not isinstance(op, (qml.MultiControlledX, qml.templates.QFT, qml.tape.QuantumScript))
-            for op in tape.operations
-        )
-
         dev = qml.device("default.qubit", wires=wires + estimation_wires)
-        res = dev.execute(tape)
 
         @qml.qnode(dev)
         def circuit():
+            qmc_circuit()
+            return qml.probs(estimation_wires)
+
+        @qml.qnode(dev)
+        def circuit_expected():
             qml.templates.QuantumMonteCarlo(
                 probs, func, target_wires=wires, estimation_wires=estimation_wires
             )
             return qml.probs(estimation_wires)
 
-        res_expected = circuit()
+        res = circuit()
+        res_expected = circuit_expected()
+
         assert np.allclose(res, res_expected)