Repository: PennyLaneAI/pennylane
Pull request #6838 (2025-01-15 21:40:45)

[Capture] `defer_measurements` is plxpr compatible
[sc-82053]

This PR adds support to capture and apply `defer_measurements` natively in plxpr:
* Add a `DeferMeasurementsInterpreter` that applies the transform to plxpr. This transform can:
  * Allow qubit reuse/reset, classical control, MCM statistics, postselection.
  * True arbitrary classical processing of MCM values. Users can apply _any_ `jax.numpy` arithmetic functions that take a scalar input to MCM values, which is not available with program capture disabled.
  * Allows using `elif` functions with `qml.cond` and MCM predicates, which is not available with program capture disabled.
* Add a `plxpr_transform` attribute to `defer_measurements`. Thus, applying `qml.defer_measurements` on a qfunc can be expanded using `qml.capture.expand_plxpr_transforms`.

* Some other miscellaneous changes to support the implementation.
  * I made `MeasurementValue._items()` -> `items()`, and `_postselected_items()` -> `postselected_items()`
  * Updated `PlxprInterpreter`'s custom primitive handlers to use the consts of the transformed jaxpr rather than the original consts

github-actions[bot]: Hello. You may have forgotten to update the changelog!
Please edit [`doc/releases/changelog-dev.md`](/PennyLaneAI/pennylane/blob/master/doc/releases/changelog-dev.md) with:
* A one-to-two sentence description of the change. You may include a small working example for new features.
* A link back to this PR.
* Your name (or GitHub username) in the contributors section.
codecov[bot]: ## [Codecov](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/6838?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
Attention: Patch coverage is `21.73913%` with `72 lines` in your changes missing coverage. Please review.
> Project coverage is 99.43%. Comparing base [(`807bc4c`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/807bc4c4d1df248951550fe7b7264c9fa83c2668?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) to head [(`ae67a0b`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/ae67a0bf1d412ad0dfe48e4c78b35f9224b1607e?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).
> Report is 11 commits behind head on master.

| [Files with missing lines](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/6838?dropdown=coverage&src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | Patch % | Lines |
|---|---|---|
| [pennylane/transforms/defer\_measurements.py](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/6838?src=pr&el=tree&filepath=pennylane%2Ftransforms%2Fdefer_measurements.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3RyYW5zZm9ybXMvZGVmZXJfbWVhc3VyZW1lbnRzLnB5) | 21.73% | [72 Missing :warning: ](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/6838?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) |

<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##           master    #6838      +/-   ##
==========================================
- Coverage   99.60%   99.43%   -0.17%     
==========================================
  Files         476      476              
  Lines       45237    45289      +52     
==========================================
- Hits        45060    45035      -25     
- Misses        177      254      +77     
```

</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/6838?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).


Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 0f74ea41d..1651966ad 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -4,6 +4,31 @@
 
 <h3>New features since last release</h3>
 
+* `qml.defer_measurements` can now be used with program capture enabled.
+  [(#6838)](https://github.com/PennyLaneAI/pennylane/pull/6838)
+
+  Using `qml.defer_measurements` with program capture enables many new features, including:
+  * Significantly richer variety of classical processing on mid-circuit measurement values.
+  * Using mid-circuit measurement values as gate parameters.
+
+  Functions such as the following can now be captured:
+
+  ```python
+  import jax.numpy as jnp
+
+  qml.capture.enable()
+
+  def f(x):
+      m0 = qml.measure(0)
+      m1 = qml.measure(0)
+      a = jnp.sin(0.5 * jnp.pi * m0)
+      phi = a - (m1 + 1) ** 4
+
+      qml.s_prod(x, qml.RZ(phi, 0))
+
+      return qml.expval(qml.Z(0))
+  ```
+
 * Added class `qml.capture.transforms.UnitaryToRotInterpreter` that decomposes `qml.QubitUnitary` operators 
   following the same API as `qml.transforms.unitary_to_rot` when experimental program capture is enabled.
   [(#6916)](https://github.com/PennyLaneAI/pennylane/pull/6916)
diff --git a/pennylane/measurements/__init__.py b/pennylane/measurements/__init__.py
index b7eb1fcfb..3ffbd3fc2 100644
--- a/pennylane/measurements/__init__.py
+++ b/pennylane/measurements/__init__.py
@@ -289,7 +289,13 @@ from .measurements import (
     Variance,
     VnEntropy,
 )
-from .mid_measure import MeasurementValue, MidMeasureMP, measure, find_post_processed_mcms
+from .mid_measure import (
+    MeasurementValue,
+    MidMeasureMP,
+    measure,
+    find_post_processed_mcms,
+    get_mcm_predicates,
+)
 from .mutual_info import MutualInfoMP, mutual_info
 from .probs import ProbabilityMP, probs
 from .purity import PurityMP, purity
diff --git a/pennylane/measurements/counts.py b/pennylane/measurements/counts.py
index 549aafb9b..31672fdb2 100644
--- a/pennylane/measurements/counts.py
+++ b/pennylane/measurements/counts.py
@@ -138,7 +138,11 @@ def counts(
         return CountsMP(obs=op, all_outcomes=all_outcomes)
 
     if isinstance(op, Sequence):
-        if not all(isinstance(o, MeasurementValue) and len(o.measurements) == 1 for o in op):
+        if not all(
+            qml.math.is_abstract(o)
+            or (isinstance(o, MeasurementValue) and len(o.measurements) == 1)
+            for o in op
+        ):
             raise qml.QuantumFunctionError(
                 "Only sequences of single MeasurementValues can be passed with the op argument. "
                 "MeasurementValues manipulated using arithmetic operators cannot be used when "
diff --git a/pennylane/measurements/measurements.py b/pennylane/measurements/measurements.py
index c9e6bde44..b5c05022c 100644
--- a/pennylane/measurements/measurements.py
+++ b/pennylane/measurements/measurements.py
@@ -423,9 +423,10 @@ class MeasurementProcess(ABC, metaclass=qml.capture.ABCCaptureMeta):
         """
         if self.mv is not None:
             if getattr(self.mv, "name", None) == "MeasurementValue":
-                # Indexing a MeasurementValue gives the output of the processing function
-                # for the binary number corresponding to the index.
-                return qml.math.asarray([self.mv[i] for i in range(2 ** len(self.wires))])
+                # "Eigvals" should be the processed values for all branches of a MeasurementValue
+                _, processed_values = tuple(zip(*self.mv.items()))
+                interface = qml.math.get_deep_interface(processed_values)
+                return qml.math.asarray(processed_values, like=interface)
             return qml.math.arange(0, 2 ** len(self.wires), 1)
 
         if self.obs is not None:
diff --git a/pennylane/measurements/mid_measure.py b/pennylane/measurements/mid_measure.py
index f518f25ba..de315e2b2 100644
--- a/pennylane/measurements/mid_measure.py
+++ b/pennylane/measurements/mid_measure.py
@@ -394,14 +394,14 @@ class MeasurementValue(Generic[T]):
         self.measurements = measurements
         self.processing_fn = processing_fn
 
-    def _items(self):
+    def items(self):
         """A generator representing all the possible outcomes of the MeasurementValue."""
         num_meas = len(self.measurements)
         for i in range(2**num_meas):
             branch = tuple(int(b) for b in f"{i:0{num_meas}b}")
             yield branch, self.processing_fn(*branch)
 
-    def _postselected_items(self):
+    def postselected_items(self):
         """A generator representing all the possible outcomes of the MeasurementValue,
         taking postselection into account."""
         # pylint: disable=stop-iteration-return
@@ -573,6 +573,31 @@ class MeasurementValue(Generic[T]):
         return f"MeasurementValue(wires={self.wires.tolist()})"
 
 
+def get_mcm_predicates(conditions: tuple[MeasurementValue]) -> list[MeasurementValue]:
+    r"""Function to make mid-circuit measurement predicates mutually exclusive.
+
+    The ``conditions`` are predicates to the ``if`` and ``elif`` branches of ``qml.cond``.
+    This function updates all the ``MeasurementValue``\ s in ``conditions`` such that
+    reconciling the correct branch is never ambiguous.
+
+    Args:
+        conditions (Sequence[MeasurementValue]): Sequence containing predicates for ``if``
+            and all ``elif`` branches of a function decorated with :func:`~pennylane.cond`.
+
+    Returns:
+        Sequence[MeasurementValue]: Updated sequence of mutually exclusive predicates.
+    """
+    new_conds = [conditions[0]]
+    false_cond = ~conditions[0]
+
+    for c in conditions[1:]:
+        new_conds.append(false_cond & c)
+        false_cond = false_cond & ~c
+
+    new_conds.append(false_cond)
+    return new_conds
+
+
 def find_post_processed_mcms(circuit):
     """Return the subset of mid-circuit measurements which are required for post-processing.
 
diff --git a/pennylane/measurements/sample.py b/pennylane/measurements/sample.py
index fe0a18fe7..d542de95c 100644
--- a/pennylane/measurements/sample.py
+++ b/pennylane/measurements/sample.py
@@ -165,7 +165,9 @@ class SampleMP(SampleMeasurement):
             return
 
         if isinstance(obs, Sequence):
-            if not all(isinstance(o, MeasurementValue) and len(o.measurements) == 1 for o in obs):
+            if not all(
+                isinstance(o, MeasurementValue) and len(o.measurements) == 1 for o in obs
+            ) and not all(qml.math.is_abstract(o) for o in obs):
                 raise qml.QuantumFunctionError(
                     "Only sequences of single MeasurementValues can be passed with the op "
                     "argument. MeasurementValues manipulated using arithmetic operators cannot be "
diff --git a/pennylane/ops/op_math/condition.py b/pennylane/ops/op_math/condition.py
index de7958fcf..d805fa992 100644
--- a/pennylane/ops/op_math/condition.py
+++ b/pennylane/ops/op_math/condition.py
@@ -22,7 +22,7 @@ import pennylane as qml
 from pennylane import QueuingManager
 from pennylane.capture.flatfn import FlatFn
 from pennylane.compiler import compiler
-from pennylane.measurements import MeasurementValue
+from pennylane.measurements import MeasurementValue, get_mcm_predicates
 from pennylane.operation import AnyWires, Operation, Operator
 from pennylane.ops.op_math.symbolicop import SymbolicOp
 
@@ -668,19 +668,6 @@ def _validate_abstract_values(
             )
 
 
-def _get_mcm_predicates(conditions: tuple[MeasurementValue]) -> list[MeasurementValue]:
-    """Helper function to update predicates with mid-circuit measurements"""
-    new_conds = [conditions[0]]
-    false_cond = ~conditions[0]
-
-    for c in conditions[1:]:
-        new_conds.append(false_cond & c)
-        false_cond = false_cond & ~c
-
-    new_conds.append(false_cond)
-    return new_conds
-
-
 @functools.lru_cache
 def _get_cond_qfunc_prim():
     """Get the cond primitive for quantum functions."""
@@ -711,7 +698,7 @@ def _get_cond_qfunc_prim():
                     "Cannot use qml.cond with a combination of mid-circuit measurements "
                     "and other classical conditions as predicates."
                 )
-            conditions = _get_mcm_predicates(mcm_conditions)
+            conditions = get_mcm_predicates(mcm_conditions)
 
         for pred, jaxpr, const_slice in zip(conditions, jaxpr_branches, consts_slices):
             consts = all_args[const_slice]
diff --git a/pennylane/tape/plxpr_conversion.py b/pennylane/tape/plxpr_conversion.py
index 08254fa5c..00d3bc855 100644
--- a/pennylane/tape/plxpr_conversion.py
+++ b/pennylane/tape/plxpr_conversion.py
@@ -24,25 +24,11 @@ from pennylane.capture.primitives import (
     ctrl_transform_prim,
     measure_prim,
 )
-from pennylane.measurements import MeasurementValue
+from pennylane.measurements import MeasurementValue, get_mcm_predicates
 
 from .qscript import QuantumScript
 
 
-def _get_mcm_predicates(conditions: tuple[MeasurementValue]) -> list[MeasurementValue]:
-    """Helper function to update predicates with mid-circuit measurements"""
-    # copy from ops.op_math.condition.py
-    new_conds = [conditions[0]]
-    false_cond = ~conditions[0]
-
-    for c in conditions[1:]:
-        new_conds.append(false_cond & c)
-        false_cond = false_cond & ~c
-
-    new_conds.append(false_cond)
-    return new_conds
-
-
 class CollectOpsandMeas(PlxprInterpreter):
     """Collect the dropped operations and measurements in a plxpr. Used by ``convert_to_tape``.
 
@@ -163,7 +149,7 @@ def _(self, *all_args, jaxpr_branches, consts_slices, args_slice):
                 "Cannot use qml.cond with a combination of mid-circuit measurements "
                 "and other classical conditions as predicates."
             )
-        conditions = _get_mcm_predicates(mcm_conditions)
+        conditions = get_mcm_predicates(mcm_conditions)
 
     for pred, jaxpr, const_slice in zip(conditions, jaxpr_branches, consts_slices):
         consts = all_args[const_slice]
diff --git a/pennylane/transforms/defer_measurements.py b/pennylane/transforms/defer_measurements.py
index 02f299507..6210b1980 100644
--- a/pennylane/transforms/defer_measurements.py
+++ b/pennylane/transforms/defer_measurements.py
@@ -13,8 +13,20 @@
 # limitations under the License.
 """Code for the tape transform implementing the deferred measurement principle."""
 
+from functools import lru_cache, partial
+from numbers import Number
+from typing import Callable, Optional, Sequence, Union
+from warnings import warn
+
 import pennylane as qml
-from pennylane.measurements import CountsMP, MeasurementValue, MidMeasureMP, ProbabilityMP, SampleMP
+from pennylane.measurements import (
+    CountsMP,
+    MeasurementValue,
+    MidMeasureMP,
+    ProbabilityMP,
+    SampleMP,
+    get_mcm_predicates,
+)
 from pennylane.ops.op_math import ctrl
 from pennylane.queuing import QueuingManager
 from pennylane.tape import QuantumScript, QuantumScriptBatch
@@ -101,9 +113,333 @@ def null_postprocessing(results):
     return results[0]
 
 
-@transform
+@lru_cache
+def _get_plxpr_defer_measurements():
+    try:
+        # pylint: disable=import-outside-toplevel
+        import jax
+
+        from pennylane.capture import CaptureError, PlxprInterpreter
+        from pennylane.capture.primitives import cond_prim, ctrl_transform_prim, measure_prim
+    except ImportError:  # pragma: no cover
+        return None, None
+
+    # pylint: disable=redefined-outer-name
+
+    class DeferMeasurementsInterpreter(PlxprInterpreter):
+        """Interpreter for applying the defer_measurements transform to plxpr."""
+
+        # pylint: disable=unnecessary-lambda-assignment,attribute-defined-outside-init,no-self-use
+
+        def __init__(self, aux_wires):
+            super().__init__()
+            self._aux_wires = Wires(aux_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}
+
+        def cleanup(self) -> None:
+            """Perform any final steps after iterating through all equations.
+
+            This method resets the internal ``state``, specifically the ``cur_idx`` entry.
+            ``cur_idx`` tracks the index within the auxiliary wires, determining which will
+            be used for the target wire of the next mid-circuit measurement's replacement
+            :class:`~pennylane.CNOT`.
+            """
+            self.state = {"cur_idx": 0}
+
+        def interpret_dynamic_operation(self, data, struct, inds):
+            """Interpret an operation that uses mid-circuit measurement outcomes as parameters.
+
+            * This will not work if mid-circuit measurement values are used to specify
+              operator wires.
+            * This will not work if more than one parameter uses mid-circuit measurement values.
+
+            Args:
+                data (TensorLike): Flattened data of the operator
+                struct (PyTreeDef): Pytree structure of the operator
+                inds (Sequence[int]): Indices of mid-circuit measurement values in ``data``
+
+            Returns:
+                None
+            """
+            if len(inds) > 1:
+                raise CaptureError(
+                    "Cannot create operations with multiple parameters based on "
+                    "mid-circuit measurements."
+                )
+
+            idx = inds[0]
+            mv = data[idx]
+            for branch, value in mv.items():
+                data[idx] = value
+                op = jax.tree_util.tree_unflatten(struct, data)
+                qml.ctrl(op, mv.wires, control_values=branch)
+
+        def interpret_operation(self, op: "qml.operation.Operator"):
+            """Interpret a PennyLane operation instance.
+
+            Args:
+                op (Operator): a pennylane operator instance
+
+            Returns:
+                Any
+
+            This method is only called when the operator's output is a dropped variable,
+            so the output will not affect later equations in the circuit.
+
+            See also: :meth:`~.interpret_operation_eqn`.
+
+            """
+            # 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):
+                    mcm_data_inds.append(i)
+
+            if mcm_data_inds:
+                return self.interpret_dynamic_operation(data, struct, mcm_data_inds)
+
+            return jax.tree_util.tree_unflatten(struct, data)
+
+        def interpret_measurement(self, measurement: "qml.measurement.MeasurementProcess"):
+            """Interpret a measurement process instance.
+
+            Args:
+                measurement (MeasurementProcess): a measurement instance.
+
+            See also :meth:`~.interpret_measurement_eqn`.
+
+            """
+            if measurement.mv is not None:
+                kwargs = {"wires": measurement.wires, "eigvals": measurement.eigvals()}
+                if isinstance(measurement, CountsMP):
+                    kwargs["all_outcomes"] = measurement.all_outcomes
+
+                measurement = type(measurement)(**kwargs)
+
+            return super().interpret_measurement(measurement)
+
+        def resolve_mcm_values(
+            self,
+            primitive: "jax.core.Primitive",
+            subfuns: Sequence[Callable],
+            invals: Sequence[Union[MeasurementValue, Number]],
+            params: dict,
+        ) -> MeasurementValue:
+            """Create a ``MeasurementValue`` that captures all classical processing of the
+            input ``eqn`` in its ``processing_fn``.
+
+            Args:
+                primitive (jax.core.Primitive): Jax primitive
+                subfuns (Sequence[Callable]): Callable positional arguments to the primitive.
+                    These are created by pre-processing jaxpr equation parameters.
+                invals (Sequence[Union[MeasurementValue, Number]]): Inputs to the primitive
+                params (dict): Keyword arguments to the primitive
+
+            Returns:
+                MeasurementValue: ``MeasurementValue`` containing classical processing information
+                for applying the input equation to mid-circuit measurement outcomes.
+            """
+            # pylint: disable=protected-access
+            # MeasurementValue._apply is for applying a new operation to the current
+            # MeasurementValue
+            # It is used when either:
+            # 1. A unary operation is performed on one MeasurementValue, e.g., ~m0
+            # 2. A binary operation is performed on a MeasurementValue and a scalar,
+            #    e.g., m0 + 1
+            #
+            # MeasurementValue._transform_bin_op is for applying a binary operation
+            # to two MeasurementValues
+            assert len(invals) <= 2
+            processing_fn = partial(primitive.bind, *subfuns, **params)
+
+            # One MeasurementValue
+            if len(invals) == 1:
+                m0 = invals[0]
+                return m0._apply(processing_fn)
+
+            # Two MeasurementValues
+            if all(isinstance(inval, MeasurementValue) for inval in invals):
+                m0, m1 = invals
+                return m0._transform_bin_op(processing_fn, m1)
+
+            # One MeasurementValue, one number
+            [m0, other] = invals if isinstance(invals[0], MeasurementValue) else invals[::-1]
+            return m0._apply(lambda x: processing_fn(x, other))
+
+        def eval(self, jaxpr: "jax.core.Jaxpr", consts: list, *args) -> list:
+            """Evaluate a jaxpr.
+
+            Args:
+                jaxpr (jax.core.Jaxpr): the jaxpr to evaluate
+                consts (list[TensorLike]): the constant variables for the jaxpr
+                *args (tuple[TensorLike]): The arguments for the jaxpr.
+
+            Returns:
+                list[TensorLike]: the results of the execution.
+
+            """
+            self._env = {}
+            self.setup()
+
+            for arg, invar in zip(args, jaxpr.invars, strict=True):
+                self._env[invar] = arg
+            for const, constvar in zip(consts, jaxpr.constvars, strict=True):
+                self._env[constvar] = const
+
+            for eqn in jaxpr.eqns:
+                primitive = eqn.primitive
+                custom_handler = self._primitive_registrations.get(primitive, None)
+
+                if custom_handler:
+                    invals = [self.read(invar) for invar in eqn.invars]
+                    outvals = custom_handler(self, *invals, **eqn.params)
+                elif getattr(primitive, "prim_type", "") == "operator":
+                    outvals = self.interpret_operation_eqn(eqn)
+                elif getattr(primitive, "prim_type", "") == "measurement":
+                    outvals = self.interpret_measurement_eqn(eqn)
+                else:
+                    invals = [self.read(invar) for invar in eqn.invars]
+                    subfuns, params = primitive.get_bind_params(eqn.params)
+                    if any(isinstance(inval, MeasurementValue) for inval in invals):
+                        outvals = self.resolve_mcm_values(primitive, subfuns, invals, params)
+                    else:
+                        outvals = primitive.bind(*subfuns, *invals, **params)
+
+                if not primitive.multiple_results:
+                    outvals = [outvals]
+                for outvar, outval in zip(eqn.outvars, outvals, strict=True):
+                    self._env[outvar] = outval
+
+            # Read the final result of the Jaxpr from the environment
+            outvals = []
+            for var in jaxpr.outvars:
+                outval = self.read(var)
+                if isinstance(outval, qml.operation.Operator):
+                    outvals.append(self.interpret_operation(outval))
+                else:
+                    outvals.append(outval)
+            self.cleanup()
+            self._env = {}
+            return outvals
+
+    @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."
+            )
+
+        # 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"],
+        )
+
+        cnot_wires = (wires, self._aux_wires[self.state["cur_idx"]])
+        if postselect is not None:
+            qml.Projector(jax.numpy.array([postselect]), wires=wires)
+
+        qml.CNOT(wires=cnot_wires)
+        if reset:
+            if postselect is None:
+                qml.CNOT(wires=cnot_wires[::-1])
+            elif postselect == 1:
+                qml.PauliX(wires=wires)
+
+        self.state["cur_idx"] += 1
+        return MeasurementValue([meas], lambda x: x)
+
+    @DeferMeasurementsInterpreter.register_primitive(cond_prim)
+    def _(
+        self, *invals, jaxpr_branches, consts_slices, args_slice
+    ):  # pylint: disable=unused-argument
+        n_branches = len(jaxpr_branches)
+        conditions = invals[:n_branches]
+        if not any(isinstance(c, MeasurementValue) for c in conditions):
+            return PlxprInterpreter._primitive_registrations[cond_prim](
+                self,
+                *invals,
+                jaxpr_branches=jaxpr_branches,
+                consts_slices=consts_slices,
+                args_slice=args_slice,
+            )
+
+        conditions = get_mcm_predicates(conditions[:-1])
+        args = invals[args_slice]
+
+        for i, (condition, jaxpr) in enumerate(zip(conditions, jaxpr_branches, strict=True)):
+            if jaxpr is None:
+                # If a false branch isn't provided, the jaxpr corresponding to the condition
+                # for the false branch will be None. That is the only scenario where we would
+                # reach here.
+                continue
+
+            if isinstance(condition, MeasurementValue):
+                control_wires = Wires([m.wires[0] for m in condition.measurements])
+
+                for branch, value in condition.items():
+                    cur_consts = invals[consts_slices[i]]
+                    qml.cond(value, ctrl_transform_prim.bind)(
+                        *cur_consts,
+                        *args,
+                        *control_wires,
+                        jaxpr=jaxpr,
+                        n_control=len(control_wires),
+                        control_values=branch,
+                        work_wires=None,
+                        n_consts=len(cur_consts),
+                    )
+
+        return []
+
+    def defer_measurements_plxpr_to_plxpr(
+        jaxpr, consts, targs, tkwargs, *args
+    ):  # pylint: disable=unused-argument
+
+        if (aux_wires := tkwargs.get("aux_wires", None)) is None:
+            raise ValueError(
+                "'aux_wires' argument for qml.defer_measurements must be provided "
+                "when qml.capture.enabled() is True."
+            )
+        if tkwargs.pop("reduce_postselected", False):
+            warn(
+                "Cannot set 'reduce_postselected=True' with qml.capture.enabled() "
+                "when using qml.defer_measurements. Argument will be ignored.",
+                UserWarning,
+            )
+
+        interpreter = DeferMeasurementsInterpreter(aux_wires=aux_wires)
+
+        def wrapper(*inner_args):
+            return interpreter.eval(jaxpr, consts, *inner_args)
+
+        return jax.make_jaxpr(wrapper)(*args)
+
+    return DeferMeasurementsInterpreter, defer_measurements_plxpr_to_plxpr
+
+
+DeferMeasurementsInterpreter, defer_measurements_plxpr_to_plxpr = _get_plxpr_defer_measurements()
+
+
+# pylint: disable=unused-argument
+@partial(transform, plxpr_transform=defer_measurements_plxpr_to_plxpr)
 def defer_measurements(
-    tape: QuantumScript, reduce_postselected: bool = True, allow_postselect: bool = True
+    tape: QuantumScript,
+    reduce_postselected: bool = True,
+    allow_postselect: bool = True,
+    aux_wires: Optional[Union[int, Sequence[int], Wires]] = None,
 ) -> tuple[QuantumScriptBatch, PostprocessingFn]:
     """Quantum function transform that substitutes operations conditioned on
     measurement outcomes to controlled operations.
@@ -157,10 +493,13 @@ def defer_measurements(
     Args:
         tape (QNode or QuantumTape or Callable): a quantum circuit.
         reduce_postselected (bool): Whether to use postselection information to reduce the number
-            of operations and control wires in the output tape. Active by default.
+            of operations and control wires in the output tape. Active by default. This is currently
+            ignored if program capture is enabled.
         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``.
+            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
+            only used if program capture is enabled.
 
     Returns:
         qnode (QNode) or quantum function (Callable) or tuple[List[QuantumTape], function]: The
@@ -270,6 +609,85 @@ def defer_measurements(
         5: ──────────────────────╰X─╰●─╰○────────┤
 
         There is only one controlled gate with only one control wire.
+
+    .. details::
+        :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:
+
+        **New features**
+
+        * Arbitrary classical processing of mid-circuit measurement values is now possible. With
+          program capture disabled, only limited classical processing, as detailed in the
+          documentation for :func:`~pennylane.measure`. With program capture enabled, any unary
+          or binary ``jax.numpy`` functions that can be applied to scalars can be used with mid-circuit
+          measurements.
+
+        * Using mid-circuit measurements as gate parameters is now possible. This feature currently
+          has the following restrictions:
+          * Mid-circuit measurement values cannot be used for multiple parameters of the same gate.
+          * Mid-circuit measurement values cannot be used as wires.
+
+          .. code-block:: python
+
+              from functools import partial
+              import jax
+              import jax.numpy as jnp
+
+              qml.capture.enable()
+
+              @qml.capture.expand_plxpr_transforms
+              @partial(qml.defer_measurements, aux_wires=list(range(5, 10)))
+              def f():
+                  m0 = qml.measure(0)
+
+                  phi = jnp.sin(jnp.pi * m0)
+                  qml.RX(phi, 0)
+                  return qml.expval(qml.PauliZ(0))
+
+          >>> jax.make_jaxpr(f)()
+          { lambda ; . let
+              _:AbstractOperator() = CNOT[n_wires=2] 0 5
+              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
+              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
+              g:AbstractOperator() = PauliZ[n_wires=1] 0
+              h:AbstractMeasurement(n_wires=None) = expval_obs g
+            in (h,) }
+
+        The above dummy example showcases how the transform is applied when the aforementioned
+        features are used.
+
+        **What doesn't work**
+
+        * mid-circuit measurement values cannot be used in the condition for a
+          :func:`~pennylane.while_loop`.
+        * :func:`~pennylane.measure` cannot be used inside the body of loop primitives
+          (:func:`~pennylane.while_loop`, :func:`~pennylane.for_loop`).
+        * If a branch of :func:`~pennylane.cond` uses mid-circuit measurements as its
+          predicate, then all other branches must also use mid-circuit measurement values
+          as predicates.
+        * For an ``n``-parameter gate, mid-circuit measurement values can only be used
+          for 1 of the ``n`` parameters.
+        * :func:`~pennylane.measure` can only be used in the bodies of branches of
+          :func:`~pennylane.cond` if none of the branches use mid-circuit measurements
+          as predicates
+        * :func:`~pennylane.measure` cannot be used inside the body of functions
+          being transformed with :func:`~pennylane.adjoint` or :func:`~pennylane.ctrl`.
     """
     if not any(isinstance(o, MidMeasureMP) for o in tape.operations):
         return (tape,), null_postprocessing
@@ -385,10 +803,10 @@ def _add_control_gate(op, control_wires, reduce_postselected):
     """Helper function to add control gates"""
     if reduce_postselected:
         control = [control_wires[m.id] for m in op.meas_val.measurements if m.postselect is None]
-        items = op.meas_val._postselected_items()
+        items = op.meas_val.postselected_items()
     else:
         control = [control_wires[m.id] for m in op.meas_val.measurements]
-        items = op.meas_val._items()
+        items = op.meas_val.items()
 
     new_ops = []
 


Test patch:
diff --git a/tests/capture/transforms/test_capture_defer_measurements.py b/tests/capture/transforms/test_capture_defer_measurements.py
new file mode 100644
index 000000000..7f85852f2
--- /dev/null
+++ b/tests/capture/transforms/test_capture_defer_measurements.py
@@ -0,0 +1,760 @@
+# Copyright 2025 Xanadu Quantum Technologies Inc.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Unit tests for the ``DeferMeasurementsInterpreter`` class"""
+
+import numpy as np
+
+# pylint:disable=wrong-import-position, protected-access
+import pytest
+
+import pennylane as qml
+
+jax = pytest.importorskip("jax")
+jnp = pytest.importorskip("jax.numpy")
+
+from pennylane.capture.primitives import grad_prim, jacobian_prim, qnode_prim
+from pennylane.tape.plxpr_conversion import CollectOpsandMeas
+from pennylane.transforms.defer_measurements import (
+    DeferMeasurementsInterpreter,
+    defer_measurements_plxpr_to_plxpr,
+)
+from pennylane.wires import Wires
+
+pytestmark = [pytest.mark.jax, pytest.mark.usefixtures("enable_disable_plxpr")]
+
+
+class TestDeferMeasurementsInterpreter:
+    """Unit tests for DeferMeasurementsInterpreter."""
+
+    @pytest.mark.parametrize("aux_wires", [1, (), [1, 2, 3], Wires([1, 2, 3])])
+    def test_init(self, aux_wires):
+        """Test that the interpreter is initialized correctly."""
+        interpreter = DeferMeasurementsInterpreter(aux_wires)
+        assert interpreter._aux_wires == Wires(aux_wires)
+
+    @pytest.mark.parametrize("reset", [True, False])
+    def test_single_mcm(self, reset):
+        """Test that a function with a single MCM is transformed correctly."""
+
+        @DeferMeasurementsInterpreter(aux_wires=list(range(5, 10)))
+        def f():
+            qml.measure(0, reset=reset)
+
+        jaxpr = jax.make_jaxpr(f)()
+        expected_len = 2 if reset else 1  # CNOT for measure, CNOT for reset
+        assert len(jaxpr.eqns) == expected_len
+        assert jaxpr.eqns[0].primitive == qml.CNOT._primitive
+        invals = [invar.val for invar in jaxpr.eqns[0].invars]
+        assert invals == [0, 5]
+        if reset:
+            assert jaxpr.eqns[1].primitive == qml.CNOT._primitive
+            invals = [invar.val for invar in jaxpr.eqns[1].invars]
+            assert invals == [5, 0]
+
+    @pytest.mark.parametrize("reset", [True, False])
+    @pytest.mark.parametrize("postselect", [0, 1])
+    def test_single_mcm_postselect(self, reset, postselect):
+        """Test that a function with a single MCM with postselection is transformed correctly."""
+
+        @DeferMeasurementsInterpreter(aux_wires=list(range(5, 10)))
+        def f():
+            qml.measure(0, reset=reset, postselect=postselect)
+
+        jaxpr = jax.make_jaxpr(f)()
+        expected_len = 3 if reset and postselect == 1 else 2  # CNOT + Projector + optional(X)
+        assert len(jaxpr.consts) == 1  # Input to Projector for postselect
+
+        assert len(jaxpr.eqns) == expected_len
+        # Projector for postselect
+        assert jaxpr.eqns[0].primitive == qml.Projector._primitive
+        assert jaxpr.eqns[0].invars[0] == jaxpr.jaxpr.constvars[0]
+
+        # CNOT for measure
+        assert jaxpr.eqns[1].primitive == qml.CNOT._primitive
+        invals = [invar.val for invar in jaxpr.eqns[1].invars]
+        assert invals == [0, 5]
+
+        if reset:
+            if postselect == 1:
+                # PauliX since we know the state is |1>
+                assert jaxpr.eqns[2].primitive == qml.PauliX._primitive
+                assert jaxpr.eqns[2].invars[0].val == 0
+
+    def test_multiple_mcms(self):
+        """Test that applying multiple MCMs is transformed correctly."""
+
+        @DeferMeasurementsInterpreter(aux_wires=list(range(5, 10)))
+        def f(w):
+            qml.measure(w)
+            qml.measure(w + 1)
+            qml.measure(w + 2)
+
+        jaxpr = jax.make_jaxpr(f)(0)  # wires will be 0, 1, 2
+        assert jaxpr.eqns[0].primitive == qml.CNOT._primitive
+        assert jaxpr.eqns[0].invars[1].val == 5
+        assert jaxpr.eqns[2].primitive == qml.CNOT._primitive
+        assert jaxpr.eqns[2].invars[1].val == 6
+        assert jaxpr.eqns[4].primitive == qml.CNOT._primitive
+        assert jaxpr.eqns[4].invars[1].val == 7
+
+    def test_too_many_mcms(self):
+        """Test that an error is raised if more MCMs are present than the number of aux_wires."""
+
+        @DeferMeasurementsInterpreter(aux_wires=[5])
+        def f(w):
+            qml.measure(w)
+            qml.measure(w)
+
+        with pytest.raises(ValueError, match="Not enough auxiliary wires"):
+            f(0)
+
+    def test_mcms_as_gate_parameters(self):
+        """Test that MCMs can be used as gate parameters."""
+
+        @DeferMeasurementsInterpreter(aux_wires=list(range(5, 10)))
+        def f(x):
+            m = qml.measure(0)
+            qml.RX(x * m, 0)
+
+        x = 1.5
+        jaxpr = jax.make_jaxpr(f)(x)
+        collector = CollectOpsandMeas()
+        collector.eval(jaxpr.jaxpr, jaxpr.consts, x)
+
+        ops = collector.state["ops"]
+        expected_ops = [
+            qml.CNOT([0, 5]),
+            qml.ops.Controlled(qml.RX(jnp.array(0), 0), 5, control_values=[0]),
+            qml.ops.Controlled(qml.RX(jnp.array(x), 0), 5, control_values=[1]),
+        ]
+        assert ops == expected_ops
+
+    def test_multiple_mcms_as_gate_parameters_error(self):
+        """Test that multiple MCM parameters for a single operator raises an error."""
+
+        @DeferMeasurementsInterpreter(aux_wires=list(range(5, 10)))
+        def f(x):
+            m1 = qml.measure(0)
+            m2 = qml.measure(0)
+            qml.Rot(x, m1, m2, 0)
+
+        with pytest.raises(
+            qml.capture.CaptureError,
+            match="Cannot create operations with multiple parameters based on",
+        ):
+            _ = jax.make_jaxpr(f)(1.5)
+
+    def test_mcms_as_nested_gate_parameters(self):
+        """Test that MCMs can be used as gate parameters."""
+
+        @DeferMeasurementsInterpreter(aux_wires=list(range(5, 10)))
+        def f(x, y):
+            m = qml.measure(0)
+            qml.s_prod(y, qml.RX(x * m, 0))
+
+        args = (1.5, 2.5)
+        jaxpr = jax.make_jaxpr(f)(*args)
+        collector = CollectOpsandMeas()
+        collector.eval(jaxpr.jaxpr, jaxpr.consts, *args)
+
+        ops = collector.state["ops"]
+        expected_ops = [
+            qml.CNOT([0, 5]),
+            qml.ops.Controlled(qml.s_prod(args[1], qml.RX(jnp.array(0), 0)), 5, control_values=[0]),
+            qml.ops.Controlled(
+                qml.s_prod(args[1], qml.RX(jnp.array(args[0]), 0)), 5, control_values=[1]
+            ),
+        ]
+        assert ops == expected_ops
+
+    def test_simple_cond(self):
+        """Test that a qml.cond using a single MCM predicate is transformed correctly."""
+
+        @DeferMeasurementsInterpreter(aux_wires=list(range(5, 10)))
+        def f(x):
+            m = qml.measure(0)
+
+            @qml.cond(m)
+            def true_fn(phi):
+                qml.RX(phi, 0)
+
+            true_fn(x)
+
+        x = 1.5
+        jaxpr = jax.make_jaxpr(f)(x)
+        collector = CollectOpsandMeas()
+        collector.eval(jaxpr.jaxpr, jaxpr.consts, x)
+
+        ops = collector.state["ops"]
+        expected_ops = [qml.CNOT([0, 5]), qml.CRX(x, [5, 0])]
+        assert ops == expected_ops
+
+    def test_non_trivial_cond_predicate(self):
+        """Test that a qml.cond using processed MCMs as a predicate is transformed correctly."""
+
+        @DeferMeasurementsInterpreter(aux_wires=list(range(5, 10)))
+        def f(x):
+            m0 = qml.measure(0)
+            m1 = qml.measure(0)
+
+            @qml.cond(2 * m0 + m1)
+            def true_fn(phi):
+                qml.RX(phi, 0)
+
+            true_fn(x)
+
+        x = 1.5
+        jaxpr = jax.make_jaxpr(f)(x)
+        collector = CollectOpsandMeas()
+        collector.eval(jaxpr.jaxpr, jaxpr.consts, x)
+
+        ops = collector.state["ops"]
+        expected_ops = [
+            qml.CNOT([0, 5]),
+            qml.CNOT([0, 6]),
+            qml.ctrl(qml.RX(x, 0), [5, 6], [0, 1]),
+            qml.ctrl(qml.RX(x, 0), [5, 6], [1, 0]),
+            qml.ctrl(qml.RX(x, 0), [5, 6]),
+        ]
+        assert ops == expected_ops
+
+    def test_cond_elif_false_fn(self):
+        """Test that a qml.cond with elif and false branches is transformed correctly."""
+
+        @DeferMeasurementsInterpreter(aux_wires=list(range(5, 10)))
+        def f(x):
+            m0 = qml.measure(0)
+            m1 = qml.measure(0)
+
+            @qml.cond(m0)
+            def cond_fn(phi):
+                qml.RX(phi, 0)
+
+            @cond_fn.else_if(m1)
+            def _(phi):
+                qml.RY(phi, 0)
+
+            @cond_fn.otherwise
+            def _(phi):
+                qml.RZ(phi, 0)
+
+            cond_fn(x)
+
+        x = 1.5
+        jaxpr = jax.make_jaxpr(f)(x)
+        collector = CollectOpsandMeas()
+        collector.eval(jaxpr.jaxpr, jaxpr.consts, x)
+
+        ops = collector.state["ops"]
+        expected_ops = [
+            qml.CNOT([0, 5]),
+            qml.CNOT([0, 6]),
+            qml.CRX(x, [5, 0]),
+            qml.ctrl(qml.RY(x, 0), [5, 6], [0, 1]),
+            qml.ctrl(qml.RZ(x, 0), [5, 6], [0, 0]),
+        ]
+        assert ops == expected_ops
+
+    @pytest.mark.parametrize(
+        "mp_fn, mp_class",
+        [
+            (qml.expval, qml.measurements.ExpectationMP),
+            (qml.var, qml.measurements.VarianceMP),
+            (qml.sample, qml.measurements.SampleMP),
+            (qml.probs, qml.measurements.ProbabilityMP),
+            (qml.counts, qml.measurements.CountsMP),
+        ],
+    )
+    def test_mcm_statistics(self, mp_fn, mp_class):
+        """Test that collecting statistics on MCMs is handled correctly."""
+
+        @DeferMeasurementsInterpreter(aux_wires=list(range(5, 10)))
+        def f():
+            m0 = qml.measure(0)
+            m1 = qml.measure(0)
+            m2 = qml.measure(0)
+
+            outs = (mp_fn(op=m0), mp_fn(op=2.5 * m1 - m2))
+            if mp_fn not in (qml.expval, qml.var):
+                outs += (mp_fn(op=[m0, m1, m2]),)
+
+            if mp_fn == qml.counts:
+                outs += (mp_fn(m0, all_outcomes=True),)
+
+            return outs
+
+        jaxpr = jax.make_jaxpr(f)()
+        collector = CollectOpsandMeas()
+        collector.eval(jaxpr.jaxpr, jaxpr.consts)
+        ops = collector.state["ops"]
+        expected_ops = [qml.CNOT([0, 5]), qml.CNOT([0, 6]), qml.CNOT([0, 7])]
+        assert ops == expected_ops
+
+        measurements = collector.state["measurements"]
+        expected_measurements = [
+            mp_class(wires=Wires([5]), eigvals=jnp.arange(0, 2)),
+            mp_class(wires=Wires([6, 7]), eigvals=jnp.array([0.0, -1.0, 2.5, 1.5])),
+        ]
+        if mp_fn not in (qml.expval, qml.var):
+            expected_measurements.append(
+                mp_class(wires=Wires([5, 6, 7]), eigvals=qml.math.arange(0, 2**3))
+            )
+        if mp_fn == qml.counts:
+            expected_measurements.append(
+                mp_class(wires=Wires([5]), eigvals=jnp.arange(0, 2), all_outcomes=True)
+            )
+        assert measurements == expected_measurements
+
+    def test_arbitrary_mcm_processing(self):
+        """Test that arbitrary classical processing can be used with MCMs."""
+
+        def processing_fn(*ms):
+            a = jnp.sin(0.5 * jnp.pi * ms[0])
+            b = a - (ms[1] + 1) ** 4
+            c = jnp.sinh(b) ** (-ms[2] + 2)
+            return ms[0] * ms[1] * c
+
+        @DeferMeasurementsInterpreter(aux_wires=list(range(5, 10)))
+        def f():
+            m0 = qml.measure(0)
+            m1 = qml.measure(0)
+            m2 = qml.measure(0)
+
+            inval = processing_fn(m0, m1, m2)
+            return qml.expval(inval)
+
+        jaxpr = jax.make_jaxpr(f)()
+        collector = CollectOpsandMeas()
+        collector.eval(jaxpr.jaxpr, jaxpr.consts)
+        mp = collector.state["measurements"][0]
+        assert isinstance(mp, qml.measurements.ExpectationMP)
+        assert mp.wires == Wires([5, 6, 7])
+
+        expected_eigvals = []
+        n_mcms = 3
+        # Iterate through all 3-bit binary numbers
+        for i in range(2**n_mcms):
+            branch = tuple(int(b) for b in f"{i:0{n_mcms}b}")
+            expected_eigvals.append(processing_fn(*branch))
+
+        expected_eigvals = jnp.array(expected_eigvals)
+        assert jnp.allclose(mp.eigvals(), expected_eigvals)
+
+
+@pytest.mark.parametrize("postselect", [None, 0, 1])
+class TestDeferMeasurementsHigherOrderPrimitives:
+    """Unit tests for transforming higher-order primitives with DeferMeasurementsInterpreter."""
+
+    def test_for_loop(self, postselect):
+        """Test that a for_loop primitive is transformed correctly."""
+        n = jnp.array(3, dtype=int)
+
+        @DeferMeasurementsInterpreter(aux_wires=list(range(5, 10)))
+        def f(x):
+            qml.measure(0, postselect=postselect)
+
+            @qml.for_loop(n)
+            def loop_fn(i):
+                qml.RX(x, i)
+
+            loop_fn()
+            qml.measure(0)
+
+        x = 1.5
+        jaxpr = jax.make_jaxpr(f)(x)
+        collector = CollectOpsandMeas()
+        collector.eval(jaxpr.jaxpr, jaxpr.consts, x)
+
+        ops = collector.state["ops"]
+        expected_ops = [
+            qml.CNOT([0, 5]),
+            qml.RX(x, [0]),
+            qml.RX(x, [1]),
+            qml.RX(x, [2]),
+            qml.CNOT([0, 6]),
+        ]
+        if postselect is not None:
+            expected_ops.insert(0, qml.Projector(np.array([postselect]), 0))
+        assert ops == expected_ops
+
+    def test_while_loop(self, postselect):
+        """Test that a while_loop primitive is transformed correctly."""
+        n = jnp.array(3, dtype=int)
+
+        @DeferMeasurementsInterpreter(aux_wires=list(range(5, 10)))
+        def f(x):
+            qml.measure(0, postselect=postselect)
+
+            @qml.while_loop(lambda a: a < n)
+            def loop_fn(i):
+                qml.RX(x, i)
+                return i + 1
+
+            loop_fn(0)
+            qml.measure(0)
+
+        x = 1.5
+        jaxpr = jax.make_jaxpr(f)(x)
+        collector = CollectOpsandMeas()
+        collector.eval(jaxpr.jaxpr, jaxpr.consts, x)
+
+        ops = collector.state["ops"]
+        expected_ops = [
+            qml.CNOT([0, 5]),
+            qml.RX(x, [0]),
+            qml.RX(x, [1]),
+            qml.RX(x, [2]),
+            qml.CNOT([0, 6]),
+        ]
+        if postselect is not None:
+            expected_ops.insert(0, qml.Projector(np.array([postselect]), 0))
+        assert ops == expected_ops
+
+    def test_cond_non_mcm(self, postselect):
+        """Test that a qml.cond that does not use MCM predicates is transformed correctly."""
+
+        @DeferMeasurementsInterpreter(aux_wires=list(range(5, 10)))
+        def f(x):
+            qml.measure(0, postselect=postselect)
+
+            @qml.cond(x > 2.0)
+            def cond_fn(phi):
+                return qml.RX(phi, 0)
+
+            @cond_fn.otherwise
+            def _(phi):
+                return qml.RZ(phi, 0)
+
+            cond_fn(x)
+            qml.measure(0)
+
+        x = 1.5
+        jaxpr = jax.make_jaxpr(f)(x)
+        collector = CollectOpsandMeas()
+        collector.eval(jaxpr.jaxpr, jaxpr.consts, x)
+        ops = collector.state["ops"]
+        expected_ops = [qml.CNOT([0, 5]), qml.RZ(x, 0), qml.CNOT([0, 6])]
+        if postselect is not None:
+            expected_ops.insert(0, qml.Projector(np.array([postselect]), 0))
+        assert ops == expected_ops
+
+        x = 2.5
+        collector = CollectOpsandMeas()
+        collector.eval(jaxpr.jaxpr, jaxpr.consts, x)
+        ops = collector.state["ops"]
+        expected_ops = [qml.CNOT([0, 5]), qml.RX(x, 0), qml.CNOT([0, 6])]
+        if postselect is not None:
+            expected_ops.insert(0, qml.Projector(np.array([postselect]), 0))
+        assert ops == expected_ops
+
+    def test_cond_body_mcm(self, postselect):
+        """Test that a qml.cond containing MCMs in its body is transformed correctly."""
+
+        @DeferMeasurementsInterpreter(aux_wires=list(range(5, 15)))
+        def f(x):
+            qml.measure(0)
+
+            @qml.cond(x > 2.0)
+            def cond_fn(phi):
+                qml.measure(1, postselect=postselect)
+                qml.RX(phi, 0)
+                qml.measure(1)
+
+            @cond_fn.else_if(x > 1.0)
+            def _(phi):
+                qml.measure(2, postselect=postselect)
+                qml.RY(phi, 0)
+                qml.measure(2)
+
+            @cond_fn.otherwise
+            def _(phi):
+                qml.measure(3, postselect=postselect)
+                qml.RZ(phi, 0)
+                qml.measure(3)
+
+            cond_fn(x)
+            qml.measure(4)
+
+        x = 2.5
+        jaxpr = jax.make_jaxpr(f)(x)
+        collector = CollectOpsandMeas()
+        collector.eval(jaxpr.jaxpr, jaxpr.consts, x)
+        ops = collector.state["ops"]
+        expected_ops = [
+            qml.CNOT([0, 5]),
+            qml.CNOT([1, 6]),
+            qml.RX(x, 0),
+            qml.CNOT([1, 7]),
+            qml.CNOT([4, 12]),
+        ]
+        if postselect is not None:
+            expected_ops.insert(1, qml.Projector(np.array([postselect]), 1))
+        assert ops == expected_ops
+
+        x = 1.5
+        collector = CollectOpsandMeas()
+        collector.eval(jaxpr.jaxpr, jaxpr.consts, x)
+        ops = collector.state["ops"]
+        expected_ops = [
+            qml.CNOT([0, 5]),
+            qml.CNOT([2, 8]),
+            qml.RY(x, 0),
+            qml.CNOT([2, 9]),
+            qml.CNOT([4, 12]),
+        ]
+        if postselect is not None:
+            expected_ops.insert(1, qml.Projector(np.array([postselect]), 2))
+        assert ops == expected_ops
+
+        x = 0.5
+        collector = CollectOpsandMeas()
+        collector.eval(jaxpr.jaxpr, jaxpr.consts, x)
+        ops = collector.state["ops"]
+        expected_ops = [
+            qml.CNOT([0, 5]),
+            qml.CNOT([3, 10]),
+            qml.RZ(x, 0),
+            qml.CNOT([3, 11]),
+            qml.CNOT([4, 12]),
+        ]
+        if postselect is not None:
+            expected_ops.insert(1, qml.Projector(np.array([postselect]), 3))
+        assert ops == expected_ops
+
+    @pytest.mark.parametrize("lazy", [True, False])
+    def test_adjoint(self, lazy, postselect):
+        """Test that the adjoint_transform primitive is transformed correctly."""
+
+        @DeferMeasurementsInterpreter(aux_wires=list(range(5, 10)))
+        def f(x):
+            qml.measure(0, postselect=postselect)
+
+            def adjoint_fn(phi):
+                qml.RX(phi, 0)
+                qml.RY(phi, 0)
+
+            qml.adjoint(adjoint_fn, lazy=lazy)(x)
+
+        x = 1.5
+        jaxpr = jax.make_jaxpr(f)(x)
+        collector = CollectOpsandMeas()
+        collector.eval(jaxpr.jaxpr, jaxpr.consts, x)
+
+        ops = collector.state["ops"]
+        expected_ops = [
+            qml.CNOT([0, 5]),
+            qml.ops.Adjoint(qml.RY(x, 0)) if lazy else qml.RY(-x, 0),
+            qml.ops.Adjoint(qml.RX(x, 0)) if lazy else qml.RX(-x, 0),
+        ]
+        if postselect is not None:
+            expected_ops.insert(0, qml.Projector(np.array([postselect]), 0))
+        assert ops == expected_ops
+
+    def test_control(self, postselect):
+        """Test that the ctrl_transform primitive is transformed correctly."""
+        ctrl_wires = [1, 2]
+        ctrl_vals = [True, False]
+
+        @DeferMeasurementsInterpreter(aux_wires=list(range(5, 10)))
+        def f(x):
+            qml.measure(0, postselect=postselect)
+
+            def ctrl_fn(phi):
+                qml.RX(phi, 0)
+                qml.RY(phi, 0)
+
+            qml.ctrl(ctrl_fn, ctrl_wires, control_values=ctrl_vals)(x)
+
+        x = 1.5
+        jaxpr = jax.make_jaxpr(f)(x)
+        collector = CollectOpsandMeas()
+        collector.eval(jaxpr.jaxpr, jaxpr.consts, x)
+
+        ops = collector.state["ops"]
+        expected_ops = [
+            qml.CNOT([0, 5]),
+            qml.ops.Controlled(qml.RX(x, 0), ctrl_wires, control_values=ctrl_vals),
+            qml.ops.Controlled(qml.RY(x, 0), ctrl_wires, control_values=ctrl_vals),
+        ]
+        if postselect is not None:
+            expected_ops.insert(0, qml.Projector(np.array([postselect]), 0))
+        assert ops == expected_ops
+
+    def test_qnode(self, postselect):
+        """Test that a qnode primitive is transformed correctly."""
+        dev = qml.device("default.qubit", wires=10, shots=10)
+
+        @DeferMeasurementsInterpreter(aux_wires=list(range(5, 10)))
+        @qml.qnode(dev, diff_method="parameter-shift")
+        def f(x):
+            m0 = qml.measure(0, postselect=postselect)
+            m1 = qml.measure(
+                0, postselect=postselect if postselect is None else int(not postselect)
+            )
+
+            @qml.cond(2 * m0 + m1)
+            def true_fn(phi):
+                qml.RX(phi, 0)
+
+            true_fn(x)
+
+            return (
+                qml.expval(qml.Z(0)),
+                qml.probs(wires=[0, 1, 2]),
+                qml.sample(op=[m0, m1]),
+                qml.sample(op=m0 - 4.0 * m1),
+            )
+
+        x = 1.5
+        jaxpr = jax.make_jaxpr(f)(x)
+        assert jaxpr.eqns[0].primitive == qnode_prim
+        assert jaxpr.eqns[0].params["device"] == dev
+        assert jaxpr.eqns[0].params["qnode_kwargs"]["diff_method"] == "parameter-shift"
+
+        inner_jaxpr = jaxpr.eqns[0].params["qfunc_jaxpr"]
+        collector = CollectOpsandMeas()
+        collector.eval(inner_jaxpr, jaxpr.consts, x)
+
+        ops = collector.state["ops"]
+        expected_ops = [
+            qml.CNOT([0, 5]),
+            qml.CNOT([0, 6]),
+            qml.ctrl(qml.RX(x, 0), [5, 6], [0, 1]),
+            qml.ctrl(qml.RX(x, 0), [5, 6], [1, 0]),
+            qml.ctrl(qml.RX(x, 0), [5, 6]),
+        ]
+        if postselect is not None:
+            expected_ops.insert(0, qml.Projector(np.array([postselect]), 0))
+            expected_ops.insert(2, qml.Projector(np.array([int(not postselect)]), 0))
+        assert ops == expected_ops
+
+        measurements = collector.state["measurements"]
+        expected_measurements = [
+            qml.expval(qml.Z(0)),
+            qml.probs(wires=[0, 1, 2]),
+            qml.measurements.SampleMP(wires=[5, 6], eigvals=jnp.arange(0, 2**2)),
+            qml.measurements.SampleMP(wires=[5, 6], eigvals=qml.math.array([0.0, -4.0, 1.0, -3.0])),
+        ]
+        assert measurements == expected_measurements
+
+    @pytest.mark.parametrize(
+        "diff_fn, diff_prim", [(qml.grad, grad_prim), (qml.jacobian, jacobian_prim)]
+    )
+    def test_grad_jac(self, diff_fn, diff_prim, postselect):
+        """Test that differentiation primitives are transformed correctly."""
+        dev = qml.device("default.qubit", wires=4)
+
+        @qml.qnode(dev)
+        def circuit(x):
+            qml.RX(x, 0)
+            m0 = qml.measure(0, postselect=postselect)
+
+            @qml.cond(m0)
+            def true_fn(phi):
+                qml.RX(phi, 0)
+
+            true_fn(x)
+            return qml.expval(qml.PauliZ(0))
+
+        x = 1.5
+        transformed_fn = DeferMeasurementsInterpreter(aux_wires=list(range(5, 10)))(
+            diff_fn(circuit)
+        )
+        jaxpr = jax.make_jaxpr(transformed_fn)(x)
+        assert jaxpr.eqns[0].primitive == diff_prim
+        inner_jaxpr = jaxpr.eqns[0].params["jaxpr"]
+        assert inner_jaxpr.eqns[0].primitive == qnode_prim
+        qfunc_jaxpr = inner_jaxpr.eqns[0].params["qfunc_jaxpr"]
+
+        collector = CollectOpsandMeas()
+        collector.eval(qfunc_jaxpr, jaxpr.consts, x)
+
+        ops = collector.state["ops"]
+        expected_ops = [qml.RX(x, 0), qml.CNOT([0, 5]), qml.CRX(x, [5, 0])]
+        if postselect is not None:
+            expected_ops.insert(1, qml.Projector(np.array([postselect]), 0))
+        assert ops == expected_ops
+
+        measurements = collector.state["measurements"]
+        expected_measurements = [qml.expval(qml.Z(0))]
+        assert measurements == expected_measurements
+
+
+def test_defer_measurements_plxpr_to_plxpr():
+    """Test that transforming plxpr works."""
+
+    def f(x):
+        m = qml.measure(0)
+
+        @qml.cond(m)
+        def true_fn(phi):
+            qml.RX(phi, 0)
+
+        true_fn(x)
+
+    args = (1.5,)
+    targs = ()
+    tkwargs = {"aux_wires": list(range(5, 10))}
+    jaxpr = jax.make_jaxpr(f)(*args)
+    transformed_jaxpr = defer_measurements_plxpr_to_plxpr(
+        jaxpr.jaxpr, jaxpr.consts, targs, tkwargs, *args
+    )
+    collector = CollectOpsandMeas()
+    collector.eval(transformed_jaxpr.jaxpr, transformed_jaxpr.consts, *args)
+
+    ops = collector.state["ops"]
+    expected_ops = [qml.CNOT([0, 5]), qml.CRX(args[0], [5, 0])]
+    assert ops == expected_ops
+
+
+def test_defer_measurements_plxpr_to_plxpr_no_aux_wires_error():
+    """Test that an error is raised if ``aux_wires`` are not provided."""
+
+    def f(x):
+        m = qml.measure(0)
+
+        @qml.cond(m)
+        def true_fn(phi):
+            qml.RX(phi, 0)
+
+        true_fn(x)
+
+    args = (1.5,)
+    targs = ()
+    tkwargs = {}
+    jaxpr = jax.make_jaxpr(f)(*args)
+
+    with pytest.raises(ValueError, match="'aux_wires' argument for qml.defer_measurements must be"):
+        defer_measurements_plxpr_to_plxpr(jaxpr.jaxpr, jaxpr.consts, targs, tkwargs, *args)
+
+
+def test_defer_measurements_plxpr_to_plxpr_reduce_postselected_warning():
+    """Test that a warning is raised if ``reduce_postselected=True``."""
+
+    def f(x):
+        m = qml.measure(0)
+
+        @qml.cond(m)
+        def true_fn(phi):
+            qml.RX(phi, 0)
+
+        true_fn(x)
+
+    args = (1.5,)
+    targs = ()
+    tkwargs = {"aux_wires": list(range(5, 10)), "reduce_postselected": True}
+    jaxpr = jax.make_jaxpr(f)(*args)
+
+    with pytest.warns(UserWarning, match="Cannot set 'reduce_postselected=True'"):
+        defer_measurements_plxpr_to_plxpr(jaxpr.jaxpr, jaxpr.consts, targs, tkwargs, *args)
diff --git a/tests/measurements/test_mid_measure.py b/tests/measurements/test_mid_measure.py
index ed11f8b16..84fd3fb5c 100644
--- a/tests/measurements/test_mid_measure.py
+++ b/tests/measurements/test_mid_measure.py
@@ -557,7 +557,7 @@ class TestMeasurementValueItems:
         output at all."""
         mp = MidMeasureMP(0, postselect=postselect)
         mv = MeasurementValue([mp], func)
-        items = list(mv._items())
+        items = list(mv.items())
         assert items == [((0,), expected[0]), ((1,), expected[1])]
 
     funcs_and_expected_multi = [
@@ -582,7 +582,7 @@ class TestMeasurementValueItems:
             (1, 1, 1),
         ]
         mv = MeasurementValue([MP0, MP1, MP2], func)
-        items = list(mv._items())
+        items = list(mv.items())
         assert len(items) == len(branches3) == len(expected)
         for item, branch, exp in zip(items, branches3, expected):
             assert item == (branch, exp)
@@ -593,7 +593,7 @@ class TestMeasurementValueItems:
         """Test the full items."""
         mp = MidMeasureMP(0, postselect=postselect)
         mv = MeasurementValue([mp], func)
-        items = list(mv._postselected_items())
+        items = list(mv.postselected_items())
         if postselect is None:
             assert items == [((0,), expected[0]), ((1,), expected[1])]
         else:
@@ -628,7 +628,7 @@ class TestMeasurementValueItems:
         MP2 = MidMeasureMP(2, postselect=postselects[2])
 
         mv = MeasurementValue([MP0, MP1, MP2], func)
-        items = list(mv._postselected_items())
+        items = list(mv.postselected_items())
         assert len(items) == len(branches)
         for item, branch in zip(items, branches):
             pruned_branch = tuple(b for i, b in enumerate(branch) if postselects[i] is None)