Allow postselecting on mid-circuit measurement results

doc/releases/changelog-dev.md

@@ -4,6 +4,10 @@
 
 <h3>New features since last release</h3>
 
+* `qml.measure` now includes a boolean keyword argument `reset` to reset a wire to the
+  $|0\rangle$ computational basis state after measurement.
+  [(#4402)](https://github.com/PennyLaneAI/pennylane/pull/4402/)
+
 * `DefaultQubit2` accepts a `max_workers` argument which controls multiprocessing. 
   A `ProcessPoolExecutor` executes tapes asynchronously
   using a pool of at most `max_workers` processes. If `max_workers` is `None`
@@ -63,6 +67,9 @@ array([False, False])
 
 <h3>Improvements 🛠</h3>
 
+* Wires can now be reused after making a mid-circuit measurement on them.
+  [(#4402)](https://github.com/PennyLaneAI/pennylane/pull/4402/)
+
 * Transform Programs, `qml.transforms.core.TransformProgram`, can now be called on a batch of circuits
   and return a new batch of circuits and a single post processing function.
   [(#4364)](https://github.com/PennyLaneAI/pennylane/pull/4364)

pennylane/_qubit_device.py

@@ -136,6 +136,7 @@ class QubitDevice(Device):
     _real = staticmethod(np.real)
     _size = staticmethod(np.size)
     _ndim = staticmethod(np.ndim)
+    _norm = staticmethod(np.linalg.norm)
 
     @staticmethod
     def _scatter(indices, array, new_dimensions):

pennylane/devices/default_qubit.py

@@ -270,6 +270,7 @@ def apply(self, operations, rotations=None, **kwargs):
 
         # apply the circuit operations
         for i, operation in enumerate(operations):
+            print(self._state.shape)
             if i > 0 and isinstance(operation, (QubitStateVector, BasisState)):
                 raise DeviceError(
                     f"Operation {operation.name} cannot be used after other Operations have already been applied "
@@ -333,12 +334,17 @@ def _apply_operation(self, state, operation):
         matrix = self._asarray(self._get_unitary_matrix(operation), dtype=self.C_DTYPE)
 
         if operation in diagonal_in_z_basis:
-            return self._apply_diagonal_unitary(state, matrix, wires)
-        if len(wires) <= 2:
+            new_state = self._apply_diagonal_unitary(state, matrix, wires)
+        elif len(wires) <= 2:
             # Einsum is faster for small gates
-            return self._apply_unitary_einsum(state, matrix, wires)
+            new_state = self._apply_unitary_einsum(state, matrix, wires)
+        else:
+            new_state = self._apply_unitary(state, matrix, wires)
+
+        if operation.__class__.__name__ in {"Projector", "_BasisStateProjector"}:
+            new_state = new_state / self._norm(new_state)
 
-        return self._apply_unitary(state, matrix, wires)
+        return new_state
 
     def _apply_x(self, state, axes, **kwargs):
         """Applies a PauliX gate by rolling 1 unit along the axis specified in ``axes``.

pennylane/devices/default_qubit_jax.py

@@ -162,6 +162,7 @@ def circuit():
     _const_mul = staticmethod(jnp.multiply)
     _size = staticmethod(jnp.size)
     _ndim = staticmethod(jnp.ndim)
+    _norm = staticmethod(jnp.linalg.norm)
 
     operations = DefaultQubit.operations.union({"ParametrizedEvolution"})
 

pennylane/devices/default_qubit_tf.py

@@ -136,6 +136,7 @@ class DefaultQubitTF(DefaultQubit):
     _stack = staticmethod(tf.stack)
     _size = staticmethod(tf.size)
     _ndim = staticmethod(_ndim_tf)
+    _norm = staticmethod(tf.norm)
 
     @staticmethod
     def _const_mul(constant, array):

pennylane/devices/qubit/apply_operation.py

@@ -190,8 +190,12 @@ def _(op: type_op, state):
         len(op.wires) < EINSUM_OP_WIRECOUNT_PERF_THRESHOLD
         and math.ndim(state) < EINSUM_STATE_WIRECOUNT_PERF_THRESHOLD
     ) or (op.batch_size and is_state_batched):
-        return apply_operation_einsum(op, state, is_state_batched=is_state_batched)
-    return apply_operation_tensordot(op, state, is_state_batched=is_state_batched)
+        new_state = apply_operation_einsum(op, state, is_state_batched=is_state_batched)
+    else:
+        new_state = apply_operation_tensordot(op, state, is_state_batched=is_state_batched)
+
+    if op.__class__.__name__ in {"Projector", "_BasisStateProjector"}:
+        new_state = new_state / math.norm(new_state)
 
 
 @apply_operation.register

pennylane/drawer/drawable_layers.py

@@ -97,11 +97,6 @@ def drawable_layers(ops, wire_map=None):
     # loop over operations
     for op in ops:
         is_mid_measure = is_conditional = False
-        if set(measured_wires.values()).intersection({wire_map[w] for w in op.wires}):
-            raise ValueError(
-                f"Cannot apply operations on {op.wires} as some wires have been measured already."
-            )
-
         if isinstance(op, MidMeasureMP):
             if len(op.wires) > 1:
                 raise ValueError("Cannot draw mid-circuit measurements with more than one wire.")

pennylane/measurements/mid_measure.py

@@ -24,8 +24,10 @@
 from .measurements import MeasurementProcess, MidMeasure
 
 
-def measure(wires):  # TODO: Change name to mid_measure
-    """Perform a mid-circuit measurement in the computational basis on the
+def measure(
+    wires: Wires, reset: Optional[bool] = False, postselect: Optional[int] = None
+):  # TODO: Change name to mid_measure
+    r"""Perform a mid-circuit measurement in the computational basis on the
     supplied qubit.
 
     Measurement outcomes can be obtained and used to conditionally apply
@@ -38,7 +40,7 @@ def measure(wires):  # TODO: Change name to mid_measure
 
     .. code-block:: python3
 
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit", wires=3)
 
         @qml.qnode(dev)
         def func(x, y):
@@ -55,23 +57,48 @@ def func(x, y):
     >>> func(*pars)
     tensor([0.90165331, 0.09834669], requires_grad=True)
 
+    Wires can be reused after measurement. Moreover, measured wires can be reset
+    to the :math:`|0 \rangle` by setting ``reset=True``.
+
+    .. code-block:: python3
+
+        dev = qml.device("default.qubit", wires=3)
+
+        @qml.qnode(dev)
+        def func():
+            qml.PauliX(1)
+            m_0 = qml.measure(1, reset=True)
+            return qml.probs(wires=[1])
+
+    Executing this QNode:
+
+    >>> func()
+    tensor([1., 0.], requires_grad=True)
+
     Mid circuit measurements can be manipulated using the following dunder methods
     ``+``, ``-``, ``*``, ``/``, ``~`` (not), ``&`` (and), ``|`` (or), ``==``, ``<=``,
     ``>=``, ``<``, ``>`` with other mid-circuit measurements or scalars.
 
-    Note:
-        python ``not``, ``and``, ``or``, do not work since these do not have dunder
-        methods. Instead use ``~``, ``&``, ``|``.
+    .. Note ::
+
+        Python ``not``, ``and``, ``or``, do not work since these do not have dunder methods.
+        Instead use ``~``, ``&``, ``|``.
 
     Args:
         wires (Wires): The wire of the qubit the measurement process applies to.
+        reset (Optional[bool]): Whether to reset the wire to the :math:`|0 \rangle`
+            state after measurement.
+        postselect (Optional[int]): The measured computational basis state on which to
+            optionally postselect the circuit. Must be ``0`` or ``1`` if postselection
+            is requested.
 
     Returns:
         MidMeasureMP: measurement process instance
 
     Raises:
         QuantumFunctionError: if multiple wires were specified
     """
+
     wire = Wires(wires)
     if len(wire) > 1:
         raise qml.QuantumFunctionError(
@@ -80,7 +107,7 @@ def func(x, y):
 
     # Create a UUID and a map between MP and MV to support serialization
     measurement_id = str(uuid.uuid4())[:8]
-    mp = MidMeasureMP(wires=wire, id=measurement_id)
+    mp = MidMeasureMP(wires=wire, reset=reset, postselect=postselect, id=measurement_id)
     return MeasurementValue([mp], processing_fn=lambda v: v)
 
 
@@ -90,17 +117,31 @@ def func(x, y):
 class MidMeasureMP(MeasurementProcess):
     """Mid-circuit measurement.
 
+    This class additionally stores information about unknown measurement outcomes in the qubit model.
+    Measurements on a single qubit in the computational basis are assumed.
+
     Please refer to :func:`measure` for detailed documentation.
 
     Args:
         wires (.Wires): The wires the measurement process applies to.
             This can only be specified if an observable was not provided.
-        id (str): custom label given to a measurement instance, can be useful for some applications
-            where the instance has to be identified
+        reset (bool): Whether to reset the wire after measurement.
+        postselect (Optional[int]): The measured computational basis state on which to
+            optionally postselect the circuit. Must be ``0`` or ``1`` if postselection
+            is requested.
+        id (str): Custom label given to a measurement instance.
     """
 
-    def __init__(self, wires: Optional[Wires] = None, id: Optional[str] = None):
-        super().__init__(wires=wires, id=id)
+    def __init__(
+        self,
+        wires: Optional[Wires] = None,
+        reset: Optional[bool] = False,
+        postselect: Optional[int] = None,
+        id: Optional[str] = None,
+    ):
+        super().__init__(wires=Wires(wires), id=id)
+        self.reset = reset
+        self.postselect = postselect
 
     @property
     def return_type(self):

pennylane/ops/qubit/observables.py

@@ -23,7 +23,7 @@
 from scipy.sparse import csr_matrix
 
 import pennylane as qml
-from pennylane.operation import AnyWires, Observable
+from pennylane.operation import AnyWires, Observable, Operation
 from pennylane.wires import Wires
 
 from .matrix_ops import QubitUnitary
@@ -428,7 +428,7 @@ def __copy__(self):
         return copied_op
 
 
-class _BasisStateProjector(Observable):
+class _BasisStateProjector(Observable, Operation):
     # The call signature should be the same as Projector.__new__ for the positional
     # arguments, but with free key word arguments.
     def __init__(self, state, wires, id=None):

pennylane/transforms/condition.py

@@ -67,7 +67,7 @@ def cond(condition, true_fn, false_fn=None):
     :func:`defer_measurements` transform.
 
     Args:
-        condition (.MeasurementValue[bool]): a conditional expression involving a mid-circuit
+        condition (.MeasurementValue): a conditional expression involving a mid-circuit
            measurement value (see :func:`.pennylane.measure`)
         true_fn (callable): The quantum function of PennyLane operation to
             apply if ``condition`` is ``True``
@@ -114,6 +114,8 @@ def qnode(x, y):
         Expressions with boolean logic flow using operators like ``and``,
         ``or`` and ``not`` are not supported as the ``condition`` argument.
 
+        While such statements may not result in errors, they may result in
+        incorrect behaviour.
 
     .. details::
         :title: Usage Details

pennylane/transforms/defer_measurements.py

@@ -16,12 +16,13 @@
 from pennylane.measurements import MidMeasureMP
 from pennylane.ops.op_math import ctrl
 from pennylane.queuing import apply
+from pennylane.tape import QuantumTape
 from pennylane.transforms import qfunc_transform
 from pennylane.wires import Wires
 
 
 @qfunc_transform
-def defer_measurements(tape):
+def defer_measurements(tape: QuantumTape):
     """Quantum function transform that substitutes operations conditioned on
     measurement outcomes to controlled operations.
 
@@ -40,6 +41,13 @@ def defer_measurements(tape):
         that can be controlled as such depends on the set of operations
         supported by the chosen device.
 
+    .. note::
+
+        Devices that inherit `QubitDevice` **must** be initialized with an additional
+        wire for each mid-circuit measurement for `defer_measurements` to transform
+        the quantum tape correctly. Such devices should also be initialized without
+        custom wire labels for correct behaviour.
+
     .. note::
 
         This transform does not change the list of terminal measurements returned by
@@ -54,7 +62,7 @@ def defer_measurements(tape):
         post-measurement states are considered.
 
     Args:
-        qfunc (function): a quantum function
+        tape (.QuantumTape): a quantum tape
 
     **Example**
 
@@ -84,34 +92,65 @@ def qfunc(par):
 
     >>> qml.grad(qnode)(par)
     tensor(-0.49622252, requires_grad=True)
+
+    Reusing and reseting measured wires will work as expected with the
+    ``defer_measurements`` transform:
+
+    .. code-block:: python3
+
+        dev = qml.device("default.qubit", wires=3)
+
+        @qml.qnode(dev)
+        def func(x, y):
+            qml.RY(x, wires=0)
+            qml.CNOT(wires=[0, 1])
+            m_0 = qml.measure(1, reset=True)
+
+            qml.cond(m_0, qml.RY)(y, wires=0)
+            qml.RX(np.pi/4, wires=1)
+            return qml.probs(wires=[0, 1])
+
+    Executing this QNode:
+
+    >>> pars = np.array([0.643, 0.246], requires_grad=True)
+    >>> func(*pars)
+    tensor([0.76960924, 0.13204407, 0.08394415, 0.01440254], requires_grad=True)
     """
-    measured_wires = {}
 
     cv_types = (qml.operation.CVOperation, qml.operation.CVObservable)
     ops_cv = any(isinstance(op, cv_types) for op in tape.operations)
     obs_cv = any(isinstance(getattr(op, "obs", None), cv_types) for op in tape.measurements)
     if ops_cv or obs_cv:
         raise ValueError("Continuous variable operations and observables are not supported.")
 
-    for op in tape:
-        op_wires_measured = set(wire for wire in op.wires if wire in measured_wires.values())
-        if len(op_wires_measured) > 0:
-            raise ValueError(
-                f"Cannot apply operations on {op.wires} as the following wires have been measured already: {op_wires_measured}."
-            )
+    # Current wire in which pre-measurement state will be saved if dev_wires not specified
+    cur_wire = max(tape.wires) + 1
+    new_wires = {}
 
+    for op in tape:
         if isinstance(op, MidMeasureMP):
-            measured_wires[op.id] = op.wires[0]
+            new_wires[op.id] = cur_wire
+            qml.CNOT([op.wires[0], cur_wire])
+
+            if op.postselect is not None:
+                qml.Projector([op.postselect], wires=op.wires[0])
+
+            if op.reset:
+                qml.CNOT([cur_wire, op.wires[0]])
+
+            cur_wire += 1
 
         elif op.__class__.__name__ == "Conditional":
-            _add_control_gate(op, measured_wires)
+            _add_control_gate(op, new_wires)
         else:
             apply(op)
 
+    return tape._qfunc_output  # pylint: disable=protected-access
+
 
-def _add_control_gate(op, measured_wires):
+def _add_control_gate(op, control_wires):
     """Helper function to add control gates"""
-    control = [measured_wires[m.id] for m in op.meas_val.measurements]
+    control = [control_wires[m.id] for m in op.meas_val.measurements]
     for branch, value in op.meas_val._items():  # pylint: disable=protected-access
         if value:
             ctrl(

tests/devices/qubit/test_preprocess.py

@@ -230,7 +230,8 @@ def test_expand_fn_defer_measurement(self):
         expanded_tape = expand_fn(tape)
         expected = [
             qml.Hadamard(0),
-            qml.ops.Controlled(qml.RX(0.123, wires=1), 0),
+            qml.CNOT([0, 2]),
+            qml.ops.Controlled(qml.RX(0.123, wires=1), 2),
         ]
 
         for op, exp in zip(expanded_tape, expected + measurements):