Allow postselecting on mid-circuit measurements

doc/introduction/measurements.rst

@@ -260,6 +260,9 @@ outcome of such mid-circuit measurements:
         qml.cond(m_0, qml.RY)(y, wires=0)
         return qml.probs(wires=[0])
 
+Deferred measurements
+*********************
+
 A quantum function with mid-circuit measurements (defined using
 :func:`~.pennylane.measure`) and conditional operations (defined using
 :func:`~.pennylane.cond`) can be executed by applying the `deferred measurement
@@ -269,8 +272,12 @@ measurement on qubit 1 yielded ``1`` as an outcome, otherwise doing nothing
 for the ``0`` measurement outcome.
 
 PennyLane implements the deferred measurement principle to transform
-conditional operations with the :func:`~.defer_measurements` quantum
-function transform.
+conditional operations with the :func:`~.pennylane.defer_measurements` quantum
+function transform. The deferred measurement principle provides a natural method
+to simulate the application of mid-circuit measurements and conditional operations
+in a differentiable and device-independent way. Performing true mid-circuit
+measurements and conditional operations is dependent on the quantum hardware and
+PennyLane device capabilities.
 
 .. code-block:: python
 
@@ -290,7 +297,35 @@ The decorator syntax applies equally well:
     def qnode(x, y):
         (...)
 
-Note that we can also specify an outcome when defining a conditional operation:
+Resetting wires
+***************
+
+Wires can be reused as normal after making mid-circuit measurements. Moreover, a measured wire can also be
+reset to the :math:`|0 \rangle` state by setting the ``reset`` keyword argument of :func:`~.pennylane.measure`
+to ``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)
+        qml.PauliX(1)
+        return qml.probs(wires=[1])
+
+Executing this QNode:
+
+>>> func()
+tensor([0., 1.], requires_grad=True)
+
+Conditional operators
+*********************
+
+Users can create conditional operators controlled on mid-circuit measurements using
+:func:`~.pennylane.cond`. We can also specify an outcome when defining a conditional
+operation:
 
 .. code-block:: python
 
@@ -309,30 +344,50 @@ Note that we can also specify an outcome when defining a conditional operation:
 >>> qnode_conditional_op_on_zero(*pars)
 tensor([0.88660045, 0.11339955], requires_grad=True)
 
-Wires can be reused as normal after making mid-circuit measurements. Moreover, a measured wire can also be
-reset to the :math:`|0 \rangle` state by setting the ``reset`` keyword argument of :func:`~.pennylane.measure`
-to ``True``.
+For more examples on applying quantum functions conditionally, refer to the
+:func:`~.pennylane.cond` documentation.
+
+Postselecting mid-circuit measurements
+**************************************
+
+PennyLane also supports postselecting on mid-circuit measurement outcomes by specifying the ``postselect``
+keyword argument of :func:`~.pennylane.measure`. Postselection discards outcomes that do not meet the
+criteria provided by the ``postselect`` argument. For example, specifying ``postselect=1`` on wire 0 would
+be equivalent to projecting the state vector onto the :math:`|1\rangle` state on wire 0:
 
 .. code-block:: python3
 
-    dev = qml.device("default.qubit", wires=3)
+    dev = qml.device("default.qubit")
 
     @qml.qnode(dev)
-    def func():
-        qml.PauliX(1)
-        m_0 = qml.measure(1, reset=True)
-        qml.PauliX(1)
-        return qml.probs(wires=[1])
+    def func(x):
+        qml.RX(x, wires=0)
+        m0 = qml.measure(0, postselect=1)
+        qml.cond(m0, qml.PauliX)(wires=1)
+        return qml.sample(wires=1)
 
-Executing this QNode:
+By postselecting on ``1``, we only consider the ``1`` measurement outcome on wire 0. So, the probability of
+measuring ``1`` on wire 1 after postselection should also be 1. Executing this QNode with 10 shots:
 
->>> func()
-tensor([0., 1.], requires_grad=True)
+>>> func(np.pi / 2, shots=10)
+array([1, 1, 1, 1, 1, 1, 1])
+
+Note that only 7 samples are returned. This is because samples that do not meet the postselection criteria are
+discarded.
+
+.. note::
 
-Statistics can also be collected on mid-circuit measurements along with terminal measurement statistics.
+    Currently, postselection support is only available on :class:`~.pennylane.devices.DefaultQubit`. Using
+    postselection on other devices will raise an error.
+
+Mid-circuit measurement statistics
+**********************************
+
+Statistics can be collected on mid-circuit measurements along with terminal measurement statistics.
 Currently, ``qml.probs``, ``qml.sample``, ``qml.expval``, ``qml.var``, and ``qml.counts`` are supported,
 and can be requested along with other measurements. The devices that currently support collecting such
-statistics are ``"default.qubit"``, ``"default.mixed"``, and ``"default.qubit.legacy"``.
+statistics are :class:`~.pennylane.devices.DefaultQubit`, :class:`~.pennylane.devices.DefaultMixed`, and
+:class:`~.pennylane.devices.DefaultQubitLegacy`.
 
 .. code-block:: python3
 
@@ -351,19 +406,11 @@ Executing this QNode:
 (tensor([0.9267767, 0.0732233], requires_grad=True),
  tensor([0.5, 0.5], requires_grad=True))
 
-Currently, statistics can only be collected for single mid-circuit measurement values. Moreover, any
-measurement values manipulated using boolean or arithmetic operators cannot be used. These can lead to
-unexpected/incorrect behaviour.
-
-The deferred measurement principle provides a natural method to simulate the
-application of mid-circuit measurements and conditional operations in a
-differentiable and device-independent way. Performing true mid-circuit
-measurements and conditional operations is dependent on the
-quantum hardware and PennyLane device capabilities.
-
-For more examples on applying quantum functions conditionally, refer to the
-:func:`~.pennylane.cond` transform.
+.. warning::
 
+    Currently, statistics can only be collected for single mid-circuit measurement values. Moreover, any
+    measurement values manipulated using boolean or arithmetic operators cannot be used. These can lead to
+    unexpected/incorrect behaviour.
 
 Changing the number of shots
 ----------------------------

doc/releases/changelog-dev.md

@@ -91,6 +91,30 @@
   (array(0.6), array([1, 1, 1, 0, 1]))
   ```
 
+* Users can now request postselection after making mid-circuit measurements. They can do so
+  by specifying the `postselect` keyword argument for `qml.measure` as either `0` or `1`,
+  corresponding to the basis states.
+  [(#4604)](https://github.com/PennyLaneAI/pennylane/pull/4604)
+
+  ```python
+  dev = qml.device("default.qubit", wires=3)
+
+  @qml.qnode(dev)
+  def circuit(phi):
+      qml.RX(phi, wires=0)
+      m = qml.measure(0, postselect=1)
+      qml.cond(m, qml.PauliX)(wires=1)
+      return qml.probs(wires=1)
+  ```
+  ```pycon
+  >>> circuit(np.pi)
+  tensor([0., 1.], requires_grad=True)
+  ```
+
+  Here, we measure a probability of one on wire 1 as we postselect on the $|1\rangle$ state on wire
+  0, thus resulting in the circuit being projected onto the state corresponding to the measurement
+  outcome $|1\rangle$ on wire 0.
+
 * Operator transforms `qml.matrix`, `qml.eigvals`, `qml.generator`, and `qml.transforms.to_zx` are updated
   to the new transform program system.
   [(#4573)](https://github.com/PennyLaneAI/pennylane/pull/4573)

pennylane/_device.py

@@ -977,6 +977,9 @@ def check_validity(self, queue, observables):
                     "simulate the application of mid-circuit measurements on this device."
                 )
 
+            if isinstance(o, qml.Projector):
+                raise ValueError(f"Postselection is not supported on the {self.name} device.")
+
             if not self.stopping_condition(o):
                 raise DeviceError(
                     f"Gate {operation_name} not supported on device {self.short_name}"

pennylane/devices/default_qubit.py

@@ -387,7 +387,7 @@ def preprocess(
         config = self._setup_execution_config(execution_config)
         transform_program = TransformProgram()
 
-        transform_program.add_transform(qml.defer_measurements)
+        transform_program.add_transform(qml.defer_measurements, device=self)
         transform_program.add_transform(validate_device_wires, self.wires, name=self.name)
         transform_program.add_transform(
             decompose, stopping_condition=stopping_condition, name=self.name

pennylane/devices/qubit/sampling.py

@@ -179,6 +179,7 @@ def measure_with_samples(
     Returns:
         List[TensorLike[Any]]: Sample measurement results
     """
+
     groups, indices = _group_measurements(mps)
 
     all_res = []
@@ -264,27 +265,37 @@ def _process_single_shot(samples):
             # currently we call sample_state for each shot entry, but it may be
             # better to call sample_state just once with total_shots, then use
             # the shot_range keyword argument
-            samples = sample_state(
-                state,
-                shots=s,
-                is_state_batched=is_state_batched,
-                wires=wires,
-                rng=rng,
-                prng_key=prng_key,
-            )
+            try:
+                samples = sample_state(
+                    state,
+                    shots=s,
+                    is_state_batched=is_state_batched,
+                    wires=wires,
+                    rng=rng,
+                    prng_key=prng_key,
+                )
+            except ValueError as e:
+                if str(e) != "probabilities contain NaN":
+                    raise e
+                samples = qml.math.full((s, len(wires)), 0)
 
             processed_samples.append(_process_single_shot(samples))
 
         return tuple(zip(*processed_samples))
 
-    samples = sample_state(
-        state,
-        shots=shots.total_shots,
-        is_state_batched=is_state_batched,
-        wires=wires,
-        rng=rng,
-        prng_key=prng_key,
-    )
+    try:
+        samples = sample_state(
+            state,
+            shots=shots.total_shots,
+            is_state_batched=is_state_batched,
+            wires=wires,
+            rng=rng,
+            prng_key=prng_key,
+        )
+    except ValueError as e:
+        if str(e) != "probabilities contain NaN":
+            raise e
+        samples = qml.math.full((shots.total_shots, len(wires)), 0)
 
     return _process_single_shot(samples)
 
@@ -352,7 +363,7 @@ def _sum_for_single_shot(s):
         )
         return sum(c * res for c, res in zip(mp.obs.terms()[0], results))
 
-    unsqueezed_results = tuple(_sum_for_single_shot(Shots(s)) for s in shots)
+    unsqueezed_results = tuple(_sum_for_single_shot(type(shots)(s)) for s in shots)
     return [unsqueezed_results] if shots.has_partitioned_shots else [unsqueezed_results[0]]
 
 
@@ -380,7 +391,7 @@ def _sum_for_single_shot(s):
         )
         return sum(results)
 
-    unsqueezed_results = tuple(_sum_for_single_shot(Shots(s)) for s in shots)
+    unsqueezed_results = tuple(_sum_for_single_shot(type(shots)(s)) for s in shots)
     return [unsqueezed_results] if shots.has_partitioned_shots else [unsqueezed_results[0]]
 
 

pennylane/devices/qubit/simulate.py

@@ -14,6 +14,7 @@
 """Simulate a quantum script."""
 # pylint: disable=protected-access
 from numpy.random import default_rng
+import numpy as np
 
 import pennylane as qml
 from pennylane.typing import Result
@@ -46,6 +47,20 @@
 }
 
 
+class _FlexShots(qml.measurements.Shots):
+    """Shots class that allows zero shots."""
+
+    # pylint: disable=super-init-not-called
+    def __init__(self, shots=None):
+        if isinstance(shots, int):
+            self.total_shots = shots
+            self.shot_vector = (qml.measurements.ShotCopies(shots, 1),)
+        else:
+            self.__all_tuple_init__([s if isinstance(s, tuple) else (s, 1) for s in shots])
+
+        self._frozen = True
+
+
 def expand_state_over_wires(state, state_wires, all_wires, is_state_batched):
     """
     Expand and re-order a state given some initial and target wire orders, setting
@@ -83,6 +98,39 @@ def expand_state_over_wires(state, state_wires, all_wires, is_state_batched):
     return qml.math.transpose(state, desired_axes)
 
 
+def _postselection_postprocess(state, is_state_batched, shots):
+    """Update state after projector is applied."""
+    if is_state_batched:
+        raise ValueError(
+            "Cannot postselect on circuits with broadcasting. Use the "
+            "qml.transforms.broadcast_expand transform to split a broadcasted "
+            "tape into multiple non-broadcasted tapes before executing if "
+            "postselection is used."
+        )
+
+    # The floor function is being used here so that a norm very close to zero becomes exactly
+    # equal to zero so that the state can become invalid. This way, execution can continue, and
+    # bad postselection gives results that are invalid rather than results that look valid but
+    # are incorrect.
+    norm = qml.math.floor(qml.math.real(qml.math.norm(state)) * 1e15) * 1e-15
+
+    if shots:
+        # Clip the number of shots using a binomial distribution using the probability of
+        # measuring the postselected state.
+        postselected_shots = (
+            [np.random.binomial(s, float(norm)) for s in shots]
+            if not qml.math.is_abstract(norm)
+            else shots
+        )
+
+        # _FlexShots is used here since the binomial distribution could result in zero
+        # valid samples
+        shots = _FlexShots(postselected_shots)
+
+    state = state / qml.math.cast_like(norm, state)
+    return state, shots
+
+
 def get_final_state(circuit, debugger=None, interface=None):
     """
     Get the final state that results from executing the given quantum script.
@@ -112,6 +160,12 @@ def get_final_state(circuit, debugger=None, interface=None):
     for op in circuit.operations[bool(prep) :]:
         state = apply_operation(op, state, is_state_batched=is_state_batched, debugger=debugger)
 
+        # Handle postselection on mid-circuit measurements
+        if isinstance(op, qml.Projector):
+            state, circuit._shots = _postselection_postprocess(
+                state, is_state_batched, circuit.shots
+            )
+
         # new state is batched if i) the old state is batched, or ii) the new op adds a batch dim
         is_state_batched = is_state_batched or op.batch_size is not None
 
@@ -147,6 +201,7 @@ def measure_final_state(circuit, state, is_state_batched, rng=None, prng_key=Non
     Returns:
         Tuple[TensorLike]: The measurement results
     """
+
     circuit = circuit.map_to_standard_wires()
 
     if not circuit.shots:

pennylane/drawer/draw.py

@@ -267,7 +267,15 @@ def wrapper(*args, **kwargs):
                 tapes = qnode.construct(args, kwargs)
                 if isinstance(qnode.device, qml.devices.Device):
                     program = qnode.transform_program
-                    tapes = program([qnode.tape])
+                    if any(
+                        isinstance(op, qml.measurements.MidMeasureMP)
+                        for op in qnode.tape.operations
+                    ):
+                        tapes, _ = qml.defer_measurements(qnode.tape, device=qnode.device)
+                    else:
+                        tapes = [qnode.tape]
+
+                    tapes = program(tapes)
 
             finally:
                 qnode.expansion_strategy = original_expansion_strategy