Repository: PennyLaneAI/pennylane
Pull request #5088 (2024-01-22 17:57:39)

Mid-circuit measurements are natively supported (phase 1)
### Before submitting

Please complete the following checklist when submitting a PR:

- [x] All new features must include a unit test.
      If you've fixed a bug or added code that should be tested, add a test to the
      test directory!

- [x] All new functions and code must be clearly commented and documented.
      If you do make documentation changes, make sure that the docs build and
      render correctly by running `make docs`.

- [x] Ensure that the test suite passes, by running `make test`.

- [x] Add a new entry to the `doc/releases/changelog-dev.md` file, summarizing the
      change, and including a link back to the PR.

- [x] The PennyLane source code conforms to
      [PEP8 standards](https://www.python.org/dev/peps/pep-0008/).
      We check all of our code against [Pylint](https://www.pylint.org/).
      To lint modified files, simply `pip install pylint`, and then
      run `pylint pennylane/path/to/file.py`.

When all the above are checked, delete everything above the dashed
line and fill in the pull request template.

------------------------------------------------------------------------------------------------------------

**Context:**
[Epic 50811](https://app.shortcut.com/xanaduai/epic/50811)  states:

- [x] The deferred measurement principle is not used
- [x] Each shot is a new exploration of the tree of possible mid-circuit measurement results
- [x] The mid-circuit measurement results are correctly exposed to PennyLane's frontend
- [x] Compatible with qubit reuse
- [x] Compatible with qubit reset
- [ ] Compatible with qubit postselection
- [ ] Compatible with existing support for conditioning in PennyLane
  - [ ] Addition, subtraction and multiplication of one or more mid-circuit measurement results

The current PR tackles the requirements above as indicated by the task list.

**Description of the Change:**

- The `qml.defer_measurements` `preprocess`-transform is replaced by `mid_circuit_measurements` which includes `qml.defer_measurements` if and only if the device or tape has shots.
- `Conditional` and `MidMeasureMP` are now supported when running in shot-mode.
- `get_final_state` returns a third value, a dictionary of mid-circuit measurement values.
- `apply_operation` gets a new implementation which has a `MidMeasureMP` argument and returns the branched state and measurement value.
- `simulate` branches into another workflow when the tape has `MidMeasureMP` and shots. The new workflow implement native mid-circuit measurements.
  - It loops over shot_vector.
  - Creates an auxiliary tape adapted to a one-shot circuit evaluation (`init_auxiliary_circuit`).
  - Accumulates mid-circuit and terminal measurements repeatedly evaluating one-shot circuits (`accumulate_native_mcm`).
  - The mid-circuit measurements are recorded in a dictionary which keys are the hash of the measurement value objects. These keys are used to gather the results as mentioned below.
  - Combines the measurements (`gather_native_mid_circuit_measurements`).

**Benefits:**
Mid-circuit measurements are natively supported.

**Possible Drawbacks:**
Slower in several cases (e.g. when few MCMs).

**Missing features:**
To do in follow-up PRs:

- [ ] postselection support
- [ ] composite mid-circuit measurements
- [ ] probabilities of lists of measurement values
- [ ] parameter-shift support
- [ ] optimisation

**Related GitHub Issues:**
[sc-54987]
[sc-55665]

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/5088?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
Attention: `5 lines` in your changes are missing coverage. Please review.
> Comparison is base [(`c01cb4b`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/c01cb4b96f758a9f1a0c4e422115829b80e113cb?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) 99.67% compared to head [(`8451282`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5088?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) 99.66%.
> Report is 41 commits behind head on master.

| [Files](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5088?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | Patch % | Lines |
|---|---|---|
| [pennylane/devices/qubit/simulate.py](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5088?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2RldmljZXMvcXViaXQvc2ltdWxhdGUucHk=) | 97.47% | [3 Missing :warning: ](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5088?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) |
| [pennylane/devices/qubit/apply\_operation.py](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5088?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2RldmljZXMvcXViaXQvYXBwbHlfb3BlcmF0aW9uLnB5) | 90.00% | [2 Missing :warning: ](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5088?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    #5088      +/-   ##
==========================================
- Coverage   99.67%   99.66%   -0.01%     
==========================================
  Files         394      392       -2     
  Lines       35670    35728      +58     
==========================================
+ Hits        35554    35610      +56     
- Misses        116      118       +2     
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5088?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).

vincentmr: sc-55665-mid-circuit-measurements-polish-merge-mvp
vincentmr: > Sorry about all the nitpicks.

NP, that's part of reviews :)

> Two main questions:
> 
> 1. The measurement values dictionary. It may allow us to combine a lot of functionality by making it an optional, mutable keyword argument.  We are already relying on it being mutable, so let's lean into that and take advantage of it.

Agreed, I'll refactor in that sense.

> 2. The result shape for `all_shot_meas` versus `one_shot_meas`.  Can we standardize the contents of `all_shot_meas` in such a way that we can eliminate all the:
> 
> ```
> if not isistance(obj, (tuple, list)):
>     obj = [obj]
> ```
> 
> checks?

I've attempted something, putting all the logic in `accumulate_native_mcm`. Let me know if this fits the bill. 


vincentmr: > my biggest concern at this point is if we can clean up the aux-circuit stuff enough that it can be refactored into a pre-processing transform, and it's not a huge concern

Discussion [here](https://github.com/PennyLaneAI/pennylane/pull/5088#discussion_r1478469290)
 
> my other question is about the language. I know we say "native" a lot when describing things, but I wonder if it will start to sound silly within a year if it's more-or-less an assumed capability in PennyLane. I might be thinking too far ahead with this, but it's just something I was wondering given our long-term plan here

@trbromley what do you think? We could call this "dynamical evaluation", "device emulation", ...

codecov-commenter: ## [Codecov](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5088?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
Attention: `1 lines` in your changes are missing coverage. Please review.
> Comparison is base [(`85dd14f`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/85dd14f3b37bba4fbf0153e0f533ac18ba46163f?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) 99.69% compared to head [(`d6f2a0f`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5088?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) 99.68%.
> Report is 1 commits behind head on master.

| [Files](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5088?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | Patch % | Lines |
|---|---|---|
| [pennylane/devices/qubit/simulate.py](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5088?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2RldmljZXMvcXViaXQvc2ltdWxhdGUucHk=) | 99.13% | [1 Missing :warning: ](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5088?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    #5088      +/-   ##
==========================================
- Coverage   99.69%   99.68%   -0.01%     
==========================================
  Files         394      394              
  Lines       36118    35994     -124     
==========================================
- Hits        36007    35881     -126     
- Misses        111      113       +2     
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5088?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).

vincentmr: > 🎉 Thanks for all the patience with the feedback. Good to go from my side ;)

Thanks for all the feedback, this was quite formative. 
albi3ro: @vincentmr  Looks like this line is still uncovered.

![Screenshot 2024-02-07 at 10 09 35 AM](https://github.com/PennyLaneAI/pennylane/assets/6364575/d45f4fac-2d02-480c-b455-d8035b5400e1)

vincentmr: > @vincentmr Looks like this line is still uncovered.
> 
> ![Screenshot 2024-02-07 at 10 09 35 AM](https://private-user-images.githubusercontent.com/6364575/303029432-d45f4fac-2d02-480c-b455-d8035b5400e1.png?jwt=eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJpc3MiOiJnaXRodWIuY29tIiwiYXVkIjoicmF3LmdpdGh1YnVzZXJjb250ZW50LmNvbSIsImtleSI6ImtleTUiLCJleHAiOjE3MDczMTkzODAsIm5iZiI6MTcwNzMxOTA4MCwicGF0aCI6Ii82MzY0NTc1LzMwMzAyOTQzMi1kNDVmNGZhYy0yZDAyLTQ4MGMtYjQ1NS1kODAzNWI1NDAwZTEucG5nP1gtQW16LUFsZ29yaXRobT1BV1M0LUhNQUMtU0hBMjU2JlgtQW16LUNyZWRlbnRpYWw9QUtJQVZDT0RZTFNBNTNQUUs0WkElMkYyMDI0MDIwNyUyRnVzLWVhc3QtMSUyRnMzJTJGYXdzNF9yZXF1ZXN0JlgtQW16LURhdGU9MjAyNDAyMDdUMTUxODAwWiZYLUFtei1FeHBpcmVzPTMwMCZYLUFtei1TaWduYXR1cmU9MTU3YTYxY2Y4MGU0NGU3ZDVlYjdmZGRmNTc3Y2E4M2I3ZjhlMGZhNmRlMjM5ZDk2ZTIwYzExYzFlYmRmYmMxNiZYLUFtei1TaWduZWRIZWFkZXJzPWhvc3QmYWN0b3JfaWQ9MCZrZXlfaWQ9MCZyZXBvX2lkPTAifQ.MPdcQmm1cBnULmf9c_oHcflSv1loZTMX-k1Svvp1joI)

Right. Given how the initialization moved to `accumulate_native_mcm`, I think we should never hit those lines so let's remove them. 

Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 800815331..e4ac1f558 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -6,6 +6,16 @@
 
 <h4>Native mid-circuit measurements on default qubit 💡</h4>
 
+* The `default.qubit` device treats mid-circuit measurements natively when operating in
+  shots-mode (i.e. `shots is not None`). Previously, circuits with mid-circuit measurements
+  would be decomposed using the `@qml.defer_measurements` transform (which is still a valid
+  decorator), requiring one extra wire for each mid-circuit measurement. The new
+  behavior is to evaluate the circuit `shots` times, "collapsing" the circuit state
+  stochastically along the way. While this is oftentimes slower, it requires much less
+  memory for circuits with several mid-circuit measurements, or circuits which have already
+  several wires.
+  [(#5088)](https://github.com/PennyLaneAI/pennylane/pull/5088)
+
 <h4>Work easily and efficiently with Pauli operators 🔧</h4>
 
 * New `qml.commutator` function that allows to compute commutators between
diff --git a/pennylane/_qubit_device.py b/pennylane/_qubit_device.py
index 3fd50d03f..d75807240 100644
--- a/pennylane/_qubit_device.py
+++ b/pennylane/_qubit_device.py
@@ -328,7 +328,7 @@ class QubitDevice(Device):
             QuantumFunctionError: if the value of :attr:`~.Observable.return_type` is not supported
 
         Returns:
-            tuple: stastics for each shot item from the shot vector
+            tuple: statistics for each shot item from the shot vector
         """
         results = []
         s1 = 0
diff --git a/pennylane/devices/default_qubit.py b/pennylane/devices/default_qubit.py
index ade6ec283..268bdb500 100644
--- a/pennylane/devices/default_qubit.py
+++ b/pennylane/devices/default_qubit.py
@@ -25,6 +25,8 @@ import logging
 import numpy as np
 
 import pennylane as qml
+from pennylane.ops.op_math.condition import Conditional
+from pennylane.measurements.mid_measure import MidMeasureMP
 from pennylane.tape import QuantumTape, QuantumScript
 from pennylane.typing import Result, ResultBatch
 from pennylane.transforms import convert_to_numpy_parameters
@@ -33,6 +35,7 @@ from pennylane.transforms.core import TransformProgram
 from . import Device
 from .preprocess import (
     decompose,
+    mid_circuit_measurements,
     validate_observables,
     validate_measurements,
     validate_multiprocessing_workers,
@@ -92,6 +95,11 @@ def stopping_condition(op: qml.operation.Operator) -> bool:
     return op.has_matrix
 
 
+def stopping_condition_shots(op: qml.operation.Operator) -> bool:
+    """Specify whether or not an Operator object is supported by the device with shots."""
+    return isinstance(op, (Conditional, MidMeasureMP)) or stopping_condition(op)
+
+
 def accepted_sample_measurement(m: qml.measurements.MeasurementProcess) -> bool:
     """Specifies whether or not a measurement is accepted when sampling."""
     return isinstance(
@@ -451,9 +459,12 @@ class DefaultQubit(Device):
         transform_program = TransformProgram()
 
         transform_program.add_transform(validate_device_wires, self.wires, name=self.name)
-        transform_program.add_transform(qml.defer_measurements, device=self)
+        transform_program.add_transform(mid_circuit_measurements, device=self)
         transform_program.add_transform(
-            decompose, stopping_condition=stopping_condition, name=self.name
+            decompose,
+            stopping_condition=stopping_condition,
+            stopping_condition_shots=stopping_condition_shots,
+            name=self.name,
         )
         transform_program.add_transform(
             validate_measurements, sample_measurements=accepted_sample_measurement, name=self.name
diff --git a/pennylane/devices/preprocess.py b/pennylane/devices/preprocess.py
index 3f2d7e04e..5e0f4df12 100644
--- a/pennylane/devices/preprocess.py
+++ b/pennylane/devices/preprocess.py
@@ -147,6 +147,19 @@ def validate_device_wires(
     return (tape,), null_postprocessing
 
 
+@transform
+def mid_circuit_measurements(
+    tape: qml.tape.QuantumTape, device
+) -> (Sequence[qml.tape.QuantumTape], Callable):
+    """Provide the transform to handle mid-circuit measurements.
+
+    If the tape of device uses finite-shot, use the native implementation (i.e. no transform), and use defer measurements transforms otherwise.
+    """
+    if tape.shots and tape.batch_size is None:
+        return (tape,), null_postprocessing
+    return qml.defer_measurements(tape, device=device)
+
+
 @transform
 def validate_multiprocessing_workers(
     tape: qml.tape.QuantumTape, max_workers: int, device
@@ -240,6 +253,7 @@ def validate_adjoint_trainable_params(
 def decompose(
     tape: qml.tape.QuantumTape,
     stopping_condition: Callable[[qml.operation.Operator], bool],
+    stopping_condition_shots: Callable[[qml.operation.Operator], bool] = None,
     skip_initial_state_prep: bool = True,
     decomposer: Optional[
         Callable[[qml.operation.Operator], Sequence[qml.operation.Operator]]
@@ -254,6 +268,9 @@ def decompose(
         stopping_condition (Callable): a function from an operator to a boolean. If ``False``, the operator
             should be decomposed. If an operator cannot be decomposed and is not accepted by ``stopping_condition``,
             a ``DecompositionUndefinedError`` will be raised.
+        stopping_condition_shots (Callable): a function from an operator to a boolean. If ``False``, the operator
+            should be decomposed. If an operator cannot be decomposed and is not accepted by ``stopping_condition``,
+            a ``DecompositionUndefinedError`` will be raised. This replaces stopping_condition if and only if the tape as shots.
         skip_initial_state_prep=True (bool): If ``True``, the first operator will not be decomposed if it inherits from :class:`~.StatePrepBase`.
         decomposer (Callable): an optional callable that takes an operator and implements the relevant decomposition.
             If None, defaults to using a callable returning ``op.decomposition()`` for any :class:`~.Operator` .
@@ -312,6 +329,9 @@ def decompose(
         def decomposer(op):
             return op.decomposition()
 
+    if stopping_condition_shots is not None and tape.shots:
+        stopping_condition = stopping_condition_shots
+
     if not all(stopping_condition(op) for op in tape.operations):
         try:
             # don't decompose initial operations if its StatePrepBase
diff --git a/pennylane/devices/qubit/apply_operation.py b/pennylane/devices/qubit/apply_operation.py
index 8f7c9a7a1..a55bc57f3 100644
--- a/pennylane/devices/qubit/apply_operation.py
+++ b/pennylane/devices/qubit/apply_operation.py
@@ -21,6 +21,8 @@ import numpy as np
 import pennylane as qml
 
 from pennylane import math
+from pennylane.measurements import MidMeasureMP
+from pennylane.ops import Conditional
 
 SQRT2INV = 1 / math.sqrt(2)
 
@@ -149,7 +151,11 @@ def apply_operation_tensordot(op: qml.operation.Operator, state, is_state_batche
 
 @singledispatch
 def apply_operation(
-    op: qml.operation.Operator, state, is_state_batched: bool = False, debugger=None
+    op: qml.operation.Operator,
+    state,
+    is_state_batched: bool = False,
+    debugger=None,
+    **_,
 ):
     """Apply and operator to a given state.
 
@@ -210,26 +216,92 @@ def _apply_operation_default(op, state, is_state_batched, debugger):
 
 
 @apply_operation.register
-def apply_identity(op: qml.Identity, state, is_state_batched: bool = False, debugger=None):
+def apply_conditional(
+    op: Conditional, state, is_state_batched: bool = False, debugger=None, mid_measurements=None
+):
+    """Applies a conditional operation.
+
+    Args:
+        op (Operator): The operation to apply to ``state``
+        state (TensorLike): The starting state.
+        is_state_batched (bool): Boolean representing whether the state is batched or not
+        debugger (_Debugger): The debugger to use
+        mid_measurements (dict, None): Mid-circuit measurement dictionary mutated to record the sampled value
+
+    Returns:
+        ndarray: output state
+    """
+    if mid_measurements[op.meas_val.measurements[0]]:
+        return apply_operation(
+            op.then_op,
+            state,
+            is_state_batched=is_state_batched,
+            debugger=debugger,
+            mid_measurements=mid_measurements,
+        )
+    return state
+
+
+@apply_operation.register
+def apply_mid_measure(
+    op: MidMeasureMP, state, is_state_batched: bool = False, debugger=None, mid_measurements=None
+):
+    """Applies a native mid-circuit measurement.
+
+    Args:
+        op (Operator): The operation to apply to ``state``
+        state (TensorLike): The starting state.
+        is_state_batched (bool): Boolean representing whether the state is batched or not
+        debugger (_Debugger): The debugger to use
+        mid_measurements (dict, None): Mid-circuit measurement dictionary mutated to record the sampled value
+
+    Returns:
+        ndarray: output state
+    """
+    if is_state_batched:
+        raise ValueError("MidMeasureMP cannot be applied to batched states.")
+    wire = op.wires
+    probs = qml.devices.qubit.measure(qml.probs(wire), state)
+    sample = np.random.binomial(1, probs[1])
+    mid_measurements[op] = sample
+    axis = wire.toarray()[0]
+    slices = [slice(None)] * qml.math.ndim(state)
+    slices[axis] = int(not sample)
+    state[tuple(slices)] = 0.0
+    state_norm = np.linalg.norm(state)
+    if state_norm < 1.0e-15:
+        raise ValueError("Cannot normalize projected state.")
+    state = state / state_norm
+    if op.reset and sample == 1:
+        state = apply_operation(
+            qml.PauliX(wire), state, is_state_batched=is_state_batched, debugger=debugger
+        )
+    return state
+
+
+@apply_operation.register
+def apply_identity(op: qml.Identity, state, is_state_batched: bool = False, debugger=None, **_):
     """Applies a :class:`~.Identity` operation by just returning the input state."""
     return state
 
 
 @apply_operation.register
-def apply_global_phase(op: qml.GlobalPhase, state, is_state_batched: bool = False, debugger=None):
+def apply_global_phase(
+    op: qml.GlobalPhase, state, is_state_batched: bool = False, debugger=None, **_
+):
     """Applies a :class:`~.GlobalPhase` operation by multiplying the state by ``exp(1j * op.data[0])``"""
     return qml.math.exp(-1j * qml.math.cast(op.data[0], complex)) * state
 
 
 @apply_operation.register
-def apply_paulix(op: qml.PauliX, state, is_state_batched: bool = False, debugger=None):
+def apply_paulix(op: qml.PauliX, state, is_state_batched: bool = False, debugger=None, **_):
     """Apply :class:`pennylane.PauliX` operator to the quantum state"""
     axis = op.wires[0] + is_state_batched
     return math.roll(state, 1, axis)
 
 
 @apply_operation.register
-def apply_pauliz(op: qml.PauliZ, state, is_state_batched: bool = False, debugger=None):
+def apply_pauliz(op: qml.PauliZ, state, is_state_batched: bool = False, debugger=None, **_):
     """Apply pauliz to state."""
 
     axis = op.wires[0] + is_state_batched
@@ -247,7 +319,7 @@ def apply_pauliz(op: qml.PauliZ, state, is_state_batched: bool = False, debugger
 
 
 @apply_operation.register
-def apply_cnot(op: qml.CNOT, state, is_state_batched: bool = False, debugger=None):
+def apply_cnot(op: qml.CNOT, state, is_state_batched: bool = False, debugger=None, **_):
     """Apply cnot gate to state."""
     target_axes = (op.wires[1] - 1 if op.wires[1] > op.wires[0] else op.wires[1]) + is_state_batched
     control_axes = op.wires[0] + is_state_batched
@@ -265,7 +337,11 @@ def apply_cnot(op: qml.CNOT, state, is_state_batched: bool = False, debugger=Non
 
 @apply_operation.register
 def apply_multicontrolledx(
-    op: qml.MultiControlledX, state, is_state_batched: bool = False, debugger=None
+    op: qml.MultiControlledX,
+    state,
+    is_state_batched: bool = False,
+    debugger=None,
+    **_,
 ):
     r"""Apply MultiControlledX to a state with the default einsum/tensordot choice
     for 8 operation wires or less. Otherwise, apply a custom kernel based on
@@ -314,7 +390,13 @@ def apply_multicontrolledx(
 
 
 @apply_operation.register
-def apply_grover(op: qml.GroverOperator, state, is_state_batched: bool = False, debugger=None):
+def apply_grover(
+    op: qml.GroverOperator,
+    state,
+    is_state_batched: bool = False,
+    debugger=None,
+    **_,
+):
     """Apply GroverOperator either via a custom matrix-free method (more than 8 operation
     wires) or via standard matrix based methods (else)."""
     if len(op.wires) < 9:
@@ -356,7 +438,7 @@ def _apply_grover_without_matrix(state, op_wires, is_state_batched):
 
 
 @apply_operation.register
-def apply_snapshot(op: qml.Snapshot, state, is_state_batched: bool = False, debugger=None):
+def apply_snapshot(op: qml.Snapshot, state, is_state_batched: bool = False, debugger=None, **_):
     """Take a snapshot of the state"""
     if debugger is not None and debugger.active:
         measurement = op.hyperparameters["measurement"]
@@ -375,7 +457,11 @@ def apply_snapshot(op: qml.Snapshot, state, is_state_batched: bool = False, debu
 # pylint:disable = no-value-for-parameter, import-outside-toplevel
 @apply_operation.register
 def apply_parametrized_evolution(
-    op: qml.pulse.ParametrizedEvolution, state, is_state_batched: bool = False, debugger=None
+    op: qml.pulse.ParametrizedEvolution,
+    state,
+    is_state_batched: bool = False,
+    debugger=None,
+    **_,
 ):
     """Apply ParametrizedEvolution by evolving the state rather than the operator matrix
     if we are operating on more than half of the subsystem"""
diff --git a/pennylane/devices/qubit/simulate.py b/pennylane/devices/qubit/simulate.py
index 36b9207c6..5b2d00fad 100644
--- a/pennylane/devices/qubit/simulate.py
+++ b/pennylane/devices/qubit/simulate.py
@@ -13,12 +13,22 @@
 # limitations under the License.
 """Simulate a quantum script."""
 # pylint: disable=protected-access
+from collections import Counter
+from functools import singledispatch
 from typing import Optional
 
 from numpy.random import default_rng
 import numpy as np
 
 import pennylane as qml
+from pennylane.measurements import (
+    CountsMP,
+    ExpectationMP,
+    MidMeasureMP,
+    ProbabilityMP,
+    SampleMP,
+    VarianceMP,
+)
 from pennylane.typing import Result
 
 from .initialize_state import create_initial_state
@@ -98,7 +108,7 @@ def _postselection_postprocess(state, is_state_batched, shots):
     return state, shots
 
 
-def get_final_state(circuit, debugger=None, interface=None):
+def get_final_state(circuit, debugger=None, interface=None, mid_measurements=None):
     """
     Get the final state that results from executing the given quantum script.
 
@@ -108,6 +118,7 @@ def get_final_state(circuit, debugger=None, interface=None):
         circuit (.QuantumScript): The single circuit to simulate
         debugger (._Debugger): The debugger to use
         interface (str): The machine learning interface to create the initial state with
+        mid_measurements (None, dict): Dictionary of mid-circuit measurements
 
     Returns:
         Tuple[TensorLike, bool]: A tuple containing the final state of the quantum script and
@@ -125,8 +136,13 @@ def get_final_state(circuit, debugger=None, interface=None):
     # initial state is batched only if the state preparation (if it exists) is batched
     is_state_batched = bool(prep and prep.batch_size is not None)
     for op in circuit.operations[bool(prep) :]:
-        state = apply_operation(op, state, is_state_batched=is_state_batched, debugger=debugger)
-
+        state = apply_operation(
+            op,
+            state,
+            is_state_batched=is_state_batched,
+            debugger=debugger,
+            mid_measurements=mid_measurements,
+        )
         # Handle postselection on mid-circuit measurements
         if isinstance(op, qml.Projector):
             state, circuit._shots = _postselection_postprocess(
@@ -134,7 +150,7 @@ def get_final_state(circuit, debugger=None, interface=None):
             )
 
         # new state is batched if i) the old state is batched, or ii) the new op adds a batch dim
-        is_state_batched = is_state_batched or op.batch_size is not None
+        is_state_batched = is_state_batched or (op.batch_size is not None)
 
     for _ in range(len(circuit.wires) - len(circuit.op_wires)):
         # if any measured wires are not operated on, we pad the state with zeros.
@@ -237,7 +253,296 @@ def simulate(
     tensor([0.68117888, 0.        , 0.31882112, 0.        ], requires_grad=True))
 
     """
+    if circuit.shots and has_mid_circuit_measurements(circuit):
+        return simulate_native_mcm(
+            circuit, rng=rng, prng_key=prng_key, debugger=debugger, interface=interface
+        )
     state, is_state_batched = get_final_state(circuit, debugger=debugger, interface=interface)
     if state_cache is not None:
         state_cache[circuit.hash] = state
     return measure_final_state(circuit, state, is_state_batched, rng=rng, prng_key=prng_key)
+
+
+# pylint: disable=too-many-arguments
+def simulate_native_mcm(
+    circuit: qml.tape.QuantumScript,
+    rng=None,
+    prng_key=None,
+    debugger=None,
+    interface=None,
+) -> Result:
+    """Simulate a single quantum script with native mid-circuit measurements.
+
+    Args:
+        circuit (QuantumTape): The single circuit to simulate
+        rng (Union[None, int, array_like[int], SeedSequence, BitGenerator, Generator]): A
+            seed-like parameter matching that of ``seed`` for ``numpy.random.default_rng``.
+            If no value is provided, a default RNG will be used.
+        prng_key (Optional[jax.random.PRNGKey]): An optional ``jax.random.PRNGKey``. This is
+            the key to the JAX pseudo random number generator. If None, a random key will be
+            generated. Only for simulation using JAX.
+        debugger (_Debugger): The debugger to use
+        interface (str): The machine learning interface to create the initial state with
+        state_cache=None (Optional[dict]): A dictionary mapping the hash of a circuit to the pre-rotated state. Used to pass the state between forward passes and vjp calculations.
+
+    Returns:
+        tuple(TensorLike): The results of the simulation
+    """
+    if circuit.shots.has_partitioned_shots:
+        results = []
+        for s in circuit.shots:
+            aux_circuit = qml.tape.QuantumScript(
+                circuit.operations,
+                circuit.measurements,
+                shots=s,
+                trainable_params=circuit.trainable_params,
+            )
+            results.append(simulate_native_mcm(aux_circuit, rng, prng_key, debugger, interface))
+        return tuple(results)
+    aux_circuit = init_auxiliary_circuit(circuit)
+    all_shot_meas, list_mcm_values_dict = None, []
+    for _ in range(circuit.shots.total_shots):
+        one_shot_meas, mcm_values_dict = simulate_one_shot_native_mcm(
+            aux_circuit, rng, prng_key, debugger, interface
+        )
+        all_shot_meas = accumulate_native_mcm(aux_circuit, all_shot_meas, one_shot_meas)
+        list_mcm_values_dict.append(mcm_values_dict)
+    return parse_native_mid_circuit_measurements(circuit, all_shot_meas, list_mcm_values_dict)
+
+
+def init_auxiliary_circuit(circuit: qml.tape.QuantumScript):
+    """Creates an auxiliary circuit to perform one-shot mid-circuit measurement calculations.
+
+    Measurements with non-trivial measurement values are removed from the script. VarianceMP
+    measurements are also replaced by SampleMP measurements, which are necessary to evaluate
+    the variance.
+
+    Args:
+        circuit (QuantumTape): The original QuantumScript
+
+    Returns:
+        QuantumScript: A copy of the circuit with modified measurements
+    """
+    new_measurements = []
+    for m in circuit.measurements:
+        if not m.mv:
+            if isinstance(m, VarianceMP):
+                new_measurements.append(SampleMP(obs=m.obs))
+            else:
+                new_measurements.append(m)
+    return qml.tape.QuantumScript(
+        circuit.operations, new_measurements, shots=1, trainable_params=circuit.trainable_params
+    )
+
+
+def simulate_one_shot_native_mcm(
+    circuit: qml.tape.QuantumScript,
+    rng=None,
+    prng_key=None,
+    debugger=None,
+    interface=None,
+) -> Result:
+    """Simulate a single shot of a single quantum script with native mid-circuit measurements.
+
+    Args:
+        circuit (QuantumTape): The single circuit to simulate
+        rng (Union[None, int, array_like[int], SeedSequence, BitGenerator, Generator]): A
+            seed-like parameter matching that of ``seed`` for ``numpy.random.default_rng``.
+            If no value is provided, a default RNG will be used.
+        prng_key (Optional[jax.random.PRNGKey]): An optional ``jax.random.PRNGKey``. This is
+            the key to the JAX pseudo random number generator. If None, a random key will be
+            generated. Only for simulation using JAX.
+        debugger (_Debugger): The debugger to use
+        interface (str): The machine learning interface to create the initial state with
+
+    Returns:
+        tuple(TensorLike): The results of the simulation
+        dict: The mid-circuit measurement results of the simulation
+    """
+    mcm_dict = {}
+    state, is_state_batched = get_final_state(
+        circuit, debugger=debugger, interface=interface, mid_measurements=mcm_dict
+    )
+    return (
+        measure_final_state(circuit, state, is_state_batched, rng=rng, prng_key=prng_key),
+        mcm_dict,
+    )
+
+
+def accumulate_native_mcm(circuit: qml.tape.QuantumScript, all_shot_meas, one_shot_meas):
+    """Incorporates new measurements in current measurement sequence.
+
+    Args:
+        circuit (QuantumTape): A one-shot (auxiliary) QuantumScript
+        all_shot_meas (Sequence[Any]): List of accumulated measurement results
+        one_shot_meas (Sequence[Any]): List of measurement results
+
+    Returns:
+        tuple(TensorLike): The results of the simulation
+    """
+    if len(circuit.measurements) == 1:
+        one_shot_meas = [one_shot_meas]
+    if all_shot_meas is None:
+        new_shot_meas = one_shot_meas
+        for i, (m, s) in enumerate(zip(circuit.measurements, new_shot_meas)):
+            if isinstance(m, SampleMP) and isinstance(s, np.ndarray):
+                new_shot_meas[i] = [s]
+        return new_shot_meas
+    new_shot_meas = all_shot_meas
+    for i, m in enumerate(circuit.measurements):
+        if isinstance(m, CountsMP):
+            tmp = Counter(all_shot_meas[i])
+            tmp.update(Counter(one_shot_meas[i]))
+            new_shot_meas[i] = tmp
+        elif isinstance(m, SampleMP):
+            new_shot_meas[i].append(one_shot_meas[i])
+        else:
+            new_shot_meas[i] = all_shot_meas[i] + one_shot_meas[i]
+    return new_shot_meas
+
+
+def has_mid_circuit_measurements(
+    circuit: qml.tape.QuantumScript,
+):
+    """Returns True if the circuit contains a MidMeasureMP object and False otherwise.
+
+    Args:
+        circuit (QuantumTape): A QuantumScript
+
+    Returns:
+        bool: Whether the circuit contains a MidMeasureMP object
+    """
+    return any(isinstance(op, MidMeasureMP) for op in circuit.operations)
+
+
+def find_measurement_values(
+    circuit: qml.tape.QuantumScript,
+):
+    """Returns the indices of measurements with a non-trivial measurement value.
+
+    Args:
+        circuit (QuantumTape): A QuantumScript
+
+    Returns:
+        List[int]: Indices of measurements with a non-trivial measurement value.
+    """
+    return [i for i, m in enumerate(circuit.measurements) if m.mv]
+
+
+def find_not_measurement_values(
+    circuit: qml.tape.QuantumScript,
+):
+    """Returns the indices of measurements with a trivial measurement value.
+
+    Args:
+        circuit (QuantumTape): A QuantumScript
+
+    Returns:
+        List[int]: Indices of measurements with a trivial measurement value.
+    """
+    return [i for i, m in enumerate(circuit.measurements) if not m.mv]
+
+
+def parse_native_mid_circuit_measurements(
+    circuit: qml.tape.QuantumScript, all_shot_meas, mcm_shot_meas
+):
+    """Combines, gathers and normalizes the results of native mid-circuit measurement runs.
+
+    Args:
+        circuit (QuantumTape): A one-shot (auxiliary) QuantumScript
+        all_shot_meas (Sequence[Any]): List of accumulated measurement results
+        mcm_shot_meas (Sequence[dict]): List of dictionaries containing the mid-circuit measurement results of each shot
+
+    Returns:
+        tuple(TensorLike): The results of the simulation
+    """
+    idx_one_shot = find_not_measurement_values(circuit)
+    normalized_meas = [None] * len(circuit.measurements)
+    for i, m in zip(idx_one_shot, all_shot_meas):
+        normalized_meas[i] = gather_non_mcm(circuit.measurements[i], m, mcm_shot_meas)
+
+    idx_sample = find_measurement_values(circuit)
+    counter = Counter()
+    if any(isinstance(m, (CountsMP, ProbabilityMP)) for m in circuit.measurements):
+        for d in mcm_shot_meas:
+            counter.update(d)
+    for i in idx_sample:
+        normalized_meas[i] = gather_mcm(
+            circuit.measurements[i], circuit.measurements[i].mv, mcm_shot_meas, counter
+        )
+    return tuple(normalized_meas) if len(normalized_meas) > 1 else normalized_meas[0]
+
+
+def gather_non_mcm(circuit_measurement, measurement, samples):
+    """Combines, gathers and normalizes several measurements with trivial measurement values.
+
+    Args:
+        circuit_measurement (MeasurementProcess): measurement
+        measurement (TensorLike): measurement results
+        samples (List[dict]): Mid-circuit measurement samples
+
+    Returns:
+        TensorLike: The combined measurement outcome
+    """
+    if isinstance(circuit_measurement, (ExpectationMP, ProbabilityMP)):
+        new_meas = measurement / len(samples)
+    elif isinstance(circuit_measurement, SampleMP):
+        new_meas = np.concatenate(tuple(s.ravel() for s in measurement))
+    elif isinstance(circuit_measurement, CountsMP):
+        new_meas = dict(sorted(measurement.items()))
+    elif isinstance(circuit_measurement, VarianceMP):
+        new_meas = qml.math.var(np.concatenate(tuple(s.ravel() for s in measurement)))
+    else:
+        raise ValueError(
+            f"Native mid-circuit measurement mode does not support {circuit_measurement.__class__.__name__} measurements."
+        )
+    return new_meas
+
+
+@singledispatch
+def gather_mcm(measurement, mv, samples, counter):  # pylint: disable=unused-argument
+    """Combines, gathers and normalizes several measurements with non-trivial measurement values.
+
+    Args:
+        measurement (MeasurementProcess): measurement
+        mv (MeasurementValue): measurement value
+        samples (List[dict]): Mid-circuit measurement samples
+        counter (Counter): Measurement value counts
+
+    Returns:
+        TensorLike: The combined measurement outcome
+    """
+    raise ValueError(
+        f"Native mid-circuit measurement mode does not support {measurement.__class__.__name__} measurements."
+    )
+
+
+@gather_mcm.register
+def _(measurement: CountsMP, mv, samples, counter):  # pylint: disable=unused-argument
+    new_samples = {}
+    if len(samples) - counter[mv.measurements[0]] > 0:
+        new_samples[0] = len(samples) - counter[mv.measurements[0]]
+    if counter[mv.measurements[0]] > 0:
+        new_samples[1] = counter[mv.measurements[0]]
+    return new_samples
+
+
+@gather_mcm.register
+def _(measurement: ExpectationMP, mv, samples, counter):  # pylint: disable=unused-argument
+    return np.mean(np.array([dct[mv.measurements[0]] for dct in samples]))
+
+
+@gather_mcm.register
+def _(measurement: ProbabilityMP, mv, samples, counter):  # pylint: disable=unused-argument
+    p1 = counter[mv.measurements[0]] / float(len(samples))
+    return np.array([1 - p1, p1])
+
+
+@gather_mcm.register
+def _(measurement: SampleMP, mv, samples, counter):  # pylint: disable=unused-argument
+    return np.array([dct[mv.measurements[0]] for dct in samples])
+
+
+@gather_mcm.register
+def _(measurement: VarianceMP, mv, samples, counter):  # pylint: disable=unused-argument
+    return qml.math.var(np.array([dct[mv.measurements[0]] for dct in samples]))
diff --git a/pennylane/measurements/mid_measure.py b/pennylane/measurements/mid_measure.py
index 1f603f5a8..ea2ff4a4e 100644
--- a/pennylane/measurements/mid_measure.py
+++ b/pennylane/measurements/mid_measure.py
@@ -225,6 +225,7 @@ class MidMeasureMP(MeasurementProcess):
         postselect: Optional[int] = None,
         id: Optional[str] = None,
     ):
+        self.batch_size = None
         super().__init__(wires=Wires(wires), id=id)
         self.reset = reset
         self.postselect = postselect
diff --git a/pennylane/transforms/tape_expand.py b/pennylane/transforms/tape_expand.py
index a349d94d6..cfe9c6d08 100644
--- a/pennylane/transforms/tape_expand.py
+++ b/pennylane/transforms/tape_expand.py
@@ -408,14 +408,15 @@ def _create_decomp_preprocessing(custom_decomps, dev, decomp_depth=10):
                 container.kwargs["decomposer"] = decomposer
                 container.kwargs["max_expansion"] = decomp_depth
 
-                original_stopping_condition = container.kwargs["stopping_condition"]
+                for cond in ["stopping_condition", "stopping_condition_shots"]:
+                    original_stopping_condition = container.kwargs[cond]
 
-                def stopping_condition(obj):
-                    if obj.name in custom_decomps or type(obj) in custom_decomps:
-                        return False
-                    return original_stopping_condition(obj)
+                    def stopping_condition(obj):
+                        if obj.name in custom_decomps or type(obj) in custom_decomps:
+                            return False
+                        return original_stopping_condition(obj)
 
-                container.kwargs["stopping_condition"] = stopping_condition
+                    container.kwargs[cond] = stopping_condition
 
                 break
 


Test patch:
diff --git a/tests/devices/default_qubit/test_default_qubit.py b/tests/devices/default_qubit/test_default_qubit.py
index 4c63a9790..06925409e 100644
--- a/tests/devices/default_qubit/test_default_qubit.py
+++ b/tests/devices/default_qubit/test_default_qubit.py
@@ -1680,6 +1680,7 @@ class TestPostselection:
         dev = qml.device("default.qubit")
         param = qml.math.asarray(param, like=interface)
 
+        @qml.defer_measurements
         @qml.qnode(dev, interface=interface)
         def circ_postselect(theta):
             qml.RX(theta, 0)
@@ -1687,6 +1688,7 @@ class TestPostselection:
             qml.measure(0, postselect=1)
             return qml.apply(mp)
 
+        @qml.defer_measurements
         @qml.qnode(dev, interface=interface)
         def circ_expected():
             qml.RX(np.pi, 0)
@@ -1730,6 +1732,7 @@ class TestPostselection:
         dev = qml.device("default.qubit")
         param = qml.math.asarray(param, like=interface)
 
+        @qml.defer_measurements
         @qml.qnode(dev, interface=interface)
         def circ_postselect(theta):
             qml.RX(theta, 0)
@@ -1737,6 +1740,7 @@ class TestPostselection:
             qml.measure(0, postselect=1)
             return qml.apply(mp)
 
+        @qml.defer_measurements
         @qml.qnode(dev, interface=interface)
         def circ_expected():
             qml.RX(np.pi, 0)
@@ -1781,6 +1785,7 @@ class TestPostselection:
 
         with mock.patch("numpy.random.binomial", lambda *args, **kwargs: 5):
 
+            @qml.defer_measurements
             @qml.qnode(dev, interface=interface)
             def circ_postselect(theta):
                 qml.RX(theta, 0)
@@ -1834,8 +1839,9 @@ class TestPostselection:
             pytest.skip("Jitting tested in different test.")
 
         # Wires are specified so that the shape for measurements can be determined correctly
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit")
 
+        @qml.defer_measurements
         @qml.qnode(dev, interface=interface)
         def circ():
             qml.RX(np.pi, 0)
@@ -1876,9 +1882,10 @@ class TestPostselection:
         import jax
 
         # Wires are specified so that the shape for measurements can be determined correctly
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit")
 
         @jax.jit
+        @qml.defer_measurements
         @qml.qnode(dev, interface=interface)
         def circ():
             qml.RX(np.pi, 0)
@@ -1915,6 +1922,7 @@ class TestPostselection:
 
         dev = qml.device("default.qubit")
 
+        @qml.defer_measurements
         @qml.qnode(dev, interface=interface)
         def circ():
             qml.RX(np.pi, 0)
diff --git a/tests/devices/default_qubit/test_default_qubit_native_mcm.py b/tests/devices/default_qubit/test_default_qubit_native_mcm.py
new file mode 100644
index 000000000..3da8f71bc
--- /dev/null
+++ b/tests/devices/default_qubit/test_default_qubit_native_mcm.py
@@ -0,0 +1,195 @@
+# Copyright 2023 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.
+"""Tests for default qubit preprocessing."""
+
+from flaky import flaky
+import numpy as np
+import pytest
+
+import pennylane as qml
+from pennylane.devices.qubit.apply_operation import apply_mid_measure, MidMeasureMP
+from pennylane.devices.qubit.simulate import gather_mcm
+
+
+def validate_counts(shots, results1, results2):
+    if isinstance(results1, (list, tuple)):
+        assert len(results1) == len(results2)
+        for r1, r2 in zip(results1, results2):
+            validate_counts(shots, r1, r2)
+        return
+    ncounts = 0.5 * (sum(results1.values()) + sum(results2.values()))
+    for r1, r2 in zip(sorted(results1.items()), sorted(results2.items())):
+        assert abs(sum(r1) - sum(r2)) < (ncounts // 10)
+
+
+def validate_samples(shots, results1, results2):
+    for res in [results1, results2]:
+        if isinstance(shots, (list, tuple)):
+            assert len(res) == len(shots)
+            assert all(r.shape == (s,) for r, s in zip(res, shots))
+            assert all(
+                abs(sum(r1) - sum(r2)) < (s // 10) for r1, r2, s in zip(results1, results2, shots)
+            )
+        else:
+            assert res.shape == (shots,)
+            assert abs(sum(results1) - sum(results2)) < (shots // 10)
+
+
+def validate_expval(shots, results1, results2):
+    if shots is None:
+        assert np.allclose(results1, results2)
+        return
+    assert np.allclose(results1, results2, atol=0, rtol=0.3)
+
+
+def validate_measurements(func, shots, results1, results2):
+    if func is qml.counts:
+        validate_counts(shots, results1, results2)
+        return
+
+    if func is qml.sample:
+        validate_samples(shots, results1, results2)
+        return
+
+    validate_expval(shots, results1, results2)
+
+
+def test_apply_mid_measure():
+    with pytest.raises(ValueError, match="MidMeasureMP cannot be applied to batched states."):
+        _ = apply_mid_measure(
+            MidMeasureMP(0), np.zeros((2, 2)), is_state_batched=True, mid_measurements={}
+        )
+    with pytest.raises(ValueError, match="Cannot normalize projected state."):
+        _ = apply_mid_measure(MidMeasureMP(0), np.zeros(2), mid_measurements={})
+
+
+@pytest.mark.parametrize(
+    "measurement",
+    [
+        qml.state(),
+        qml.density_matrix(0),
+        qml.vn_entropy(0),
+        qml.mutual_info(0, 1),
+        qml.purity(0),
+        qml.classical_shadow(0),
+        qml.shadow_expval(0),
+    ],
+)
+def test_gather_mcm(measurement):
+    with pytest.raises(ValueError, match="Native mid-circuit measurement mode does not support"):
+        gather_mcm(measurement, None, None, None)
+
+
+def test_unsupported_measurement():
+    dev = qml.device("default.qubit", shots=1000)
+    params = np.pi / 4 * np.ones(2)
+
+    @qml.qnode(dev)
+    def func(x, y):
+        qml.RX(x, wires=0)
+        m0 = qml.measure(0)
+        qml.cond(m0, qml.RY)(y, wires=1)
+        return qml.classical_shadow(wires=0)
+
+    with pytest.raises(
+        ValueError,
+        match="Native mid-circuit measurement mode does not support ClassicalShadowMP measurements.",
+    ):
+        func(*params)
+
+
+@flaky(max_runs=5)
+@pytest.mark.parametrize("shots", [None, 1000, [1000, 1001]])
+@pytest.mark.parametrize("reset", [False, True])
+@pytest.mark.parametrize("measure_f", [qml.expval, qml.probs, qml.sample, qml.counts, qml.var])
+def test_single_mcm_single_measure_mcm(shots, reset, measure_f):
+    """Tests that DefaultQubit handles a circuit with a single mid-circuit measurement and a
+    conditional gate. A single measurement of the mid-circuit measurement value is performed at
+    the end."""
+    dev = qml.device("default.qubit", shots=shots)
+    params = np.pi / 4 * np.ones(2)
+
+    @qml.qnode(dev)
+    def func1(x, y):
+        qml.RX(x, wires=0)
+        m0 = qml.measure(0, reset=reset)
+        qml.cond(m0, qml.RY)(y, wires=1)
+        return measure_f(op=m0)
+
+    func2 = qml.defer_measurements(func1)
+
+    if shots is None and measure_f in (qml.counts, qml.sample):
+        return
+
+    results1 = func1(*params)
+    results2 = func2(*params)
+
+    validate_measurements(measure_f, shots, results1, results2)
+
+
+@flaky(max_runs=5)
+@pytest.mark.parametrize("shots", [None, 1000, [1000, 1001]])
+@pytest.mark.parametrize("reset", [False, True])
+@pytest.mark.parametrize("measure_f", [qml.expval, qml.probs, qml.sample, qml.counts, qml.var])
+def test_single_mcm_single_measure_obs(shots, reset, measure_f):
+    """Tests that DefaultQubit handles a circuit with a single mid-circuit measurement and a
+    conditional gate. A single measurement of a common observable is is performed at the end."""
+    dev = qml.device("default.qubit", shots=shots)
+    params = np.pi / 4 * np.ones(2)
+    obs = qml.PauliZ(0)
+
+    @qml.qnode(dev)
+    def func1(x, y):
+        qml.RX(x, wires=0)
+        m0 = qml.measure(0, reset=reset)
+        qml.cond(m0, qml.RY)(y, wires=1)
+        return measure_f(op=obs)
+
+    func2 = qml.defer_measurements(func1)
+
+    if shots is None and measure_f in (qml.counts, qml.sample):
+        return
+
+    results1 = func1(*params)
+    results2 = func2(*params)
+
+    validate_measurements(measure_f, shots, results1, results2)
+
+
+@flaky(max_runs=5)
+@pytest.mark.parametrize("shots", [1000])
+@pytest.mark.parametrize("reset", [False, True])
+@pytest.mark.parametrize("measure_f", [qml.expval, qml.probs, qml.sample, qml.counts, qml.var])
+def test_single_mcm_multiple_measurements(shots, reset, measure_f):
+    """Tests that DefaultQubit handles a circuit with a single mid-circuit measurement with reset
+    and a conditional gate. Multiple measurements of the mid-circuit measurement value are
+    performed."""
+    dev = qml.device("default.qubit", shots=shots)
+    params = np.pi / 4 * np.ones(2)
+    obs = qml.PauliZ(0)
+
+    @qml.qnode(dev)
+    def func1(x, y):
+        qml.RX(x, wires=0)
+        m0 = qml.measure(0, reset=reset)
+        qml.cond(m0, qml.RY)(y, wires=1)
+        return measure_f(op=obs), measure_f(op=m0)
+
+    func2 = qml.defer_measurements(func1)
+
+    results1 = func1(*params)
+    results2 = func2(*params)
+
+    for r1, r2 in zip(results1, results2):
+        validate_measurements(measure_f, shots, r1, r2)
diff --git a/tests/devices/default_qubit/test_default_qubit_preprocessing.py b/tests/devices/default_qubit/test_default_qubit_preprocessing.py
index e6dcfc05e..0e67d55d5 100644
--- a/tests/devices/default_qubit/test_default_qubit_preprocessing.py
+++ b/tests/devices/default_qubit/test_default_qubit_preprocessing.py
@@ -618,13 +618,6 @@ class TestPreprocessingIntegration:
         expected_expval = np.cos(y)
         assert qml.math.allclose(expected_expval, processed_results[1])
 
-    def test_preprocess_defer_measurements(self):
-        """Test preprocessing contains the defer measurement transform."""
-        dev = DefaultQubit()
-
-        program, _ = dev.preprocess()
-        assert qml.defer_measurements.transform in [t.transform for t in program]
-
 
 class TestAdjointDiffTapeValidation:
     """Unit tests for validate_and_expand_adjoint"""
diff --git a/tests/measurements/test_expval.py b/tests/measurements/test_expval.py
index 68577c718..f25c51a39 100644
--- a/tests/measurements/test_expval.py
+++ b/tests/measurements/test_expval.py
@@ -127,6 +127,7 @@ class TestExpval:
         are correct for a composite measurement value."""
         dev = qml.device("default.qubit")
 
+        @qml.defer_measurements
         @qml.qnode(dev)
         def circuit(phi):
             qml.RX(phi, 0)
diff --git a/tests/measurements/test_probs.py b/tests/measurements/test_probs.py
index 1e75208f0..48c1aa1ab 100644
--- a/tests/measurements/test_probs.py
+++ b/tests/measurements/test_probs.py
@@ -282,6 +282,7 @@ class TestProbs:
         are correct for a measurement value list."""
         dev = qml.device("default.qubit")
 
+        @qml.defer_measurements
         @qml.qnode(dev)
         def circuit(phi):
             qml.RX(phi, 0)
diff --git a/tests/measurements/test_sample.py b/tests/measurements/test_sample.py
index e510a556a..c6df2fc2f 100644
--- a/tests/measurements/test_sample.py
+++ b/tests/measurements/test_sample.py
@@ -178,8 +178,9 @@ class TestSample:
     def test_observable_is_composite_measurement_value(self, shots, phi):
         """Test that samples for mid-circuit measurement values
         are correct for a composite measurement value."""
-        dev = qml.device("default.qubit", wires=2, shots=shots)
+        dev = qml.device("default.qubit", shots=shots)
 
+        @qml.defer_measurements
         @qml.qnode(dev)
         def circuit(phi):
             qml.RX(phi, 0)
@@ -201,8 +202,9 @@ class TestSample:
     def test_observable_is_measurement_value_list(self, shots, phi):
         """Test that samples for mid-circuit measurement values
         are correct for a measurement value list."""
-        dev = qml.device("default.qubit", wires=2, shots=shots)
+        dev = qml.device("default.qubit", shots=shots)
 
+        @qml.defer_measurements
         @qml.qnode(dev)
         def circuit(phi):
             qml.RX(phi, 0)
diff --git a/tests/measurements/test_var.py b/tests/measurements/test_var.py
index f94f4299e..19b802c4a 100644
--- a/tests/measurements/test_var.py
+++ b/tests/measurements/test_var.py
@@ -101,6 +101,7 @@ class TestVar:
         are correct for a composite measurement value."""
         dev = qml.device("default.qubit")
 
+        @qml.defer_measurements
         @qml.qnode(dev)
         def circuit(phi):
             qml.RX(phi, 0)
diff --git a/tests/ops/functions/test_iterative_qpe.py b/tests/ops/functions/test_iterative_qpe.py
index a11fad161..0f1faf6ef 100644
--- a/tests/ops/functions/test_iterative_qpe.py
+++ b/tests/ops/functions/test_iterative_qpe.py
@@ -30,6 +30,7 @@ class TestIQPE:
         # TODO: When we have general statistics on measurements we can calculate it exactly with qml.probs
         dev = qml.device("default.qubit", shots=10000000)
 
+        @qml.defer_measurements
         @qml.qnode(dev)
         def circuit_iterative():
             # Initial state
diff --git a/tests/test_qnode.py b/tests/test_qnode.py
index 7976f98ec..ced345496 100644
--- a/tests/test_qnode.py
+++ b/tests/test_qnode.py
@@ -949,7 +949,29 @@ class TestIntegration:
         r1 = cry_qnode(first_par)
         r2 = conditional_ry_qnode(first_par)
         assert np.allclose(r1, r2)
-        assert spy.call_count == 2  # once per device preprocessing
+        assert spy.call_count == 0
+
+        @qml.defer_measurements
+        @qml.qnode(dev)
+        def cry_qnode_deferred(x):
+            """QNode where we apply a controlled Y-rotation."""
+            qml.BasisStatePreparation(basis_state, wires=[0, 1])
+            qml.CRY(x, wires=[0, 1])
+            return qml.sample(qml.PauliZ(1))
+
+        @qml.defer_measurements
+        @qml.qnode(dev)
+        def conditional_ry_qnode_deferred(x):
+            """QNode where the defer measurements transform is applied by
+            default under the hood."""
+            qml.BasisStatePreparation(basis_state, wires=[0, 1])
+            m_0 = qml.measure(0)
+            qml.cond(m_0, qml.RY)(x, wires=1)
+            return qml.sample(qml.PauliZ(1))
+
+        r1 = cry_qnode_deferred(first_par)
+        r2 = conditional_ry_qnode_deferred(first_par)
+        assert np.allclose(r1, r2)
 
     @pytest.mark.tf
     @pytest.mark.parametrize("interface", ["tf", "auto"])
diff --git a/tests/transforms/test_defer_measurements.py b/tests/transforms/test_defer_measurements.py
index ffb32192a..513d7795e 100644
--- a/tests/transforms/test_defer_measurements.py
+++ b/tests/transforms/test_defer_measurements.py
@@ -140,6 +140,7 @@ class TestQNode:
 
         dev = qml.device("default.qubit")
 
+        @qml.defer_measurements
         @qml.qnode(dev)
         def circ():
             qml.PauliX(0)
@@ -154,7 +155,6 @@ class TestQNode:
         measurement yields the correct results and is transformed correctly."""
         dev = qml.device("default.qubit", wires=3)
 
-        @qml.defer_measurements
         @qml.qnode(dev)
         def qnode1():
             return qml.expval(qml.PauliZ(0))
@@ -182,8 +182,8 @@ class TestQNode:
         """Test that wires can be reused after measurement."""
         dev = qml.device("default.qubit", wires=2)
 
-        @qml.qnode(dev)
         @qml.defer_measurements
+        @qml.qnode(dev)
         def qnode():
             qml.Hadamard(0)
             qml.measure(0)
@@ -237,6 +237,7 @@ class TestQNode:
         dev = qml.device("default.qubit", wires=4)
         spy = mocker.spy(qml.defer_measurements, "_transform")
 
+        @qml.defer_measurements
         @qml.qnode(dev)
         def qnode1(phi, theta):
             qml.RX(phi, 0)
@@ -249,6 +250,7 @@ class TestQNode:
 
         res1 = qnode1(np.pi / 4, 3 * np.pi / 4)
 
+        @qml.defer_measurements
         @qml.qnode(dev)
         def qnode2(phi, theta):
             qml.RX(phi, 0)
@@ -261,7 +263,7 @@ class TestQNode:
 
         res2 = qnode2(np.pi / 4, 3 * np.pi / 4)
 
-        assert spy.call_count == 2
+        assert spy.call_count == 4
 
         deferred_tapes1, _ = qml.defer_measurements(qnode1.qtape)
         deferred_tape1 = deferred_tapes1[0]
@@ -377,6 +379,7 @@ class TestQNode:
         a single mid-circuit measurement."""
         dev = DefaultQubit(seed=10)
 
+        @qml.defer_measurements
         @qml.qnode(dev)
         def circ1(x):
             qml.RX(x, 0)
@@ -434,6 +437,7 @@ class TestQNode:
         # Using DefaultQubit to allow non-commuting measurements
         dev = DefaultQubit(seed=10)
 
+        @qml.defer_measurements
         @qml.qnode(dev)
         def circ1(x, y):
             qml.RX(x, 0)
@@ -529,8 +533,8 @@ class TestQNode:
         """Test that CV operations are not supported."""
         dev = qml.device("default.gaussian", wires=3)
 
-        @qml.qnode(dev)
         @qml.defer_measurements
+        @qml.qnode(dev)
         def qnode():
             qml.measure(0)
             qml.Rotation(0.123, wires=[0])
@@ -545,8 +549,8 @@ class TestQNode:
         """Test that CV observables are not supported."""
         dev = qml.device("default.gaussian", wires=3)
 
-        @qml.qnode(dev)
         @qml.defer_measurements
+        @qml.qnode(dev)
         def qnode():
             qml.measure(0)
             return qml.expval(qml.NumberOperator(1))
@@ -752,8 +756,8 @@ class TestConditionalOperations:
             controlled_op(rads, wires=[0, 1])
             return qml.probs(wires=1)
 
-        @qml.qnode(dev)
         @qml.defer_measurements
+        @qml.qnode(dev)
         def quantum_control_circuit(rads):
             qml.Hadamard(0)
             m_0 = qml.measure(0)
@@ -850,8 +854,8 @@ class TestConditionalOperations:
             controlled_op(rads, wires=[0, 1])
             return qml.probs(wires=1)
 
-        @qml.qnode(dev)
         @qml.defer_measurements
+        @qml.qnode(dev)
         def quantum_control_circuit(rads):
             qml.Hadamard(0)
             qml.PauliX(0)
@@ -883,8 +887,8 @@ class TestConditionalOperations:
             qml.PauliX(0)
             return qml.probs(wires=1)
 
-        @qml.qnode(dev)
         @qml.defer_measurements
+        @qml.qnode(dev)
         def quantum_control_circuit(rads):
             qml.Hadamard(0)
             m_0 = qml.measure(0)
@@ -909,8 +913,8 @@ class TestConditionalOperations:
             qml.ctrl(op, control=0)(phi=par, wires=1)
             return qml.expval(qml.PauliZ(1))
 
-        @qml.qnode(dev)
         @qml.defer_measurements
+        @qml.qnode(dev)
         def qnode2(par):
             qml.Hadamard(0)
             m_0 = qml.measure(0)
@@ -927,6 +931,7 @@ class TestConditionalOperations:
         outcome."""
         dev = qml.device("default.qubit", wires=3)
 
+        @qml.defer_measurements
         @qml.qnode(dev)
         def qnode():
             m_0 = qml.measure(0)
@@ -956,8 +961,8 @@ class TestConditionalOperations:
             qml.RY(x, wires=1)
             qml.PauliZ(1)
 
-        @qml.qnode(dev)
         @qml.defer_measurements
+        @qml.qnode(dev)
         def quantum_control_circuit(r):
             qml.Hadamard(0)
             m_0 = qml.measure(0)
@@ -999,6 +1004,7 @@ class TestConditionalOperations:
             qml.RX(a, wires=0)
             qml.PhaseShift(a, wires=0)
 
+        @qml.defer_measurements
         @qml.qnode(dev)
         def cond_qnode(x, y):
             qml.RY(x, wires=1)
@@ -1013,8 +1019,8 @@ class TestConditionalOperations:
         that is measured works as expected."""
         dev = qml.device("default.qubit", wires=2)
 
-        @qml.qnode(dev)
         @qml.defer_measurements
+        @qml.qnode(dev)
         def qnode():
             qml.Hadamard(0)
             m = qml.measure(0)
@@ -1045,8 +1051,8 @@ class TestExpressionConditionals:
             qml.ctrl(op, control=(0, 1), control_values=[True, True])(rads, wires=2)
             return qml.probs(wires=2)
 
-        @qml.qnode(dev)
         @qml.defer_measurements
+        @qml.qnode(dev)
         def quantum_control_circuit(rads):
             qml.RX(2.4, wires=0)
             qml.RY(1.3, wires=1)
@@ -1065,8 +1071,8 @@ class TestExpressionConditionals:
         """Test that combining the results of three mid-circuit measurements works as expected."""
         dev = qml.device("default.qubit", wires=7)
 
-        @qml.qnode(dev)
         @qml.defer_measurements
+        @qml.qnode(dev)
         def normal_circuit(rads):
             qml.RX(2.4, wires=0)
             qml.RY(1.3, wires=1)
@@ -1079,8 +1085,8 @@ class TestExpressionConditionals:
 
             return qml.probs(wires=3)
 
-        @qml.qnode(dev)
         @qml.defer_measurements
+        @qml.qnode(dev)
         def quantum_control_circuit(rads):
             qml.RX(2.4, wires=0)
             qml.RY(1.3, wires=1)
@@ -1107,8 +1113,8 @@ class TestExpressionConditionals:
         """
         dev = qml.device("default.qubit", wires=7)
 
-        @qml.qnode(dev)
         @qml.defer_measurements
+        @qml.qnode(dev)
         def normal_circuit(rads):
             qml.RX(2.4, wires=0)
             qml.RY(1.3, wires=1)
@@ -1124,8 +1130,8 @@ class TestExpressionConditionals:
 
             return qml.probs(wires=3)
 
-        @qml.qnode(dev)
         @qml.defer_measurements
+        @qml.qnode(dev)
         def quantum_control_circuit(rads):
             qml.RX(2.4, wires=0)
             qml.RY(1.3, wires=1)
@@ -1150,8 +1156,8 @@ class TestExpressionConditionals:
         """Test that a complex nested expression gets resolved correctly to the corresponding correct control gates."""
         dev = qml.device("default.qubit", wires=7)
 
-        @qml.qnode(dev)
         @qml.defer_measurements
+        @qml.qnode(dev)
         def normal_circuit(rads):
             qml.RX(2.4, wires=0)
             qml.RY(1.3, wires=1)
@@ -1168,8 +1174,8 @@ class TestExpressionConditionals:
 
             return qml.probs(wires=3)
 
-        @qml.qnode(dev)
         @qml.defer_measurements
+        @qml.qnode(dev)
         def quantum_control_circuit(rads):
             qml.RX(2.4, wires=0)
             qml.RY(1.3, wires=1)
@@ -1183,7 +1189,6 @@ class TestExpressionConditionals:
             expr2 = (3 * m_1 + 2) * (4 * expr1 + 2)
             m_2 = qml.measure(2)
             expr3 = expr2 / (m_2 + 3)
-            print(expr3)
             qml.cond(expr3 <= 6, qml.RX)(rads, wires=3)
             return qml.probs(wires=3)
 
@@ -1401,8 +1406,8 @@ class TestQubitReuseAndReset:
         after measurement if reset is requested."""
         dev = qml.device("default.qubit", wires=3)
 
-        @qml.qnode(dev)
         @qml.defer_measurements
+        @qml.qnode(dev)
         def qnode(x):
             qml.Hadamard(0)
             qml.PhaseShift(np.pi / 4, 0)
diff --git a/tests/transforms/test_tape_expand.py b/tests/transforms/test_tape_expand.py
index 3eff75985..7dfd56ca3 100644
--- a/tests/transforms/test_tape_expand.py
+++ b/tests/transforms/test_tape_expand.py
@@ -912,3 +912,38 @@ class TestCreateCustomDecompExpandFn:
             res.append(cost(x))
 
         assert res[0] == res[1]
+
+    @pytest.mark.parametrize("shots", [None, 100])
+    def test_custom_decomp_with_mcm(self, shots):
+        """Test that specifying a single custom decomposition works as expected."""
+
+        custom_decomps = {"Hadamard": custom_hadamard}
+        decomp_dev = qml.device("default.qubit", shots=shots, custom_decomps=custom_decomps)
+
+        @qml.qnode(decomp_dev, expansion_strategy="device")
+        def circuit():
+            qml.Hadamard(wires=0)
+            _ = qml.measure(0)
+            qml.CNOT(wires=[0, 1])
+            _ = qml.measure(1)
+            return qml.expval(qml.PauliZ(0))
+
+        _ = circuit()
+        decomp_ops = circuit.tape.operations
+
+        print(decomp_ops)
+        assert len(decomp_ops) == 4 if shots is None else 5
+
+        assert decomp_ops[0].name == "RZ"
+        assert np.isclose(decomp_ops[0].parameters[0], np.pi)
+
+        assert decomp_ops[1].name == "RY"
+        assert np.isclose(decomp_ops[1].parameters[0], np.pi / 2)
+
+        if shots is None:
+            assert decomp_ops[2].name == "CNOT"
+            assert decomp_ops[3].name == "CNOT"
+        else:
+            assert decomp_ops[2].name == "MidMeasureMP"
+            assert decomp_ops[3].name == "CNOT"
+            assert decomp_ops[4].name == "MidMeasureMP"