[Capture] default.qubit defer_measurements finite shots execution

doc/releases/changelog-dev.md

@@ -8,6 +8,7 @@
   `qml.defer_measurements` can be executed on `default.qubit`.
   [(#6838)](https://github.com/PennyLaneAI/pennylane/pull/6838)
   [(#6937)](https://github.com/PennyLaneAI/pennylane/pull/6937)
+  [(#6961)](https://github.com/PennyLaneAI/pennylane/pull/6961)
 
   Using `qml.defer_measurements` with program capture enables many new features, including:
   * Significantly richer variety of classical processing on mid-circuit measurement values.

pennylane/devices/default_qubit.py

@@ -955,7 +955,10 @@ def eval_jaxpr(
             key = jax.random.PRNGKey(self._rng.integers(100000))
 
         interpreter = DefaultQubitInterpreter(
-            num_wires=len(self.wires), shots=self.shots.total_shots, key=key
+            num_wires=len(self.wires),
+            shots=self.shots.total_shots,
+            key=key,
+            execution_config=execution_config,
         )
         return interpreter.eval(jaxpr, consts, *args)
 

pennylane/devices/qubit/dq_interpreter.py

@@ -21,6 +21,7 @@
 from pennylane.capture import pause
 from pennylane.capture.base_interpreter import FlattenedHigherOrderPrimitives, PlxprInterpreter
 from pennylane.capture.primitives import adjoint_transform_prim, ctrl_transform_prim, measure_prim
+from pennylane.devices import ExecutionConfig
 from pennylane.measurements import MidMeasureMP, Shots
 from pennylane.ops import adjoint, ctrl
 from pennylane.ops.qubit import Projector
@@ -33,7 +34,7 @@
 from .simulate import _postselection_postprocess  # pylint: disable=protected-access
 
 
-# pylint: disable=attribute-defined-outside-init, access-member-before-definition
+# pylint: disable=attribute-defined-outside-init, access-member-before-definition,too-many-instance-attributes
 class DefaultQubitInterpreter(PlxprInterpreter):
     """Implements a class for interpreting plxpr using python simulation tools.
 
@@ -77,11 +78,15 @@ def __copy__(self):
         return inst
 
     def __init__(
-        self, num_wires: int, shots: int | None = None, key: None | jax.numpy.ndarray = None
+        self,
+        num_wires: int,
+        shots: int | None = None,
+        key: None | jax.numpy.ndarray = None,
+        execution_config: None | ExecutionConfig = None,
     ):
         self.num_wires = num_wires
-        self.shots = Shots(shots)
-        if self.shots.has_partitioned_shots:
+        self.original_shots = Shots(shots)
+        if self.original_shots.has_partitioned_shots:
             raise NotImplementedError(
                 "DefaultQubitInterpreter does not yet support partitioned shots."
             )
@@ -90,6 +95,8 @@ def __init__(
 
         self.initial_key = key
         self.stateref = None
+        self.execution_config = execution_config or ExecutionConfig()
+
         super().__init__()
 
     @property
@@ -107,6 +114,21 @@ def state(self, new_val):
         except TypeError as e:
             raise AttributeError("execution not yet initialized.") from e
 
+    @property
+    def shots(self):
+        """The shots"""
+        try:
+            return self.stateref["shots"]
+        except TypeError as e:
+            raise AttributeError("execution not yet initialized.") from e
+
+    @shots.setter
+    def shots(self, new_val):
+        try:
+            self.stateref["shots"] = new_val
+        except TypeError as e:
+            raise AttributeError("execution not yet initialized.") from e
+
     @property
     def key(self):
         """A jax PRNGKey for random number generation."""
@@ -141,6 +163,7 @@ def setup(self) -> None:
         if self.stateref is None:
             self.stateref = {
                 "state": create_initial_state(range(self.num_wires), like="jax"),
+                "shots": self.original_shots,
                 "key": self.initial_key,
                 "is_state_batched": False,
             }
@@ -156,8 +179,13 @@ def interpret_operation(self, op):
             self.is_state_batched = True
 
         if isinstance(op, Projector):
-            self.state, _ = _postselection_postprocess(
-                self.state, self.is_state_batched, self.shots
+            self.key, new_key = jax.random.split(self.key, 2)
+            self.state, self.shots = _postselection_postprocess(
+                self.state,
+                self.is_state_batched,
+                self.shots,
+                prng_key=new_key,
+                postselect_mode=self.execution_config.mcm_config.postselect_mode,
             )
 
         return op

pennylane/transforms/defer_measurements.py

@@ -30,7 +30,7 @@
 from pennylane.ops.op_math import ctrl
 from pennylane.queuing import QueuingManager
 from pennylane.tape import QuantumScript, QuantumScriptBatch
-from pennylane.transforms import transform
+from pennylane.transforms import TransformError, transform
 from pennylane.typing import PostprocessingFn
 from pennylane.wires import Wires
 
@@ -119,7 +119,7 @@ def _get_plxpr_defer_measurements():
         # pylint: disable=import-outside-toplevel
         import jax
 
-        from pennylane.capture import CaptureError, PlxprInterpreter
+        from pennylane.capture import PlxprInterpreter
         from pennylane.capture.primitives import cond_prim, ctrl_transform_prim, measure_prim
     except ImportError:  # pragma: no cover
         return None, None
@@ -131,13 +131,15 @@ class DeferMeasurementsInterpreter(PlxprInterpreter):
 
         # pylint: disable=unnecessary-lambda-assignment,attribute-defined-outside-init,no-self-use
 
-        def __init__(self, aux_wires):
+        def __init__(self, num_wires):
             super().__init__()
-            self._aux_wires = Wires(aux_wires)
+            self._num_wires = num_wires
 
             # We use a dict here instead of a normal int variable because we want the state to mutate
             # when we interpret higher-order primitives
-            self.state = {"cur_idx": 0}
+            # We store all used wires rather than just the max because if the wires are tracers, then
+            # we can't do comparisons to find the max wire
+            self.state = {"cur_target": num_wires - 1, "used_wires": set()}
 
         def cleanup(self) -> None:
             """Perform any final steps after iterating through all equations.
@@ -147,7 +149,27 @@ def cleanup(self) -> None:
             be used for the target wire of the next mid-circuit measurement's replacement
             :class:`~pennylane.CNOT`.
             """
-            self.state = {"cur_idx": 0}
+            self.state = {"cur_target": self._num_wires - 1, "used_wires": set()}
+
+        def _update_used_wires(self, wires: qml.wires.Wires, cur_target: int):
+            """Update the state with the number of wires that have been used and validate that
+            there is no overlap between the used circuit wires and the mid-circuit measurement
+            target wires.
+
+            Args:
+                wires (pennylane.wires.Wires): wires to add to the set of used wires
+                cur_target (int): target wire to be used for a mid-circuit measurement
+
+            Raises:
+                TransformError: if there is overlap between the used circuit wires and mid-circuit
+                measurement target wires
+            """
+            self.state["used_wires"] |= wires.toset()
+            if self.state["used_wires"].intersection(range(cur_target, self._num_wires)):
+                raise TransformError(
+                    "Too many mid-circuit measurements for the specified number of wires "
+                    "with 'defer_measurements'."
+                )
 
         def interpret_dynamic_operation(self, data, struct, inds):
             """Interpret an operation that uses mid-circuit measurement outcomes as parameters.
@@ -161,13 +183,14 @@ def interpret_dynamic_operation(self, data, struct, inds):
                 struct (PyTreeDef): Pytree structure of the operator
                 inds (Sequence[int]): Indices of mid-circuit measurement values in ``data``
 
-            Returns:
-                None
+            Raises:
+                TransformError: if there is overlap between the used circuit wires and mid-circuit
+                measurement target wires
             """
             if len(inds) > 1:
-                raise CaptureError(
+                raise TransformError(
                     "Cannot create operations with multiple parameters based on "
-                    "mid-circuit measurements."
+                    "mid-circuit measurements with 'defer_measurements'."
                 )
 
             idx = inds[0]
@@ -192,10 +215,14 @@ def interpret_operation(self, op: "qml.operation.Operator"):
             See also: :meth:`~.interpret_operation_eqn`.
 
             """
+            # Range for comparison is [cur_target + 1, num_wires) because cur_target
+            # is the _next_ wire to be used for an MCM. We want to check if the used
+            # wires overlap with the already applied MCMs.
+            self._update_used_wires(op.wires, self.state["cur_target"] + 1)
+
             # We treat operators with operators based on mid-circuit measurement values
             # separately, and otherwise default to the standard behaviour
             data, struct = jax.tree_util.tree_flatten(op)
-
             mcm_data_inds = []
             for i, d in enumerate(data):
                 if isinstance(d, MeasurementValue):
@@ -219,9 +246,14 @@ def interpret_measurement(self, measurement: "qml.measurement.MeasurementProcess
                 kwargs = {"wires": measurement.wires, "eigvals": measurement.eigvals()}
                 if isinstance(measurement, CountsMP):
                     kwargs["all_outcomes"] = measurement.all_outcomes
-
                 measurement = type(measurement)(**kwargs)
 
+            else:
+                # Range for comparison is [cur_target + 1, num_wires) because cur_target
+                # is the _next_ wire to be used for an MCM. We want to check if the used
+                # wires overlap with the already applied MCMs.
+                self._update_used_wires(measurement.wires, self.state["cur_target"] + 1)
+
             return super().interpret_measurement(measurement)
 
         def resolve_mcm_values(
@@ -330,24 +362,19 @@ def eval(self, jaxpr: "jax.core.Jaxpr", consts: list, *args) -> list:
 
     @DeferMeasurementsInterpreter.register_primitive(measure_prim)
     def _(self, wires, reset=False, postselect=None):
-        if self.state["cur_idx"] >= len(self._aux_wires):
-            raise ValueError(
-                "Not enough auxiliary wires provided to apply specified number of mid-circuit "
-                "measurements using qml.defer_measurements."
-            )
+        cur_target = self.state["cur_target"]
+        # Range for comparison is [cur_target, num_wires) because cur_target
+        # is the _current_ wire to be used for an MCM.
+        self._update_used_wires(Wires(wires), cur_target)
 
         # Using type.__call__ instead of normally constructing the class prevents
         # the primitive corresponding to the class to get binded. We do not want the
         # MidMeasureMP's primitive to get recorded.
         meas = type.__call__(
-            MidMeasureMP,
-            Wires(self._aux_wires[self.state["cur_idx"]]),
-            reset=reset,
-            postselect=postselect,
-            id=self.state["cur_idx"],
+            MidMeasureMP, Wires(cur_target), reset=reset, postselect=postselect, id=str(cur_target)
         )
 
-        cnot_wires = (wires, self._aux_wires[self.state["cur_idx"]])
+        cnot_wires = (wires, cur_target)
         if postselect is not None:
             qml.Projector(jax.numpy.array([postselect]), wires=wires)
 
@@ -358,7 +385,7 @@ def _(self, wires, reset=False, postselect=None):
             elif postselect == 1:
                 qml.PauliX(wires=wires)
 
-        self.state["cur_idx"] += 1
+        self.state["cur_target"] -= 1
         return MeasurementValue([meas], lambda x: x)
 
     @DeferMeasurementsInterpreter.register_primitive(cond_prim)
@@ -390,6 +417,7 @@ def _(
                 control_wires = Wires([m.wires[0] for m in condition.measurements])
 
                 for branch, value in condition.items():
+                    # When reduce_postselected is True, some branches can be ()
                     cur_consts = invals[consts_slices[i]]
                     qml.cond(value, ctrl_transform_prim.bind)(
                         *cur_consts,
@@ -407,9 +435,9 @@ def _(
     def defer_measurements_plxpr_to_plxpr(jaxpr, consts, targs, tkwargs, *args):
         """Function for applying the ``defer_measurements`` transform on plxpr."""
 
-        if not tkwargs.get("aux_wires", None):
+        if not tkwargs.get("num_wires", None):
             raise ValueError(
-                "'aux_wires' argument for qml.defer_measurements must be provided "
+                "'num_wires' argument for qml.defer_measurements must be provided "
                 "when qml.capture.enabled() is True."
             )
         if tkwargs.pop("reduce_postselected", False):
@@ -444,7 +472,7 @@ def defer_measurements(
     tape: QuantumScript,
     reduce_postselected: bool = True,
     allow_postselect: bool = True,
-    aux_wires: Optional[Union[int, Sequence[int], Wires]] = None,
+    num_wires: Optional[int] = None,
 ) -> tuple[QuantumScriptBatch, PostprocessingFn]:
     """Quantum function transform that substitutes operations conditioned on
     measurement outcomes to controlled operations.
@@ -503,7 +531,7 @@ def defer_measurements(
         allow_postselect (bool): Whether postselection is allowed. In order to perform postselection
             with ``defer_measurements``, the device must support the :class:`~.Projector` operation.
             Defaults to ``True``. This is currently ignored if program capture is enabled.
-        aux_wires (Sequence): Optional sequence of wires to use to map mid-circuit measurements. This is
+        num_wires (int): Optional argument to specify the total number of circuit wires. This is
             only used if program capture is enabled.
 
     Returns:
@@ -619,9 +647,41 @@ def node(x):
         :title: Deferred measurements with program capture
 
         ``qml.defer_measurements`` can be applied to callables when program capture is enabled. To do so,
-        the ``aux_wires`` argument must be provided, which should be a sequence of integers to be used
-        as the target wires for transforming mid-circuit measurements. With program capture enabled, some
-        new features, as well as new restrictions are introduced, that are detailed below:
+        the ``num_wires`` argument must be provided, which should be an integer corresponding to the total
+        number of available wires. For ``m`` mid-circuit measurements, ``range(num_wires - m, num_wires)``
+        will be the range of wires used to map mid-circuit measurements to ``CNOT`` gates.
+
+        .. warning::
+
+            While the transform includes validation to avoid overlap between wires of the original
+            circuit and mid-circuit measurement target wires, if any wires of the original ciruit
+            are traced, i.e. dependent on dynamic arguments to the transformed workflow, the
+            validation may not catch overlaps. Consider the following example:
+
+            .. code-block:: python
+
+                from functools import partial
+                import jax
+
+                qml.capture.enable()
+
+                @qml.capture.expand_plxpr_transforms
+                @partial(qml.defer_measurements, num_wires=1)
+                def f(n):
+                    qml.measure(n)
+
+            >>> jax.make_jaxpr(f)(0)
+            { lambda ; a:i64[]. let _:AbstractOperator() = CNOT[n_wires=2] a 0 in () }
+
+            The circuit gets transformed without issue because the concrete value of the measured wire
+            is unknown. However, execution with n = 0 would raise an error, as the CNOT wires would
+            be (0, 0).
+
+            Thus, users must by cautious when transforming a circuit. **For ``n`` total wires and
+            ``c`` circuit wires, the number of mid-circuit measurements allowed is ``n - c``.**
+
+        Using ``defer_measurements`` with program capture enabled introduces new features and
+        restrictions:
 
         **New features**
 
@@ -645,7 +705,7 @@ def node(x):
               qml.capture.enable()
 
               @qml.capture.expand_plxpr_transforms
-              @partial(qml.defer_measurements, aux_wires=list(range(5, 10)))
+              @partial(qml.defer_measurements, num_wires=10)
               def f():
                   m0 = qml.measure(0)
 
@@ -655,21 +715,21 @@ def f():
 
           >>> jax.make_jaxpr(f)()
           { lambda ; . let
-              _:AbstractOperator() = CNOT[n_wires=2] 0 5
+              _:AbstractOperator() = CNOT[n_wires=2] 0 9
               a:f64[] = mul 0.0 3.141592653589793
               b:f64[] = sin a
               c:AbstractOperator() = RX[n_wires=1] b 0
               _:AbstractOperator() = Controlled[
                 control_values=(False,)
                 work_wires=Wires([])
-              ] c 5
+              ] c 9
               d:f64[] = mul 1.0 3.141592653589793
               e:f64[] = sin d
               f:AbstractOperator() = RX[n_wires=1] e 0
               _:AbstractOperator() = Controlled[
                 control_values=(True,)
                 work_wires=Wires([])
-              ] f 5
+              ] f 9
               g:AbstractOperator() = PauliZ[n_wires=1] 0
               h:AbstractMeasurement(n_wires=None) = expval_obs g
             in (h,) }