Add QNode config for mid-circuit measurement options

doc/introduction/measurements.rst

@@ -524,6 +524,86 @@ Collecting statistics for sequences of mid-circuit measurements is supported wit
     When collecting statistics for a list of mid-circuit measurements, values manipulated using
     arithmetic operators should not be used as this behaviour is not supported.
 
+Configuring mid-circuit measurements
+************************************
+
+As seen above, there are multiple ways in which circuits with mid-circuit measurements can be executed with
+PennyLane. For ease of use, we provide the following configuration options to users when initializing a
+:class:`~pennylane.QNode`:
+
+* ``mcm_method``: To set the method used for applying mid-circuit measurements. Use ``mcm_method="deferred"``
+  to use the deferred measurements principle or ``mcm_method="one-shot"`` to use the one-shot transform as
+  described above. When executing with finite shots, ``mcm_method="one-shot"`` will be the default, and
+  ``mcm_method="deferred"`` otherwise.
+
+  .. warning::
+
+      If the ``mcm_method`` argument is provided, the :func:`~pennylane.defer_measurements` or
+      :func:`~pennylane.dynamic_one_shot` transforms must not be applied directly to the :class:`~pennylane.QNode`
+      as it can lead to incorrect behaviour.
+
+* ``postselect_mode``: To configure how invalid shots are handled when postselecting mid-circuit measurements
+  with finite-shot circuits. Use ``postselect_mode="hw-like"`` to discard invalid samples. In this case, the number
+  of samples that are used for processing results can be less than the total number of shots. If
+  ``postselect_mode="fill-shots"`` is used, then the postselected value will be sampled unconditionally, and all
+  samples will be valid. This is equivalent to sampling until the number of valid samples matches the total number
+  of shots. The default behaviour is ``postselect_mode="hw-like"``.
+
+  .. code-block:: python3
+
+      import pennylane as qml
+      import numpy as np
+
+      dev = qml.device("default.qubit", wires=3, shots=10)
+
+      def circuit(x):
+          qml.RX(x, 0)
+          m0 = qml.measure(0, postselect=1)
+          qml.CNOT([0, 1])
+          return qml.sample(qml.PauliZ(0))
+
+      fill_shots_qnode = qml.QNode(circuit, dev, mcm_method="one-shot", postselect_mode="fill-shots")
+      hw_like_qnode = qml.QNode(circuit, dev, mcm_method="one-shot", postselect_mode="hw-like")
+
+  >>> fill_shots_qnode(np.pi / 2)
+  array([-1., -1., -1., -1., -1., -1., -1., -1., -1., -1.])
+  >>> hw_like_qnode(np.pi / 2)
+  array([-1., -1., -1., -1., -1., -1., -1.])
+
+  .. note::
+
+      When using the ``jax`` interface, ``postselect_mode="hw-like"`` will have different behaviour based on the
+      chosen ``mcm_method``.
+
+      * If ``mcm_method="one-shot"``, invalid shots will not be discarded. Instead, invalid samples will be replaced
+        by ``np.iinfo(np.int32).min``. These invalid samples will not be used for processing final results (like
+        expectation values), but will appear in the ``QNode`` output if samples are requested directly. Consider
+        the circuit below:
+
+        .. code-block:: python3
+
+            import pennylane as qml
+            import jax
+            import jax.numpy as jnp
+
+            dev = qml.device("default.qubit", wires=3, shots=10, seed=jax.random.PRNGKey(123))
+
+            @qml.qnode(dev, postselect_mode="hw-like", mcm_method="one-shot")
+            def circuit(x):
+                qml.RX(x, 0)
+                qml.measure(0, postselect=1)
+                return qml.sample(qml.PauliZ(0))
+
+        >>> x = jnp.array(1.8)
+        >>> f(x)
+        Array([-2.1474836e+09, -1.0000000e+00, -2.1474836e+09, -2.1474836e+09,
+               -1.0000000e+00, -2.1474836e+09, -1.0000000e+00, -2.1474836e+09,
+               -1.0000000e+00, -1.0000000e+00], dtype=float32, weak_type=True)
+
+      * If ``mcm_method="deferred"``, then using ``postselect_mode="hw-like"`` will have the same behaviour as when
+        ``postselect_mode="fill-shots"``. This is due to the limitations of the :func:`~pennylane.defer_measurements`
+        transform, and this behaviour will change in the future to be more consistent with ``mcm_method="one-shot"``.
+
 Changing the number of shots
 ----------------------------
 

doc/releases/changelog-dev.md

@@ -4,6 +4,19 @@
 
 <h3>New features since last release</h3>
 
+* `qml.QNode` and `qml.qnode` now accept two new keyword arguments: `postselect_mode` and `mcm_method`.
+  These keyword arguments can be used to configure how the device should behave when running circuits with
+  mid-circuit measurements.
+  [(#5679)](https://github.com/PennyLaneAI/pennylane/pull/5679)
+
+  * `postselect_mode="hw-like"` will indicate to devices to discard invalid shots when postselecting
+    mid-circuit measurements. Use `postselect_mode="fill-shots"` to unconditionally sample the postselected
+    value, thus making all samples valid. This is equivalent to sampling until the number of valid samples
+    matches the total number of shots.
+  * `mcm_method` will indicate which strategy to use for running circuits with mid-circuit measurements.
+    Use `mcm_method="deferred"` to use the deferred measurements principle, or `mcm_method="one-shot"`
+    to execute once for each shot.
+
 * The `default.tensor` device is introduced to perform tensor network simulation of a quantum circuit.
   [(#5699)](https://github.com/PennyLaneAI/pennylane/pull/5699)
 
@@ -31,7 +44,7 @@
 
 * The `dynamic_one_shot` transform can be compiled with `jax.jit`.
   [(#5557)](https://github.com/PennyLaneAI/pennylane/pull/5557)
-  
+
 * When using `defer_measurements` with postselecting mid-circuit measurements, operations
   that will never be active due to the postselected state are skipped in the transformed
   quantum circuit. In addition, postselected controls are skipped, as they are evaluated

pennylane/devices/__init__.py

@@ -54,6 +54,7 @@
     :toctree: api
 
     ExecutionConfig
+    MCMConfig
     Device
     DefaultQubit
     NullQubit
@@ -146,7 +147,7 @@ def execute(self, circuits, execution_config = qml.devices.DefaultExecutionConfi
 
 """
 
-from .execution_config import ExecutionConfig, DefaultExecutionConfig
+from .execution_config import ExecutionConfig, DefaultExecutionConfig, MCMConfig
 from .device_api import Device
 from .default_qubit import DefaultQubit
 

pennylane/devices/default_qubit.py

@@ -500,7 +500,9 @@ def preprocess(
         transform_program = TransformProgram()
 
         transform_program.add_transform(validate_device_wires, self.wires, name=self.name)
-        transform_program.add_transform(mid_circuit_measurements, device=self)
+        transform_program.add_transform(
+            mid_circuit_measurements, device=self, mcm_config=config.mcm_config
+        )
         transform_program.add_transform(
             decompose,
             stopping_condition=stopping_condition,
@@ -596,14 +598,22 @@ def execute(
                         "interface": interface,
                         "state_cache": self._state_cache,
                         "prng_key": _key,
+                        "postselect_mode": execution_config.mcm_config.postselect_mode,
                     },
                 )
                 for c, _key in zip(circuits, prng_keys)
             )
 
         vanilla_circuits = convert_to_numpy_parameters(circuits)[0]
         seeds = self._rng.integers(2**31 - 1, size=len(vanilla_circuits))
-        simulate_kwargs = [{"rng": _rng, "prng_key": _key} for _rng, _key in zip(seeds, prng_keys)]
+        simulate_kwargs = [
+            {
+                "rng": _rng,
+                "prng_key": _key,
+                "postselect_mode": execution_config.mcm_config.postselect_mode,
+            }
+            for _rng, _key in zip(seeds, prng_keys)
+        ]
 
         with concurrent.futures.ProcessPoolExecutor(max_workers=max_workers) as executor:
             exec_map = executor.map(_simulate_wrapper, vanilla_circuits, simulate_kwargs)
@@ -847,20 +857,23 @@ def _simulate_wrapper(circuit, kwargs):
 
 
 def _adjoint_jac_wrapper(c, debugger=None):
+    c = c.map_to_standard_wires()
     state, is_state_batched = get_final_state(c, debugger=debugger)
     jac = adjoint_jacobian(c, state=state)
     res = measure_final_state(c, state, is_state_batched)
     return res, jac
 
 
 def _adjoint_jvp_wrapper(c, t, debugger=None):
+    c = c.map_to_standard_wires()
     state, is_state_batched = get_final_state(c, debugger=debugger)
     jvp = adjoint_jvp(c, t, state=state)
     res = measure_final_state(c, state, is_state_batched)
     return res, jvp
 
 
 def _adjoint_vjp_wrapper(c, t, debugger=None):
+    c = c.map_to_standard_wires()
     state, is_state_batched = get_final_state(c, debugger=debugger)
     vjp = adjoint_vjp(c, t, state=state)
     res = measure_final_state(c, state, is_state_batched)

pennylane/devices/execution_config.py

@@ -15,11 +15,28 @@
 Contains the :class:`ExecutionConfig` data class.
 """
 from dataclasses import dataclass
-from typing import Optional
+from typing import Optional, Union
 
 from pennylane.workflow import SUPPORTED_INTERFACES
 
 
+@dataclass
+class MCMConfig:
+    """A class to store mid-circuit measurement configurations."""
+
+    mcm_method: Optional[str] = None
+    """Which mid-circuit measurement strategy to use. Use ``deferred`` for the deferred
+    measurements principle and "one-shot" if using finite shots to execute the circuit
+    for each shot separately. If not specified, the device will decide which method to
+    use."""
+
+    postselect_mode: str = "hw-like"
+    """Configuration for handling shots with mid-circuit measurement postselection. If
+    ``"hw-like"``, invalid shots will be discarded and only results for valid shots will
+    be returned. If ``"fill-shots"``, results corresponding to the original number of
+    shots will be returned."""
+
+
 # pylint: disable=too-many-instance-attributes
 @dataclass
 class ExecutionConfig:
@@ -67,6 +84,9 @@ class ExecutionConfig:
     derivative_order: int = 1
     """The derivative order to compute while evaluating a gradient"""
 
+    mcm_config: Union[MCMConfig, dict] = MCMConfig()
+    """Configuration options for handling mid-circuit measurements"""
+
     def __post_init__(self):
         """
         Validate the configured execution options.
@@ -89,5 +109,10 @@ def __post_init__(self):
         if self.gradient_keyword_arguments is None:
             self.gradient_keyword_arguments = {}
 
+        if isinstance(self.mcm_config, dict):
+            self.mcm_config = MCMConfig(**self.mcm_config)
+        elif not isinstance(self.mcm_config, MCMConfig):
+            raise ValueError(f"Got invalid type {type(self.mcm_config)} for 'mcm_config'")
+
 
 DefaultExecutionConfig = ExecutionConfig()

pennylane/devices/preprocess.py

@@ -28,6 +28,8 @@
 from pennylane.typing import Result, ResultBatch
 from pennylane.wires import WireError
 
+from .execution_config import MCMConfig
+
 PostprocessingFn = Callable[[ResultBatch], Union[Result, ResultBatch]]
 
 
@@ -80,7 +82,7 @@ def _operator_decomposition_gen(
 @transform
 def no_sampling(
     tape: qml.tape.QuantumTape, name: str = "device"
-) -> (Sequence[qml.tape.QuantumTape], Callable):
+) -> tuple[Sequence[qml.tape.QuantumTape], Callable]:
     """Raises an error if the tape has finite shots.
 
     Args:
@@ -104,7 +106,7 @@ def no_sampling(
 @transform
 def validate_device_wires(
     tape: qml.tape.QuantumTape, wires: Optional[qml.wires.Wires] = None, name: str = "device"
-) -> (Sequence[qml.tape.QuantumTape], Callable):
+) -> tuple[Sequence[qml.tape.QuantumTape], Callable]:
     """Validates that all wires present in the tape are in the set of provided wires. Adds the
     device wires to measurement processes like :class:`~.measurements.StateMP` that are broadcasted
     across all available wires.
@@ -145,23 +147,29 @@ def validate_device_wires(
 
 @transform
 def mid_circuit_measurements(
-    tape: qml.tape.QuantumTape, device
-) -> (Sequence[qml.tape.QuantumTape], Callable):
+    tape: qml.tape.QuantumTape, device, mcm_config=MCMConfig()
+) -> tuple[Sequence[qml.tape.QuantumTape], Callable]:
     """Provide the transform to handle mid-circuit measurements.
 
     If the tape or device uses finite-shot, use the native implementation (i.e. no transform),
     and use the ``qml.defer_measurements`` transform otherwise.
     """
 
-    if tape.shots:
+    if isinstance(mcm_config, dict):
+        mcm_config = MCMConfig(**mcm_config)
+    mcm_method = mcm_config.mcm_method
+    if mcm_method is None:
+        mcm_method = "one-shot" if tape.shots else "deferred"
+
+    if mcm_method == "one-shot":
         return qml.dynamic_one_shot(tape)
     return qml.defer_measurements(tape, device=device)
 
 
 @transform
 def validate_multiprocessing_workers(
     tape: qml.tape.QuantumTape, max_workers: int, device
-) -> (Sequence[qml.tape.QuantumTape], Callable):
+) -> tuple[Sequence[qml.tape.QuantumTape], Callable]:
     """Validates the number of workers for multiprocessing.
 
     Checks that the CPU is not oversubscribed and warns user if it is,
@@ -220,7 +228,7 @@ def validate_multiprocessing_workers(
 @transform
 def validate_adjoint_trainable_params(
     tape: qml.tape.QuantumTape,
-) -> (Sequence[qml.tape.QuantumTape], Callable):
+) -> tuple[Sequence[qml.tape.QuantumTape], Callable]:
     """Raises a warning if any of the observables is trainable, and raises an error if any
     trainable parameters belong to state-prep operations. Can be used in validating circuits
     for adjoint differentiation.
@@ -256,7 +264,7 @@ def decompose(
     max_expansion: Union[int, None] = None,
     name: str = "device",
     error: Exception = DeviceError,
-) -> (Sequence[qml.tape.QuantumTape], Callable):
+) -> tuple[Sequence[qml.tape.QuantumTape], Callable]:
     """Decompose operations until the stopping condition is met.
 
     Args:
@@ -370,7 +378,7 @@ def validate_observables(
     tape: qml.tape.QuantumTape,
     stopping_condition: Callable[[qml.operation.Operator], bool],
     name: str = "device",
-) -> (Sequence[qml.tape.QuantumTape], Callable):
+) -> tuple[Sequence[qml.tape.QuantumTape], Callable]:
     """Validates the observables and measurements for a circuit.
 
     Args:
@@ -412,7 +420,7 @@ def validate_observables(
 @transform
 def validate_measurements(
     tape: qml.tape.QuantumTape, analytic_measurements=None, sample_measurements=None, name="device"
-) -> (Sequence[qml.tape.QuantumTape], Callable):
+) -> tuple[Sequence[qml.tape.QuantumTape], Callable]:
     """Validates the supported state and sample based measurement processes.
 
     Args: