From 64a3deb4aca68921691a72865e4e6903941af7d3 Mon Sep 17 00:00:00 2001
From: Mudit Pandey <mudit.pandey@xanadu.ai>
Date: Mon, 28 Aug 2023 09:18:54 -0400
Subject: [PATCH 001/131] Updated hash/eq

---
 pennylane/measurements/measurements.py  | 24 ++------------
 pennylane/measurements/mid_measure.py   | 15 ++-------
 pennylane/operation.py                  | 23 ++-----------
 tests/measurements/test_measurements.py | 19 ++++++-----
 tests/test_operation.py                 | 43 ++++++-------------------
 5 files changed, 24 insertions(+), 100 deletions(-)

diff --git a/pennylane/measurements/measurements.py b/pennylane/measurements/measurements.py
index 09ea33fcdf5..4e9a57f479d 100644
--- a/pennylane/measurements/measurements.py
+++ b/pennylane/measurements/measurements.py
@@ -256,31 +256,11 @@ def diagonalizing_gates(self):
         except qml.operation.DecompositionUndefinedError:
             return []
 
-    # pylint: disable=useless-super-delegation
     def __eq__(self, other):
-        warnings.warn(
-            "The behaviour of measurement process equality will be updated soon. Currently, "
-            "mp1 == mp2 is True if mp1 and mp2 are the same object. Soon, mp1 == mp2 will be "
-            "equivalent to qml.equal(mp1, mp2). To continue using measurement process equality "
-            "in its current state, use 'mp1 is mp2'.",
-            UserWarning,
-        )
-
-        return super().__eq__(other)
+        return qml.equal(self, other)
 
-    # pylint: disable=useless-super-delegation
     def __hash__(self):
-        warnings.warn(
-            "The behaviour of measurement process hashing will be updated soon. Currently, each "
-            "measurement process instance has a unique hash. Soon, a measurement process's hash "
-            "will be determined by the combined hash of the name, wires, observable and/or "
-            "eigenvalues of the measurement process. To continue using measurement process hashing "
-            "in its current state, wrap the measurement process inside a qml.queuing.WrappedObj "
-            "instance.",
-            UserWarning,
-        )
-
-        return super().__hash__()
+        return self.hash
 
     def __repr__(self):
         """Representation of this class."""
diff --git a/pennylane/measurements/mid_measure.py b/pennylane/measurements/mid_measure.py
index e98a7513104..61e143eeed3 100644
--- a/pennylane/measurements/mid_measure.py
+++ b/pennylane/measurements/mid_measure.py
@@ -270,19 +270,8 @@ def _merge(self, other: "MeasurementValue"):
 
         # create a new function that selects the correct indices for each sub function
         def merged_fn(*x):
-            with warnings.catch_warnings():
-                # Using a filter because the new behaviour of MP equality will be valid here
-                warnings.filterwarnings(
-                    "ignore",
-                    message="The behaviour of measurement process equality",
-                    category=UserWarning,
-                )
-                sub_args_1 = (
-                    x[i] for i in [merged_measurements.index(m) for m in self.measurements]
-                )
-                sub_args_2 = (
-                    x[i] for i in [merged_measurements.index(m) for m in other.measurements]
-                )
+            sub_args_1 = (x[i] for i in [merged_measurements.index(m) for m in self.measurements])
+            sub_args_2 = (x[i] for i in [merged_measurements.index(m) for m in other.measurements])
 
             out_1 = self.processing_fn(*sub_args_1)
             out_2 = other.processing_fn(*sub_args_2)
diff --git a/pennylane/operation.py b/pennylane/operation.py
index 7f752c82c7a..ad9d16ecc85 100644
--- a/pennylane/operation.py
+++ b/pennylane/operation.py
@@ -719,30 +719,11 @@ def hash(self):
             )
         )
 
-    # pylint: disable=useless-super-delegation
     def __eq__(self, other):
-        warnings.warn(
-            "The behaviour of operator equality will be updated soon. Currently, op1 == op2 is "
-            "True if op1 and op2 are the same object. Soon, op1 == op2 will be equivalent to "
-            "qml.equal(op1, op2). To continue using operator equality in its current state, "
-            "use 'op1 is op2'.",
-            UserWarning,
-        )
-
-        return super().__eq__(other)
+        return qml.equal(self, other)
 
-    # pylint: disable=useless-super-delegation
     def __hash__(self):
-        warnings.warn(
-            "The behaviour of operator hashing will be updated soon. Currently, each operator "
-            "instance has a unique hash. Soon, an operator's hash will be determined by the "
-            "combined hash of the name, wires, parameters and hyperparameters of the operator. "
-            "To continue using operator hashing in its current state, wrap the operator inside "
-            "a qml.queuing.WrappedObj instance.",
-            UserWarning,
-        )
-
-        return super().__hash__()
+        return self.hash
 
     @staticmethod
     def compute_matrix(*params, **hyperparams):  # pylint:disable=unused-argument
diff --git a/tests/measurements/test_measurements.py b/tests/measurements/test_measurements.py
index 3b8160419a5..1b0c9b16b9a 100644
--- a/tests/measurements/test_measurements.py
+++ b/tests/measurements/test_measurements.py
@@ -125,8 +125,8 @@ class NoReturnTypeMeasurement(MeasurementProcess):
 
 
 def test_eq_correctness():
-    """Test that using `==` on two equivalent operators is True when both measurement
-    processes are the same object and False otherwise."""
+    """Test that using `==` on measurement processes behaves the same as
+    `qml.equal`."""
 
     class DummyMP(MeasurementProcess):
         """Dummy measurement process with no return type."""
@@ -134,14 +134,12 @@ class DummyMP(MeasurementProcess):
     mp1 = DummyMP(0)
     mp2 = DummyMP(0)
 
-    with pytest.warns(UserWarning, match="The behaviour of measurement process equality"):
-        assert mp1 == mp1  # pylint: disable=comparison-with-itself
-        assert mp1 != mp2
+    assert mp1 == mp1  # pylint: disable=comparison-with-itself
+    assert mp1 == mp2
 
 
 def test_hash_correctness():
-    """Test that the hash of two equivalent measurement processes is the same when
-    both are the same object and different otherwise."""
+    """Test that the hash of two equivalent measurement processes is the same."""
 
     class DummyMP(MeasurementProcess):
         """Dummy measurement process with no return type."""
@@ -149,9 +147,10 @@ class DummyMP(MeasurementProcess):
     mp1 = DummyMP(0)
     mp2 = DummyMP(0)
 
-    with pytest.warns(UserWarning, match="The behaviour of measurement process hashing"):
-        assert len({mp1, mp1}) == 1
-        assert len({mp1, mp2}) == 2
+    assert len({mp1, mp2}) == 1
+    assert hash(mp1) == mp1.hash
+    assert hash(mp2) == mp2.hash
+    assert hash(mp1) == hash(mp2)
 
 
 @pytest.mark.parametrize(
diff --git a/tests/test_operation.py b/tests/test_operation.py
index 2485edf76d0..7c69f4579b1 100644
--- a/tests/test_operation.py
+++ b/tests/test_operation.py
@@ -272,22 +272,9 @@ def __init__(self, wires, basis_state=None):  # pylint:disable=super-init-not-ca
         state = [0, 1, 0]
         assert MyOp(wires=1, basis_state=state).hyperparameters["basis_state"] == state
 
-    def test_eq_warning(self):
-        """Test that a warning is raised when two operators are compared for equality
-        using `==`."""
-
-        class DummyOp(qml.operation.Operator):
-            num_wires = 1
-
-        op1 = DummyOp(0)
-        op2 = DummyOp(0)
-
-        with pytest.warns(UserWarning, match="The behaviour of operator equality"):
-            _ = op1 == op2
-
     def test_eq_correctness(self):
-        """Test that using `==` on two equivalent operators is True when both operators
-        are the same object and False otherwise."""
+        """Test that using `==` on operators behaves the same as
+        `qml.equal`."""
 
         class DummyOp(qml.operation.Operator):
             num_wires = 1
@@ -295,24 +282,11 @@ class DummyOp(qml.operation.Operator):
         op1 = DummyOp(0)
         op2 = DummyOp(0)
 
-        with pytest.warns(UserWarning, match="The behaviour of operator equality"):
-            assert op1 == op1  # pylint: disable=comparison-with-itself
-            assert op1 != op2
-
-    def test_hash_warning(self):
-        """Test that a warning is raised when an operator's hash is used."""
-
-        class DummyOp(qml.operation.Operator):
-            num_wires = 1
-
-        op = DummyOp(0)
-
-        with pytest.warns(UserWarning, match="The behaviour of operator hashing"):
-            _ = hash(op)
+        assert op1 == op1  # pylint: disable=comparison-with-itself
+        assert op1 == op2
 
     def test_hash_correctness(self):
-        """Test that the hash of two equivalent operators is the same when both operators
-        are the same object and different otherwise."""
+        """Test that the hash of two equivalent operators is the same."""
 
         class DummyOp(qml.operation.Operator):
             num_wires = 1
@@ -320,9 +294,10 @@ class DummyOp(qml.operation.Operator):
         op1 = DummyOp(0)
         op2 = DummyOp(0)
 
-        with pytest.warns(UserWarning, match="The behaviour of operator hash"):
-            assert len({op1, op1}) == 1
-            assert len({op1, op2}) == 2
+        assert len({op1, op2}) == 1
+        assert hash(op1) == op1.hash
+        assert hash(op2) == op2.hash
+        assert hash(op1) == hash(op2)
 
 
 class TestPytreeMethods:

From ed9d1f9ccdf8ce903c127ab93685b077312af4b6 Mon Sep 17 00:00:00 2001
From: Mudit Pandey <mudit.pandey@xanadu.ai>
Date: Mon, 28 Aug 2023 10:00:49 -0400
Subject: [PATCH 002/131] Updated pytest.ini

---
 tests/pytest.ini | 2 --
 1 file changed, 2 deletions(-)

diff --git a/tests/pytest.ini b/tests/pytest.ini
index 58ec919029f..0171b47641d 100644
--- a/tests/pytest.ini
+++ b/tests/pytest.ini
@@ -19,5 +19,3 @@ filterwarnings =
     ignore:Casting complex values to real discards the imaginary part:UserWarning:torch.autograd
     ignore:Call to deprecated create function:DeprecationWarning
     ignore:the imp module is deprecated:DeprecationWarning
-    error:The behaviour of operator:UserWarning
-    error:The behaviour of measurement process:UserWarning

From 1c9a2563cab80af6cf9860ddea8bd8b37b414b2a Mon Sep 17 00:00:00 2001
From: Mudit Pandey <mudit.pandey@xanadu.ai>
Date: Mon, 28 Aug 2023 15:28:18 -0400
Subject: [PATCH 003/131] Updated changelog

---
 doc/releases/changelog-dev.md | 10 ++++++++--
 1 file changed, 8 insertions(+), 2 deletions(-)

diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 6e674c8597a..a1e5abf0d3d 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -8,11 +8,17 @@
 
 <h3>Breaking changes 💔</h3>
 
+* The `__eq__` and `__hash__` methods of `Operator` and `MeasurementProcess` no longer rely on the
+  object's address is memory. Using `==` with operators and measurement processes will now behave the
+  same as `qml.equal`, and objects of the same type with the same data and hyperparameters will have
+  the same hash.
+  [(#)]()
+
 * The old return type and associated functions ``qml.enable_return`` and ``qml.disable_return`` are removed.
- [#4503](https://github.com/PennyLaneAI/pennylane/pull/4503)
+ [(#4503)](https://github.com/PennyLaneAI/pennylane/pull/4503)
 
 * The ``mode`` keyword argument in ``QNode`` is removed. Please use ``grad_on_execution`` instead.
- [#4503](https://github.com/PennyLaneAI/pennylane/pull/4503)
+ [(#4503)](https://github.com/PennyLaneAI/pennylane/pull/4503)
 
 <h3>Deprecations 👋</h3>
 

From 593115d8cc9ab28fce0d38ace354725f56f7b162 Mon Sep 17 00:00:00 2001
From: Mudit Pandey <mudit.pandey@xanadu.ai>
Date: Mon, 28 Aug 2023 15:32:22 -0400
Subject: [PATCH 004/131] Update doc/releases/changelog-dev.md

---
 doc/releases/changelog-dev.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index a1e5abf0d3d..1faf5b9a428 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -12,7 +12,7 @@
   object's address is memory. Using `==` with operators and measurement processes will now behave the
   same as `qml.equal`, and objects of the same type with the same data and hyperparameters will have
   the same hash.
-  [(#)]()
+  [(#4536)](https://github.com/PennyLaneAI/pennylane/pull/4536)
 
 * The old return type and associated functions ``qml.enable_return`` and ``qml.disable_return`` are removed.
  [(#4503)](https://github.com/PennyLaneAI/pennylane/pull/4503)

From 4380982e1492aafca394a0dd04d8940b03a33904 Mon Sep 17 00:00:00 2001
From: Mudit Pandey <mudit.pandey@xanadu.ai>
Date: Mon, 28 Aug 2023 15:40:24 -0400
Subject: [PATCH 005/131] Removed warning imports

---
 pennylane/measurements/measurements.py |  1 -
 pennylane/measurements/mid_measure.py  | 13 ++-----------
 2 files changed, 2 insertions(+), 12 deletions(-)

diff --git a/pennylane/measurements/measurements.py b/pennylane/measurements/measurements.py
index 4e9a57f479d..c195e201187 100644
--- a/pennylane/measurements/measurements.py
+++ b/pennylane/measurements/measurements.py
@@ -22,7 +22,6 @@
 from abc import ABC, abstractmethod
 from enum import Enum
 from typing import Sequence, Tuple, Optional
-import warnings
 
 import pennylane as qml
 from pennylane.operation import Operator
diff --git a/pennylane/measurements/mid_measure.py b/pennylane/measurements/mid_measure.py
index 61e143eeed3..92d7ce411fd 100644
--- a/pennylane/measurements/mid_measure.py
+++ b/pennylane/measurements/mid_measure.py
@@ -16,7 +16,6 @@
 """
 import uuid
 from typing import Generic, TypeVar, Optional
-import warnings
 
 import pennylane as qml
 import pennylane.numpy as np
@@ -256,16 +255,8 @@ def _apply(self, fn):
     def _merge(self, other: "MeasurementValue"):
         """Merge two measurement values"""
 
-        with warnings.catch_warnings():
-            # Using a filter because the new behaviour of MP hash will be valid here
-            warnings.filterwarnings(
-                "ignore",
-                message="The behaviour of measurement process hashing",
-                category=UserWarning,
-            )
-            # create a new merged list with no duplicates and in lexical ordering
-            merged_measurements = list(set(self.measurements).union(set(other.measurements)))
-
+        # create a new merged list with no duplicates and in lexical ordering
+        merged_measurements = list(set(self.measurements).union(set(other.measurements)))
         merged_measurements.sort(key=lambda m: m.id)
 
         # create a new function that selects the correct indices for each sub function

From 12e3fb4abcf0888aac402bf2336bfc7fee8944c7 Mon Sep 17 00:00:00 2001
From: Mudit Pandey <mudit.pandey@xanadu.ai>
Date: Mon, 28 Aug 2023 16:01:35 -0400
Subject: [PATCH 006/131] Update tests

---
 tests/measurements/test_measurements.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/tests/measurements/test_measurements.py b/tests/measurements/test_measurements.py
index 1b0c9b16b9a..3c0ee019cfa 100644
--- a/tests/measurements/test_measurements.py
+++ b/tests/measurements/test_measurements.py
@@ -131,8 +131,8 @@ def test_eq_correctness():
     class DummyMP(MeasurementProcess):
         """Dummy measurement process with no return type."""
 
-    mp1 = DummyMP(0)
-    mp2 = DummyMP(0)
+    mp1 = DummyMP(wires=0)
+    mp2 = DummyMP(wires=0)
 
     assert mp1 == mp1  # pylint: disable=comparison-with-itself
     assert mp1 == mp2
@@ -144,8 +144,8 @@ def test_hash_correctness():
     class DummyMP(MeasurementProcess):
         """Dummy measurement process with no return type."""
 
-    mp1 = DummyMP(0)
-    mp2 = DummyMP(0)
+    mp1 = DummyMP(wires=0)
+    mp2 = DummyMP(wires=0)
 
     assert len({mp1, mp2}) == 1
     assert hash(mp1) == mp1.hash

From a052b7efba495c27536f6aa5307a91c3a6cc5aed Mon Sep 17 00:00:00 2001
From: Mudit Pandey <mudit.pandey@xanadu.ai>
Date: Tue, 29 Aug 2023 12:09:57 -0400
Subject: [PATCH 007/131] Fixed test

---
 tests/measurements/test_measurements.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/tests/measurements/test_measurements.py b/tests/measurements/test_measurements.py
index 3c0ee019cfa..8f1b1d532a4 100644
--- a/tests/measurements/test_measurements.py
+++ b/tests/measurements/test_measurements.py
@@ -131,8 +131,8 @@ def test_eq_correctness():
     class DummyMP(MeasurementProcess):
         """Dummy measurement process with no return type."""
 
-    mp1 = DummyMP(wires=0)
-    mp2 = DummyMP(wires=0)
+    mp1 = DummyMP(wires=qml.wires.Wires(0))
+    mp2 = DummyMP(wires=qml.wires.Wires(0))
 
     assert mp1 == mp1  # pylint: disable=comparison-with-itself
     assert mp1 == mp2
@@ -144,8 +144,8 @@ def test_hash_correctness():
     class DummyMP(MeasurementProcess):
         """Dummy measurement process with no return type."""
 
-    mp1 = DummyMP(wires=0)
-    mp2 = DummyMP(wires=0)
+    mp1 = DummyMP(wires=qml.wires.Wires(0))
+    mp2 = DummyMP(wires=qml.wires.Wires(0))
 
     assert len({mp1, mp2}) == 1
     assert hash(mp1) == mp1.hash

From 664772995fb09df1bfb2b411f830bf2584a9f569 Mon Sep 17 00:00:00 2001
From: Mudit Pandey <mudit.pandey@xanadu.ai>
Date: Tue, 29 Aug 2023 13:33:59 -0400
Subject: [PATCH 008/131] Added single MV support for MPs

---
 pennylane/_device.py                          |  3 +-
 pennylane/_qubit_device.py                    | 55 +++++++++++++-----
 pennylane/_qutrit_device.py                   |  4 +-
 pennylane/devices/qubit/measure.py            | 15 ++++-
 pennylane/interfaces/jax.py                   |  2 +-
 pennylane/measurements/counts.py              | 17 ++++--
 pennylane/measurements/expval.py              | 35 ++++++++---
 pennylane/measurements/measurements.py        | 33 ++++++++---
 pennylane/measurements/probs.py               | 14 +++--
 pennylane/measurements/sample.py              | 58 ++++++++++++++-----
 pennylane/measurements/var.py                 | 34 ++++++++---
 pennylane/tape/qscript.py                     | 14 +++--
 .../transforms/convert_to_numpy_parameters.py |  7 ++-
 13 files changed, 214 insertions(+), 77 deletions(-)

diff --git a/pennylane/_device.py b/pennylane/_device.py
index bb46b77bf04..b3552f8f49a 100644
--- a/pennylane/_device.py
+++ b/pennylane/_device.py
@@ -30,6 +30,7 @@
     CountsMP,
     Expectation,
     ExpectationMP,
+    MeasurementValue,
     MidMeasureMP,
     Probability,
     ProbabilityMP,
@@ -996,7 +997,7 @@ def check_validity(self, queue, observables):
                         raise DeviceError(
                             f"Observable {i.name} not supported on device {self.short_name}"
                         )
-            else:
+            elif not isinstance(o, MeasurementValue):
                 observable_name = o.name
 
                 if not self.supports_observable(observable_name):
diff --git a/pennylane/_qubit_device.py b/pennylane/_qubit_device.py
index da1857ad86f..bb9c187014b 100644
--- a/pennylane/_qubit_device.py
+++ b/pennylane/_qubit_device.py
@@ -41,6 +41,7 @@
     ExpectationMP,
     MeasurementProcess,
     MeasurementTransform,
+    MeasurementValue,
     MutualInfoMP,
     ProbabilityMP,
     SampleMeasurement,
@@ -612,7 +613,7 @@ def statistics(
 
         for m in measurements:
             # TODO: Remove this when all overriden measurements support the `MeasurementProcess` class
-            if m.obs is not None:
+            if m.obs is not None and not isinstance(m.obs, MeasurementValue):
                 obs = m.obs
                 obs.return_type = m.return_type
             else:
@@ -1316,7 +1317,12 @@ def expval(self, observable, shot_range=None, bin_size=None):
         # exact expectation value
         if self.shots is None:
             try:
-                eigvals = self._asarray(observable.eigvals(), dtype=self.R_DTYPE)
+                eigvals = self._asarray(
+                    observable.eigvals()
+                    if not isinstance(observable, MeasurementProcess)
+                    else np.arange(2 ** len(observable.wires)),
+                    dtype=self.R_DTYPE,
+                )
             except qml.operation.EigvalsUndefinedError as e:
                 raise qml.operation.EigvalsUndefinedError(
                     f"Cannot compute analytic expectations of {observable.name}."
@@ -1327,7 +1333,10 @@ def expval(self, observable, shot_range=None, bin_size=None):
             return self._dot(prob, eigvals)
 
         # estimate the ev
-        samples = self.sample(observable, shot_range=shot_range, bin_size=bin_size)
+        # Get samples as decimal integer values if computing ev for a MeasurementValue,
+        # otherwise the returned samples would be boolean lists
+        decimal = isinstance(observable, MeasurementProcess)
+        samples = self.sample(observable, shot_range=shot_range, bin_size=bin_size, decimal=decimal)
         # With broadcasting, we want to take the mean over axis 1, which is the -1st/-2nd with/
         # without bin_size. Without broadcasting, axis 0 is the -1st/-2nd with/without bin_size
         axis = -1 if bin_size is None else -2
@@ -1346,7 +1355,12 @@ def var(self, observable, shot_range=None, bin_size=None):
         # exact variance value
         if self.shots is None:
             try:
-                eigvals = self._asarray(observable.eigvals(), dtype=self.R_DTYPE)
+                eigvals = self._asarray(
+                    observable.eigvals()
+                    if not isinstance(observable, MeasurementProcess)
+                    else np.arange(2 ** len(observable.wires)),
+                    dtype=self.R_DTYPE,
+                )
             except qml.operation.EigvalsUndefinedError as e:
                 # if observable has no info on eigenvalues, we cannot return this measurement
                 raise qml.operation.EigvalsUndefinedError(
@@ -1358,7 +1372,11 @@ def var(self, observable, shot_range=None, bin_size=None):
             return self._dot(prob, (eigvals**2)) - self._dot(prob, eigvals) ** 2
 
         # estimate the variance
-        samples = self.sample(observable, shot_range=shot_range, bin_size=bin_size)
+        # Get samples as decimal integer values if computing ev for a MeasurementValue,
+        # otherwise the returned samples would be boolean lists
+        decimal = isinstance(observable, MeasurementProcess)
+        samples = self.sample(observable, shot_range=shot_range, bin_size=bin_size, decimal=decimal)
+
         # With broadcasting, we want to take the variance over axis 1, which is the -1st/-2nd with/
         # without bin_size. Without broadcasting, axis 0 is the -1st/-2nd with/without bin_size
         axis = -1 if bin_size is None else -2
@@ -1445,7 +1463,9 @@ def circuit(x):
 
         return outcome_dicts if batched else outcome_dicts[0]
 
-    def sample(self, observable, shot_range=None, bin_size=None, counts=False):
+    def sample(
+        self, observable, shot_range=None, bin_size=None, counts=False, decimal=False
+    ):  # pylint: disable=too-many-arguments
         """Return samples of an observable.
 
         Args:
@@ -1457,6 +1477,8 @@ def sample(self, observable, shot_range=None, bin_size=None, counts=False):
                 provided, the entire shot range is treated as a single bin.
             counts (bool): whether counts (``True``) or raw samples (``False``)
                 should be returned
+            decimal (bool): Whether to return samples as base-10 integers or boolean lists.
+                Setting ``decimal`` does not do anything if ``counts`` is ``True``.
 
         Raises:
             EigvalsUndefinedError: if no information is available about the
@@ -1485,7 +1507,7 @@ def sample(self, observable, shot_range=None, bin_size=None, counts=False):
             # Process samples for observables with eigenvalues {1, -1}
             samples = 1 - 2 * sub_samples[..., device_wires[0]]
 
-        elif no_observable_provided:
+        elif no_observable_provided and (not decimal or counts):
             # if no observable was provided then return the raw samples
             if len(observable.wires) != 0:
                 # if wires are provided, then we only return samples from those wires
@@ -1500,13 +1522,16 @@ def sample(self, observable, shot_range=None, bin_size=None, counts=False):
             powers_of_two = 2 ** np.arange(samples.shape[-1])[::-1]
             indices = samples @ powers_of_two
             indices = np.array(indices)  # Add np.array here for Jax support.
-            try:
-                samples = observable.eigvals()[indices]
-            except qml.operation.EigvalsUndefinedError as e:
-                # if observable has no info on eigenvalues, we cannot return this measurement
-                raise qml.operation.EigvalsUndefinedError(
-                    f"Cannot compute samples of {observable.name}."
-                ) from e
+            if decimal:
+                samples = indices
+            else:
+                try:
+                    samples = observable.eigvals()[indices]
+                except qml.operation.EigvalsUndefinedError as e:
+                    # if observable has no info on eigenvalues, we cannot return this measurement
+                    raise qml.operation.EigvalsUndefinedError(
+                        f"Cannot compute samples of {observable.name}."
+                    ) from e
 
         num_wires = len(device_wires) if len(device_wires) > 0 else self.num_wires
         if bin_size is None:
@@ -1523,7 +1548,7 @@ def sample(self, observable, shot_range=None, bin_size=None, counts=False):
 
         return (
             samples.T.reshape((num_wires, bin_size, -1))
-            if no_observable_provided
+            if no_observable_provided and not decimal
             else samples.reshape((bin_size, -1))
         )
 
diff --git a/pennylane/_qutrit_device.py b/pennylane/_qutrit_device.py
index d9513a33d06..7739878f752 100644
--- a/pennylane/_qutrit_device.py
+++ b/pennylane/_qutrit_device.py
@@ -357,7 +357,9 @@ def marginal_prob(self, prob, wires=None):
         prob = self._transpose(prob, np.argsort(np.argsort(device_wires)))
         return self._flatten(prob)
 
-    def sample(self, observable, shot_range=None, bin_size=None, counts=False):
+    def sample(
+        self, observable, shot_range=None, bin_size=None, counts=False, decimal=False
+    ):  # pylint: disable=too-many-arguments, unused-argument
         def _samples_to_counts(samples, no_observable_provided):
             """Group the obtained samples into a dictionary.
 
diff --git a/pennylane/devices/qubit/measure.py b/pennylane/devices/qubit/measure.py
index e30a0c809b0..fef0ac5fd4c 100644
--- a/pennylane/devices/qubit/measure.py
+++ b/pennylane/devices/qubit/measure.py
@@ -20,7 +20,12 @@
 
 from pennylane import math
 from pennylane.ops import Sum, Hamiltonian
-from pennylane.measurements import StateMeasurement, MeasurementProcess, ExpectationMP
+from pennylane.measurements import (
+    StateMeasurement,
+    MeasurementProcess,
+    MeasurementValue,
+    ExpectationMP,
+)
 from pennylane.typing import TensorLike
 from pennylane.wires import Wires
 
@@ -132,6 +137,8 @@ def get_measurement_function(
     """
     if isinstance(measurementprocess, StateMeasurement):
         if isinstance(measurementprocess, ExpectationMP):
+            if isinstance(measurementprocess.obs, MeasurementValue):
+                return state_diagonalizing_gates
             if measurementprocess.obs.name == "SparseHamiltonian":
                 return csr_dot_products
 
@@ -153,7 +160,11 @@ def get_measurement_function(
 
                     return csr_dot_products
 
-        if measurementprocess.obs is None or measurementprocess.obs.has_diagonalizing_gates:
+        if (
+            measurementprocess.obs is None
+            or isinstance(measurementprocess.obs, MeasurementValue)
+            or measurementprocess.obs.has_diagonalizing_gates
+        ):
             return state_diagonalizing_gates
 
     raise NotImplementedError
diff --git a/pennylane/interfaces/jax.py b/pennylane/interfaces/jax.py
index 632be8a51a7..8b7bd1cafc6 100644
--- a/pennylane/interfaces/jax.py
+++ b/pennylane/interfaces/jax.py
@@ -67,7 +67,7 @@ def get_jax_interface_name(tapes):
         for op in t:
             # Unwrap the observable from a MeasurementProcess
             op = op.obs if hasattr(op, "obs") else op
-            if op is not None:
+            if op is not None and not isinstance(op, qml.measurements.MeasurementValue):
                 # Some MeasurementProcess objects have op.obs=None
                 for param in op.data:
                     if qml.math.is_abstract(param):
diff --git a/pennylane/measurements/counts.py b/pennylane/measurements/counts.py
index fc8d629716a..e317785fed7 100644
--- a/pennylane/measurements/counts.py
+++ b/pennylane/measurements/counts.py
@@ -22,6 +22,7 @@
 from pennylane.wires import Wires
 
 from .measurements import AllCounts, Counts, SampleMeasurement
+from .mid_measure import MeasurementValue
 
 
 def _sample_to_str(sample):
@@ -39,7 +40,8 @@ def counts(op=None, wires=None, all_outcomes=False) -> "CountsMP":
     specified on the device.
 
     Args:
-        op (Observable or None): a quantum observable object
+        op (Observable or MeasurementValue or None): a quantum observable object. To get counts
+            for mid-circuit measurements, ``op`` should be a ``MeasurementValue``.
         wires (Sequence[int] or int or None): the wires we wish to sample from, ONLY set wires if
             op is None
         all_outcomes(bool): determines whether the returned dict will contain only the observed
@@ -131,6 +133,9 @@ def circuit():
     {'00': 0, '01': 0, '10': 4, '11': 0}
 
     """
+    if isinstance(op, MeasurementValue):
+        return CountsMP(obs=op, all_outcomes=all_outcomes)
+
     if op is not None and not op.is_hermitian:  # None type is also allowed for op
         warnings.warn(f"{op.name} might not be hermitian.")
 
@@ -152,9 +157,9 @@ class CountsMP(SampleMeasurement):
     Please refer to :func:`counts` for detailed documentation.
 
     Args:
-        obs (.Operator): The observable that is to be measured as part of the
-            measurement process. Not all measurement processes require observables (for
-            example ``Probability``); this argument is optional.
+        obs (Union[.Operator, .MeasurementValue]): The observable that is to be measured
+            as part of the measurement process. Not all measurement processes require observables
+            (for example ``Probability``); this argument is optional.
         wires (.Wires): The wires the measurement process applies to.
             This can only be specified if an observable was not provided.
         eigvals (array): A flat array representing the eigenvalues of the measurement.
@@ -219,7 +224,7 @@ def process_samples(
         num_wires = len(self.wires) if self.wires else len(wire_order)
         samples = (
             samples.reshape((num_wires, -1)).T.reshape(-1, bin_size, num_wires)
-            if self.obs is None
+            if self.obs is None or isinstance(self.obs, MeasurementValue)
             else samples.reshape((-1, bin_size))
         )
 
@@ -277,7 +282,7 @@ def circuit(x):
         batched_ndims = 2
         shape = qml.math.shape(samples)
 
-        if self.obs is None:
+        if self.obs is None or isinstance(self.obs, MeasurementValue):
             # convert samples and outcomes (if using) from arrays to str for dict keys
             samples = qml.math.cast_like(samples, qml.math.int8(0))
             samples = qml.math.apply_along_axis(_sample_to_str, -1, samples)
diff --git a/pennylane/measurements/expval.py b/pennylane/measurements/expval.py
index 3e380cc1cde..5df3d142b24 100644
--- a/pennylane/measurements/expval.py
+++ b/pennylane/measurements/expval.py
@@ -15,7 +15,7 @@
 This module contains the qml.expval measurement.
 """
 import warnings
-from typing import Sequence, Tuple
+from typing import Sequence, Tuple, Union
 
 import pennylane as qml
 from pennylane.operation import Operator
@@ -23,9 +23,10 @@
 from pennylane.wires import Wires
 
 from .measurements import Expectation, SampleMeasurement, StateMeasurement
+from .mid_measure import MeasurementValue
 
 
-def expval(op: Operator):
+def expval(op: Union[Operator, MeasurementValue]):
     r"""Expectation value of the supplied observable.
 
     **Example:**
@@ -47,11 +48,16 @@ def circuit(x):
     -0.4794255386042029
 
     Args:
-        op (Observable): a quantum observable object
+        op (Union[Observable, MeasurementValue]): a quantum observable object. To
+            get expectation values for mid-circuit measurements, ``op`` should be
+            a ``MeasurementValue``.
 
     Returns:
         ExpectationMP: measurement process instance
     """
+    if isinstance(op, MeasurementValue):
+        return ExpectationMP(obs=op)
+
     if not op.is_hermitian:
         warnings.warn(f"{op.name} might not be hermitian.")
 
@@ -64,9 +70,9 @@ class ExpectationMP(SampleMeasurement, StateMeasurement):
     Please refer to :func:`expval` for detailed documentation.
 
     Args:
-        obs (.Operator): The observable that is to be measured as part of the
-            measurement process. Not all measurement processes require observables (for
-            example ``Probability``); this argument is optional.
+        obs (Union[.Operator, .MeasurementValue]): The observable that is to be measured
+            as part of the measurement process. Not all measurement processes require observables
+            (for example ``Probability``); this argument is optional.
         wires (.Wires): The wires the measurement process applies to.
             This can only be specified if an observable was not provided.
         eigvals (array): A flat array representing the eigenvalues of the measurement.
@@ -103,10 +109,19 @@ def process_samples(
                 samples=samples, wire_order=wire_order, shot_range=shot_range, bin_size=bin_size
             )
             return probs[idx]
+
         # estimate the ev
+        # Get samples as decimal integer values if computing ev for a MeasurementValue,
+        # otherwise the returned samples would be boolean lists
+        decimal = isinstance(self.obs, MeasurementValue)
         samples = qml.sample(op=self.obs).process_samples(
-            samples=samples, wire_order=wire_order, shot_range=shot_range, bin_size=bin_size
+            samples=samples,
+            wire_order=wire_order,
+            shot_range=shot_range,
+            bin_size=bin_size,
+            decimal=decimal,
         )
+
         # With broadcasting, we want to take the mean over axis 1, which is the -1st/-2nd with/
         # without bin_size. Without broadcasting, axis 0 is the -1st/-2nd with/without bin_size
         axis = -1 if bin_size is None else -2
@@ -119,7 +134,11 @@ def process_state(self, state: Sequence[complex], wire_order: Wires):
             idx = int("".join(str(i) for i in self.obs.parameters[0]), 2)
             probs = qml.probs(wires=self.wires).process_state(state=state, wire_order=wire_order)
             return probs[idx]
-        eigvals = qml.math.asarray(self.obs.eigvals(), dtype="float64")
+        eigvals = (
+            qml.math.asarray(self.obs.eigvals(), dtype="float64")
+            if not isinstance(self.obs, MeasurementValue)
+            else qml.math.asarray(qml.math.arange(0, 2 ** len(self.wires)), dtype="float64")
+        )
         # we use ``self.wires`` instead of ``self.obs`` because the observable was
         # already applied to the state
         prob = qml.probs(wires=self.wires).process_state(state=state, wire_order=wire_order)
diff --git a/pennylane/measurements/measurements.py b/pennylane/measurements/measurements.py
index 09ea33fcdf5..ddc11dfa7b1 100644
--- a/pennylane/measurements/measurements.py
+++ b/pennylane/measurements/measurements.py
@@ -116,7 +116,7 @@ class MeasurementProcess(ABC):
     quantum variational circuit.
 
     Args:
-        obs (.Operator): The observable that is to be measured as part of the
+        obs (Union[.Operator, .MeasurementValue]): The observable that is to be measured as part of the
             measurement process. Not all measurement processes require observables (for
             example ``Probability``); this argument is optional.
         wires (.Wires): The wires the measurement process applies to.
@@ -146,7 +146,11 @@ def __init__(
 
         # _wires = None indicates broadcasting across all available wires.
         # It translates to the public property wires = Wires([])
-        self._wires = wires
+        self._wires = (
+            self.obs.measurements[0].wires
+            if self.obs is not None and self.obs.__class__.__name__ == "MeasurementValue"
+            else wires
+        )
         self._eigvals = None
 
         if eigvals is not None:
@@ -313,7 +317,7 @@ def wires(self):
 
         This is the union of all the Wires objects of the measurement.
         """
-        if self.obs is not None:
+        if self.obs is not None and self.obs.__class__.__name__ != "MeasurementValue":
             return self.obs.wires
 
         return (
@@ -352,6 +356,9 @@ def eigvals(self):
         Returns:
             array: eigvals representation
         """
+        if self.obs.__class__.__name__ == "MeasurementValue":
+            return qml.math.arange(0, 2 ** len(self.wires), 1)
+
         if self.obs is not None:
             with contextlib.suppress(qml.operation.EigvalsUndefinedError):
                 return self.obs.eigvals()
@@ -365,12 +372,16 @@ def has_decomposition(self):
         # If self.obs is not None, `expand` queues the diagonalizing gates of self.obs,
         # which we have to check to be defined. The subsequent creation of the new
         # `MeasurementProcess` within `expand` should never fail with the given parameters.
-        return False if self.obs is None else self.obs.has_diagonalizing_gates
+        return (
+            self.obs.has_diagonalizing_gates
+            if self.obs is not None and self.obs.__class__.__name__ != "MeasurementValue"
+            else False
+        )
 
     @property
     def samples_computational_basis(self):
         r"""Bool: Whether or not the MeasurementProcess measures in the computational basis."""
-        return self.obs is None
+        return self.obs is None or self.obs.__class__.__name__ == "MeasurementValue"
 
     def expand(self):
         """Expand the measurement of an observable to a unitary
@@ -405,7 +416,7 @@ def expand(self):
         >>> print(tape.measurements[0].obs)
         None
         """
-        if self.obs is None:
+        if self.obs is None or self.obs.__class__.__name__ == "MeasurementValue":
             raise qml.operation.DecompositionUndefinedError
 
         with qml.queuing.AnnotatedQueue() as q:
@@ -416,7 +427,7 @@ def expand(self):
 
     def queue(self, context=qml.QueuingManager):
         """Append the measurement process to an annotated queue."""
-        if self.obs is not None:
+        if self.obs is not None and self.obs.__class__.__name__ != "MeasurementValue":
             context.remove(self.obs)
         context.append(self)
 
@@ -450,7 +461,11 @@ def simplify(self):
         Returns:
             .MeasurementProcess: A measurement process with a simplified observable.
         """
-        return self if self.obs is None else self.__class__(obs=self.obs.simplify())
+        return (
+            self
+            if self.obs is None or self.obs.__class__.__name__ == "MeasurementValue"
+            else self.__class__(obs=self.obs.simplify())
+        )
 
     # pylint: disable=protected-access
     def map_wires(self, wire_map: dict):
@@ -464,7 +479,7 @@ def map_wires(self, wire_map: dict):
             .MeasurementProcess: new measurement process
         """
         new_measurement = copy.copy(self)
-        if self.obs is not None:
+        if self.obs is not None and self.obs.__class__.__name__ != "MeasurementValue":
             new_measurement.obs = self.obs.map_wires(wire_map=wire_map)
         else:
             new_measurement._wires = Wires([wire_map.get(wire, wire) for wire in self.wires])
diff --git a/pennylane/measurements/probs.py b/pennylane/measurements/probs.py
index ab79326e612..950a0c51fa6 100644
--- a/pennylane/measurements/probs.py
+++ b/pennylane/measurements/probs.py
@@ -21,6 +21,7 @@
 from pennylane.wires import Wires
 
 from .measurements import Probability, SampleMeasurement, StateMeasurement
+from .mid_measure import MeasurementValue
 
 
 def probs(wires=None, op=None) -> "ProbabilityMP":
@@ -42,8 +43,9 @@ def probs(wires=None, op=None) -> "ProbabilityMP":
 
     Args:
         wires (Sequence[int] or int): the wire the operation acts on
-        op (Observable): Observable (with a ``diagonalizing_gates`` attribute) that rotates
-            the computational basis
+        op (Observable or MeasurementValue]): Observable (with a ``diagonalizing_gates``
+            attribute) that rotates the computational basis, or a  ``MeasurementValue``
+            corresponding to mid-circuit measurements.
 
     Returns:
         ProbabilityMP: Measurement process instance
@@ -90,6 +92,8 @@ def circuit():
     Note that the output shape of this measurement process depends on whether
     the device simulates qubit or continuous variable quantum systems.
     """
+    if isinstance(op, MeasurementValue):
+        return ProbabilityMP(obs=op)
 
     if isinstance(op, qml.Hamiltonian):
         raise qml.QuantumFunctionError("Hamiltonians are not supported for rotating probabilities.")
@@ -120,9 +124,9 @@ class ProbabilityMP(SampleMeasurement, StateMeasurement):
     Please refer to :func:`probs` for detailed documentation.
 
     Args:
-        obs (.Operator): The observable that is to be measured as part of the
-            measurement process. Not all measurement processes require observables (for
-            example ``Probability``); this argument is optional.
+        obs (Union[.Operator, .MeasurementValue]): The observable that is to be measured
+            as part of the measurement process. Not all measurement processes require observables
+            (for example ``Probability``); this argument is optional.
         wires (.Wires): The wires the measurement process applies to.
             This can only be specified if an observable was not provided.
         eigvals (array): A flat array representing the eigenvalues of the measurement.
diff --git a/pennylane/measurements/sample.py b/pennylane/measurements/sample.py
index 83bf6298a64..695b86b73b4 100644
--- a/pennylane/measurements/sample.py
+++ b/pennylane/measurements/sample.py
@@ -16,16 +16,19 @@
 """
 import functools
 import warnings
-from typing import Sequence, Tuple, Optional
+from typing import Sequence, Tuple, Optional, Union
 
 import pennylane as qml
 from pennylane.operation import Operator
 from pennylane.wires import Wires
 
 from .measurements import MeasurementShapeError, Sample, SampleMeasurement
+from .mid_measure import MeasurementValue
 
 
-def sample(op: Optional[Operator] = None, wires=None) -> "SampleMP":
+def sample(
+    op: Optional[Union[Operator, Sequence[MeasurementValue]]] = None, wires=None
+) -> "SampleMP":
     r"""Sample from the supplied observable, with the number of shots
     determined from the ``dev.shots`` attribute of the corresponding device,
     returning raw samples. If no observable is provided then basis state samples are returned
@@ -35,7 +38,8 @@ def sample(op: Optional[Operator] = None, wires=None) -> "SampleMP":
     specified on the device.
 
     Args:
-        op (Observable or None): a quantum observable object
+        op (Observable or MeasurementValue): a quantum observable object. To get samples
+            for mid-circuit measurements, ``op`` should be a``MeasurementValue``.
         wires (Sequence[int] or int or None): the wires we wish to sample from; ONLY set wires if
             op is ``None``
 
@@ -101,6 +105,9 @@ def circuit(x):
            [0, 0]])
 
     """
+    if isinstance(op, MeasurementValue):
+        return SampleMP(obs=op)
+
     if op is not None and not op.is_hermitian:  # None type is also allowed for op
         warnings.warn(f"{op.name} might not be hermitian.")
 
@@ -122,9 +129,9 @@ class SampleMP(SampleMeasurement):
     Please refer to :func:`sample` for detailed documentation.
 
     Args:
-        obs (.Operator): The observable that is to be measured as part of the
-            measurement process. Not all measurement processes require observables (for
-            example ``Probability``); this argument is optional.
+        obs (Union[.Operator, .MeasurementValue]): The observable that is to be measured
+            as part of the measurement process. Not all measurement processes require observables
+            (for example ``Probability``); this argument is optional.
         wires (.Wires): The wires the measurement process applies to.
             This can only be specified if an observable was not provided.
         eigvals (array): A flat array representing the eigenvalues of the measurement.
@@ -183,10 +190,26 @@ def process_samples(
         wire_order: Wires,
         shot_range: Tuple[int] = None,
         bin_size: int = None,
-    ):
+        decimal: bool = False,
+    ):  # pylint: disable=too-many-arguments, arguments-differ
+        """Process the given samples.
+        Args:
+            samples (Sequence[complex]): computational basis samples generated for all wires
+            wire_order (Wires): wires determining the subspace that ``samples`` acts on
+            shot_range (tuple[int]): 2-tuple of integers specifying the range of samples
+                to use. If not specified, all samples are used.
+            bin_size (int): Divides the shot range into bins of size ``bin_size``, and
+                returns the measurement statistic separately over each bin. If not
+                provided, the entire shot range is treated as a single bin.
+            decimal (bool): Whether to return samples as base-10 integers or boolean lists.
+        """
         wire_map = dict(zip(wire_order, range(len(wire_order))))
         mapped_wires = [wire_map[w] for w in self.wires]
-        name = self.obs.name if self.obs is not None else None
+        name = (
+            self.obs.name
+            if self.obs is not None and not isinstance(self.obs, MeasurementValue)
+            else None
+        )
         # Select the samples from samples that correspond to ``shot_range`` if provided
         if shot_range is not None:
             # Indexing corresponds to: (potential broadcasting, shots, wires). Note that the last
@@ -200,7 +223,7 @@ def process_samples(
 
         num_wires = samples.shape[-1]  # wires is the last dimension
 
-        if self.obs is None:
+        if (self.obs is None or isinstance(self.obs, MeasurementValue)) and not decimal:
             # if no observable was provided then return the raw samples
             return samples if bin_size is None else samples.T.reshape(num_wires, bin_size, -1)
 
@@ -215,12 +238,15 @@ def process_samples(
             powers_of_two = 2 ** qml.math.arange(num_wires)[::-1]
             indices = samples @ powers_of_two
             indices = qml.math.array(indices)  # Add np.array here for Jax support.
-            try:
-                samples = self.obs.eigvals()[indices]
-            except qml.operation.EigvalsUndefinedError as e:
-                # if observable has no info on eigenvalues, we cannot return this measurement
-                raise qml.operation.EigvalsUndefinedError(
-                    f"Cannot compute samples of {self.obs.name}."
-                ) from e
+            if decimal:
+                samples = indices
+            else:
+                try:
+                    samples = self.obs.eigvals()[indices]
+                except qml.operation.EigvalsUndefinedError as e:
+                    # if observable has no info on eigenvalues, we cannot return this measurement
+                    raise qml.operation.EigvalsUndefinedError(
+                        f"Cannot compute samples of {self.obs.name}."
+                    ) from e
 
         return samples if bin_size is None else samples.reshape((bin_size, -1))
diff --git a/pennylane/measurements/var.py b/pennylane/measurements/var.py
index 212f17cb6d1..bce2647cd67 100644
--- a/pennylane/measurements/var.py
+++ b/pennylane/measurements/var.py
@@ -16,7 +16,7 @@
 This module contains the qml.var measurement.
 """
 import warnings
-from typing import Sequence, Tuple
+from typing import Sequence, Tuple, Union
 
 import pennylane as qml
 from pennylane.operation import Operator
@@ -24,13 +24,16 @@
 from pennylane.wires import Wires
 
 from .measurements import SampleMeasurement, StateMeasurement, Variance
+from .mid_measure import MeasurementValue
 
 
-def var(op: Operator) -> "VarianceMP":
+def var(op: Union[Operator, MeasurementValue]) -> "VarianceMP":
     r"""Variance of the supplied observable.
 
     Args:
-        op (Operator): a quantum observable object
+        op (Union[Operator, MeasurementValue]): a quantum observable object.
+            To get variances for mid-circuit measurements, ``op`` should be a
+            ``MeasurementValue``.
 
     Returns:
         VarianceMP: Measurement process instance
@@ -53,6 +56,9 @@ def circuit(x):
     >>> circuit(0.5)
     0.7701511529340698
     """
+    if isinstance(op, MeasurementValue):
+        return VarianceMP(obs=op)
+
     if not op.is_hermitian:
         warnings.warn(f"{op.name} might not be hermitian.")
     return VarianceMP(obs=op)
@@ -64,9 +70,9 @@ class VarianceMP(SampleMeasurement, StateMeasurement):
     Please refer to :func:`var` for detailed documentation.
 
     Args:
-        obs (.Operator): The observable that is to be measured as part of the
-            measurement process. Not all measurement processes require observables (for
-            example ``Probability``); this argument is optional.
+        obs (Union[.Operator, .MeasurementValue]): The observable that is to be measured
+            as part of the measurement process. Not all measurement processes require observables
+            (for example ``Probability``); this argument is optional.
         wires (.Wires): The wires the measurement process applies to.
             This can only be specified if an observable was not provided.
         eigvals (array): A flat array representing the eigenvalues of the measurement.
@@ -107,9 +113,17 @@ def process_samples(
             return probs[idx] - probs[idx] ** 2
 
         # estimate the variance
+        # Get samples as decimal integer values if computing ev for a MeasurementValue,
+        # otherwise the returned samples would be boolean lists
+        decimal = isinstance(self.obs, MeasurementValue)
         samples = qml.sample(op=self.obs).process_samples(
-            samples=samples, wire_order=wire_order, shot_range=shot_range, bin_size=bin_size
+            samples=samples,
+            wire_order=wire_order,
+            shot_range=shot_range,
+            bin_size=bin_size,
+            decimal=decimal,
         )
+
         # With broadcasting, we want to take the variance over axis 1, which is the -1st/-2nd with/
         # without bin_size. Without broadcasting, axis 0 is the -1st/-2nd with/without bin_size
         axis = -1 if bin_size is None else -2
@@ -125,7 +139,11 @@ def process_state(self, state: Sequence[complex], wire_order: Wires):
             probs = qml.probs(wires=self.wires).process_state(state=state, wire_order=wire_order)
             return probs[idx] - probs[idx] ** 2
 
-        eigvals = qml.math.asarray(self.obs.eigvals(), dtype=float)
+        eigvals = (
+            qml.math.asarray(self.obs.eigvals(), dtype="float64")
+            if not isinstance(self.obs, MeasurementValue)
+            else qml.math.asarray(qml.math.arange(0, 2 ** len(self.wires)), dtype="float64")
+        )
 
         # we use ``wires`` instead of ``op`` because the observable was
         # already applied to the state
diff --git a/pennylane/tape/qscript.py b/pennylane/tape/qscript.py
index 27b60aece85..88629067a5c 100644
--- a/pennylane/tape/qscript.py
+++ b/pennylane/tape/qscript.py
@@ -28,6 +28,7 @@
     ClassicalShadowMP,
     CountsMP,
     MeasurementProcess,
+    MeasurementValue,
     ProbabilityMP,
     SampleMP,
     ShadowExpvalMP,
@@ -308,7 +309,7 @@ def observables(self) -> List[Union[MeasurementProcess, Observable]]:
         obs = []
 
         for m in self.measurements:
-            if m.obs is not None:
+            if m.obs is not None and not isinstance(m.obs, MeasurementValue):
                 m.obs.return_type = m.return_type
                 obs.append(m.obs)
             else:
@@ -448,7 +449,7 @@ def _update_par_info(self):
 
         n_ops = len(self.operations)
         for idx, m in enumerate(self.measurements):
-            if m.obs is not None:
+            if m.obs is not None and not isinstance(m.obs, MeasurementValue):
                 self._par_info.extend(
                     {"op": m.obs, "op_idx": idx + n_ops, "p_idx": i}
                     for i, d in enumerate(m.obs.data)
@@ -474,7 +475,12 @@ def _update_observables(self):
             _obs_sharing_wires_id (list[int]): Indices of the measurements that contain
                 the observables in _obs_sharing_wires
         """
-        obs_wires = [wire for m in self.measurements for wire in m.wires if m.obs is not None]
+        obs_wires = [
+            wire
+            for m in self.measurements
+            for wire in m.wires
+            if m.obs is not None and not isinstance(m.obs, MeasurementValue)
+        ]
         self._obs_sharing_wires = []
         self._obs_sharing_wires_id = []
 
@@ -673,7 +679,7 @@ def get_parameters(
         if operations_only:
             return params
         for m in self.measurements:
-            if m.obs is not None:
+            if m.obs is not None and not isinstance(m.obs, MeasurementValue):
                 params.extend(m.obs.data)
         return params
 
diff --git a/pennylane/transforms/convert_to_numpy_parameters.py b/pennylane/transforms/convert_to_numpy_parameters.py
index d16454c1eea..75323789043 100644
--- a/pennylane/transforms/convert_to_numpy_parameters.py
+++ b/pennylane/transforms/convert_to_numpy_parameters.py
@@ -19,6 +19,7 @@
 
 import pennylane as qml
 from pennylane import math
+from pennylane.measurements import MeasurementValue
 from pennylane.tape import QuantumScript
 
 
@@ -34,7 +35,11 @@ def _convert_op_to_numpy_data(op: qml.operation.Operator) -> qml.operation.Opera
 def _convert_measurement_to_numpy_data(
     m: qml.measurements.MeasurementProcess,
 ) -> qml.measurements.MeasurementProcess:
-    if m.obs is None or math.get_interface(*m.obs.data) == "numpy":
+    if (
+        m.obs is None
+        or isinstance(m.obs, MeasurementValue)
+        or math.get_interface(*m.obs.data) == "numpy"
+    ):
         return m
     # Use measurement method to change parameters when it becomes available
     copied_m = copy.copy(m)

From 31bef6e0819fcd02252d22e586b20473f8b52aec Mon Sep 17 00:00:00 2001
From: Mudit Pandey <mudit.pandey@xanadu.ai>
Date: Tue, 29 Aug 2023 17:19:10 -0400
Subject: [PATCH 009/131] Added docs; updated defer_measurements

---
 pennylane/measurements/__init__.py         | 47 ++++++++++++++++++++++
 pennylane/transforms/defer_measurements.py | 24 +++++++----
 2 files changed, 63 insertions(+), 8 deletions(-)

diff --git a/pennylane/measurements/__init__.py b/pennylane/measurements/__init__.py
index bb11974ccf2..2d070b09840 100644
--- a/pennylane/measurements/__init__.py
+++ b/pennylane/measurements/__init__.py
@@ -73,6 +73,53 @@
 that parameter is continuous. When using the analytic method of differentiation the output of the
 measurement process must be a real scalar value for it to be differentiable.
 
+Working with mid-circuit measurements
+-------------------------------------
+Mid-circuit measurements can be made using :func:`qml.measure`. The measurement value is returned by ``qml.measure``
+and can be used as a condition for classical control. Moreover, multiple measurement values can be combined
+using arithmetic operators for more complex conditioning:
+
+.. code-block:: python
+
+    import pennylane as qml
+
+    dev = qml.device("default.qubit", wires=3)
+
+    @qml.qnode(dev)
+    def circ(x, y):
+        qml.RX(x, wires=0)
+        qml.RY(y, wires=1)
+
+        m0 = qml.measure(0)
+        m1 = qml.measure(1)
+        qml.cond(~m0 & m1 == 0, qml.PauliX)(wires=2)
+        return qml.expval(qml.PauliZ(wires=2))
+
+Wires can be reused as normal after making mid-circuit measurements. Moreover, a measured wire can also be
+reset to the :math:`|0 \rangle` state by setting the ``reset`` keyword argument of ``qml.measure`` to ``True``.
+
+Users can also collect statistics on mid-circuit measurements along with other terminal measurements. Currently,
+``qml.expval``, ``qml.probs``, ``qml.sample``, ``qml.counts``, and ``qml.var`` are supported. Users have the
+ability to collect statistics on single measurement values.
+
+.. code-block:: python
+
+import pennylane as qml
+
+dev = qml.device("default.qubit", wires=3)
+
+@qml.qnode(dev)
+def circ(x, y):
+    qml.RX(x, wires=0)
+    qml.RY(y, wires=1)
+    m0 = qml.measure(1)
+    return qml.expval(qml.PauliZ(0)), qml.sample(m0)
+
+QNodes can be executed as usual when collecting mid-circuit measurement statistics:
+
+>>> circ(1.0, 2.0, shots=5)
+(array(0.6), array([1, 1, 1, 0, 1]))
+
 Creating custom measurements
 ----------------------------
 A custom measurement process can be created by inheriting from any of the classes mentioned above.
diff --git a/pennylane/transforms/defer_measurements.py b/pennylane/transforms/defer_measurements.py
index b14e3ba4b81..ccb5469df5b 100644
--- a/pennylane/transforms/defer_measurements.py
+++ b/pennylane/transforms/defer_measurements.py
@@ -13,7 +13,7 @@
 # limitations under the License.
 """Code for the tape transform implementing the deferred measurement principle."""
 import pennylane as qml
-from pennylane.measurements import MidMeasureMP
+from pennylane.measurements import MidMeasureMP, MeasurementValue
 from pennylane.ops.op_math import ctrl
 from pennylane.queuing import apply
 from pennylane.tape import QuantumTape
@@ -119,6 +119,7 @@ def func(x, y):
     >>> func(*pars)
     tensor([0.76960924, 0.13204407, 0.08394415, 0.01440254], requires_grad=True)
     """
+    # pylint: disable=protected-access
 
     cv_types = (qml.operation.CVOperation, qml.operation.CVObservable)
     ops_cv = any(isinstance(op, cv_types) for op in tape.operations)
@@ -127,18 +128,18 @@ def func(x, y):
         raise ValueError("Continuous variable operations and observables are not supported.")
 
     # Find wires that are reused after measurement
-    measured_wires = set()
+    measured_wires = []
     reused_measurement_wires = set()
 
     for op in tape.operations:
         if isinstance(op, MidMeasureMP):
-            if op.wires[0] in measured_wires or op.reset is True:
+            if op.reset is True:
                 reused_measurement_wires.add(op.wires[0])
-            measured_wires.add(op.wires[0])
+            measured_wires.append(op.wires[0])
 
         else:
             reused_measurement_wires = reused_measurement_wires.union(
-                measured_wires.intersection(op.wires.toset())
+                set(measured_wires).intersection(op.wires.toset())
             )
 
     # Apply controlled operations to store measurement outcomes and replace
@@ -146,10 +147,12 @@ def func(x, y):
     control_wires = {}
     cur_wire = max(tape.wires) + 1 if reused_measurement_wires else None
 
-    for op in tape:
+    for op in tape.operations:
         if isinstance(op, MidMeasureMP):
-            # Only store measurement outcome in new wire if wire gets reused
-            if op.wires[0] in reused_measurement_wires:
+            _ = measured_wires.pop(0)
+
+            # Store measurement outcome in new wire if wire gets reused
+            if op.wires[0] in reused_measurement_wires or op.wires[0] in measured_wires:
                 control_wires[op.id] = cur_wire
 
                 qml.CNOT([op.wires[0], cur_wire])
@@ -165,6 +168,11 @@ def func(x, y):
         else:
             apply(op)
 
+    for mp in tape.measurements:
+        if isinstance(mp.obs, MeasurementValue):
+            mp._wires = [control_wires[mp.obs.measurements[0].id]]
+        apply(mp)
+
     return tape._qfunc_output  # pylint: disable=protected-access
 
 

From 855a1943814c252105088e8cc60cb3a38fb9ad8f Mon Sep 17 00:00:00 2001
From: Mudit Pandey <mudit.pandey@xanadu.ai>
Date: Wed, 30 Aug 2023 12:39:30 -0400
Subject: [PATCH 010/131] Added mp tests

---
 tests/measurements/test_counts.py       | 34 ++++++++++++
 tests/measurements/test_expval.py       | 27 +++++++++
 tests/measurements/test_measurements.py |  9 +++
 tests/measurements/test_probs.py        | 31 +++++++++++
 tests/measurements/test_sample.py       | 74 ++++++++++++++++++++++++-
 tests/measurements/test_var.py          | 29 +++++++++-
 6 files changed, 202 insertions(+), 2 deletions(-)

diff --git a/tests/measurements/test_counts.py b/tests/measurements/test_counts.py
index eeee4a24e5f..625f264b1ea 100644
--- a/tests/measurements/test_counts.py
+++ b/tests/measurements/test_counts.py
@@ -167,6 +167,20 @@ def test_counts_obs(self):
         assert result[1] == np.count_nonzero(samples[:, 0] == 0)
         assert result[-1] == np.count_nonzero(samples[:, 0] == 1)
 
+    def test_counts_shape_single_measurement_value(self):
+        """Test that the counts output is correct for single mid-circuit measurement
+        values."""
+        shots = 1000
+        samples = np.random.choice([0, 1], size=(shots, 2)).astype(np.bool8)
+        mv = qml.measure(0)
+
+        result = qml.counts(mv).process_samples(samples, wire_order=[0])
+
+        assert len(result) == 2
+        assert set(result.keys()) == {"0", "1"}
+        assert result["0"] == np.count_nonzero(samples[:, 0] == 0)
+        assert result["1"] == np.count_nonzero(samples[:, 0] == 1)
+
     def test_counts_all_outcomes_wires(self):
         """Test that the counts output is correct when all_outcomes is passed"""
         shots = 1000
@@ -207,6 +221,26 @@ def test_counts_all_outcomes_obs(self):
         assert result2[1] == shots
         assert result2[-1] == 0
 
+    def test_counts_all_outcomes_measurement_value(self):
+        """Test that the counts output is correct when all_outcomes is passed
+        for mid-circuit measurement values."""
+        shots = 1000
+        samples = np.zeros((shots, 2)).astype(np.bool8)
+        mv = qml.measure(0)
+
+        result1 = qml.counts(mv, all_outcomes=False).process_samples(samples, wire_order=[0])
+
+        assert len(result1) == 1
+        assert set(result1.keys()) == {"0"}
+        assert result1["0"] == shots
+
+        result2 = qml.counts(mv, all_outcomes=True).process_samples(samples, wire_order=[0])
+
+        assert len(result2) == 2
+        assert set(result2.keys()) == {"0", "1"}
+        assert result2["0"] == shots
+        assert result2["1"] == 0
+
 
 class TestCountsIntegration:
     # pylint:disable=too-many-public-methods,not-an-iterable
diff --git a/tests/measurements/test_expval.py b/tests/measurements/test_expval.py
index 3d0e5c009b2..ea7fb94ce66 100644
--- a/tests/measurements/test_expval.py
+++ b/tests/measurements/test_expval.py
@@ -37,6 +37,9 @@ def custom_measurement_process(device, spy):
         # no need to use op, because the observable has already been applied to ``self.dev._state``
         meas = qml.expval(op=obs)
         old_res = device.expval(obs, shot_range=shot_range, bin_size=bin_size)
+        if isinstance(obs, qml.measurements.MeasurementProcess):
+            obs = obs.obs
+        meas = qml.expval(op=obs)
         if device.shots is None:
             new_res = meas.process_state(state=state, wire_order=device.wires)
         else:
@@ -116,6 +119,30 @@ def circuit():
 
         custom_measurement_process(new_dev, spy)
 
+    @pytest.mark.parametrize("shots", [None, 10000, [10000, 10000]])
+    @pytest.mark.parametrize("phi", np.arange(0, 2 * np.pi, np.pi / 3))
+    def test_observable_is_measurement_value(
+        self, shots, phi, mocker, tol, tol_stochastic
+    ):  # pylint: disable=too-many-arguments
+        """Test that expectation values for mid-circuit measurement values
+        are correct for a single measurement value."""
+        dev = qml.device("default.qubit", wires=2, shots=shots)
+
+        @qml.qnode(dev)
+        def circuit(phi):
+            qml.RX(phi, 0)
+            m0 = qml.measure(0)
+            return qml.expval(m0)
+
+        new_dev = circuit.device
+        spy = mocker.spy(qml.QubitDevice, "expval")
+
+        res = circuit(phi)
+
+        atol = tol if shots is None else tol_stochastic
+        assert np.allclose(np.array(res), np.sin(phi / 2) ** 2, atol=atol, rtol=0)
+        custom_measurement_process(new_dev, spy)
+
     @pytest.mark.parametrize(
         "obs",
         [qml.PauliZ(0), qml.Hermitian(np.diag([1, 2]), 0), qml.Hermitian(np.diag([1.0, 2.0]), 0)],
diff --git a/tests/measurements/test_measurements.py b/tests/measurements/test_measurements.py
index 3b8160419a5..6b2ee8815b6 100644
--- a/tests/measurements/test_measurements.py
+++ b/tests/measurements/test_measurements.py
@@ -279,6 +279,15 @@ def test_eigvals_match_observable(self):
         obs.data = [np.diag([5, 6, 7, 8])]
         assert np.all(m.eigvals() == np.array([5, 6, 7, 8]))
 
+    def test_measurement_value_eigvals(self):
+        """Test that eigenvalues of the measurement process
+        are correct if the internal observable is a
+        MeasurementValue."""
+        m0 = qml.measure(0)
+
+        m = qml.expval(m0)
+        assert np.all(m.eigvals() == [0, 1])
+
     def test_error_obs_and_eigvals(self):
         """Test that providing both eigenvalues and an observable
         results in an error"""
diff --git a/tests/measurements/test_probs.py b/tests/measurements/test_probs.py
index a07fd3da1b9..0c171568712 100644
--- a/tests/measurements/test_probs.py
+++ b/tests/measurements/test_probs.py
@@ -280,6 +280,37 @@ def circuit():
 
         custom_measurement_process(dev, spy)
 
+    @pytest.mark.parametrize("shots", [None, 10000, [10000, 10000]])
+    @pytest.mark.parametrize("phi", np.arange(0, 2 * np.pi, np.pi / 3))
+    def test_observable_is_measurement_value(
+        self, shots, phi, mocker, tol, tol_stochastic
+    ):  # pylint: disable=too-many-arguments
+        """Test that expectation values for mid-circuit measurement values
+        are correct for a single measurement value."""
+        dev = qml.device("default.qubit", wires=2, shots=shots)
+
+        @qml.qnode(dev)
+        def circuit(phi):
+            qml.RX(phi, 0)
+            m0 = qml.measure(0)
+            return qml.probs(op=m0)
+
+        new_dev = circuit.device
+        spy = mocker.spy(qml.QubitDevice, "probability")
+
+        res = circuit(phi)
+
+        atol = tol if shots is None else tol_stochastic
+        expected = np.array([np.cos(phi / 2) ** 2, np.sin(phi / 2) ** 2])
+
+        if not isinstance(shots, list):
+            assert np.allclose(np.array(res), expected, atol=atol, rtol=0)
+        else:
+            for r in res:  # pylint: disable=not-an-iterable
+                assert np.allclose(r, expected, atol=atol, rtol=0)
+
+        custom_measurement_process(new_dev, spy)
+
     @pytest.mark.parametrize("shots", [None, 100])
     def test_batch_size(self, mocker, shots):
         """Test the probability is correct for a batched input."""
diff --git a/tests/measurements/test_sample.py b/tests/measurements/test_sample.py
index 18b9d3f5ad3..066910f8f80 100644
--- a/tests/measurements/test_sample.py
+++ b/tests/measurements/test_sample.py
@@ -44,7 +44,7 @@ def custom_measurement_process(device, spy):
         )
 
 
-class TestSample:
+class TestSample:  # pylint: disable=too-many-public-methods
     """Tests for the sample function"""
 
     @pytest.mark.parametrize("n_sample", (1, 10))
@@ -196,6 +196,78 @@ def circuit():
 
         custom_measurement_process(dev, spy)
 
+    @pytest.mark.parametrize("shots", [5, [5, 5]])
+    @pytest.mark.parametrize("phi", np.arange(0, 2 * np.pi, np.pi / 2))
+    def test_observable_is_measurement_value(self, shots, phi, mocker):
+        """Test that expectation values for mid-circuit measurement values
+        are correct for a single measurement value."""
+        dev = qml.device("default.qubit", wires=2, shots=shots)
+
+        @qml.qnode(dev)
+        def circuit(phi):
+            qml.RX(phi, 0)
+            m0 = qml.measure(0)
+            return qml.sample(m0)
+
+        new_dev = circuit.device
+        spy = mocker.spy(qml.QubitDevice, "sample")
+
+        res = circuit(phi)
+
+        if isinstance(shots, list):
+            assert len(res) == len(shots)
+            assert all(r.shape == (s,) for r, s in zip(res, shots))
+        else:
+            assert res.shape == (shots,)
+        custom_measurement_process(new_dev, spy)
+
+    @pytest.mark.parametrize("shots", [8, [4, 4]])
+    def test_decimal_true(self, shots):
+        """Test that if `decimal==True`, then samples are returned as base-10 integers."""
+        dev = qml.device("default.qubit", wires=3, shots=shots)
+
+        samples = np.array(
+            [
+                [0, 0, 0],
+                [0, 0, 1],
+                [0, 1, 0],
+                [0, 1, 1],
+                [1, 0, 0],
+                [1, 0, 1],
+                [1, 1, 0],
+                [1, 1, 1],
+            ]
+        )
+        samples_decimal = np.arange(0, 8, 1)
+
+        dev._samples = samples
+        meas = qml.sample(wires=dev.wires)
+        shot_range = None if isinstance(shots, int) else (0, sum(shots))
+        bin_size = None if isinstance(shots, int) else shots[0]
+
+        # Using `meas` as the observable because `QubitDevice.sample` expects a
+        # measurement process for that argument when no observable is provided
+        old_samples = dev.sample(meas, shot_range=shot_range, bin_size=bin_size, decimal=True)
+        new_samples = meas.process_samples(
+            samples, wire_order=dev.wires, shot_range=shot_range, bin_size=bin_size, decimal=True
+        )
+
+        assert np.allclose(old_samples, new_samples)
+
+        if isinstance(shots, int):
+            assert new_samples.shape == old_samples.shape == samples_decimal.shape
+            # Enough to check that new samples are correct as we already know old and
+            # new samples are the same
+            assert np.allclose(new_samples, samples_decimal)
+
+        else:
+            # Samples are reshaped to dimension (bin_size, -1) for decimal values
+            expected_shape = (bin_size, len(shots))
+            assert old_samples.shape == new_samples.shape == expected_shape
+            # Enough to check that new samples are correct as we already know old and
+            # new samples are the same
+            np.allclose(new_samples, samples_decimal.reshape(expected_shape))
+
     def test_providing_observable_and_wires(self):
         """Test that a ValueError is raised if both an observable is provided and wires are specified"""
         dev = qml.device("default.qubit", wires=2)
diff --git a/tests/measurements/test_var.py b/tests/measurements/test_var.py
index 191e6937cfa..99dc3d4df4d 100644
--- a/tests/measurements/test_var.py
+++ b/tests/measurements/test_var.py
@@ -33,8 +33,10 @@ def custom_measurement_process(device, spy):
             call_args.kwargs["shot_range"],
             call_args.kwargs["bin_size"],
         )
+        old_res = device.var(obs, shot_range=shot_range, bin_size=bin_size)
+        if isinstance(obs, qml.measurements.MeasurementProcess):
+            obs = obs.obs
         meas = qml.var(op=obs)
-        old_res = device.var(obs, shot_range, bin_size)
         if samples is not None:
             new_res = meas.process_samples(
                 samples=samples, wire_order=device.wires, shot_range=shot_range, bin_size=bin_size
@@ -107,6 +109,31 @@ def circuit():
 
         custom_measurement_process(dev, spy)
 
+    @pytest.mark.parametrize("shots", [None, 10000, [10000, 10000]])
+    @pytest.mark.parametrize("phi", np.arange(0, 2 * np.pi, np.pi / 3))
+    def test_observable_is_measurement_value(
+        self, shots, phi, mocker, tol, tol_stochastic
+    ):  # pylint: disable=too-many-arguments
+        """Test that variances for mid-circuit measurement values
+        are correct for a single measurement value."""
+        dev = qml.device("default.qubit", wires=2, shots=shots)
+
+        @qml.qnode(dev)
+        def circuit(phi):
+            qml.RX(phi, 0)
+            m0 = qml.measure(0)
+            return qml.var(m0)
+
+        new_dev = circuit.device
+        spy = mocker.spy(qml.QubitDevice, "var")
+
+        res = circuit(phi)
+
+        atol = tol if shots is None else tol_stochastic
+        expected = np.sin(phi / 2) ** 2 - np.sin(phi / 2) ** 4
+        assert np.allclose(np.array(res), expected, atol=atol, rtol=0)
+        custom_measurement_process(new_dev, spy)
+
     @pytest.mark.parametrize(
         "obs",
         [qml.PauliZ(0), qml.Hermitian(np.diag([1, 2]), 0), qml.Hermitian(np.diag([1.0, 2.0]), 0)],

From a003eb372dce530990d938b900db0767bb9c06d3 Mon Sep 17 00:00:00 2001
From: Mudit Pandey <mudit.pandey@xanadu.ai>
Date: Wed, 30 Aug 2023 12:43:27 -0400
Subject: [PATCH 011/131] Added qnode test

---
 tests/test_qnode.py | 22 ++++++++++++++++++----
 1 file changed, 18 insertions(+), 4 deletions(-)

diff --git a/tests/test_qnode.py b/tests/test_qnode.py
index eca820b3953..b7e0d208320 100644
--- a/tests/test_qnode.py
+++ b/tests/test_qnode.py
@@ -1073,16 +1073,28 @@ def circuit():
         assert len(circuit.tape.operations) == 2
         assert isinstance(circuit.tape.operations[1], qml.measurements.MidMeasureMP)
 
+    @pytest.mark.parametrize(
+        "dev", [qml.device("default.qubit", wires=3), qml.devices.experimental.DefaultQubit2()]
+    )
     @pytest.mark.parametrize("first_par", np.linspace(0.15, np.pi - 0.3, 3))
     @pytest.mark.parametrize("sec_par", np.linspace(0.15, np.pi - 0.3, 3))
     @pytest.mark.parametrize(
         "return_type", [qml.expval(qml.PauliZ(1)), qml.var(qml.PauliZ(1)), qml.probs(wires=[1])]
     )
-    def test_defer_meas_if_mcm_unsupported(self, first_par, sec_par, return_type, mocker):
+    @pytest.mark.parametrize(
+        "mv_return, mv_res",
+        [
+            (qml.expval, lambda x: np.sin(x / 2) ** 2),
+            (qml.var, lambda x: np.sin(x / 2) ** 2 - np.sin(x / 2) ** 4),
+            (qml.probs, lambda x: [np.cos(x / 2) ** 2, np.sin(x / 2) ** 2]),
+        ],
+    )
+    def test_defer_meas_if_mcm_unsupported(
+        self, dev, first_par, sec_par, return_type, mv_return, mv_res, mocker
+    ):  # pylint: disable=too-many-arguments
         """Tests that the transform using the deferred measurement principle is
         applied if the device doesn't support mid-circuit measurements
         natively."""
-        dev = qml.device("default.qubit", wires=3)
 
         @qml.qnode(dev)
         def cry_qnode(x, y):
@@ -1100,12 +1112,14 @@ def conditional_ry_qnode(x, y):
             qml.RY(x, wires=0)
             m_0 = qml.measure(0)
             qml.cond(m_0, qml.RY)(y, wires=1)
-            return qml.apply(return_type)
+            return qml.apply(return_type), mv_return(op=m_0)
 
         spy = mocker.spy(qml, "defer_measurements")
         r1 = cry_qnode(first_par, sec_par)
         r2 = conditional_ry_qnode(first_par, sec_par)
-        assert np.allclose(r1, r2)
+
+        assert np.allclose(r1, r2[0])
+        assert np.allclose(r2[1], mv_res(first_par))
         spy.assert_called_once()
 
     def test_drawing_has_deferred_measurements(self):

From 29f93d93c9ba7b6dfb331729dad1ce3cc065b6c6 Mon Sep 17 00:00:00 2001
From: Mudit Pandey <mudit.pandey@xanadu.ai>
Date: Wed, 30 Aug 2023 12:47:59 -0400
Subject: [PATCH 012/131] Added deferred_measurements test

---
 tests/transforms/test_defer_measurements.py | 85 ++++++++++++++++++++-
 1 file changed, 83 insertions(+), 2 deletions(-)

diff --git a/tests/transforms/test_defer_measurements.py b/tests/transforms/test_defer_measurements.py
index 717c4fbd6b0..311d3b77758 100644
--- a/tests/transforms/test_defer_measurements.py
+++ b/tests/transforms/test_defer_measurements.py
@@ -20,6 +20,7 @@
 
 import pennylane as qml
 import pennylane.numpy as np
+from pennylane.devices.experimental import DefaultQubit2
 
 
 class TestQNode:
@@ -149,6 +150,61 @@ def qnode2(phi, theta):
         assert len(deferred_tape2.wires) == 3
         assert len(deferred_tape2.operations) == 4
 
+    @pytest.mark.parametrize("shots", [None, 1000, [1000, 1000]])
+    def test_measurement_statistics_single_wire(self, shots):
+        """Test that users can collect measurement statistics on
+        a single mid-circuit measurement."""
+        dev = DefaultQubit2(seed=10)
+
+        @qml.qnode(dev)
+        def circ1(x):
+            qml.RX(x, 0)
+            m0 = qml.measure(0)
+            return qml.probs(op=m0)
+
+        dev = DefaultQubit2(seed=10)
+
+        @qml.qnode(dev)
+        def circ2(x):
+            qml.RX(x, 0)
+            return qml.probs(wires=[0])
+
+        param = 1.5
+        assert np.allclose(circ1(param, shots=shots), circ2(param, shots=shots))
+
+    @pytest.mark.parametrize("shots", [None, 1000, [1000, 1000]])
+    def test_terminal_measurements(self, shots):
+        """Test that mid-circuit measurement statistics and terminal measurements
+        can be made together."""
+        # Using DefaultQubit2 to allow non-commuting measurements
+        dev = DefaultQubit2(seed=10)
+
+        @qml.qnode(dev)
+        def circ1(x, y):
+            qml.RX(x, 0)
+            m0 = qml.measure(0)
+            qml.RY(y, 1)
+            return qml.expval(qml.PauliX(1)), qml.probs(op=m0)
+
+        dev = DefaultQubit2(seed=10)
+
+        @qml.qnode(dev)
+        def circ2(x, y):
+            qml.RX(x, 0)
+            qml.RY(y, 1)
+            return qml.expval(qml.PauliX(1)), qml.probs(wires=[0])
+
+        params = [1.5, 2.5]
+        if isinstance(shots, list):
+            for out1, out2 in zip(circ1(*params, shots=shots), circ2(*params, shots=shots)):
+                for o1, o2 in zip(out1, out2):
+                    assert np.allclose(o1, o2)
+        else:
+            assert all(
+                np.allclose(out1, out2)
+                for out1, out2 in zip(circ1(*params, shots=shots), circ2(*params, shots=shots))
+            )
+
     def test_measure_between_ops(self):
         """Test that a quantum function that contains one operation before and
         after a mid-circuit measurement yields the correct results and is
@@ -995,8 +1051,33 @@ def qnode2(parameters):
             assert np.allclose(op1.data, op2.data)
 
 
-class TestQubitReset:
-    """Tests for the qubit reset functionality of `qml.measure`."""
+class TestQubitReuseAndReset:
+    """Tests for the qubit reuse/reset functionality of `qml.measure`."""
+
+    def test_new_wire_for_multiple_measurements(self):
+        """Test that a new wire is added if there are multiple mid-circuit measurements
+        on the same wire."""
+        dev = qml.device("default.qubit", wires=3)
+
+        @qml.qnode(dev)
+        def circ(x, y):
+            qml.RX(x, 0)
+            qml.measure(0)
+            qml.RY(y, 1)
+            qml.measure(0)
+            return qml.expval(qml.PauliZ(1))
+
+        _ = circ(1.0, 2.0)
+
+        expected = [
+            qml.RX(1.0, 0),
+            qml.CNOT([0, 2]),
+            qml.RY(2.0, 1),
+        ]
+
+        assert len(circ.qtape.operations) == 3
+        for op, exp in zip(circ.qtape.operations, expected):
+            assert qml.equal(op, exp)
 
     def test_correct_cnot_for_reset(self):
         """Test that a CNOT is applied from the wire that stores the measurement

From e9f5200f77494417a2bb917a3fff08e0e7710a2d Mon Sep 17 00:00:00 2001
From: Mudit Pandey <mudit.pandey@xanadu.ai>
Date: Wed, 30 Aug 2023 12:51:37 -0400
Subject: [PATCH 013/131] Added changelog entry

---
 doc/releases/changelog-dev.md | 23 +++++++++++++++++++++++
 1 file changed, 23 insertions(+)

diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 2ffd4a3a3b9..268cd2ee5d4 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -4,6 +4,29 @@
 
 <h3>New features since last release</h3>
 
+* Measurement statistics can now be collected for mid-circuit measurements. Currently,
+  `qml.expval`, `qml.var`, `qml.probs`, `qml.sample`, and `qml.counts` are supported on
+  `default.qubit`, `default.mixed`, and the new `DefaultQubit2` devices.
+  [(#4544)](https://github.com/PennyLaneAI/pennylane/pull/4544)
+
+  ```python
+  dev = qml.device("default.qubit", wires=3)
+
+  @qml.qnode(dev)
+  def circ(x, y):
+      qml.RX(x, wires=0)
+      qml.RY(y, wires=1)
+      m0 = qml.measure(1)
+      return qml.expval(qml.PauliZ(0)), qml.sample(m0)
+  ```
+
+  QNodes can be executed as usual when collecting mid-circuit measurement statistics:
+
+  ```pycon
+  >>> circ(1.0, 2.0, shots=5)
+  (array(0.6), array([1, 1, 1, 0, 1]))
+  ```
+
 <h3>Improvements 🛠</h3>
 
 * `qml.sample()` in the new device API now returns a `np.int64` array instead of `np.bool8`.

From 096f90f7c6cb7faaff1b93f089a3e677e9a54c93 Mon Sep 17 00:00:00 2001
From: Mudit Pandey <mudit.pandey@xanadu.ai>
Date: Wed, 30 Aug 2023 12:52:16 -0400
Subject: [PATCH 014/131] forgot to add name to changelog

---
 doc/releases/changelog-dev.md | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 268cd2ee5d4..008801c8c0f 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -70,4 +70,5 @@
 This release contains contributions from (in alphabetical order):
 
 Romain Moyard,
-Matthew Silverman
+Mudit Pandey,
+Matthew Silverman,

From 6baf8eafd4bfbe758a987493fccd323dc4e600cf Mon Sep 17 00:00:00 2001
From: Mudit Pandey <mudit.pandey@xanadu.ai>
Date: Wed, 30 Aug 2023 12:55:43 -0400
Subject: [PATCH 015/131] Fixed docs

---
 pennylane/measurements/__init__.py | 16 ++++++++--------
 pennylane/measurements/sample.py   |  2 ++
 2 files changed, 10 insertions(+), 8 deletions(-)

diff --git a/pennylane/measurements/__init__.py b/pennylane/measurements/__init__.py
index 2d070b09840..476d9427663 100644
--- a/pennylane/measurements/__init__.py
+++ b/pennylane/measurements/__init__.py
@@ -104,16 +104,16 @@ def circ(x, y):
 
 .. code-block:: python
 
-import pennylane as qml
+    import pennylane as qml
 
-dev = qml.device("default.qubit", wires=3)
+    dev = qml.device("default.qubit", wires=3)
 
-@qml.qnode(dev)
-def circ(x, y):
-    qml.RX(x, wires=0)
-    qml.RY(y, wires=1)
-    m0 = qml.measure(1)
-    return qml.expval(qml.PauliZ(0)), qml.sample(m0)
+    @qml.qnode(dev)
+    def circ(x, y):
+        qml.RX(x, wires=0)
+        qml.RY(y, wires=1)
+        m0 = qml.measure(1)
+        return qml.expval(qml.PauliZ(0)), qml.sample(m0)
 
 QNodes can be executed as usual when collecting mid-circuit measurement statistics:
 
diff --git a/pennylane/measurements/sample.py b/pennylane/measurements/sample.py
index 695b86b73b4..d7baaf93e50 100644
--- a/pennylane/measurements/sample.py
+++ b/pennylane/measurements/sample.py
@@ -193,7 +193,9 @@ def process_samples(
         decimal: bool = False,
     ):  # pylint: disable=too-many-arguments, arguments-differ
         """Process the given samples.
+
         Args:
+
             samples (Sequence[complex]): computational basis samples generated for all wires
             wire_order (Wires): wires determining the subspace that ``samples`` acts on
             shot_range (tuple[int]): 2-tuple of integers specifying the range of samples

From 362ac1c65f35f7cd473502d41fe3ff7ad7099a9c Mon Sep 17 00:00:00 2001
From: Mudit Pandey <mudit.pandey@xanadu.ai>
Date: Wed, 30 Aug 2023 14:56:01 -0400
Subject: [PATCH 016/131] Fixed test validation function

---
 tests/measurements/test_expval.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/tests/measurements/test_expval.py b/tests/measurements/test_expval.py
index ea7fb94ce66..c5a42690dfc 100644
--- a/tests/measurements/test_expval.py
+++ b/tests/measurements/test_expval.py
@@ -35,11 +35,10 @@ def custom_measurement_process(device, spy):
             call_args.kwargs["bin_size"],
         )
         # no need to use op, because the observable has already been applied to ``self.dev._state``
-        meas = qml.expval(op=obs)
-        old_res = device.expval(obs, shot_range=shot_range, bin_size=bin_size)
         if isinstance(obs, qml.measurements.MeasurementProcess):
             obs = obs.obs
         meas = qml.expval(op=obs)
+        old_res = device.expval(obs, shot_range=shot_range, bin_size=bin_size)
         if device.shots is None:
             new_res = meas.process_state(state=state, wire_order=device.wires)
         else:

From 83fdd6604ab82f154cd52e417057e2239794a8f8 Mon Sep 17 00:00:00 2001
From: Mudit Pandey <mudit.pandey@xanadu.ai>
Date: Wed, 30 Aug 2023 15:04:27 -0400
Subject: [PATCH 017/131] Fixed defer_measurementes

---
 pennylane/transforms/defer_measurements.py  | 8 +++++++-
 tests/transforms/test_defer_measurements.py | 4 ++--
 2 files changed, 9 insertions(+), 3 deletions(-)

diff --git a/pennylane/transforms/defer_measurements.py b/pennylane/transforms/defer_measurements.py
index ccb5469df5b..3fe92266c40 100644
--- a/pennylane/transforms/defer_measurements.py
+++ b/pennylane/transforms/defer_measurements.py
@@ -130,11 +130,15 @@ def func(x, y):
     # Find wires that are reused after measurement
     measured_wires = []
     reused_measurement_wires = set()
+    repeated_measurement_wire = False
 
     for op in tape.operations:
         if isinstance(op, MidMeasureMP):
             if op.reset is True:
                 reused_measurement_wires.add(op.wires[0])
+
+            if op.wires[0] in measured_wires:
+                repeated_measurement_wire = True
             measured_wires.append(op.wires[0])
 
         else:
@@ -145,7 +149,9 @@ def func(x, y):
     # Apply controlled operations to store measurement outcomes and replace
     # classically controlled operations
     control_wires = {}
-    cur_wire = max(tape.wires) + 1 if reused_measurement_wires else None
+    cur_wire = (
+        max(tape.wires) + 1 if reused_measurement_wires or repeated_measurement_wire else None
+    )
 
     for op in tape.operations:
         if isinstance(op, MidMeasureMP):
diff --git a/tests/transforms/test_defer_measurements.py b/tests/transforms/test_defer_measurements.py
index 311d3b77758..efd4d3f17aa 100644
--- a/tests/transforms/test_defer_measurements.py
+++ b/tests/transforms/test_defer_measurements.py
@@ -472,7 +472,7 @@ def test_quantum_teleportation(self, rads):
         assert ctrl_op2.wires == qml.wires.Wires([0, 2])
 
         # Check the measurement
-        assert tape.measurements[0] == terminal_measurement
+        assert qml.equal(tape.measurements[0], terminal_measurement)
 
     @pytest.mark.parametrize("r", np.linspace(0.1, 2 * np.pi - 0.1, 4))
     @pytest.mark.parametrize("device", ["default.qubit", "default.mixed", "lightning.qubit"])
@@ -538,7 +538,7 @@ def test_hermitian_queued(self):
         assert qml.equal(first_ctrl_op.base, qml.RY(rads, 4))
 
         assert len(tape.measurements) == 1
-        assert tape.measurements[0] == measurement
+        assert qml.equal(tape.measurements[0], measurement)
 
     def test_hamiltonian_queued(self):
         """Test that the defer_measurements transform works with

From 08d6dc218cb1cf9221c23ec22838e1a9550e7d4b Mon Sep 17 00:00:00 2001
From: Mudit Pandey <mudit.pandey@xanadu.ai>
Date: Wed, 30 Aug 2023 17:49:47 -0400
Subject: [PATCH 018/131] Refactoring

---
 pennylane/_qubit_device.py                 | 14 +++++++-------
 pennylane/devices/default_qubit.py         |  7 +++++--
 pennylane/measurements/measurements.py     |  8 ++------
 pennylane/measurements/mid_measure.py      |  6 ++++++
 pennylane/measurements/sample.py           |  2 +-
 pennylane/transforms/defer_measurements.py |  2 +-
 tests/measurements/test_expval.py          |  2 --
 tests/measurements/test_var.py             |  4 +---
 8 files changed, 23 insertions(+), 22 deletions(-)

diff --git a/pennylane/_qubit_device.py b/pennylane/_qubit_device.py
index bb9c187014b..8928506e422 100644
--- a/pennylane/_qubit_device.py
+++ b/pennylane/_qubit_device.py
@@ -613,7 +613,7 @@ def statistics(
 
         for m in measurements:
             # TODO: Remove this when all overriden measurements support the `MeasurementProcess` class
-            if m.obs is not None and not isinstance(m.obs, MeasurementValue):
+            if m.obs is not None:
                 obs = m.obs
                 obs.return_type = m.return_type
             else:
@@ -1319,7 +1319,7 @@ def expval(self, observable, shot_range=None, bin_size=None):
             try:
                 eigvals = self._asarray(
                     observable.eigvals()
-                    if not isinstance(observable, MeasurementProcess)
+                    if not isinstance(observable, MeasurementValue)
                     else np.arange(2 ** len(observable.wires)),
                     dtype=self.R_DTYPE,
                 )
@@ -1335,7 +1335,7 @@ def expval(self, observable, shot_range=None, bin_size=None):
         # estimate the ev
         # Get samples as decimal integer values if computing ev for a MeasurementValue,
         # otherwise the returned samples would be boolean lists
-        decimal = isinstance(observable, MeasurementProcess)
+        decimal = isinstance(observable, MeasurementValue)
         samples = self.sample(observable, shot_range=shot_range, bin_size=bin_size, decimal=decimal)
         # With broadcasting, we want to take the mean over axis 1, which is the -1st/-2nd with/
         # without bin_size. Without broadcasting, axis 0 is the -1st/-2nd with/without bin_size
@@ -1357,7 +1357,7 @@ def var(self, observable, shot_range=None, bin_size=None):
             try:
                 eigvals = self._asarray(
                     observable.eigvals()
-                    if not isinstance(observable, MeasurementProcess)
+                    if not isinstance(observable, MeasurementValue)
                     else np.arange(2 ** len(observable.wires)),
                     dtype=self.R_DTYPE,
                 )
@@ -1374,7 +1374,7 @@ def var(self, observable, shot_range=None, bin_size=None):
         # estimate the variance
         # Get samples as decimal integer values if computing ev for a MeasurementValue,
         # otherwise the returned samples would be boolean lists
-        decimal = isinstance(observable, MeasurementProcess)
+        decimal = isinstance(observable, MeasurementValue)
         samples = self.sample(observable, shot_range=shot_range, bin_size=bin_size, decimal=decimal)
 
         # With broadcasting, we want to take the variance over axis 1, which is the -1st/-2nd with/
@@ -1491,7 +1491,7 @@ def sample(
 
         # translate to wire labels used by device
         device_wires = self.map_wires(observable.wires)
-        name = observable.name
+        name = observable.name if not isinstance(observable, MeasurementValue) else None
         # Select the samples from self._samples that correspond to ``shot_range`` if provided
         if shot_range is None:
             sub_samples = self._samples
@@ -1501,7 +1501,7 @@ def sample(
             # Ellipsis (...) otherwise would take up broadcasting and shots axes.
             sub_samples = self._samples[..., slice(*shot_range), :]
 
-        no_observable_provided = isinstance(observable, MeasurementProcess)
+        no_observable_provided = isinstance(observable, (MeasurementProcess, MeasurementValue))
 
         if isinstance(name, str) and name in {"PauliX", "PauliY", "PauliZ", "Hadamard"}:
             # Process samples for observables with eigenvalues {1, -1}
diff --git a/pennylane/devices/default_qubit.py b/pennylane/devices/default_qubit.py
index c71eb5c3263..2c22411f21d 100644
--- a/pennylane/devices/default_qubit.py
+++ b/pennylane/devices/default_qubit.py
@@ -33,7 +33,7 @@
 from pennylane.ops import Sum
 from pennylane.ops.qubit.attributes import diagonal_in_z_basis
 from pennylane.pulse import ParametrizedEvolution
-from pennylane.measurements import ExpectationMP
+from pennylane.measurements import ExpectationMP, MeasurementValue
 from pennylane.typing import TensorLike
 from pennylane.wires import WireError
 
@@ -583,7 +583,10 @@ def expval(self, observable, shot_range=None, bin_size=None):
         # intercept other Hamiltonians
         # TODO: Ideally, this logic should not live in the Device, but be moved
         # to a component that can be re-used by devices as needed.
-        if observable.name not in ("Hamiltonian", "SparseHamiltonian"):
+        if isinstance(observable, MeasurementValue) or observable.name not in (
+            "Hamiltonian",
+            "SparseHamiltonian",
+        ):
             return super().expval(observable, shot_range=shot_range, bin_size=bin_size)
 
         assert self.shots is None, f"{observable.name} must be used with shots=None"
diff --git a/pennylane/measurements/measurements.py b/pennylane/measurements/measurements.py
index ddc11dfa7b1..54b746ffc39 100644
--- a/pennylane/measurements/measurements.py
+++ b/pennylane/measurements/measurements.py
@@ -146,11 +146,7 @@ def __init__(
 
         # _wires = None indicates broadcasting across all available wires.
         # It translates to the public property wires = Wires([])
-        self._wires = (
-            self.obs.measurements[0].wires
-            if self.obs is not None and self.obs.__class__.__name__ == "MeasurementValue"
-            else wires
-        )
+        self._wires = wires
         self._eigvals = None
 
         if eigvals is not None:
@@ -317,7 +313,7 @@ def wires(self):
 
         This is the union of all the Wires objects of the measurement.
         """
-        if self.obs is not None and self.obs.__class__.__name__ != "MeasurementValue":
+        if self.obs is not None:
             return self.obs.wires
 
         return (
diff --git a/pennylane/measurements/mid_measure.py b/pennylane/measurements/mid_measure.py
index 9ae9720f87c..8f260ed97e1 100644
--- a/pennylane/measurements/mid_measure.py
+++ b/pennylane/measurements/mid_measure.py
@@ -168,6 +168,7 @@ class MeasurementValue(Generic[T]):
     def __init__(self, measurements, processing_fn):
         self.measurements = measurements
         self.processing_fn = processing_fn
+        self._wires = Wires([m.wires[0] for m in self.measurements])
 
     def _items(self):
         """A generator representing all the possible outcomes of the MeasurementValue."""
@@ -175,6 +176,11 @@ def _items(self):
             branch = tuple(int(b) for b in np.binary_repr(i, width=len(self.measurements)))
             yield branch, self.processing_fn(*branch)
 
+    @property
+    def wires(self):
+        """Returns a list of wires corresponding to the mid-circuit measurements."""
+        return self._wires
+
     @property
     def branches(self):
         """A dictionary representing all possible outcomes of the MeasurementValue."""
diff --git a/pennylane/measurements/sample.py b/pennylane/measurements/sample.py
index d7baaf93e50..8487a1e244b 100644
--- a/pennylane/measurements/sample.py
+++ b/pennylane/measurements/sample.py
@@ -149,7 +149,7 @@ def return_type(self):
     def numeric_type(self):
         # Note: we only assume an integer numeric type if the observable is a
         # built-in observable with integer eigenvalues or a tensor product thereof
-        if self.obs is None:
+        if self.obs is None or isinstance(self.obs, MeasurementValue):
             # Computational basis samples
             return int
         int_eigval_obs = {qml.PauliX, qml.PauliY, qml.PauliZ, qml.Hadamard, qml.Identity}
diff --git a/pennylane/transforms/defer_measurements.py b/pennylane/transforms/defer_measurements.py
index 3fe92266c40..066bfd88c85 100644
--- a/pennylane/transforms/defer_measurements.py
+++ b/pennylane/transforms/defer_measurements.py
@@ -176,7 +176,7 @@ def func(x, y):
 
     for mp in tape.measurements:
         if isinstance(mp.obs, MeasurementValue):
-            mp._wires = [control_wires[mp.obs.measurements[0].id]]
+            mp.obs._wires = Wires([control_wires[mp.obs.measurements[0].id]])
         apply(mp)
 
     return tape._qfunc_output  # pylint: disable=protected-access
diff --git a/tests/measurements/test_expval.py b/tests/measurements/test_expval.py
index c5a42690dfc..25a890c4d4c 100644
--- a/tests/measurements/test_expval.py
+++ b/tests/measurements/test_expval.py
@@ -35,8 +35,6 @@ def custom_measurement_process(device, spy):
             call_args.kwargs["bin_size"],
         )
         # no need to use op, because the observable has already been applied to ``self.dev._state``
-        if isinstance(obs, qml.measurements.MeasurementProcess):
-            obs = obs.obs
         meas = qml.expval(op=obs)
         old_res = device.expval(obs, shot_range=shot_range, bin_size=bin_size)
         if device.shots is None:
diff --git a/tests/measurements/test_var.py b/tests/measurements/test_var.py
index 99dc3d4df4d..f89f9a31477 100644
--- a/tests/measurements/test_var.py
+++ b/tests/measurements/test_var.py
@@ -33,10 +33,8 @@ def custom_measurement_process(device, spy):
             call_args.kwargs["shot_range"],
             call_args.kwargs["bin_size"],
         )
-        old_res = device.var(obs, shot_range=shot_range, bin_size=bin_size)
-        if isinstance(obs, qml.measurements.MeasurementProcess):
-            obs = obs.obs
         meas = qml.var(op=obs)
+        old_res = device.var(obs, shot_range=shot_range, bin_size=bin_size)
         if samples is not None:
             new_res = meas.process_samples(
                 samples=samples, wire_order=device.wires, shot_range=shot_range, bin_size=bin_size

From e46e7200068e411cda7b508d6968c0a042b61fa4 Mon Sep 17 00:00:00 2001
From: Mudit Pandey <mudit.pandey@xanadu.ai>
Date: Wed, 30 Aug 2023 18:07:45 -0400
Subject: [PATCH 019/131] Reverted QNode changes

---
 pennylane/qnode.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pennylane/qnode.py b/pennylane/qnode.py
index 580b8f2e8e8..321956881c7 100644
--- a/pennylane/qnode.py
+++ b/pennylane/qnode.py
@@ -872,7 +872,7 @@ def construct(self, args, kwargs):  # pylint: disable=too-many-branches
         terminal_measurements = [
             m for m in self.tape.measurements if not isinstance(m, MidMeasureMP)
         ]
-        if any(ret is not m for ret, m in zip(measurement_processes, terminal_measurements)):
+        if any(ret != m for ret, m in zip(measurement_processes, terminal_measurements)):
             raise qml.QuantumFunctionError(
                 "All measurements must be returned in the order they are measured."
             )

From 23d8fd32a9371c614b14757981d48789dfaba41b Mon Sep 17 00:00:00 2001
From: Mudit Pandey <mudit.pandey@xanadu.ai>
Date: Wed, 30 Aug 2023 18:27:54 -0400
Subject: [PATCH 020/131] Added example to changelog

---
 doc/releases/changelog-dev.md | 47 ++++++++++++++++++++++++++++++++++-
 1 file changed, 46 insertions(+), 1 deletion(-)

diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index fbc75022b9b..cdf82142294 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -14,6 +14,50 @@
   the same hash.
   [(#4536)](https://github.com/PennyLaneAI/pennylane/pull/4536)
 
+  In the following scenario, the second and third code blocks show the previous and current behaviour
+  of operator and measurement process equality, determined by the `__eq__` dunder method:
+
+  ```python
+  op1 = qml.PauliX(0)
+  op2 = qml.PauliX(0)
+  op3 = op1
+  ```
+  Old behaviour:
+  ```pycon
+  >>> op1 == op2
+  False
+  >>> op1 == op3
+  True
+  ```
+  New behaviour:
+  ```pycon
+  >>> op1 == op2
+  True
+  >>> op1 == op3
+  True
+  ```
+
+  The `__hash__` dunder method defines the hash of an object, which is important for data
+  structures like sets and dictionaries to work correctly. The default hash of an object
+  is determined by the objects memory address. However, the hash for different objects should
+  be the same if they have the same properties. Consider the scenario below. The second and third
+  code blocks show the previous and current behaviour.
+
+  ```python
+  op1 = qml.PauliX(0)
+  op2 = qml.PauliX(0)
+  ```
+  Old behaviour:
+  ```pycon
+  >>> print({op1, op2})
+  {PauliX(wires=[0]), PauliX(wires=[0])}
+  ```
+  New behaviour:
+  ```pycon
+  >>> print({op1, op2})
+  {PauliX(wires=[0])}
+  ```
+
 * The old return type and associated functions ``qml.enable_return`` and ``qml.disable_return`` are removed.
  [(#4503)](https://github.com/PennyLaneAI/pennylane/pull/4503)
 
@@ -38,4 +82,5 @@
 
 This release contains contributions from (in alphabetical order):
 
-Romain Moyard
\ No newline at end of file
+Romain Moyard,
+Mudit Pandey
\ No newline at end of file

From dc418ab4572d67706a9f85f247f1dd498728be2e Mon Sep 17 00:00:00 2001
From: Mudit Pandey <mudit.pandey@xanadu.ai>
Date: Wed, 30 Aug 2023 18:30:26 -0400
Subject: [PATCH 021/131] Fixed example to changelog

---
 doc/releases/changelog-dev.md | 9 ++++-----
 1 file changed, 4 insertions(+), 5 deletions(-)

diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index cdf82142294..5f80ef79870 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -37,11 +37,10 @@
   True
   ```
 
-  The `__hash__` dunder method defines the hash of an object, which is important for data
-  structures like sets and dictionaries to work correctly. The default hash of an object
-  is determined by the objects memory address. However, the hash for different objects should
-  be the same if they have the same properties. Consider the scenario below. The second and third
-  code blocks show the previous and current behaviour.
+  The `__hash__` dunder method defines the hash of an object. The default hash of an object
+  is determined by the objects memory address. However, the new hash is determined by the
+  properties and attributes of operators and measurement processes. Consider the scenario below.
+  The second and third code blocks show the previous and current behaviour.
 
   ```python
   op1 = qml.PauliX(0)

From 655748a3366fd11e9e8a204f9e8ad1a2441607ab Mon Sep 17 00:00:00 2001
From: Mudit Pandey <mudit.pandey@xanadu.ai>
Date: Thu, 31 Aug 2023 10:35:38 -0400
Subject: [PATCH 022/131] Updated Hamiltonian.compare

---
 pennylane/ops/qubit/hamiltonian.py | 13 ++++++-------
 1 file changed, 6 insertions(+), 7 deletions(-)

diff --git a/pennylane/ops/qubit/hamiltonian.py b/pennylane/ops/qubit/hamiltonian.py
index 1a2ff7ac6b2..8991b31497c 100644
--- a/pennylane/ops/qubit/hamiltonian.py
+++ b/pennylane/ops/qubit/hamiltonian.py
@@ -613,16 +613,15 @@ def compare(self, other):
         >>> ob1.compare(ob2)
         False
         """
+        # pylint: disable=protected-access
         if isinstance(other, Hamiltonian):
-            self.simplify()
-            other.simplify()
-            return self._obs_data() == other._obs_data()  # pylint: disable=protected-access
+            H1 = copy(self).simplify()
+            H2 = copy(other).simplify()
+            return H1._obs_data() == H2._obs_data()
 
         if isinstance(other, (Tensor, Observable)):
-            self.simplify()
-            return self._obs_data() == {
-                (1, frozenset(other._obs_data()))  # pylint: disable=protected-access
-            }
+            H1 = copy(self).simplify()
+            return H1._obs_data() == {(1, frozenset(other._obs_data()))}
 
         raise ValueError("Can only compare a Hamiltonian, and a Hamiltonian/Observable/Tensor.")
 

From 5e9249572e2becc6afd05c04b329b5de6cbd1674 Mon Sep 17 00:00:00 2001
From: Mudit Pandey <mudit.pandey@xanadu.ai>
Date: Thu, 31 Aug 2023 10:51:31 -0400
Subject: [PATCH 023/131] Updated doc

---
 pennylane/_qubit_device.py        | 9 +++++----
 pennylane/measurements/expval.py  | 2 +-
 pennylane/measurements/sample.py  | 3 ++-
 pennylane/measurements/var.py     | 2 +-
 tests/measurements/test_counts.py | 4 ++--
 5 files changed, 11 insertions(+), 9 deletions(-)

diff --git a/pennylane/_qubit_device.py b/pennylane/_qubit_device.py
index 8928506e422..3e81169b5f9 100644
--- a/pennylane/_qubit_device.py
+++ b/pennylane/_qubit_device.py
@@ -1334,7 +1334,7 @@ def expval(self, observable, shot_range=None, bin_size=None):
 
         # estimate the ev
         # Get samples as decimal integer values if computing ev for a MeasurementValue,
-        # otherwise the returned samples would be boolean lists
+        # otherwise the returned samples would be boolean integer lists
         decimal = isinstance(observable, MeasurementValue)
         samples = self.sample(observable, shot_range=shot_range, bin_size=bin_size, decimal=decimal)
         # With broadcasting, we want to take the mean over axis 1, which is the -1st/-2nd with/
@@ -1373,7 +1373,7 @@ def var(self, observable, shot_range=None, bin_size=None):
 
         # estimate the variance
         # Get samples as decimal integer values if computing ev for a MeasurementValue,
-        # otherwise the returned samples would be boolean lists
+        # otherwise the returned samples would be boolean integer lists
         decimal = isinstance(observable, MeasurementValue)
         samples = self.sample(observable, shot_range=shot_range, bin_size=bin_size, decimal=decimal)
 
@@ -1477,8 +1477,9 @@ def sample(
                 provided, the entire shot range is treated as a single bin.
             counts (bool): whether counts (``True``) or raw samples (``False``)
                 should be returned
-            decimal (bool): Whether to return samples as base-10 integers or boolean lists.
-                Setting ``decimal`` does not do anything if ``counts`` is ``True``.
+            decimal (bool): Whether to return samples as base-10 integers or boolean
+                integer lists. Setting ``decimal`` does not do anything if ``counts``
+                is ``True``.
 
         Raises:
             EigvalsUndefinedError: if no information is available about the
diff --git a/pennylane/measurements/expval.py b/pennylane/measurements/expval.py
index 5df3d142b24..81700b60230 100644
--- a/pennylane/measurements/expval.py
+++ b/pennylane/measurements/expval.py
@@ -112,7 +112,7 @@ def process_samples(
 
         # estimate the ev
         # Get samples as decimal integer values if computing ev for a MeasurementValue,
-        # otherwise the returned samples would be boolean lists
+        # otherwise the returned samples would be boolean integer lists
         decimal = isinstance(self.obs, MeasurementValue)
         samples = qml.sample(op=self.obs).process_samples(
             samples=samples,
diff --git a/pennylane/measurements/sample.py b/pennylane/measurements/sample.py
index 8487a1e244b..423ff0bb903 100644
--- a/pennylane/measurements/sample.py
+++ b/pennylane/measurements/sample.py
@@ -203,7 +203,8 @@ def process_samples(
             bin_size (int): Divides the shot range into bins of size ``bin_size``, and
                 returns the measurement statistic separately over each bin. If not
                 provided, the entire shot range is treated as a single bin.
-            decimal (bool): Whether to return samples as base-10 integers or boolean lists.
+            decimal (bool): Whether to return samples as base-10 integers or boolean
+                integer lists.
         """
         wire_map = dict(zip(wire_order, range(len(wire_order))))
         mapped_wires = [wire_map[w] for w in self.wires]
diff --git a/pennylane/measurements/var.py b/pennylane/measurements/var.py
index bce2647cd67..29c83d65dff 100644
--- a/pennylane/measurements/var.py
+++ b/pennylane/measurements/var.py
@@ -114,7 +114,7 @@ def process_samples(
 
         # estimate the variance
         # Get samples as decimal integer values if computing ev for a MeasurementValue,
-        # otherwise the returned samples would be boolean lists
+        # otherwise the returned samples would be boolean integer lists
         decimal = isinstance(self.obs, MeasurementValue)
         samples = qml.sample(op=self.obs).process_samples(
             samples=samples,
diff --git a/tests/measurements/test_counts.py b/tests/measurements/test_counts.py
index 9ce21835e1a..aa2cee5a973 100644
--- a/tests/measurements/test_counts.py
+++ b/tests/measurements/test_counts.py
@@ -171,7 +171,7 @@ def test_counts_shape_single_measurement_value(self):
         """Test that the counts output is correct for single mid-circuit measurement
         values."""
         shots = 1000
-        samples = np.random.choice([0, 1], size=(shots, 2)).astype(np.bool8)
+        samples = np.random.choice([0, 1], size=(shots, 2)).astype(np.int64)
         mv = qml.measure(0)
 
         result = qml.counts(mv).process_samples(samples, wire_order=[0])
@@ -225,7 +225,7 @@ def test_counts_all_outcomes_measurement_value(self):
         """Test that the counts output is correct when all_outcomes is passed
         for mid-circuit measurement values."""
         shots = 1000
-        samples = np.zeros((shots, 2)).astype(np.bool8)
+        samples = np.zeros((shots, 2)).astype(np.int64)
         mv = qml.measure(0)
 
         result1 = qml.counts(mv, all_outcomes=False).process_samples(samples, wire_order=[0])

From b1f389601dda68eddf49b7c97b82f9b282b417e4 Mon Sep 17 00:00:00 2001
From: Mudit Pandey <mudit.pandey@xanadu.ai>
Date: Thu, 31 Aug 2023 10:56:59 -0400
Subject: [PATCH 024/131] Changed QNode back

---
 pennylane/qnode.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pennylane/qnode.py b/pennylane/qnode.py
index 321956881c7..580b8f2e8e8 100644
--- a/pennylane/qnode.py
+++ b/pennylane/qnode.py
@@ -872,7 +872,7 @@ def construct(self, args, kwargs):  # pylint: disable=too-many-branches
         terminal_measurements = [
             m for m in self.tape.measurements if not isinstance(m, MidMeasureMP)
         ]
-        if any(ret != m for ret, m in zip(measurement_processes, terminal_measurements)):
+        if any(ret is not m for ret, m in zip(measurement_processes, terminal_measurements)):
             raise qml.QuantumFunctionError(
                 "All measurements must be returned in the order they are measured."
             )

From a4e80746d2fd745a51d944479bb69553a103ab69 Mon Sep 17 00:00:00 2001
From: Mudit Pandey <mudit.pandey@xanadu.ai>
Date: Thu, 31 Aug 2023 12:51:32 -0400
Subject: [PATCH 025/131] Fixed sample tests

---
 tests/measurements/test_sample.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tests/measurements/test_sample.py b/tests/measurements/test_sample.py
index 066910f8f80..426ab11f4cf 100644
--- a/tests/measurements/test_sample.py
+++ b/tests/measurements/test_sample.py
@@ -16,7 +16,7 @@
 import pytest
 
 import pennylane as qml
-from pennylane.measurements import MeasurementShapeError, Sample, Shots
+from pennylane.measurements import MeasurementShapeError, Sample, Shots, MeasurementValue
 from pennylane.operation import EigvalsUndefinedError, Operator
 
 # pylint: disable=protected-access, no-member
@@ -31,7 +31,7 @@ def custom_measurement_process(device, spy):
     for call_args in call_args_list:
         meas = call_args.args[1]
         shot_range, bin_size = (call_args.kwargs["shot_range"], call_args.kwargs["bin_size"])
-        if isinstance(meas, Operator):
+        if isinstance(meas, (Operator, MeasurementValue)):
             meas = qml.sample(op=meas)
         assert qml.math.allequal(
             device.sample(call_args.args[1], **call_args.kwargs),

From 802c7a0f837f859112adc693b33c27816be092a6 Mon Sep 17 00:00:00 2001
From: Mudit Pandey <mudit.pandey@xanadu.ai>
Date: Thu, 31 Aug 2023 17:18:41 -0400
Subject: [PATCH 026/131] Removed decimal

---
 pennylane/_qubit_device.py        | 28 ++++++------------
 pennylane/_qutrit_device.py       |  4 +--
 pennylane/measurements/expval.py  |  4 ---
 pennylane/measurements/sample.py  | 23 +++------------
 pennylane/measurements/var.py     |  4 ---
 tests/measurements/test_sample.py | 47 -------------------------------
 6 files changed, 13 insertions(+), 97 deletions(-)

diff --git a/pennylane/_qubit_device.py b/pennylane/_qubit_device.py
index 3e81169b5f9..fa9d86fd587 100644
--- a/pennylane/_qubit_device.py
+++ b/pennylane/_qubit_device.py
@@ -1333,10 +1333,7 @@ def expval(self, observable, shot_range=None, bin_size=None):
             return self._dot(prob, eigvals)
 
         # estimate the ev
-        # Get samples as decimal integer values if computing ev for a MeasurementValue,
-        # otherwise the returned samples would be boolean integer lists
-        decimal = isinstance(observable, MeasurementValue)
-        samples = self.sample(observable, shot_range=shot_range, bin_size=bin_size, decimal=decimal)
+        samples = self.sample(observable, shot_range=shot_range, bin_size=bin_size)
         # With broadcasting, we want to take the mean over axis 1, which is the -1st/-2nd with/
         # without bin_size. Without broadcasting, axis 0 is the -1st/-2nd with/without bin_size
         axis = -1 if bin_size is None else -2
@@ -1372,11 +1369,7 @@ def var(self, observable, shot_range=None, bin_size=None):
             return self._dot(prob, (eigvals**2)) - self._dot(prob, eigvals) ** 2
 
         # estimate the variance
-        # Get samples as decimal integer values if computing ev for a MeasurementValue,
-        # otherwise the returned samples would be boolean integer lists
-        decimal = isinstance(observable, MeasurementValue)
-        samples = self.sample(observable, shot_range=shot_range, bin_size=bin_size, decimal=decimal)
-
+        samples = self.sample(observable, shot_range=shot_range, bin_size=bin_size)
         # With broadcasting, we want to take the variance over axis 1, which is the -1st/-2nd with/
         # without bin_size. Without broadcasting, axis 0 is the -1st/-2nd with/without bin_size
         axis = -1 if bin_size is None else -2
@@ -1463,9 +1456,7 @@ def circuit(x):
 
         return outcome_dicts if batched else outcome_dicts[0]
 
-    def sample(
-        self, observable, shot_range=None, bin_size=None, counts=False, decimal=False
-    ):  # pylint: disable=too-many-arguments
+    def sample(self, observable, shot_range=None, bin_size=None, counts=False):
         """Return samples of an observable.
 
         Args:
@@ -1477,9 +1468,6 @@ def sample(
                 provided, the entire shot range is treated as a single bin.
             counts (bool): whether counts (``True``) or raw samples (``False``)
                 should be returned
-            decimal (bool): Whether to return samples as base-10 integers or boolean
-                integer lists. Setting ``decimal`` does not do anything if ``counts``
-                is ``True``.
 
         Raises:
             EigvalsUndefinedError: if no information is available about the
@@ -1502,13 +1490,13 @@ def sample(
             # Ellipsis (...) otherwise would take up broadcasting and shots axes.
             sub_samples = self._samples[..., slice(*shot_range), :]
 
-        no_observable_provided = isinstance(observable, (MeasurementProcess, MeasurementValue))
+        no_observable_provided = isinstance(observable, MeasurementProcess)
 
         if isinstance(name, str) and name in {"PauliX", "PauliY", "PauliZ", "Hadamard"}:
             # Process samples for observables with eigenvalues {1, -1}
             samples = 1 - 2 * sub_samples[..., device_wires[0]]
 
-        elif no_observable_provided and (not decimal or counts):
+        elif no_observable_provided:
             # if no observable was provided then return the raw samples
             if len(observable.wires) != 0:
                 # if wires are provided, then we only return samples from those wires
@@ -1523,8 +1511,8 @@ def sample(
             powers_of_two = 2 ** np.arange(samples.shape[-1])[::-1]
             indices = samples @ powers_of_two
             indices = np.array(indices)  # Add np.array here for Jax support.
-            if decimal:
-                samples = indices
+            if isinstance(observable, MeasurementValue):
+                samples = np.arange(0, 2 ** len(observable.wires), 1)[indices]
             else:
                 try:
                     samples = observable.eigvals()[indices]
@@ -1549,7 +1537,7 @@ def sample(
 
         return (
             samples.T.reshape((num_wires, bin_size, -1))
-            if no_observable_provided and not decimal
+            if no_observable_provided
             else samples.reshape((bin_size, -1))
         )
 
diff --git a/pennylane/_qutrit_device.py b/pennylane/_qutrit_device.py
index 7739878f752..d9513a33d06 100644
--- a/pennylane/_qutrit_device.py
+++ b/pennylane/_qutrit_device.py
@@ -357,9 +357,7 @@ def marginal_prob(self, prob, wires=None):
         prob = self._transpose(prob, np.argsort(np.argsort(device_wires)))
         return self._flatten(prob)
 
-    def sample(
-        self, observable, shot_range=None, bin_size=None, counts=False, decimal=False
-    ):  # pylint: disable=too-many-arguments, unused-argument
+    def sample(self, observable, shot_range=None, bin_size=None, counts=False):
         def _samples_to_counts(samples, no_observable_provided):
             """Group the obtained samples into a dictionary.
 
diff --git a/pennylane/measurements/expval.py b/pennylane/measurements/expval.py
index 81700b60230..015b8c46c2d 100644
--- a/pennylane/measurements/expval.py
+++ b/pennylane/measurements/expval.py
@@ -111,15 +111,11 @@ def process_samples(
             return probs[idx]
 
         # estimate the ev
-        # Get samples as decimal integer values if computing ev for a MeasurementValue,
-        # otherwise the returned samples would be boolean integer lists
-        decimal = isinstance(self.obs, MeasurementValue)
         samples = qml.sample(op=self.obs).process_samples(
             samples=samples,
             wire_order=wire_order,
             shot_range=shot_range,
             bin_size=bin_size,
-            decimal=decimal,
         )
 
         # With broadcasting, we want to take the mean over axis 1, which is the -1st/-2nd with/
diff --git a/pennylane/measurements/sample.py b/pennylane/measurements/sample.py
index 423ff0bb903..28b00626249 100644
--- a/pennylane/measurements/sample.py
+++ b/pennylane/measurements/sample.py
@@ -190,22 +190,7 @@ def process_samples(
         wire_order: Wires,
         shot_range: Tuple[int] = None,
         bin_size: int = None,
-        decimal: bool = False,
-    ):  # pylint: disable=too-many-arguments, arguments-differ
-        """Process the given samples.
-
-        Args:
-
-            samples (Sequence[complex]): computational basis samples generated for all wires
-            wire_order (Wires): wires determining the subspace that ``samples`` acts on
-            shot_range (tuple[int]): 2-tuple of integers specifying the range of samples
-                to use. If not specified, all samples are used.
-            bin_size (int): Divides the shot range into bins of size ``bin_size``, and
-                returns the measurement statistic separately over each bin. If not
-                provided, the entire shot range is treated as a single bin.
-            decimal (bool): Whether to return samples as base-10 integers or boolean
-                integer lists.
-        """
+    ):
         wire_map = dict(zip(wire_order, range(len(wire_order))))
         mapped_wires = [wire_map[w] for w in self.wires]
         name = (
@@ -226,7 +211,7 @@ def process_samples(
 
         num_wires = samples.shape[-1]  # wires is the last dimension
 
-        if (self.obs is None or isinstance(self.obs, MeasurementValue)) and not decimal:
+        if self.obs is None:
             # if no observable was provided then return the raw samples
             return samples if bin_size is None else samples.T.reshape(num_wires, bin_size, -1)
 
@@ -241,8 +226,8 @@ def process_samples(
             powers_of_two = 2 ** qml.math.arange(num_wires)[::-1]
             indices = samples @ powers_of_two
             indices = qml.math.array(indices)  # Add np.array here for Jax support.
-            if decimal:
-                samples = indices
+            if isinstance(self.obs, MeasurementValue):
+                samples = qml.math.arange(0, 2 ** len(self.obs.wires), 1)[indices]
             else:
                 try:
                     samples = self.obs.eigvals()[indices]
diff --git a/pennylane/measurements/var.py b/pennylane/measurements/var.py
index 29c83d65dff..818f241a91d 100644
--- a/pennylane/measurements/var.py
+++ b/pennylane/measurements/var.py
@@ -113,15 +113,11 @@ def process_samples(
             return probs[idx] - probs[idx] ** 2
 
         # estimate the variance
-        # Get samples as decimal integer values if computing ev for a MeasurementValue,
-        # otherwise the returned samples would be boolean integer lists
-        decimal = isinstance(self.obs, MeasurementValue)
         samples = qml.sample(op=self.obs).process_samples(
             samples=samples,
             wire_order=wire_order,
             shot_range=shot_range,
             bin_size=bin_size,
-            decimal=decimal,
         )
 
         # With broadcasting, we want to take the variance over axis 1, which is the -1st/-2nd with/
diff --git a/tests/measurements/test_sample.py b/tests/measurements/test_sample.py
index 426ab11f4cf..1a0a4bb1e98 100644
--- a/tests/measurements/test_sample.py
+++ b/tests/measurements/test_sample.py
@@ -221,53 +221,6 @@ def circuit(phi):
             assert res.shape == (shots,)
         custom_measurement_process(new_dev, spy)
 
-    @pytest.mark.parametrize("shots", [8, [4, 4]])
-    def test_decimal_true(self, shots):
-        """Test that if `decimal==True`, then samples are returned as base-10 integers."""
-        dev = qml.device("default.qubit", wires=3, shots=shots)
-
-        samples = np.array(
-            [
-                [0, 0, 0],
-                [0, 0, 1],
-                [0, 1, 0],
-                [0, 1, 1],
-                [1, 0, 0],
-                [1, 0, 1],
-                [1, 1, 0],
-                [1, 1, 1],
-            ]
-        )
-        samples_decimal = np.arange(0, 8, 1)
-
-        dev._samples = samples
-        meas = qml.sample(wires=dev.wires)
-        shot_range = None if isinstance(shots, int) else (0, sum(shots))
-        bin_size = None if isinstance(shots, int) else shots[0]
-
-        # Using `meas` as the observable because `QubitDevice.sample` expects a
-        # measurement process for that argument when no observable is provided
-        old_samples = dev.sample(meas, shot_range=shot_range, bin_size=bin_size, decimal=True)
-        new_samples = meas.process_samples(
-            samples, wire_order=dev.wires, shot_range=shot_range, bin_size=bin_size, decimal=True
-        )
-
-        assert np.allclose(old_samples, new_samples)
-
-        if isinstance(shots, int):
-            assert new_samples.shape == old_samples.shape == samples_decimal.shape
-            # Enough to check that new samples are correct as we already know old and
-            # new samples are the same
-            assert np.allclose(new_samples, samples_decimal)
-
-        else:
-            # Samples are reshaped to dimension (bin_size, -1) for decimal values
-            expected_shape = (bin_size, len(shots))
-            assert old_samples.shape == new_samples.shape == expected_shape
-            # Enough to check that new samples are correct as we already know old and
-            # new samples are the same
-            np.allclose(new_samples, samples_decimal.reshape(expected_shape))
-
     def test_providing_observable_and_wires(self):
         """Test that a ValueError is raised if both an observable is provided and wires are specified"""
         dev = qml.device("default.qubit", wires=2)

From 708baa51598e9bd13ee4176e4784ae92d5f75e4e Mon Sep 17 00:00:00 2001
From: Mudit Pandey <mudit.pandey@xanadu.ai>
Date: Fri, 1 Sep 2023 11:42:41 -0400
Subject: [PATCH 027/131] Fixed counts

---
 pennylane/measurements/counts.py | 11 ++++++++---
 1 file changed, 8 insertions(+), 3 deletions(-)

diff --git a/pennylane/measurements/counts.py b/pennylane/measurements/counts.py
index e317785fed7..83e17a43b4b 100644
--- a/pennylane/measurements/counts.py
+++ b/pennylane/measurements/counts.py
@@ -214,9 +214,14 @@ def process_samples(
         shot_range: Tuple[int] = None,
         bin_size: int = None,
     ):
-        samples = qml.sample(op=self.obs, wires=self._wires).process_samples(
-            samples, wire_order, shot_range, bin_size
-        )
+        if not isinstance(self.obs, MeasurementValue):
+            samples = qml.sample(op=self.obs, wires=self._wires).process_samples(
+                samples, wire_order, shot_range, bin_size
+            )
+        else:
+            samples = qml.sample(wires=self.obs.wires).process_samples(
+                samples, wire_order, shot_range, bin_size
+            )
 
         if bin_size is None:
             return self._samples_to_counts(samples)

From 6c2a232b70ebe0c49ecdb8fcd76a8598bf6691c9 Mon Sep 17 00:00:00 2001
From: Mudit Pandey <mudit.pandey@xanadu.ai>
Date: Tue, 5 Sep 2023 15:20:03 -0400
Subject: [PATCH 028/131] Update changelog

---
 doc/releases/changelog-dev.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 1458c3844bb..5446e6a78d8 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -6,7 +6,7 @@
 
 * Measurement statistics can now be collected for mid-circuit measurements. Currently,
   `qml.expval`, `qml.var`, `qml.probs`, `qml.sample`, and `qml.counts` are supported on
-  `default.qubit` and the new `DefaultQubit2` device.
+  `default.qubit`, `default.mixed`, and the new `DefaultQubit2` device.
   [(#4544)](https://github.com/PennyLaneAI/pennylane/pull/4544)
 
   ```python

From 5ebbc2ee27fe045122c92f083749dc8b9da3ac94 Mon Sep 17 00:00:00 2001
From: Matthew Silverman <matthews@xanadu.ai>
Date: Fri, 1 Sep 2023 12:15:07 -0400
Subject: [PATCH 029/131] separate DensityMatrixMP into its own measurement
 process (#4558)

**Context:**
Today, `qml.state()` returns a `StateMP` with no wires set, while
`qml.density_matrix()` returns a `StateMP` with the wires of the
subsystem being traced out. The new device implementation would benefit
greatly from being able to assign wires to a state measurement, and also
it just makes sense for these to be separate things.

**Description of the Change:**
Adds a new `DensityMatrixMP` that inherits from `StateMP` (not the
parent `StateMeasurement` so that `isinstance(mp, StateMP)` still
returns True for density matrix measurements).

I literally copy-pasted the `StateMP` class, then deleted any
density-matrix-related things from the `StateMP` class, and deleted and
state-related things from the `DensityMatrixMP` class.

**Benefits:**
- Density matrix is its own measurement process!
- We can assign wires to a `StateMP` without it thinking it represents a
density matrix.

**Possible Drawbacks:**
This can also be viewed as a benefit (see above, no code changes outside
module needed), but `isinstance(my_density_mp, StateMP)` returns true.
If you want to know that your measurement _is_ a state measurement, you
can do `type(mp) is StateMP`.

**UPDATE for reviewers**:
Responding to a comment, I've changed `StateMP` to not accept `wires` in
its init, so now `DensityMatrixMP` directly calls
`StateMeasurement.__init__(self, wires=wires, id=id)` to skip its parent
(StateMP) which would now fail if given wires. I call it directly
because this is the most pythonic way to call a grandparent's method
(the other option being `super(StateMP, self).__init__` but I think
`StateMeasurement` is a more reliable and responsible choice than
"whoever StateMP's parent is"). Plus, I'll be removing that change when
StateMP supports wires.
---
 doc/releases/changelog-dev.md      |  5 ++-
 pennylane/measurements/__init__.py |  2 +-
 pennylane/measurements/state.py    | 67 ++++++++++++++++++------------
 tests/measurements/test_state.py   | 31 ++++++++++----
 4 files changed, 70 insertions(+), 35 deletions(-)

diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 5446e6a78d8..291159a6df2 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -45,6 +45,10 @@
 * The new device API now has a `repr()`
   [(#4562)](https://github.com/PennyLaneAI/pennylane/pull/4562)
 
+* The density matrix aspects of `StateMP` have been split into their own measurement
+  process, `DensityMatrixMP`.
+  [(#4558)](https://github.com/PennyLaneAI/pennylane/pull/4558)
+
 <h3>Breaking changes 💔</h3>
 
 * The `__eq__` and `__hash__` methods of `Operator` and `MeasurementProcess` no longer rely on the
@@ -123,7 +127,6 @@
   been removed. Please use ``QuantumScript.bind_new_parameters`` instead.
   [(#4548)](https://github.com/PennyLaneAI/pennylane/pull/4548)
 
-  
 <h3>Deprecations 👋</h3>
 
 * The ``prep`` keyword argument in ``QuantumScript`` is deprecated and will be removed from `QuantumScript`.
diff --git a/pennylane/measurements/__init__.py b/pennylane/measurements/__init__.py
index 476d9427663..db26c173a17 100644
--- a/pennylane/measurements/__init__.py
+++ b/pennylane/measurements/__init__.py
@@ -266,6 +266,6 @@ def circuit(x):
 from .probs import ProbabilityMP, probs
 from .sample import SampleMP, sample
 from .shots import Shots, ShotCopies
-from .state import StateMP, density_matrix, state
+from .state import StateMP, DensityMatrixMP, density_matrix, state
 from .var import VarianceMP, var
 from .vn_entropy import VnEntropyMP, vn_entropy
diff --git a/pennylane/measurements/state.py b/pennylane/measurements/state.py
index 8ca9c2cba9a..b8d0c212593 100644
--- a/pennylane/measurements/state.py
+++ b/pennylane/measurements/state.py
@@ -81,7 +81,7 @@ def circuit():
     return StateMP()
 
 
-def density_matrix(wires) -> "StateMP":
+def density_matrix(wires) -> "DensityMatrixMP":
     r"""Quantum density matrix in the computational basis.
 
     This function accepts no observables and instead instructs the QNode to return its density
@@ -93,7 +93,7 @@ def density_matrix(wires) -> "StateMP":
         wires (Sequence[int] or int): the wires of the subsystem
 
     Returns:
-        StateMP: Measurement process instance
+        DensityMatrixMP: Measurement process instance
 
     **Example:**
 
@@ -122,23 +122,21 @@ def circuit():
         with a compatible device.
     """
     wires = Wires(wires)
-    return StateMP(wires=wires)
+    return DensityMatrixMP(wires=wires)
 
 
 class StateMP(StateMeasurement):
     """Measurement process that returns the quantum state in the computational basis.
 
-    Please refer to :func:`state` and :func:`density_matrix` for detailed documentation.
+    Please refer to :func:`state` for detailed documentation.
 
     Args:
-        wires (.Wires): The wires the measurement process applies to.
-            This can only be specified if an observable was not provided.
         id (str): custom label given to a measurement instance, can be useful for some applications
             where the instance has to be identified
     """
 
-    def __init__(self, wires: Optional[Wires] = None, id: Optional[str] = None):
-        super().__init__(wires=wires, id=id)
+    def __init__(self, *, id: Optional[str] = None):
+        super().__init__(wires=None, id=id)
 
     @property
     def return_type(self):
@@ -152,26 +150,43 @@ def shape(self, device, shots):
         num_shot_elements = (
             sum(s.copies for s in shots.shot_vector) if shots.has_partitioned_shots else 1
         )
-
-        if self.wires:
-            # qml.density_matrix()
-            dim = 2 ** len(self.wires)
-            return (
-                (dim, dim)
-                if num_shot_elements == 1
-                else tuple((dim, dim) for _ in range(num_shot_elements))
-            )
-
-        # qml.state()
         dim = 2 ** len(device.wires)
         return (dim,) if num_shot_elements == 1 else tuple((dim,) for _ in range(num_shot_elements))
 
-    # pylint: disable=redefined-outer-name
     def process_state(self, state: Sequence[complex], wire_order: Wires):
-        if self.wires:
-            # qml.density_matrix
-            wire_map = dict(zip(wire_order, range(len(wire_order))))
-            mapped_wires = [wire_map[w] for w in self.wires]
-            return qml.math.reduce_statevector(state, indices=mapped_wires, c_dtype=state.dtype)
-        # qml.state
+        # pylint:disable=redefined-outer-name
         return state
+
+
+class DensityMatrixMP(StateMP):
+    """Measurement process that returns the quantum state in the computational basis.
+
+    Please refer to :func:`density_matrix` for detailed documentation.
+
+    Args:
+        wires (.Wires): The wires the measurement process applies to.
+        id (str): custom label given to a measurement instance, can be useful for some applications
+            where the instance has to be identified
+    """
+
+    def __init__(self, wires: Wires, id: Optional[str] = None):
+        # pylint:disable=non-parent-init-called,super-init-not-called
+        StateMeasurement.__init__(self, wires=wires, id=id)
+
+    def shape(self, device, shots):
+        num_shot_elements = (
+            sum(s.copies for s in shots.shot_vector) if shots.has_partitioned_shots else 1
+        )
+
+        dim = 2 ** len(self.wires)
+        return (
+            (dim, dim)
+            if num_shot_elements == 1
+            else tuple((dim, dim) for _ in range(num_shot_elements))
+        )
+
+    def process_state(self, state: Sequence[complex], wire_order: Wires):
+        # pylint:disable=redefined-outer-name
+        wire_map = dict(zip(wire_order, range(len(wire_order))))
+        mapped_wires = [wire_map[w] for w in self.wires]
+        return qml.math.reduce_statevector(state, indices=mapped_wires, c_dtype=state.dtype)
diff --git a/tests/measurements/test_state.py b/tests/measurements/test_state.py
index 5198c77ecaa..eebfb128ad1 100644
--- a/tests/measurements/test_state.py
+++ b/tests/measurements/test_state.py
@@ -18,12 +18,18 @@
 import pennylane as qml
 from pennylane import numpy as pnp
 from pennylane.devices import DefaultQubit
-from pennylane.measurements import State, StateMP, Shots, density_matrix, expval, state
+from pennylane.measurements import (
+    State,
+    StateMP,
+    DensityMatrixMP,
+    Shots,
+    density_matrix,
+    expval,
+    state,
+)
 from pennylane.math.quantum import reduce_statevector, reduce_dm
 from pennylane.math.matrix_manipulation import _permute_dense_matrix
 
-# pylint: disable=no-member, comparison-with-callable, import-outside-toplevel
-
 
 class TestStateMP:
     """Tests for the State measurement process."""
@@ -39,13 +45,22 @@ class TestStateMP:
     def test_process_state_vector(self, vec):
         """Test the processing of a state vector."""
 
-        mp = StateMP(wires=None)
+        mp = StateMP()
         assert mp.return_type == State
         assert mp.numeric_type is complex
 
         processed = mp.process_state(vec, None)
         assert qml.math.allclose(processed, vec)
 
+    def test_state_does_not_accept_wires(self):
+        """Test that StateMP does not accept wires."""
+        with pytest.raises(TypeError, match="unexpected keyword argument 'wires'"):
+            StateMP(wires=[0])
+
+
+class TestDensityMatrixMP:
+    """Tests for the DensityMatrix measurement process"""
+
     @pytest.mark.parametrize(
         "vec, wires",
         [
@@ -58,7 +73,7 @@ def test_process_state_vector(self, vec):
     def test_process_state_matrix_from_vec(self, vec, wires):
         """Test the processing of a state vector into a matrix."""
 
-        mp = StateMP(wires=wires)
+        mp = DensityMatrixMP(wires=wires)
         assert mp.return_type == State
         assert mp.numeric_type is complex
 
@@ -71,7 +86,9 @@ def test_process_state_matrix_from_vec(self, vec, wires):
             exp = reduce_statevector(vec, wires)
         assert qml.math.allclose(processed, exp)
 
-    @pytest.mark.xfail(reason="StateMP.process_state no longer supports density matrix parameters")
+    @pytest.mark.xfail(
+        reason="DensityMatrixMP.process_state no longer supports density matrix parameters"
+    )
     @pytest.mark.parametrize(
         "mat, wires",
         [
@@ -86,7 +103,7 @@ def test_process_state_matrix_from_vec(self, vec, wires):
     def test_process_state_matrix_from_matrix(self, mat, wires):
         """Test the processing of a density matrix into a matrix."""
 
-        mp = StateMP(wires=wires)
+        mp = DensityMatrixMP(wires=wires)
         assert mp.return_type == State
         assert mp.numeric_type is complex
 

From f10bfe4ad405a6d03eae925e3138e6080d56820c Mon Sep 17 00:00:00 2001
From: Matthew Silverman <matthews@xanadu.ai>
Date: Tue, 5 Sep 2023 09:43:37 -0400
Subject: [PATCH 030/131] add default.qubit.legacy and rename some test devices
 to prepare for DQ2 (#4568)

**Context:**
DefaultQubit2 will replace DefaultQubit when users call
`qml.device("default.qubit")` soon. DefaultQubit will still be available
with the name `"default.qubit.legacy"`. My goal in this PR is to ease
that transition, because a lot of tests need changing. There are 2
reasons a test would be changed to use `default.qubit.legacy`.
1. It tests the old device
2. It fails with the new device (for better or for worse)

This PR should _only_ change tests of the first kind.

**Description of the Change:**
- Add `default.qubit.legacy` as an alternate entrypoint for
`qml.devices.DefaultQubit` (which I may now sometimes call DQL). Note
that `default.qubit` also points to DQL (for now)
- Change tests that need the old device to use `default.qubit.legacy` as
the name to help with the transition in the future.

**Benefits:**
Future PRs will be easier to review when completing the migration.

**Possible Drawbacks:**
N/A

**Related GitHub Issues:**
#4436, #4534
---
 setup.py                                      |   1 +
 tests/conftest.py                             |  14 ++-
 tests/devices/test_default_qubit_autograd.py  |   4 +-
 ..._qubit.py => test_default_qubit_legacy.py} | 117 +++++++++---------
 ...test_default_qubit_legacy_broadcasting.py} |  72 +++++------
 tests/devices/test_default_qubit_tf.py        |   4 +-
 tests/devices/test_default_qubit_torch.py     |   2 +-
 tests/gradients/core/test_adjoint_diff.py     |  10 +-
 tests/interfaces/test_autograd.py             |  70 +++++------
 tests/interfaces/test_autograd_qnode.py       |  70 +++++------
 .../test_autograd_qnode_shot_vector.py        |   6 +-
 tests/interfaces/test_jax.py                  |  50 ++++----
 tests/interfaces/test_jax_jit.py              |  52 ++++----
 tests/interfaces/test_jax_jit_qnode.py        |  22 ++--
 tests/interfaces/test_jax_qnode.py            |  24 ++--
 .../interfaces/test_jax_qnode_shot_vector.py  |  14 +--
 tests/interfaces/test_tensorflow.py           |  48 +++----
 ..._tensorflow_autograph_qnode_shot_vector.py |   6 +-
 tests/interfaces/test_tensorflow_qnode.py     |  56 ++++-----
 .../test_tensorflow_qnode_shot_vector.py      |   6 +-
 tests/interfaces/test_torch.py                |  62 +++++-----
 tests/interfaces/test_torch_qnode.py          |  52 ++++----
 .../test_transform_program_integration.py     |   4 +-
 tests/measurements/test_classical_shadow.py   |  38 +++---
 tests/measurements/test_counts.py             |  56 ++++-----
 tests/measurements/test_expval.py             |  14 +--
 tests/measurements/test_measurements.py       |  22 ++--
 tests/measurements/test_mutual_info.py        |  30 ++---
 tests/measurements/test_probs.py              |  44 +++----
 tests/measurements/test_purity_measurement.py |   6 +-
 tests/measurements/test_sample.py             |  36 +++---
 tests/measurements/test_state.py              |  64 +++++-----
 tests/measurements/test_var.py                |  14 +--
 tests/measurements/test_vn_entropy.py         |  22 ++--
 tests/qinfo/test_reduced_dm.py                |   2 +-
 tests/test_debugging.py                       |   6 +-
 tests/test_device.py                          |   6 +-
 tests/test_operation.py                       |   4 +-
 tests/test_qubit_device.py                    |  12 +-
 tests/transforms/test_batch_transform.py      |   8 +-
 tests/transforms/test_tape_expand.py          |  34 ++---
 41 files changed, 596 insertions(+), 588 deletions(-)
 rename tests/devices/{test_default_qubit.py => test_default_qubit_legacy.py} (95%)
 rename tests/devices/{test_default_qubit_broadcasting.py => test_default_qubit_legacy_broadcasting.py} (96%)

diff --git a/setup.py b/setup.py
index ac4ebedfcef..669a115e131 100644
--- a/setup.py
+++ b/setup.py
@@ -48,6 +48,7 @@
         # This requires a rename in the setup file of all devices, and is best done during another refactor
         "pennylane.plugins": [
             "default.qubit = pennylane.devices:DefaultQubit",
+            "default.qubit.legacy = pennylane.devices:DefaultQubit",
             "default.gaussian = pennylane.devices:DefaultGaussian",
             "default.qubit.tf = pennylane.devices.default_qubit_tf:DefaultQubitTF",
             "default.qubit.torch = pennylane.devices.default_qubit_torch:DefaultQubitTorch",
diff --git a/tests/conftest.py b/tests/conftest.py
index 2a8d0bc558b..26eb11463b5 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -70,22 +70,28 @@ def n_subsystems_fixture(request):
 
 @pytest.fixture(scope="session")
 def qubit_device(n_subsystems):
-    return qml.device("default.qubit", wires=n_subsystems)
+    return qml.device("default.qubit.legacy", wires=n_subsystems)
 
 
 @pytest.fixture(scope="function", params=[(np.float32, np.complex64), (np.float64, np.complex128)])
 def qubit_device_1_wire(request):
-    return qml.device("default.qubit", wires=1, r_dtype=request.param[0], c_dtype=request.param[1])
+    return qml.device(
+        "default.qubit.legacy", wires=1, r_dtype=request.param[0], c_dtype=request.param[1]
+    )
 
 
 @pytest.fixture(scope="function", params=[(np.float32, np.complex64), (np.float64, np.complex128)])
 def qubit_device_2_wires(request):
-    return qml.device("default.qubit", wires=2, r_dtype=request.param[0], c_dtype=request.param[1])
+    return qml.device(
+        "default.qubit.legacy", wires=2, r_dtype=request.param[0], c_dtype=request.param[1]
+    )
 
 
 @pytest.fixture(scope="function", params=[(np.float32, np.complex64), (np.float64, np.complex128)])
 def qubit_device_3_wires(request):
-    return qml.device("default.qubit", wires=3, r_dtype=request.param[0], c_dtype=request.param[1])
+    return qml.device(
+        "default.qubit.legacy", wires=3, r_dtype=request.param[0], c_dtype=request.param[1]
+    )
 
 
 # The following 3 fixtures are for default.qutrit devices to be used
diff --git a/tests/devices/test_default_qubit_autograd.py b/tests/devices/test_default_qubit_autograd.py
index 21b8d139a62..78094c23191 100644
--- a/tests/devices/test_default_qubit_autograd.py
+++ b/tests/devices/test_default_qubit_autograd.py
@@ -407,8 +407,8 @@ def circuit(x):
                 qml.CNOT(wires=[i, i + 1])
             return qml.expval(qml.PauliZ(0)), qml.var(qml.PauliZ(1))
 
-        dev1 = qml.device("default.qubit", wires=3)
-        dev2 = qml.device("default.qubit", wires=3)
+        dev1 = qml.device("default.qubit.legacy", wires=3)
+        dev2 = qml.device("default.qubit.legacy", wires=3)
 
         def cost(x):
             return qml.math.stack(circuit(x))
diff --git a/tests/devices/test_default_qubit.py b/tests/devices/test_default_qubit_legacy.py
similarity index 95%
rename from tests/devices/test_default_qubit.py
rename to tests/devices/test_default_qubit_legacy.py
index 33cdafc86c1..fe5742255cd 100644
--- a/tests/devices/test_default_qubit.py
+++ b/tests/devices/test_default_qubit_legacy.py
@@ -98,17 +98,17 @@ def test_analytic_deprecation():
         DeviceError,
         match=msg,
     ):
-        qml.device("default.qubit", wires=1, shots=1, analytic=True)
+        qml.device("default.qubit.legacy", wires=1, shots=1, analytic=True)
 
 
 def test_dtype_errors():
     """Test that if an incorrect dtype is provided to the device then an error is raised."""
     with pytest.raises(DeviceError, match="Real datatype must be a floating point type."):
-        qml.device("default.qubit", wires=1, r_dtype=np.complex128)
+        qml.device("default.qubit.legacy", wires=1, r_dtype=np.complex128)
     with pytest.raises(
         DeviceError, match="Complex datatype must be a complex floating point type."
     ):
-        qml.device("default.qubit", wires=1, c_dtype=np.float64)
+        qml.device("default.qubit.legacy", wires=1, c_dtype=np.float64)
 
 
 def test_custom_op_with_matrix():
@@ -117,7 +117,7 @@ def test_custom_op_with_matrix():
     class DummyOp(qml.operation.Operation):
         num_wires = 1
 
-        def compute_matrix(self):
+        def compute_matrix(self):  # pylint:disable=arguments-differ
             return np.eye(2)
 
     with qml.queuing.AnnotatedQueue() as q:
@@ -125,7 +125,7 @@ def compute_matrix(self):
         qml.state()
 
     tape = qml.tape.QuantumScript.from_queue(q)
-    dev = qml.device("default.qubit", wires=1)
+    dev = qml.device("default.qubit.legacy", wires=1)
     assert qml.math.allclose(dev.execute(tape), np.array([1, 0]))
 
 
@@ -792,7 +792,7 @@ def test_expval_two_wires_with_parameters(
     def test_expval_estimate(self):
         """Test that the expectation value is not analytically calculated"""
 
-        dev = qml.device("default.qubit", wires=1, shots=3)
+        dev = qml.device("default.qubit.legacy", wires=1, shots=3)
 
         @qml.qnode(dev, diff_method="parameter-shift")
         def circuit():
@@ -919,7 +919,7 @@ def test_var_two_wires_with_parameters(
     def test_var_estimate(self):
         """Test that the variance is not analytically calculated"""
 
-        dev = qml.device("default.qubit", wires=1, shots=3)
+        dev = qml.device("default.qubit.legacy", wires=1, shots=3)
 
         @qml.qnode(dev, diff_method="parameter-shift")
         def circuit():
@@ -943,7 +943,7 @@ def test_sample_dimensions(self):
         # Explicitly resetting is necessary as the internal
         # state is set to None in __init__ and only properly
         # initialized during reset
-        dev = qml.device("default.qubit", wires=2, shots=1000)
+        dev = qml.device("default.qubit.legacy", wires=2, shots=1000)
 
         dev.apply([qml.RX(1.5708, wires=[0]), qml.RX(1.5708, wires=[1])])
 
@@ -975,7 +975,7 @@ def test_sample_values(self, tol):
         # Explicitly resetting is necessary as the internal
         # state is set to None in __init__ and only properly
         # initialized during reset
-        dev = qml.device("default.qubit", wires=2, shots=1000)
+        dev = qml.device("default.qubit.legacy", wires=2, shots=1000)
 
         dev.apply([qml.RX(1.5708, wires=[0])])
         dev._wires_measured = {0}
@@ -995,7 +995,7 @@ class TestDefaultQubitIntegration:
     def test_defines_correct_capabilities(self):
         """Test that the device defines the right capabilities"""
 
-        dev = qml.device("default.qubit", wires=1)
+        dev = qml.device("default.qubit.legacy", wires=1)
         cap = dev.capabilities()
         capabilities = {
             "model": "qubit",
@@ -1033,7 +1033,7 @@ def circuit(x):
 
         res = circuit(p)
         assert np.isclose(res, expected, atol=tol, rtol=0)
-        assert res.dtype == r_dtype
+        assert res.dtype == r_dtype  # pylint:disable=no-member
 
     def test_qubit_identity(self, qubit_device_1_wire, tol):
         """Test that the default qubit plugin provides correct result for the Identity expectation"""
@@ -1052,7 +1052,7 @@ def test_nonzero_shots(self, tol):
         """Test that the default qubit plugin provides correct result for high shot number"""
 
         shots = 10**5
-        dev = qml.device("default.qubit", wires=1, shots=shots)
+        dev = qml.device("default.qubit.legacy", wires=1, shots=shots)
 
         p = 0.543
 
@@ -1195,7 +1195,7 @@ def test_multi_samples_return_correlated_results(self):
         the correct dimensions
         """
 
-        dev = qml.device("default.qubit", wires=2, shots=1000)
+        dev = qml.device("default.qubit.legacy", wires=2, shots=1000)
 
         @qml.qnode(dev, diff_method="parameter-shift")
         def circuit():
@@ -1215,7 +1215,7 @@ def test_multi_samples_return_correlated_results_more_wires_than_size_of_observa
         the correct dimensions
         """
 
-        dev = qml.device("default.qubit", wires=num_wires, shots=1000)
+        dev = qml.device("default.qubit.legacy", wires=num_wires, shots=1000)
 
         @qml.qnode(dev, diff_method="parameter-shift")
         def circuit():
@@ -1235,7 +1235,7 @@ class TestTensorExpval:
 
     def test_paulix_pauliy(self, theta, phi, varphi, tol):
         """Test that a tensor product involving PauliX and PauliY works correctly"""
-        dev = qml.device("default.qubit", wires=3)
+        dev = qml.device("default.qubit.legacy", wires=3)
         dev.reset()
 
         obs = qml.PauliX(0) @ qml.PauliY(2)
@@ -1259,7 +1259,7 @@ def test_paulix_pauliy(self, theta, phi, varphi, tol):
 
     def test_pauliz_identity(self, theta, phi, varphi, tol):
         """Test that a tensor product involving PauliZ and Identity works correctly"""
-        dev = qml.device("default.qubit", wires=3)
+        dev = qml.device("default.qubit.legacy", wires=3)
         dev.reset()
 
         obs = qml.PauliZ(0) @ qml.Identity(1) @ qml.PauliZ(2)
@@ -1283,7 +1283,7 @@ def test_pauliz_identity(self, theta, phi, varphi, tol):
 
     def test_pauliz_hadamard(self, theta, phi, varphi, tol):
         """Test that a tensor product involving PauliZ and PauliY and hadamard works correctly"""
-        dev = qml.device("default.qubit", wires=3)
+        dev = qml.device("default.qubit.legacy", wires=3)
         obs = qml.PauliZ(0) @ qml.Hadamard(1) @ qml.PauliY(2)
 
         dev.reset()
@@ -1306,7 +1306,7 @@ def test_pauliz_hadamard(self, theta, phi, varphi, tol):
 
     def test_hermitian(self, theta, phi, varphi, tol):
         """Test that a tensor product involving qml.Hermitian works correctly"""
-        dev = qml.device("default.qubit", wires=3)
+        dev = qml.device("default.qubit.legacy", wires=3)
         dev.reset()
 
         A = np.array(
@@ -1344,7 +1344,7 @@ def test_hermitian(self, theta, phi, varphi, tol):
 
     def test_hermitian_hermitian(self, theta, phi, varphi, tol):
         """Test that a tensor product involving two Hermitian matrices works correctly"""
-        dev = qml.device("default.qubit", wires=3)
+        dev = qml.device("default.qubit.legacy", wires=3)
 
         A1 = np.array([[1, 2], [2, 4]])
 
@@ -1393,7 +1393,7 @@ def test_hermitian_hermitian(self, theta, phi, varphi, tol):
 
     def test_hermitian_identity_expectation(self, theta, phi, varphi, tol):
         """Test that a tensor product involving an Hermitian matrix and the identity works correctly"""
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
 
         A = np.array(
             [[1.02789352, 1.61296440 - 0.3498192j], [1.61296440 + 0.3498192j, 1.23920938 + 0j]]
@@ -1417,7 +1417,7 @@ def test_hermitian_identity_expectation(self, theta, phi, varphi, tol):
 
     def test_hermitian_two_wires_identity_expectation(self, theta, phi, varphi, tol):
         """Test that a tensor product involving an Hermitian matrix for two wires and the identity works correctly"""
-        dev = qml.device("default.qubit", wires=3)
+        dev = qml.device("default.qubit.legacy", wires=3)
 
         A = np.array(
             [[1.02789352, 1.61296440 - 0.3498192j], [1.61296440 + 0.3498192j, 1.23920938 + 0j]]
@@ -1446,7 +1446,7 @@ class TestTensorVar:
 
     def test_paulix_pauliy(self, theta, phi, varphi, tol):
         """Test that a tensor product involving PauliX and PauliY works correctly"""
-        dev = qml.device("default.qubit", wires=3)
+        dev = qml.device("default.qubit.legacy", wires=3)
 
         obs = qml.PauliX(0) @ qml.PauliY(2)
 
@@ -1476,7 +1476,7 @@ def test_paulix_pauliy(self, theta, phi, varphi, tol):
 
     def test_pauliz_hadamard(self, theta, phi, varphi, tol):
         """Test that a tensor product involving PauliZ and PauliY and hadamard works correctly"""
-        dev = qml.device("default.qubit", wires=3)
+        dev = qml.device("default.qubit.legacy", wires=3)
         obs = qml.PauliZ(0) @ qml.Hadamard(1) @ qml.PauliY(2)
 
         dev.reset()
@@ -1504,7 +1504,7 @@ def test_pauliz_hadamard(self, theta, phi, varphi, tol):
 
     def test_hermitian(self, theta, phi, varphi, tol):
         """Test that a tensor product involving qml.Hermitian works correctly"""
-        dev = qml.device("default.qubit", wires=3)
+        dev = qml.device("default.qubit.legacy", wires=3)
 
         A = np.array(
             [
@@ -1570,7 +1570,7 @@ class TestTensorSample:
 
     def test_paulix_pauliy(self, theta, phi, varphi, tol_stochastic):
         """Test that a tensor product involving PauliX and PauliY works correctly"""
-        dev = qml.device("default.qubit", wires=3, shots=int(1e6))
+        dev = qml.device("default.qubit.legacy", wires=3, shots=int(1e6))
 
         obs = qml.PauliX(0) @ qml.PauliY(2)
 
@@ -1612,7 +1612,7 @@ def test_paulix_pauliy(self, theta, phi, varphi, tol_stochastic):
 
     def test_pauliz_hadamard(self, theta, phi, varphi, tol_stochastic):
         """Test that a tensor product involving PauliZ and PauliY and hadamard works correctly"""
-        dev = qml.device("default.qubit", wires=3, shots=int(1e6))
+        dev = qml.device("default.qubit.legacy", wires=3, shots=int(1e6))
         obs = qml.PauliZ(0) @ qml.Hadamard(1) @ qml.PauliY(2)
         dev.apply(
             [
@@ -1650,7 +1650,7 @@ def test_pauliz_hadamard(self, theta, phi, varphi, tol_stochastic):
 
     def test_hermitian(self, theta, phi, varphi, tol_stochastic):
         """Test that a tensor product involving qml.Hermitian works correctly"""
-        dev = qml.device("default.qubit", wires=3, shots=int(1e6))
+        dev = qml.device("default.qubit.legacy", wires=3, shots=int(1e6))
 
         A = 0.1 * np.array(
             [
@@ -1765,7 +1765,7 @@ def test_state_dtype_after_op(self, r_dtype, c_dtype, op):
         examples.
         """
 
-        dev = qml.device("default.qubit", wires=4, r_dtype=r_dtype, c_dtype=c_dtype)
+        dev = qml.device("default.qubit.legacy", wires=4, r_dtype=r_dtype, c_dtype=c_dtype)
 
         n_wires = op.num_wires
         n_params = op.num_params
@@ -1781,7 +1781,7 @@ def circuit():
             return qml.state()
 
         res = circuit()
-        assert res.dtype == c_dtype
+        assert res.dtype == c_dtype  # pylint:disable=no-member
 
     @pytest.mark.parametrize(
         "measurement",
@@ -1796,7 +1796,7 @@ def test_measurement_real_dtype(self, r_dtype, c_dtype, measurement):
         """Test that the default qubit plugin provides correct result for a simple circuit"""
         p = 0.543
 
-        dev = qml.device("default.qubit", wires=3, r_dtype=r_dtype, c_dtype=c_dtype)
+        dev = qml.device("default.qubit.legacy", wires=3, r_dtype=r_dtype, c_dtype=c_dtype)
 
         @qml.qnode(dev, diff_method="parameter-shift")
         def circuit(x):
@@ -1814,7 +1814,7 @@ def test_measurement_complex_dtype(self, r_dtype, c_dtype, measurement):
         """Test that the default qubit plugin provides correct result for a simple circuit"""
         p = 0.543
 
-        dev = qml.device("default.qubit", wires=3, r_dtype=r_dtype, c_dtype=c_dtype)
+        dev = qml.device("default.qubit.legacy", wires=3, r_dtype=r_dtype, c_dtype=c_dtype)
 
         @qml.qnode(dev, diff_method="parameter-shift")
         def circuit(x):
@@ -1836,8 +1836,8 @@ def mock_analytic_counter(self, wires=None):
     @pytest.mark.parametrize("x", [[0.2, 0.5], [0.4, 0.9], [0.8, 0.3]])
     def test_probability(self, x, tol):
         """Test that the probability function works for finite and infinite shots"""
-        dev = qml.device("default.qubit", wires=2, shots=1000)
-        dev_analytic = qml.device("default.qubit", wires=2, shots=None)
+        dev = qml.device("default.qubit.legacy", wires=2, shots=1000)
+        dev_analytic = qml.device("default.qubit.legacy", wires=2, shots=None)
 
         def circuit(x):
             qml.RX(x[0], wires=0)
@@ -1857,7 +1857,7 @@ def test_call_generate_samples(self, monkeypatch):
         """Test analytic_probability call when generating samples"""
         self.analytic_counter = False
 
-        dev = qml.device("default.qubit", wires=2, shots=1000)
+        dev = qml.device("default.qubit.legacy", wires=2, shots=1000)
         monkeypatch.setattr(dev, "analytic_probability", self.mock_analytic_counter)
 
         # generate samples through `generate_samples` (using 'analytic_probability')
@@ -1868,7 +1868,7 @@ def test_call_generate_samples(self, monkeypatch):
 
     def test_stateless_analytic_return(self):
         """Test that analytic_probability returns None if device is stateless"""
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
         dev._state = None
 
         assert dev.analytic_probability() is None
@@ -1879,7 +1879,7 @@ class TestWiresIntegration:
 
     def make_circuit_probs(self, wires):
         """Factory for a qnode returning probabilities using arbitrary wire labels."""
-        dev = qml.device("default.qubit", wires=wires)
+        dev = qml.device("default.qubit.legacy", wires=wires)
         n_wires = len(wires)
 
         @qml.qnode(dev, diff_method="parameter-shift")
@@ -1912,7 +1912,7 @@ def test_wires_probs(self, wires1, wires2, tol):
 
     def test_wires_not_found_exception(self):
         """Tests that an exception is raised when wires not present on the device are adressed."""
-        dev = qml.device("default.qubit", wires=["a", "b"])
+        dev = qml.device("default.qubit.legacy", wires=["a", "b"])
 
         with qml.queuing.AnnotatedQueue() as q:
             qml.RX(0.5, wires="c")
@@ -1931,7 +1931,7 @@ def test_wires_not_found_exception(self):
     @pytest.mark.parametrize("dev_wires, wires_to_map", wires_to_try)
     def test_map_wires_caches(self, dev_wires, wires_to_map):
         """Test that multiple calls to map_wires will use caching."""
-        dev = qml.device("default.qubit", wires=dev_wires)
+        dev = qml.device("default.qubit.legacy", wires=dev_wires)
 
         original_hits = dev.map_wires.cache_info().hits
         original_misses = dev.map_wires.cache_info().misses
@@ -2002,7 +2002,7 @@ class TestApplyOps:
     gates in DefaultQubit."""
 
     state = np.arange(2**4, dtype=np.complex128).reshape((2, 2, 2, 2))
-    dev = qml.device("default.qubit", wires=4)
+    dev = qml.device("default.qubit.legacy", wires=4)
 
     single_qubit_ops = [
         (qml.PauliX, dev._apply_x),
@@ -2155,7 +2155,7 @@ def test_internal_apply_ops_case(self, monkeypatch):
         This test provides a new internal function that `default.qubit` uses to
         apply `PauliX` (rather than redefining the gate itself).
         """
-        dev = qml.device("default.qubit", wires=1)
+        dev = qml.device("default.qubit.legacy", wires=1)
 
         # Create a dummy operation
         expected_test_output = np.ones(1)
@@ -2174,7 +2174,7 @@ def test_internal_apply_ops_case(self, monkeypatch):
     def test_diagonal_operation_case(self, monkeypatch):
         """Tests the case when the operation to be applied is
         diagonal in the computational basis and the _apply_diagonal_unitary method is used."""
-        dev = qml.device("default.qubit", wires=1)
+        dev = qml.device("default.qubit.legacy", wires=1)
         par = 0.3
 
         test_state = np.array([1, 0])
@@ -2200,7 +2200,7 @@ def test_diagonal_operation_case(self, monkeypatch):
     def test_apply_einsum_case(self, monkeypatch):
         """Tests the case when np.einsum is used to apply an operation in
         default.qubit."""
-        dev = qml.device("default.qubit", wires=1)
+        dev = qml.device("default.qubit.legacy", wires=1)
 
         test_state = np.array([1, 0])
         wires = 0
@@ -2238,7 +2238,7 @@ def compute_matrix(*params, **hyperparams):
     def test_apply_tensordot_case(self, monkeypatch):
         """Tests the case when np.tensordot is used to apply an operation in
         default.qubit."""
-        dev = qml.device("default.qubit", wires=3)
+        dev = qml.device("default.qubit.legacy", wires=3)
 
         test_state = np.array([1, 0])
         wires = [0, 1, 2]
@@ -2277,7 +2277,7 @@ def compute_matrix(*params, **hyperparams):
     def test_apply_unitary_tensordot_double_broadcasting_error(self):
         """Tests that an error is raised if attempting to use _apply_unitary
         with a broadcasted matrix and a broadcasted state simultaneously."""
-        dev = qml.device("default.qubit", wires=3)
+        dev = qml.device("default.qubit.legacy", wires=3)
 
         class BroadcastedToffoli(qml.operation.Operation):
             num_wires = 3
@@ -2298,7 +2298,7 @@ def compute_matrix(*params, **hyperparams):
 
     def test_identity_skipped(self, mocker):
         """Test that applying the identity operation does not perform any additional computations."""
-        dev = qml.device("default.qubit", wires=1)
+        dev = qml.device("default.qubit.legacy", wires=1)
 
         starting_state = np.array([1, 0])
         op = qml.Identity(0)
@@ -2320,7 +2320,7 @@ class TestHamiltonianSupport:
 
     def test_do_not_split_analytic(self, mocker):
         """Tests that the Hamiltonian is not split for shots=None."""
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
         H = qml.Hamiltonian(np.array([0.1, 0.2]), [qml.PauliX(0), qml.PauliZ(1)])
 
         @qml.qnode(dev, diff_method="parameter-shift", interface=None)
@@ -2335,7 +2335,7 @@ def circuit():
 
     def test_split_finite_shots(self, mocker):
         """Tests that the Hamiltonian is split for finite shots."""
-        dev = qml.device("default.qubit", wires=2, shots=10)
+        dev = qml.device("default.qubit.legacy", wires=2, shots=10)
         spy = mocker.spy(dev, "expval")
 
         H = qml.Hamiltonian(np.array([0.1, 0.2]), [qml.PauliX(0), qml.PauliZ(1)])
@@ -2351,7 +2351,7 @@ def circuit():
 
     def test_error_hamiltonian_expval_finite_shots(self):
         """Tests that the Hamiltonian is split for finite shots."""
-        dev = qml.device("default.qubit", wires=2, shots=10)
+        dev = qml.device("default.qubit.legacy", wires=2, shots=10)
         H = qml.Hamiltonian([0.1, 0.2], [qml.PauliX(0), qml.PauliZ(1)])
 
         with pytest.raises(AssertionError, match="Hamiltonian must be used with shots=None"):
@@ -2359,7 +2359,7 @@ def test_error_hamiltonian_expval_finite_shots(self):
 
     def test_error_hamiltonian_expval_wrong_wires(self):
         """Tests that expval fails if Hamiltonian uses non-device wires."""
-        dev = qml.device("default.qubit", wires=2, shots=None)
+        dev = qml.device("default.qubit.legacy", wires=2, shots=None)
         H = qml.Hamiltonian([0.1, 0.2, 0.3], [qml.PauliX(0), qml.PauliZ(1), qml.PauliY(2)])
 
         with pytest.raises(
@@ -2370,7 +2370,7 @@ def test_error_hamiltonian_expval_wrong_wires(self):
 
     def test_Hamiltonian_filtered_from_rotations(self, mocker):
         """Tests that the device does not attempt to get rotations for Hamiltonians."""
-        dev = qml.device("default.qubit", wires=2, shots=10)
+        dev = qml.device("default.qubit.legacy", wires=2, shots=10)
         H = qml.Hamiltonian([0.1, 0.2], [qml.PauliX(0), qml.PauliZ(1)])
 
         spy = mocker.spy(qml.QubitDevice, "_get_diagonalizing_gates")
@@ -2410,7 +2410,7 @@ def circuit(y, z, is_state_batched):
 
     def test_super_expval_not_called(self, is_state_batched, mocker):
         """Tests basic expval result, and ensures QubitDevice.expval is not called."""
-        dev = qml.device("default.qubit", wires=1)
+        dev = qml.device("default.qubit.legacy", wires=1)
         spy = mocker.spy(qml.QubitDevice, "expval")
         obs = qml.sum(qml.s_prod(0.1, qml.PauliX(0)), qml.s_prod(0.2, qml.PauliZ(0)))
         assert np.isclose(dev.expval(obs), 0.2)
@@ -2421,7 +2421,7 @@ def test_trainable_autograd(self, is_state_batched):
         """Tests that coeffs passed to a sum are trainable with autograd."""
         if is_state_batched:
             pytest.skip(msg="Broadcasting, qml.jacobian and new return types do not work together")
-        dev = qml.device("default.qubit", wires=1)
+        dev = qml.device("default.qubit.legacy", wires=1)
         qnode = qml.QNode(self.circuit, dev, interface="autograd")
         y, z = np.array([1.1, 2.2])
         actual = qml.grad(qnode, argnum=[0, 1])(y, z, is_state_batched)
@@ -2432,7 +2432,7 @@ def test_trainable_torch(self, is_state_batched):
         """Tests that coeffs passed to a sum are trainable with torch."""
         import torch
 
-        dev = qml.device("default.qubit", wires=1)
+        dev = qml.device("default.qubit.legacy", wires=1)
         qnode = qml.QNode(self.circuit, dev, interface="torch")
         y, z = torch.tensor(1.1, requires_grad=True), torch.tensor(2.2, requires_grad=True)
         _qnode = partial(qnode, is_state_batched=is_state_batched)
@@ -2444,7 +2444,7 @@ def test_trainable_tf(self, is_state_batched):
         """Tests that coeffs passed to a sum are trainable with tf."""
         import tensorflow as tf
 
-        dev = qml.device("default.qubit", wires=1)
+        dev = qml.device("default.qubit.legacy", wires=1)
         qnode = qml.QNode(self.circuit, dev, interface="tensorflow")
         y, z = tf.Variable(1.1, dtype=tf.float64), tf.Variable(2.2, dtype=tf.float64)
         with tf.GradientTape() as tape:
@@ -2457,7 +2457,7 @@ def test_trainable_jax(self, is_state_batched):
         """Tests that coeffs passed to a sum are trainable with jax."""
         import jax
 
-        dev = qml.device("default.qubit", wires=1)
+        dev = qml.device("default.qubit.legacy", wires=1)
         qnode = qml.QNode(self.circuit, dev, interface="jax")
         y, z = jax.numpy.array([1.1, 2.2])
         actual = jax.jacobian(qnode, argnums=[0, 1])(y, z, is_state_batched)
@@ -2470,7 +2470,7 @@ class TestGetBatchSize:
     @pytest.mark.parametrize("shape", [(4, 4), (1, 8), (4,)])
     def test_batch_size_None(self, shape):
         """Test that a ``batch_size=None`` is reported correctly."""
-        dev = qml.device("default.qubit", wires=1)
+        dev = qml.device("default.qubit.legacy", wires=1)
         tensor0 = np.ones(shape, dtype=complex)
         assert dev._get_batch_size(tensor0, shape, qml.math.prod(shape)) is None
         tensor1 = np.arange(np.prod(shape)).reshape(shape)
@@ -2480,7 +2480,7 @@ def test_batch_size_None(self, shape):
     @pytest.mark.parametrize("batch_size", [1, 3])
     def test_batch_size_int(self, shape, batch_size):
         """Test that an integral ``batch_size`` is reported correctly."""
-        dev = qml.device("default.qubit", wires=1)
+        dev = qml.device("default.qubit.legacy", wires=1)
         full_shape = (batch_size,) + shape
         tensor0 = np.ones(full_shape, dtype=complex)
         assert dev._get_batch_size(tensor0, shape, qml.math.prod(shape)) == batch_size
@@ -2491,7 +2491,7 @@ def test_batch_size_int(self, shape, batch_size):
     def test_invalid_tensor(self):
         """Test that an error is raised if a tensor is provided that does not
         have a proper shape/ndim."""
-        dev = qml.device("default.qubit", wires=1)
+        dev = qml.device("default.qubit.legacy", wires=1)
         with pytest.raises(ValueError, match="could not broadcast"):
             dev._get_batch_size([qml.math.ones((2, 3)), qml.math.ones((2, 2))], (2, 2, 2), 8)
 
@@ -2511,4 +2511,5 @@ class TestDenseMatrixDecompositionThreshold:
     def test_threshold(self, op, n_wires, condition):
         wires = np.linspace(0, n_wires - 1, n_wires, dtype=int)
         op = op(wires=wires)
+        # pylint:disable=no-member
         assert DefaultQubit.stopping_condition.__get__(op)(op) == condition
diff --git a/tests/devices/test_default_qubit_broadcasting.py b/tests/devices/test_default_qubit_legacy_broadcasting.py
similarity index 96%
rename from tests/devices/test_default_qubit_broadcasting.py
rename to tests/devices/test_default_qubit_legacy_broadcasting.py
index e72be460804..594a228cd05 100644
--- a/tests/devices/test_default_qubit_broadcasting.py
+++ b/tests/devices/test_default_qubit_legacy_broadcasting.py
@@ -591,7 +591,7 @@ def test_expval_two_wires_with_parameters_broadcasted(
     def test_expval_estimate_broadcasted(self):
         """Test that the expectation value is not analytically calculated"""
 
-        dev = qml.device("default.qubit", wires=1, shots=3)
+        dev = qml.device("default.qubit.legacy", wires=1, shots=3)
 
         @qml.qnode(dev, diff_method="parameter-shift")
         def circuit():
@@ -691,7 +691,7 @@ def test_var_two_wires_with_parameters_broadcasted(
     def test_var_estimate_broadcasted(self):
         """Test that the variance is not analytically calculated"""
 
-        dev = qml.device("default.qubit", wires=1, shots=3)
+        dev = qml.device("default.qubit.legacy", wires=1, shots=3)
 
         @qml.qnode(dev, diff_method="parameter-shift")
         def circuit():
@@ -716,7 +716,7 @@ def test_sample_dimensions_broadcasted(self):
         # Explicitly resetting is necessary as the internal
         # state is set to None in __init__ and only properly
         # initialized during reset
-        dev = qml.device("default.qubit", wires=2, shots=1000)
+        dev = qml.device("default.qubit.legacy", wires=2, shots=1000)
 
         dev.apply([qml.RX(np.array([np.pi / 2, 0.0]), 0), qml.RX(np.array([np.pi / 2, 0.0]), 1)])
 
@@ -752,7 +752,7 @@ def test_sample_values_broadcasted(self, tol):
         # Explicitly resetting is necessary as the internal
         # state is set to None in __init__ and only properly
         # initialized during reset
-        dev = qml.device("default.qubit", wires=2, shots=1000)
+        dev = qml.device("default.qubit.legacy", wires=2, shots=1000)
 
         dev.apply([qml.RX(np.ones(3), wires=[0])])
         dev._wires_measured = {0}
@@ -787,7 +787,7 @@ def circuit(x):
 
         res = circuit(p)
         assert np.allclose(res, expected, atol=tol, rtol=0)
-        assert res.dtype == r_dtype
+        assert res.dtype == r_dtype  # pylint:disable=no-member
 
     def test_qubit_identity_broadcasted(self, qubit_device_1_wire, tol):
         """Test that the default qubit plugin provides correct result for the Identity expectation"""
@@ -806,7 +806,7 @@ def test_nonzero_shots_broadcasted(self, tol):
         """Test that the default qubit plugin provides correct result for high shot number"""
 
         shots = 10**5
-        dev = qml.device("default.qubit", wires=1, shots=shots)
+        dev = qml.device("default.qubit.legacy", wires=1, shots=shots)
 
         p = np.array([0.543, np.pi / 2, 0.0, 1.0])
 
@@ -912,7 +912,7 @@ def test_multi_samples_return_correlated_results_broadcasted(self):
         correctly for correlated observables.
         """
 
-        dev = qml.device("default.qubit", wires=2, shots=1000)
+        dev = qml.device("default.qubit.legacy", wires=2, shots=1000)
 
         @qml.qnode(dev, diff_method="parameter-shift")
         def circuit():
@@ -933,7 +933,7 @@ def test_multi_samples_correlated_results_more_wires_than_observable_broadcasted
         correctly for correlated observables on larger devices than the observables
         """
 
-        dev = qml.device("default.qubit", wires=num_wires, shots=1000)
+        dev = qml.device("default.qubit.legacy", wires=num_wires, shots=1000)
 
         @qml.qnode(dev, diff_method="parameter-shift")
         def circuit():
@@ -956,7 +956,7 @@ class TestTensorExpvalBroadcasted:
 
     def test_paulix_pauliy_broadcasted(self, theta, phi, varphi, tol):
         """Test that a tensor product involving PauliX and PauliY works correctly"""
-        dev = qml.device("default.qubit", wires=3)
+        dev = qml.device("default.qubit.legacy", wires=3)
         dev.reset()
 
         obs = qml.PauliX(0) @ qml.PauliY(2)
@@ -980,7 +980,7 @@ def test_paulix_pauliy_broadcasted(self, theta, phi, varphi, tol):
 
     def test_pauliz_identity_broadcasted(self, theta, phi, varphi, tol):
         """Test that a tensor product involving PauliZ and Identity works correctly"""
-        dev = qml.device("default.qubit", wires=3)
+        dev = qml.device("default.qubit.legacy", wires=3)
         dev.reset()
 
         obs = qml.PauliZ(0) @ qml.Identity(1) @ qml.PauliZ(2)
@@ -1004,7 +1004,7 @@ def test_pauliz_identity_broadcasted(self, theta, phi, varphi, tol):
 
     def test_pauliz_hadamard_broadcasted(self, theta, phi, varphi, tol):
         """Test that a tensor product involving PauliZ and PauliY and hadamard works correctly"""
-        dev = qml.device("default.qubit", wires=3)
+        dev = qml.device("default.qubit.legacy", wires=3)
         obs = qml.PauliZ(0) @ qml.Hadamard(1) @ qml.PauliY(2)
 
         dev.reset()
@@ -1027,7 +1027,7 @@ def test_pauliz_hadamard_broadcasted(self, theta, phi, varphi, tol):
 
     def test_hermitian_broadcasted(self, theta, phi, varphi, tol):
         """Test that a tensor product involving qml.Hermitian works correctly"""
-        dev = qml.device("default.qubit", wires=3)
+        dev = qml.device("default.qubit.legacy", wires=3)
         dev.reset()
 
         A = np.array(
@@ -1065,7 +1065,7 @@ def test_hermitian_broadcasted(self, theta, phi, varphi, tol):
 
     def test_hermitian_hermitian_broadcasted(self, theta, phi, varphi, tol):
         """Test that a tensor product involving two Hermitian matrices works correctly"""
-        dev = qml.device("default.qubit", wires=3)
+        dev = qml.device("default.qubit.legacy", wires=3)
 
         A1 = np.array([[1, 2], [2, 4]])
 
@@ -1114,7 +1114,7 @@ def test_hermitian_hermitian_broadcasted(self, theta, phi, varphi, tol):
 
     def test_hermitian_identity_expectation_broadcasted(self, theta, phi, varphi, tol):
         """Test that a tensor product involving an Hermitian matrix and the identity works correctly"""
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
 
         A = np.array(
             [[1.02789352, 1.61296440 - 0.3498192j], [1.61296440 + 0.3498192j, 1.23920938 + 0j]]
@@ -1138,7 +1138,7 @@ def test_hermitian_identity_expectation_broadcasted(self, theta, phi, varphi, to
 
     def test_hermitian_two_wires_identity_expectation_broadcasted(self, theta, phi, varphi, tol):
         """Test that a tensor product involving an Hermitian matrix for two wires and the identity works correctly"""
-        dev = qml.device("default.qubit", wires=3)
+        dev = qml.device("default.qubit.legacy", wires=3)
 
         A = np.array(
             [[1.02789352, 1.61296440 - 0.3498192j], [1.61296440 + 0.3498192j, 1.23920938 + 0j]]
@@ -1169,7 +1169,7 @@ class TestTensorVarBroadcasted:
 
     def test_paulix_pauliy_broadcasted(self, theta, phi, varphi, tol):
         """Test that a tensor product involving PauliX and PauliY works correctly"""
-        dev = qml.device("default.qubit", wires=3)
+        dev = qml.device("default.qubit.legacy", wires=3)
 
         obs = qml.PauliX(0) @ qml.PauliY(2)
 
@@ -1199,7 +1199,7 @@ def test_paulix_pauliy_broadcasted(self, theta, phi, varphi, tol):
 
     def test_pauliz_hadamard_broadcasted(self, theta, phi, varphi, tol):
         """Test that a tensor product involving PauliZ and PauliY and hadamard works correctly"""
-        dev = qml.device("default.qubit", wires=3)
+        dev = qml.device("default.qubit.legacy", wires=3)
         obs = qml.PauliZ(0) @ qml.Hadamard(1) @ qml.PauliY(2)
 
         dev.reset()
@@ -1227,7 +1227,7 @@ def test_pauliz_hadamard_broadcasted(self, theta, phi, varphi, tol):
 
     def test_hermitian_broadcasted(self, theta, phi, varphi, tol):
         """Test that a tensor product involving qml.Hermitian works correctly"""
-        dev = qml.device("default.qubit", wires=3)
+        dev = qml.device("default.qubit.legacy", wires=3)
 
         A = np.array(
             [
@@ -1295,7 +1295,7 @@ class TestTensorSampleBroadcasted:
 
     def test_paulix_pauliy_broadcasted(self, theta, phi, varphi, tol_stochastic):
         """Test that a tensor product involving PauliX and PauliY works correctly"""
-        dev = qml.device("default.qubit", wires=3, shots=int(1e6))
+        dev = qml.device("default.qubit.legacy", wires=3, shots=int(1e6))
 
         obs = qml.PauliX(0) @ qml.PauliY(2)
 
@@ -1337,7 +1337,7 @@ def test_paulix_pauliy_broadcasted(self, theta, phi, varphi, tol_stochastic):
 
     def test_pauliz_hadamard_broadcasted(self, theta, phi, varphi, tol_stochastic):
         """Test that a tensor product involving PauliZ and PauliY and hadamard works correctly"""
-        dev = qml.device("default.qubit", wires=3, shots=int(1e6))
+        dev = qml.device("default.qubit.legacy", wires=3, shots=int(1e6))
         obs = qml.PauliZ(0) @ qml.Hadamard(1) @ qml.PauliY(2)
         dev.apply(
             [
@@ -1375,7 +1375,7 @@ def test_pauliz_hadamard_broadcasted(self, theta, phi, varphi, tol_stochastic):
 
     def test_hermitian_broadcasted(self, theta, phi, varphi, tol_stochastic):
         """Test that a tensor product involving qml.Hermitian works correctly"""
-        dev = qml.device("default.qubit", wires=3, shots=int(1e6))
+        dev = qml.device("default.qubit.legacy", wires=3, shots=int(1e6))
 
         A = 0.1 * np.array(
             [
@@ -1490,7 +1490,7 @@ def test_state_dtype_after_op_broadcasted(self, r_dtype, c_dtype, op):
         examples.
         """
 
-        dev = qml.device("default.qubit", wires=4, r_dtype=r_dtype, c_dtype=c_dtype)
+        dev = qml.device("default.qubit.legacy", wires=4, r_dtype=r_dtype, c_dtype=c_dtype)
 
         n_wires = op.num_wires
         n_params = op.num_params
@@ -1507,7 +1507,7 @@ def circuit():
             return qml.state()
 
         res = circuit()
-        assert res.dtype == c_dtype
+        assert res.dtype == c_dtype  # pylint:disable=no-member
 
     @pytest.mark.parametrize(
         "measurement",
@@ -1522,7 +1522,7 @@ def test_measurement_real_dtype_broadcasted(self, r_dtype, c_dtype, measurement)
         """Test that the default qubit plugin provides correct result for a simple circuit"""
         p = np.array([0.543, 0.622, 1.3])
 
-        dev = qml.device("default.qubit", wires=3, r_dtype=r_dtype, c_dtype=c_dtype)
+        dev = qml.device("default.qubit.legacy", wires=3, r_dtype=r_dtype, c_dtype=c_dtype)
 
         @qml.qnode(dev, diff_method="parameter-shift")
         def circuit(x):
@@ -1537,7 +1537,7 @@ def test_measurement_complex_dtype_broadcasted(self, r_dtype, c_dtype):
         p = np.array([0.543, 0.622, 1.3])
         m = qml.state()
 
-        dev = qml.device("default.qubit", wires=3, r_dtype=r_dtype, c_dtype=c_dtype)
+        dev = qml.device("default.qubit.legacy", wires=3, r_dtype=r_dtype, c_dtype=c_dtype)
 
         @qml.qnode(dev, diff_method="parameter-shift")
         def circuit(x):
@@ -1558,8 +1558,8 @@ def mock_analytic_counter(self, wires=None):
 
     def test_probability_broadcasted(self, tol):
         """Test that the probability function works for finite and infinite shots"""
-        dev = qml.device("default.qubit", wires=2, shots=1000)
-        dev_analytic = qml.device("default.qubit", wires=2, shots=None)
+        dev = qml.device("default.qubit.legacy", wires=2, shots=1000)
+        dev_analytic = qml.device("default.qubit.legacy", wires=2, shots=None)
 
         x = np.array([[0.2, 0.5, 0.4], [0.9, 0.8, 0.3]])
 
@@ -1582,7 +1582,7 @@ class TestWiresIntegrationBroadcasted:
 
     def make_circuit_probs(self, wires):
         """Factory for a qnode returning probabilities using arbitrary wire labels."""
-        dev = qml.device("default.qubit", wires=wires)
+        dev = qml.device("default.qubit.legacy", wires=wires)
         n_wires = len(wires)
 
         @qml.qnode(dev, diff_method="parameter-shift")
@@ -1619,7 +1619,7 @@ class TestApplyOpsBroadcasted:
     gates in DefaultQubit."""
 
     broadcasted_state = np.arange(2**4 * 3, dtype=np.complex128).reshape((3, 2, 2, 2, 2))
-    dev = qml.device("default.qubit", wires=4)
+    dev = qml.device("default.qubit.legacy", wires=4)
 
     single_qubit_ops = [
         (qml.PauliX, dev._apply_x),
@@ -1746,7 +1746,7 @@ def test_internal_apply_ops_case_broadcasted(self, monkeypatch):
         This test provides a new internal function that `default.qubit` uses to
         apply `PauliX` (rather than redefining the gate itself).
         """
-        dev = qml.device("default.qubit", wires=1)
+        dev = qml.device("default.qubit.legacy", wires=1)
 
         test_state = np.array([[1, 0], [INVSQ2, INVSQ2], [0, 1]])
         # Create a dummy operation
@@ -1765,7 +1765,7 @@ def test_internal_apply_ops_case_broadcasted(self, monkeypatch):
     def test_diagonal_operation_case_broadcasted(self, monkeypatch):
         """Tests the case when the operation to be applied is
         diagonal in the computational basis and the _apply_diagonal_unitary method is used."""
-        dev = qml.device("default.qubit", wires=1)
+        dev = qml.device("default.qubit.legacy", wires=1)
         par = 0.3
 
         test_state = np.array([[1, 0], [INVSQ2, INVSQ2], [0, 1]])
@@ -1791,7 +1791,7 @@ def test_diagonal_operation_case_broadcasted(self, monkeypatch):
     def test_apply_einsum_case_broadcasted(self, monkeypatch):
         """Tests the case when np.einsum is used to apply an operation in
         default.qubit."""
-        dev = qml.device("default.qubit", wires=1)
+        dev = qml.device("default.qubit.legacy", wires=1)
 
         test_state = np.array([[1, 0], [INVSQ2, INVSQ2], [0, 1]])
         wires = 0
@@ -1830,7 +1830,7 @@ def compute_matrix(*params, **hyperparams):
     def test_apply_tensordot_case_broadcasted(self, monkeypatch):
         """Tests the case when np.tensordot is used to apply an operation in
         default.qubit."""
-        dev = qml.device("default.qubit", wires=3)
+        dev = qml.device("default.qubit.legacy", wires=3)
 
         test_state = np.array([[1, 0], [INVSQ2, INVSQ2], [0, 1]])
         wires = [0, 1, 2]
@@ -1869,7 +1869,7 @@ def compute_matrix(*params, **hyperparams):
 
     def test_identity_skipped_broadcasted(self, mocker):
         """Test that applying the identity operation does not perform any additional computations."""
-        dev = qml.device("default.qubit", wires=1)
+        dev = qml.device("default.qubit.legacy", wires=1)
 
         starting_state = np.array([[1, 0], [INVSQ2, INVSQ2], [0, 1]])
         op = qml.Identity(0)
@@ -1891,7 +1891,7 @@ class TestHamiltonianSupportBroadcasted:
 
     def test_do_not_split_analytic_broadcasted(self, mocker):
         """Tests that the Hamiltonian is not split for shots=None."""
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
         Ham = qml.Hamiltonian(np.array([0.1, 0.2]), [qml.PauliX(0), qml.PauliZ(1)])
 
         @qml.qnode(dev, diff_method="parameter-shift", interface=None)
@@ -1907,7 +1907,7 @@ def circuit():
 
     def test_split_finite_shots_broadcasted(self, mocker):
         """Tests that the Hamiltonian is split for finite shots."""
-        dev = qml.device("default.qubit", wires=2, shots=10)
+        dev = qml.device("default.qubit.legacy", wires=2, shots=10)
         spy = mocker.spy(dev, "expval")
 
         ham = qml.Hamiltonian(np.array([0.1, 0.2]), [qml.PauliX(0), qml.PauliZ(1)])
diff --git a/tests/devices/test_default_qubit_tf.py b/tests/devices/test_default_qubit_tf.py
index 1443d5926e1..f30d4d70f17 100644
--- a/tests/devices/test_default_qubit_tf.py
+++ b/tests/devices/test_default_qubit_tf.py
@@ -1795,8 +1795,8 @@ def circuit(x):
                 qml.CNOT(wires=[i, i + 1])
             return qml.expval(qml.PauliZ(0)), qml.var(qml.PauliZ(1))
 
-        dev1 = qml.device("default.qubit", wires=3)
-        dev2 = qml.device("default.qubit", wires=3)
+        dev1 = qml.device("default.qubit.legacy", wires=3)
+        dev2 = qml.device("default.qubit.legacy", wires=3)
 
         def cost(x):
             return qml.math.stack(circuit(x))
diff --git a/tests/devices/test_default_qubit_torch.py b/tests/devices/test_default_qubit_torch.py
index f0cb3ee284a..9c2a4c53801 100644
--- a/tests/devices/test_default_qubit_torch.py
+++ b/tests/devices/test_default_qubit_torch.py
@@ -1962,7 +1962,7 @@ def circuit(x):
             return qml.expval(qml.PauliZ(0))  # , qml.var(qml.PauliZ(1))
 
         dev1 = qml.device("default.qubit.torch", wires=3, torch_device=torch_device)
-        dev2 = qml.device("default.qubit", wires=3)
+        dev2 = qml.device("default.qubit.legacy", wires=3)
 
         circuit1 = qml.QNode(circuit, dev1, diff_method="backprop", interface="torch")
         circuit2 = qml.QNode(circuit, dev2, diff_method="parameter-shift")
diff --git a/tests/gradients/core/test_adjoint_diff.py b/tests/gradients/core/test_adjoint_diff.py
index f42ad56a6bc..7e85b140598 100644
--- a/tests/gradients/core/test_adjoint_diff.py
+++ b/tests/gradients/core/test_adjoint_diff.py
@@ -26,7 +26,7 @@ class TestAdjointJacobian:
     @pytest.fixture
     def dev(self):
         """Fixture that creates a device with two wires."""
-        return qml.device("default.qubit", wires=2)
+        return qml.device("default.qubit.legacy", wires=2)
 
     def test_not_expval(self, dev):
         """Test if a QuantumFunctionError is raised for a tape with measurements that are not
@@ -43,7 +43,7 @@ def test_not_expval(self, dev):
     def test_finite_shots_warns(self):
         """Tests warning raised when finite shots specified"""
 
-        dev = qml.device("default.qubit", wires=1, shots=10)
+        dev = qml.device("default.qubit.legacy", wires=1, shots=10)
 
         with qml.queuing.AnnotatedQueue() as q:
             qml.expval(qml.PauliZ(0))
@@ -88,7 +88,7 @@ def test_trainable_hermitian_warns(self):
         """Test attempting to compute the gradient of a tape that obtains the
         expectation value of a Hermitian operator emits a warning if the
         parameters to Hermitian are trainable."""
-        dev = qml.device("default.qubit", wires=3)
+        dev = qml.device("default.qubit.legacy", wires=3)
 
         mx = qml.matrix(qml.PauliX(0) @ qml.PauliY(2))
         with qml.queuing.AnnotatedQueue() as q:
@@ -193,7 +193,7 @@ def test_rx_gradient(self, tol, dev):
 
     def test_multiple_rx_gradient(self, tol):
         """Tests that the gradient of multiple RX gates in a circuit yields the correct result."""
-        dev = qml.device("default.qubit", wires=3)
+        dev = qml.device("default.qubit.legacy", wires=3)
         params = np.array([np.pi, np.pi / 2, np.pi / 3])
 
         with qml.queuing.AnnotatedQueue() as q:
@@ -331,7 +331,7 @@ def test_provide_starting_state(self, tol, dev):
     def test_gradient_of_tape_with_hermitian(self, tol):
         """Test that computing the gradient of a tape that obtains the
         expectation value of a Hermitian operator works correctly."""
-        dev = qml.device("default.qubit", wires=3)
+        dev = qml.device("default.qubit.legacy", wires=3)
 
         a, b, c = [0.5, 0.3, -0.7]
 
diff --git a/tests/interfaces/test_autograd.py b/tests/interfaces/test_autograd.py
index 6f32aec02bb..6732eb9316b 100644
--- a/tests/interfaces/test_autograd.py
+++ b/tests/interfaces/test_autograd.py
@@ -41,7 +41,7 @@ def test_import_error(self, mocker):
         except KeyError:
             pass
 
-        dev = qml.device("default.qubit", wires=2, shots=None)
+        dev = qml.device("default.qubit.legacy", wires=2, shots=None)
 
         with qml.queuing.AnnotatedQueue() as q:
             qml.expval(qml.PauliY(1))
@@ -60,7 +60,7 @@ def test_jacobian_options(self, mocker):
 
         a = np.array([0.1, 0.2], requires_grad=True)
 
-        dev = qml.device("default.qubit", wires=1)
+        dev = qml.device("default.qubit.legacy", wires=1)
 
         def cost(a, device):
             with qml.queuing.AnnotatedQueue() as q:
@@ -86,7 +86,7 @@ def test_incorrect_grad_on_execution(self):
         is used with grad_on_execution=True"""
         a = np.array([0.1, 0.2], requires_grad=True)
 
-        dev = qml.device("default.qubit", wires=1)
+        dev = qml.device("default.qubit.legacy", wires=1)
 
         def cost(a, device):
             with qml.queuing.AnnotatedQueue() as q:
@@ -106,7 +106,7 @@ def test_unknown_interface(self):
         """Test that an error is raised if the interface is unknown"""
         a = np.array([0.1, 0.2], requires_grad=True)
 
-        dev = qml.device("default.qubit", wires=1)
+        dev = qml.device("default.qubit.legacy", wires=1)
 
         def cost(a, device):
             with qml.queuing.AnnotatedQueue() as q:
@@ -122,7 +122,7 @@ def cost(a, device):
 
     def test_grad_on_execution(self, mocker):
         """Test that grad on execution uses the `device.execute_and_gradients` pathway"""
-        dev = qml.device("default.qubit", wires=1)
+        dev = qml.device("default.qubit.legacy", wires=1)
         spy = mocker.spy(dev, "execute_and_gradients")
 
         def cost(a):
@@ -148,7 +148,7 @@ def cost(a):
 
     def test_no_gradients_on_execution(self, mocker):
         """Test that no grad on execution uses the `device.batch_execute` and `device.gradients` pathway"""
-        dev = qml.device("default.qubit", wires=1)
+        dev = qml.device("default.qubit.legacy", wires=1)
         spy_execute = mocker.spy(qml.devices.DefaultQubit, "batch_execute")
         spy_gradients = mocker.spy(qml.devices.DefaultQubit, "gradients")
 
@@ -184,7 +184,7 @@ class TestBatchTransformExecution:
 
     def test_no_batch_transform(self, mocker):
         """Test that batch transforms can be disabled and enabled"""
-        dev = qml.device("default.qubit", wires=2, shots=100000)
+        dev = qml.device("default.qubit.legacy", wires=2, shots=100000)
 
         H = qml.PauliZ(0) @ qml.PauliZ(1) - qml.PauliX(0)
         x = 0.6
@@ -214,7 +214,7 @@ def test_no_batch_transform(self, mocker):
     def test_batch_transform_dynamic_shots(self):
         """Tests that the batch transform considers the number of shots for the execution, not those
         statically on the device."""
-        dev = qml.device("default.qubit", wires=1)
+        dev = qml.device("default.qubit.legacy", wires=1)
         H = 2.0 * qml.PauliZ(0)
         qscript = qml.tape.QuantumScript(measurements=[qml.expval(H)])
         res = qml.execute([qscript], dev, interface=None, override_shots=10)
@@ -226,7 +226,7 @@ class TestCaching:
 
     def test_cache_maxsize(self, mocker):
         """Test the cachesize property of the cache"""
-        dev = qml.device("default.qubit", wires=1)
+        dev = qml.device("default.qubit.legacy", wires=1)
         spy = mocker.spy(qml.interfaces, "cache_execute")
 
         def cost(a, cachesize):
@@ -248,7 +248,7 @@ def cost(a, cachesize):
 
     def test_custom_cache(self, mocker):
         """Test the use of a custom cache object"""
-        dev = qml.device("default.qubit", wires=1)
+        dev = qml.device("default.qubit.legacy", wires=1)
         spy = mocker.spy(qml.interfaces, "cache_execute")
 
         def cost(a, cache):
@@ -270,7 +270,7 @@ def cost(a, cache):
     def test_caching_param_shift(self, tol):
         """Test that, when using parameter-shift transform,
         caching reduces the number of evaluations to their optimum."""
-        dev = qml.device("default.qubit", wires=1)
+        dev = qml.device("default.qubit.legacy", wires=1)
 
         def cost(a, cache):
             with qml.queuing.AnnotatedQueue() as q:
@@ -315,7 +315,7 @@ def test_caching_param_shift_hessian(self, num_params, tol):
         """Test that, when using parameter-shift transform,
         caching reduces the number of evaluations to their optimum
         when computing Hessians."""
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
         params = np.arange(1, num_params + 1) / 10
 
         N = len(params)
@@ -378,7 +378,7 @@ def cost(x, cache):
     def test_caching_adjoint_no_grad_on_execution(self):
         """Test that caching reduces the number of adjoint evaluations
         when the grads is not on execution."""
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
         params = np.array([0.1, 0.2, 0.3])
 
         def cost(a, cache):
@@ -418,7 +418,7 @@ def test_single_backward_pass_batch(self):
         """Tests that the backward pass is one single batch, not a bunch of batches, when parameter shift
         is requested for multiple tapes."""
 
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
 
         def f(x):
             tape1 = qml.tape.QuantumScript([qml.RX(x, 0)], [qml.probs(wires=0)])
@@ -440,7 +440,7 @@ def test_single_backward_pass_split_hamiltonian(self):
         """Tests that the backward pass is one single batch, not a bunch of batches, when parameter shift
         derivatives are requested for a a tape that the device split into batches."""
 
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
 
         H = qml.Hamiltonian([1, 1], [qml.PauliY(0), qml.PauliZ(0)], grouping_type="qwc")
 
@@ -480,7 +480,7 @@ class TestAutogradExecuteIntegration:
 
     def test_execution(self, execute_kwargs):
         """Test execution"""
-        dev = qml.device("default.qubit", wires=1)
+        dev = qml.device("default.qubit.legacy", wires=1)
 
         def cost(a, b):
             with qml.queuing.AnnotatedQueue() as q1:
@@ -508,7 +508,7 @@ def cost(a, b):
     def test_scalar_jacobian(self, execute_kwargs, tol):
         """Test scalar jacobian calculation"""
         a = np.array(0.1, requires_grad=True)
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
 
         def cost(a):
             with qml.queuing.AnnotatedQueue() as q:
@@ -548,7 +548,7 @@ def cost(a, b, device):
             tape = qml.tape.QuantumScript.from_queue(q)
             return autograd.numpy.hstack(qml.execute([tape], device, **execute_kwargs)[0])
 
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
 
         res = cost(a, b, device=dev)
         expected = [np.cos(a), -np.cos(a) * np.sin(b)]
@@ -570,7 +570,7 @@ def test_tape_no_parameters(self, execute_kwargs, tol):
         if execute_kwargs["gradient_fn"] == "device":
             pytest.skip("Adjoint differentiation does not yet support probabilities")
 
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
 
         def cost(params):
             with qml.queuing.AnnotatedQueue() as q1:
@@ -614,7 +614,7 @@ def cost(params):
     @pytest.mark.filterwarnings("ignore:Attempted to compute the gradient")
     def test_tapes_with_different_return_size(self, execute_kwargs):
         """Test that tapes wit different can be executed and differentiated."""
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
 
         def cost(params):
             with qml.queuing.AnnotatedQueue() as q1:
@@ -652,7 +652,7 @@ def test_reusing_quantum_tape(self, execute_kwargs, tol):
         a = np.array(0.1, requires_grad=True)
         b = np.array(0.2, requires_grad=True)
 
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
 
         with qml.queuing.AnnotatedQueue() as q:
             qml.RY(a, wires=0)
@@ -705,7 +705,7 @@ def cost(a, b, c, device):
             tape = qml.tape.QuantumScript.from_queue(q)
             return qml.execute([tape], device, **execute_kwargs)[0]
 
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
         res = qml.jacobian(cost)(a, b, c, device=dev)
 
         # Only two arguments are trainable
@@ -728,7 +728,7 @@ def cost(a, b, device):
             tape = qml.tape.QuantumScript.from_queue(q)
             return autograd.numpy.hstack(qml.execute([tape], device, **execute_kwargs)[0])
 
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
         res = cost(a, b, device=dev)
         assert res.shape == (2,)
 
@@ -758,7 +758,7 @@ def cost(a, U, device):
             tape = qml.tape.QuantumScript.from_queue(q)
             return qml.execute([tape], device, **execute_kwargs)[0]
 
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
         res = cost(a, U, device=dev)
         assert isinstance(res, np.ndarray)
         assert np.allclose(res, -np.cos(a), atol=tol, rtol=0)
@@ -794,7 +794,7 @@ def cost_fn(a, p, device):
         a = np.array(0.1, requires_grad=False)
         p = np.array([0.1, 0.2, 0.3], requires_grad=True)
 
-        dev = qml.device("default.qubit", wires=1)
+        dev = qml.device("default.qubit.legacy", wires=1)
         res = cost_fn(a, p, device=dev)
         expected = np.cos(a) * np.cos(p[1]) * np.sin(p[0]) + np.sin(a) * (
             np.cos(p[2]) * np.sin(p[1]) + np.cos(p[0]) * np.cos(p[1]) * np.sin(p[2])
@@ -833,7 +833,7 @@ def cost(x, y, device):
             tape = qml.tape.QuantumScript.from_queue(q)
             return autograd.numpy.hstack(qml.execute([tape], device, **execute_kwargs)[0])
 
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
         x = np.array(0.543, requires_grad=True)
         y = np.array(-0.654, requires_grad=True)
 
@@ -894,7 +894,7 @@ def cost(x, y, device):
             tape = qml.tape.QuantumScript.from_queue(q)
             return autograd.numpy.hstack(qml.execute([tape], device, **execute_kwargs)[0])
 
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
         x = np.array(0.543, requires_grad=True)
         y = np.array(-0.654, requires_grad=True)
 
@@ -936,7 +936,7 @@ def cost(device):
             return qml.execute([tape], device, **execute_kwargs)[0]
 
         shots = 10
-        dev = qml.device("default.qubit", wires=2, shots=shots)
+        dev = qml.device("default.qubit.legacy", wires=2, shots=shots)
         res = cost(device=dev)
         assert isinstance(res, tuple)
         assert len(res) == 2
@@ -1056,7 +1056,7 @@ def cost_fn(x):
     def test_max_diff(self, tol):
         """Test that setting the max_diff parameter blocks higher-order
         derivatives"""
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
         params = np.array([0.543, -0.654], requires_grad=True)
 
         def cost_fn(x):
@@ -1100,7 +1100,7 @@ class TestOverridingShots:
 
     def test_changing_shots(self, mocker, tol):
         """Test that changing shots works on execution"""
-        dev = qml.device("default.qubit", wires=2, shots=None)
+        dev = qml.device("default.qubit.legacy", wires=2, shots=None)
         a, b = np.array([0.543, -0.654], requires_grad=True)
 
         with qml.queuing.AnnotatedQueue() as q:
@@ -1130,7 +1130,7 @@ def test_changing_shots(self, mocker, tol):
 
     def test_overriding_shots_with_same_value(self, mocker):
         """Overriding shots with the same value as the device will have no effect"""
-        dev = qml.device("default.qubit", wires=2, shots=123)
+        dev = qml.device("default.qubit.legacy", wires=2, shots=123)
         a, b = np.array([0.543, -0.654], requires_grad=True)
 
         with qml.queuing.AnnotatedQueue() as q:
@@ -1161,7 +1161,7 @@ def test_overriding_shots_with_same_value(self, mocker):
     def test_overriding_device_with_shot_vector(self):
         """Overriding a device that has a batch of shots set
         results in original shots being returned after execution"""
-        dev = qml.device("default.qubit", wires=2, shots=[10, (1, 3), 5])
+        dev = qml.device("default.qubit.legacy", wires=2, shots=[10, (1, 3), 5])
 
         assert dev.shots == 18
         assert dev._shot_vector == [(10, 1), (1, 3), (5, 1)]
@@ -1192,7 +1192,7 @@ def test_gradient_integration(self):
         """Test that temporarily setting the shots works
         for gradient computations"""
         # TODO: Update here when shot vectors are supported
-        dev = qml.device("default.qubit", wires=2, shots=None)
+        dev = qml.device("default.qubit.legacy", wires=2, shots=None)
         a, b = np.array([0.543, -0.654], requires_grad=True)
 
         def cost_fn(a, b, shots):
@@ -1285,7 +1285,7 @@ def test_multiple_hamiltonians_not_trainable(self, cost_fn, execute_kwargs, tol)
         coeffs1 = np.array([0.1, 0.2, 0.3], requires_grad=False)
         coeffs2 = np.array([0.7], requires_grad=False)
         weights = np.array([0.4, 0.5], requires_grad=True)
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
 
         res = cost_fn(weights, coeffs1, coeffs2, dev=dev)
         expected = self.cost_fn_expected(weights, coeffs1, coeffs2)
@@ -1301,7 +1301,7 @@ def test_multiple_hamiltonians_trainable(self, cost_fn, execute_kwargs, tol):
         coeffs1 = np.array([0.1, 0.2, 0.3], requires_grad=True)
         coeffs2 = np.array([0.7], requires_grad=True)
         weights = np.array([0.4, 0.5], requires_grad=True)
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
 
         res = cost_fn(weights, coeffs1, coeffs2, dev=dev)
         expected = self.cost_fn_expected(weights, coeffs1, coeffs2)
diff --git a/tests/interfaces/test_autograd_qnode.py b/tests/interfaces/test_autograd_qnode.py
index 476e694f6f0..fdaf469980c 100644
--- a/tests/interfaces/test_autograd_qnode.py
+++ b/tests/interfaces/test_autograd_qnode.py
@@ -22,30 +22,30 @@
 from pennylane import qnode
 
 qubit_device_and_diff_method = [
-    ["default.qubit", "finite-diff", False],
-    ["default.qubit", "parameter-shift", False],
-    ["default.qubit", "backprop", True],
-    ["default.qubit", "adjoint", True],
-    ["default.qubit", "adjoint", False],
-    ["default.qubit", "spsa", False],
-    ["default.qubit", "hadamard", False],
+    ["default.qubit.legacy", "finite-diff", False],
+    ["default.qubit.legacy", "parameter-shift", False],
+    ["default.qubit.legacy", "backprop", True],
+    ["default.qubit.legacy", "adjoint", True],
+    ["default.qubit.legacy", "adjoint", False],
+    ["default.qubit.legacy", "spsa", False],
+    ["default.qubit.legacy", "hadamard", False],
 ]
 
 interface_qubit_device_and_diff_method = [
-    ["autograd", "default.qubit", "finite-diff", False],
-    ["autograd", "default.qubit", "parameter-shift", False],
-    ["autograd", "default.qubit", "backprop", True],
-    ["autograd", "default.qubit", "adjoint", True],
-    ["autograd", "default.qubit", "adjoint", False],
-    ["autograd", "default.qubit", "spsa", False],
-    ["autograd", "default.qubit", "hadamard", False],
-    ["auto", "default.qubit", "finite-diff", False],
-    ["auto", "default.qubit", "parameter-shift", False],
-    ["auto", "default.qubit", "backprop", True],
-    ["auto", "default.qubit", "adjoint", True],
-    ["auto", "default.qubit", "adjoint", False],
-    ["auto", "default.qubit", "spsa", False],
-    ["auto", "default.qubit", "hadamard", False],
+    ["autograd", "default.qubit.legacy", "finite-diff", False],
+    ["autograd", "default.qubit.legacy", "parameter-shift", False],
+    ["autograd", "default.qubit.legacy", "backprop", True],
+    ["autograd", "default.qubit.legacy", "adjoint", True],
+    ["autograd", "default.qubit.legacy", "adjoint", False],
+    ["autograd", "default.qubit.legacy", "spsa", False],
+    ["autograd", "default.qubit.legacy", "hadamard", False],
+    ["auto", "default.qubit.legacy", "finite-diff", False],
+    ["auto", "default.qubit.legacy", "parameter-shift", False],
+    ["auto", "default.qubit.legacy", "backprop", True],
+    ["auto", "default.qubit.legacy", "adjoint", True],
+    ["auto", "default.qubit.legacy", "adjoint", False],
+    ["auto", "default.qubit.legacy", "spsa", False],
+    ["auto", "default.qubit.legacy", "hadamard", False],
 ]
 
 pytestmark = pytest.mark.autograd
@@ -71,7 +71,7 @@ def test_nondiff_param_unwrapping(
         if diff_method != "parameter-shift":
             pytest.skip("Test only supports parameter-shift")
 
-        dev = qml.device("default.qubit", wires=1)
+        dev = qml.device("default.qubit.legacy", wires=1)
 
         @qnode(dev, interface=interface, diff_method=diff_method)
         def circuit(x, y):
@@ -301,7 +301,7 @@ def test_jacobian_options(self, interface, dev_name, diff_method, grad_on_execut
 
         a = np.array([0.1, 0.2], requires_grad=True)
 
-        dev = qml.device("default.qubit", wires=wires)
+        dev = qml.device("default.qubit.legacy", wires=wires)
 
         @qnode(dev, interface=interface, diff_method=diff_method, **kwargs)
         def circuit(a):
@@ -329,7 +329,7 @@ def test_changing_trainability(
         a = np.array(0.1, requires_grad=True)
         b = np.array(0.2, requires_grad=True)
 
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
 
         @qnode(dev, interface=interface, diff_method=diff_method)
         def circuit(a, b):
@@ -554,7 +554,7 @@ class TestShotsIntegration:
 
     def test_changing_shots(self, mocker, tol):
         """Test that changing shots works on execution"""
-        dev = qml.device("default.qubit", wires=2, shots=None)
+        dev = qml.device("default.qubit.legacy", wires=2, shots=None)
         a, b = np.array([0.543, -0.654], requires_grad=True)
 
         @qnode(dev, diff_method=qml.gradients.param_shift)
@@ -586,7 +586,7 @@ def circuit(a, b):
     def test_gradient_integration(self):
         """Test that temporarily setting the shots works
         for gradient computations"""
-        dev = qml.device("default.qubit", wires=2, shots=None)
+        dev = qml.device("default.qubit.legacy", wires=2, shots=None)
         a, b = np.array([0.543, -0.654], requires_grad=True)
 
         @qnode(dev, diff_method=qml.gradients.param_shift)
@@ -609,7 +609,7 @@ def cost_fn(a, b):
 
     def test_update_diff_method(self, mocker):
         """Test that temporarily setting the shots updates the diff method"""
-        dev = qml.device("default.qubit", wires=2, shots=100)
+        dev = qml.device("default.qubit.legacy", wires=2, shots=100)
         a, b = np.array([0.543, -0.654], requires_grad=True)
 
         spy = mocker.spy(qml, "execute")
@@ -934,7 +934,7 @@ def circuit1(a, b, c):
             qml.CNOT(wires=[1, 2])
             return qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliY(2))
 
-        dev2 = qml.device("default.qubit", wires=num_wires)
+        dev2 = qml.device("default.qubit.legacy", wires=num_wires)
 
         @qnode(dev2, interface=interface, diff_method=diff_method)
         def circuit2(data, weights):
@@ -1562,7 +1562,7 @@ def circuit(n, a):
 
 def test_adjoint_reuse_device_state(mocker):
     """Tests that the autograd interface reuses the device state for adjoint differentiation"""
-    dev = qml.device("default.qubit", wires=1)
+    dev = qml.device("default.qubit.legacy", wires=1)
 
     @qnode(dev, diff_method="adjoint")
     def circ(x):
@@ -1773,7 +1773,7 @@ class TestSample:
 
     def test_backprop_error(self):
         """Test that sampling in backpropagation grad_on_execution raises an error"""
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
 
         @qnode(dev, diff_method="backprop")
         def circuit():
@@ -1787,7 +1787,7 @@ def test_sample_dimension(self):
         """Test that the sample function outputs samples of the right size"""
         n_sample = 10
 
-        dev = qml.device("default.qubit", wires=2, shots=n_sample)
+        dev = qml.device("default.qubit.legacy", wires=2, shots=n_sample)
 
         @qnode(dev)
         def circuit():
@@ -1810,7 +1810,7 @@ def test_sample_combination(self):
 
         n_sample = 10
 
-        dev = qml.device("default.qubit", wires=3, shots=n_sample)
+        dev = qml.device("default.qubit.legacy", wires=3, shots=n_sample)
 
         @qnode(dev, diff_method="parameter-shift")
         def circuit():
@@ -1832,7 +1832,7 @@ def test_single_wire_sample(self):
         """Test the return type and shape of sampling a single wire"""
         n_sample = 10
 
-        dev = qml.device("default.qubit", wires=1, shots=n_sample)
+        dev = qml.device("default.qubit.legacy", wires=1, shots=n_sample)
 
         @qnode(dev)
         def circuit():
@@ -1850,7 +1850,7 @@ def test_multi_wire_sample_regular_shape(self):
         where a rectangular array is expected"""
         n_sample = 10
 
-        dev = qml.device("default.qubit", wires=3, shots=n_sample)
+        dev = qml.device("default.qubit.legacy", wires=3, shots=n_sample)
 
         @qnode(dev)
         def circuit():
@@ -2388,7 +2388,7 @@ def cost(x):
         assert hess.shape == (3, 2, 2)
 
 
-@pytest.mark.parametrize("dev_name", ["default.qubit", "default.mixed"])
+@pytest.mark.parametrize("dev_name", ["default.qubit.legacy", "default.mixed"])
 def test_no_ops(dev_name):
     """Test that the return value of the QNode matches in the interface
     even if there are no ops"""
diff --git a/tests/interfaces/test_autograd_qnode_shot_vector.py b/tests/interfaces/test_autograd_qnode_shot_vector.py
index 1a835aba39e..fb0ba04eaf0 100644
--- a/tests/interfaces/test_autograd_qnode_shot_vector.py
+++ b/tests/interfaces/test_autograd_qnode_shot_vector.py
@@ -29,9 +29,9 @@
 spsa_kwargs = {"h": 0.05, "num_directions": 20, "sampler_rng": np.random.default_rng(SEED_FOR_SPSA)}
 
 qubit_device_and_diff_method = [
-    ["default.qubit", "finite-diff", {"h": 0.05}],
-    ["default.qubit", "parameter-shift", {}],
-    ["default.qubit", "spsa", spsa_kwargs],
+    ["default.qubit.legacy", "finite-diff", {"h": 0.05}],
+    ["default.qubit.legacy", "parameter-shift", {}],
+    ["default.qubit.legacy", "spsa", spsa_kwargs],
 ]
 
 TOLS = {
diff --git a/tests/interfaces/test_jax.py b/tests/interfaces/test_jax.py
index b726d810fd5..9f519b733fe 100644
--- a/tests/interfaces/test_jax.py
+++ b/tests/interfaces/test_jax.py
@@ -35,7 +35,7 @@ def test_jacobian_options(self, mocker):
 
         a = jax.numpy.array([0.1, 0.2])
 
-        dev = qml.device("default.qubit", wires=1)
+        dev = qml.device("default.qubit.legacy", wires=1)
 
         def cost(a, device):
             with qml.queuing.AnnotatedQueue() as q:
@@ -61,7 +61,7 @@ def test_incorrect_grad_on_execution(self):
         is used with grad_on_execution=True"""
         a = jax.numpy.array([0.1, 0.2])
 
-        dev = qml.device("default.qubit", wires=1)
+        dev = qml.device("default.qubit.legacy", wires=1)
 
         def cost(a, device):
             with qml.queuing.AnnotatedQueue() as q:
@@ -86,7 +86,7 @@ def test_unknown_interface(self):
         """Test that an error is raised if the interface is unknown"""
         a = jax.numpy.array([0.1, 0.2])
 
-        dev = qml.device("default.qubit", wires=1)
+        dev = qml.device("default.qubit.legacy", wires=1)
 
         def cost(a, device):
             with qml.queuing.AnnotatedQueue() as q:
@@ -107,7 +107,7 @@ def cost(a, device):
 
     def test_grad_on_execution(self, mocker):
         """Test that grad_on_execution uses the `device.execute_and_gradients` pathway"""
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
         spy = mocker.spy(dev, "execute_and_gradients")
 
         def cost(params):
@@ -150,7 +150,7 @@ def cost(params):
 
     def test_no_grad_on_execution(self, mocker):
         """Test that no grad on execution uses the `device.batch_execute` and `device.gradients` pathway"""
-        dev = qml.device("default.qubit", wires=1)
+        dev = qml.device("default.qubit.legacy", wires=1)
         spy_execute = mocker.spy(qml.devices.DefaultQubit, "batch_execute")
         spy_gradients = mocker.spy(qml.devices.DefaultQubit, "gradients")
 
@@ -185,7 +185,7 @@ class TestCaching:
 
     def test_cache_maxsize(self, mocker):
         """Test the cachesize property of the cache"""
-        dev = qml.device("default.qubit", wires=1)
+        dev = qml.device("default.qubit.legacy", wires=1)
         spy = mocker.spy(qml.interfaces, "cache_execute")
 
         def cost(a, cachesize):
@@ -212,7 +212,7 @@ def cost(a, cachesize):
 
     def test_custom_cache(self, mocker):
         """Test the use of a custom cache object"""
-        dev = qml.device("default.qubit", wires=1)
+        dev = qml.device("default.qubit.legacy", wires=1)
         spy = mocker.spy(qml.interfaces, "cache_execute")
 
         def cost(a, cache):
@@ -238,7 +238,7 @@ def cost(a, cache):
 
     def test_custom_cache_multiple(self, mocker):
         """Test the use of a custom cache object with multiple tapes"""
-        dev = qml.device("default.qubit", wires=1)
+        dev = qml.device("default.qubit.legacy", wires=1)
         spy = mocker.spy(qml.interfaces, "cache_execute")
 
         a = jax.numpy.array(0.1)
@@ -274,7 +274,7 @@ def cost(a, b, cache):
     def test_caching_param_shift(self, tol):
         """Test that, when using parameter-shift transform,
         caching produces the optimum number of evaluations."""
-        dev = qml.device("default.qubit", wires=1)
+        dev = qml.device("default.qubit.legacy", wires=1)
 
         def cost(a, cache):
             with qml.queuing.AnnotatedQueue() as q:
@@ -319,7 +319,7 @@ def cost(a, cache):
     def test_caching_adjoint_backward(self):
         """Test that caching produces the optimum number of adjoint evaluations
         when no grad on execution."""
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
         params = jax.numpy.array([0.1, 0.2, 0.3])
 
         def cost(a, cache):
@@ -375,7 +375,7 @@ class TestJaxExecuteIntegration:
 
     def test_execution(self, execute_kwargs):
         """Test execution"""
-        dev = qml.device("default.qubit", wires=1)
+        dev = qml.device("default.qubit.legacy", wires=1)
 
         def cost(a, b):
             with qml.queuing.AnnotatedQueue() as q1:
@@ -404,7 +404,7 @@ def cost(a, b):
     def test_scalar_jacobian(self, execute_kwargs, tol):
         """Test scalar jacobian calculation"""
         a = jax.numpy.array(0.1)
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
 
         def cost(a):
             with qml.queuing.AnnotatedQueue() as q:
@@ -436,7 +436,7 @@ def test_reusing_quantum_tape(self, execute_kwargs, tol):
         a = jax.numpy.array(0.1)
         b = jax.numpy.array(0.2)
 
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
 
         with qml.queuing.AnnotatedQueue() as q:
             qml.RY(a, wires=0)
@@ -477,7 +477,7 @@ def cost(a, b):
 
     def test_grad_with_different_grad_on_execution(self, execute_kwargs):
         """Test jax grad for adjoint diff method with different execution kwargs."""
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
         params = jax.numpy.array([0.1, 0.2, 0.3])
         expected_results = jax.numpy.array([-0.3875172, -0.18884787, -0.38355705])
 
@@ -513,14 +513,14 @@ def cost(a, b, c, device):
 
             return execute([tape], device, **execute_kwargs)[0]
 
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
         res = jax.grad(cost, argnums=(0, 1, 2))(a, b, c, device=dev)
         assert len(res) == 3
 
     def test_classical_processing_multiple_tapes(self, execute_kwargs):
         """Test classical processing within the quantum tape for multiple
         tapes"""
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
         params = jax.numpy.array([0.3, 0.2])
 
         def cost_fn(x):
@@ -546,7 +546,7 @@ def cost_fn(x):
 
     def test_multiple_tapes_output(self, execute_kwargs):
         """Test the output types for the execution of multiple quantum tapes"""
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
         params = jax.numpy.array([0.3, 0.2])
 
         def cost_fn(x):
@@ -588,7 +588,7 @@ def cost(a, U, device):
             tape.trainable_params = [0]
             return execute([tape], device, **execute_kwargs)[0]
 
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
         res = cost(a, U, device=dev)
         assert np.allclose(res, -np.cos(a), atol=tol, rtol=0)
 
@@ -622,7 +622,7 @@ def cost_fn(a, p, device):
         a = jax.numpy.array(0.1)
         p = jax.numpy.array([0.1, 0.2, 0.3])
 
-        dev = qml.device("default.qubit", wires=1)
+        dev = qml.device("default.qubit.legacy", wires=1)
         res = cost_fn(a, p, device=dev)
         expected = np.cos(a) * np.cos(p[1]) * np.sin(p[0]) + np.sin(a) * (
             np.cos(p[2]) * np.sin(p[1]) + np.cos(p[0]) * np.cos(p[1]) * np.sin(p[2])
@@ -646,7 +646,7 @@ def cost_fn(a, p, device):
     def test_independent_expval(self, execute_kwargs):
         """Tests computing an expectation value that is independent of trainable
         parameters."""
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
         params = jax.numpy.array([0.1, 0.2, 0.3])
 
         def cost(a, cache):
@@ -672,7 +672,7 @@ class TestVectorValued:
     def test_multiple_expvals(self, execute_kwargs):
         """Tests computing multiple expectation values in a tape."""
 
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
         params = jax.numpy.array([0.1, 0.2, 0.3])
 
         def cost(a, cache):
@@ -701,7 +701,7 @@ def cost(a, cache):
     def test_multiple_expvals_single_par(self, execute_kwargs):
         """Tests computing multiple expectation values in a tape with a single
         trainable parameter."""
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
         params = jax.numpy.array([0.1])
 
         def cost(a, cache):
@@ -727,7 +727,7 @@ def cost(a, cache):
     def test_multi_tape_fwd(self, execute_kwargs):
         """Test the forward evaluation of a cost function that uses the output
         of multiple tapes that be vector-valued."""
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
         params = jax.numpy.array([0.3, 0.2])
 
         def cost_fn(x):
@@ -773,7 +773,7 @@ def cost(x, y, device, interface, ek):
             tape2 = qml.tape.QuantumScript.from_queue(q2)
             return qml.execute([tape1, tape2], device, **ek, interface=interface)
 
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
         x = jax.numpy.array(0.543)
         y = jax.numpy.array(-0.654)
 
@@ -830,7 +830,7 @@ def cost(x, y, device, interface, ek):
             tape2 = qml.tape.QuantumScript.from_queue(q2)
             return qml.execute([tape1, tape2], device, **ek, interface=interface)
 
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
         x = jax.numpy.array(0.543)
         y = jax.numpy.array(-0.654)
 
diff --git a/tests/interfaces/test_jax_jit.py b/tests/interfaces/test_jax_jit.py
index a9f9e494757..5e375f8e912 100644
--- a/tests/interfaces/test_jax_jit.py
+++ b/tests/interfaces/test_jax_jit.py
@@ -36,7 +36,7 @@ def test_jacobian_options(self, mocker):
 
         a = jax.numpy.array([0.1, 0.2])
 
-        dev = qml.device("default.qubit", wires=1)
+        dev = qml.device("default.qubit.legacy", wires=1)
 
         def cost(a, device):
             with qml.queuing.AnnotatedQueue() as q:
@@ -63,7 +63,7 @@ def test_incorrect_gradients_on_execution(self):
         is used with grad_on_execution=True"""
         a = jax.numpy.array([0.1, 0.2])
 
-        dev = qml.device("default.qubit", wires=1)
+        dev = qml.device("default.qubit.legacy", wires=1)
 
         def cost(a, device):
             with qml.queuing.AnnotatedQueue() as q:
@@ -89,7 +89,7 @@ def test_unknown_interface(self):
         """Test that an error is raised if the interface is unknown"""
         a = jax.numpy.array([0.1, 0.2])
 
-        dev = qml.device("default.qubit", wires=1)
+        dev = qml.device("default.qubit.legacy", wires=1)
 
         def cost(a, device):
             with qml.queuing.AnnotatedQueue() as q:
@@ -111,7 +111,7 @@ def cost(a, device):
 
     def test_grad_on_execution(self, mocker):
         """Test that grad on execution uses the `device.execute_and_gradients` pathway"""
-        dev = qml.device("default.qubit", wires=1)
+        dev = qml.device("default.qubit.legacy", wires=1)
         spy = mocker.spy(dev, "execute_and_gradients")
 
         def cost(a):
@@ -141,7 +141,7 @@ def cost(a):
 
     def test_no_gradients_on_execution(self, mocker):
         """Test that no grad on execution uses the `device.batch_execute` and `device.gradients` pathway"""
-        dev = qml.device("default.qubit", wires=1)
+        dev = qml.device("default.qubit.legacy", wires=1)
         spy_execute = mocker.spy(qml.devices.DefaultQubit, "batch_execute")
         spy_gradients = mocker.spy(qml.devices.DefaultQubit, "gradients")
 
@@ -177,7 +177,7 @@ class TestCaching:
 
     def test_cache_maxsize(self, mocker):
         """Test the cachesize property of the cache"""
-        dev = qml.device("default.qubit", wires=1)
+        dev = qml.device("default.qubit.legacy", wires=1)
         spy = mocker.spy(qml.interfaces, "cache_execute")
 
         def cost(a, cachesize):
@@ -205,7 +205,7 @@ def cost(a, cachesize):
 
     def test_custom_cache(self, mocker):
         """Test the use of a custom cache object"""
-        dev = qml.device("default.qubit", wires=1)
+        dev = qml.device("default.qubit.legacy", wires=1)
         spy = mocker.spy(qml.interfaces, "cache_execute")
 
         def cost(a, cache):
@@ -232,7 +232,7 @@ def cost(a, cache):
 
     def test_custom_cache_multiple(self, mocker):
         """Test the use of a custom cache object with multiple tapes"""
-        dev = qml.device("default.qubit", wires=1)
+        dev = qml.device("default.qubit.legacy", wires=1)
         spy = mocker.spy(qml.interfaces, "cache_execute")
 
         a = jax.numpy.array(0.1)
@@ -270,7 +270,7 @@ def cost(a, b, cache):
     def test_caching_param_shift(self, tol):
         """Test that, when using parameter-shift transform,
         caching produces the optimum number of evaluations."""
-        dev = qml.device("default.qubit", wires=1)
+        dev = qml.device("default.qubit.legacy", wires=1)
 
         def cost(a, cache):
             with qml.queuing.AnnotatedQueue() as q:
@@ -316,7 +316,7 @@ def cost(a, cache):
     def test_caching_adjoint_backward(self):
         """Test that caching produces the optimum number of adjoint evaluations
         when mode=backward"""
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
         params = jax.numpy.array([0.1, 0.2, 0.3])
 
         def cost(a, cache):
@@ -373,7 +373,7 @@ class TestJaxExecuteIntegration:
 
     def test_execution(self, execute_kwargs):
         """Test execution"""
-        dev = qml.device("default.qubit", wires=1)
+        dev = qml.device("default.qubit.legacy", wires=1)
 
         def cost(a, b):
             with qml.queuing.AnnotatedQueue() as q1:
@@ -403,7 +403,7 @@ def cost(a, b):
     def test_scalar_jacobian(self, execute_kwargs, tol):
         """Test scalar jacobian calculation"""
         a = jax.numpy.array(0.1)
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
 
         def cost(a):
             with qml.queuing.AnnotatedQueue() as q:
@@ -436,7 +436,7 @@ def test_reusing_quantum_tape(self, execute_kwargs, tol):
         a = jax.numpy.array(0.1)
         b = jax.numpy.array(0.2)
 
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
 
         with qml.queuing.AnnotatedQueue() as q:
             qml.RY(a, wires=0)
@@ -478,7 +478,7 @@ def cost(a, b):
 
     def test_grad_with_backward_mode(self, execute_kwargs):
         """Test jax grad for adjoint diff method in backward mode"""
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
         params = jax.numpy.array([0.1, 0.2, 0.3])
         expected_results = jax.numpy.array([-0.3875172, -0.18884787, -0.38355705])
 
@@ -517,14 +517,14 @@ def cost(a, b, c, device):
 
             return execute([tape], device, **execute_kwargs)[0]
 
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
         res = jax.jit(jax.grad(cost, argnums=(0, 1, 2)), static_argnums=3)(a, b, c, device=dev)
         assert len(res) == 3
 
     def test_classical_processing_multiple_tapes(self, execute_kwargs):
         """Test classical processing within the quantum tape for multiple
         tapes"""
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
         params = jax.numpy.array([0.3, 0.2])
 
         def cost_fn(x):
@@ -552,7 +552,7 @@ def cost_fn(x):
 
     def test_multiple_tapes_output(self, execute_kwargs):
         """Test the output types for the execution of multiple quantum tapes"""
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
         params = jax.numpy.array([0.3, 0.2])
 
         def cost_fn(x):
@@ -596,7 +596,7 @@ def cost(a, U, device):
             tape.trainable_params = [0]
             return execute([tape], device, **execute_kwargs)[0]
 
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
         res = jax.jit(cost, static_argnums=2)(a, U, device=dev)
         assert np.allclose(res, -np.cos(a), atol=tol, rtol=0)
 
@@ -627,7 +627,7 @@ def cost_fn(a, p, device):
         a = jax.numpy.array(0.1)
         p = jax.numpy.array([0.1, 0.2, 0.3])
 
-        dev = qml.device("default.qubit", wires=1)
+        dev = qml.device("default.qubit.legacy", wires=1)
         res = jax.jit(cost_fn, static_argnums=2)(a, p, device=dev)
         expected = np.cos(a) * np.cos(p[1]) * np.sin(p[0]) + np.sin(a) * (
             np.cos(p[2]) * np.sin(p[1]) + np.cos(p[0]) * np.cos(p[1]) * np.sin(p[2])
@@ -651,7 +651,7 @@ def cost_fn(a, p, device):
     def test_independent_expval(self, execute_kwargs):
         """Tests computing an expectation value that is independent of trainable
         parameters."""
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
         params = jax.numpy.array([0.1, 0.2, 0.3])
 
         def cost(a, cache):
@@ -687,7 +687,7 @@ def test_shapes(self, execute_kwargs, ret_type, shape, expected_type):
         if adjoint:
             pytest.skip("The adjoint diff method doesn't support probabilities.")
 
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
         params = jax.numpy.array([0.1, 0.2, 0.3])
 
         def cost(a, cache):
@@ -715,7 +715,7 @@ def cost(a, cache):
     def test_independent_expval(self, execute_kwargs):
         """Tests computing an expectation value that is independent of trainable
         parameters."""
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
         params = jax.numpy.array([0.1, 0.2, 0.3])
 
         def cost(a, cache):
@@ -747,7 +747,7 @@ def cost(a, cache):
     def test_vector_valued_qnode(self, execute_kwargs, ret, out_dim, expected_type):
         """Tests the shape of vector-valued QNode results."""
 
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
         params = jax.numpy.array([0.1, 0.2, 0.3])
         grad_meth = (
             execute_kwargs["gradient_kwargs"]["method"]
@@ -786,7 +786,7 @@ def cost(a, cache):
 
     def test_qnode_sample(self, execute_kwargs):
         """Tests computing multiple expectation values in a tape."""
-        dev = qml.device("default.qubit", wires=2, shots=10)
+        dev = qml.device("default.qubit.legacy", wires=2, shots=10)
         params = jax.numpy.array([0.1, 0.2, 0.3])
 
         grad_meth = (
@@ -814,7 +814,7 @@ def cost(a, cache):
 
     def test_multiple_expvals_grad(self, execute_kwargs):
         """Tests computing multiple expectation values in a tape."""
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
         params = jax.numpy.array([0.1, 0.2, 0.3])
 
         def cost(a, cache):
@@ -861,7 +861,7 @@ def cost(x, y, device, interface, ek):
 
             return qml.execute([tape1, tape2], device, **ek, interface=interface)[0]
 
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
         x = jax.numpy.array(0.543)
         y = jax.numpy.array(-0.654)
 
diff --git a/tests/interfaces/test_jax_jit_qnode.py b/tests/interfaces/test_jax_jit_qnode.py
index 01dcaa6c52b..6fb29cec859 100644
--- a/tests/interfaces/test_jax_jit_qnode.py
+++ b/tests/interfaces/test_jax_jit_qnode.py
@@ -20,13 +20,13 @@
 from pennylane import qnode
 
 qubit_device_and_diff_method = [
-    ["default.qubit", "backprop", True],
-    ["default.qubit", "finite-diff", False],
-    ["default.qubit", "parameter-shift", False],
-    ["default.qubit", "adjoint", True],
-    ["default.qubit", "adjoint", False],
-    ["default.qubit", "spsa", False],
-    ["default.qubit", "hadamard", False],
+    ["default.qubit.legacy", "backprop", True],
+    ["default.qubit.legacy", "finite-diff", False],
+    ["default.qubit.legacy", "parameter-shift", False],
+    ["default.qubit.legacy", "adjoint", True],
+    ["default.qubit.legacy", "adjoint", False],
+    ["default.qubit.legacy", "spsa", False],
+    ["default.qubit.legacy", "hadamard", False],
 ]
 interface_and_qubit_device_and_diff_method = [
     ["auto"] + inner_list for inner_list in qubit_device_and_diff_method
@@ -818,7 +818,7 @@ def circuit(x):
 
     def test_changing_shots(self, interface, mocker, tol):
         """Test that changing shots works on execution"""
-        dev = qml.device("default.qubit", wires=2, shots=None)
+        dev = qml.device("default.qubit.legacy", wires=2, shots=None)
         a, b = jax.numpy.array([0.543, -0.654])
 
         @qnode(dev, diff_method=qml.gradients.param_shift, interface=interface)
@@ -849,7 +849,7 @@ def circuit(a, b):
     def test_gradient_integration(self, interface):
         """Test that temporarily setting the shots works
         for gradient computations"""
-        dev = qml.device("default.qubit", wires=2, shots=1)
+        dev = qml.device("default.qubit.legacy", wires=2, shots=1)
         a, b = jax.numpy.array([0.543, -0.654])
 
         @qnode(dev, diff_method=qml.gradients.param_shift, interface=interface)
@@ -869,7 +869,7 @@ def cost_fn(a, b):
     def test_update_diff_method(self, mocker, interface):
         """Test that temporarily setting the shots updates the diff method"""
         # pylint: disable=unused-argument
-        dev = qml.device("default.qubit", wires=2, shots=100)
+        dev = qml.device("default.qubit.legacy", wires=2, shots=100)
         a, b = jax.numpy.array([0.543, -0.654])
 
         spy = mocker.spy(qml, "execute")
@@ -1504,7 +1504,7 @@ def circuit(n, a):
 @pytest.mark.parametrize("interface", ["auto", "jax-jit"])
 def test_adjoint_reuse_device_state(mocker, interface):
     """Tests that the jax interface reuses the device state for adjoint differentiation"""
-    dev = qml.device("default.qubit", wires=1)
+    dev = qml.device("default.qubit.legacy", wires=1)
 
     @qnode(dev, interface=interface, diff_method="adjoint")
     def circ(x):
diff --git a/tests/interfaces/test_jax_qnode.py b/tests/interfaces/test_jax_qnode.py
index 2a55910e10e..91260e2d90c 100644
--- a/tests/interfaces/test_jax_qnode.py
+++ b/tests/interfaces/test_jax_qnode.py
@@ -21,13 +21,13 @@
 from pennylane.tape import QuantumScript
 
 qubit_device_and_diff_method = [
-    ["default.qubit", "backprop", True],
-    ["default.qubit", "finite-diff", False],
-    ["default.qubit", "parameter-shift", False],
-    ["default.qubit", "adjoint", True],
-    ["default.qubit", "adjoint", False],
-    ["default.qubit", "spsa", False],
-    ["default.qubit", "hadamard", False],
+    ["default.qubit.legacy", "backprop", True],
+    ["default.qubit.legacy", "finite-diff", False],
+    ["default.qubit.legacy", "parameter-shift", False],
+    ["default.qubit.legacy", "adjoint", True],
+    ["default.qubit.legacy", "adjoint", False],
+    ["default.qubit.legacy", "spsa", False],
+    ["default.qubit.legacy", "hadamard", False],
 ]
 
 interface_and_qubit_device_and_diff_method = [
@@ -782,7 +782,7 @@ def circuit(x):
 
     def test_changing_shots(self, interface, mocker, tol):
         """Test that changing shots works on execution"""
-        dev = qml.device("default.qubit", wires=2, shots=None)
+        dev = qml.device("default.qubit.legacy", wires=2, shots=None)
         a, b = jax.numpy.array([0.543, -0.654])
 
         @qnode(dev, diff_method=qml.gradients.param_shift, interface=interface)
@@ -813,7 +813,7 @@ def circuit(a, b):
     def test_gradient_integration(self, interface):
         """Test that temporarily setting the shots works
         for gradient computations"""
-        dev = qml.device("default.qubit", wires=2, shots=1)
+        dev = qml.device("default.qubit.legacy", wires=2, shots=1)
         a, b = jax.numpy.array([0.543, -0.654])
 
         @qnode(dev, diff_method=qml.gradients.param_shift, interface=interface)
@@ -832,7 +832,7 @@ def cost_fn(a, b):
     def test_update_diff_method(self, mocker, interface):
         """Test that temporarily setting the shots updates the diff method"""
         # pylint: disable=unused-argument
-        dev = qml.device("default.qubit", wires=2, shots=100)
+        dev = qml.device("default.qubit.legacy", wires=2, shots=100)
         a, b = jax.numpy.array([0.543, -0.654])
 
         spy = mocker.spy(qml, "execute")
@@ -1448,7 +1448,7 @@ def circuit(n, a):
 @pytest.mark.parametrize("interface", ["auto", "jax", "jax-python"])
 def test_adjoint_reuse_device_state(mocker, interface):
     """Tests that the jax interface reuses the device state for adjoint differentiation"""
-    dev = qml.device("default.qubit", wires=1)
+    dev = qml.device("default.qubit.legacy", wires=1)
 
     @qnode(dev, interface=interface, diff_method="adjoint")
     def circ(x):
@@ -2549,7 +2549,7 @@ def circuit(x):
         assert hess[1].shape == (2, 2, 2)
 
 
-@pytest.mark.parametrize("dev_name", ["default.qubit", "default.mixed"])
+@pytest.mark.parametrize("dev_name", ["default.qubit.legacy", "default.mixed"])
 def test_no_ops(dev_name):
     """Test that the return value of the QNode matches in the interface
     even if there are no ops"""
diff --git a/tests/interfaces/test_jax_qnode_shot_vector.py b/tests/interfaces/test_jax_qnode_shot_vector.py
index 5fd6356c314..6645e008c31 100644
--- a/tests/interfaces/test_jax_qnode_shot_vector.py
+++ b/tests/interfaces/test_jax_qnode_shot_vector.py
@@ -29,9 +29,9 @@
 all_shots = [(1, 20, 100), (1, (20, 1), 100), (1, (5, 4), 100)]
 
 qubit_device_and_diff_method = [
-    ["default.qubit", "finite-diff", {"h": 10e-2}],
-    ["default.qubit", "parameter-shift", {}],
-    ["default.qubit", "spsa", {"h": 10e-2, "num_directions": 20}],
+    ["default.qubit.legacy", "finite-diff", {"h": 10e-2}],
+    ["default.qubit.legacy", "parameter-shift", {}],
+    ["default.qubit.legacy", "spsa", {"h": 10e-2, "num_directions": 20}],
 ]
 
 interface_and_qubit_device_and_diff_method = [
@@ -773,7 +773,7 @@ class TestReturnShotVectorsDevice:
 
     def test_jac_adjoint_fwd_error(self, shots):
         """Test that an error is raised for adjoint forward."""
-        dev = qml.device("default.qubit", wires=1, shots=shots)
+        dev = qml.device("default.qubit.legacy", wires=1, shots=shots)
 
         with pytest.warns(
             UserWarning, match="Requested adjoint differentiation to be computed with finite shots."
@@ -796,7 +796,7 @@ def circuit(a):
 
     def test_jac_adjoint_bwd_error(self, shots):
         """Test that an error is raised for adjoint backward."""
-        dev = qml.device("default.qubit", wires=1, shots=shots)
+        dev = qml.device("default.qubit.legacy", wires=1, shots=shots)
 
         with pytest.warns(
             UserWarning, match="Requested adjoint differentiation to be computed with finite shots."
@@ -818,8 +818,8 @@ def circuit(a):
 
 
 qubit_device_and_diff_method = [
-    ["default.qubit", "finite-diff", {"h": 10e-2}],
-    ["default.qubit", "parameter-shift", {}],
+    ["default.qubit.legacy", "finite-diff", {"h": 10e-2}],
+    ["default.qubit.legacy", "parameter-shift", {}],
 ]
 
 shots_large = [(1000000, 900000, 800000), (1000000, (900000, 2))]
diff --git a/tests/interfaces/test_tensorflow.py b/tests/interfaces/test_tensorflow.py
index f5e26461799..d0ce671989c 100644
--- a/tests/interfaces/test_tensorflow.py
+++ b/tests/interfaces/test_tensorflow.py
@@ -34,7 +34,7 @@ def test_jacobian_options(self, mocker):
 
         a = tf.Variable([0.1, 0.2], dtype=tf.float64)
 
-        dev = qml.device("default.qubit", wires=1)
+        dev = qml.device("default.qubit.legacy", wires=1)
 
         with tf.GradientTape() as t:
             with qml.queuing.AnnotatedQueue() as q:
@@ -61,7 +61,7 @@ def test_incorrect_grad_on_execution(self):
         is used with grad_on_execution"""
         a = tf.Variable([0.1, 0.2])
 
-        dev = qml.device("default.qubit", wires=1)
+        dev = qml.device("default.qubit.legacy", wires=1)
 
         with tf.GradientTape():
             with qml.queuing.AnnotatedQueue() as q:
@@ -77,7 +77,7 @@ def test_incorrect_grad_on_execution(self):
 
     def test_grad_on_execution(self, mocker):
         """Test that grad on execution uses the `device.execute_and_gradients` pathway"""
-        dev = qml.device("default.qubit", wires=1)
+        dev = qml.device("default.qubit.legacy", wires=1)
         a = tf.Variable([0.1, 0.2])
         spy = mocker.spy(dev, "execute_and_gradients")
 
@@ -102,7 +102,7 @@ def test_grad_on_execution(self, mocker):
 
     def test_no_grad_execution(self, mocker):
         """Test that no grad on execution uses the `device.batch_execute` and `device.gradients` pathway"""
-        dev = qml.device("default.qubit", wires=1)
+        dev = qml.device("default.qubit.legacy", wires=1)
         spy_execute = mocker.spy(qml.devices.DefaultQubit, "batch_execute")
         spy_gradients = mocker.spy(qml.devices.DefaultQubit, "gradients")
         a = tf.Variable([0.1, 0.2])
@@ -136,7 +136,7 @@ class TestCaching:
 
     def test_cache_maxsize(self, mocker):
         """Test the cachesize property of the cache"""
-        dev = qml.device("default.qubit", wires=1)
+        dev = qml.device("default.qubit.legacy", wires=1)
         spy = mocker.spy(qml.interfaces, "cache_execute")
         a = tf.Variable([0.1, 0.2])
 
@@ -158,7 +158,7 @@ def test_cache_maxsize(self, mocker):
 
     def test_custom_cache(self, mocker):
         """Test the use of a custom cache object"""
-        dev = qml.device("default.qubit", wires=1)
+        dev = qml.device("default.qubit.legacy", wires=1)
         spy = mocker.spy(qml.interfaces, "cache_execute")
         a = tf.Variable([0.1, 0.2])
         custom_cache = {}
@@ -188,7 +188,7 @@ def test_custom_cache(self, mocker):
     def test_caching_param_shift(self):
         """Test that, when using parameter-shift transform,
         caching reduces the number of evaluations to their optimum."""
-        dev = qml.device("default.qubit", wires=1)
+        dev = qml.device("default.qubit.legacy", wires=1)
         a = tf.Variable([0.1, 0.2], dtype=tf.float64)
 
         def cost(a, cache):
@@ -228,7 +228,7 @@ def test_caching_param_shift_hessian(self, num_params, tol):
         """Test that, when using parameter-shift transform,
         caching reduces the number of evaluations to their optimum
         when computing Hessians."""
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
         params = tf.Variable(np.arange(1, num_params + 1) / 10, dtype=tf.float64)
 
         N = params.shape[0]
@@ -324,7 +324,7 @@ class TestTensorFlowExecuteIntegration:
 
     def test_execution(self, execute_kwargs):
         """Test execution"""
-        dev = qml.device("default.qubit", wires=1)
+        dev = qml.device("default.qubit.legacy", wires=1)
         a = tf.Variable(0.1)
         b = tf.Variable(0.2)
 
@@ -352,7 +352,7 @@ def test_execution(self, execute_kwargs):
     def test_scalar_jacobian(self, execute_kwargs, tol):
         """Test scalar jacobian calculation"""
         a = tf.Variable(0.1, dtype=tf.float64)
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
 
         with tf.GradientTape() as t:
             with qml.queuing.AnnotatedQueue() as q:
@@ -381,7 +381,7 @@ def test_jacobian(self, execute_kwargs, tol):
         """Test jacobian calculation"""
         a = tf.Variable(0.1, dtype=tf.float64)
         b = tf.Variable(0.2, dtype=tf.float64)
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
 
         with tf.GradientTape() as t:
             with qml.queuing.AnnotatedQueue() as q:
@@ -407,7 +407,7 @@ def test_jacobian(self, execute_kwargs, tol):
     def test_tape_no_parameters(self, execute_kwargs, tol):
         """Test that a tape with no parameters is correctly
         ignored during the gradient computation"""
-        dev = qml.device("default.qubit", wires=1)
+        dev = qml.device("default.qubit.legacy", wires=1)
         params = tf.Variable([0.1, 0.2], dtype=tf.float64)
         x, y = 1.0 * params
 
@@ -443,7 +443,7 @@ def test_reusing_quantum_tape(self, execute_kwargs, tol):
         a = tf.Variable(0.1, dtype=tf.float64)
         b = tf.Variable(0.2, dtype=tf.float64)
 
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
 
         with tf.GradientTape() as t:
             with qml.queuing.AnnotatedQueue() as q:
@@ -486,7 +486,7 @@ def test_reusing_pre_constructed_quantum_tape(self, execute_kwargs, tol):
         a = tf.Variable(0.1, dtype=tf.float64)
         b = tf.Variable(0.2, dtype=tf.float64)
 
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
 
         with qml.queuing.AnnotatedQueue() as q:
             qml.RY(a, wires=0)
@@ -528,7 +528,7 @@ def test_classical_processing(self, execute_kwargs):
         b = tf.constant(0.2, dtype=tf.float64)
         c = tf.Variable(0.3, dtype=tf.float64)
 
-        dev = qml.device("default.qubit", wires=1)
+        dev = qml.device("default.qubit.legacy", wires=1)
 
         with tf.GradientTape() as t:
             with qml.queuing.AnnotatedQueue() as q:
@@ -550,7 +550,7 @@ def test_classical_processing(self, execute_kwargs):
     def test_no_trainable_parameters(self, execute_kwargs):
         """Test evaluation and Jacobian if there are no trainable parameters"""
         b = tf.constant(0.2, dtype=tf.float64)
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
 
         with tf.GradientTape() as t:
             with qml.queuing.AnnotatedQueue() as q:
@@ -576,7 +576,7 @@ def test_matrix_parameter(self, execute_kwargs, U, tol):
         with a matrix parameter"""
         a = tf.Variable(0.1, dtype=tf.float64)
 
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
 
         with tf.GradientTape() as t:
             with qml.queuing.AnnotatedQueue() as q:
@@ -606,7 +606,7 @@ def decomposition(self):
                     qml.PhaseShift(phi + lam, wires=wires),
                 ]
 
-        dev = qml.device("default.qubit", wires=1)
+        dev = qml.device("default.qubit.legacy", wires=1)
         a = np.array(0.1)
         p = tf.Variable([0.1, 0.2, 0.3], dtype=tf.float64)
 
@@ -644,7 +644,7 @@ def test_probability_differentiation(self, execute_kwargs, tol):
         if execute_kwargs["gradient_fn"] == "device":
             pytest.skip("Adjoint differentiation does not yet support probabilities")
 
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
         x = tf.Variable(0.543, dtype=tf.float64)
         y = tf.Variable(-0.654, dtype=tf.float64)
 
@@ -689,7 +689,7 @@ def test_ragged_differentiation(self, execute_kwargs, tol):
         if execute_kwargs["gradient_fn"] == "device":
             pytest.skip("Adjoint differentiation does not yet support probabilities")
 
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
         x = tf.Variable(0.543, dtype=tf.float64)
         y = tf.Variable(-0.654, dtype=tf.float64)
 
@@ -727,7 +727,7 @@ def test_sampling(self, execute_kwargs):
         ):
             pytest.skip("Adjoint differentiation does not support samples")
 
-        dev = qml.device("default.qubit", wires=2, shots=10)
+        dev = qml.device("default.qubit.legacy", wires=2, shots=10)
 
         with tf.GradientTape():
             with qml.queuing.AnnotatedQueue() as q:
@@ -855,7 +855,7 @@ def test_hessian_vector_valued(self, tol, interface):
     def test_adjoint_hessian(self, interface):
         """Since the adjoint hessian is not a differentiable transform,
         higher-order derivatives are not supported."""
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
         params = tf.Variable([0.543, -0.654])
 
         with tf.GradientTape() as t2:
@@ -994,7 +994,7 @@ def test_multiple_hamiltonians_not_trainable(self, cost_fn, execute_kwargs, tol)
         weights = tf.Variable([0.4, 0.5], dtype=tf.float64)
         coeffs1 = tf.constant([0.1, 0.2, 0.3], dtype=tf.float64)
         coeffs2 = tf.constant([0.7], dtype=tf.float64)
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
 
         with tf.GradientTape() as tape:
             res = cost_fn(weights, coeffs1, coeffs2, dev=dev)
@@ -1013,7 +1013,7 @@ def test_multiple_hamiltonians_trainable(self, cost_fn, execute_kwargs, tol):
         weights = tf.Variable([0.4, 0.5], dtype=tf.float64)
         coeffs1 = tf.Variable([0.1, 0.2, 0.3], dtype=tf.float64)
         coeffs2 = tf.Variable([0.7], dtype=tf.float64)
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
 
         with tf.GradientTape() as tape:
             res = cost_fn(weights, coeffs1, coeffs2, dev=dev)
diff --git a/tests/interfaces/test_tensorflow_autograph_qnode_shot_vector.py b/tests/interfaces/test_tensorflow_autograph_qnode_shot_vector.py
index 05c6786c4b4..7500f21d5f6 100644
--- a/tests/interfaces/test_tensorflow_autograph_qnode_shot_vector.py
+++ b/tests/interfaces/test_tensorflow_autograph_qnode_shot_vector.py
@@ -28,9 +28,9 @@
 
 spsa_kwargs = {"h": 10e-2, "num_directions": 30, "sampler_rng": np.random.default_rng(42)}
 qubit_device_and_diff_method = [
-    ["default.qubit", "finite-diff", {"h": 10e-2}],
-    ["default.qubit", "parameter-shift", {}],
-    ["default.qubit", "spsa", spsa_kwargs],
+    ["default.qubit.legacy", "finite-diff", {"h": 10e-2}],
+    ["default.qubit.legacy", "parameter-shift", {}],
+    ["default.qubit.legacy", "spsa", spsa_kwargs],
 ]
 
 TOLS = {
diff --git a/tests/interfaces/test_tensorflow_qnode.py b/tests/interfaces/test_tensorflow_qnode.py
index 1905a87e578..7773c86d27d 100644
--- a/tests/interfaces/test_tensorflow_qnode.py
+++ b/tests/interfaces/test_tensorflow_qnode.py
@@ -24,13 +24,13 @@
 
 
 qubit_device_and_diff_method = [
-    ["default.qubit", "finite-diff", False],
-    ["default.qubit", "parameter-shift", False],
-    ["default.qubit", "backprop", True],
-    ["default.qubit", "adjoint", True],
-    ["default.qubit", "adjoint", False],
-    ["default.qubit", "spsa", False],
-    ["default.qubit", "hadamard", False],
+    ["default.qubit.legacy", "finite-diff", False],
+    ["default.qubit.legacy", "parameter-shift", False],
+    ["default.qubit.legacy", "backprop", True],
+    ["default.qubit.legacy", "adjoint", True],
+    ["default.qubit.legacy", "adjoint", False],
+    ["default.qubit.legacy", "spsa", False],
+    ["default.qubit.legacy", "hadamard", False],
 ]
 
 TOL_FOR_SPSA = 1.0
@@ -534,7 +534,7 @@ class TestShotsIntegration:
 
     def test_changing_shots(self, mocker, tol, interface):
         """Test that changing shots works on execution"""
-        dev = qml.device("default.qubit", wires=2, shots=None)
+        dev = qml.device("default.qubit.legacy", wires=2, shots=None)
         a, b = [0.543, -0.654]
         weights = tf.Variable([a, b], dtype=tf.float64)
 
@@ -569,7 +569,7 @@ def test_gradient_integration(self, interface):
         """Test that temporarily setting the shots works
         for gradient computations"""
         # pylint: disable=unexpected-keyword-arg
-        dev = qml.device("default.qubit", wires=2, shots=None)
+        dev = qml.device("default.qubit.legacy", wires=2, shots=None)
         a, b = [0.543, -0.654]
         weights = tf.Variable([a, b], dtype=tf.float64)
 
@@ -595,7 +595,7 @@ def test_multiple_gradient_integration(self, tol, interface):
         """Test that temporarily setting the shots works
         for gradient computations, even if the QNode has been re-evaluated
         with a different number of shots in the meantime."""
-        dev = qml.device("default.qubit", wires=2, shots=None)
+        dev = qml.device("default.qubit.legacy", wires=2, shots=None)
         a, b = [0.543, -0.654]
         weights = tf.Variable([a, b], dtype=tf.float64)
 
@@ -620,7 +620,7 @@ def circuit(weights):
 
     def test_update_diff_method(self, mocker, interface):
         """Test that temporarily setting the shots updates the diff method"""
-        dev = qml.device("default.qubit", wires=2, shots=100)
+        dev = qml.device("default.qubit.legacy", wires=2, shots=100)
         weights = tf.Variable([0.543, -0.654], dtype=tf.float64)
 
         spy = mocker.spy(qml, "execute")
@@ -656,7 +656,7 @@ class TestAdjoint:
 
     def test_reuse_state(self, mocker, interface):
         """Tests that the TF interface reuses the device state for adjoint differentiation"""
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
 
         @qnode(dev, diff_method="adjoint", interface=interface)
         def circ(x):
@@ -683,7 +683,7 @@ def circ(x):
     def test_resuse_state_multiple_evals(self, mocker, tol, interface):
         """Tests that the TF interface reuses the device state for adjoint differentiation,
         even where there are intermediate evaluations."""
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
 
         x_val = 0.543
         y_val = -0.654
@@ -1527,7 +1527,7 @@ class TestSample:
 
     def test_sample_dimension(self):
         """Test sampling works as expected"""
-        dev = qml.device("default.qubit", wires=2, shots=10)
+        dev = qml.device("default.qubit.legacy", wires=2, shots=10)
 
         @qnode(dev, diff_method="parameter-shift", interface="tf")
         def circuit():
@@ -1547,7 +1547,7 @@ def circuit():
 
     def test_sampling_expval(self):
         """Test sampling works as expected if combined with expectation values"""
-        dev = qml.device("default.qubit", wires=2, shots=10)
+        dev = qml.device("default.qubit.legacy", wires=2, shots=10)
 
         @qnode(dev, diff_method="parameter-shift", interface="tf")
         def circuit():
@@ -1568,7 +1568,7 @@ def test_sample_combination(self):
         """Test the output of combining expval, var and sample"""
         n_sample = 10
 
-        dev = qml.device("default.qubit", wires=3, shots=n_sample)
+        dev = qml.device("default.qubit.legacy", wires=3, shots=n_sample)
 
         @qnode(dev, diff_method="parameter-shift", interface="tf")
         def circuit():
@@ -1593,7 +1593,7 @@ def test_single_wire_sample(self):
         """Test the return type and shape of sampling a single wire"""
         n_sample = 10
 
-        dev = qml.device("default.qubit", wires=1, shots=n_sample)
+        dev = qml.device("default.qubit.legacy", wires=1, shots=n_sample)
 
         @qnode(dev, diff_method="parameter-shift", interface="tf")
         def circuit():
@@ -1611,7 +1611,7 @@ def test_multi_wire_sample_regular_shape(self):
         where a rectangular array is expected"""
         n_sample = 10
 
-        dev = qml.device("default.qubit", wires=3, shots=n_sample)
+        dev = qml.device("default.qubit.legacy", wires=3, shots=n_sample)
 
         @qnode(dev, diff_method="parameter-shift", interface="tf")
         def circuit():
@@ -1627,7 +1627,7 @@ def circuit():
 
     def test_counts(self):
         """Test counts works as expected for TF"""
-        dev = qml.device("default.qubit", wires=2, shots=100)
+        dev = qml.device("default.qubit.legacy", wires=2, shots=100)
 
         @qnode(dev, interface="tf")
         def circuit():
@@ -1665,7 +1665,7 @@ class TestAutograph:
     def test_autograph_gradients(self, decorator, interface, tol):
         """Test that a parameter-shift QNode can be compiled
         using @tf.function, and differentiated"""
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
         x = tf.Variable(0.543, dtype=tf.float64)
         y = tf.Variable(-0.654, dtype=tf.float64)
 
@@ -1691,7 +1691,7 @@ def circuit(x, y):
     def test_autograph_jacobian(self, decorator, interface, tol):
         """Test that a parameter-shift vector-valued QNode can be compiled
         using @tf.function, and differentiated"""
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
         x = tf.Variable(0.543, dtype=tf.float64)
         y = tf.Variable(-0.654, dtype=tf.float64)
 
@@ -1734,7 +1734,7 @@ def circuit(x, y):
     def test_autograph_adjoint_single_param(self, grad_on_execution, decorator, interface, tol):
         """Test that a parameter-shift QNode can be compiled
         using @tf.function, and differentiated to second order"""
-        dev = qml.device("default.qubit", wires=1)
+        dev = qml.device("default.qubit.legacy", wires=1)
 
         @decorator
         @qnode(dev, diff_method="adjoint", interface=interface, grad_on_execution=grad_on_execution)
@@ -1758,7 +1758,7 @@ def circuit(x):
     def test_autograph_adjoint_multi_params(self, grad_on_execution, decorator, interface, tol):
         """Test that a parameter-shift QNode can be compiled
         using @tf.function, and differentiated to second order"""
-        dev = qml.device("default.qubit", wires=1)
+        dev = qml.device("default.qubit.legacy", wires=1)
 
         @decorator
         @qnode(dev, diff_method="adjoint", interface=interface, grad_on_execution=grad_on_execution)
@@ -1783,7 +1783,7 @@ def circuit(x):
     def test_autograph_ragged_differentiation(self, decorator, interface, tol):
         """Tests correct output shape and evaluation for a tape
         with prob and expval outputs"""
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
 
         x = tf.Variable(0.543, dtype=tf.float64)
         y = tf.Variable(-0.654, dtype=tf.float64)
@@ -1821,7 +1821,7 @@ def circuit(x, y):
     def test_autograph_hessian(self, decorator, interface, tol):
         """Test that a parameter-shift QNode can be compiled
         using @tf.function, and differentiated to second order"""
-        dev = qml.device("default.qubit", wires=1)
+        dev = qml.device("default.qubit.legacy", wires=1)
         a = tf.Variable(0.543, dtype=tf.float64)
         b = tf.Variable(-0.654, dtype=tf.float64)
 
@@ -1855,7 +1855,7 @@ def circuit(x, y):
     def test_autograph_state(self, decorator, interface, tol):
         """Test that a parameter-shift QNode returning a state can be compiled
         using @tf.function"""
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
         x = tf.Variable(0.543, dtype=tf.float64)
         y = tf.Variable(-0.654, dtype=tf.float64)
 
@@ -1878,7 +1878,7 @@ def circuit(x, y):
 
     def test_autograph_dimension(self, decorator, interface):
         """Test sampling works as expected"""
-        dev = qml.device("default.qubit", wires=2, shots=10)
+        dev = qml.device("default.qubit.legacy", wires=2, shots=10)
 
         @decorator
         @qnode(dev, diff_method="parameter-shift", interface=interface)
@@ -2552,7 +2552,7 @@ def circuit(x):
         assert hess.shape == (3, 2, 2)
 
 
-@pytest.mark.parametrize("dev_name", ["default.qubit", "default.mixed"])
+@pytest.mark.parametrize("dev_name", ["default.qubit.legacy", "default.mixed"])
 def test_no_ops(dev_name):
     """Test that the return value of the QNode matches in the interface
     even if there are no ops"""
diff --git a/tests/interfaces/test_tensorflow_qnode_shot_vector.py b/tests/interfaces/test_tensorflow_qnode_shot_vector.py
index 882bd89bb6c..113e4f04412 100644
--- a/tests/interfaces/test_tensorflow_qnode_shot_vector.py
+++ b/tests/interfaces/test_tensorflow_qnode_shot_vector.py
@@ -27,9 +27,9 @@
 shots_and_num_copies_hess = [((10, (5, 1)), 2)]
 
 qubit_device_and_diff_method = [
-    ["default.qubit", "finite-diff", {"h": 10e-2}],
-    ["default.qubit", "parameter-shift", {}],
-    ["default.qubit", "spsa", {"h": 10e-2, "num_directions": 20}],
+    ["default.qubit.legacy", "finite-diff", {"h": 10e-2}],
+    ["default.qubit.legacy", "parameter-shift", {}],
+    ["default.qubit.legacy", "spsa", {"h": 10e-2, "num_directions": 20}],
 ]
 
 TOLS = {
diff --git a/tests/interfaces/test_torch.py b/tests/interfaces/test_torch.py
index 84a0a08b7a7..5a2a0a93f84 100644
--- a/tests/interfaces/test_torch.py
+++ b/tests/interfaces/test_torch.py
@@ -36,7 +36,7 @@ def test_jacobian_options(self, interface, mocker):
 
         a = torch.tensor([0.1, 0.2], requires_grad=True)
 
-        dev = qml.device("default.qubit", wires=1)
+        dev = qml.device("default.qubit.legacy", wires=1)
 
         with qml.queuing.AnnotatedQueue() as q:
             qml.RY(a[0], wires=0)
@@ -63,7 +63,7 @@ def test_incorrect_grad_on_execution(self, interface):
         is used with grad_on_execution=True"""
         a = torch.tensor([0.1, 0.2], requires_grad=True)
 
-        dev = qml.device("default.qubit", wires=1)
+        dev = qml.device("default.qubit.legacy", wires=1)
 
         with qml.queuing.AnnotatedQueue() as q:
             qml.RY(a[0], wires=0)
@@ -82,7 +82,7 @@ def test_incorrect_grad_on_execution(self, interface):
     def test_grad_on_execution_reuse_state(self, interface, mocker):
         """Test that grad_on_execution uses the `device.execute_and_gradients` pathway
         while reusing the quantum state."""
-        dev = qml.device("default.qubit", wires=1)
+        dev = qml.device("default.qubit.legacy", wires=1)
         spy = mocker.spy(dev, "execute_and_gradients")
 
         a = torch.tensor([0.1, 0.2], requires_grad=True)
@@ -108,7 +108,7 @@ def test_grad_on_execution_reuse_state(self, interface, mocker):
 
     def test_grad_on_execution(self, interface, mocker):
         """Test that grad on execution uses the `device.execute_and_gradients` pathway"""
-        dev = qml.device("default.qubit", wires=1)
+        dev = qml.device("default.qubit.legacy", wires=1)
         spy = mocker.spy(dev, "execute_and_gradients")
 
         a = torch.tensor([0.1, 0.2], requires_grad=True)
@@ -134,7 +134,7 @@ def test_grad_on_execution(self, interface, mocker):
 
     def test_no_grad_on_execution(self, interface, mocker):
         """Test that no grad on execution uses the `device.batch_execute` and `device.gradients` pathway"""
-        dev = qml.device("default.qubit", wires=1)
+        dev = qml.device("default.qubit.legacy", wires=1)
         spy_execute = mocker.spy(qml.devices.DefaultQubit, "batch_execute")
         spy_gradients = mocker.spy(qml.devices.DefaultQubit, "gradients")
 
@@ -169,7 +169,7 @@ class TestCaching:
 
     def test_cache_maxsize(self, mocker):
         """Test the cachesize property of the cache"""
-        dev = qml.device("default.qubit", wires=1)
+        dev = qml.device("default.qubit.legacy", wires=1)
         spy = mocker.spy(qml.interfaces, "cache_execute")
 
         def cost(a, cachesize):
@@ -195,7 +195,7 @@ def cost(a, cachesize):
 
     def test_custom_cache(self, mocker):
         """Test the use of a custom cache object"""
-        dev = qml.device("default.qubit", wires=1)
+        dev = qml.device("default.qubit.legacy", wires=1)
         spy = mocker.spy(qml.interfaces, "cache_execute")
 
         def cost(a, cache):
@@ -221,7 +221,7 @@ def cost(a, cache):
     def test_caching_param_shift(self):
         """Test that, with the parameter-shift transform,
         Torch always uses the optimum number of evals when computing the Jacobian."""
-        dev = qml.device("default.qubit", wires=1)
+        dev = qml.device("default.qubit.legacy", wires=1)
 
         def cost(a, cache):
             with qml.queuing.AnnotatedQueue() as q:
@@ -252,7 +252,7 @@ def test_caching_param_shift_hessian(self, num_params, tol):
         """Test that, with the parameter-shift transform,
         caching reduces the number of evaluations to their optimum
         when computing Hessians."""
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
         params = torch.tensor(np.arange(1, num_params + 1) / 10, requires_grad=True)
 
         N = len(params)
@@ -309,7 +309,7 @@ def cost(x, cache):
     def test_caching_adjoint_no_grad_on_execution(self):
         """Test that caching reduces the number of adjoint evaluations
         when grad_on_execution=False"""
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
         params = torch.tensor([0.1, 0.2, 0.3])
 
         def cost(a, cache):
@@ -402,7 +402,7 @@ class TestTorchExecuteIntegration:
 
     def test_execution(self, torch_device, execute_kwargs):
         """Test that the execute function produces results with the expected shapes"""
-        dev = qml.device("default.qubit", wires=1)
+        dev = qml.device("default.qubit.legacy", wires=1)
         a = torch.tensor(0.1, requires_grad=True, device=torch_device)
         b = torch.tensor(0.2, requires_grad=False, device=torch_device)
 
@@ -430,7 +430,7 @@ def test_execution(self, torch_device, execute_kwargs):
     def test_scalar_jacobian(self, torch_device, execute_kwargs, tol):
         """Test scalar jacobian calculation by comparing two types of pipelines"""
         a = torch.tensor(0.1, requires_grad=True, dtype=torch.float64, device=torch_device)
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
 
         with qml.queuing.AnnotatedQueue() as q:
             qml.RY(a, wires=0)
@@ -463,7 +463,7 @@ def test_jacobian(self, torch_device, execute_kwargs, tol):
         a = torch.tensor(a_val, requires_grad=True, device=torch_device)
         b = torch.tensor(b_val, requires_grad=True, device=torch_device)
 
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
 
         with qml.queuing.AnnotatedQueue() as q:
             qml.RZ(torch.tensor(0.543, device=torch_device), wires=0)
@@ -506,7 +506,7 @@ def test_jacobian(self, torch_device, execute_kwargs, tol):
     def test_tape_no_parameters(self, torch_device, execute_kwargs, tol):
         """Test that a tape with no parameters is correctly
         ignored during the gradient computation"""
-        dev = qml.device("default.qubit", wires=1)
+        dev = qml.device("default.qubit.legacy", wires=1)
         params = torch.tensor([0.1, 0.2], requires_grad=True, device=torch_device)
         x, y = params.detach()
 
@@ -549,7 +549,7 @@ def test_reusing_quantum_tape(self, torch_device, execute_kwargs, tol):
         a = torch.tensor(0.1, requires_grad=True, device=torch_device)
         b = torch.tensor(0.2, requires_grad=True, device=torch_device)
 
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
 
         with qml.queuing.AnnotatedQueue() as q:
             qml.RY(a, wires=0)
@@ -602,7 +602,7 @@ def test_classical_processing(self, torch_device, execute_kwargs, tol):
         p_val = [0.1, 0.2]
         params = torch.tensor(p_val, requires_grad=True, device=torch_device)
 
-        dev = qml.device("default.qubit", wires=1)
+        dev = qml.device("default.qubit.legacy", wires=1)
 
         with qml.queuing.AnnotatedQueue() as q:
             qml.RY(params[0] * params[1], wires=0)
@@ -637,7 +637,7 @@ def test_classical_processing(self, torch_device, execute_kwargs, tol):
 
     def test_no_trainable_parameters(self, torch_device, execute_kwargs):
         """Test evaluation and Jacobian if there are no trainable parameters"""
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
 
         with qml.queuing.AnnotatedQueue() as q:
             qml.RY(0.2, wires=0)
@@ -684,7 +684,7 @@ def test_matrix_parameter(self, torch_device, U, execute_kwargs, tol):
         if isinstance(U, torch.Tensor) and torch_device is not None:
             U = U.to(torch_device)
 
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
 
         with qml.queuing.AnnotatedQueue() as q:
             qml.QubitUnitary(U, wires=0)
@@ -716,7 +716,7 @@ def decomposition(self):
                     qml.PhaseShift(phi + lam, wires=wires),
                 ]
 
-        dev = qml.device("default.qubit", wires=1)
+        dev = qml.device("default.qubit.legacy", wires=1)
         a = np.array(0.1)
         p_val = [0.1, 0.2, 0.3]
         p = torch.tensor(p_val, requires_grad=True, device=torch_device)
@@ -768,7 +768,7 @@ def test_probability_differentiation(self, torch_device, execute_kwargs, tol):
         if execute_kwargs["gradient_fn"] == "device":
             pytest.skip("Adjoint differentiation does not yet support probabilities")
 
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
         x_val = 0.543
         y_val = -0.654
         x = torch.tensor(x_val, requires_grad=True, device=torch_device)
@@ -835,7 +835,7 @@ def test_ragged_differentiation(self, torch_device, execute_kwargs, tol):
         if execute_kwargs["gradient_fn"] == "device":
             pytest.skip("Adjoint differentiation does not yet support probabilities")
 
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
         x_val = 0.543
         y_val = -0.654
         x = torch.tensor(x_val, requires_grad=True, device=torch_device)
@@ -908,7 +908,7 @@ def test_sampling(self, torch_device, execute_kwargs):
         if execute_kwargs["interface"] == "auto":
             pytest.skip("Can't detect interface without a parametrized gate in the tape")
 
-        dev = qml.device("default.qubit", wires=2, shots=10)
+        dev = qml.device("default.qubit.legacy", wires=2, shots=10)
 
         with qml.queuing.AnnotatedQueue() as q:
             qml.Hadamard(wires=[0])
@@ -940,7 +940,7 @@ def test_sampling_expval(self, torch_device, execute_kwargs):
         if execute_kwargs["interface"] == "auto":
             pytest.skip("Can't detect interface without a parametrized gate in the tape")
 
-        dev = qml.device("default.qubit", wires=2, shots=10)
+        dev = qml.device("default.qubit.legacy", wires=2, shots=10)
 
         with qml.queuing.AnnotatedQueue() as q:
             qml.Hadamard(wires=[0])
@@ -968,7 +968,7 @@ def test_sampling_gradient_error(self, torch_device, execute_kwargs):
         ):
             pytest.skip("Adjoint differentiation does not support samples")
 
-        dev = qml.device("default.qubit", wires=1, shots=10)
+        dev = qml.device("default.qubit.legacy", wires=1, shots=10)
 
         x = torch.tensor(0.65, requires_grad=True)
 
@@ -991,7 +991,7 @@ def test_repeated_application_after_expand(self, torch_device, execute_kwargs):
         tape expansions"""
         # pylint: disable=unused-argument
         n_qubits = 2
-        dev = qml.device("default.qubit", wires=n_qubits)
+        dev = qml.device("default.qubit.legacy", wires=n_qubits)
 
         weights = torch.ones((3,))
 
@@ -1021,7 +1021,7 @@ def test_parameter_shift_hessian(self, torch_device, params, tol):
         """Tests that the output of the parameter-shift transform
         can be differentiated using torch, yielding second derivatives."""
         # pylint: disable=unused-argument
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
         params = torch.tensor([0.543, -0.654], requires_grad=True, dtype=torch.float64)
 
         def cost_fn(x):
@@ -1068,7 +1068,7 @@ def cost_fn(x):
 
     def test_hessian_vector_valued(self, torch_device, tol):
         """Test hessian calculation of a vector valued tape"""
-        dev = qml.device("default.qubit", wires=1)
+        dev = qml.device("default.qubit.legacy", wires=1)
 
         def circuit(x):
             with qml.queuing.AnnotatedQueue() as q:
@@ -1135,7 +1135,7 @@ def circuit(x):
     def test_adjoint_hessian(self, torch_device, tol):
         """Since the adjoint hessian is not a differentiable transform,
         higher-order derivatives are not supported."""
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
         params = torch.tensor(
             [0.543, -0.654], requires_grad=True, dtype=torch.float64, device=torch_device
         )
@@ -1164,7 +1164,7 @@ def cost_fn(x):
     def test_max_diff(self, torch_device, tol):
         """Test that setting the max_diff parameter blocks higher-order
         derivatives"""
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
         params = torch.tensor([0.543, -0.654], requires_grad=True, dtype=torch.float64)
 
         def cost_fn(x):
@@ -1275,7 +1275,7 @@ def test_multiple_hamiltonians_not_trainable(self, cost_fn, execute_kwargs, tol)
         coeffs1 = torch.tensor([0.1, 0.2, 0.3], requires_grad=False, dtype=torch.float64)
         coeffs2 = torch.tensor([0.7], requires_grad=False, dtype=torch.float64)
         weights = torch.tensor([0.4, 0.5], requires_grad=True, dtype=torch.float64)
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
 
         res = cost_fn(weights, coeffs1, coeffs2, dev=dev)
         expected = self.cost_fn_expected(weights, coeffs1, coeffs2)
@@ -1293,7 +1293,7 @@ def test_multiple_hamiltonians_trainable(self, cost_fn, execute_kwargs, tol):
         coeffs1 = torch.tensor([0.1, 0.2, 0.3], requires_grad=True, dtype=torch.float64)
         coeffs2 = torch.tensor([0.7], requires_grad=True, dtype=torch.float64)
         weights = torch.tensor([0.4, 0.5], requires_grad=True, dtype=torch.float64)
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
 
         res = cost_fn(weights, coeffs1, coeffs2, dev=dev)
         expected = self.cost_fn_expected(weights, coeffs1, coeffs2)
diff --git a/tests/interfaces/test_torch_qnode.py b/tests/interfaces/test_torch_qnode.py
index 5f295f996fe..77382a6597c 100644
--- a/tests/interfaces/test_torch_qnode.py
+++ b/tests/interfaces/test_torch_qnode.py
@@ -26,13 +26,13 @@
 hessian = torch.autograd.functional.hessian
 
 qubit_device_and_diff_method = [
-    ["default.qubit", "finite-diff", False],
-    ["default.qubit", "parameter-shift", False],
-    ["default.qubit", "backprop", True],
-    ["default.qubit", "adjoint", True],
-    ["default.qubit", "adjoint", False],
-    ["default.qubit", "spsa", False],
-    ["default.qubit", "hadamard", False],
+    ["default.qubit.legacy", "finite-diff", False],
+    ["default.qubit.legacy", "parameter-shift", False],
+    ["default.qubit.legacy", "backprop", True],
+    ["default.qubit.legacy", "adjoint", True],
+    ["default.qubit.legacy", "adjoint", False],
+    ["default.qubit.legacy", "spsa", False],
+    ["default.qubit.legacy", "hadamard", False],
 ]
 
 interface_and_qubit_device_and_diff_method = [
@@ -554,7 +554,7 @@ class TestShotsIntegration:
 
     def test_changing_shots(self, mocker, tol):
         """Test that changing shots works on execution"""
-        dev = qml.device("default.qubit", wires=2, shots=None)
+        dev = qml.device("default.qubit.legacy", wires=2, shots=None)
         a, b = torch.tensor([0.543, -0.654], requires_grad=True, dtype=torch.float64)
 
         @qnode(dev, interface="torch", diff_method=qml.gradients.param_shift)
@@ -588,7 +588,7 @@ def test_gradient_integration(self):
         """Test that temporarily setting the shots works
         for gradient computations"""
         # pylint: disable=unexpected-keyword-arg
-        dev = qml.device("default.qubit", wires=2, shots=None)
+        dev = qml.device("default.qubit.legacy", wires=2, shots=None)
         a, b = torch.tensor([0.543, -0.654], requires_grad=True)
 
         @qnode(dev, interface="torch", diff_method=qml.gradients.param_shift)
@@ -610,7 +610,7 @@ def test_multiple_gradient_integration(self, tol):
         """Test that temporarily setting the shots works
         for gradient computations, even if the QNode has been re-evaluated
         with a different number of shots in the meantime."""
-        dev = qml.device("default.qubit", wires=2, shots=None)
+        dev = qml.device("default.qubit.legacy", wires=2, shots=None)
         weights = torch.tensor([0.543, -0.654], requires_grad=True)
         a, b = weights
 
@@ -634,7 +634,7 @@ def circuit(a, b):
 
     def test_update_diff_method(self, mocker):
         """Test that temporarily setting the shots updates the diff method"""
-        dev = qml.device("default.qubit", wires=2, shots=100)
+        dev = qml.device("default.qubit.legacy", wires=2, shots=100)
         a, b = torch.tensor([0.543, -0.654], requires_grad=True)
 
         spy = mocker.spy(qml, "execute")
@@ -668,7 +668,7 @@ class TestAdjoint:
 
     def test_reuse_state(self, mocker):
         """Tests that the Torch interface reuses the device state for adjoint differentiation"""
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
 
         @qnode(dev, diff_method="adjoint", interface="torch")
         def circ(x):
@@ -692,7 +692,7 @@ def circ(x):
     def test_resuse_state_multiple_evals(self, mocker, tol):
         """Tests that the Torch interface reuses the device state for adjoint differentiation,
         even where there are intermediate evaluations."""
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
 
         x_val = 0.543
         y_val = -0.654
@@ -1616,7 +1616,7 @@ class TestSample:
 
     def test_sample_dimension(self):
         """Test sampling works as expected"""
-        dev = qml.device("default.qubit", wires=2, shots=10)
+        dev = qml.device("default.qubit.legacy", wires=2, shots=10)
 
         @qnode(dev, diff_method="parameter-shift", interface="torch")
         def circuit():
@@ -1638,7 +1638,7 @@ def circuit():
     def test_sampling_expval(self):
         """Test sampling works as expected if combined with expectation values"""
         shots = 10
-        dev = qml.device("default.qubit", wires=2, shots=shots)
+        dev = qml.device("default.qubit.legacy", wires=2, shots=shots)
 
         @qnode(dev, diff_method="parameter-shift", interface="torch")
         def circuit():
@@ -1659,7 +1659,7 @@ def circuit():
     def test_counts_expval(self):
         """Test counts works as expected if combined with expectation values"""
         shots = 10
-        dev = qml.device("default.qubit", wires=2, shots=shots)
+        dev = qml.device("default.qubit.legacy", wires=2, shots=shots)
 
         @qnode(dev, diff_method="parameter-shift", interface="torch")
         def circuit():
@@ -1680,7 +1680,7 @@ def test_sample_combination(self):
         """Test the output of combining expval, var and sample"""
         n_sample = 10
 
-        dev = qml.device("default.qubit", wires=3, shots=n_sample)
+        dev = qml.device("default.qubit.legacy", wires=3, shots=n_sample)
 
         @qnode(dev, diff_method="parameter-shift", interface="torch")
         def circuit():
@@ -1704,7 +1704,7 @@ def test_single_wire_sample(self):
         """Test the return type and shape of sampling a single wire"""
         n_sample = 10
 
-        dev = qml.device("default.qubit", wires=1, shots=n_sample)
+        dev = qml.device("default.qubit.legacy", wires=1, shots=n_sample)
 
         @qnode(dev, diff_method="parameter-shift", interface="torch")
         def circuit():
@@ -1721,7 +1721,7 @@ def test_multi_wire_sample_regular_shape(self):
         where a rectangular array is expected"""
         n_sample = 10
 
-        dev = qml.device("default.qubit", wires=3, shots=n_sample)
+        dev = qml.device("default.qubit.legacy", wires=3, shots=n_sample)
 
         @qnode(dev, diff_method="parameter-shift", interface="torch")
         def circuit():
@@ -1740,12 +1740,12 @@ def circuit():
 
 
 qubit_device_and_diff_method_and_grad_on_execution = [
-    ["default.qubit", "backprop", True],
-    ["default.qubit", "finite-diff", False],
-    ["default.qubit", "parameter-shift", False],
-    ["default.qubit", "adjoint", True],
-    ["default.qubit", "adjoint", False],
-    ["default.qubit", "hadamard", False],
+    ["default.qubit.legacy", "backprop", True],
+    ["default.qubit.legacy", "finite-diff", False],
+    ["default.qubit.legacy", "parameter-shift", False],
+    ["default.qubit.legacy", "adjoint", True],
+    ["default.qubit.legacy", "adjoint", False],
+    ["default.qubit.legacy", "hadamard", False],
 ]
 
 
@@ -2581,7 +2581,7 @@ def circuit_stack(x):
         assert tuple(hess.shape) == (3, 2, 2)
 
 
-@pytest.mark.parametrize("dev_name", ["default.qubit", "default.mixed"])
+@pytest.mark.parametrize("dev_name", ["default.qubit.legacy", "default.mixed"])
 def test_no_ops(dev_name):
     """Test that the return value of the QNode matches in the interface
     even if there are no ops"""
diff --git a/tests/interfaces/test_transform_program_integration.py b/tests/interfaces/test_transform_program_integration.py
index 39c84a02c48..3708c52552b 100644
--- a/tests/interfaces/test_transform_program_integration.py
+++ b/tests/interfaces/test_transform_program_integration.py
@@ -25,7 +25,7 @@
 
 
 device_suite = (
-    qml.device("default.qubit", wires=5),
+    qml.device("default.qubit.legacy", wires=5),
     qml.devices.experimental.DefaultQubit2(),
     qml.device("lightning.qubit", wires=5),
 )
@@ -168,7 +168,7 @@ def split_sum_terms(tape):
     def test_chained_preprocessing(self):
         """Test a transform program with two transforms where their order affects the output."""
 
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
 
         def null_postprocessing(results):
             return results[0]
diff --git a/tests/measurements/test_classical_shadow.py b/tests/measurements/test_classical_shadow.py
index 8b523526b40..06f0b3a7ca2 100644
--- a/tests/measurements/test_classical_shadow.py
+++ b/tests/measurements/test_classical_shadow.py
@@ -26,7 +26,7 @@
 # pylint: disable=dangerous-default-value, too-many-arguments, comparison-with-callable, no-member
 
 
-def get_circuit(wires, shots, seed_recipes, interface="autograd", device="default.qubit"):
+def get_circuit(wires, shots, seed_recipes, interface="autograd", device="default.qubit.legacy"):
     """
     Return a QNode that prepares the state (|00...0> + |11...1>) / sqrt(2)
         and performs the classical shadow measurement
@@ -34,7 +34,7 @@ def get_circuit(wires, shots, seed_recipes, interface="autograd", device="defaul
     if device is not None:
         dev = qml.device(device, wires=wires, shots=shots)
     else:
-        dev = qml.device("default.qubit", wires=wires, shots=shots)
+        dev = qml.device("default.qubit.legacy", wires=wires, shots=shots)
 
         # make the device call the superclass method to switch between the general qubit device and device specific implementations (i.e. for default qubit)
         dev.classical_shadow = super(type(dev), dev).classical_shadow
@@ -51,14 +51,14 @@ def circuit():
     return circuit
 
 
-def get_x_basis_circuit(wires, shots, interface="autograd", device="default.qubit"):
+def get_x_basis_circuit(wires, shots, interface="autograd", device="default.qubit.legacy"):
     """
     Return a QNode that prepares the |++..+> state and performs a classical shadow measurement
     """
     if device is not None:
         dev = qml.device(device, wires=wires, shots=shots)
     else:
-        dev = qml.device("default.qubit", wires=wires, shots=shots)
+        dev = qml.device("default.qubit.legacy", wires=wires, shots=shots)
 
         # make the device call the superclass method to switch between the general qubit device and device specific implementations (i.e. for default qubit)
         dev.classical_shadow = super(type(dev), dev).classical_shadow
@@ -72,14 +72,14 @@ def circuit():
     return circuit
 
 
-def get_y_basis_circuit(wires, shots, interface="autograd", device="default.qubit"):
+def get_y_basis_circuit(wires, shots, interface="autograd", device="default.qubit.legacy"):
     """
     Return a QNode that prepares the |+i>|+i>...|+i> state and performs a classical shadow measurement
     """
     if device is not None:
         dev = qml.device(device, wires=wires, shots=shots)
     else:
-        dev = qml.device("default.qubit", wires=wires, shots=shots)
+        dev = qml.device("default.qubit.legacy", wires=wires, shots=shots)
 
         # make the device call the superclass method to switch between the general qubit device and device specific implementations (i.e. for default qubit)
         dev.classical_shadow = super(type(dev), dev).classical_shadow
@@ -94,14 +94,14 @@ def circuit():
     return circuit
 
 
-def get_z_basis_circuit(wires, shots, interface="autograd", device="default.qubit"):
+def get_z_basis_circuit(wires, shots, interface="autograd", device="default.qubit.legacy"):
     """
     Return a QNode that prepares the |00..0> state and performs a classical shadow measurement
     """
     if device is not None:
         dev = qml.device(device, wires=wires, shots=shots)
     else:
-        dev = qml.device("default.qubit", wires=wires, shots=shots)
+        dev = qml.device("default.qubit.legacy", wires=wires, shots=shots)
 
         # make the device call the superclass method to switch between the general qubit device and device specific implementations (i.e. for default qubit)
         dev.classical_shadow = super(type(dev), dev).classical_shadow
@@ -263,7 +263,7 @@ def test_measurement_process_numeric_type(self, wires, seed):
     @pytest.mark.parametrize("seed", seed_recipes_list)
     def test_measurement_process_shape(self, wires, shots, seed):
         """Test that the shape of the MeasurementProcess instance is correct"""
-        dev = qml.device("default.qubit", wires=wires, shots=shots)
+        dev = qml.device("default.qubit.legacy", wires=wires, shots=shots)
         shots_obj = Shots(shots)
         res = qml.classical_shadow(wires=range(wires), seed=seed)
         assert res.shape(dev, shots_obj) == (2, shots, wires)
@@ -304,7 +304,7 @@ def test_measurement_process_copy(self, wires, seed):
     @pytest.mark.parametrize("shots", shots_list)
     @pytest.mark.parametrize("seed", seed_recipes_list)
     @pytest.mark.parametrize("interface", ["autograd", "jax", "tf", "torch"])
-    @pytest.mark.parametrize("device", ["default.qubit", "default.mixed", None])
+    @pytest.mark.parametrize("device", ["default.qubit.legacy", "default.mixed", None])
     def test_format(self, wires, shots, seed, interface, device):
         """Test that the format of the returned classical shadow
         measurement is correct"""
@@ -334,7 +334,7 @@ def test_format(self, wires, shots, seed, interface, device):
 
     @pytest.mark.all_interfaces
     @pytest.mark.parametrize("interface", ["autograd", "jax", "tf", "torch"])
-    @pytest.mark.parametrize("device", ["default.qubit", None])
+    @pytest.mark.parametrize("device", ["default.qubit.legacy", None])
     @pytest.mark.parametrize(
         "circuit_fn, basis_recipe",
         [(get_x_basis_circuit, 0), (get_y_basis_circuit, 1), (get_z_basis_circuit, 2)],
@@ -390,7 +390,7 @@ def test_shots_none_error(self, wires, seed):
     def test_multi_measurement_error(self, wires, shots):
         """Test that an error is raised when classical shadows is returned
         with other measurement processes"""
-        dev = qml.device("default.qubit", wires=wires, shots=shots)
+        dev = qml.device("default.qubit.legacy", wires=wires, shots=shots)
 
         @qml.qnode(dev)
         def circuit():
@@ -407,7 +407,7 @@ def circuit():
 
 
 def hadamard_circuit(wires, shots=10000, interface="autograd"):
-    dev = qml.device("default.qubit", wires=wires, shots=shots)
+    dev = qml.device("default.qubit.legacy", wires=wires, shots=shots)
 
     @qml.qnode(dev, interface=interface)
     def circuit(obs, k=1):
@@ -419,7 +419,7 @@ def circuit(obs, k=1):
 
 
 def max_entangled_circuit(wires, shots=10000, interface="autograd"):
-    dev = qml.device("default.qubit", wires=wires, shots=shots)
+    dev = qml.device("default.qubit.legacy", wires=wires, shots=shots)
 
     @qml.qnode(dev, interface=interface)
     def circuit(obs, k=1):
@@ -432,7 +432,7 @@ def circuit(obs, k=1):
 
 
 def qft_circuit(wires, shots=10000, interface="autograd"):
-    dev = qml.device("default.qubit", wires=wires, shots=shots)
+    dev = qml.device("default.qubit.legacy", wires=wires, shots=shots)
 
     one_state = np.zeros(wires)
     one_state[-1] = 1
@@ -458,7 +458,7 @@ def test_measurement_process_numeric_type(self):
     @pytest.mark.parametrize("shots", [1, 10])
     def test_measurement_process_shape(self, wires, shots):
         """Test that the shape of the MeasurementProcess instance is correct"""
-        dev = qml.device("default.qubit", wires=wires, shots=shots)
+        dev = qml.device("default.qubit.legacy", wires=wires, shots=shots)
         H = qml.PauliZ(0)
         res = qml.shadow_expval(H)
         assert len(res.shape(dev, Shots(shots))) == 0
@@ -504,7 +504,7 @@ def test_shots_none_error(self):
     def test_multi_measurement_error(self):
         """Test that an error is raised when classical shadows is returned
         with other measurement processes"""
-        dev = qml.device("default.qubit", wires=2, shots=100)
+        dev = qml.device("default.qubit.legacy", wires=2, shots=100)
 
         @qml.qnode(dev)
         def circuit():
@@ -647,7 +647,7 @@ def test_qft_expval(self, interface, k=1, obs=obs_qft, expected=expected_qft):
 
 
 def strongly_entangling_circuit(wires, shots=10000, interface="autograd"):
-    dev = qml.device("default.qubit", wires=wires, shots=shots)
+    dev = qml.device("default.qubit.legacy", wires=wires, shots=shots)
 
     @qml.qnode(dev, interface=interface)
     def circuit(x, obs, k):
@@ -658,7 +658,7 @@ def circuit(x, obs, k):
 
 
 def strongly_entangling_circuit_exact(wires, interface="autograd"):
-    dev = qml.device("default.qubit", wires=wires)
+    dev = qml.device("default.qubit.legacy", wires=wires)
 
     @qml.qnode(dev, interface=interface)
     def circuit(x, obs):
diff --git a/tests/measurements/test_counts.py b/tests/measurements/test_counts.py
index aa2cee5a973..f936561995f 100644
--- a/tests/measurements/test_counts.py
+++ b/tests/measurements/test_counts.py
@@ -249,7 +249,7 @@ def test_counts_dimension(self, mocker):
         """Test that the counts function outputs counts of the right size"""
         n_sample = 10
 
-        dev = qml.device("default.qubit", wires=2, shots=n_sample)
+        dev = qml.device("default.qubit.legacy", wires=2, shots=n_sample)
         spy = mocker.spy(qml.QubitDevice, "sample")
 
         @qml.qnode(dev)
@@ -268,7 +268,7 @@ def test_batched_counts_dimension(self, mocker):
         """Test that the counts function outputs counts of the right size with batching"""
         n_sample = 10
 
-        dev = qml.device("default.qubit", wires=2, shots=n_sample)
+        dev = qml.device("default.qubit.legacy", wires=2, shots=n_sample)
         spy = mocker.spy(qml.QubitDevice, "sample")
 
         @qml.qnode(dev)
@@ -288,7 +288,7 @@ def test_counts_combination(self, mocker):
         """Test the output of combining expval, var and counts"""
         n_sample = 10
 
-        dev = qml.device("default.qubit", wires=3, shots=n_sample)
+        dev = qml.device("default.qubit.legacy", wires=3, shots=n_sample)
         spy = mocker.spy(qml.QubitDevice, "sample")
 
         @qml.qnode(dev, diff_method="parameter-shift")
@@ -312,7 +312,7 @@ def test_single_wire_counts(self, mocker):
         """Test the return type and shape of sampling counts from a single wire"""
         n_sample = 10
 
-        dev = qml.device("default.qubit", wires=1, shots=n_sample)
+        dev = qml.device("default.qubit.legacy", wires=1, shots=n_sample)
         spy = mocker.spy(qml.QubitDevice, "sample")
 
         @qml.qnode(dev)
@@ -333,7 +333,7 @@ def test_multi_wire_counts_regular_shape(self, mocker):
         where a rectangular array is expected"""
         n_sample = 10
 
-        dev = qml.device("default.qubit", wires=3, shots=n_sample)
+        dev = qml.device("default.qubit.legacy", wires=3, shots=n_sample)
         spy = mocker.spy(qml.QubitDevice, "sample")
 
         @qml.qnode(dev)
@@ -357,7 +357,7 @@ def circuit():
     def test_observable_return_type_is_counts(self):
         """Test that the return type of the observable is :attr:`ObservableReturnTypes.Counts`"""
         n_shots = 10
-        dev = qml.device("default.qubit", wires=1, shots=n_shots)
+        dev = qml.device("default.qubit.legacy", wires=1, shots=n_shots)
 
         @qml.qnode(dev)
         def circuit():
@@ -369,7 +369,7 @@ def circuit():
 
     def test_providing_no_observable_and_no_wires_counts(self, mocker):
         """Test that we can provide no observable and no wires to sample function"""
-        dev = qml.device("default.qubit", wires=2, shots=1000)
+        dev = qml.device("default.qubit.legacy", wires=2, shots=1000)
         spy = mocker.spy(qml.QubitDevice, "sample")
 
         @qml.qnode(dev)
@@ -388,7 +388,7 @@ def test_providing_no_observable_and_wires_counts(self, mocker):
         """Test that we can provide no observable but specify wires to the sample function"""
         wires = [0, 2]
         wires_obj = qml.wires.Wires(wires)
-        dev = qml.device("default.qubit", wires=3, shots=1000)
+        dev = qml.device("default.qubit.legacy", wires=3, shots=1000)
         spy = mocker.spy(qml.QubitDevice, "sample")
 
         @qml.qnode(dev)
@@ -407,7 +407,7 @@ def circuit():
     def test_batched_counts_work_individually(self, mocker):
         """Test that each counts call operates independently"""
         n_shots = 10
-        dev = qml.device("default.qubit", wires=1, shots=n_shots)
+        dev = qml.device("default.qubit.legacy", wires=1, shots=n_shots)
         spy = mocker.spy(qml.QubitDevice, "sample")
 
         @qml.qnode(dev)
@@ -425,7 +425,7 @@ def test_counts_no_op_finite_shots(self, interface, wires, basis_state, mocker):
         """Check all interfaces with computational basis state counts and
         finite shot"""
         n_shots = 10
-        dev = qml.device("default.qubit", wires=3, shots=n_shots)
+        dev = qml.device("default.qubit.legacy", wires=3, shots=n_shots)
         spy = mocker.spy(qml.QubitDevice, "sample")
 
         @qml.qnode(dev, interface=interface)
@@ -445,7 +445,7 @@ def test_counts_operator_finite_shots(self, interface, mocker):
         """Check all interfaces with observable measurement counts and finite
         shot"""
         n_shots = 10
-        dev = qml.device("default.qubit", wires=3, shots=n_shots)
+        dev = qml.device("default.qubit.legacy", wires=3, shots=n_shots)
         spy = mocker.spy(qml.QubitDevice, "sample")
 
         @qml.qnode(dev, interface=interface)
@@ -466,7 +466,7 @@ def test_counts_binned(
     ):  # pylint:disable=too-many-arguments
         """Check all interfaces with computational basis state counts and
         different shot vectors"""
-        dev = qml.device("default.qubit", wires=3, shots=shot_vec)
+        dev = qml.device("default.qubit.legacy", wires=3, shots=shot_vec)
         spy = mocker.spy(qml.QubitDevice, "sample")
 
         @qml.qnode(dev, interface=interface)
@@ -491,7 +491,7 @@ def circuit():
     def test_counts_operator_binned(self, shot_vec, interface, mocker):
         """Check all interfaces with observable measurement counts and different
         shot vectors"""
-        dev = qml.device("default.qubit", wires=3, shots=shot_vec)
+        dev = qml.device("default.qubit.legacy", wires=3, shots=shot_vec)
         spy = mocker.spy(qml.QubitDevice, "sample")
 
         @qml.qnode(dev, interface=interface)
@@ -513,7 +513,7 @@ def circuit():
     def test_counts_binned_4_wires(self, shot_vec, mocker):
         """Check the autograd interface with computational basis state counts and
         different shot vectors on a device with 4 wires"""
-        dev = qml.device("default.qubit", wires=4, shots=shot_vec)
+        dev = qml.device("default.qubit.legacy", wires=4, shots=shot_vec)
         spy = mocker.spy(qml.QubitDevice, "sample")
 
         @qml.qnode(dev, interface="autograd")
@@ -539,7 +539,7 @@ def circuit():
     def test_counts_operator_binned_4_wires(self, shot_vec, mocker):
         """Check the autograd interface with observable samples to obtain
         counts from and different shot vectors on a device with 4 wires"""
-        dev = qml.device("default.qubit", wires=4, shots=shot_vec)
+        dev = qml.device("default.qubit.legacy", wires=4, shots=shot_vec)
         spy = mocker.spy(qml.QubitDevice, "sample")
 
         @qml.qnode(dev, interface="autograd")
@@ -575,7 +575,7 @@ def circuit():
     def test_counts_observable_finite_shots(self, interface, meas2, shots, mocker):
         """Check all interfaces with observable measurement counts and finite
         shot"""
-        dev = qml.device("default.qubit", wires=3, shots=shots)
+        dev = qml.device("default.qubit.legacy", wires=3, shots=shots)
         spy = mocker.spy(qml.QubitDevice, "sample")
 
         if isinstance(shots, tuple) and interface == "torch":
@@ -610,7 +610,7 @@ def test_all_outcomes_kwarg_providing_observable(self, mocker):
         including 0 count values, if observable is given and all_outcomes=True"""
 
         n_shots = 10
-        dev = qml.device("default.qubit", wires=1, shots=n_shots)
+        dev = qml.device("default.qubit.legacy", wires=1, shots=n_shots)
         spy = mocker.spy(qml.QubitDevice, "sample")
 
         @qml.qnode(dev)
@@ -630,7 +630,7 @@ def test_all_outcomes_kwarg_no_observable_no_wires(self, mocker):
         count and no observable are given and all_outcomes=True"""
 
         n_shots = 10
-        dev = qml.device("default.qubit", wires=2, shots=n_shots)
+        dev = qml.device("default.qubit.legacy", wires=2, shots=n_shots)
         spy = mocker.spy(qml.QubitDevice, "sample")
 
         @qml.qnode(dev)
@@ -649,7 +649,7 @@ def test_all_outcomes_kwarg_providing_wires_and_no_observable(self, mocker):
         if wire count is given and all_outcomes=True"""
 
         n_shots = 10
-        dev = qml.device("default.qubit", wires=4, shots=n_shots)
+        dev = qml.device("default.qubit.legacy", wires=4, shots=n_shots)
         spy = mocker.spy(qml.QubitDevice, "sample")
 
         @qml.qnode(dev)
@@ -667,7 +667,7 @@ def test_all_outcomes_hermitian(self, mocker):
         qml.Hermitian observable"""
 
         n_shots = 10
-        dev = qml.device("default.qubit", wires=2, shots=n_shots)
+        dev = qml.device("default.qubit.legacy", wires=2, shots=n_shots)
         spy = mocker.spy(qml.QubitDevice, "sample")
 
         A = np.array([[1, 0], [0, -1]])
@@ -686,7 +686,7 @@ def test_all_outcomes_multiple_measurements(self, mocker):
         """Tests that the all_outcomes=True option for counts works when
         multiple measurements are performed"""
 
-        dev = qml.device("default.qubit", wires=2, shots=10)
+        dev = qml.device("default.qubit.legacy", wires=2, shots=10)
         spy = mocker.spy(qml.QubitDevice, "sample")
 
         @qml.qnode(dev)
@@ -703,7 +703,7 @@ def circuit():
     def test_batched_all_outcomes(self, mocker):
         """Tests that all_outcomes=True works with broadcasting."""
         n_shots = 10
-        dev = qml.device("default.qubit", wires=1, shots=n_shots)
+        dev = qml.device("default.qubit.legacy", wires=1, shots=n_shots)
         spy = mocker.spy(qml.QubitDevice, "sample")
 
         @qml.qnode(dev)
@@ -723,7 +723,7 @@ def test_counts_empty_wires(self):
     @pytest.mark.parametrize("shots", [1, 100])
     def test_counts_no_arguments(self, shots):
         """Test that using ``qml.counts`` with no arguments returns the counts of all wires."""
-        dev = qml.device("default.qubit", wires=3, shots=shots)
+        dev = qml.device("default.qubit.legacy", wires=3, shots=shots)
 
         @qml.qnode(dev)
         def circuit():
@@ -740,7 +740,7 @@ def circuit():
 def test_counts_no_op_finite_shots(interface, wires, basis_state, mocker):
     """Check all interfaces with computational basis state counts and finite shot"""
     n_shots = 10
-    dev = qml.device("default.qubit", wires=3, shots=n_shots)
+    dev = qml.device("default.qubit.legacy", wires=3, shots=n_shots)
     spy = mocker.spy(qml.QubitDevice, "sample")
 
     @qml.qnode(dev, interface=interface)
@@ -762,7 +762,7 @@ def test_batched_counts_no_op_finite_shots(interface, wires, basis_states, mocke
     """Check all interfaces with computational basis state counts and
     finite shot"""
     n_shots = 10
-    dev = qml.device("default.qubit", wires=3, shots=n_shots)
+    dev = qml.device("default.qubit.legacy", wires=3, shots=n_shots)
     spy = mocker.spy(qml.QubitDevice, "sample")
 
     @qml.qnode(dev, interface=interface)
@@ -782,7 +782,7 @@ def test_batched_counts_and_expval_no_op_finite_shots(interface, wires, basis_st
     """Check all interfaces with computational basis state counts and
     finite shot"""
     n_shots = 10
-    dev = qml.device("default.qubit", wires=3, shots=n_shots)
+    dev = qml.device("default.qubit.legacy", wires=3, shots=n_shots)
     spy = mocker.spy(qml.QubitDevice, "sample")
 
     @qml.qnode(dev, interface=interface)
@@ -803,7 +803,7 @@ def circuit():
 def test_batched_counts_operator_finite_shots(interface, mocker):
     """Check all interfaces with observable measurement counts, batching and finite shots"""
     n_shots = 10
-    dev = qml.device("default.qubit", wires=3, shots=n_shots)
+    dev = qml.device("default.qubit.legacy", wires=3, shots=n_shots)
     spy = mocker.spy(qml.QubitDevice, "sample")
 
     @qml.qnode(dev, interface=interface)
@@ -821,7 +821,7 @@ def circuit():
 def test_batched_counts_and_expval_operator_finite_shots(interface, mocker):
     """Check all interfaces with observable measurement counts, batching and finite shots"""
     n_shots = 10
-    dev = qml.device("default.qubit", wires=3, shots=n_shots)
+    dev = qml.device("default.qubit.legacy", wires=3, shots=n_shots)
     spy = mocker.spy(qml.QubitDevice, "sample")
 
     @qml.qnode(dev, interface=interface)
diff --git a/tests/measurements/test_expval.py b/tests/measurements/test_expval.py
index 25a890c4d4c..c5a6850d4ae 100644
--- a/tests/measurements/test_expval.py
+++ b/tests/measurements/test_expval.py
@@ -53,7 +53,7 @@ class TestExpval:
     @pytest.mark.parametrize("r_dtype", [np.float32, np.float64])
     def test_value(self, tol, r_dtype, mocker, shots):
         """Test that the expval interface works"""
-        dev = qml.device("default.qubit", wires=2, shots=shots)
+        dev = qml.device("default.qubit.legacy", wires=2, shots=shots)
         dev.R_DTYPE = r_dtype
 
         @qml.qnode(dev, diff_method="parameter-shift")
@@ -84,7 +84,7 @@ def circuit(x):
     def test_not_an_observable(self, mocker):
         """Test that a warning is raised if the provided
         argument might not be hermitian."""
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
 
         @qml.qnode(dev)
         def circuit():
@@ -101,7 +101,7 @@ def circuit():
 
     def test_observable_return_type_is_expectation(self, mocker):
         """Test that the return type of the observable is :attr:`ObservableReturnTypes.Expectation`"""
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
 
         @qml.qnode(dev)
         def circuit():
@@ -155,7 +155,7 @@ def test_numeric_type(self, obs):
     )
     def test_shape(self, obs):
         """Test that the shape is correct."""
-        dev = qml.device("default.qubit", wires=1)
+        dev = qml.device("default.qubit.legacy", wires=1)
 
         res = qml.expval(obs)
         # pylint: disable=use-implicit-booleaness-not-comparison
@@ -170,7 +170,7 @@ def test_shape_shot_vector(self, obs):
         """Test that the shape is correct with the shot vector too."""
         res = qml.expval(obs)
         shot_vector = (1, 2, 3)
-        dev = qml.device("default.qubit", wires=3, shots=shot_vector)
+        dev = qml.device("default.qubit.legacy", wires=3, shots=shot_vector)
         assert res.shape(dev, Shots(shot_vector)) == ((), (), ())
 
     @pytest.mark.parametrize("state", [np.array([0, 0, 0]), np.array([1, 0, 0, 0, 0, 0, 0, 0])])
@@ -178,7 +178,7 @@ def test_shape_shot_vector(self, obs):
     def test_projector_expval(self, state, shots, mocker):
         """Tests that the expectation of a ``Projector`` object is computed correctly for both of
         its subclasses."""
-        dev = qml.device("default.qubit", wires=3, shots=shots)
+        dev = qml.device("default.qubit.legacy", wires=3, shots=shots)
         np.random.seed(42)
 
         @qml.qnode(dev)
@@ -202,7 +202,7 @@ def test_permuted_wires(self, mocker):
             qml.s_prod(3, qml.PauliZ("h")), qml.PauliZ(8), qml.s_prod(2, qml.PauliZ(10))
         )
 
-        dev = qml.device("default.qubit", wires=["h", 8, 10])
+        dev = qml.device("default.qubit.legacy", wires=["h", 8, 10])
         spy = mocker.spy(qml.QubitDevice, "expval")
 
         @qml.qnode(dev)
diff --git a/tests/measurements/test_measurements.py b/tests/measurements/test_measurements.py
index 0ff5b408bb4..a88b72f6391 100644
--- a/tests/measurements/test_measurements.py
+++ b/tests/measurements/test_measurements.py
@@ -78,7 +78,7 @@ def test_ObservableReturnTypes(return_type, value):
 def test_no_measure():
     """Test that failing to specify a measurement
     raises an exception"""
-    dev = qml.device("default.qubit", wires=2)
+    dev = qml.device("default.qubit.legacy", wires=2)
 
     @qml.qnode(dev)
     def circuit(x):
@@ -104,7 +104,7 @@ def test_numeric_type_unrecognized_error():
 def test_shape_unrecognized_error():
     """Test that querying the shape of a measurement process with an
     unrecognized return type raises an error."""
-    dev = qml.device("default.qubit", wires=2)
+    dev = qml.device("default.qubit.legacy", wires=2)
     mp = NotValidMeasurement()
     with pytest.raises(
         qml.QuantumFunctionError,
@@ -244,7 +244,7 @@ def test_not_an_observable(self, stat_func, return_type):  # pylint: disable=unu
         if stat_func is sample:
             pytest.skip("Sampling is not yet supported with symbolic operators.")
 
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
 
         @qml.qnode(dev)
         def circuit():
@@ -492,7 +492,7 @@ class MyMeasurement(SampleMeasurement):
             def process_samples(self, samples, wire_order, shot_range, bin_size):
                 return qml.math.sum(samples[..., self.wires])
 
-        dev = qml.device("default.qubit", wires=2, shots=1000)
+        dev = qml.device("default.qubit.legacy", wires=2, shots=1000)
 
         @qml.qnode(dev)
         def circuit():
@@ -509,7 +509,7 @@ class MyMeasurement(SampleMeasurement):
             def process_samples(self, samples, wire_order, shot_range, bin_size):
                 return qml.math.sum(samples[..., self.wires])
 
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
 
         @qml.qnode(dev)
         def circuit():
@@ -524,7 +524,7 @@ def circuit():
     def test_method_overriden_by_device(self):
         """Test that the device can override a measurement process."""
 
-        dev = qml.device("default.qubit", wires=2, shots=1000)
+        dev = qml.device("default.qubit.legacy", wires=2, shots=1000)
 
         @qml.qnode(dev)
         def circuit():
@@ -547,7 +547,7 @@ class MyMeasurement(StateMeasurement):
             def process_state(self, state, wire_order):
                 return qml.math.sum(state)
 
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
 
         @qml.qnode(dev)
         def circuit():
@@ -562,7 +562,7 @@ class MyMeasurement(StateMeasurement):
             def process_state(self, state, wire_order):
                 return qml.math.sum(state)
 
-        dev = qml.device("default.qubit", wires=2, shots=1000)
+        dev = qml.device("default.qubit.legacy", wires=2, shots=1000)
 
         @qml.qnode(dev)
         def circuit():
@@ -577,7 +577,7 @@ def circuit():
     def test_method_overriden_by_device(self):
         """Test that the device can override a measurement process."""
 
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
 
         @qml.qnode(dev, interface="autograd")
         def circuit():
@@ -599,7 +599,7 @@ class MyMeasurement(MeasurementTransform):
             def process(self, tape, device):
                 return {device.shots: len(tape)}
 
-        dev = qml.device("default.qubit", wires=2, shots=1000)
+        dev = qml.device("default.qubit.legacy", wires=2, shots=1000)
 
         @qml.qnode(dev)
         def circuit():
@@ -610,7 +610,7 @@ def circuit():
     def test_method_overriden_by_device(self):
         """Test that the device can override a measurement process."""
 
-        dev = qml.device("default.qubit", wires=2, shots=1000)
+        dev = qml.device("default.qubit.legacy", wires=2, shots=1000)
 
         @qml.qnode(dev)
         def circuit():
diff --git a/tests/measurements/test_mutual_info.py b/tests/measurements/test_mutual_info.py
index 79eedc1d44e..d75ce6de280 100644
--- a/tests/measurements/test_mutual_info.py
+++ b/tests/measurements/test_mutual_info.py
@@ -39,7 +39,7 @@ def test_queue(self):
     @pytest.mark.parametrize("shots, shape", [(None, ()), (10, ()), ([1, 10], ((), ()))])
     def test_shape(self, shots, shape):
         """Test that the shape is correct."""
-        dev = qml.device("default.qubit", wires=3, shots=shots)
+        dev = qml.device("default.qubit.legacy", wires=3, shots=shots)
         res = qml.mutual_info(wires0=[0], wires1=[1])
         assert res.shape(dev, Shots(shots)) == shape
 
@@ -88,7 +88,7 @@ class TestIntegration:
     )
     def test_mutual_info_output(self, interface, state, expected):
         """Test the output of qml.mutual_info"""
-        dev = qml.device("default.qubit", wires=4)
+        dev = qml.device("default.qubit.legacy", wires=4)
 
         @qml.qnode(dev, interface=interface)
         def circuit():
@@ -106,7 +106,7 @@ def circuit():
 
     def test_shot_vec_error(self):
         """Test an error is raised when using shot vectors with mutual_info."""
-        dev = qml.device("default.qubit", wires=2, shots=[1, 10, 10, 1000])
+        dev = qml.device("default.qubit.legacy", wires=2, shots=[1, 10, 10, 1000])
 
         @qml.qnode(device=dev)
         def circuit(x):
@@ -122,7 +122,7 @@ def circuit(x):
     diff_methods = ["backprop", "finite-diff"]
 
     @pytest.mark.all_interfaces
-    @pytest.mark.parametrize("device", ["default.qubit", "default.mixed", "lightning.qubit"])
+    @pytest.mark.parametrize("device", ["default.qubit.legacy", "default.mixed", "lightning.qubit"])
     @pytest.mark.parametrize("interface", ["autograd", "jax", "tensorflow", "torch"])
     @pytest.mark.parametrize("params", np.linspace(0, 2 * np.pi, 8))
     def test_qnode_state(self, device, interface, params):
@@ -148,7 +148,7 @@ def circuit(params):
         assert np.allclose(actual, expected)
 
     @pytest.mark.all_interfaces
-    @pytest.mark.parametrize("device", ["default.qubit", "default.mixed", "lightning.qubit"])
+    @pytest.mark.parametrize("device", ["default.qubit.legacy", "default.mixed", "lightning.qubit"])
     @pytest.mark.parametrize("interface", ["autograd", "jax", "tensorflow", "torch"])
     @pytest.mark.parametrize("params", zip(np.linspace(0, np.pi, 8), np.linspace(0, 2 * np.pi, 8)))
     def test_qnode_mutual_info(self, device, interface, params):
@@ -179,7 +179,7 @@ def circuit_state(params):
 
         assert np.allclose(actual, expected)
 
-    @pytest.mark.parametrize("device", ["default.qubit", "default.mixed", "lightning.qubit"])
+    @pytest.mark.parametrize("device", ["default.qubit.legacy", "default.mixed", "lightning.qubit"])
     def test_mutual_info_wire_labels(self, device):
         """Test that mutual_info is correct with custom wire labels"""
         param = np.array([0.678, 1.234])
@@ -209,7 +209,7 @@ def test_qnode_state_jax_jit(self, params):
         import jax
         import jax.numpy as jnp
 
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
 
         params = jnp.array(params)
 
@@ -237,7 +237,7 @@ def test_qnode_mutual_info_jax_jit(self, params, interface):
         import jax
         import jax.numpy as jnp
 
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
 
         params = jnp.array(params)
 
@@ -269,7 +269,7 @@ def circuit_state(params):
     def test_qnode_grad(self, param, diff_method, interface):
         """Test that the gradient of mutual information works for QNodes
         with the autograd interface"""
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
 
         @qml.qnode(dev, interface=interface, diff_method=diff_method)
         def circuit(param):
@@ -301,7 +301,7 @@ def test_qnode_grad_jax(self, param, diff_method, interface):
         import jax
         import jax.numpy as jnp
 
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
 
         param = jnp.array(param)
 
@@ -335,7 +335,7 @@ def test_qnode_grad_jax_jit(self, param, diff_method, interface):
         import jax
         import jax.numpy as jnp
 
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
 
         param = jnp.array(param)
 
@@ -368,7 +368,7 @@ def test_qnode_grad_tf(self, param, diff_method, interface):
         with the tensorflow interface"""
         import tensorflow as tf
 
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
 
         param = tf.Variable(param)
 
@@ -404,7 +404,7 @@ def test_qnode_grad_torch(self, param, diff_method, interface):
         with the torch interface"""
         import torch
 
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
 
         @qml.qnode(dev, interface=interface, diff_method=diff_method)
         def circuit(param):
@@ -431,7 +431,7 @@ def circuit(param):
         assert np.allclose(actual, expected, atol=tol)
 
     @pytest.mark.all_interfaces
-    @pytest.mark.parametrize("device", ["default.qubit", "default.mixed", "lightning.qubit"])
+    @pytest.mark.parametrize("device", ["default.qubit.legacy", "default.mixed", "lightning.qubit"])
     @pytest.mark.parametrize("interface", ["autograd", "jax", "tensorflow", "torch"])
     @pytest.mark.parametrize(
         "params", [np.array([0.0, 0.0]), np.array([0.3, 0.4]), np.array([0.6, 0.8])]
@@ -455,7 +455,7 @@ def circuit(params):
             circuit(params)
 
     @pytest.mark.all_interfaces
-    @pytest.mark.parametrize("device", ["default.qubit", "default.mixed", "lightning.qubit"])
+    @pytest.mark.parametrize("device", ["default.qubit.legacy", "default.mixed", "lightning.qubit"])
     @pytest.mark.parametrize("interface", ["autograd", "jax", "tensorflow", "torch"])
     @pytest.mark.parametrize(
         "params", [np.array([0.0, 0.0]), np.array([0.3, 0.4]), np.array([0.6, 0.8])]
diff --git a/tests/measurements/test_probs.py b/tests/measurements/test_probs.py
index 45f06e5fc21..fd3fc698920 100644
--- a/tests/measurements/test_probs.py
+++ b/tests/measurements/test_probs.py
@@ -78,7 +78,7 @@ class TestProbs:
 
     def test_queue(self):
         """Test that the right measurement class is queued."""
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
 
         @qml.qnode(dev)
         def circuit():
@@ -97,7 +97,7 @@ def test_numeric_type(self):
     @pytest.mark.parametrize("shots", [None, 10])
     def test_shape(self, wires, shots):
         """Test that the shape is correct."""
-        dev = qml.device("default.qubit", wires=3, shots=shots)
+        dev = qml.device("default.qubit.legacy", wires=3, shots=shots)
         res = qml.probs(wires=wires)
         assert res.shape(dev, Shots(shots)) == (2 ** len(wires),)
 
@@ -106,7 +106,7 @@ def test_shape_shot_vector(self, wires):
         """Test that the shape is correct with the shot vector too."""
         res = qml.probs(wires=wires)
         shot_vector = (1, 2, 3)
-        dev = qml.device("default.qubit", wires=3, shots=shot_vector)
+        dev = qml.device("default.qubit.legacy", wires=3, shots=shot_vector)
         assert res.shape(dev, Shots(shot_vector)) == (
             (2 ** len(wires),),
             (2 ** len(wires),),
@@ -133,7 +133,7 @@ def test_probs_empty_wires(self):
     @pytest.mark.parametrize("shots", [None, 100])
     def test_probs_no_arguments(self, shots):
         """Test that using ``qml.probs`` with no arguments returns the probabilities of all wires."""
-        dev = qml.device("default.qubit", wires=3, shots=shots)
+        dev = qml.device("default.qubit.legacy", wires=3, shots=shots)
 
         @qml.qnode(dev)
         def circuit():
@@ -145,7 +145,7 @@ def circuit():
 
     def test_full_prob(self, init_state, tol, mocker):
         """Test that the correct probability is returned."""
-        dev = qml.device("default.qubit", wires=4)
+        dev = qml.device("default.qubit.legacy", wires=4)
         spy = mocker.spy(qml.QubitDevice, "probability")
 
         state = init_state(4)
@@ -163,7 +163,7 @@ def circuit():
 
     def test_marginal_prob(self, init_state, tol, mocker):
         """Test that the correct marginal probability is returned."""
-        dev = qml.device("default.qubit", wires=4)
+        dev = qml.device("default.qubit.legacy", wires=4)
         spy = mocker.spy(qml.QubitDevice, "probability")
 
         state = init_state(4)
@@ -183,7 +183,7 @@ def circuit():
     def test_marginal_prob_more_wires(self, init_state, mocker, tol):
         """Test that the correct marginal probability is returned, when the
         states_to_binary method is used for probability computations."""
-        dev = qml.device("default.qubit", wires=4)
+        dev = qml.device("default.qubit.legacy", wires=4)
         spy_probs = mocker.spy(qml.QubitDevice, "probability")
         state = init_state(4)
 
@@ -244,7 +244,7 @@ def test_process_state_batched(self, interface, subset_wires, expected):
 
     def test_integration(self, tol, mocker):
         """Test the probability is correct for a known state preparation."""
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
         spy = mocker.spy(qml.QubitDevice, "probability")
 
         @qml.qnode(dev)
@@ -266,7 +266,7 @@ def circuit():
     def test_integration_analytic_false(self, tol, mocker, shots):
         """Test the probability is correct for a known state preparation when the
         analytic attribute is set to False."""
-        dev = qml.device("default.qubit", wires=3, shots=shots)
+        dev = qml.device("default.qubit.legacy", wires=3, shots=shots)
         spy = mocker.spy(qml.QubitDevice, "probability")
 
         @qml.qnode(dev)
@@ -314,7 +314,7 @@ def circuit(phi):
     @pytest.mark.parametrize("shots", [None, 100])
     def test_batch_size(self, mocker, shots):
         """Test the probability is correct for a batched input."""
-        dev = qml.device("default.qubit", wires=1, shots=shots)
+        dev = qml.device("default.qubit.legacy", wires=1, shots=shots)
         spy = mocker.spy(qml.QubitDevice, "probability")
 
         @qml.qnode(dev)
@@ -355,7 +355,7 @@ def test_numerical_analytic_diff_agree(self, tol, mocker):
         """Test that the finite difference and parameter shift rule
         provide the same Jacobian."""
         w = 4
-        dev = qml.device("default.qubit", wires=w)
+        dev = qml.device("default.qubit.legacy", wires=w)
         spy = mocker.spy(qml.QubitDevice, "probability")
 
         def circuit(x, y, z):
@@ -391,7 +391,7 @@ def circuit(x, y, z):
     @pytest.mark.parametrize("hermitian", [1 / np.sqrt(2) * np.array([[1, 1], [1, -1]])])
     def test_prob_generalize_param_one_qubit(self, hermitian, tol, mocker):
         """Test that the correct probability is returned."""
-        dev = qml.device("default.qubit", wires=1)
+        dev = qml.device("default.qubit.legacy", wires=1)
         spy = mocker.spy(qml.QubitDevice, "probability")
 
         @qml.qnode(dev)
@@ -418,7 +418,7 @@ def circuit_rotated(x):
     @pytest.mark.parametrize("hermitian", [1 / np.sqrt(2) * np.array([[1, 1], [1, -1]])])
     def test_prob_generalize_param(self, hermitian, tol, mocker):
         """Test that the correct probability is returned."""
-        dev = qml.device("default.qubit", wires=3)
+        dev = qml.device("default.qubit.legacy", wires=3)
         spy = mocker.spy(qml.QubitDevice, "probability")
 
         @qml.qnode(dev)
@@ -450,7 +450,7 @@ def circuit_rotated(x, y):
     @pytest.mark.parametrize("hermitian", [1 / np.sqrt(2) * np.array([[1, 1], [1, -1]])])
     def test_prob_generalize_param_multiple(self, hermitian, tol, mocker):
         """Test that the correct probability is returned."""
-        dev = qml.device("default.qubit", wires=3)
+        dev = qml.device("default.qubit.legacy", wires=3)
         spy = mocker.spy(qml.QubitDevice, "probability")
 
         @qml.qnode(dev)
@@ -495,7 +495,7 @@ def circuit_rotated(x, y):
     def test_prob_generalize_initial_state(self, hermitian, wire, init_state, tol, mocker):
         """Test that the correct probability is returned."""
         # pylint:disable=too-many-arguments
-        dev = qml.device("default.qubit", wires=4)
+        dev = qml.device("default.qubit.legacy", wires=4)
         spy = mocker.spy(qml.QubitDevice, "probability")
         state = init_state(4)
 
@@ -539,7 +539,7 @@ def circuit_rotated():
     def test_operation_prob(self, operation, wire, init_state, tol, mocker):
         "Test the rotated probability with different wires and rotating operations."
         # pylint:disable=too-many-arguments
-        dev = qml.device("default.qubit", wires=4)
+        dev = qml.device("default.qubit.legacy", wires=4)
         spy = mocker.spy(qml.QubitDevice, "probability")
         state = init_state(4)
 
@@ -581,7 +581,7 @@ def circuit_rotated():
     @pytest.mark.parametrize("observable", [(qml.PauliX, qml.PauliY)])
     def test_observable_tensor_prob(self, observable, init_state, tol, mocker):
         "Test the rotated probability with a tensor observable."
-        dev = qml.device("default.qubit", wires=4)
+        dev = qml.device("default.qubit.legacy", wires=4)
         spy = mocker.spy(qml.QubitDevice, "probability")
         state = init_state(4)
 
@@ -619,7 +619,7 @@ def test_hamiltonian_error(self, coeffs, obs, init_state):
         "Test that an error is returned for hamiltonians."
         H = qml.Hamiltonian(coeffs, obs)
 
-        dev = qml.device("default.qubit", wires=4)
+        dev = qml.device("default.qubit.legacy", wires=4)
         state = init_state(4)
 
         @qml.qnode(dev)
@@ -644,7 +644,7 @@ def test_generalize_prob_not_hermitian(self, operation):
         """Test that Operators that do not have a diagonalizing_gates representation cannot
         be used in probability measurements."""
 
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
 
         @qml.qnode(dev)
         def circuit():
@@ -662,7 +662,7 @@ def circuit():
     def test_prob_wires_and_hermitian(self, hermitian):
         """Test that we can cannot give simultaneously wires and a hermitian."""
 
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
 
         @qml.qnode(dev)
         def circuit():
@@ -714,7 +714,7 @@ def test_estimate_probability_with_binsize_with_broadcasting(self, wires, expect
         assert np.allclose(res, expected)
 
     def test_non_commuting_probs_raises_error(self):
-        dev = qml.device("default.qubit", wires=5)
+        dev = qml.device("default.qubit.legacy", wires=5)
 
         @qml.qnode(dev)
         def circuit(x, y):
@@ -730,7 +730,7 @@ def circuit(x, y):
 
     def test_commuting_probs_in_computational_basis(self):
         """Test that `qml.probs` can be used in the computational basis with other commuting observables."""
-        dev = qml.device("default.qubit", wires=5)
+        dev = qml.device("default.qubit.legacy", wires=5)
 
         @qml.qnode(dev)
         def circuit(x, y):
diff --git a/tests/measurements/test_purity_measurement.py b/tests/measurements/test_purity_measurement.py
index ccaab4fb5bc..969f9a16c14 100644
--- a/tests/measurements/test_purity_measurement.py
+++ b/tests/measurements/test_purity_measurement.py
@@ -71,15 +71,15 @@ def test_numeric_type(self):
     def test_shape_new(self, shots, shape):
         """Test the ``shape_new`` method."""
         meas = qml.purity(wires=0)
-        dev = qml.device("default.qubit", wires=1, shots=shots)
+        dev = qml.device("default.qubit.legacy", wires=1, shots=shots)
         assert meas.shape(dev, Shots(shots)) == shape
 
 
 class TestPurityIntegration:
     """Test the purity meausrement with qnodes and devices."""
 
-    devices = ["default.qubit", "lightning.qubit", "default.mixed"]
-    grad_supported_devices = ["default.qubit", "default.mixed"]
+    devices = ["default.qubit.legacy", "lightning.qubit", "default.mixed"]
+    grad_supported_devices = ["default.qubit.legacy", "default.mixed"]
     mix_supported_devices = ["default.mixed"]
 
     diff_methods = ["backprop", "finite-diff"]
diff --git a/tests/measurements/test_sample.py b/tests/measurements/test_sample.py
index 1a0a4bb1e98..3dfcf34280b 100644
--- a/tests/measurements/test_sample.py
+++ b/tests/measurements/test_sample.py
@@ -51,7 +51,7 @@ class TestSample:  # pylint: disable=too-many-public-methods
     def test_sample_dimension(self, mocker, n_sample):
         """Test that the sample function outputs samples of the right size"""
 
-        dev = qml.device("default.qubit", wires=2, shots=n_sample)
+        dev = qml.device("default.qubit.legacy", wires=2, shots=n_sample)
         spy = mocker.spy(qml.QubitDevice, "sample")
 
         @qml.qnode(dev)
@@ -76,7 +76,7 @@ def test_sample_combination(self, mocker):
         """Test the output of combining expval, var and sample"""
         n_sample = 10
 
-        dev = qml.device("default.qubit", wires=3, shots=n_sample)
+        dev = qml.device("default.qubit.legacy", wires=3, shots=n_sample)
         spy = mocker.spy(qml.QubitDevice, "sample")
 
         @qml.qnode(dev, diff_method="parameter-shift")
@@ -99,7 +99,7 @@ def test_single_wire_sample(self, mocker):
         """Test the return type and shape of sampling a single wire"""
         n_sample = 10
 
-        dev = qml.device("default.qubit", wires=1, shots=n_sample)
+        dev = qml.device("default.qubit.legacy", wires=1, shots=n_sample)
         spy = mocker.spy(qml.QubitDevice, "sample")
 
         @qml.qnode(dev)
@@ -120,7 +120,7 @@ def test_multi_wire_sample_regular_shape(self, mocker):
         where a rectangular array is expected"""
         n_sample = 10
 
-        dev = qml.device("default.qubit", wires=3, shots=n_sample)
+        dev = qml.device("default.qubit.legacy", wires=3, shots=n_sample)
         spy = mocker.spy(qml.QubitDevice, "sample")
 
         @qml.qnode(dev)
@@ -146,7 +146,7 @@ def test_sample_output_type_in_combination(self, mocker):
         in combination with expvals and vars"""
         n_sample = 10
 
-        dev = qml.device("default.qubit", wires=3, shots=n_sample)
+        dev = qml.device("default.qubit.legacy", wires=3, shots=n_sample)
         spy = mocker.spy(qml.QubitDevice, "sample")
 
         @qml.qnode(dev, diff_method="parameter-shift")
@@ -167,7 +167,7 @@ def circuit():
     def test_not_an_observable(self, mocker):
         """Test that a UserWarning is raised if the provided
         argument might not be hermitian."""
-        dev = qml.device("default.qubit", wires=2, shots=10)
+        dev = qml.device("default.qubit.legacy", wires=2, shots=10)
         spy = mocker.spy(qml.QubitDevice, "sample")
 
         @qml.qnode(dev)
@@ -183,7 +183,7 @@ def circuit():
     def test_observable_return_type_is_sample(self, mocker):
         """Test that the return type of the observable is :attr:`ObservableReturnTypes.Sample`"""
         n_shots = 10
-        dev = qml.device("default.qubit", wires=1, shots=n_shots)
+        dev = qml.device("default.qubit.legacy", wires=1, shots=n_shots)
         spy = mocker.spy(qml.QubitDevice, "sample")
 
         @qml.qnode(dev)
@@ -223,7 +223,7 @@ def circuit(phi):
 
     def test_providing_observable_and_wires(self):
         """Test that a ValueError is raised if both an observable is provided and wires are specified"""
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
 
         @qml.qnode(dev)
         def circuit():
@@ -239,7 +239,7 @@ def circuit():
 
     def test_providing_no_observable_and_no_wires(self, mocker):
         """Test that we can provide no observable and no wires to sample function"""
-        dev = qml.device("default.qubit", wires=2, shots=1000)
+        dev = qml.device("default.qubit.legacy", wires=2, shots=1000)
         spy = mocker.spy(qml.QubitDevice, "sample")
 
         @qml.qnode(dev)
@@ -262,7 +262,7 @@ def test_providing_no_observable_and_no_wires_shot_vector(self, mocker):
         shots1 = 1
         shots2 = 10
         shots3 = 1000
-        dev = qml.device("default.qubit", wires=num_wires, shots=[shots1, shots2, shots3])
+        dev = qml.device("default.qubit.legacy", wires=num_wires, shots=[shots1, shots2, shots3])
         spy = mocker.spy(qml.QubitDevice, "sample")
 
         @qml.qnode(dev)
@@ -291,7 +291,7 @@ def test_providing_no_observable_and_wires(self, mocker):
         """Test that we can provide no observable but specify wires to the sample function"""
         wires = [0, 2]
         wires_obj = qml.wires.Wires(wires)
-        dev = qml.device("default.qubit", wires=3, shots=1000)
+        dev = qml.device("default.qubit.legacy", wires=3, shots=1000)
         spy = mocker.spy(qml.QubitDevice, "sample")
 
         @qml.qnode(dev)
@@ -343,7 +343,7 @@ def test_numeric_type(self, obs, exp):
 
     def test_shape_no_shots_error(self):
         """Test that the appropriate error is raised with no shots are specified"""
-        dev = qml.device("default.qubit", wires=2, shots=None)
+        dev = qml.device("default.qubit.legacy", wires=2, shots=None)
         shots = Shots(None)
         mp = qml.sample()
 
@@ -364,7 +364,7 @@ def test_shape_no_shots_error(self):
     def test_shape(self, obs):
         """Test that the shape is correct."""
         shots = 10
-        dev = qml.device("default.qubit", wires=3, shots=shots)
+        dev = qml.device("default.qubit.legacy", wires=3, shots=shots)
         res = qml.sample(obs) if obs is not None else qml.sample()
         expected = (shots,) if obs is not None else (shots, 3)
         assert res.shape(dev, Shots(shots)) == expected
@@ -372,7 +372,7 @@ def test_shape(self, obs):
     @pytest.mark.parametrize("n_samples", (1, 10))
     def test_shape_wires(self, n_samples):
         """Test that the shape is correct when wires are provided."""
-        dev = qml.device("default.qubit", wires=3, shots=n_samples)
+        dev = qml.device("default.qubit.legacy", wires=3, shots=n_samples)
         mp = qml.sample(wires=(0, 1))
         assert mp.shape(dev, Shots(n_samples)) == (n_samples, 2) if n_samples != 1 else (2,)
 
@@ -388,7 +388,7 @@ def test_shape_wires(self, n_samples):
     def test_shape_shot_vector(self, obs):
         """Test that the shape is correct with the shot vector too."""
         shot_vector = (1, 2, 3)
-        dev = qml.device("default.qubit", wires=3, shots=shot_vector)
+        dev = qml.device("default.qubit.legacy", wires=3, shots=shot_vector)
         res = qml.sample(obs) if obs is not None else qml.sample()
         expected = ((), (2,), (3,)) if obs is not None else ((3,), (2, 3), (3, 3))
         assert res.shape(dev, Shots(shot_vector)) == expected
@@ -396,7 +396,7 @@ def test_shape_shot_vector(self, obs):
     def test_shape_shot_vector_obs(self):
         """Test that the shape is correct with the shot vector and a observable too."""
         shot_vec = (2, 2)
-        dev = qml.device("default.qubit", wires=3, shots=shot_vec)
+        dev = qml.device("default.qubit.legacy", wires=3, shots=shot_vec)
 
         @qml.qnode(dev)
         def circuit():
@@ -419,7 +419,7 @@ def test_sample_empty_wires(self):
     @pytest.mark.parametrize("shots", [2, 100])
     def test_sample_no_arguments(self, shots):
         """Test that using ``qml.sample`` with no arguments returns the samples of all wires."""
-        dev = qml.device("default.qubit", wires=3, shots=shots)
+        dev = qml.device("default.qubit.legacy", wires=3, shots=shots)
 
         @qml.qnode(dev)
         def circuit():
@@ -450,7 +450,7 @@ def test_jitting_with_sampling_on_subset_of_wires(samples):
 
     jax.config.update("jax_enable_x64", True)
 
-    dev = qml.device("default.qubit", wires=3, shots=samples)
+    dev = qml.device("default.qubit.legacy", wires=3, shots=samples)
 
     @qml.qnode(dev, interface="jax")
     def circuit(x):
diff --git a/tests/measurements/test_state.py b/tests/measurements/test_state.py
index eebfb128ad1..811d0ad0a54 100644
--- a/tests/measurements/test_state.py
+++ b/tests/measurements/test_state.py
@@ -125,7 +125,7 @@ class TestState:
     def test_state_shape_and_dtype(self, wires):
         """Test that the state is of correct size and dtype for a trivial circuit"""
 
-        dev = qml.device("default.qubit", wires=wires)
+        dev = qml.device("default.qubit.legacy", wires=wires)
 
         @qml.qnode(dev)
         def func():
@@ -138,7 +138,7 @@ def func():
     def test_return_type_is_state(self):
         """Test that the return type of the observable is State"""
 
-        dev = qml.device("default.qubit", wires=1)
+        dev = qml.device("default.qubit.legacy", wires=1)
 
         @qml.qnode(dev)
         def func():
@@ -154,7 +154,7 @@ def func():
     def test_state_correct_ghz(self, wires):
         """Test that the correct state is returned when the circuit prepares a GHZ state"""
 
-        dev = qml.device("default.qubit", wires=wires)
+        dev = qml.device("default.qubit.legacy", wires=wires)
 
         @qml.qnode(dev)
         def func():
@@ -175,7 +175,7 @@ def test_return_with_other_types(self):
         """Test that an exception is raised when a state is returned along with another return
         type"""
 
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
 
         @qml.qnode(dev)
         def func():
@@ -193,7 +193,7 @@ def test_state_equal_to_dev_state(self, wires):
         """Test that the returned state is equal to the one stored in dev.state for a template
         circuit"""
 
-        dev = qml.device("default.qubit", wires=wires)
+        dev = qml.device("default.qubit.legacy", wires=wires)
 
         weights = np.random.random(
             qml.templates.StronglyEntanglingLayers.shape(n_layers=3, n_wires=wires)
@@ -212,7 +212,7 @@ def test_interface_tf(self):
         """Test that the state correctly outputs in the tensorflow interface"""
         import tensorflow as tf
 
-        dev = qml.device("default.qubit", wires=4)
+        dev = qml.device("default.qubit.legacy", wires=4)
 
         @qml.qnode(dev, interface="tf")
         def func():
@@ -233,7 +233,7 @@ def test_interface_torch(self):
         """Test that the state correctly outputs in the torch interface"""
         import torch
 
-        dev = qml.device("default.qubit", wires=4)
+        dev = qml.device("default.qubit.legacy", wires=4)
 
         @qml.qnode(dev, interface="torch")
         def func():
@@ -252,7 +252,7 @@ def func():
     @pytest.mark.autograd
     def test_jacobian_not_supported(self):
         """Test if an error is raised if the jacobian method is called via qml.grad"""
-        dev = qml.device("default.qubit", wires=4)
+        dev = qml.device("default.qubit.legacy", wires=4)
 
         @qml.qnode(dev, diff_method="parameter-shift")
         def func(x):
@@ -274,7 +274,7 @@ def func(x):
     def test_no_state_capability(self, monkeypatch):
         """Test if an error is raised for devices that are not capable of returning the state.
         This is tested by changing the capability of default.qubit"""
-        dev = qml.device("default.qubit", wires=1)
+        dev = qml.device("default.qubit.legacy", wires=1)
         capabilities = dev.capabilities().copy()
         capabilities["returns_state"] = False
 
@@ -304,7 +304,7 @@ def test_default_qubit(self, diff_method):
         """Test that the returned state is equal to the expected returned state for all of
         PennyLane's built in statevector devices"""
 
-        dev = qml.device("default.qubit", wires=4)
+        dev = qml.device("default.qubit.legacy", wires=4)
 
         @qml.qnode(dev, diff_method=diff_method)
         def func():
@@ -408,7 +408,7 @@ def loss_fn(x):
     @pytest.mark.parametrize("wires", [[0, 2, 3, 1], ["a", -1, "b", 1000]])
     def test_custom_wire_labels(self, wires):
         """Test the state when custom wire labels are used"""
-        dev = qml.device("default.qubit", wires=wires)
+        dev = qml.device("default.qubit.legacy", wires=wires)
 
         @qml.qnode(dev, diff_method="parameter-shift")
         def func():
@@ -424,14 +424,14 @@ def func():
     @pytest.mark.parametrize("shots", [None, 1, 10])
     def test_shape(self, shots):
         """Test that the shape is correct for qml.state."""
-        dev = qml.device("default.qubit", wires=3, shots=shots)
+        dev = qml.device("default.qubit.legacy", wires=3, shots=shots)
         res = qml.state()
         assert res.shape(dev, Shots(shots)) == (2**3,)
 
     @pytest.mark.parametrize("s_vec", [(3, 2, 1), (1, 5, 10), (3, 1, 20)])
     def test_shape_shot_vector(self, s_vec):
         """Test that the shape is correct for qml.state with the shot vector too."""
-        dev = qml.device("default.qubit", wires=3, shots=s_vec)
+        dev = qml.device("default.qubit.legacy", wires=3, shots=s_vec)
         res = qml.state()
         assert res.shape(dev, Shots(s_vec)) == ((2**3,), (2**3,), (2**3,))
 
@@ -447,7 +447,7 @@ class TestDensityMatrix:
     # pylint: disable=too-many-public-methods
 
     @pytest.mark.parametrize("wires", range(2, 5))
-    @pytest.mark.parametrize("dev_name", ["default.qubit", "default.mixed"])
+    @pytest.mark.parametrize("dev_name", ["default.qubit.legacy", "default.mixed"])
     def test_density_matrix_shape_and_dtype(self, dev_name, wires):
         """Test that the density matrix is of correct size and dtype for a
         trivial circuit"""
@@ -463,7 +463,7 @@ def circuit():
         assert state_val.shape == (2, 2)
         assert state_val.dtype == np.complex128
 
-    @pytest.mark.parametrize("dev_name", ["default.qubit", "default.mixed"])
+    @pytest.mark.parametrize("dev_name", ["default.qubit.legacy", "default.mixed"])
     def test_return_type_is_state(self, dev_name):
         """Test that the return type of the observable is State"""
 
@@ -480,7 +480,7 @@ def func():
         assert obs[0].return_type is State
 
     @pytest.mark.torch
-    @pytest.mark.parametrize("dev_name", ["default.qubit", "default.mixed"])
+    @pytest.mark.parametrize("dev_name", ["default.qubit.legacy", "default.mixed"])
     @pytest.mark.parametrize("diff_method", [None, "backprop"])
     def test_correct_density_matrix_torch(self, dev_name, diff_method):
         """Test that the correct density matrix is returned using torch interface."""
@@ -504,7 +504,7 @@ def func():
             )
 
     @pytest.mark.jax
-    @pytest.mark.parametrize("dev_name", ["default.qubit", "default.mixed"])
+    @pytest.mark.parametrize("dev_name", ["default.qubit.legacy", "default.mixed"])
     @pytest.mark.parametrize("diff_method", [None, "backprop"])
     def test_correct_density_matrix_jax(self, dev_name, diff_method):
         """Test that the correct density matrix is returned using JAX interface."""
@@ -527,7 +527,7 @@ def func():
             )
 
     @pytest.mark.tf
-    @pytest.mark.parametrize("dev_name", ["default.qubit", "default.mixed"])
+    @pytest.mark.parametrize("dev_name", ["default.qubit.legacy", "default.mixed"])
     @pytest.mark.parametrize("diff_method", [None, "backprop"])
     def test_correct_density_matrix_tf(self, dev_name, diff_method):
         """Test that the correct density matrix is returned using the TensorFlow interface."""
@@ -549,7 +549,7 @@ def func():
                 qml.density_matrix(wires=0).process_state(state=dev.state, wire_order=dev.wires),
             )
 
-    @pytest.mark.parametrize("dev_name", ["default.qubit", "default.mixed"])
+    @pytest.mark.parametrize("dev_name", ["default.qubit.legacy", "default.mixed"])
     def test_correct_density_matrix_product_state_first(self, dev_name):
         """Test that the correct density matrix is returned when
         tracing out a product state"""
@@ -573,7 +573,7 @@ def func():
                 qml.density_matrix(wires=0).process_state(state=dev.state, wire_order=dev.wires),
             )
 
-    @pytest.mark.parametrize("dev_name", ["default.qubit", "default.mixed"])
+    @pytest.mark.parametrize("dev_name", ["default.qubit.legacy", "default.mixed"])
     def test_correct_density_matrix_product_state_second(self, dev_name):
         """Test that the correct density matrix is returned when
         tracing out a product state"""
@@ -596,7 +596,7 @@ def func():
                 qml.density_matrix(wires=1).process_state(state=dev.state, wire_order=dev.wires),
             )
 
-    @pytest.mark.parametrize("dev_name", ["default.qubit", "default.mixed"])
+    @pytest.mark.parametrize("dev_name", ["default.qubit.legacy", "default.mixed"])
     @pytest.mark.parametrize("return_wire_order", ([0, 1], [1, 0]))
     def test_correct_density_matrix_product_state_both(self, dev_name, return_wire_order):
         """Test that the correct density matrix is returned
@@ -625,7 +625,7 @@ def func():
                 ),
             )
 
-    @pytest.mark.parametrize("dev_name", ["default.qubit", "default.mixed"])
+    @pytest.mark.parametrize("dev_name", ["default.qubit.legacy", "default.mixed"])
     def test_correct_density_matrix_three_wires_first_two(self, dev_name):
         """Test that the correct density matrix is returned for an example with three wires,
         and tracing out the third wire."""
@@ -657,7 +657,7 @@ def func():
                 ),
             )
 
-    @pytest.mark.parametrize("dev_name", ["default.qubit", "default.mixed"])
+    @pytest.mark.parametrize("dev_name", ["default.qubit.legacy", "default.mixed"])
     def test_correct_density_matrix_three_wires_last_two(self, dev_name):
         """Test that the correct density matrix is returned for an example with three wires,
         and tracing out the first wire."""
@@ -693,7 +693,7 @@ def func():
                 ),
             )
 
-    @pytest.mark.parametrize("dev_name", ["default.qubit", "default.mixed"])
+    @pytest.mark.parametrize("dev_name", ["default.qubit.legacy", "default.mixed"])
     @pytest.mark.parametrize(
         "return_wire_order", ([0], [1], [2], [0, 1], [1, 0], [0, 2], [2, 0], [1, 2, 0], [2, 1, 0])
     )
@@ -733,7 +733,7 @@ def func():
                 ),
             )
 
-    @pytest.mark.parametrize("dev_name", ["default.qubit", "default.mixed"])
+    @pytest.mark.parametrize("dev_name", ["default.qubit.legacy", "default.mixed"])
     def test_correct_density_matrix_mixed_state(self, dev_name):
         """Test that the correct density matrix for an example with a mixed state"""
 
@@ -749,7 +749,7 @@ def func():
 
         assert np.allclose(np.array([[0.5 + 0.0j, 0.0 + 0.0j], [0.0 + 0.0j, 0.5 + 0.0j]]), density)
 
-    @pytest.mark.parametrize("dev_name", ["default.qubit", "default.mixed"])
+    @pytest.mark.parametrize("dev_name", ["default.qubit.legacy", "default.mixed"])
     def test_correct_density_matrix_all_wires(self, dev_name):
         """Test that the correct density matrix is returned when all wires are given"""
 
@@ -781,7 +781,7 @@ def func():
                 ),
             )
 
-    @pytest.mark.parametrize("dev_name", ["default.qubit", "default.mixed"])
+    @pytest.mark.parametrize("dev_name", ["default.qubit.legacy", "default.mixed"])
     def test_return_with_other_types(self, dev_name):
         """Test that an exception is raised when a state is returned along with another return
         type"""
@@ -804,7 +804,7 @@ def func():
     def test_no_state_capability(self, monkeypatch):
         """Test if an error is raised for devices that are not capable of returning
         the density matrix. This is tested by changing the capability of default.qubit"""
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
         capabilities = dev.capabilities().copy()
         capabilities["returns_state"] = False
 
@@ -833,7 +833,7 @@ def func():
             func()
 
     @pytest.mark.parametrize("wires", [[0, 2], ["a", -1]])
-    @pytest.mark.parametrize("dev_name", ["default.qubit", "default.mixed"])
+    @pytest.mark.parametrize("dev_name", ["default.qubit.legacy", "default.mixed"])
     def test_custom_wire_labels(self, wires, dev_name):
         """Test that the correct density matrix for an example with a mixed
         state when using custom wires"""
@@ -860,7 +860,7 @@ def func():
             )
 
     @pytest.mark.parametrize("wires", [[3, 1], ["b", 1000]])
-    @pytest.mark.parametrize("dev_name", ["default.qubit", "default.mixed"])
+    @pytest.mark.parametrize("dev_name", ["default.qubit.legacy", "default.mixed"])
     def test_custom_wire_labels_all_wires(self, wires, dev_name):
         """Test that the correct density matrix for an example with a mixed
         state when using custom wires"""
@@ -889,14 +889,14 @@ def func():
     @pytest.mark.parametrize("shots", [None, 1, 10])
     def test_shape(self, shots):
         """Test that the shape is correct for qml.density_matrix."""
-        dev = qml.device("default.qubit", wires=3, shots=shots)
+        dev = qml.device("default.qubit.legacy", wires=3, shots=shots)
         res = qml.density_matrix(wires=[0, 1])
         assert res.shape(dev, Shots(shots)) == (2**2, 2**2)
 
     @pytest.mark.parametrize("s_vec", [(3, 2, 1), (1, 5, 10), (3, 1, 20)])
     def test_shape_shot_vector(self, s_vec):
         """Test that the shape is correct for qml.density_matrix with the shot vector too."""
-        dev = qml.device("default.qubit", wires=3, shots=s_vec)
+        dev = qml.device("default.qubit.legacy", wires=3, shots=s_vec)
         res = qml.density_matrix(wires=[0, 1])
         assert res.shape(dev, Shots(s_vec)) == (
             (2**2, 2**2),
diff --git a/tests/measurements/test_var.py b/tests/measurements/test_var.py
index f89f9a31477..331defbae60 100644
--- a/tests/measurements/test_var.py
+++ b/tests/measurements/test_var.py
@@ -51,7 +51,7 @@ class TestVar:
     @pytest.mark.parametrize("r_dtype", [np.float32, np.float64])
     def test_value(self, tol, r_dtype, mocker, shots):
         """Test that the var function works"""
-        dev = qml.device("default.qubit", wires=2, shots=shots)
+        dev = qml.device("default.qubit.legacy", wires=2, shots=shots)
         spy = mocker.spy(qml.QubitDevice, "var")
         dev.R_DTYPE = r_dtype
 
@@ -79,7 +79,7 @@ def circuit(x):
     def test_not_an_observable(self, mocker):
         """Test that a UserWarning is raised if the provided
         argument might not be hermitian."""
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
         spy = mocker.spy(qml.QubitDevice, "var")
 
         @qml.qnode(dev)
@@ -94,7 +94,7 @@ def circuit():
 
     def test_observable_return_type_is_variance(self, mocker):
         """Test that the return type of the observable is :attr:`ObservableReturnTypes.Variance`"""
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
         spy = mocker.spy(qml.QubitDevice, "var")
 
         @qml.qnode(dev)
@@ -147,7 +147,7 @@ def test_numeric_type(self, obs):
     )
     def test_shape(self, obs):
         """Test that the shape is correct."""
-        dev = qml.device("default.qubit", wires=1)
+        dev = qml.device("default.qubit.legacy", wires=1)
         res = qml.var(obs)
         # pylint: disable=use-implicit-booleaness-not-comparison
         assert res.shape(dev, Shots(None)) == ()
@@ -161,14 +161,14 @@ def test_shape_shot_vector(self, obs):
         """Test that the shape is correct with the shot vector too."""
         res = qml.var(obs)
         shot_vector = (1, 2, 3)
-        dev = qml.device("default.qubit", wires=3, shots=shot_vector)
+        dev = qml.device("default.qubit.legacy", wires=3, shots=shot_vector)
         assert res.shape(dev, Shots(shot_vector)) == ((), (), ())
 
     @pytest.mark.parametrize("state", [np.array([0, 0, 0]), np.array([1, 0, 0, 0, 0, 0, 0, 0])])
     @pytest.mark.parametrize("shots", [None, 1000, [1000, 10000]])
     def test_projector_var(self, state, shots, mocker):
         """Tests that the variance of a ``Projector`` object is computed correctly."""
-        dev = qml.device("default.qubit", wires=3, shots=shots)
+        dev = qml.device("default.qubit.legacy", wires=3, shots=shots)
         spy = mocker.spy(qml.QubitDevice, "var")
 
         @qml.qnode(dev)
@@ -190,7 +190,7 @@ def test_permuted_wires(self, mocker):
             qml.s_prod(3, qml.PauliZ("h")), qml.PauliZ(8), qml.s_prod(2, qml.PauliZ(10))
         )
 
-        dev = qml.device("default.qubit", wires=["h", 8, 10])
+        dev = qml.device("default.qubit.legacy", wires=["h", 8, 10])
         spy = mocker.spy(qml.QubitDevice, "var")
 
         @qml.qnode(dev)
diff --git a/tests/measurements/test_vn_entropy.py b/tests/measurements/test_vn_entropy.py
index d9551f3b34a..e69e88e5256 100644
--- a/tests/measurements/test_vn_entropy.py
+++ b/tests/measurements/test_vn_entropy.py
@@ -76,7 +76,7 @@ class TestInitialization:
     @pytest.mark.parametrize("interface", ["autograd", "jax", "tf", "torch"])
     def test_vn_entropy(self, interface, state_vector, expected):
         """Tests the output of qml.vn_entropy"""
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
 
         @qml.qnode(dev, interface=interface)
         def circuit():
@@ -94,7 +94,7 @@ def circuit():
 
     def test_queue(self):
         """Test that the right measurement class is queued."""
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
 
         @qml.qnode(dev)
         def circuit():
@@ -121,7 +121,7 @@ def test_properties(self):
     def test_shape(self, shots, shape):
         """Test the ``shape`` method."""
         meas = qml.vn_entropy(wires=0)
-        dev = qml.device("default.qubit", wires=1, shots=shots)
+        dev = qml.device("default.qubit.legacy", wires=1, shots=shots)
 
         assert meas.shape(dev, Shots(shots)) == shape
 
@@ -131,7 +131,7 @@ class TestIntegration:
 
     parameters = np.linspace(0, 2 * np.pi, 10)
 
-    devices = ["default.qubit", "default.mixed", "lightning.qubit"]
+    devices = ["default.qubit.legacy", "default.mixed", "lightning.qubit"]
 
     single_wires_list = [
         [0],
@@ -142,12 +142,12 @@ class TestIntegration:
 
     check_state = [True, False]
 
-    devices = ["default.qubit", "default.mixed"]
+    devices = ["default.qubit.legacy", "default.mixed"]
     diff_methods = ["backprop", "finite-diff"]
 
     def test_shot_vec_error(self):
         """Test an error is raised when using shot vectors with vn_entropy."""
-        dev = qml.device("default.qubit", wires=2, shots=[1, 10, 10, 1000])
+        dev = qml.device("default.qubit.legacy", wires=2, shots=[1, 10, 10, 1000])
 
         @qml.qnode(device=dev)
         def circuit(x):
@@ -187,7 +187,7 @@ def circuit_entropy(x):
     def test_IsingXX_qnode_entropy_grad(self, param, wires, base, diff_method):
         """Test entropy for a QNode gradient with autograd."""
 
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
 
         @qml.qnode(dev, diff_method=diff_method)
         def circuit_entropy(x):
@@ -234,7 +234,7 @@ def test_IsingXX_qnode_entropy_grad_torch(self, param, wires, base, diff_method)
         """Test entropy for a QNode gradient with torch."""
         import torch
 
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
 
         @qml.qnode(dev, interface="torch", diff_method=diff_method)
         def circuit_entropy(x):
@@ -286,7 +286,7 @@ def test_IsingXX_qnode_entropy_grad_tf(self, param, wires, base, diff_method, in
         """Test entropy for a QNode gradient with tf."""
         import tensorflow as tf
 
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
 
         @qml.qnode(dev, interface=interface, diff_method=diff_method)
         def circuit_entropy(x):
@@ -339,7 +339,7 @@ def test_IsingXX_qnode_entropy_grad_jax(self, param, wires, base, diff_method, i
         """Test entropy for a QNode gradient with Jax."""
         import jax
 
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
 
         @qml.qnode(dev, interface=interface, diff_method=diff_method)
         def circuit_entropy(x):
@@ -388,7 +388,7 @@ def test_IsingXX_qnode_entropy_grad_jax_jit(self, param, wires, base, diff_metho
         """Test entropy for a QNode gradient with Jax-jit."""
         import jax
 
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
 
         @qml.qnode(dev, interface=interface, diff_method=diff_method)
         def circuit_entropy(x):
diff --git a/tests/qinfo/test_reduced_dm.py b/tests/qinfo/test_reduced_dm.py
index da76fd5a256..3427bf1e066 100644
--- a/tests/qinfo/test_reduced_dm.py
+++ b/tests/qinfo/test_reduced_dm.py
@@ -158,7 +158,7 @@ def circuit(x):
     def test_density_matrix_c_dtype(self, wires, c_dtype):
         """Test different complex dtype."""
 
-        dev = qml.device("default.qubit", wires=2, c_dtype=c_dtype)
+        dev = qml.device("default.qubit.legacy", wires=2, c_dtype=c_dtype)
 
         @qml.qnode(dev, diff_method=None)
         def circuit(x):
diff --git a/tests/test_debugging.py b/tests/test_debugging.py
index 2bcec651d36..42f59ca5f8e 100644
--- a/tests/test_debugging.py
+++ b/tests/test_debugging.py
@@ -24,9 +24,9 @@ class TestSnapshot:
 
     # pylint: disable=protected-access
     @pytest.mark.parametrize("method", [None, "backprop", "parameter-shift", "adjoint"])
-    def test_default_qubit(self, method):
+    def test_default_qubit_legacy(self, method):
         """Test that multiple snapshots are returned correctly on the state-vector simulator."""
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
 
         @qml.qnode(dev, diff_method=method)
         def circuit():
@@ -190,7 +190,7 @@ def circuit():
 
     def test_unsupported_device(self):
         """Test that an error is raised on unsupported devices."""
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
         # remove attributes to simulate unsupported device
         delattr(dev, "_debugger")
         dev.operations.remove("Snapshot")
diff --git a/tests/test_device.py b/tests/test_device.py
index d153cd21d1a..c3c8b0c253d 100644
--- a/tests/test_device.py
+++ b/tests/test_device.py
@@ -561,7 +561,7 @@ def test_default_expand_with_initial_state(self, op, decomp):
         prep = [op]
         ops = [qml.AngleEmbedding(features=[0.1], wires=[0], rotation="Z"), op, qml.PauliZ(wires=2)]
 
-        dev = qml.device("default.qubit", wires=3)
+        dev = qml.device("default.qubit.legacy", wires=3)
         tape = qml.tape.QuantumTape(ops=ops, measurements=[], prep=prep, shots=100)
         new_tape = dev.default_expand_fn(tape)
 
@@ -918,7 +918,7 @@ def test_outdated_API(self, monkeypatch):
         with monkeypatch.context() as m:
             m.setattr(qml, "version", lambda: "0.0.1")
             with pytest.raises(DeviceError, match="plugin requires PennyLane versions"):
-                qml.device("default.qubit", wires=0)
+                qml.device("default.qubit.legacy", wires=0)
 
     @pytest.mark.skip(reason="Reloading PennyLane messes with tape mode")
     def test_refresh_entrypoints(self, monkeypatch):
@@ -975,7 +975,7 @@ def test_hot_refresh_entrypoints(self, monkeypatch):
 
     def test_shot_vector_property(self):
         """Tests shot vector initialization."""
-        dev = qml.device("default.qubit", wires=1, shots=[1, 3, 3, 4, 4, 4, 3])
+        dev = qml.device("default.qubit.legacy", wires=1, shots=[1, 3, 3, 4, 4, 4, 3])
         shot_vector = dev.shot_vector
         assert len(shot_vector) == 4
         assert shot_vector[0].shots == 1
diff --git a/tests/test_operation.py b/tests/test_operation.py
index 7c69f4579b1..d484ef26e6e 100644
--- a/tests/test_operation.py
+++ b/tests/test_operation.py
@@ -1010,7 +1010,7 @@ def test_all_wires_defined_but_init_with_one(self):
         """Test that an exception is raised if the class is defined with ALL wires,
         but then instantiated with only one"""
 
-        dev1 = qml.device("default.qubit", wires=2)
+        dev1 = qml.device("default.qubit.legacy", wires=2)
 
         class DummyOp(qml.operation.Operation):
             r"""Dummy custom operator"""
@@ -2686,7 +2686,7 @@ class CustomOperator(qml.operation.Operator):
     assert qml.equal(new_op, CustomOperator(2.3, wires=0))
 
 
-# pylint: disable=unused-import
+# pylint: disable=unused-import,no-name-in-module
 def test_get_attr():
     """Test that importing attributes of operation work as expected"""
 
diff --git a/tests/test_qubit_device.py b/tests/test_qubit_device.py
index 18b93e37858..37f53e26cbd 100644
--- a/tests/test_qubit_device.py
+++ b/tests/test_qubit_device.py
@@ -1148,7 +1148,7 @@ def test_device_executions(self):
         """Test the number of times a qubit device is executed over a QNode's
         lifetime is tracked by `num_executions`"""
 
-        dev_1 = qml.device("default.qubit", wires=2)
+        dev_1 = qml.device("default.qubit.legacy", wires=2)
 
         def circuit_1(x, y):
             qml.RX(x, wires=[0])
@@ -1164,7 +1164,7 @@ def circuit_1(x, y):
         assert dev_1.num_executions == num_evals_1
 
         # test a second instance of a default qubit device
-        dev_2 = qml.device("default.qubit", wires=2)
+        dev_2 = qml.device("default.qubit.legacy", wires=2)
 
         def circuit_2(x):
             qml.RX(x, wires=[0])
@@ -1209,7 +1209,7 @@ def test_device_executions(self):
         """Test the number of times a qubit device is executed over a QNode's
         lifetime is tracked by `num_executions`"""
 
-        dev_1 = qml.device("default.qubit", wires=2)
+        dev_1 = qml.device("default.qubit.legacy", wires=2)
 
         def circuit_1(x, y):
             qml.RX(x, wires=[0])
@@ -1225,7 +1225,7 @@ def circuit_1(x, y):
         assert dev_1.num_executions == num_evals_1
 
         # test a second instance of a default qubit device
-        dev_2 = qml.device("default.qubit", wires=2)
+        dev_2 = qml.device("default.qubit.legacy", wires=2)
 
         assert dev_2.num_executions == 0
 
@@ -1392,7 +1392,7 @@ class TestResourcesTracker:
     )  # Resources(wires, gates, gate_types, gate_sizes, depth, shots)
 
     devices = (
-        "default.qubit",
+        "default.qubit.legacy",
         "default.qubit.autograd",
         "default.qubit.jax",
         "default.qubit.torch",
@@ -1441,7 +1441,7 @@ def test_tracker_multi_execution(self, dev_name):
     @pytest.mark.all_interfaces
     def test_tracker_grad(self):
         """Test that the tracker can track resources through a gradient computation"""
-        dev = qml.device("default.qubit", wires=1, shots=100)
+        dev = qml.device("default.qubit.legacy", wires=1, shots=100)
 
         @qml.qnode(dev, diff_method="parameter-shift")
         def circuit(x):
diff --git a/tests/transforms/test_batch_transform.py b/tests/transforms/test_batch_transform.py
index 448eaf2b7bd..742e15bc78b 100644
--- a/tests/transforms/test_batch_transform.py
+++ b/tests/transforms/test_batch_transform.py
@@ -340,7 +340,7 @@ def test_parametrized_transform_device(self, mocker):
         b = 0.4
         x = 0.543
 
-        dev = qml.device("default.qubit", wires=1)
+        dev = qml.device("default.qubit.legacy", wires=1)
         dev = self.my_transform(dev, a, b)
 
         @qml.qnode(dev, interface="autograd")
@@ -371,7 +371,7 @@ def test_parametrized_transform_device_decorator(self, mocker):
         b = 0.4
         x = 0.543
 
-        dev = qml.device("default.qubit", wires=1)
+        dev = qml.device("default.qubit.legacy", wires=1)
         dev = self.my_transform(a, b)(dev)  # pylint: disable=no-value-for-parameter
 
         @qml.qnode(dev, interface="autograd")
@@ -679,7 +679,7 @@ class TestMapBatchTransform:
 
     def test_result(self, mocker):
         """Test that it correctly applies the transform to be mapped"""
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
         H = qml.PauliZ(0) @ qml.PauliZ(1) - qml.PauliX(0)
         x = 0.6
         y = 0.7
@@ -713,7 +713,7 @@ def test_result(self, mocker):
 
     def test_differentiation(self):
         """Test that an execution using map_batch_transform can be differentiated"""
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
         H = qml.PauliZ(0) @ qml.PauliZ(1) - qml.PauliX(0)
 
         weights = np.array([0.6, 0.8], requires_grad=True)
diff --git a/tests/transforms/test_tape_expand.py b/tests/transforms/test_tape_expand.py
index 258acb743c1..82f5ecdfb3a 100644
--- a/tests/transforms/test_tape_expand.py
+++ b/tests/transforms/test_tape_expand.py
@@ -64,7 +64,7 @@ def test_device_and_stopping_expansion(self):
         """Test that passing a device alongside a stopping condition ensures
         that all operations are expanded to match the devices default gate
         set"""
-        dev = qml.device("default.qubit", wires=1)
+        dev = qml.device("default.qubit.legacy", wires=1)
         expand_fn = qml.transforms.create_expand_fn(device=dev, depth=10, stop_at=self.crit_0)
 
         with qml.queuing.AnnotatedQueue() as q:
@@ -80,7 +80,7 @@ def test_device_and_stopping_expansion(self):
     def test_device_only_expansion(self):
         """Test that passing a device ensures that all operations are expanded
         to match the devices default gate set"""
-        dev = qml.device("default.qubit", wires=1)
+        dev = qml.device("default.qubit.legacy", wires=1)
         expand_fn = qml.transforms.create_expand_fn(device=dev, depth=10)
 
         with qml.queuing.AnnotatedQueue() as q:
@@ -432,8 +432,8 @@ def circuit():
             qml.CNOT(wires=[0, 1])
             return qml.expval(qml.PauliZ(0))
 
-        original_dev = qml.device("default.qubit", wires=3)
-        decomp_dev = qml.device("default.qubit", wires=3, custom_decomps={})
+        original_dev = qml.device("default.qubit.legacy", wires=3)
+        decomp_dev = qml.device("default.qubit.legacy", wires=3, custom_decomps={})
 
         original_qnode = qml.QNode(circuit, original_dev, expansion_strategy="device")
         decomp_qnode = qml.QNode(circuit, decomp_dev, expansion_strategy="device")
@@ -457,8 +457,8 @@ def circuit():
             qml.BasicEntanglerLayers([[0.1, 0.2]], wires=[0, 1])
             return qml.expval(qml.PauliZ(0))
 
-        original_dev = qml.device("default.qubit", wires=3)
-        decomp_dev = qml.device("default.qubit", wires=3, custom_decomps={})
+        original_dev = qml.device("default.qubit.legacy", wires=3)
+        decomp_dev = qml.device("default.qubit.legacy", wires=3, custom_decomps={})
 
         original_qnode = qml.QNode(circuit, original_dev, expansion_strategy="device")
         decomp_qnode = qml.QNode(circuit, decomp_dev, expansion_strategy="device")
@@ -474,7 +474,7 @@ def circuit():
             )
         ]
 
-    @pytest.mark.parametrize("device_name", ["default.qubit", "lightning.qubit"])
+    @pytest.mark.parametrize("device_name", ["default.qubit.legacy", "lightning.qubit"])
     def test_one_custom_decomp(self, device_name):
         """Test that specifying a single custom decomposition works as expected."""
 
@@ -509,7 +509,7 @@ def circuit():
 
         custom_decomps = {"Hadamard": custom_hadamard, "CNOT": custom_cnot}
         decomp_dev = qml.device(
-            "default.qubit", wires=2, custom_decomps=custom_decomps, decomp_depth=0
+            "default.qubit.legacy", wires=2, custom_decomps=custom_decomps, decomp_depth=0
         )
         decomp_qnode = qml.QNode(circuit, decomp_dev, expansion_strategy="device")
         _ = decomp_qnode()
@@ -529,8 +529,8 @@ def circuit(x):
             qml.Hadamard(wires=0)
             return qml.expval(qml.PauliZ(0))
 
-        original_dev = qml.device("default.qubit", wires=3)
-        decomp_dev = qml.device("default.qubit", wires=3, custom_decomps={"Rot": custom_rot})
+        original_dev = qml.device("default.qubit.legacy", wires=3)
+        decomp_dev = qml.device("default.qubit.legacy", wires=3, custom_decomps={"Rot": custom_rot})
 
         original_qnode = qml.QNode(circuit, original_dev, expansion_strategy="device")
         decomp_qnode = qml.QNode(circuit, decomp_dev, expansion_strategy="device")
@@ -558,7 +558,7 @@ def circuit():
             return qml.expval(qml.PauliZ(0))
 
         custom_decomps = {"Hadamard": custom_hadamard, qml.CNOT: custom_cnot}
-        decomp_dev = qml.device("default.qubit", wires=2, custom_decomps=custom_decomps)
+        decomp_dev = qml.device("default.qubit.legacy", wires=2, custom_decomps=custom_decomps)
         decomp_qnode = qml.QNode(circuit, decomp_dev, expansion_strategy="device")
         _ = decomp_qnode()
         decomp_ops = decomp_qnode.qtape.operations
@@ -596,7 +596,7 @@ def circuit():
             return qml.expval(qml.PauliZ(0))
 
         custom_decomps = {"Hadamard": custom_hadamard, qml.CNOT: custom_cnot}
-        decomp_dev = qml.device("default.qubit", wires=2, custom_decomps=custom_decomps)
+        decomp_dev = qml.device("default.qubit.legacy", wires=2, custom_decomps=custom_decomps)
         decomp_qnode = qml.QNode(circuit, decomp_dev, expansion_strategy="device")
         _ = decomp_qnode()
         decomp_ops = decomp_qnode.qtape.operations
@@ -640,7 +640,7 @@ def circuit():
 
         # BasicEntanglerLayers custom decomposition involves AngleEmbedding
         custom_decomps = {"BasicEntanglerLayers": custom_basic_entangler_layers, "RX": custom_rx}
-        decomp_dev = qml.device("default.qubit", wires=2, custom_decomps=custom_decomps)
+        decomp_dev = qml.device("default.qubit.legacy", wires=2, custom_decomps=custom_decomps)
         decomp_qnode = qml.QNode(circuit, decomp_dev, expansion_strategy="device")
         _ = decomp_qnode()
         decomp_ops = decomp_qnode.qtape.operations
@@ -677,7 +677,7 @@ def circuit():
         # not be further decomposed even though the custom decomposition is specified.
         custom_decomps = {"BasicEntanglerLayers": custom_basic_entangler_layers, "RX": custom_rx}
         decomp_dev = qml.device(
-            "default.qubit", wires=2, custom_decomps=custom_decomps, decomp_depth=2
+            "default.qubit.legacy", wires=2, custom_decomps=custom_decomps, decomp_depth=2
         )
         decomp_qnode = qml.QNode(circuit, decomp_dev, expansion_strategy="device")
         _ = decomp_qnode()
@@ -706,7 +706,7 @@ def circuit():
             return qml.expval(qml.PauliZ("a"))
 
         custom_decomps = {qml.RX: custom_rx}
-        decomp_dev = qml.device("default.qubit", wires="a", custom_decomps=custom_decomps)
+        decomp_dev = qml.device("default.qubit.legacy", wires="a", custom_decomps=custom_decomps)
         decomp_qnode = qml.QNode(circuit, decomp_dev, expansion_strategy="device")
         _ = decomp_qnode()
         decomp_ops = decomp_qnode.qtape.operations
@@ -732,7 +732,7 @@ def compute_decomposition(wire):
                 return [qml.S(wire)]
 
         custom_decomps = {CustomOp: lambda wires: [qml.T(wires), qml.T(wires)]}
-        decomp_dev = qml.device("default.qubit", wires=2, custom_decomps=custom_decomps)
+        decomp_dev = qml.device("default.qubit.legacy", wires=2, custom_decomps=custom_decomps)
 
         @qml.qnode(decomp_dev, expansion_strategy="device")
         def circuit():
@@ -751,7 +751,7 @@ def circuit():
     def test_custom_decomp_in_separate_context(self):
         """Test that the set_decomposition context manager works."""
 
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
 
         @qml.qnode(dev, expansion_strategy="device")
         def circuit():

From da51a8310de4f85a34d5edcc91bbb7fd37589898 Mon Sep 17 00:00:00 2001
From: Matthew Silverman <matthews@xanadu.ai>
Date: Tue, 5 Sep 2023 17:27:29 -0400
Subject: [PATCH 031/131] StateMP accepts wires (#4570)

**Context:**
In the new device interface, a device need not have wires. In the case
when it does provide wires, it would be best if measurement processes
that don't require wires at all (`state()`, `sample()`, `probs()`) know
to expand over the extra wires. To be able to do this, they need to save
the wires somewhere, and `self.wires` kinda makes sense for this:)

**Description of the Change:**
- Allow initialization of a `StateMP` with wires (now that it doesn't
indicate a density-matrix measurement).
- When given a strict subset of wires, `process_state` pads the state
with zeros for all other wire axes
- When given a different `wire_order` from `self.wires`, `process_state`
will transpose the state to match the MP wires

**Note on wire ordering:**
I use the convention that `StateMP.wires` is the global wire order we'd
like to return when used (so we can swap out measurement processes
[here](https://github.com/PennyLaneAI/pennylane/blob/078c4503b87e988f0ca68344ff57687d36f3e7fa/pennylane/devices/experimental/default_qubit_2.py#L217)),
and the `wire_order` given to `process_state` is the wire order of the
`state` being passed in (which happens
[here](https://github.com/PennyLaneAI/pennylane/blob/078c4503b87e988f0ca68344ff57687d36f3e7fa/pennylane/devices/qubit/simulate.py#L91),
at a point where we do not know about the device or its wires).

**Changes for `pad` dispatching:**
1. I'm not sure what `qml.math.pad(x, like="torch")` was even calling,
but the docs suggest to use `torch.nn.functional.pad` so I registered it
2. The tensorflow `pad` assumes padding over more than one dimension, so
I "hacked" it to add an empty broadcast dimension if needed (in other
words, turn `[pad]` into `[[pad]]`)

**Benefits:**
- `StateMP.process_state` is more powerful!
- Allows me to open a follow-up PR where DQ2 expands results for the
above listed MPs

**Possible Drawbacks:**
- Doesn't support batching (yet, I can add in this PR if necessary, just
wanted something that worked)

---------

Co-authored-by: Romain Moyard <rmoyard@gmail.com>
---
 doc/releases/changelog-dev.md     |   5 ++
 pennylane/math/single_dispatch.py |   9 +++
 pennylane/measurements/state.py   |  42 +++++++++--
 tests/measurements/test_state.py  | 115 ++++++++++++++++++++++++++++--
 4 files changed, 160 insertions(+), 11 deletions(-)

diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 291159a6df2..a2e9a186306 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -49,6 +49,11 @@
   process, `DensityMatrixMP`.
   [(#4558)](https://github.com/PennyLaneAI/pennylane/pull/4558)
 
+* The `StateMP` measurement now accepts a wire order (eg. a device wire order). The `process_state`
+  method will re-order the given state to go from the inputted wire-order to the process's wire-order.
+  If the process's wire-order contains extra wires, it will assume those are in the zero-state.
+  [(#4570)](https://github.com/PennyLaneAI/pennylane/pull/4570)
+
 <h3>Breaking changes 💔</h3>
 
 * The `__eq__` and `__hash__` methods of `Operator` and `MeasurementProcess` no longer rely on the
diff --git a/pennylane/math/single_dispatch.py b/pennylane/math/single_dispatch.py
index 1ca2a96dfde..b8722551daf 100644
--- a/pennylane/math/single_dispatch.py
+++ b/pennylane/math/single_dispatch.py
@@ -456,6 +456,15 @@ def _cond_tf(pred, true_fn, false_fn, args):
 )
 
 
+def _pad_tf(tensor, paddings, mode="CONSTANT", constant_values=0, name=None):
+    if ar.ndim(paddings) == 1:
+        paddings = (paddings,)
+    return _i("tf").pad(tensor, paddings, mode=mode, constant_values=constant_values, name=name)
+
+
+ar.register_function("tensorflow", "pad", _pad_tf)
+
+
 # -------------------------------- Torch --------------------------------- #
 
 ar.autoray._FUNC_ALIASES["torch", "unstack"] = "unbind"
diff --git a/pennylane/measurements/state.py b/pennylane/measurements/state.py
index b8d0c212593..e268d4840b9 100644
--- a/pennylane/measurements/state.py
+++ b/pennylane/measurements/state.py
@@ -17,7 +17,7 @@
 from typing import Sequence, Optional
 
 import pennylane as qml
-from pennylane.wires import Wires
+from pennylane.wires import Wires, WireError
 
 from .measurements import State, StateMeasurement
 
@@ -131,12 +131,13 @@ class StateMP(StateMeasurement):
     Please refer to :func:`state` for detailed documentation.
 
     Args:
+        wires (.Wires): The wires the measurement process applies to.
         id (str): custom label given to a measurement instance, can be useful for some applications
             where the instance has to be identified
     """
 
-    def __init__(self, *, id: Optional[str] = None):
-        super().__init__(wires=None, id=id)
+    def __init__(self, wires: Optional[Wires] = None, id: Optional[str] = None):
+        super().__init__(wires=wires, id=id)
 
     @property
     def return_type(self):
@@ -155,7 +156,37 @@ def shape(self, device, shots):
 
     def process_state(self, state: Sequence[complex], wire_order: Wires):
         # pylint:disable=redefined-outer-name
-        return state
+        wires = self.wires
+        if not wires or wire_order == wires:
+            return state
+
+        if not wires.contains_wires(wire_order):
+            raise WireError(
+                f"Unexpected wires {set(wire_order) - set(wires)} found in wire order. Expected wire order to be a subset of {wires}"
+            )
+
+        # pad with zeros, put existing wires last
+        is_torch = qml.math.get_interface(state) == "torch"
+        is_state_batched = qml.math.ndim(state) == 2
+        pad_width = 2 ** len(wires) - 2 ** len(wire_order)
+        pad = (pad_width, 0) if is_torch else (0, pad_width)
+        shape = (2,) * len(wires)
+        if is_state_batched:
+            pad = ((0, 0), pad)
+            shape = (qml.math.shape(state)[0],) + shape
+        elif is_torch:
+            pad = (pad,)
+
+        state = qml.math.pad(state, pad)
+        state = qml.math.reshape(state, shape)
+
+        # re-order
+        new_wire_order = Wires.unique_wires([wires, wire_order]) + wire_order
+        desired_axes = [new_wire_order.index(w) for w in wires]
+        if is_state_batched:
+            desired_axes = [0] + [i + 1 for i in desired_axes]
+        state = qml.math.transpose(state, desired_axes)
+        return qml.math.flatten(state)
 
 
 class DensityMatrixMP(StateMP):
@@ -170,8 +201,7 @@ class DensityMatrixMP(StateMP):
     """
 
     def __init__(self, wires: Wires, id: Optional[str] = None):
-        # pylint:disable=non-parent-init-called,super-init-not-called
-        StateMeasurement.__init__(self, wires=wires, id=id)
+        super().__init__(wires=wires, id=id)
 
     def shape(self, device, shots):
         num_shot_elements = (
diff --git a/tests/measurements/test_state.py b/tests/measurements/test_state.py
index 811d0ad0a54..909d1276d59 100644
--- a/tests/measurements/test_state.py
+++ b/tests/measurements/test_state.py
@@ -29,6 +29,7 @@
 )
 from pennylane.math.quantum import reduce_statevector, reduce_dm
 from pennylane.math.matrix_manipulation import _permute_dense_matrix
+from pennylane.wires import Wires, WireError
 
 
 class TestStateMP:
@@ -45,17 +46,121 @@ class TestStateMP:
     def test_process_state_vector(self, vec):
         """Test the processing of a state vector."""
 
-        mp = StateMP()
+        mp = StateMP(wires=None)
         assert mp.return_type == State
         assert mp.numeric_type is complex
 
         processed = mp.process_state(vec, None)
         assert qml.math.allclose(processed, vec)
 
-    def test_state_does_not_accept_wires(self):
-        """Test that StateMP does not accept wires."""
-        with pytest.raises(TypeError, match="unexpected keyword argument 'wires'"):
-            StateMP(wires=[0])
+    @pytest.mark.all_interfaces
+    @pytest.mark.parametrize("interface", ["numpy", "autograd", "jax", "torch", "tensorflow"])
+    def test_state_returns_itself_if_wires_match(self, interface):
+        """Test that when wire_order matches the StateMP, the state is returned."""
+        ket = qml.math.array([0.48j, 0.48, -0.64j, 0.36], like=interface)
+        assert StateMP(wires=[1, 0]).process_state(ket, wire_order=Wires([1, 0])) is ket
+
+    @pytest.mark.all_interfaces
+    @pytest.mark.parametrize("interface", ["numpy", "autograd", "jax", "torch", "tensorflow"])
+    @pytest.mark.parametrize("wires, wire_order", [([1, 0], [0, 1]), (["b", "a"], ["a", "b"])])
+    def test_reorder_state(self, interface, wires, wire_order):
+        """Test that a state can be re-ordered."""
+        ket = qml.math.array([0.48j, 0.48, -0.64j, 0.36], like=interface)
+        result = StateMP(wires=wires).process_state(ket, wire_order=Wires(wire_order))
+        assert qml.math.allclose(result, np.array([0.48j, -0.64j, 0.48, 0.36]))
+        assert qml.math.get_interface(ket) == interface
+
+    @pytest.mark.parametrize(
+        "mp_wires, expected_state",
+        [
+            ([0, 1, 2], [1, 0, 2, 0, 3, 0, 4, 0]),
+            ([2, 0, 1], [1, 2, 3, 4, 0, 0, 0, 0]),
+            ([1, 0, 2], [1, 0, 3, 0, 2, 0, 4, 0]),
+            ([1, 2, 0], [1, 3, 0, 0, 2, 4, 0, 0]),
+        ],
+    )
+    @pytest.mark.parametrize("custom_wire_labels", [False, True])
+    def test_expand_state_over_wires(self, mp_wires, expected_state, custom_wire_labels):
+        """Test the expanded state is correctly ordered with extra wires being zero."""
+        wire_order = [0, 1]
+        if custom_wire_labels:
+            # non-lexicographical-ordered
+            wire_map = {0: "b", 1: "c", 2: "a"}
+            mp_wires = [wire_map[w] for w in mp_wires]
+            wire_order = ["b", "c"]
+        mp = StateMP(wires=mp_wires)
+        ket = np.arange(1, 5)
+        result = mp.process_state(ket, wire_order=Wires(wire_order))
+        assert np.array_equal(result, expected_state)
+
+    @pytest.mark.all_interfaces
+    @pytest.mark.parametrize("interface", ["numpy", "autograd", "jax", "torch", "tensorflow"])
+    def test_expand_state_all_interfaces(self, interface):
+        """Test that expanding the state over wires preserves interface."""
+        mp = StateMP(wires=[4, 2, 0, 1])
+        ket = qml.math.array([0.48j, 0.48, -0.64j, 0.36], like=interface)
+        result = mp.process_state(ket, wire_order=Wires([1, 2]))
+        reshaped = qml.math.reshape(result, (2, 2, 2, 2))
+        assert qml.math.all(reshaped[1, :, 1, :] == 0)
+        assert qml.math.allclose(reshaped[0, :, 0, :], np.array([[0.48j, -0.64j], [0.48, 0.36]]))
+        if interface != "autograd":
+            # autograd.numpy.pad drops pennylane tensor for some reason
+            assert qml.math.get_interface(result) == interface
+
+    @pytest.mark.all_interfaces
+    @pytest.mark.parametrize("interface", ["numpy", "autograd", "jax", "torch", "tensorflow"])
+    def test_expand_state_batched_all_interfaces(self, interface):
+        """Test that expanding the state over wires preserves interface."""
+        mp = StateMP(wires=[4, 2, 0, 1])
+        ket = qml.math.array(
+            [
+                [0.48j, 0.48, -0.64j, 0.36],
+                [0.3, 0.4, 0.5, 1 / np.sqrt(2)],
+                [-0.3, -0.4, -0.5, -1 / np.sqrt(2)],
+            ],
+            like=interface,
+        )
+        result = mp.process_state(ket, wire_order=Wires([1, 2]))
+        reshaped = qml.math.reshape(result, (3, 2, 2, 2, 2))
+        assert qml.math.all(reshaped[:, 1, :, 1, :] == 0)
+        assert qml.math.allclose(
+            reshaped[:, 0, :, 0, :],
+            np.array(
+                [
+                    [[0.48j, -0.64j], [0.48, 0.36]],
+                    [[0.3, 0.5], [0.4, 1 / np.sqrt(2)]],
+                    [[-0.3, -0.5], [-0.4, -1 / np.sqrt(2)]],
+                ],
+            ),
+        )
+        if interface != "autograd":
+            # autograd.numpy.pad drops pennylane tensor for some reason
+            assert qml.math.get_interface(result) == interface
+
+    @pytest.mark.jax
+    @pytest.mark.parametrize(
+        "wires,expected",
+        [
+            ([1, 0], np.array([0.48j, -0.64j, 0.48, 0.36])),
+            ([2, 1, 0], np.array([0.48j, -0.64j, 0.48, 0.36, 0.0, 0.0, 0.0, 0.0])),
+        ],
+    )
+    def test_state_jax_jit(self, wires, expected):
+        """Test that re-ordering and expanding works with jax-jit."""
+        import jax
+
+        @jax.jit
+        def get_state(ket):
+            return StateMP(wires=wires).process_state(ket, wire_order=Wires([0, 1]))
+
+        result = get_state(jax.numpy.array([0.48j, 0.48, -0.64j, 0.36]))
+        assert qml.math.allclose(result, expected)
+        assert isinstance(result, jax.Array)
+
+    def test_wire_ordering_error(self):
+        """Test that a wire order error is raised when unknown wires are given."""
+        with pytest.raises(WireError, match=r"Unexpected wires \{2\} found in wire order"):
+            StateMP(wires=[0, 1]).process_state([1, 0], wire_order=[2])
 
 
 class TestDensityMatrixMP:

From 361caf3f6f08184b24df044331e085ab8acdcbc8 Mon Sep 17 00:00:00 2001
From: Edward Jiang <34989448+eddddddy@users.noreply.github.com>
Date: Wed, 6 Sep 2023 12:29:29 -0400
Subject: [PATCH 032/131] Deprecate fancy decorator syntax in batch transforms
 (#4457)

**Context:**
Currently we allow the user to do something like
```python
@qml.batch_input(argnum=0)
@qml.qnode(dev)
def circuit(inputs):
    ...
```
The only reason this works is because of some custom magic in
`qml.batch_transform` that wraps the arguments in a `partial` call and
returns that wrapper.

With the new transform program we want to switch to using the following
syntax:
```python
@partial(qml.batch_input, argnum=0)
@qml.qnode(dev)
def circuit(inputs):
    ...
```
which works identically but is easier to understand directly from the
documentation of `qml.batch_input`.

**Description of the Change:**
- Add deprecation warning for the use of the first syntax
- Allow the first syntax in the new transform API so that existing batch
transforms can be ported over without breaking things

**Benefits:**
Less confusing

**Possible Drawbacks:**
Slightly more typing for users

---------

Co-authored-by: Romain Moyard <rmoyard@gmail.com>
Co-authored-by: Christina Lee <christina@xanadu.ai>
Co-authored-by: Mudit Pandey <mudit.pandey@xanadu.ai>
---
 doc/development/deprecations.rst              | 23 +++++++++++
 doc/releases/changelog-dev.md                 | 18 +++++++++
 .../transforms/core/transform_dispatcher.py   | 26 +++++++++++-
 .../test_transform_dispatcher.py              | 40 +++++++++++++------
 4 files changed, 93 insertions(+), 14 deletions(-)

diff --git a/doc/development/deprecations.rst b/doc/development/deprecations.rst
index 4fa916327f3..03eb3009f1e 100644
--- a/doc/development/deprecations.rst
+++ b/doc/development/deprecations.rst
@@ -149,6 +149,29 @@ Completed deprecation cycles
   - Deprecated in v0.32
   - Removed in v0.33
 
+* The following decorator syntax for transforms has been deprecated:
+
+  .. code-block:: python
+
+      @transform_fn(**transform_kwargs)
+      @qml.qnode(dev)
+      def circuit():
+          ...
+
+  If you are using a transform that has supporting ``transform_kwargs``, please call the
+  transform directly using ``circuit = transform_fn(circuit, **transform_kwargs)``,
+  or use ``functools.partial``:
+
+  .. code-block:: python
+
+      @functools.partial(transform_fn, **transform_kwargs)
+      @qml.qnode(dev)
+      def circuit():
+          ...
+
+  - Deprecated in v0.33
+  - Will be removed in v0.34
+
 * The ``mode`` keyword argument in ``QNode`` has been removed, as it was only used in the old return
   system (which has also been removed). Please use ``grad_on_execution`` instead.
 
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index a2e9a186306..e3bdd2d632b 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -138,6 +138,24 @@
   ``StatePrepBase`` operations should be placed at the beginning of the `ops` list instead.
   [(#4554)](https://github.com/PennyLaneAI/pennylane/pull/4554)
 
+* The following decorator syntax for transforms has been deprecated and will raise a warning:
+  ```python
+  @transform_fn(**transform_kwargs)
+  @qml.qnode(dev)
+  def circuit():
+      ...
+  ```
+  If you are using a transform that has supporting `transform_kwargs`, please call the
+  transform directly using `circuit = transform_fn(circuit, **transform_kwargs)`,
+  or use `functools.partial`:
+  ```python
+  @functools.partial(transform_fn, **transform_kwargs)
+  @qml.qnode(dev)
+  def circuit():
+      ...
+  ```
+  [(#4457)](https://github.com/PennyLaneAI/pennylane/pull/4457/)
+
 <h3>Documentation 📝</h3>
 
 * Minor documentation improvements to the new device API. The documentation now correctly states that interface-specific
diff --git a/pennylane/transforms/core/transform_dispatcher.py b/pennylane/transforms/core/transform_dispatcher.py
index b10d5646a50..7bbcaac9bdf 100644
--- a/pennylane/transforms/core/transform_dispatcher.py
+++ b/pennylane/transforms/core/transform_dispatcher.py
@@ -15,6 +15,7 @@
 This module contains the transform function, the transform dispatcher and the transform container.
 """
 import copy
+import warnings
 import types
 
 import pennylane as qml
@@ -72,10 +73,31 @@ def __call__(self, *targs, **tkwargs):
         if callable(obj):
             return self._qfunc_transform(obj, targs, tkwargs)
 
-        raise TransformError(
-            "The object on which the transform is applied is not valid. It can only be a tape, a QNode or a qfunc."
+        # Input is not a QNode nor a quantum tape nor a device.
+        # Assume Python decorator syntax:
+        #
+        # result = some_transform(*transform_args)(qnode)(*qnode_args)
+
+        warnings.warn(
+            "Decorating a QNode with @transform_fn(**transform_kwargs) has been "
+            "deprecated and will be removed in a future version. Please decorate "
+            "with @functools.partial(transform_fn, **transform_kwargs) instead, "
+            "or call the transform directly using qnode = transform_fn(qnode, **transform_kwargs)",
+            UserWarning,
         )
 
+        if obj is not None:
+            targs = (obj, *targs)
+
+        def wrapper(obj):
+            return self(obj, *targs, **tkwargs)
+
+        wrapper.__doc__ = (
+            f"Partial of transform {self._transform} with bound arguments and keyword arguments."
+        )
+
+        return wrapper
+
     @property
     def transform(self):
         """Return the quantum transform."""
diff --git a/tests/transforms/test_experimental/test_transform_dispatcher.py b/tests/transforms/test_experimental/test_transform_dispatcher.py
index 39ced38d45f..2e5e6300fbc 100644
--- a/tests/transforms/test_experimental/test_transform_dispatcher.py
+++ b/tests/transforms/test_experimental/test_transform_dispatcher.py
@@ -174,7 +174,7 @@ def qnode_circuit(a):
         assert not dispatched_transform.is_informative
 
     @pytest.mark.parametrize("valid_transform", valid_transforms)
-    def test_integration_dispatcher_with_valid_transform_decorator(self, valid_transform):
+    def test_integration_dispatcher_with_valid_transform_decorator_partial(self, valid_transform):
         """Test that no error is raised with the transform function and that the transform dispatcher returns
         the right object."""
 
@@ -197,6 +197,33 @@ def qnode_circuit(a):
             qnode_circuit.transform_program.pop_front(), qml.transforms.core.TransformContainer
         )
 
+    @pytest.mark.parametrize("valid_transform", valid_transforms)
+    def test_integration_dispatcher_with_valid_transform_decorator(self, valid_transform):
+        """Test that a warning is raised with the transform function and that the transform dispatcher returns
+        the right object."""
+
+        dispatched_transform = transform(valid_transform)
+        targs = [0]
+
+        msg = r"Decorating a QNode with @transform_fn\(\*\*transform_kwargs\) has been deprecated"
+        with pytest.warns(UserWarning, match=msg):
+
+            @dispatched_transform(targs)
+            @qml.qnode(device=dev)
+            def qnode_circuit(a):
+                """QNode circuit."""
+                qml.Hadamard(wires=0)
+                qml.CNOT(wires=[0, 1])
+                qml.PauliX(wires=0)
+                qml.RZ(a, wires=1)
+                return qml.expval(qml.PauliZ(wires=0))
+
+        assert isinstance(qnode_circuit, qml.QNode)
+        assert isinstance(qnode_circuit.transform_program, qml.transforms.core.TransformProgram)
+        assert isinstance(
+            qnode_circuit.transform_program.pop_front(), qml.transforms.core.TransformContainer
+        )
+
     def test_queuing_qfunc_transform(self):
         """Test that queuing works with the transformed quantum function."""
 
@@ -339,17 +366,6 @@ def test_cotransform_not_implemented(self):
         ):
             transform(first_valid_transform, classical_cotransform=non_callable)
 
-    def test_apply_dispatched_transform_non_valid_obj(self):
-        """Test that applying a dispatched function on a non-valid object raises an error."""
-        dispatched_transform = transform(first_valid_transform)
-        obj = qml.RX(0.1, wires=0)
-        with pytest.raises(
-            TransformError,
-            match="The object on which the transform is applied is not valid. It can only be a tape, a QNode or a "
-            "qfunc.",
-        ):
-            dispatched_transform(obj)
-
     def test_qfunc_transform_multiple_tapes(self):
         """Test that quantum function is not compatible with multiple tapes."""
         dispatched_transform = transform(second_valid_transform)

From 126bd7b844ead70c253831e2b404757494012316 Mon Sep 17 00:00:00 2001
From: Romain Moyard <rmoyard@gmail.com>
Date: Wed, 6 Sep 2023 13:21:51 -0400
Subject: [PATCH 033/131] Batch transforms are updated (#4440)

**Description of the Change:**

All qfuncs transforms need to be updated with the new transform system.

- [x] batch_input
- [x]  batch_params
- [x] split_non_commuting
- [x] _replace_obs
- [x] broadcast_expand
- [x] algebra_commutator
- [x] fold_global
- [x] mitigate_with_zne
- [x] cut_circuit
- [x] cut_circuit_mc
- [x] make_probs
- [x] sign_expand
- [x] hamiltonian_expand
- [x] sum_expand

**Benefits:**

Batch transforms are compatible with transform program

**TODO**

1) Use the decorator on the function
2) Add the type hints
3) Add direct qnode tests

---------

Co-authored-by: Edward Jiang <34989448+eddddddy@users.noreply.github.com>
Co-authored-by: Edward Jiang <eddy3.14159@gmail.com>
Co-authored-by: Mudit Pandey <mudit.pandey@xanadu.ai>
Co-authored-by: Matthew Silverman <matthews@xanadu.ai>
---
 doc/code/qml_qinfo.rst                        |   1 +
 doc/releases/changelog-dev.md                 |   3 +
 pennylane/optimize/riemannian_gradient.py     |   5 +-
 pennylane/qinfo/transforms.py                 |  20 ++-
 pennylane/shadows/transforms.py               |   8 +-
 pennylane/transforms/batch_input.py           |  10 +-
 pennylane/transforms/batch_params.py          |  10 +-
 pennylane/transforms/broadcast_expand.py      |   8 +-
 pennylane/transforms/core/transform.py        |   6 +-
 .../transforms/core/transform_dispatcher.py   |   7 +-
 .../transforms/core/transform_program.py      |   4 +-
 pennylane/transforms/hamiltonian_expand.py    |  15 +-
 pennylane/transforms/mitigate.py              |  32 ++--
 pennylane/transforms/qcut/__init__.py         |   3 +-
 pennylane/transforms/qcut/cutcircuit.py       |  52 +++---
 pennylane/transforms/qcut/montecarlo.py       | 151 +++++++++---------
 .../transforms/sign_expand/sign_expand.py     |  14 +-
 pennylane/transforms/split_non_commuting.py   |   7 +-
 tests/devices/test_default_mixed_jax.py       |   4 +-
 tests/qinfo/test_fisher.py                    |   9 +-
 tests/shadow/test_shadow_transforms.py        |  41 +++--
 tests/transforms/test_batch_input.py          |  50 +++---
 tests/transforms/test_batch_params.py         |   4 +-
 tests/transforms/test_broadcast_expand.py     |  24 +++
 .../test_transform_dispatcher.py              |  22 ++-
 .../test_transform_program.py                 |   6 +-
 tests/transforms/test_mitigate.py             |  81 +++++++---
 tests/transforms/test_qcut.py                 |  98 ++++++------
 tests/transforms/test_sign_expand.py          |  35 ++++
 tests/transforms/test_split_non_commuting.py  |  34 +++-
 30 files changed, 460 insertions(+), 304 deletions(-)

diff --git a/doc/code/qml_qinfo.rst b/doc/code/qml_qinfo.rst
index 7235adc42fa..80501373d86 100644
--- a/doc/code/qml_qinfo.rst
+++ b/doc/code/qml_qinfo.rst
@@ -17,3 +17,4 @@ Transforms
     :no-inherited-members:
     :skip: metric_tensor
     :skip: adjoint_metric_tensor
+    :skip: transform
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index e3bdd2d632b..a4d274d5e53 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -27,6 +27,9 @@
   (array(0.6), array([1, 1, 1, 0, 1]))
   ```
 
+* All batch transforms are updated to the new transform program system.
+  [(4440)](https://github.com/PennyLaneAI/pennylane/pull/4440)
+
 <h3>Improvements 🛠</h3>
 
 * Add the method ``add_transform`` and ``insert_front_transform`` transform in the ``TransformProgram``.
diff --git a/pennylane/optimize/riemannian_gradient.py b/pennylane/optimize/riemannian_gradient.py
index 26530dbce0f..27e4c88d021 100644
--- a/pennylane/optimize/riemannian_gradient.py
+++ b/pennylane/optimize/riemannian_gradient.py
@@ -69,7 +69,6 @@ def append_time_evolution(tape, riemannian_gradient, t, n, exact=False):
         qml.apply(obj)
 
 
-@qml.batch_transform
 def algebra_commutator(tape, observables, lie_algebra_basis_names, nqubits):
     """Calculate the Riemannian gradient in the Lie algebra with the parameter shift rule
     (see :meth:`RiemannianGradientOptimizer.get_omegas`).
@@ -133,7 +132,7 @@ def algebra_commutator(tape, observables, lie_algebra_basis_names, nqubits):
                     for q, p in zip(queues_min, lie_algebra_basis_names)
                 ]
             )
-    return tapes_plus_total + tapes_min_total, None
+    return tapes_plus_total + tapes_min_total
 
 
 class RiemannianGradientOptimizer:
@@ -366,7 +365,7 @@ def get_omegas(self):
             obs_groupings,
             self.lie_algebra_basis_names,
             self.nqubits,
-        )[0]
+        )
         circuits = qml.execute(circuits, self.circuit.device, gradient_fn=None)
         circuits_plus = np.array(circuits[: len(circuits) // 2]).reshape(
             len(self.coeffs), len(self.lie_algebra_basis_names)
diff --git a/pennylane/qinfo/transforms.py b/pennylane/qinfo/transforms.py
index 07853563557..44a3e95dbc1 100644
--- a/pennylane/qinfo/transforms.py
+++ b/pennylane/qinfo/transforms.py
@@ -14,11 +14,13 @@
 """QNode transforms for the quantum information quantities."""
 # pylint: disable=import-outside-toplevel, not-callable
 import functools
+from typing import Sequence, Callable
 
 import pennylane as qml
 from pennylane.devices import DefaultQubit
 from pennylane.measurements import StateMP
-from pennylane.transforms import adjoint_metric_tensor, batch_transform, metric_tensor
+from pennylane.transforms import adjoint_metric_tensor, metric_tensor
+from pennylane.transforms.core import transform
 
 
 def reduced_dm(qnode, wires):
@@ -334,19 +336,15 @@ def _compute_cfim(p, dp):
     return dp_over_p @ dp
 
 
-@batch_transform
-def _make_probs(tape, wires=None, post_processing_fn=None):
+@transform
+def _make_probs(tape: qml.tape.QuantumTape) -> (Sequence[qml.tape.QuantumTape], Callable):
     """Ignores the return types of the provided circuit and creates a new one
     that outputs probabilities"""
-    qscript = qml.tape.QuantumScript(
-        tape.operations, [qml.probs(wires=wires or tape.wires)], shots=tape.shots
-    )
+    qscript = qml.tape.QuantumScript(tape.operations, [qml.probs(tape.wires)], shots=tape.shots)
 
-    if post_processing_fn is None:
-
-        def post_processing_fn(res):
-            # only a single probs measurement, so no stacking needed
-            return res[0]
+    def post_processing_fn(res):
+        # only a single probs measurement, so no stacking needed
+        return res[0]
 
     return [qscript], post_processing_fn
 
diff --git a/pennylane/shadows/transforms.py b/pennylane/shadows/transforms.py
index a16564504b0..43d6d535eb9 100644
--- a/pennylane/shadows/transforms.py
+++ b/pennylane/shadows/transforms.py
@@ -16,15 +16,17 @@
 import warnings
 from functools import reduce, wraps
 from itertools import product
+from typing import Sequence, Callable
 
 import pennylane as qml
 import pennylane.numpy as np
 from pennylane.measurements import ClassicalShadowMP
 from pennylane.tape import QuantumScript, QuantumTape
+from pennylane.transforms.core import transform
 
 
-@qml.batch_transform
-def _replace_obs(tape: QuantumTape, obs, *args, **kwargs):
+@transform
+def _replace_obs(tape: QuantumTape, obs, *args, **kwargs) -> (Sequence[QuantumTape], Callable):
     """
     Tape transform to replace the measurement processes with the given one
     """
@@ -88,7 +90,7 @@ def circuit(x):
     """
 
     def decorator(qnode):
-        return wraps(qnode)(_replace_obs(qnode, qml.shadow_expval, H, k=k))
+        return _replace_obs(qnode, qml.shadow_expval, H, k=k)
 
     return decorator
 
diff --git a/pennylane/transforms/batch_input.py b/pennylane/transforms/batch_input.py
index df18ec3c8f7..7effe8f5678 100644
--- a/pennylane/transforms/batch_input.py
+++ b/pennylane/transforms/batch_input.py
@@ -13,20 +13,20 @@
 """
 Batch transformation for multiple (non-trainable) input examples following issue #2037
 """
-from typing import Callable, Sequence, Tuple, Union
+from typing import Callable, Sequence, Union
 
 import pennylane as qml
 from pennylane import numpy as np
 from pennylane.tape import QuantumTape
-from pennylane.transforms.batch_transform import batch_transform
+from pennylane.transforms.core import transform
 from pennylane.transforms.batch_params import _nested_stack, _split_operations
 
 
-@batch_transform
+@transform
 def batch_input(
-    tape: Union[QuantumTape, qml.QNode],
+    tape: QuantumTape,
     argnum: Union[Sequence[int], int],
-) -> Tuple[Sequence[QuantumTape], Callable]:
+) -> (Sequence[QuantumTape], Callable):
     """
     Transform a QNode to support an initial batch dimension for gate inputs.
 
diff --git a/pennylane/transforms/batch_params.py b/pennylane/transforms/batch_params.py
index 07610a76b4f..05dfe70fb3e 100644
--- a/pennylane/transforms/batch_params.py
+++ b/pennylane/transforms/batch_params.py
@@ -15,10 +15,12 @@
 Contains the batch dimension transform.
 """
 # pylint: disable=import-outside-toplevel
+from typing import Callable, Sequence
+
 import pennylane as qml
 
 
-from .batch_transform import batch_transform
+from .core import transform
 
 
 def _nested_stack(res):
@@ -78,8 +80,10 @@ def _split_operations(ops, params, split_indices, num_tapes):
     return new_ops
 
 
-@batch_transform
-def batch_params(tape, all_operations=False):
+@transform
+def batch_params(
+    tape: qml.tape.QuantumTape, all_operations=False
+) -> (Sequence[qml.tape.QuantumTape], Callable):
     """Transform a QNode to support an initial batch dimension
     for operation parameters.
 
diff --git a/pennylane/transforms/broadcast_expand.py b/pennylane/transforms/broadcast_expand.py
index 047251f5c9e..40eebc6e137 100644
--- a/pennylane/transforms/broadcast_expand.py
+++ b/pennylane/transforms/broadcast_expand.py
@@ -13,8 +13,10 @@
 # limitations under the License.
 """This module contains the tape expansion function for expanding a
 broadcasted tape into multiple tapes."""
+from typing import Sequence, Callable
+
 import pennylane as qml
-from .batch_transform import batch_transform
+from .core import transform
 
 
 def _split_operations(ops, num_tapes):
@@ -43,8 +45,8 @@ def _split_operations(ops, num_tapes):
     return new_ops
 
 
-@batch_transform
-def broadcast_expand(tape):
+@transform
+def broadcast_expand(tape: qml.tape.QuantumTape) -> (Sequence[qml.tape.QuantumTape], Callable):
     r"""Expand a broadcasted tape into multiple tapes
     and a function that stacks and squeezes the results.
 
diff --git a/pennylane/transforms/core/transform.py b/pennylane/transforms/core/transform.py
index bf403bd929a..3fd69ce8e0e 100644
--- a/pennylane/transforms/core/transform.py
+++ b/pennylane/transforms/core/transform.py
@@ -37,7 +37,7 @@ def transform(
         expand_transform(callable): An expand transform is defined as a function that has the following requirements:
 
             * An expand transform is a function that is applied before applying the defined quantum transform. It
-              takes a quantum tape as single input and returns a single tape in a sequence with a dummy processing
+              takes the same arguments as the transform and returns a single tape in a sequence with a dummy processing
               function.
 
             * The expand transform must have the same type hinting as a quantum transform.
@@ -116,9 +116,9 @@ def qnode_circuit(a):
         # Check the signature of expand_transform to force the fn style tape - > (Sequence(tape), fn)
         _transform_signature_check(signature_expand_transform)
 
-        if len(signature_expand_transform) > 2:
+        if signature_expand_transform != signature_transform:
             raise TransformError(
-                "The expand transform does not support arg and kwargs other than tape."
+                "The expand transform must have the same signature as the transform"
             )
 
     # 3: CHeck the classical co-transform
diff --git a/pennylane/transforms/core/transform_dispatcher.py b/pennylane/transforms/core/transform_dispatcher.py
index 7bbcaac9bdf..c3f8c9f2d88 100644
--- a/pennylane/transforms/core/transform_dispatcher.py
+++ b/pennylane/transforms/core/transform_dispatcher.py
@@ -61,9 +61,8 @@ def __call__(self, *targs, **tkwargs):
             # is the object we wish to transform
             obj, *targs = targs
 
-        if isinstance(obj, qml.tape.QuantumTape):
-            new_tape = copy.deepcopy(obj)
-            return self._transform(new_tape, *targs, **tkwargs)
+        if isinstance(obj, qml.tape.QuantumScript):
+            return self._transform(obj, *targs, **tkwargs)
         if isinstance(obj, qml.QNode):
             return self._qnode_transform(
                 obj,
@@ -157,7 +156,7 @@ def default_qnode_transform(self, qnode, targs, tkwargs):
         qnode = copy.deepcopy(qnode)
 
         if self.expand_transform:
-            qnode.add_transform(TransformContainer(self._expand_transform))
+            qnode.add_transform(TransformContainer(self._expand_transform, targs, tkwargs))
         qnode.add_transform(
             TransformContainer(
                 self._transform, targs, tkwargs, self._classical_cotransform, self._is_informative
diff --git a/pennylane/transforms/core/transform_program.py b/pennylane/transforms/core/transform_program.py
index 7a260f70ce0..ef181cd9573 100644
--- a/pennylane/transforms/core/transform_program.py
+++ b/pennylane/transforms/core/transform_program.py
@@ -185,7 +185,7 @@ def add_transform(self, transform: TransformDispatcher, *targs, **tkwargs):
             raise TransformError("Only transform dispatcher can be added to the transform program.")
 
         if transform.expand_transform:
-            self.push_back(TransformContainer(transform.expand_transform))
+            self.push_back(TransformContainer(transform.expand_transform, targs, tkwargs))
         self.push_back(
             TransformContainer(
                 transform.transform,
@@ -223,7 +223,7 @@ def insert_front_transform(self, transform: TransformDispatcher, *targs, **tkwar
         )
 
         if transform.expand_transform:
-            self.insert_front(TransformContainer(transform.expand_transform))
+            self.insert_front(TransformContainer(transform.expand_transform, targs, tkwargs))
 
     def pop_front(self):
         """Pop the transform container at the beginning of the program.
diff --git a/pennylane/transforms/hamiltonian_expand.py b/pennylane/transforms/hamiltonian_expand.py
index 5d7e703d073..8c3dd49e1cd 100644
--- a/pennylane/transforms/hamiltonian_expand.py
+++ b/pennylane/transforms/hamiltonian_expand.py
@@ -15,15 +15,17 @@
 Contains the hamiltonian expand tape transform
 """
 # pylint: disable=protected-access
-from typing import List
+from typing import List, Sequence, Callable
 
 import pennylane as qml
 from pennylane.measurements import ExpectationMP, MeasurementProcess
 from pennylane.ops import SProd, Sum
 from pennylane.tape import QuantumScript, QuantumTape
+from pennylane.transforms.core import transform
 
 
-def hamiltonian_expand(tape: QuantumTape, group=True):
+@transform
+def hamiltonian_expand(tape: QuantumTape, group: bool = True) -> (Sequence[QuantumTape], Callable):
     r"""
     Splits a tape measuring a Hamiltonian expectation into mutliple tapes of Pauli expectations,
     and provides a function to recombine the results.
@@ -35,7 +37,7 @@ def hamiltonian_expand(tape: QuantumTape, group=True):
             If grouping information can be found in the Hamiltonian, it will be used even if group=False.
 
     Returns:
-        tuple[list[.QuantumTape], function]: Returns a tuple containing a list of
+        tuple[Sequence[.QuantumTape], Callable]: Returns a tuple containing a list of
         quantum tapes to be evaluated, and a function to be applied to these
         tape executions to compute the expectation value.
 
@@ -67,7 +69,7 @@ def hamiltonian_expand(tape: QuantumTape, group=True):
     Applying the processing function results in the expectation value of the Hamiltonian:
 
     >>> fn(res)
-    -0.5
+    array(-0.5)
 
     Fewer tapes can be constructed by grouping commuting observables. This can be achieved
     by the ``group`` keyword argument:
@@ -200,7 +202,8 @@ def processing_fn(res):
 
 
 # pylint: disable=too-many-branches, too-many-statements
-def sum_expand(tape: QuantumTape, group=True):
+@transform
+def sum_expand(tape: QuantumTape, group: bool = True) -> (Sequence[QuantumTape], Callable):
     """Splits a quantum tape measuring a Sum expectation into multiple tapes of summand
     expectations, and provides a function to recombine the results.
 
@@ -211,7 +214,7 @@ def sum_expand(tape: QuantumTape, group=True):
             wires, leading to fewer tapes.
 
     Returns:
-        tuple[list[.QuantumTape], function]: Returns a tuple containing a list of
+        tuple[Sequence[.QuantumTape], Callable]: Returns a tuple containing a list of
         quantum tapes to be evaluated, and a function to be applied to these
         tape executions to compute the expectation value.
 
diff --git a/pennylane/transforms/mitigate.py b/pennylane/transforms/mitigate.py
index 1069a03f886..4964b5d627a 100644
--- a/pennylane/transforms/mitigate.py
+++ b/pennylane/transforms/mitigate.py
@@ -14,19 +14,19 @@
 """Provides transforms for mitigating quantum circuits."""
 from copy import copy
 
-from typing import Any, Dict, Optional, Sequence, Union
+from typing import Any, Dict, Optional, Sequence, Callable
 
-from pennylane import QNode, apply, adjoint
+from pennylane import apply, adjoint
 from pennylane.math import mean, shape, round
 from pennylane.queuing import AnnotatedQueue
 from pennylane.tape import QuantumTape, QuantumScript
-from pennylane.transforms import batch_transform
+from pennylane.transforms.core import transform
 
 import pennylane as qml
 
 
-@batch_transform
-def fold_global(circuit, scale_factor):
+@transform
+def fold_global(tape: QuantumTape, scale_factor) -> (Sequence[QuantumTape], Callable):
     r"""Differentiable circuit folding of the global unitary ``circuit``.
 
     For a unitary circuit :math:`U = L_d .. L_1`, where :math:`L_i` can be either a gate or layer, ``fold_global`` constructs
@@ -37,7 +37,7 @@ def fold_global(circuit, scale_factor):
     The purpose of folding is to artificially increase the noise for zero noise extrapolation, see :func:`~.pennylane.transforms.mitigate_with_zne`.
 
     Args:
-        circuit (callable or QuantumTape): the circuit to be folded
+        tape (~.QuantumTape): the circuit to be folded
         scale_factor (float): Scale factor :math:`\lambda` determining :math:`n` and :math:`s`
 
     Returns:
@@ -183,7 +183,7 @@ def circuit(x):
     # The main intention for providing ``fold_global`` was for it to be used in combination with ``mitigate_with_zne``, which also works with mitiq functions.
     # To preserve the mitiq functionality, ``mitigate_with_zne`` should get a tape transform.
     # To make ``fold_global`` also user-facing and work with qnodes, this function is batch_transformed instead, and therefore applicable on qnodes.
-    return [fold_global_tape(circuit, scale_factor)], lambda x: x[0]
+    return [fold_global_tape(tape, scale_factor)], lambda x: x[0]
 
 
 def _divmod(a, b):
@@ -328,16 +328,16 @@ def richardson_extrapolate(x, y):
 
 
 # pylint: disable=too-many-arguments, protected-access
-@batch_transform
+@transform
 def mitigate_with_zne(
-    circuit: Union[QNode, QuantumTape],
+    tape: QuantumTape,
     scale_factors: Sequence[float],
     folding: callable,
     extrapolate: callable,
     folding_kwargs: Optional[Dict[str, Any]] = None,
     extrapolate_kwargs: Optional[Dict[str, Any]] = None,
     reps_per_factor=1,
-) -> float:
+) -> (Sequence[QuantumTape], Callable):
     r"""Mitigate an input circuit using zero-noise extrapolation.
 
     Error mitigation is a precursor to error correction and is compatible with near-term quantum
@@ -354,7 +354,7 @@ def mitigate_with_zne(
     see the example and usage details for further information.
 
     Args:
-        circuit (callable or QuantumTape): the circuit to be error-mitigated
+        tape (~.QuantumTape): the circuit to be error-mitigated
         scale_factors (Sequence[float]): the range of noise scale factors used
         folding (callable): a function that returns a folded circuit for a specified scale factor
         extrapolate (callable): a function that returns an extrapolated result when provided a
@@ -508,10 +508,7 @@ def circuit(w1, w2):
     folding_kwargs = folding_kwargs or {}
     extrapolate_kwargs = extrapolate_kwargs or {}
 
-    if isinstance(folding, qml.batch_transform):
-        folding = fold_global_tape
-
-    tape = circuit.expand(stop_at=lambda op: not isinstance(op, QuantumScript))
+    tape = tape.expand(stop_at=lambda op: not isinstance(op, QuantumScript))
     script_removed = QuantumScript(tape.operations[tape.num_preps :])
 
     tapes = [
@@ -520,6 +517,11 @@ def circuit(w1, w2):
     ]
 
     tapes = [tape_ for tapes_ in tapes for tape_ in tapes_]  # flattens nested list
+
+    # if folding was a batch transform, ignore the processing function
+    if isinstance(tapes[0], tuple) and isinstance(tapes[0][0], list) and callable(tapes[0][1]):
+        tapes = [t[0] for t, _ in tapes]
+
     prep_ops = tape.operations[: tape.num_preps]
     out_tapes = [QuantumScript(prep_ops + tape_.operations, tape.measurements) for tape_ in tapes]
 
diff --git a/pennylane/transforms/qcut/__init__.py b/pennylane/transforms/qcut/__init__.py
index 949f0330a92..49729b9dd9f 100644
--- a/pennylane/transforms/qcut/__init__.py
+++ b/pennylane/transforms/qcut/__init__.py
@@ -34,11 +34,10 @@
     _find_new_wire,
     _add_operator_node,
 )
-from .cutcircuit import cut_circuit, _cut_circuit_expand, qnode_execution_wrapper
+from .cutcircuit import cut_circuit, _cut_circuit_expand
 from .montecarlo import (
     cut_circuit_mc,
     _cut_circuit_mc_expand,
-    qnode_execution_wrapper_mc,
     expand_fragment_tapes_mc,
     MC_MEASUREMENTS,
     MC_STATES,
diff --git a/pennylane/transforms/qcut/cutcircuit.py b/pennylane/transforms/qcut/cutcircuit.py
index eaa532db1bf..2780b10db4b 100644
--- a/pennylane/transforms/qcut/cutcircuit.py
+++ b/pennylane/transforms/qcut/cutcircuit.py
@@ -15,14 +15,13 @@
 Function cut_circuit for cutting a quantum circuit into smaller circuit fragments.
 """
 
-
 from functools import partial
-from typing import Callable, Optional, Tuple, Union
+from typing import Callable, Optional, Union, Sequence
 
 import pennylane as qml
 from pennylane.measurements import ExpectationMP
 from pennylane.tape import QuantumTape
-from pennylane.transforms.batch_transform import batch_transform
+from pennylane.transforms.core import transform
 from pennylane.wires import Wires
 
 from .cutstrategy import CutStrategy
@@ -32,7 +31,25 @@
 from .utils import find_and_place_cuts, fragment_graph, replace_wire_cut_nodes
 
 
-@batch_transform
+def _cut_circuit_expand(
+    tape: QuantumTape,
+    use_opt_einsum: bool = False,
+    device_wires: Optional[Wires] = None,
+    max_depth: int = 1,
+    auto_cutter: Union[bool, Callable] = False,
+    **kwargs,
+) -> (Sequence[QuantumTape], Callable):
+    """Main entry point for expanding operations until reaching a depth that
+    includes :class:`~.WireCut` operations."""
+    # pylint: disable=unused-argument
+
+    def processing_fn(res):
+        return res[0]
+
+    return [_qcut_expand_fn(tape, max_depth, auto_cutter)], processing_fn
+
+
+@partial(transform, expand_transform=_cut_circuit_expand)
 def cut_circuit(
     tape: QuantumTape,
     auto_cutter: Union[bool, Callable] = False,
@@ -40,7 +57,7 @@ def cut_circuit(
     device_wires: Optional[Wires] = None,
     max_depth: int = 1,
     **kwargs,
-) -> Tuple[Tuple[QuantumTape], Callable]:
+) -> (Sequence[QuantumTape], Callable):
     """
     Cut up a quantum circuit into smaller circuit fragments.
 
@@ -387,28 +404,9 @@ def circuit(x):
     )
 
 
-def _cut_circuit_expand(
-    tape: QuantumTape,
-    use_opt_einsum: bool = False,
-    device_wires: Optional[Wires] = None,
-    max_depth: int = 1,
-    auto_cutter: Union[bool, Callable] = False,
-    **kwargs,
-):
-    """Main entry point for expanding operations until reaching a depth that
-    includes :class:`~.WireCut` operations."""
-    # pylint: disable=unused-argument
-    return _qcut_expand_fn(tape, max_depth, auto_cutter)
-
-
-cut_circuit.expand_fn = _cut_circuit_expand
-
-
-@cut_circuit.custom_qnode_wrapper
-def qnode_execution_wrapper(self, qnode, targs, tkwargs):
+@cut_circuit.custom_qnode_transform
+def _qnode_transform(self, qnode, targs, tkwargs):
     """Here, we overwrite the QNode execution wrapper in order
     to access the device wires."""
-    # pylint: disable=function-redefined
-
     tkwargs.setdefault("device_wires", qnode.device.wires)
-    return self.default_qnode_wrapper(qnode, targs, tkwargs)
+    return self.default_qnode_transform(qnode, targs, tkwargs)
diff --git a/pennylane/transforms/qcut/montecarlo.py b/pennylane/transforms/qcut/montecarlo.py
index 5096cfa5780..7ea1a76978e 100644
--- a/pennylane/transforms/qcut/montecarlo.py
+++ b/pennylane/transforms/qcut/montecarlo.py
@@ -26,7 +26,7 @@
 from pennylane.measurements import SampleMP
 from pennylane.queuing import AnnotatedQueue
 from pennylane.tape import QuantumScript, QuantumTape
-from pennylane.transforms.batch_transform import batch_transform
+from pennylane.transforms.core import transform
 from pennylane.wires import Wires
 
 from .cutstrategy import CutStrategy
@@ -46,7 +46,26 @@
 from .utils import find_and_place_cuts, fragment_graph, replace_wire_cut_nodes
 
 
-@batch_transform
+def _cut_circuit_mc_expand(
+    tape: QuantumTape,
+    classical_processing_fn: Optional[callable] = None,
+    max_depth: int = 1,
+    shots: Optional[int] = None,
+    device_wires: Optional[Wires] = None,
+    auto_cutter: Union[bool, Callable] = False,
+    **kwargs,
+) -> (Sequence[QuantumTape], Callable):
+    """Main entry point for expanding operations in sample-based tapes until
+    reaching a depth that includes :class:`~.WireCut` operations."""
+    # pylint: disable=unused-argument, too-many-arguments
+
+    def processing_fn(res):
+        return res[0]
+
+    return [_qcut_expand_fn(tape, max_depth, auto_cutter)], processing_fn
+
+
+@partial(transform, expand_transform=_cut_circuit_mc_expand)
 def cut_circuit_mc(
     tape: QuantumTape,
     classical_processing_fn: Optional[callable] = None,
@@ -55,7 +74,7 @@ def cut_circuit_mc(
     shots: Optional[int] = None,
     device_wires: Optional[Wires] = None,
     **kwargs,
-) -> Tuple[Tuple[QuantumTape], Callable]:
+) -> (Sequence[QuantumTape], Callable):
     """
     Cut up a circuit containing sample measurements into smaller fragments using a
     Monte Carlo method.
@@ -442,62 +461,39 @@ def circuit(x):
     tapes = tuple(tape for c in configurations for tape in c)
 
     if classical_processing_fn:
-        return tapes, partial(
-            qcut_processing_fn_mc,
-            communication_graph=communication_graph,
-            settings=settings,
-            shots=shots,
-            classical_processing_fn=classical_processing_fn,
-        )
-
-    return tapes, partial(
-        qcut_processing_fn_sample, communication_graph=communication_graph, shots=shots
-    )
 
+        def processing_fn(results):
+            results = qcut_processing_fn_mc(
+                results,
+                communication_graph=communication_graph,
+                settings=settings,
+                shots=shots,
+                classical_processing_fn=classical_processing_fn,
+            )
 
-def _cut_circuit_mc_expand(
-    tape: QuantumTape,
-    classical_processing_fn: Optional[callable] = None,
-    max_depth: int = 1,
-    shots: Optional[int] = None,
-    device_wires: Optional[Wires] = None,
-    auto_cutter: Union[bool, Callable] = False,
-    **kwargs,
-):
-    """Main entry point for expanding operations in sample-based tapes until
-    reaching a depth that includes :class:`~.WireCut` operations."""
-    # pylint: disable=unused-argument, too-many-arguments
-    return _qcut_expand_fn(tape, max_depth, auto_cutter)
-
+            return results
 
-cut_circuit_mc.expand_fn = _cut_circuit_mc_expand
+    else:
 
+        def processing_fn(results):
+            results = qcut_processing_fn_sample(
+                results, communication_graph=communication_graph, shots=shots
+            )
 
-@cut_circuit_mc.custom_qnode_wrapper
-def qnode_execution_wrapper_mc(self, qnode, targs, tkwargs):
-    """Here, we overwrite the QNode execution wrapper in order
-    to replace execution variables"""
+            return results[0]
 
-    transform_max_diff = tkwargs.pop("max_diff", None)
-    tkwargs.setdefault("device_wires", qnode.device.wires)
+    return tapes, processing_fn
 
-    if "shots" in inspect.signature(qnode.func).parameters:
-        raise ValueError(
-            "Detected 'shots' as an argument of the quantum function to transform. "
-            "The 'shots' argument name is reserved for overriding the number of shots "
-            "taken by the device."
-        )
 
-    def _wrapper(*args, **kwargs):
-        if tkwargs.get("shots", False):
-            raise ValueError(
-                "Cannot provide a 'shots' value directly to the cut_circuit_mc "
-                "decorator when transforming a QNode. Please provide the number of shots in "
-                "the device or when calling the QNode."
-            )
+class CustomQNode(qml.QNode):
+    """
+    A subclass with a custom __call__ method. The custom QNode transform returns an instance
+    of this class.
+    """
 
+    def __call__(self, *args, **kwargs):
         shots = kwargs.pop("shots", False)
-        shots = shots or qnode.device.shots
+        shots = shots or self.device.shots
 
         if shots is None:
             raise ValueError(
@@ -505,39 +501,44 @@ def _wrapper(*args, **kwargs):
                 "or when calling the QNode to be cut"
             )
 
-        qnode.construct(args, kwargs)
-        tapes, processing_fn = self.construct(qnode.qtape, *targs, **tkwargs, shots=shots)
+        # find the qcut transform inside the transform program and set the shots argument
+        qcut_tc = [
+            tc for tc in self.transform_program if tc.transform.__name__ == "cut_circuit_mc"
+        ][-1]
+        qcut_tc._kwargs["shots"] = shots
 
-        interface = qnode.interface
-        execute_kwargs = getattr(qnode, "execute_kwargs", {}).copy()
-        max_diff = execute_kwargs.pop("max_diff", 2)
-        max_diff = transform_max_diff or max_diff
+        kwargs["shots"] = 1
+        return super().__call__(*args, **kwargs)
 
-        gradient_fn = getattr(qnode, "gradient_fn", qnode.diff_method)
-        gradient_kwargs = getattr(qnode, "gradient_kwargs", {})
 
-        if interface is None or not self.differentiable:
-            gradient_fn = None
-
-        execute_kwargs["cache"] = False
+@cut_circuit_mc.custom_qnode_transform
+def _qnode_transform_mc(self, qnode, targs, tkwargs):
+    """Here, we overwrite the QNode execution wrapper in order
+    to access the device wires."""
+    if tkwargs.get("shots", False):
+        raise ValueError(
+            "Cannot provide a 'shots' value directly to the cut_circuit_mc "
+            "decorator when transforming a QNode. Please provide the number of shots in "
+            "the device or when calling the QNode."
+        )
 
-        res = qml.execute(
-            tapes,
-            device=qnode.device,
-            gradient_fn=gradient_fn,
-            interface=interface,
-            max_diff=max_diff,
-            override_shots=1,
-            gradient_kwargs=gradient_kwargs,
-            **execute_kwargs,
+    if "shots" in inspect.signature(qnode.func).parameters:
+        raise ValueError(
+            "Detected 'shots' as an argument of the quantum function to transform. "
+            "The 'shots' argument name is reserved for overriding the number of shots "
+            "taken by the device."
         )
 
-        out = processing_fn(res)
-        if isinstance(out, list) and len(out) == 1:
-            return out[0]
-        return out
+    tkwargs.setdefault("device_wires", qnode.device.wires)
+
+    execute_kwargs = getattr(qnode, "execute_kwargs", {}).copy()
+    execute_kwargs["cache"] = False
+
+    new_qnode = self.default_qnode_transform(qnode, targs, tkwargs)
+    new_qnode.__class__ = CustomQNode
+    new_qnode.execute_kwargs = execute_kwargs
 
-    return _wrapper
+    return new_qnode
 
 
 MC_STATES = [
diff --git a/pennylane/transforms/sign_expand/sign_expand.py b/pennylane/transforms/sign_expand/sign_expand.py
index 9201c6de7ed..312625beca9 100644
--- a/pennylane/transforms/sign_expand/sign_expand.py
+++ b/pennylane/transforms/sign_expand/sign_expand.py
@@ -15,10 +15,11 @@
 # pylint: disable=protected-access
 import json
 from os import path
-from functools import partial
+from typing import Sequence, Callable
+
 import pennylane as qml
 from pennylane import numpy as np
-from pennylane.transforms.batch_transform import batch_transform
+from pennylane.transforms.core import transform
 
 
 def controlled_pauli_evolution(theta, wires, pauli_word, controls):
@@ -194,13 +195,10 @@ def construct_sgn_circuit(  # pylint: disable=too-many-arguments
     return tapes
 
 
-batch_transform_with_diff_false = partial(batch_transform, differentiable=False)
-
-
-@batch_transform_with_diff_false
+@transform
 def sign_expand(  # pylint: disable=too-many-arguments
-    tape, circuit=False, J=10, delta=0.0, controls=("Hadamard", "Target")
-):
+    tape: qml.tape.QuantumTape, circuit=False, J=10, delta=0.0, controls=("Hadamard", "Target")
+) -> (Sequence[qml.tape.QuantumTape], Callable):
     r"""
     Splits a tape measuring a (fast-forwardable) Hamiltonian expectation into mutliple tapes of
     the Xi or sgn decomposition, and provides a function to recombine the results.
diff --git a/pennylane/transforms/split_non_commuting.py b/pennylane/transforms/split_non_commuting.py
index d7a5451c438..ef4c64b3266 100644
--- a/pennylane/transforms/split_non_commuting.py
+++ b/pennylane/transforms/split_non_commuting.py
@@ -15,16 +15,17 @@
 Contains the tape transform that splits non-commuting terms
 """
 # pylint: disable=protected-access
+from typing import Sequence, Callable
 from functools import reduce
 
 import pennylane as qml
 from pennylane.measurements import ProbabilityMP, SampleMP
 
-from .batch_transform import batch_transform
+from pennylane.transforms.core import transform
 
 
-@batch_transform
-def split_non_commuting(tape):
+@transform
+def split_non_commuting(tape: qml.tape.QuantumTape) -> (Sequence[qml.tape.QuantumTape], Callable):
     r"""
     Splits a qnode measuring non-commuting observables into groups of commuting observables.
 
diff --git a/tests/devices/test_default_mixed_jax.py b/tests/devices/test_default_mixed_jax.py
index c3df917f6e2..42ef6618098 100644
--- a/tests/devices/test_default_mixed_jax.py
+++ b/tests/devices/test_default_mixed_jax.py
@@ -15,6 +15,8 @@
 Tests for the ``default.mixed`` device for the JAX interface
 """
 # pylint: disable=protected-access
+from functools import partial
+
 import pytest
 
 import numpy as np
@@ -678,7 +680,7 @@ def test_batching(self, jacobian_fn, decorator, tol):
             # TODO: https://github.com/PennyLaneAI/pennylane/issues/2762
             pytest.xfail("Parameter broadcasting currently not supported for JAX jit")
 
-        @qml.batch_params(all_operations=True)
+        @partial(qml.batch_params, all_operations=True)
         @qml.qnode(dev, diff_method="backprop", interface="jax")
         def circuit(a, b):
             qml.RX(a, wires=0)
diff --git a/tests/qinfo/test_fisher.py b/tests/qinfo/test_fisher.py
index b8ea2f4c361..26a5456631d 100644
--- a/tests/qinfo/test_fisher.py
+++ b/tests/qinfo/test_fisher.py
@@ -39,10 +39,13 @@ def qnode(x):
             qml.CNOT(wires=[0, 1])
             return qml.expval(qml.PauliX(0))
 
-        x = pnp.array([0.5])
+        x = pnp.array(0.5)
         new_qnode = _make_probs(qnode)
-        tape, _ = new_qnode.construct(x, {})
-        assert tape[0].observables[0].return_type == qml.measurements.Probability
+        res = new_qnode(x)
+
+        assert isinstance(res, np.ndarray)
+        assert res.shape == (4,)
+        assert np.isclose(sum(res), 1)
 
     @pytest.mark.parametrize("shots", [None, 100])
     def test_make_probs(self, shots):
diff --git a/tests/shadow/test_shadow_transforms.py b/tests/shadow/test_shadow_transforms.py
index f021ce21e3c..c924a2dc65d 100644
--- a/tests/shadow/test_shadow_transforms.py
+++ b/tests/shadow/test_shadow_transforms.py
@@ -19,7 +19,7 @@
 
 import pennylane as qml
 from pennylane import numpy as np
-from pennylane.measurements.classical_shadow import ShadowExpvalMP
+from pennylane.shadows.transforms import _replace_obs
 
 
 def hadamard_circuit(wires, shots=10000, interface="autograd"):
@@ -98,6 +98,29 @@ def circuit(x, obs):
     return circuit
 
 
+class TestReplaceObs:
+    """Test that the _replace_obs transform works as expected"""
+
+    def test_replace_tape(self):
+        """Test that the transform works for tapes"""
+        tape = qml.tape.QuantumScript([], [qml.classical_shadow(wires=0)])
+        new_tapes, _ = _replace_obs(tape, qml.probs, wires=0)
+
+        assert len(new_tapes) == 1
+        assert new_tapes[0].operations == []
+        assert len(new_tapes[0].observables) == 1
+        assert isinstance(new_tapes[0].observables[0], qml.measurements.ProbabilityMP)
+
+    def test_replace_qnode(self):
+        """Test that the transform works for QNodes"""
+        circuit = hadamard_circuit(2, shots=1000)
+        circuit = _replace_obs(circuit, qml.probs, wires=[0, 1])
+        res = circuit()
+
+        assert isinstance(res, np.ndarray)
+        assert res.shape == (4,)
+
+
 @pytest.mark.autograd
 class TestStateForward:
     """Test that the state reconstruction is correct for a variety of states"""
@@ -308,22 +331,6 @@ def test_backward_torch(self):
 class TestExpvalTransform:
     """Test that the expval transform is applied correctly"""
 
-    def test_hadamard_transform(self):
-        """
-        Test that the transform is correct for a circuit that prepares
-        the uniform superposition
-        """
-        obs = qml.PauliZ(0)
-        circuit = hadamard_circuit(3, shots=100000)
-        circuit = qml.shadows.shadow_expval(obs)(circuit)
-
-        tape = circuit.construct((), {})[0][0]
-
-        assert all(qml.equal(qml.Hadamard(i), tape.operations[i]) for i in range(3))
-        assert len(tape.observables) == 1
-        assert isinstance(tape.observables[0], ShadowExpvalMP)
-        assert tape.observables[0].H == obs
-
     def test_hadamard_forward(self):
         """Test that the expval estimation is correct for a uniform
         superposition of qubits"""
diff --git a/tests/transforms/test_batch_input.py b/tests/transforms/test_batch_input.py
index e7abe9b6b32..9f02442d2ee 100644
--- a/tests/transforms/test_batch_input.py
+++ b/tests/transforms/test_batch_input.py
@@ -15,6 +15,8 @@
 Unit tests for the ``batch_inputs`` transform.
 """
 # pylint: disable=too-few-public-methods
+from functools import partial
+
 import pytest
 
 import pennylane as qml
@@ -25,7 +27,7 @@ def test_simple_circuit():
     """Test that batching works for a simple circuit"""
     dev = qml.device("default.qubit", wires=2)
 
-    @qml.batch_input(argnum=1)
+    @partial(qml.batch_input, argnum=1)
     @qml.qnode(dev, diff_method="parameter-shift")
     def circuit(inputs, weights):
         qml.RY(weights[0], wires=0)
@@ -45,7 +47,7 @@ def test_simple_circuit_one_batch():
     """Test that batching works for a simple circuit when the batch size is 1"""
     dev = qml.device("default.qubit", wires=2)
 
-    @qml.batch_input(argnum=1)
+    @partial(qml.batch_input, argnum=1)
     @qml.qnode(dev, diff_method="parameter-shift")
     def circuit(inputs, weights):
         qml.RY(weights[0], wires=0)
@@ -65,7 +67,7 @@ def test_simple_circuit_with_prep():
     """Test that batching works for a simple circuit with a state preparation"""
     dev = qml.device("default.qubit", wires=2)
 
-    @qml.batch_input(argnum=1)
+    @partial(qml.batch_input, argnum=1)
     @qml.qnode(dev)
     def circuit(inputs, weights):
         qml.StatePrep(np.array([0, 0, 1, 0]), wires=[0, 1])
@@ -86,7 +88,7 @@ def test_circuit_non_param_operator_before_batched_operator():
     """Test a circuit where a non-parametric operation is located before a batched operator."""
     dev = qml.device("default.qubit", wires=2)
 
-    @qml.batch_input(argnum=0)
+    @partial(qml.batch_input, argnum=0)
     @qml.qnode(dev)
     def circuit(input):
         qml.CNOT(wires=[0, 1])
@@ -117,7 +119,7 @@ class Embedding(qml.AngleEmbedding):
         def ndim_params(self):
             return self._ndim_params
 
-    @qml.batch_input(argnum=[0, 2])
+    @partial(qml.batch_input, argnum=[0, 2])
     @qml.qnode(dev, diff_method="parameter-shift")
     def circuit(input1, input2, weights):
         Embedding(input1, wires=range(2), rotation="Y")
@@ -140,7 +142,7 @@ def test_batch_input_with_trainable_parameters_raises_error():
     """Test that using the batch_input method with trainable parameters raises a ValueError."""
     dev = qml.device("default.qubit", wires=2)
 
-    @qml.batch_input(argnum=0)
+    @partial(qml.batch_input, argnum=0)
     @qml.qnode(dev)
     def circuit(input):
         qml.RY(input, wires=1)
@@ -165,7 +167,7 @@ def test_mottonenstate_preparation(mocker):
     """Test that batching works for MottonenStatePreparation"""
     dev = qml.device("default.qubit", wires=3)
 
-    @qml.batch_input(argnum=0)
+    @partial(qml.batch_input, argnum=0)
     @qml.qnode(dev, interface="autograd")
     def circuit(data, weights):
         qml.templates.MottonenStatePreparation(data, wires=[0, 1, 2])
@@ -204,7 +206,7 @@ def test_qubit_state_prep(mocker):
 
     dev = qml.device("default.qubit", wires=3)
 
-    @qml.batch_input(argnum=0)
+    @partial(qml.batch_input, argnum=0)
     @qml.qnode(dev, interface="autograd")
     def circuit(data, weights):
         qml.StatePrep(data, wires=[0, 1, 2])
@@ -243,7 +245,7 @@ def test_multi_returns():
     """Test that batching works for a simple circuit with multiple returns"""
     dev = qml.device("default.qubit", wires=2)
 
-    @qml.batch_input(argnum=1)
+    @partial(qml.batch_input, argnum=1)
     @qml.qnode(dev, diff_method="parameter-shift")
     def circuit(inputs, weights):
         qml.RY(weights[0], wires=0)
@@ -267,7 +269,7 @@ def test_shot_vector():
     """Test that batching works for a simple circuit with a shot vector"""
     dev = qml.device("default.qubit", wires=2, shots=(100, (200, 3), 300))
 
-    @qml.batch_input(argnum=1)
+    @partial(qml.batch_input, argnum=1)
     @qml.qnode(dev, diff_method="parameter-shift")
     def circuit(inputs, weights):
         qml.RY(weights[0], wires=0)
@@ -292,7 +294,7 @@ def test_multi_returns_shot_vector():
     and with a shot vector"""
     dev = qml.device("default.qubit", wires=2, shots=(100, (200, 3), 300))
 
-    @qml.batch_input(argnum=1)
+    @partial(qml.batch_input, argnum=1)
     @qml.qnode(dev, diff_method="parameter-shift")
     def circuit(inputs, weights):
         qml.RY(weights[0], wires=0)
@@ -323,7 +325,7 @@ def test_autograd(self, diff_method, tol):
         """Test derivatives when using autograd"""
         dev = qml.device("default.qubit", wires=2)
 
-        @qml.batch_input(argnum=0)
+        @partial(qml.batch_input, argnum=0)
         @qml.qnode(dev, diff_method=diff_method)
         def circuit(input, x):
             qml.RY(input, wires=1)
@@ -353,7 +355,7 @@ def test_jax(self, diff_method, tol, interface):
 
         dev = qml.device("default.qubit", wires=2)
 
-        @qml.batch_input(argnum=0)
+        @partial(qml.batch_input, argnum=0)
         @qml.qnode(dev, diff_method=diff_method, interface=interface)
         def circuit(input, x):
             qml.RY(input, wires=1)
@@ -386,7 +388,7 @@ def test_jax_jit(self, diff_method, tol, interface):
         dev = qml.device("default.qubit", wires=2)
 
         @jax.jit
-        @qml.batch_input(argnum=0)
+        @partial(qml.batch_input, argnum=0)
         @qml.qnode(dev, diff_method=diff_method, interface=interface)
         def circuit(input, x):
             qml.RY(input, wires=1)
@@ -415,7 +417,7 @@ def test_torch(self, diff_method, tol, interface):
 
         dev = qml.device("default.qubit", wires=2)
 
-        @qml.batch_input(argnum=0)
+        @partial(qml.batch_input, argnum=0)
         @qml.qnode(dev, diff_method=diff_method, interface=interface)
         def circuit(input, x):
             qml.RY(input, wires=1)
@@ -447,7 +449,7 @@ def test_tf(self, diff_method, tol, interface):
 
         dev = qml.device("default.qubit", wires=2)
 
-        @qml.batch_input(argnum=0)
+        @partial(qml.batch_input, argnum=0)
         @qml.qnode(dev, diff_method=diff_method, interface=interface)
         def circuit(input, x):
             qml.RY(input, wires=1)
@@ -476,7 +478,7 @@ def test_tf_autograph(self, diff_method, tol, interface):
         dev = qml.device("default.qubit", wires=2)
 
         @tf.function
-        @qml.batch_input(argnum=0)
+        @partial(qml.batch_input, argnum=0)
         @qml.qnode(dev, diff_method=diff_method, interface=interface)
         def circuit(input, x):
             qml.RY(input, wires=1)
@@ -505,7 +507,7 @@ def test_autograd(self, diff_method, tol):
         """Test derivatives when using autograd"""
         dev = qml.device("default.qubit", wires=2)
 
-        @qml.batch_input(argnum=0)
+        @partial(qml.batch_input, argnum=0)
         @qml.qnode(dev, diff_method=diff_method)
         def circuit(input, x):
             qml.RY(input, wires=0)
@@ -559,7 +561,7 @@ def test_jax(self, diff_method, tol, interface):
 
         dev = qml.device("default.qubit", wires=2)
 
-        @qml.batch_input(argnum=0)
+        @partial(qml.batch_input, argnum=0)
         @qml.qnode(dev, diff_method=diff_method, interface=interface)
         def circuit(input, x):
             qml.RY(input, wires=0)
@@ -624,7 +626,7 @@ def test_jax_jit(self, diff_method, tol, interface):
         dev = qml.device("default.qubit", wires=2)
 
         @jax.jit
-        @qml.batch_input(argnum=0)
+        @partial(qml.batch_input, argnum=0)
         @qml.qnode(dev, diff_method=diff_method, interface=interface)
         def circuit(input, x):
             qml.RY(input, wires=0)
@@ -685,7 +687,7 @@ def test_torch(self, diff_method, tol, interface):
 
         dev = qml.device("default.qubit", wires=2)
 
-        @qml.batch_input(argnum=0)
+        @partial(qml.batch_input, argnum=0)
         @qml.qnode(dev, diff_method=diff_method, interface=interface)
         def circuit(input, x):
             qml.RY(input, wires=0)
@@ -746,7 +748,7 @@ def test_tf(self, diff_method, tol, interface):
 
         dev = qml.device("default.qubit", wires=2)
 
-        @qml.batch_input(argnum=0)
+        @partial(qml.batch_input, argnum=0)
         @qml.qnode(dev, diff_method=diff_method, interface=interface)
         def circuit(input, x):
             qml.RY(input, wires=0)
@@ -799,7 +801,7 @@ def test_tf_autograph(self, diff_method, tol, interface):
         dev = qml.device("default.qubit", wires=2)
 
         @tf.function
-        @qml.batch_input(argnum=0)
+        @partial(qml.batch_input, argnum=0)
         @qml.qnode(dev, diff_method=diff_method, interface=interface)
         def circuit(input, x):
             qml.RY(input, wires=0)
@@ -859,7 +861,7 @@ def test_unbatched_not_copied():
     tape = qml.tape.QuantumScript(ops, meas)
     tape.trainable_params = [0, 2]
 
-    new_tapes = qml.batch_input(argnum=1)(tape)[0]
+    new_tapes = qml.batch_input(tape, argnum=1)[0]
     assert len(new_tapes) == batch_size
 
     for new_tape in new_tapes:
diff --git a/tests/transforms/test_batch_params.py b/tests/transforms/test_batch_params.py
index 7e8aa5a0491..f2333c59eb2 100644
--- a/tests/transforms/test_batch_params.py
+++ b/tests/transforms/test_batch_params.py
@@ -557,7 +557,7 @@ def test_jax(self, diff_method, tol, interface):
         jax.config.update("jax_enable_x64", True)
         dev = qml.device("default.qubit", wires=2)
 
-        @qml.batch_params(all_operations=True)
+        @functools.partial(qml.batch_params, all_operations=True)
         @qml.qnode(dev, diff_method=diff_method, interface=interface)
         def circuit(x):
             qml.RY(x, wires=0)
@@ -607,7 +607,7 @@ def test_jax_jit(self, diff_method, tol, interface):
         dev = qml.device("default.qubit", wires=2)
 
         @jax.jit
-        @qml.batch_params(all_operations=True)
+        @functools.partial(qml.batch_params, all_operations=True)
         @qml.qnode(dev, diff_method=diff_method, interface=interface)
         def circuit(x):
             qml.RY(x, wires=0)
diff --git a/tests/transforms/test_broadcast_expand.py b/tests/transforms/test_broadcast_expand.py
index 599acf761e1..614e0ec5e5b 100644
--- a/tests/transforms/test_broadcast_expand.py
+++ b/tests/transforms/test_broadcast_expand.py
@@ -112,6 +112,30 @@ def test_expansion(self, params, size, obs, exp_fn):
 
         assert qml.math.allclose(result, expected)
 
+    @pytest.mark.parametrize("params, size", parameters_and_size)
+    @pytest.mark.parametrize("obs, exp_fn", observables_and_exp_fns)
+    def test_expansion_qnode(self, params, size, obs, exp_fn):
+        """Test that the transform integrates correctly with the transform program"""
+
+        @qml.transforms.broadcast_expand
+        @qml.qnode(dev)
+        def circuit(x, y, z, obs):
+            qml.StatePrep(np.array([1, 0, 0, 0]), wires=[0, 1])
+            RX_broadcasted(x, wires=0)
+            qml.PauliY(0)
+            RX_broadcasted(y, wires=1)
+            RZ_broadcasted(z, wires=1)
+            qml.Hadamard(1)
+            return [qml.expval(ob) for ob in obs]
+
+        result = circuit(*params, obs)
+        expected = exp_fn(*params)
+
+        if len(obs) > 1 and size == 1:
+            expected = expected.T
+
+        assert qml.math.allclose(result, expected)
+
     def test_without_broadcasting(self):
         tape = make_tape(0.2, 0.1, 0.5, [qml.PauliZ(0)])
         with pytest.raises(ValueError, match="The provided tape is not broadcasted."):
diff --git a/tests/transforms/test_experimental/test_transform_dispatcher.py b/tests/transforms/test_experimental/test_transform_dispatcher.py
index 2e5e6300fbc..a37403d2eb0 100644
--- a/tests/transforms/test_experimental/test_transform_dispatcher.py
+++ b/tests/transforms/test_experimental/test_transform_dispatcher.py
@@ -49,6 +49,7 @@ def no_tape_transform(
     circuit: qml.tape.QuantumTape, index: int
 ) -> (Sequence[qml.tape.QuantumTape], Callable):
     """Transform without tape."""
+    circuit = circuit.copy()
     circuit._ops.pop(index)  # pylint:disable=protected-access
     return [circuit], lambda x: x
 
@@ -57,6 +58,7 @@ def no_quantum_tape_transform(
     tape: qml.operation.Operator, index: int
 ) -> (Sequence[qml.tape.QuantumTape], Callable):
     """Transform with wrong hinting."""
+    tape = tape.copy()
     tape._ops.pop(index)  # pylint:disable=protected-access
     return [tape], lambda x: x
 
@@ -97,6 +99,7 @@ def first_valid_transform(
     tape: qml.tape.QuantumTape, index: int
 ) -> (Sequence[qml.tape.QuantumTape], Callable):
     """A valid transform."""
+    tape = tape.copy()
     tape._ops.pop(index)  # pylint:disable=protected-access
     return [tape], lambda x: x
 
@@ -106,6 +109,7 @@ def second_valid_transform(
 ) -> (Sequence[qml.tape.QuantumTape], Callable):
     """A valid trasnform."""
     tape1 = tape.copy()
+    tape2 = tape.copy()
     tape2 = tape._ops.pop(index)  # pylint:disable=protected-access
 
     def fn(results):
@@ -118,8 +122,8 @@ def fn(results):
 
 
 ##########################################
-# Non-valid expand transform
-def multiple_args_expand_transform(
+# Valid expand transform
+def expand_transform(
     tape: qml.tape.QuantumTape, index: int
 ) -> (Sequence[qml.tape.QuantumTape], Callable):
     """Multiple args expand fn."""
@@ -127,8 +131,10 @@ def multiple_args_expand_transform(
     return [tape], lambda x: x
 
 
-# Valid expand transform
-def expand_transform(tape: qml.tape.QuantumTape) -> (Sequence[qml.tape.QuantumTape], Callable):
+# Non-valid expand transform
+def non_valid_expand_transform(
+    tape: qml.tape.QuantumTape,
+) -> (Sequence[qml.tape.QuantumTape], Callable):
     """A valid expand transform."""
     return [tape], lambda x: x
 
@@ -279,7 +285,7 @@ def qnode_circuit(a):
         assert isinstance(qnode_transformed.transform_program, qml.transforms.core.TransformProgram)
         expand_transform_container = qnode_transformed.transform_program.pop_front()
         assert isinstance(expand_transform_container, qml.transforms.core.TransformContainer)
-        assert expand_transform_container.args == []
+        assert expand_transform_container.args == [0]
         assert expand_transform_container.kwargs == {}
         assert expand_transform_container.classical_cotransform is None
         assert not expand_transform_container.is_informative
@@ -350,13 +356,13 @@ def test_expand_transform_not_callable(self):
             transform(first_valid_transform, expand_transform=non_callable)
 
     def test_multiple_args_expand_transform(self):
-        """Test that an expand transform must take a single argument which is the tape."""
+        """Test that an expand transform must match the signature of the transform"""
 
         with pytest.raises(
             TransformError,
-            match="The expand transform does not support arg and kwargs other than tape.",
+            match="The expand transform must have the same signature as the transform",
         ):
-            transform(first_valid_transform, expand_transform=multiple_args_expand_transform)
+            transform(first_valid_transform, expand_transform=non_valid_expand_transform)
 
     def test_cotransform_not_implemented(self):
         """Test that a co-transform must be a callable."""
diff --git a/tests/transforms/test_experimental/test_transform_program.py b/tests/transforms/test_experimental/test_transform_program.py
index 146692471dc..00c3153867c 100644
--- a/tests/transforms/test_experimental/test_transform_program.py
+++ b/tests/transforms/test_experimental/test_transform_program.py
@@ -36,11 +36,14 @@ def first_valid_transform(
     tape: qml.tape.QuantumTape, index: int
 ) -> (Sequence[qml.tape.QuantumTape], Callable):
     """A valid transform."""
+    tape = tape.copy()
     tape._ops.pop(index)  # pylint:disable=protected-access
     return [tape], lambda x: x
 
 
-def expand_transform(tape: qml.tape.QuantumTape) -> (Sequence[qml.tape.QuantumTape], Callable):
+def expand_transform(
+    tape: qml.tape.QuantumTape, index: int  # pylint:disable=unused-argument
+) -> (Sequence[qml.tape.QuantumTape], Callable):
     """A valid expand transform."""
     return [tape], lambda x: x
 
@@ -50,6 +53,7 @@ def second_valid_transform(
 ) -> (Sequence[qml.tape.QuantumTape], Callable):
     """A valid trasnform."""
     tape1 = tape.copy()
+    tape2 = tape.copy()
     tape2 = tape._ops.pop(index)  # pylint:disable=protected-access
 
     def fn(results):
diff --git a/tests/transforms/test_mitigate.py b/tests/transforms/test_mitigate.py
index 046a320b319..5a321fcc104 100644
--- a/tests/transforms/test_mitigate.py
+++ b/tests/transforms/test_mitigate.py
@@ -14,6 +14,8 @@
 """
 Tests for mitigation transforms.
 """
+# pylint:disable=no-self-use
+from functools import partial
 import pytest
 
 from packaging import version
@@ -121,7 +123,12 @@ def test_multi_returns(self):
         np.random.seed(0)
         w1, w2 = [np.random.random(s) for s in shapes]
 
-        @qml.transforms.mitigate_with_zne([1, 2, 3], fold_global, richardson_extrapolate)
+        @partial(
+            qml.transforms.mitigate_with_zne,
+            scale_factors=[1, 2, 3],
+            folding=fold_global,
+            extrapolate=richardson_extrapolate,
+        )
         @qml.qnode(dev)
         def mitigated_circuit(w1, w2):
             qml.SimplifiedTwoDesign(w1, w2, wires=range(2))
@@ -214,7 +221,12 @@ def test_multiple_returns(self):
         np.random.seed(0)
         w1, w2 = [np.random.random(s) for s in shapes]
 
-        @qml.transforms.mitigate_with_zne([1, 2, 3], fold_global, RichardsonFactory.extrapolate)
+        @partial(
+            qml.transforms.mitigate_with_zne,
+            scale_factors=[1, 2, 3],
+            folding=fold_global,
+            extrapolate=RichardsonFactory.extrapolate,
+        )
         @qml.qnode(dev)
         def mitigated_circuit(w1, w2):
             qml.SimplifiedTwoDesign(w1, w2, wires=range(2))
@@ -259,7 +271,12 @@ def test_single_return(self):
         np.random.seed(0)
         w1, w2 = [np.random.random(s) for s in shapes]
 
-        @qml.transforms.mitigate_with_zne([1, 2, 3], fold_global, RichardsonFactory.extrapolate)
+        @partial(
+            qml.transforms.mitigate_with_zne,
+            scale_factors=[1, 2, 3],
+            folding=fold_global,
+            extrapolate=RichardsonFactory.extrapolate,
+        )
         @qml.qnode(dev)
         def mitigated_circuit(w1, w2):
             qml.SimplifiedTwoDesign(w1, w2, wires=range(2))
@@ -294,8 +311,12 @@ def test_with_reps_per_factor(self):
         np.random.seed(0)
         w1, w2 = [np.random.random(s) for s in shapes]
 
-        @qml.transforms.mitigate_with_zne(
-            [1, 2, 3], fold_gates_at_random, RichardsonFactory.extrapolate, reps_per_factor=2
+        @partial(
+            qml.transforms.mitigate_with_zne,
+            scale_factors=[1, 2, 3],
+            folding=fold_gates_at_random,
+            extrapolate=RichardsonFactory.extrapolate,
+            reps_per_factor=2,
         )
         @qml.qnode(dev)
         def mitigated_circuit(w1, w2):
@@ -338,7 +359,12 @@ def circuit(w1, w2):
         exact_qnode = qml.QNode(circuit, dev_noise_free)
         noisy_qnode = qml.QNode(circuit, dev)
 
-        @qml.transforms.mitigate_with_zne([1, 2, 3], fold_global, RichardsonFactory.extrapolate)
+        @partial(
+            qml.transforms.mitigate_with_zne,
+            scale_factors=[1, 2, 3],
+            folding=fold_global,
+            extrapolate=RichardsonFactory.extrapolate,
+        )
         @qml.qnode(dev)
         def mitigated_qnode(w1, w2):
             qml.SimplifiedTwoDesign(w1, w2, wires=range(2))
@@ -383,7 +409,12 @@ def test_grad(self):
         np.random.seed(0)
         w1, w2 = [np.random.random(s, requires_grad=True) for s in shapes]
 
-        @qml.transforms.mitigate_with_zne([1, 2, 3], fold_global, RichardsonFactory.extrapolate)
+        @partial(
+            qml.transforms.mitigate_with_zne,
+            scale_factors=[1, 2, 3],
+            folding=fold_global,
+            extrapolate=RichardsonFactory.extrapolate,
+        )
         @qml.qnode(dev)
         def mitigated_circuit(w1, w2):
             qml.SimplifiedTwoDesign(w1, w2, wires=range(2))
@@ -466,8 +497,8 @@ def test_diffability_autograd(self):
 
         scale_factors = [1.0, 2.0, 3.0]
 
-        mitigated_qnode = mitigate_with_zne(scale_factors, fold_global, richardson_extrapolate)(
-            qnode_noisy
+        mitigated_qnode = mitigate_with_zne(
+            qnode_noisy, scale_factors, fold_global, richardson_extrapolate
         )
 
         theta = np.array([np.pi / 4, np.pi / 4], requires_grad=True)
@@ -492,8 +523,8 @@ def test_diffability_jax(self, interface):
 
         scale_factors = [1.0, 2.0, 3.0]
 
-        mitigated_qnode = mitigate_with_zne(scale_factors, fold_global, richardson_extrapolate)(
-            qnode_noisy
+        mitigated_qnode = mitigate_with_zne(
+            qnode_noisy, scale_factors, fold_global, richardson_extrapolate
         )
 
         theta = jnp.array(
@@ -521,7 +552,7 @@ def test_diffability_jaxjit(self, interface):
         scale_factors = [1.0, 2.0, 3.0]
 
         mitigated_qnode = jax.jit(
-            mitigate_with_zne(scale_factors, fold_global, richardson_extrapolate)(qnode_noisy)
+            mitigate_with_zne(qnode_noisy, scale_factors, fold_global, richardson_extrapolate)
         )
 
         theta = jnp.array(
@@ -547,8 +578,8 @@ def test_diffability_torch(self, interface):
 
         scale_factors = [1.0, 2.0, 3.0]
 
-        mitigated_qnode = mitigate_with_zne(scale_factors, fold_global, richardson_extrapolate)(
-            qnode_noisy
+        mitigated_qnode = mitigate_with_zne(
+            qnode_noisy, scale_factors, fold_global, richardson_extrapolate
         )
 
         theta = torch.tensor([np.pi / 4, np.pi / 4], requires_grad=True)
@@ -576,8 +607,8 @@ def test_diffability_tf(self, interface):
 
         scale_factors = [1.0, 2.0, 3.0]
 
-        mitigated_qnode = mitigate_with_zne(scale_factors, fold_global, richardson_extrapolate)(
-            qnode_noisy
+        mitigated_qnode = mitigate_with_zne(
+            qnode_noisy, scale_factors, fold_global, richardson_extrapolate
         )
 
         theta = tf.Variable([np.pi / 4, np.pi / 4])
@@ -604,8 +635,8 @@ def test_diffability_autograd_multi(self):
 
         scale_factors = [1.0, 2.0, 3.0]
 
-        mitigated_qnode = mitigate_with_zne(scale_factors, fold_global, richardson_extrapolate)(
-            qnode_noisy
+        mitigated_qnode = mitigate_with_zne(
+            qnode_noisy, scale_factors, fold_global, richardson_extrapolate
         )
 
         theta = np.array([np.pi / 4, np.pi / 6], requires_grad=True)
@@ -631,8 +662,8 @@ def test_diffability_jax_multi(self, interface):
 
         scale_factors = [1.0, 2.0, 3.0]
 
-        mitigated_qnode = mitigate_with_zne(scale_factors, fold_global, richardson_extrapolate)(
-            qnode_noisy
+        mitigated_qnode = mitigate_with_zne(
+            qnode_noisy, scale_factors, fold_global, richardson_extrapolate
         )
 
         theta = jnp.array(
@@ -661,7 +692,7 @@ def test_diffability_jaxjit_multi(self, interface):
         scale_factors = [1.0, 2.0, 3.0]
 
         mitigated_qnode = jax.jit(
-            mitigate_with_zne(scale_factors, fold_global, richardson_extrapolate)(qnode_noisy)
+            mitigate_with_zne(qnode_noisy, scale_factors, fold_global, richardson_extrapolate)
         )
 
         theta = jnp.array(
@@ -688,8 +719,8 @@ def test_diffability_torch_multi(self, interface):
 
         scale_factors = [1.0, 2.0, 3.0]
 
-        mitigated_qnode = mitigate_with_zne(scale_factors, fold_global, richardson_extrapolate)(
-            qnode_noisy
+        mitigated_qnode = mitigate_with_zne(
+            qnode_noisy, scale_factors, fold_global, richardson_extrapolate
         )
 
         theta = torch.tensor([np.pi / 4, np.pi / 6], requires_grad=True)
@@ -717,8 +748,8 @@ def test_diffability_tf_multi(self):
 
         scale_factors = [1.0, 2.0, 3.0]
 
-        mitigated_qnode = mitigate_with_zne(scale_factors, fold_global, richardson_extrapolate)(
-            qnode_noisy
+        mitigated_qnode = mitigate_with_zne(
+            qnode_noisy, scale_factors, fold_global, richardson_extrapolate
         )
 
         theta = tf.Variable([np.pi / 4, np.pi / 6])
diff --git a/tests/transforms/test_qcut.py b/tests/transforms/test_qcut.py
index b2e74c8e37d..61b1ab37eb4 100644
--- a/tests/transforms/test_qcut.py
+++ b/tests/transforms/test_qcut.py
@@ -23,6 +23,7 @@
 from itertools import product
 from os import environ
 from pathlib import Path
+from functools import partial
 
 import pytest
 from flaky import flaky
@@ -2487,7 +2488,7 @@ def target_circuit(v):
 
         dev = qml.device("default.qubit", wires=2, shots=10000)
 
-        @qml.cut_circuit_mc(fn)
+        @partial(qml.cut_circuit_mc, classical_processing_fn=fn)
         @qml.qnode(dev)
         def circuit(v):
             qml.RX(v, wires=0)
@@ -2640,27 +2641,25 @@ def test_transform_shots_error(self):
 
         dev = qml.device("default.qubit", wires=2)
 
-        @qml.cut_circuit_mc(shots=456)
-        @qml.qnode(dev)
-        def cut_circuit(x):
-            qml.RX(x, wires=0)
-            qml.RY(0.5, wires=1)
-            qml.RX(1.3, wires=2)
-
-            qml.CNOT(wires=[0, 1])
-            qml.WireCut(wires=1)
-            qml.CNOT(wires=[1, 2])
-
-            qml.RX(x, wires=0)
-            qml.RY(0.7, wires=1)
-            qml.RX(2.3, wires=2)
-            return qml.sample(wires=[0, 2])
-
-        v = 0.319
         with pytest.raises(
             ValueError, match="Cannot provide a 'shots' value directly to the cut_circuit_mc "
         ):
-            cut_circuit(v)
+
+            @partial(qml.cut_circuit_mc, shots=456)
+            @qml.qnode(dev)
+            def cut_circuit(x):  # pylint: disable=unused-variable,unused-argument
+                qml.RX(x, wires=0)
+                qml.RY(0.5, wires=1)
+                qml.RX(1.3, wires=2)
+
+                qml.CNOT(wires=[0, 1])
+                qml.WireCut(wires=1)
+                qml.CNOT(wires=[1, 2])
+
+                qml.RX(x, wires=0)
+                qml.RY(0.7, wires=1)
+                qml.RX(2.3, wires=2)
+                return qml.sample(wires=[0, 2])
 
     def test_multiple_meas_error(self):
         """
@@ -2918,7 +2917,7 @@ def test_mc_autograd(self):
         shots = 10
         dev = qml.device("default.qubit", wires=2, shots=shots)
 
-        @qml.cut_circuit_mc(fn)
+        @partial(qml.cut_circuit_mc, classical_processing_fn=fn)
         @qml.qnode(dev)
         def cut_circuit(x):
             qml.RX(x, wires=0)
@@ -2951,7 +2950,7 @@ def test_mc_tf(self):
         shots = 10
         dev = qml.device("default.qubit", wires=2, shots=shots)
 
-        @qml.cut_circuit_mc(fn)
+        @partial(qml.cut_circuit_mc, classical_processing_fn=fn)
         @qml.qnode(dev)
         def cut_circuit(x):
             qml.RX(x, wires=0)
@@ -2984,7 +2983,7 @@ def test_mc_torch(self):
         shots = 10
         dev = qml.device("default.qubit", wires=2, shots=shots)
 
-        @qml.cut_circuit_mc(fn)
+        @partial(qml.cut_circuit_mc, classical_processing_fn=fn)
         @qml.qnode(dev)
         def cut_circuit(x):
             qml.RX(x, wires=0)
@@ -3017,7 +3016,7 @@ def test_mc_jax(self):
         shots = 10
         dev = qml.device("default.qubit", wires=2, shots=shots)
 
-        @qml.cut_circuit_mc(fn)
+        @partial(qml.cut_circuit_mc, classical_processing_fn=fn)
         @qml.qnode(dev)
         def cut_circuit(x):
             qml.RX(x, wires=0)
@@ -3048,7 +3047,7 @@ def test_mc_with_mid_circuit_measurement(self, mocker):
         shots = 10
         dev = qml.device("default.qubit", wires=3, shots=shots)
 
-        @qml.cut_circuit_mc(fn)
+        @partial(qml.cut_circuit_mc, classical_processing_fn=fn)
         @qml.qnode(dev)
         def circuit(x):
             qml.RX(x, wires=0)
@@ -3074,7 +3073,7 @@ def test_mc_circuit_with_disconnected_components(self):
         shots = 10
         dev = qml.device("default.qubit", wires=3, shots=shots)
 
-        @qml.cut_circuit_mc(fn)
+        @partial(qml.cut_circuit_mc, classical_processing_fn=fn)
         @qml.qnode(dev)
         def circuit(x):
             qml.RX(x, wires=0)
@@ -3094,7 +3093,7 @@ def test_mc_circuit_with_trivial_wire_cut(self, mocker):
         shots = 10
         dev = qml.device("default.qubit", wires=2, shots=shots)
 
-        @qml.cut_circuit_mc(fn)
+        @partial(qml.cut_circuit_mc, classical_processing_fn=fn)
         @qml.qnode(dev)
         def circuit(x):
             qml.RX(x, wires=0)
@@ -3132,7 +3131,7 @@ def two_qubit_unitary(param, wires):
             qml.Hadamard(wires=[wires[0]])
             qml.CRY(param, wires=[wires[0], wires[1]])
 
-        @qml.cut_circuit_mc(fn)
+        @partial(qml.cut_circuit_mc, classical_processing_fn=fn)
         @qml.qnode(dev)
         def circuit(params):
             qml.BasisState(np.array([1]), wires=[0])
@@ -4195,6 +4194,7 @@ def circuit(x):
         spy.assert_called_once()
         assert np.isclose(res, res_expected)
 
+        jax.grad(cut_circuit_jit)(x)
         grad = jax.grad(cut_circuit_jit)(x)
         grad_expected = jax.grad(circuit)(x)
 
@@ -4211,6 +4211,7 @@ def circuit(x):
 
             assert np.isclose(res, res_expected)
 
+            jax.grad(cut_circuit_jit)(x)
             grad = jax.grad(cut_circuit_jit)(x)
             grad_expected = jax.grad(circuit)(x)
 
@@ -4262,7 +4263,7 @@ def test_circuit_with_disconnected_components(self, use_opt_einsum):
 
         dev = qml.device("default.qubit", wires=3)
 
-        @qml.transforms.cut_circuit(use_opt_einsum=use_opt_einsum)
+        @partial(qml.transforms.cut_circuit, use_opt_einsum=use_opt_einsum)
         @qml.qnode(dev)
         def circuit(x):
             qml.RX(x, wires=0)
@@ -4284,7 +4285,7 @@ def test_circuit_with_trivial_wire_cut(self, use_opt_einsum, mocker):
 
         dev = qml.device("default.qubit", wires=2)
 
-        @qml.transforms.cut_circuit(use_opt_einsum=use_opt_einsum)
+        @partial(qml.transforms.cut_circuit, use_opt_einsum=use_opt_einsum)
         @qml.qnode(dev)
         def circuit(x):
             qml.RX(x, wires=0)
@@ -4477,12 +4478,14 @@ def test_fail_import(self, monkeypatch):
     def test_no_cuts_raises(self):
         """Tests if a ValueError is raised when circuit cutting is to be applied to a circuit
         without cuts"""
-        with qml.queuing.AnnotatedQueue() as q:
-            qml.expval(qml.PauliZ(0))
+        dev = qml.device("default.qubit", wires=3)
+
+        @qml.qnode(dev)
+        def circuit():
+            return qml.expval(qml.PauliZ(0))
 
-        tape = qml.tape.QuantumScript.from_queue(q)
         with pytest.raises(ValueError, match="No WireCut operations found in the circuit."):
-            qcut.cut_circuit(tape)
+            qcut.cut_circuit(circuit)()
 
 
 class TestCutCircuitExpansion:
@@ -4497,18 +4500,20 @@ class TestCutCircuitExpansion:
     def test_no_expansion(self, mocker, cut_transform, measurement):
         """Test if no/trivial expansion occurs if WireCut operations are already present in the
         tape"""
-        with qml.queuing.AnnotatedQueue() as q:
+        dev = qml.device("default.qubit", wires=3)
+
+        @qml.qnode(dev)
+        def circuit():
             qml.RX(0.3, wires=0)
             qml.WireCut(wires=0)
             qml.RY(0.4, wires=0)
-            qml.apply(measurement)
+            return qml.apply(measurement)
 
-        tape = qml.tape.QuantumScript.from_queue(q)
         spy = mocker.spy(qcut.cutcircuit, "_qcut_expand_fn")
         spy_mc = mocker.spy(qcut.montecarlo, "_qcut_expand_fn")
 
         kwargs = {"shots": 10} if measurement.return_type is qml.measurements.Sample else {}
-        cut_transform(tape, device_wires=[0], **kwargs)
+        cut_transform(circuit, device_wires=[0])(**kwargs)
 
         assert spy.call_count == 1 or spy_mc.call_count == 1
 
@@ -4537,20 +4542,17 @@ def test_expansion(self, mocker, cut_transform, measurement):
     @pytest.mark.parametrize("cut_transform, measurement", transform_measurement_pairs)
     def test_expansion_error(self, cut_transform, measurement):
         """Test if a ValueError is raised if expansion continues beyond the maximum depth"""
-        with qml.queuing.AnnotatedQueue() as q:
+        dev = qml.device("default.qubit", wires=3)
+
+        @qml.qnode(dev)
+        def circuit():
             qml.RX(0.3, wires=0)
-            with qml.queuing.AnnotatedQueue() as q__:
-                with qml.queuing.AnnotatedQueue() as q___:
-                    qml.WireCut(wires=0)
-                __ = qml.tape.QuantumScript.from_queue(q___)
-            _ = qml.tape.QuantumScript.from_queue(q__)
             qml.RY(0.4, wires=0)
-            qml.apply(measurement)
+            return qml.apply(measurement)
 
-        tape = qml.tape.QuantumScript.from_queue(q)
         with pytest.raises(ValueError, match="No WireCut operations found in the circuit."):
             kwargs = {"shots": 10} if measurement.return_type is qml.measurements.Sample else {}
-            cut_transform(tape, device_wires=[0], **kwargs)
+            cut_transform(circuit, device_wires=[0])(**kwargs)
 
     def test_expansion_ttn(self, mocker):
         """Test if wire cutting is compatible with the tree tensor network operation"""
@@ -5276,7 +5278,7 @@ def test_circuit_with_disconnected_components(self):
 
         dev = qml.device("default.qubit", wires=3)
 
-        @qml.transforms.cut_circuit(auto_cutter=True)
+        @partial(qml.transforms.cut_circuit, auto_cutter=True)
         @qml.qnode(dev)
         def circuit(x):
             qml.RX(x, wires=0)
@@ -5296,7 +5298,7 @@ def test_circuit_with_trivial_wire_cut(self):
 
         dev = qml.device("default.qubit", wires=2)
 
-        @qml.transforms.cut_circuit(auto_cutter=True)
+        @partial(qml.transforms.cut_circuit, auto_cutter=True)
         @qml.qnode(dev)
         def circuit(x):
             qml.RX(x, wires=0)
diff --git a/tests/transforms/test_sign_expand.py b/tests/transforms/test_sign_expand.py
index d59b56022c2..2641d749f1b 100644
--- a/tests/transforms/test_sign_expand.py
+++ b/tests/transforms/test_sign_expand.py
@@ -10,7 +10,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 """This file contains unit tests for the ``sign_expand`` transform."""
+import functools
 import pytest
+
 import numpy as np
 import pennylane as qml
 import pennylane.tape
@@ -101,6 +103,21 @@ def test_hamiltonians(self, tape, output):
 
         assert np.isclose(output, expval)
 
+    @pytest.mark.parametrize(("tape", "output"), zip(TAPES, OUTPUTS))
+    def test_hamiltonians_qnode(self, tape, output):
+        """Tests that the sign_expand transform returns the correct value as a transform program"""
+
+        @qml.transforms.sign_expand
+        @qml.qnode(dev)
+        def qnode():
+            for op in tape.operations:
+                qml.apply(op)
+
+            return qml.apply(tape.measurements[0])
+
+        expval = qnode()
+        assert np.isclose(output, expval)
+
     @pytest.mark.parametrize(("tape", "output"), zip(TAPES, OUTPUTS))
     def test_hamiltonians_circuit_impl(self, tape, output):
         """Tests that the sign_expand transform returns the correct value
@@ -114,6 +131,24 @@ def test_hamiltonians_circuit_impl(self, tape, output):
         assert np.isclose(output, expval, 1e-2)
         # as these are approximations, these are only correct up to finite precision
 
+    @pytest.mark.parametrize(("tape", "output"), zip(TAPES, OUTPUTS))
+    def test_hamiltonians_circuit_impl_qnode(self, tape, output):
+        """Tests that the sign_expand transform returns the correct value as a transform program
+        if we do not calculate analytical expectation values of groups but rely on their circuit approximations
+        """
+
+        @functools.partial(qml.transforms.sign_expand, circuit=True)
+        @qml.qnode(dev)
+        def qnode():
+            for op in tape.operations:
+                qml.apply(op)
+
+            return qml.apply(tape.measurements[0])
+
+        expval = qnode()
+        assert np.isclose(output, expval, 1e-2)
+        # as these are approximations, these are only correct up to finite precision
+
     @pytest.mark.parametrize("shots", [None, 100])
     @pytest.mark.parametrize("circuit", [True, False])
     def test_shots_attribute(self, shots, circuit):
diff --git a/tests/transforms/test_split_non_commuting.py b/tests/transforms/test_split_non_commuting.py
index c7eccfc5be0..fb0399e48ba 100644
--- a/tests/transforms/test_split_non_commuting.py
+++ b/tests/transforms/test_split_non_commuting.py
@@ -68,13 +68,11 @@ def test_commuting_group_no_split(self, mocker):
 
         spy = mocker.spy(qml.math, "concatenate")
 
-        assert split == [tape]
         assert all(isinstance(t, qml.tape.QuantumScript) for t in split)
         assert fn([0.5]) == 0.5
 
         qs = qml.tape.QuantumScript(tape.operations, tape.measurements)
         split, fn = split_non_commuting(qs)
-        assert split == [qs]
         assert all(isinstance(i_qs, qml.tape.QuantumScript) for i_qs in split)
         assert fn([0.5]) == 0.5
 
@@ -202,6 +200,38 @@ def circuit():
 
         assert all(np.isclose(res, np.array([0.0, -1.0, 0.0, 0.0, 1.0, 1 / np.sqrt(2)])))
 
+    def test_expval_non_commuting_observables_qnode(self):
+        """Test expval with multiple non-commuting operators as a tranform program on the qnode."""
+        dev = qml.device("default.qubit", wires=6)
+
+        @qml.qnode(dev)
+        def circuit():
+            qml.Hadamard(1)
+            qml.Hadamard(0)
+            qml.PauliZ(0)
+            qml.Hadamard(3)
+            qml.Hadamard(5)
+            qml.T(5)
+            return (
+                qml.expval(qml.PauliZ(0) @ qml.PauliZ(1)),
+                qml.expval(qml.PauliX(0)),
+                qml.expval(qml.PauliZ(1)),
+                qml.expval(qml.PauliX(1) @ qml.PauliX(4)),
+                qml.expval(qml.PauliX(3)),
+                qml.expval(qml.PauliY(5)),
+            )
+
+        res = split_non_commuting(circuit)()
+
+        assert isinstance(res, tuple)
+        assert len(res) == 6
+        assert all(isinstance(r, np.ndarray) for r in res)
+        assert all(r.shape == () for r in res)
+
+        res = qml.math.stack(res)
+
+        assert all(np.isclose(res, np.array([0.0, -1.0, 0.0, 0.0, 1.0, 1 / np.sqrt(2)])))
+
     def test_shot_vector_support(self):
         """Test output is correct when using shot vectors"""
 

From ebb4d91d9c6c71e72e59da785b0e39b3c685cfdc Mon Sep 17 00:00:00 2001
From: Matthew Silverman <matthews@xanadu.ai>
Date: Wed, 6 Sep 2023 14:50:15 -0400
Subject: [PATCH 034/131] Update measurements tests for DQ2; create legacy
 tests folder (#4574)

**Context:**
Measurement tests are focused on legacy devices, but we should use the
new device API for these tests.

**Description of the Change:**
Here's what I did, chronologically:
1. copy-paste every measurement test file into the `legacy` sub-folder
2. change the original files to use `default.qubit` (not DQL), change
`default.qubit` to point to DQ2, fix tests so they work with DQ2 (they
all passed)
3. Put back `default.qubit` to point to DQL, re-run non-legacy tests to
see which ones fail, mark those as xfail (for now)
4. Remove all tests in the `legacy` folder that double-cover anything
(in other words, any tests that didn't use DQL)

Step 2 above involved some changes to PennyLane source code, and they
are:
1. Introduce `MutualInfoMP.map_wires`
2. Change the qinfo transforms to use tape wires if device wires aren't
present (I will probably revert this once device wires are supported for
DQ2)
3. Fix `sampling.py` to not squeeze any CountsMP result (batched results
also shouldn't be squeezed, but they are lists of dicts instead of just
dicts)
4. EDIT: Update `ProbabilityMP.marginal_prob` to just return `prob` if
`self.wires` is None. This was commented out before, and I'm not sure
why.

**Benefits:**
- Migrating to DQ2 will be much easier, and we'll have better test
coverage for measurements with the new device API
- `qml.probs()` did not work with DQ2 before this little change D:

**Possible Drawbacks:**
Slightly more tests. Think of this as effectively adding one more device
to test with whenever doing device integration tests.

**Follow-Up:**
I made [this
story](https://app.shortcut.com/xanaduai/story/44915/update-setup-py-to-point-default-qubit-to-defaultqubit2)
and mentioned the removal of xfails in it.
---
 doc/releases/changelog-dev.md                 |   5 +
 pennylane/devices/qubit/sampling.py           |   4 +-
 pennylane/measurements/mutual_info.py         |   8 +
 pennylane/measurements/probs.py               |   7 +-
 .../legacy/test_classical_shadow_legacy.py    | 571 ++++++++++++
 .../measurements/legacy/test_counts_legacy.py | 641 ++++++++++++++
 .../measurements/legacy/test_expval_legacy.py | 210 +++++
 .../legacy/test_measurements_legacy.py        | 220 +++++
 .../legacy/test_mutual_info_legacy.py         | 433 ++++++++++
 .../measurements/legacy/test_probs_legacy.py  | 634 ++++++++++++++
 .../legacy/test_purity_measurement_legacy.py  | 303 +++++++
 .../measurements/legacy/test_sample_legacy.py | 418 +++++++++
 .../measurements/legacy/test_state_legacy.py  | 809 ++++++++++++++++++
 tests/measurements/legacy/test_var_legacy.py  | 200 +++++
 .../legacy/test_vn_entropy_legacy.py          | 394 +++++++++
 tests/measurements/test_classical_shadow.py   | 113 +--
 tests/measurements/test_counts.py             | 225 ++---
 tests/measurements/test_expval.py             |  75 +-
 tests/measurements/test_measurements.py       |  94 +-
 tests/measurements/test_mutual_info.py        |  80 +-
 tests/measurements/test_probs.py              | 162 +---
 tests/measurements/test_purity_measurement.py |   6 +-
 tests/measurements/test_sample.py             | 123 +--
 tests/measurements/test_state.py              | 431 ++++++----
 tests/measurements/test_var.py                |  73 +-
 tests/measurements/test_vn_entropy.py         |  53 +-
 26 files changed, 5443 insertions(+), 849 deletions(-)
 create mode 100644 tests/measurements/legacy/test_classical_shadow_legacy.py
 create mode 100644 tests/measurements/legacy/test_counts_legacy.py
 create mode 100644 tests/measurements/legacy/test_expval_legacy.py
 create mode 100644 tests/measurements/legacy/test_measurements_legacy.py
 create mode 100644 tests/measurements/legacy/test_mutual_info_legacy.py
 create mode 100644 tests/measurements/legacy/test_probs_legacy.py
 create mode 100644 tests/measurements/legacy/test_purity_measurement_legacy.py
 create mode 100644 tests/measurements/legacy/test_sample_legacy.py
 create mode 100644 tests/measurements/legacy/test_state_legacy.py
 create mode 100644 tests/measurements/legacy/test_var_legacy.py
 create mode 100644 tests/measurements/legacy/test_vn_entropy_legacy.py

diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index a4d274d5e53..59afddf3409 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -57,6 +57,11 @@
   If the process's wire-order contains extra wires, it will assume those are in the zero-state.
   [(#4570)](https://github.com/PennyLaneAI/pennylane/pull/4570)
 
+* Various changes to measurements to improve feature parity between the legacy `default.qubit` and
+  the new `DefaultQubit2`. This includes not trying to squeeze batched `CountsMP` results and implementing
+  `MutualInfoMP.map_wires`.
+  [(#4574)](https://github.com/PennyLaneAI/pennylane/pull/4574)
+
 <h3>Breaking changes 💔</h3>
 
 * The `__eq__` and `__hash__` methods of `Operator` and `MeasurementProcess` no longer rely on the
diff --git a/pennylane/devices/qubit/sampling.py b/pennylane/devices/qubit/sampling.py
index 2084c533376..0a2acd9dad1 100644
--- a/pennylane/devices/qubit/sampling.py
+++ b/pennylane/devices/qubit/sampling.py
@@ -23,6 +23,7 @@
     ExpectationMP,
     ClassicalShadowMP,
     ShadowExpvalMP,
+    CountsMP,
 )
 from pennylane.typing import TensorLike
 from .apply_operation import apply_operation
@@ -194,7 +195,8 @@ def _measure_with_samples_diagonalizing_gates(
     def _process_single_shot(samples):
         processed = []
         for mp in mps:
-            if not isinstance(res := mp.process_samples(samples, wires), dict):
+            res = mp.process_samples(samples, wires)
+            if not isinstance(mp, CountsMP):
                 res = qml.math.squeeze(res)
 
             processed.append(res)
diff --git a/pennylane/measurements/mutual_info.py b/pennylane/measurements/mutual_info.py
index 2659cd2296e..e50f1bb9586 100644
--- a/pennylane/measurements/mutual_info.py
+++ b/pennylane/measurements/mutual_info.py
@@ -15,6 +15,7 @@
 """
 This module contains the qml.mutual_info measurement.
 """
+from copy import copy
 from typing import Sequence, Optional
 
 import pennylane as qml
@@ -131,6 +132,13 @@ def return_type(self):
     def numeric_type(self):
         return float
 
+    def map_wires(self, wire_map: dict):
+        new_measurement = copy(self)
+        new_measurement._wires = [
+            Wires([wire_map.get(wire, wire) for wire in wires]) for wires in self.raw_wires
+        ]
+        return new_measurement
+
     def shape(self, device, shots):
         if not shots.has_partitioned_shots:
             return ()
diff --git a/pennylane/measurements/probs.py b/pennylane/measurements/probs.py
index 1dcb95ebf99..5aa6b21adf7 100644
--- a/pennylane/measurements/probs.py
+++ b/pennylane/measurements/probs.py
@@ -268,10 +268,9 @@ def marginal_prob(self, prob, wire_order, batch_size):
         Returns:
             array[float]: array of the resulting marginal probabilities.
         """
-        # TODO: Add when ``qml.probs()`` is supported
-        # if self.wires == Wires([]):
-        #     # no need to marginalize
-        #     return prob
+        if self.wires == Wires([]):
+            # no need to marginalize
+            return prob
 
         # determine which subsystems are to be summed over
         inactive_wires = Wires.unique_wires([wire_order, self.wires])
diff --git a/tests/measurements/legacy/test_classical_shadow_legacy.py b/tests/measurements/legacy/test_classical_shadow_legacy.py
new file mode 100644
index 00000000000..7066f26f068
--- /dev/null
+++ b/tests/measurements/legacy/test_classical_shadow_legacy.py
@@ -0,0 +1,571 @@
+# Copyright 2022 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 classical shadows measurement processes"""
+
+import autograd.numpy
+import pytest
+
+import pennylane as qml
+from pennylane import numpy as np
+from pennylane.measurements import Shots
+
+# pylint: disable=dangerous-default-value, too-many-arguments
+
+
+def get_circuit(wires, shots, seed_recipes, interface="autograd", device="default.qubit.legacy"):
+    """
+    Return a QNode that prepares the state (|00...0> + |11...1>) / sqrt(2)
+        and performs the classical shadow measurement
+    """
+    if device is not None:
+        dev = qml.device(device, wires=wires, shots=shots)
+    else:
+        dev = qml.device("default.qubit.legacy", wires=wires, shots=shots)
+
+        # make the device call the superclass method to switch between the general qubit device and device specific implementations (i.e. for default qubit)
+        dev.classical_shadow = super(type(dev), dev).classical_shadow
+
+    @qml.qnode(dev, interface=interface)
+    def circuit():
+        qml.Hadamard(wires=0)
+
+        for target in range(1, wires):
+            qml.CNOT(wires=[0, target])
+
+        return qml.classical_shadow(wires=range(wires), seed=seed_recipes)
+
+    return circuit
+
+
+def get_x_basis_circuit(wires, shots, interface="autograd", device="default.qubit.legacy"):
+    """
+    Return a QNode that prepares the |++..+> state and performs a classical shadow measurement
+    """
+    if device is not None:
+        dev = qml.device(device, wires=wires, shots=shots)
+    else:
+        dev = qml.device("default.qubit.legacy", wires=wires, shots=shots)
+
+        # make the device call the superclass method to switch between the general qubit device and device specific implementations (i.e. for default qubit)
+        dev.classical_shadow = super(type(dev), dev).classical_shadow
+
+    @qml.qnode(dev, interface=interface)
+    def circuit():
+        for wire in range(wires):
+            qml.Hadamard(wire)
+        return qml.classical_shadow(wires=range(wires))
+
+    return circuit
+
+
+def get_y_basis_circuit(wires, shots, interface="autograd", device="default.qubit.legacy"):
+    """
+    Return a QNode that prepares the |+i>|+i>...|+i> state and performs a classical shadow measurement
+    """
+    if device is not None:
+        dev = qml.device(device, wires=wires, shots=shots)
+    else:
+        dev = qml.device("default.qubit.legacy", wires=wires, shots=shots)
+
+        # make the device call the superclass method to switch between the general qubit device and device specific implementations (i.e. for default qubit)
+        dev.classical_shadow = super(type(dev), dev).classical_shadow
+
+    @qml.qnode(dev, interface=interface)
+    def circuit():
+        for wire in range(wires):
+            qml.Hadamard(wire)
+            qml.RZ(np.pi / 2, wire)
+        return qml.classical_shadow(wires=range(wires))
+
+    return circuit
+
+
+def get_z_basis_circuit(wires, shots, interface="autograd", device="default.qubit.legacy"):
+    """
+    Return a QNode that prepares the |00..0> state and performs a classical shadow measurement
+    """
+    if device is not None:
+        dev = qml.device(device, wires=wires, shots=shots)
+    else:
+        dev = qml.device("default.qubit.legacy", wires=wires, shots=shots)
+
+        # make the device call the superclass method to switch between the general qubit device and device specific implementations (i.e. for default qubit)
+        dev.classical_shadow = super(type(dev), dev).classical_shadow
+
+    @qml.qnode(dev, interface=interface)
+    def circuit():
+        return qml.classical_shadow(wires=range(wires))
+
+    return circuit
+
+
+wires_list = [1, 3]
+
+
+@pytest.mark.parametrize("wires", wires_list)
+class TestClassicalShadow:
+    """Unit tests for classical_shadow measurement"""
+
+    shots_list = [1, 100]
+    seed_recipes_list = [None, 74]  # random seed
+
+    @pytest.mark.parametrize("shots", shots_list)
+    @pytest.mark.parametrize("seed", seed_recipes_list)
+    def test_measurement_process_shape(self, wires, shots, seed):
+        """Test that the shape of the MeasurementProcess instance is correct"""
+        dev = qml.device("default.qubit.legacy", wires=wires, shots=shots)
+        shots_obj = Shots(shots)
+        res = qml.classical_shadow(wires=range(wires), seed=seed)
+        assert res.shape(dev, shots_obj) == (2, shots, wires)
+
+        # test an error is raised when device is None
+        msg = "Shots must be specified to obtain the shape of a classical shadow measurement"
+        with pytest.raises(qml.measurements.MeasurementShapeError, match=msg):
+            res.shape(dev, Shots(None))
+
+    def test_shape_matches(self, wires):
+        """Test that the shape of the MeasurementProcess matches the shape
+        of the tape execution"""
+        shots = 100
+
+        circuit = get_circuit(wires, shots, True)
+        circuit.construct((), {})
+
+        res = qml.execute([circuit.tape], circuit.device, None)[0]
+        expected_shape = qml.classical_shadow(wires=range(wires)).shape(
+            circuit.device, Shots(shots)
+        )
+
+        assert res.shape == expected_shape
+
+    @pytest.mark.all_interfaces
+    @pytest.mark.parametrize("shots", shots_list)
+    @pytest.mark.parametrize("seed", seed_recipes_list)
+    @pytest.mark.parametrize("interface", ["autograd", "jax", "tf", "torch"])
+    @pytest.mark.parametrize("device", ["default.qubit.legacy", None])
+    def test_format(self, wires, shots, seed, interface, device):
+        """Test that the format of the returned classical shadow
+        measurement is correct"""
+        import tensorflow as tf
+        import torch
+
+        circuit = get_circuit(wires, shots, seed, interface, device)
+        shadow = circuit()
+
+        # test shape is correct
+        assert shadow.shape == (2, shots, wires)
+
+        # test dtype is correct
+        expected_dtype = np.int8
+        if interface == "tf":
+            expected_dtype = tf.int8
+        elif interface == "torch":
+            expected_dtype = torch.int8
+
+        assert shadow.dtype == expected_dtype
+
+        bits, recipes = shadow  # pylint: disable=unpacking-non-sequence
+
+        # test allowed values of bits and recipes
+        assert qml.math.all(np.logical_or(bits == 0, bits == 1))
+        assert qml.math.all(np.logical_or(recipes == 0, np.logical_or(recipes == 1, recipes == 2)))
+
+    @pytest.mark.all_interfaces
+    @pytest.mark.parametrize("interface", ["autograd", "jax", "tf", "torch"])
+    @pytest.mark.parametrize("device", ["default.qubit.legacy", None])
+    @pytest.mark.parametrize(
+        "circuit_fn, basis_recipe",
+        [(get_x_basis_circuit, 0), (get_y_basis_circuit, 1), (get_z_basis_circuit, 2)],
+    )
+    def test_return_distribution(self, wires, interface, device, circuit_fn, basis_recipe):
+        """Test that the distribution of the bits and recipes are correct for a circuit
+        that prepares all qubits in a Pauli basis"""
+        # high number of shots to prevent true negatives
+        np.random.seed(42)
+        shots = 1000
+
+        circuit = circuit_fn(wires, shots=shots, interface=interface, device=device)
+        bits, recipes = circuit()
+        new_bits, new_recipes = circuit.tape.measurements[0].process(circuit.tape, circuit.device)
+
+        # test that the recipes follow a rough uniform distribution
+        ratios = np.unique(recipes, return_counts=True)[1] / (wires * shots)
+        assert np.allclose(ratios, 1 / 3, atol=1e-1)
+        new_ratios = np.unique(new_recipes, return_counts=True)[1] / (wires * shots)
+        assert np.allclose(new_ratios, 1 / 3, atol=1e-1)
+
+        # test that the bit is 0 for all X measurements
+        assert qml.math.allequal(bits[recipes == basis_recipe], 0)
+        assert qml.math.allequal(new_bits[new_recipes == basis_recipe], 0)
+
+        # test that the bits are uniformly distributed for all Y and Z measurements
+        bits1 = bits[recipes == (basis_recipe + 1) % 3]
+        ratios1 = np.unique(bits1, return_counts=True)[1] / bits1.shape[0]
+        assert np.allclose(ratios1, 1 / 2, atol=1e-1)
+        new_bits1 = new_bits[new_recipes == (basis_recipe + 1) % 3]
+        new_ratios1 = np.unique(new_bits1, return_counts=True)[1] / new_bits1.shape[0]
+        assert np.allclose(new_ratios1, 1 / 2, atol=1e-1)
+
+        bits2 = bits[recipes == (basis_recipe + 2) % 3]
+        ratios2 = np.unique(bits2, return_counts=True)[1] / bits2.shape[0]
+        assert np.allclose(ratios2, 1 / 2, atol=1e-1)
+
+        new_bits2 = new_bits[new_recipes == (basis_recipe + 2) % 3]
+        new_ratios2 = np.unique(new_bits2, return_counts=True)[1] / new_bits2.shape[0]
+        assert np.allclose(new_ratios2, 1 / 2, atol=1e-1)
+
+    @pytest.mark.parametrize("shots", shots_list)
+    def test_multi_measurement_error(self, wires, shots):
+        """Test that an error is raised when classical shadows is returned
+        with other measurement processes"""
+        dev = qml.device("default.qubit.legacy", wires=wires, shots=shots)
+
+        @qml.qnode(dev)
+        def circuit():
+            qml.Hadamard(wires=0)
+
+            for target in range(1, wires):
+                qml.CNOT(wires=[0, target])
+
+            return qml.classical_shadow(wires=range(wires)), qml.expval(qml.PauliZ(0))
+
+        msg = "Classical shadows cannot be returned in combination with other return types"
+        with pytest.raises(qml.QuantumFunctionError, match=msg):
+            circuit()
+
+
+def hadamard_circuit(wires, shots=10000, interface="autograd"):
+    dev = qml.device("default.qubit.legacy", wires=wires, shots=shots)
+
+    @qml.qnode(dev, interface=interface)
+    def circuit(obs, k=1):
+        for i in range(wires):
+            qml.Hadamard(wires=i)
+        return qml.shadow_expval(obs, k=k)
+
+    return circuit
+
+
+def max_entangled_circuit(wires, shots=10000, interface="autograd"):
+    dev = qml.device("default.qubit.legacy", wires=wires, shots=shots)
+
+    @qml.qnode(dev, interface=interface)
+    def circuit(obs, k=1):
+        qml.Hadamard(wires=0)
+        for i in range(1, wires):
+            qml.CNOT(wires=[0, i])
+        return qml.shadow_expval(obs, k=k)
+
+    return circuit
+
+
+def qft_circuit(wires, shots=10000, interface="autograd"):
+    dev = qml.device("default.qubit.legacy", wires=wires, shots=shots)
+
+    one_state = np.zeros(wires)
+    one_state[-1] = 1
+
+    @qml.qnode(dev, interface=interface)
+    def circuit(obs, k=1):
+        qml.BasisState(one_state, wires=range(wires))
+        qml.QFT(wires=range(wires))
+        return qml.shadow_expval(obs, k=k)
+
+    return circuit
+
+
+@pytest.mark.autograd
+class TestExpvalMeasurement:
+    @pytest.mark.parametrize("wires", [1, 2])
+    @pytest.mark.parametrize("shots", [1, 10])
+    def test_measurement_process_shape(self, wires, shots):
+        """Test that the shape of the MeasurementProcess instance is correct"""
+        dev = qml.device("default.qubit.legacy", wires=wires, shots=shots)
+        H = qml.PauliZ(0)
+        res = qml.shadow_expval(H)
+        assert len(res.shape(dev, Shots(shots))) == 0
+
+    def test_shape_matches(self):
+        """Test that the shape of the MeasurementProcess matches the shape
+        of the tape execution"""
+        wires = 2
+        shots = 100
+        H = qml.PauliZ(0)
+
+        circuit = hadamard_circuit(wires, shots)
+        circuit.construct((H,), {})
+
+        res = qml.execute([circuit.tape], circuit.device, None)[0]
+        expected_shape = qml.shadow_expval(H).shape(circuit.device, Shots(shots))
+
+        assert res.shape == expected_shape
+
+    def test_shots_none_error(self):
+        """Test that an error is raised when a device with shots=None is used
+        to obtain classical shadows"""
+        circuit = hadamard_circuit(2, None)
+        H = qml.PauliZ(0)
+
+        msg = "The number of shots has to be explicitly set on the device when using sample-based measurements"
+        with pytest.raises(qml.QuantumFunctionError, match=msg):
+            _ = circuit(H, k=10)
+
+    def test_multi_measurement_error(self):
+        """Test that an error is raised when classical shadows is returned
+        with other measurement processes"""
+        dev = qml.device("default.qubit.legacy", wires=2, shots=100)
+
+        @qml.qnode(dev)
+        def circuit():
+            qml.Hadamard(wires=0)
+            qml.CNOT(wires=[0, 1])
+            return qml.shadow_expval(qml.PauliZ(0)), qml.expval(qml.PauliZ(0))
+
+        msg = "Classical shadows cannot be returned in combination with other return types"
+        with pytest.raises(qml.QuantumFunctionError, match=msg):
+            _ = circuit()
+
+
+obs_hadamard = [
+    qml.PauliX(1),
+    qml.PauliX(0) @ qml.PauliX(2),
+    qml.PauliX(0) @ qml.Identity(1) @ qml.PauliX(2),
+    qml.PauliY(2),
+    qml.PauliY(1) @ qml.PauliZ(2),
+    qml.PauliX(0) @ qml.PauliY(1),
+    qml.PauliX(0) @ qml.PauliY(1) @ qml.Identity(2),
+]
+expected_hadamard = [1, 1, 1, 0, 0, 0, 0]
+
+obs_max_entangled = [
+    qml.PauliX(1),
+    qml.PauliX(0) @ qml.PauliX(2),
+    qml.PauliZ(2),
+    qml.Identity(1) @ qml.PauliZ(2),
+    qml.PauliZ(1) @ qml.PauliZ(2),
+    qml.PauliX(0) @ qml.PauliY(1),
+    qml.PauliX(0) @ qml.PauliY(1) @ qml.Identity(2),
+    qml.PauliY(0) @ qml.PauliX(1) @ qml.PauliY(2),
+]
+expected_max_entangled = [0, 0, 0, 0, 1, 0, 0, -1]
+
+obs_qft = [
+    qml.PauliX(0),
+    qml.PauliX(0) @ qml.PauliX(1),
+    qml.PauliX(0) @ qml.PauliX(2),
+    qml.PauliX(0) @ qml.Identity(1) @ qml.PauliX(2),
+    qml.PauliZ(2),
+    qml.PauliX(1) @ qml.PauliY(2),
+    qml.PauliY(1) @ qml.PauliX(2),
+    qml.Identity(0) @ qml.PauliY(1) @ qml.PauliX(2),
+    qml.PauliX(0) @ qml.PauliY(1) @ qml.PauliY(2),
+    qml.PauliY(0) @ qml.PauliX(1) @ qml.PauliX(2),
+]
+expected_qft = [
+    -1,
+    0,
+    -1 / np.sqrt(2),
+    -1 / np.sqrt(2),
+    0,
+    0,
+    1 / np.sqrt(2),
+    1 / np.sqrt(2),
+    -1 / np.sqrt(2),
+    0,
+]
+
+
+@pytest.mark.autograd
+class TestExpvalForward:
+    """Test the shadow_expval measurement process forward pass"""
+
+    def test_hadamard_expval(self, k=1, obs=obs_hadamard, expected=expected_hadamard):
+        """Test that the expval estimation is correct for a uniform
+        superposition of qubits"""
+        circuit = hadamard_circuit(3, shots=100000)
+        actual = circuit(obs, k=k)
+        new_actual = circuit.tape.measurements[0].process(circuit.tape, circuit.device)
+
+        assert actual.shape == (len(obs_hadamard),)
+        assert actual.dtype == np.float64
+        assert qml.math.allclose(actual, expected, atol=1e-1)
+        assert qml.math.allclose(new_actual, expected, atol=1e-1)
+
+    def test_max_entangled_expval(
+        self, k=1, obs=obs_max_entangled, expected=expected_max_entangled
+    ):
+        """Test that the expval estimation is correct for a maximally
+        entangled state"""
+        circuit = max_entangled_circuit(3, shots=100000)
+        actual = circuit(obs, k=k)
+        new_actual = circuit.tape.measurements[0].process(circuit.tape, circuit.device)
+
+        assert actual.shape == (len(obs_max_entangled),)
+        assert actual.dtype == np.float64
+        assert qml.math.allclose(actual, expected, atol=1e-1)
+        assert qml.math.allclose(new_actual, expected, atol=1e-1)
+
+    def test_non_pauli_error(self):
+        """Test that an error is raised when a non-Pauli observable is passed"""
+        circuit = hadamard_circuit(3)
+
+        msg = "Observable must be a linear combination of Pauli observables"
+        with pytest.raises(ValueError, match=msg):
+            circuit(qml.Hadamard(0) @ qml.Hadamard(2))
+
+
+# pylint: disable=too-few-public-methods
+@pytest.mark.all_interfaces
+class TestExpvalForwardInterfaces:
+    @pytest.mark.parametrize("interface", ["autograd", "jax", "tf", "torch"])
+    def test_qft_expval(self, interface, k=1, obs=obs_qft, expected=expected_qft):
+        """Test that the expval estimation is correct for a QFT state"""
+        import torch
+
+        circuit = qft_circuit(3, shots=100000, interface=interface)
+        actual = circuit(obs, k=k)
+        new_actual = circuit.tape.measurements[0].process(circuit.tape, circuit.device)
+
+        assert actual.shape == (len(obs_qft),)
+        assert actual.dtype == torch.float64 if interface == "torch" else np.float64
+        assert qml.math.allclose(actual, expected, atol=1e-1)
+        assert qml.math.allclose(new_actual, expected, atol=1e-1)
+
+
+obs_strongly_entangled = [
+    qml.PauliX(1),
+    qml.PauliX(0) @ qml.PauliX(2),
+    qml.PauliX(0) @ qml.Identity(1) @ qml.PauliX(2),
+    qml.PauliY(2),
+    qml.PauliY(1) @ qml.PauliZ(2),
+    qml.PauliX(0) @ qml.PauliY(1),
+    qml.PauliX(0) @ qml.PauliY(1) @ qml.Identity(2),
+]
+
+
+def strongly_entangling_circuit(wires, shots=10000, interface="autograd"):
+    dev = qml.device("default.qubit.legacy", wires=wires, shots=shots)
+
+    @qml.qnode(dev, interface=interface)
+    def circuit(x, obs, k):
+        qml.StronglyEntanglingLayers(weights=x, wires=range(wires))
+        return qml.shadow_expval(obs, k=k)
+
+    return circuit
+
+
+def strongly_entangling_circuit_exact(wires, interface="autograd"):
+    dev = qml.device("default.qubit.legacy", wires=wires)
+
+    @qml.qnode(dev, interface=interface)
+    def circuit(x, obs):
+        qml.StronglyEntanglingLayers(weights=x, wires=range(wires))
+        return [qml.expval(o) for o in obs]
+
+    return circuit
+
+
+class TestExpvalBackward:
+    """Test the shadow_expval measurement process backward pass"""
+
+    @pytest.mark.autograd
+    def test_backward_autograd(self, obs=obs_strongly_entangled):
+        """Test that the gradient of the expval estimation is correct for
+        the autograd interface"""
+        shadow_circuit = strongly_entangling_circuit(3, shots=20000, interface="autograd")
+        exact_circuit = strongly_entangling_circuit_exact(3, "autograd")
+
+        def cost_exact(x, obs):
+            return autograd.numpy.hstack(exact_circuit(x, obs))
+
+        # make rotations close to pi / 2 to ensure gradients are not too small
+        x = np.random.uniform(
+            0.8, 2, size=qml.StronglyEntanglingLayers.shape(n_layers=2, n_wires=3)
+        )
+        actual = qml.jacobian(shadow_circuit)(x, obs, k=1)
+        expected = qml.jacobian(cost_exact, argnum=0)(x, obs)
+
+        assert qml.math.allclose(actual, expected, atol=1e-1)
+
+    @pytest.mark.jax
+    def test_backward_jax(self, obs=obs_strongly_entangled):
+        """Test that the gradient of the expval estimation is correct for
+        the jax interface"""
+        import jax
+        from jax import numpy as jnp
+
+        shadow_circuit = strongly_entangling_circuit(3, shots=20000, interface="jax")
+        exact_circuit = strongly_entangling_circuit_exact(3, "jax")
+
+        # make rotations close to pi / 2 to ensure gradients are not too small
+        x = jnp.array(
+            np.random.uniform(
+                0.8, 2, size=qml.StronglyEntanglingLayers.shape(n_layers=2, n_wires=3)
+            )
+        )
+
+        actual = jax.jacrev(shadow_circuit)(x, obs, k=1)
+        expected = jax.jacrev(exact_circuit)(x, obs)
+
+        assert qml.math.allclose(actual, expected, atol=1e-1)
+
+    @pytest.mark.tf
+    def test_backward_tf(self, obs=obs_strongly_entangled):
+        """Test that the gradient of the expval estimation is correct for
+        the tensorflow interface"""
+        import tensorflow as tf
+
+        shadow_circuit = strongly_entangling_circuit(3, shots=20000, interface="tf")
+        exact_circuit = strongly_entangling_circuit_exact(3, "tf")
+
+        # make rotations close to pi / 2 to ensure gradients are not too small
+        x = tf.Variable(
+            np.random.uniform(
+                0.8, 2, size=qml.StronglyEntanglingLayers.shape(n_layers=2, n_wires=3)
+            )
+        )
+
+        with tf.GradientTape() as tape:
+            out = shadow_circuit(x, obs, k=10)
+
+        actual = tape.jacobian(out, x)
+
+        with tf.GradientTape() as tape2:
+            out2 = qml.math.hstack(exact_circuit(x, obs))
+
+        expected = tape2.jacobian(out2, x)
+
+        assert qml.math.allclose(actual, expected, atol=1e-1)
+
+    @pytest.mark.torch
+    def test_backward_torch(self, obs=obs_strongly_entangled):
+        """Test that the gradient of the expval estimation is correct for
+        the pytorch interface"""
+        import torch
+
+        shadow_circuit = strongly_entangling_circuit(3, shots=20000, interface="torch")
+        exact_circuit = strongly_entangling_circuit_exact(3, "torch")
+
+        # make rotations close to pi / 2 to ensure gradients are not too small
+        x = torch.tensor(
+            np.random.uniform(
+                0.8, 2, size=qml.StronglyEntanglingLayers.shape(n_layers=2, n_wires=3)
+            ),
+            requires_grad=True,
+        )
+
+        actual = torch.autograd.functional.jacobian(lambda x: shadow_circuit(x, obs, k=10), x)
+        expected = torch.autograd.functional.jacobian(lambda x: tuple(exact_circuit(x, obs)), x)
+
+        assert qml.math.allclose(actual, qml.math.stack(expected), atol=1e-1)
diff --git a/tests/measurements/legacy/test_counts_legacy.py b/tests/measurements/legacy/test_counts_legacy.py
new file mode 100644
index 00000000000..58b0a95966d
--- /dev/null
+++ b/tests/measurements/legacy/test_counts_legacy.py
@@ -0,0 +1,641 @@
+# Copyright 2018-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 the qml.counts measurement process."""
+import numpy as np
+import pytest
+
+import pennylane as qml
+from pennylane.measurements import AllCounts, Counts
+from pennylane.operation import Operator
+
+
+# TODO: Remove this when new CustomMP are the default
+def custom_measurement_process(device, spy):
+    assert len(spy.call_args_list) > 0  # make sure method is mocked properly
+
+    samples = device._samples  # pylint:disable=protected-access
+    call_args_list = list(spy.call_args_list)
+    for call_args in call_args_list:
+        if not call_args.kwargs.get("counts", False):
+            continue
+        meas = call_args.args[1]
+        shot_range, bin_size = (call_args.kwargs["shot_range"], call_args.kwargs["bin_size"])
+        if isinstance(meas, Operator):
+            all_outcomes = meas.return_type is AllCounts
+            meas = qml.counts(op=meas, all_outcomes=all_outcomes)
+        old_res = device.sample(call_args.args[1], **call_args.kwargs)
+        new_res = meas.process_samples(
+            samples=samples, wire_order=device.wires, shot_range=shot_range, bin_size=bin_size
+        )
+        if isinstance(old_res, dict):
+            old_res = [old_res]
+            new_res = [new_res]
+        for old, new in zip(old_res, new_res):
+            assert old.keys() == new.keys()
+            assert qml.math.allequal(list(old.values()), list(new.values()))
+
+
+class TestCountsIntegration:
+    # pylint:disable=too-many-public-methods,not-an-iterable
+
+    def test_counts_dimension(self, mocker):
+        """Test that the counts function outputs counts of the right size"""
+        n_sample = 10
+
+        dev = qml.device("default.qubit.legacy", wires=2, shots=n_sample)
+        spy = mocker.spy(qml.QubitDevice, "sample")
+
+        @qml.qnode(dev)
+        def circuit():
+            qml.RX(0.54, wires=0)
+            return qml.counts(qml.PauliZ(0)), qml.counts(qml.PauliX(1))
+
+        sample = circuit()
+
+        assert len(sample) == 2
+        assert np.all([sum(s.values()) == n_sample for s in sample])
+
+        custom_measurement_process(dev, spy)
+
+    def test_batched_counts_dimension(self, mocker):
+        """Test that the counts function outputs counts of the right size with batching"""
+        n_sample = 10
+
+        dev = qml.device("default.qubit.legacy", wires=2, shots=n_sample)
+        spy = mocker.spy(qml.QubitDevice, "sample")
+
+        @qml.qnode(dev)
+        def circuit():
+            qml.RX([0.54, 0.65], wires=0)
+            return qml.counts(qml.PauliZ(0)), qml.counts(qml.PauliX(1))
+
+        sample = circuit()
+
+        assert isinstance(sample, tuple)
+        assert len(sample) == 2
+        assert np.all([sum(s.values()) == n_sample for batch in sample for s in batch])
+
+        custom_measurement_process(dev, spy)
+
+    def test_counts_combination(self, mocker):
+        """Test the output of combining expval, var and counts"""
+        n_sample = 10
+
+        dev = qml.device("default.qubit.legacy", wires=3, shots=n_sample)
+        spy = mocker.spy(qml.QubitDevice, "sample")
+
+        @qml.qnode(dev, diff_method="parameter-shift")
+        def circuit():
+            qml.RX(0.54, wires=0)
+
+            return (
+                qml.counts(qml.PauliZ(0)),
+                qml.expval(qml.PauliX(1)),
+                qml.var(qml.PauliY(2)),
+            )
+
+        result = circuit()
+
+        assert len(result) == 3
+        assert sum(result[0].values()) == n_sample
+
+        custom_measurement_process(dev, spy)
+
+    def test_single_wire_counts(self, mocker):
+        """Test the return type and shape of sampling counts from a single wire"""
+        n_sample = 10
+
+        dev = qml.device("default.qubit.legacy", wires=1, shots=n_sample)
+        spy = mocker.spy(qml.QubitDevice, "sample")
+
+        @qml.qnode(dev)
+        def circuit():
+            qml.RX(0.54, wires=0)
+
+            return qml.counts(qml.PauliZ(0))
+
+        result = circuit()
+
+        assert isinstance(result, dict)
+        assert sum(result.values()) == n_sample
+
+        custom_measurement_process(dev, spy)
+
+    def test_multi_wire_counts_regular_shape(self, mocker):
+        """Test the return type and shape of sampling multiple wires
+        where a rectangular array is expected"""
+        n_sample = 10
+
+        dev = qml.device("default.qubit.legacy", wires=3, shots=n_sample)
+        spy = mocker.spy(qml.QubitDevice, "sample")
+
+        @qml.qnode(dev)
+        def circuit():
+            return (
+                qml.counts(qml.PauliZ(0)),
+                qml.counts(qml.PauliZ(1)),
+                qml.counts(qml.PauliZ(2)),
+            )
+
+        result = circuit()
+
+        # If all the dimensions are equal the result will end up to be a proper rectangular array
+        assert isinstance(result, tuple)
+        assert len(result) == 3
+        assert all(sum(r.values()) == n_sample for r in result)
+        assert all(all(v.dtype == np.dtype("int") for v in r.values()) for r in result)
+
+        custom_measurement_process(dev, spy)
+
+    def test_observable_return_type_is_counts(self):
+        """Test that the return type of the observable is :attr:`ObservableReturnTypes.Counts`"""
+        n_shots = 10
+        dev = qml.device("default.qubit.legacy", wires=1, shots=n_shots)
+
+        @qml.qnode(dev)
+        def circuit():
+            res = qml.counts(qml.PauliZ(0))
+            return res
+
+        circuit()
+        assert circuit._qfunc_output.return_type is Counts  # pylint: disable=protected-access
+
+    def test_providing_no_observable_and_no_wires_counts(self, mocker):
+        """Test that we can provide no observable and no wires to sample function"""
+        dev = qml.device("default.qubit.legacy", wires=2, shots=1000)
+        spy = mocker.spy(qml.QubitDevice, "sample")
+
+        @qml.qnode(dev)
+        def circuit():
+            qml.Hadamard(wires=0)
+            res = qml.counts()
+            assert res.obs is None
+            assert res.wires == qml.wires.Wires([])
+            return res
+
+        circuit()
+
+        custom_measurement_process(dev, spy)
+
+    def test_providing_no_observable_and_wires_counts(self, mocker):
+        """Test that we can provide no observable but specify wires to the sample function"""
+        wires = [0, 2]
+        wires_obj = qml.wires.Wires(wires)
+        dev = qml.device("default.qubit.legacy", wires=3, shots=1000)
+        spy = mocker.spy(qml.QubitDevice, "sample")
+
+        @qml.qnode(dev)
+        def circuit():
+            qml.Hadamard(wires=0)
+            res = qml.counts(wires=wires)
+
+            assert res.obs is None
+            assert res.wires == wires_obj
+            return res
+
+        circuit()
+
+        custom_measurement_process(dev, spy)
+
+    def test_batched_counts_work_individually(self, mocker):
+        """Test that each counts call operates independently"""
+        n_shots = 10
+        dev = qml.device("default.qubit.legacy", wires=1, shots=n_shots)
+        spy = mocker.spy(qml.QubitDevice, "sample")
+
+        @qml.qnode(dev)
+        def circuit():
+            qml.pow(qml.PauliX(0), z=[1, 2])
+            return qml.counts()
+
+        assert circuit() == [{"1": 10}, {"0": 10}]
+        custom_measurement_process(dev, spy)
+
+    @pytest.mark.all_interfaces
+    @pytest.mark.parametrize("wires, basis_state", [(None, "010"), ([2, 1], "01")])
+    @pytest.mark.parametrize("interface", ["autograd", "jax", "tensorflow", "torch"])
+    def test_counts_no_op_finite_shots(self, interface, wires, basis_state, mocker):
+        """Check all interfaces with computational basis state counts and
+        finite shot"""
+        n_shots = 10
+        dev = qml.device("default.qubit.legacy", wires=3, shots=n_shots)
+        spy = mocker.spy(qml.QubitDevice, "sample")
+
+        @qml.qnode(dev, interface=interface)
+        def circuit():
+            qml.PauliX(1)
+            return qml.counts(wires=wires)
+
+        res = circuit()
+        assert res == {basis_state: n_shots}
+        assert qml.math.get_interface(res[basis_state]) == interface
+
+        custom_measurement_process(dev, spy)
+
+    @pytest.mark.all_interfaces
+    @pytest.mark.parametrize("interface", ["autograd", "jax", "tensorflow", "torch"])
+    def test_counts_operator_finite_shots(self, interface, mocker):
+        """Check all interfaces with observable measurement counts and finite
+        shot"""
+        n_shots = 10
+        dev = qml.device("default.qubit.legacy", wires=3, shots=n_shots)
+        spy = mocker.spy(qml.QubitDevice, "sample")
+
+        @qml.qnode(dev, interface=interface)
+        def circuit():
+            return qml.counts(qml.PauliZ(0))
+
+        res = circuit()
+        assert res == {1: n_shots}
+
+        custom_measurement_process(dev, spy)
+
+    @pytest.mark.all_interfaces
+    @pytest.mark.parametrize("shot_vec", [(1, 10, 10), (1, 10, 1000)])
+    @pytest.mark.parametrize("wires, basis_state", [(None, "010"), ([2, 1], "01")])
+    @pytest.mark.parametrize("interface", ["autograd", "jax", "tensorflow", "torch"])
+    def test_counts_binned(
+        self, shot_vec, interface, wires, basis_state, mocker
+    ):  # pylint:disable=too-many-arguments
+        """Check all interfaces with computational basis state counts and
+        different shot vectors"""
+        dev = qml.device("default.qubit.legacy", wires=3, shots=shot_vec)
+        spy = mocker.spy(qml.QubitDevice, "sample")
+
+        @qml.qnode(dev, interface=interface)
+        def circuit():
+            qml.PauliX(1)
+            return qml.counts(wires=wires)
+
+        res = circuit()
+
+        assert isinstance(res, tuple)
+        assert res[0] == {basis_state: shot_vec[0]}
+        assert res[1] == {basis_state: shot_vec[1]}
+        assert res[2] == {basis_state: shot_vec[2]}
+        assert len(res) == len(shot_vec)
+        assert sum(sum(res_bin.values()) for res_bin in res) == sum(shot_vec)
+
+        custom_measurement_process(dev, spy)
+
+    @pytest.mark.all_interfaces
+    @pytest.mark.parametrize("shot_vec", [(1, 10, 10), (1, 10, 1000)])
+    @pytest.mark.parametrize("interface", ["autograd", "jax", "tensorflow", "torch"])
+    def test_counts_operator_binned(self, shot_vec, interface, mocker):
+        """Check all interfaces with observable measurement counts and different
+        shot vectors"""
+        dev = qml.device("default.qubit.legacy", wires=3, shots=shot_vec)
+        spy = mocker.spy(qml.QubitDevice, "sample")
+
+        @qml.qnode(dev, interface=interface)
+        def circuit():
+            return qml.counts(qml.PauliZ(0))
+
+        res = circuit()
+
+        assert isinstance(res, tuple)
+        assert res[0] == {1: shot_vec[0]}
+        assert res[1] == {1: shot_vec[1]}
+        assert res[2] == {1: shot_vec[2]}
+        assert len(res) == len(shot_vec)
+        assert sum(sum(res_bin.values()) for res_bin in res) == sum(shot_vec)
+
+        custom_measurement_process(dev, spy)
+
+    @pytest.mark.parametrize("shot_vec", [(1, 10, 10), (1, 10, 1000)])
+    def test_counts_binned_4_wires(self, shot_vec, mocker):
+        """Check the autograd interface with computational basis state counts and
+        different shot vectors on a device with 4 wires"""
+        dev = qml.device("default.qubit.legacy", wires=4, shots=shot_vec)
+        spy = mocker.spy(qml.QubitDevice, "sample")
+
+        @qml.qnode(dev, interface="autograd")
+        def circuit():
+            qml.PauliX(1)
+            qml.PauliX(2)
+            qml.PauliX(3)
+            return qml.counts()
+
+        res = circuit()
+        basis_state = "0111"
+
+        assert isinstance(res, tuple)
+        assert res[0][basis_state] == shot_vec[0]
+        assert res[1][basis_state] == shot_vec[1]
+        assert res[2][basis_state] == shot_vec[2]
+        assert len(res) == len(shot_vec)
+        assert sum(sum(res_bin.values()) for res_bin in res) == sum(shot_vec)
+
+        custom_measurement_process(dev, spy)
+
+    @pytest.mark.parametrize("shot_vec", [(1, 10, 10), (1, 10, 1000)])
+    def test_counts_operator_binned_4_wires(self, shot_vec, mocker):
+        """Check the autograd interface with observable samples to obtain
+        counts from and different shot vectors on a device with 4 wires"""
+        dev = qml.device("default.qubit.legacy", wires=4, shots=shot_vec)
+        spy = mocker.spy(qml.QubitDevice, "sample")
+
+        @qml.qnode(dev, interface="autograd")
+        def circuit():
+            qml.PauliX(1)
+            qml.PauliX(2)
+            qml.PauliX(3)
+            return qml.counts(qml.PauliZ(0))
+
+        res = circuit()
+        sample = 1
+
+        assert isinstance(res, tuple)
+        assert res[0][sample] == shot_vec[0]
+        assert res[1][sample] == shot_vec[1]
+        assert res[2][sample] == shot_vec[2]
+        assert len(res) == len(shot_vec)
+        assert sum(sum(res_bin.values()) for res_bin in res) == sum(shot_vec)
+
+        custom_measurement_process(dev, spy)
+
+    meas2 = [
+        qml.expval(qml.PauliZ(0)),
+        qml.var(qml.PauliZ(0)),
+        qml.probs(wires=[1, 0]),
+        qml.sample(wires=1),
+    ]
+
+    @pytest.mark.all_interfaces
+    @pytest.mark.parametrize("interface", ["autograd", "jax", "tensorflow", "torch"])
+    @pytest.mark.parametrize("meas2", meas2)
+    @pytest.mark.parametrize("shots", [1000, (1, 10)])
+    def test_counts_observable_finite_shots(self, interface, meas2, shots, mocker):
+        """Check all interfaces with observable measurement counts and finite
+        shot"""
+        dev = qml.device("default.qubit.legacy", wires=3, shots=shots)
+        spy = mocker.spy(qml.QubitDevice, "sample")
+
+        if isinstance(shots, tuple) and interface == "torch":
+            pytest.skip("Torch needs to be updated for shot vectors.")
+
+        @qml.qnode(dev, interface=interface)
+        def circuit():
+            qml.PauliX(0)
+            return qml.counts(wires=0), qml.apply(meas2)
+
+        res = circuit()
+        assert isinstance(res, tuple)
+
+        num_shot_bins = 1 if isinstance(shots, int) else len(shots)
+
+        if num_shot_bins == 1:
+            counts_term_indices = [i * 2 for i in range(num_shot_bins)]
+            for ind in counts_term_indices:
+                assert isinstance(res[ind], dict)
+        else:
+            assert len(res) == 2
+
+            assert isinstance(res[0], tuple)
+            assert isinstance(res[0][0], dict)
+            assert isinstance(res[1], tuple)
+            assert isinstance(res[1][0], dict)
+
+        custom_measurement_process(dev, spy)
+
+    def test_all_outcomes_kwarg_providing_observable(self, mocker):
+        """Test that the dictionary keys *all* eigenvalues of the observable,
+        including 0 count values, if observable is given and all_outcomes=True"""
+
+        n_shots = 10
+        dev = qml.device("default.qubit.legacy", wires=1, shots=n_shots)
+        spy = mocker.spy(qml.QubitDevice, "sample")
+
+        @qml.qnode(dev)
+        def circuit():
+            res = qml.counts(qml.PauliZ(0), all_outcomes=True)
+            return res
+
+        res = circuit()
+
+        assert res == {1: n_shots, -1: 0}
+
+        custom_measurement_process(dev, spy)
+
+    def test_all_outcomes_kwarg_no_observable_no_wires(self, mocker):
+        """Test that the dictionary keys are *all* the possible combinations
+        of basis states for the device, including 0 count values, if no wire
+        count and no observable are given and all_outcomes=True"""
+
+        n_shots = 10
+        dev = qml.device("default.qubit.legacy", wires=2, shots=n_shots)
+        spy = mocker.spy(qml.QubitDevice, "sample")
+
+        @qml.qnode(dev)
+        def circuit():
+            return qml.counts(all_outcomes=True)
+
+        res = circuit()
+
+        assert res == {"00": n_shots, "01": 0, "10": 0, "11": 0}
+
+        custom_measurement_process(dev, spy)
+
+    def test_all_outcomes_kwarg_providing_wires_and_no_observable(self, mocker):
+        """Test that the dictionary keys are *all* possible combinations
+        of basis states for the specified wires, including 0 count values,
+        if wire count is given and all_outcomes=True"""
+
+        n_shots = 10
+        dev = qml.device("default.qubit.legacy", wires=4, shots=n_shots)
+        spy = mocker.spy(qml.QubitDevice, "sample")
+
+        @qml.qnode(dev)
+        def circuit():
+            return qml.counts(wires=[0, 2], all_outcomes=True)
+
+        res = circuit()
+
+        assert res == {"00": n_shots, "01": 0, "10": 0, "11": 0}
+
+        custom_measurement_process(dev, spy)
+
+    def test_all_outcomes_hermitian(self, mocker):
+        """Tests that the all_outcomes=True option for counts works with the
+        qml.Hermitian observable"""
+
+        n_shots = 10
+        dev = qml.device("default.qubit.legacy", wires=2, shots=n_shots)
+        spy = mocker.spy(qml.QubitDevice, "sample")
+
+        A = np.array([[1, 0], [0, -1]])
+
+        @qml.qnode(dev)
+        def circuit(x):
+            return qml.counts(qml.Hermitian(x, wires=0), all_outcomes=True)
+
+        res = circuit(A)
+
+        assert res == {-1.0: 0, 1.0: n_shots}
+
+        custom_measurement_process(dev, spy)
+
+    def test_all_outcomes_multiple_measurements(self, mocker):
+        """Tests that the all_outcomes=True option for counts works when
+        multiple measurements are performed"""
+
+        dev = qml.device("default.qubit.legacy", wires=2, shots=10)
+        spy = mocker.spy(qml.QubitDevice, "sample")
+
+        @qml.qnode(dev)
+        def circuit():
+            return qml.sample(qml.PauliZ(0)), qml.counts(), qml.counts(all_outcomes=True)
+
+        res = circuit()
+
+        assert len(res[0]) == 10
+        assert res[1] == {"00": 10}
+        assert res[2] == {"00": 10, "01": 0, "10": 0, "11": 0}
+        custom_measurement_process(dev, spy)
+
+    def test_batched_all_outcomes(self, mocker):
+        """Tests that all_outcomes=True works with broadcasting."""
+        n_shots = 10
+        dev = qml.device("default.qubit.legacy", wires=1, shots=n_shots)
+        spy = mocker.spy(qml.QubitDevice, "sample")
+
+        @qml.qnode(dev)
+        def circuit():
+            qml.pow(qml.PauliX(0), z=[1, 2])
+            return qml.counts(qml.PauliZ(0), all_outcomes=True)
+
+        assert circuit() == [{1: 0, -1: n_shots}, {1: n_shots, -1: 0}]
+
+        custom_measurement_process(dev, spy)
+
+    def test_counts_empty_wires(self):
+        """Test that using ``qml.counts`` with an empty wire list raises an error."""
+        with pytest.raises(ValueError, match="Cannot set an empty list of wires."):
+            qml.counts(wires=[])
+
+    @pytest.mark.parametrize("shots", [1, 100])
+    def test_counts_no_arguments(self, shots):
+        """Test that using ``qml.counts`` with no arguments returns the counts of all wires."""
+        dev = qml.device("default.qubit.legacy", wires=3, shots=shots)
+
+        @qml.qnode(dev)
+        def circuit():
+            return qml.counts()
+
+        res = circuit()
+
+        assert qml.math.allequal(res, {"000": shots})
+
+
+@pytest.mark.all_interfaces
+@pytest.mark.parametrize("wires, basis_state", [(None, "010"), ([2, 1], "01")])
+@pytest.mark.parametrize("interface", ["autograd", "jax", "tensorflow", "torch"])
+def test_counts_no_op_finite_shots(interface, wires, basis_state, mocker):
+    """Check all interfaces with computational basis state counts and finite shot"""
+    n_shots = 10
+    dev = qml.device("default.qubit.legacy", wires=3, shots=n_shots)
+    spy = mocker.spy(qml.QubitDevice, "sample")
+
+    @qml.qnode(dev, interface=interface)
+    def circuit():
+        qml.PauliX(1)
+        return qml.counts(wires=wires)
+
+    res = circuit()
+    assert res == {basis_state: n_shots}
+    assert qml.math.get_interface(res[basis_state]) == interface
+
+    custom_measurement_process(dev, spy)
+
+
+@pytest.mark.all_interfaces
+@pytest.mark.parametrize("wires, basis_states", [(None, ("010", "000")), ([2, 1], ("01", "00"))])
+@pytest.mark.parametrize("interface", ["autograd", "jax", "tensorflow", "torch"])
+def test_batched_counts_no_op_finite_shots(interface, wires, basis_states, mocker):
+    """Check all interfaces with computational basis state counts and
+    finite shot"""
+    n_shots = 10
+    dev = qml.device("default.qubit.legacy", wires=3, shots=n_shots)
+    spy = mocker.spy(qml.QubitDevice, "sample")
+
+    @qml.qnode(dev, interface=interface)
+    def circuit():
+        qml.pow(qml.PauliX(1), z=[1, 2])
+        return qml.counts(wires=wires)
+
+    assert circuit() == [{basis_state: n_shots} for basis_state in basis_states]
+
+    custom_measurement_process(dev, spy)
+
+
+@pytest.mark.all_interfaces
+@pytest.mark.parametrize("wires, basis_states", [(None, ("010", "000")), ([2, 1], ("01", "00"))])
+@pytest.mark.parametrize("interface", ["autograd", "jax", "tensorflow", "torch"])
+def test_batched_counts_and_expval_no_op_finite_shots(interface, wires, basis_states, mocker):
+    """Check all interfaces with computational basis state counts and
+    finite shot"""
+    n_shots = 10
+    dev = qml.device("default.qubit.legacy", wires=3, shots=n_shots)
+    spy = mocker.spy(qml.QubitDevice, "sample")
+
+    @qml.qnode(dev, interface=interface)
+    def circuit():
+        qml.pow(qml.PauliX(1), z=[1, 2])
+        return qml.counts(wires=wires), qml.expval(qml.PauliZ(0))
+
+    res = circuit()
+    assert isinstance(res, tuple) and len(res) == 2
+    assert res[0] == [{basis_state: n_shots} for basis_state in basis_states]
+    assert len(res[1]) == 2 and qml.math.allequal(res[1], 1)
+
+    custom_measurement_process(dev, spy)
+
+
+@pytest.mark.all_interfaces
+@pytest.mark.parametrize("interface", ["autograd", "jax", "tensorflow", "torch"])
+def test_batched_counts_operator_finite_shots(interface, mocker):
+    """Check all interfaces with observable measurement counts, batching and finite shots"""
+    n_shots = 10
+    dev = qml.device("default.qubit.legacy", wires=3, shots=n_shots)
+    spy = mocker.spy(qml.QubitDevice, "sample")
+
+    @qml.qnode(dev, interface=interface)
+    def circuit():
+        qml.pow(qml.PauliX(0), z=[1, 2])
+        return qml.counts(qml.PauliZ(0))
+
+    assert circuit() == [{-1: n_shots}, {1: n_shots}]
+
+    custom_measurement_process(dev, spy)
+
+
+@pytest.mark.all_interfaces
+@pytest.mark.parametrize("interface", ["autograd", "jax", "tensorflow", "torch"])
+def test_batched_counts_and_expval_operator_finite_shots(interface, mocker):
+    """Check all interfaces with observable measurement counts, batching and finite shots"""
+    n_shots = 10
+    dev = qml.device("default.qubit.legacy", wires=3, shots=n_shots)
+    spy = mocker.spy(qml.QubitDevice, "sample")
+
+    @qml.qnode(dev, interface=interface)
+    def circuit():
+        qml.pow(qml.PauliX(0), z=[1, 2])
+        return qml.counts(qml.PauliZ(0)), qml.expval(qml.PauliZ(0))
+
+    res = circuit()
+    assert isinstance(res, tuple) and len(res) == 2
+    assert res[0] == [{-1: n_shots}, {1: n_shots}]
+    assert len(res[1]) == 2 and qml.math.allequal(res[1], [-1, 1])
+
+    custom_measurement_process(dev, spy)
diff --git a/tests/measurements/legacy/test_expval_legacy.py b/tests/measurements/legacy/test_expval_legacy.py
new file mode 100644
index 00000000000..436b6b2e121
--- /dev/null
+++ b/tests/measurements/legacy/test_expval_legacy.py
@@ -0,0 +1,210 @@
+# Copyright 2018-2020 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 expval module"""
+import numpy as np
+import pytest
+
+import pennylane as qml
+from pennylane.measurements import Expectation, Shots
+
+
+# TODO: Remove this when new CustomMP are the default
+def custom_measurement_process(device, spy):
+    assert len(spy.call_args_list) > 0  # make sure method is mocked properly
+
+    samples = device._samples  # pylint: disable=protected-access
+    state = device._state  # pylint: disable=protected-access
+    call_args_list = list(spy.call_args_list)
+    for call_args in call_args_list:
+        obs = call_args.args[1]
+        shot_range, bin_size = (
+            call_args.kwargs["shot_range"],
+            call_args.kwargs["bin_size"],
+        )
+        # no need to use op, because the observable has already been applied to ``self.dev._state``
+        meas = qml.expval(op=obs)
+        old_res = device.expval(obs, shot_range=shot_range, bin_size=bin_size)
+        if device.shots is None:
+            new_res = meas.process_state(state=state, wire_order=device.wires)
+        else:
+            new_res = meas.process_samples(
+                samples=samples, wire_order=device.wires, shot_range=shot_range, bin_size=bin_size
+            )
+        assert qml.math.allclose(old_res, new_res)
+
+
+class TestExpval:
+    """Tests for the expval function"""
+
+    @pytest.mark.parametrize("shots", [None, 10000, [10000, 10000]])
+    @pytest.mark.parametrize("r_dtype", [np.float32, np.float64])
+    def test_value(self, tol, r_dtype, mocker, shots):
+        """Test that the expval interface works"""
+        dev = qml.device("default.qubit.legacy", wires=2, shots=shots)
+        dev.R_DTYPE = r_dtype
+
+        @qml.qnode(dev, diff_method="parameter-shift")
+        def circuit(x):
+            qml.RX(x, wires=0)
+            return qml.expval(qml.PauliY(0))
+
+        new_dev = circuit.device
+        spy = mocker.spy(qml.QubitDevice, "expval")
+
+        x = 0.54
+        res = circuit(x)
+        expected = -np.sin(x)
+
+        atol = tol if shots is None else 0.05
+        rtol = 0 if shots is None else 0.05
+        assert np.allclose(res, expected, atol=atol, rtol=rtol)
+
+        # pylint: disable=no-member, unsubscriptable-object
+        if isinstance(res, tuple):
+            assert res[0].dtype == r_dtype
+            assert res[1].dtype == r_dtype
+        else:
+            assert res.dtype == r_dtype
+
+        custom_measurement_process(new_dev, spy)
+
+    def test_not_an_observable(self, mocker):
+        """Test that a warning is raised if the provided
+        argument might not be hermitian."""
+        dev = qml.device("default.qubit.legacy", wires=2)
+
+        @qml.qnode(dev)
+        def circuit():
+            qml.RX(0.52, wires=0)
+            return qml.expval(qml.prod(qml.PauliX(0), qml.PauliZ(0)))
+
+        new_dev = circuit.device
+        spy = mocker.spy(qml.QubitDevice, "expval")
+
+        with pytest.warns(UserWarning, match="Prod might not be hermitian."):
+            _ = circuit()
+
+        custom_measurement_process(new_dev, spy)
+
+    def test_observable_return_type_is_expectation(self, mocker):
+        """Test that the return type of the observable is :attr:`ObservableReturnTypes.Expectation`"""
+        dev = qml.device("default.qubit.legacy", wires=2)
+
+        @qml.qnode(dev)
+        def circuit():
+            res = qml.expval(qml.PauliZ(0))
+            assert res.return_type is Expectation
+            return res
+
+        new_dev = circuit.device
+        spy = mocker.spy(qml.QubitDevice, "expval")
+
+        circuit()
+
+        custom_measurement_process(new_dev, spy)
+
+    @pytest.mark.parametrize("shots", [None, 10000, [10000, 10000]])
+    @pytest.mark.parametrize("phi", np.arange(0, 2 * np.pi, np.pi / 3))
+    def test_observable_is_measurement_value(
+        self, shots, phi, mocker, tol, tol_stochastic
+    ):  # pylint: disable=too-many-arguments
+        """Test that expectation values for mid-circuit measurement values
+        are correct for a single measurement value."""
+        dev = qml.device("default.qubit", wires=2, shots=shots)
+
+        @qml.qnode(dev)
+        def circuit(phi):
+            qml.RX(phi, 0)
+            m0 = qml.measure(0)
+            return qml.expval(m0)
+
+        new_dev = circuit.device
+        spy = mocker.spy(qml.QubitDevice, "expval")
+
+        res = circuit(phi)
+
+        atol = tol if shots is None else tol_stochastic
+        assert np.allclose(np.array(res), np.sin(phi / 2) ** 2, atol=atol, rtol=0)
+        custom_measurement_process(new_dev, spy)
+
+    @pytest.mark.parametrize(
+        "obs",
+        [qml.PauliZ(0), qml.Hermitian(np.diag([1, 2]), 0), qml.Hermitian(np.diag([1.0, 2.0]), 0)],
+    )
+    def test_shape(self, obs):
+        """Test that the shape is correct."""
+        dev = qml.device("default.qubit.legacy", wires=1)
+
+        res = qml.expval(obs)
+        # pylint: disable=use-implicit-booleaness-not-comparison
+        assert res.shape(dev, Shots(None)) == ()
+        assert res.shape(dev, Shots(100)) == ()
+
+    @pytest.mark.parametrize(
+        "obs",
+        [qml.PauliZ(0), qml.Hermitian(np.diag([1, 2]), 0), qml.Hermitian(np.diag([1.0, 2.0]), 0)],
+    )
+    def test_shape_shot_vector(self, obs):
+        """Test that the shape is correct with the shot vector too."""
+        res = qml.expval(obs)
+        shot_vector = (1, 2, 3)
+        dev = qml.device("default.qubit.legacy", wires=3, shots=shot_vector)
+        assert res.shape(dev, Shots(shot_vector)) == ((), (), ())
+
+    @pytest.mark.parametrize("state", [np.array([0, 0, 0]), np.array([1, 0, 0, 0, 0, 0, 0, 0])])
+    @pytest.mark.parametrize("shots", [None, 1000, [1000, 10000]])
+    def test_projector_expval(self, state, shots, mocker):
+        """Tests that the expectation of a ``Projector`` object is computed correctly for both of
+        its subclasses."""
+        dev = qml.device("default.qubit.legacy", wires=3, shots=shots)
+        np.random.seed(42)
+
+        @qml.qnode(dev)
+        def circuit():
+            qml.Hadamard(0)
+            return qml.expval(qml.Projector(state, wires=range(3)))
+
+        new_dev = circuit.device
+        spy = mocker.spy(qml.QubitDevice, "expval")
+
+        res = circuit()
+        expected = [0.5, 0.5] if isinstance(shots, list) else 0.5
+        assert np.allclose(res, expected, atol=0.02, rtol=0.02)
+
+        custom_measurement_process(new_dev, spy)
+
+    def test_permuted_wires(self, mocker):
+        """Test that the expectation value of an operator with permuted wires is the same."""
+        obs = qml.prod(qml.PauliZ(8), qml.s_prod(2, qml.PauliZ(10)), qml.s_prod(3, qml.PauliZ("h")))
+        obs_2 = qml.prod(
+            qml.s_prod(3, qml.PauliZ("h")), qml.PauliZ(8), qml.s_prod(2, qml.PauliZ(10))
+        )
+
+        dev = qml.device("default.qubit.legacy", wires=["h", 8, 10])
+        spy = mocker.spy(qml.QubitDevice, "expval")
+
+        @qml.qnode(dev)
+        def circuit():
+            qml.RX(1.23, wires=["h"])
+            qml.RY(2.34, wires=[8])
+            return qml.expval(obs)
+
+        @qml.qnode(dev)
+        def circuit2():
+            qml.RX(1.23, wires=["h"])
+            qml.RY(2.34, wires=[8])
+            return qml.expval(obs_2)
+
+        assert circuit() == circuit2()
+        custom_measurement_process(dev, spy)
diff --git a/tests/measurements/legacy/test_measurements_legacy.py b/tests/measurements/legacy/test_measurements_legacy.py
new file mode 100644
index 00000000000..749dc1b5b15
--- /dev/null
+++ b/tests/measurements/legacy/test_measurements_legacy.py
@@ -0,0 +1,220 @@
+# Copyright 2018-2020 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 measurements module"""
+import pytest
+
+import pennylane as qml
+from pennylane.measurements import (
+    ClassicalShadowMP,
+    MeasurementProcess,
+    MeasurementTransform,
+    SampleMeasurement,
+    SampleMP,
+    Shots,
+    StateMeasurement,
+    StateMP,
+    expval,
+    var,
+)
+
+# pylint: disable=too-few-public-methods, unused-argument
+
+
+class NotValidMeasurement(MeasurementProcess):
+    @property
+    def return_type(self):
+        return "NotValidReturnType"
+
+
+def test_no_measure():
+    """Test that failing to specify a measurement
+    raises an exception"""
+    dev = qml.device("default.qubit.legacy", wires=2)
+
+    @qml.qnode(dev)
+    def circuit(x):
+        qml.RX(x, wires=0)
+        return qml.PauliY(0)
+
+    with pytest.raises(qml.QuantumFunctionError, match="must return either a single measurement"):
+        _ = circuit(0.65)
+
+
+def test_shape_unrecognized_error():
+    """Test that querying the shape of a measurement process with an
+    unrecognized return type raises an error."""
+    dev = qml.device("default.qubit.legacy", wires=2)
+    mp = NotValidMeasurement()
+    with pytest.raises(
+        qml.QuantumFunctionError,
+        match="The shape of the measurement NotValidMeasurement is not defined",
+    ):
+        mp.shape(dev, Shots(None))
+
+
+@pytest.mark.parametrize("stat_func", [expval, var])
+def test_not_an_observable(stat_func):
+    """Test that a UserWarning is raised if the provided
+    argument might not be hermitian."""
+
+    dev = qml.device("default.qubit.legacy", wires=2)
+
+    @qml.qnode(dev)
+    def circuit():
+        qml.RX(0.52, wires=0)
+        return stat_func(qml.prod(qml.PauliX(0), qml.PauliZ(0)))
+
+    with pytest.warns(UserWarning, match="Prod might not be hermitian."):
+        _ = circuit()
+
+
+class TestSampleMeasurement:
+    """Tests for the SampleMeasurement class."""
+
+    def test_custom_sample_measurement(self):
+        """Test the execution of a custom sampled measurement."""
+
+        class MyMeasurement(SampleMeasurement):
+            # pylint: disable=signature-differs
+            def process_samples(self, samples, wire_order, shot_range, bin_size):
+                return qml.math.sum(samples[..., self.wires])
+
+        dev = qml.device("default.qubit.legacy", wires=2, shots=1000)
+
+        @qml.qnode(dev)
+        def circuit():
+            qml.PauliX(0)
+            return MyMeasurement(wires=[0]), MyMeasurement(wires=[1])
+
+        assert qml.math.allequal(circuit(), [1000, 0])
+
+    def test_sample_measurement_without_shots(self):
+        """Test that executing a sampled measurement with ``shots=None`` raises an error."""
+
+        class MyMeasurement(SampleMeasurement):
+            # pylint: disable=signature-differs
+            def process_samples(self, samples, wire_order, shot_range, bin_size):
+                return qml.math.sum(samples[..., self.wires])
+
+        dev = qml.device("default.qubit.legacy", wires=2)
+
+        @qml.qnode(dev)
+        def circuit():
+            qml.PauliX(0)
+            return MyMeasurement(wires=[0]), MyMeasurement(wires=[1])
+
+        with pytest.raises(
+            ValueError, match="Shots must be specified in the device to compute the measurement "
+        ):
+            circuit()
+
+    def test_method_overriden_by_device(self):
+        """Test that the device can override a measurement process."""
+
+        dev = qml.device("default.qubit.legacy", wires=2, shots=1000)
+
+        @qml.qnode(dev)
+        def circuit():
+            qml.PauliX(0)
+            return qml.sample(wires=[0]), qml.sample(wires=[1])
+
+        circuit.device.measurement_map[SampleMP] = "test_method"
+        circuit.device.test_method = lambda obs, shot_range=None, bin_size=None: 2
+
+        assert qml.math.allequal(circuit(), [2, 2])
+
+
+class TestStateMeasurement:
+    """Tests for the SampleMeasurement class."""
+
+    def test_custom_state_measurement(self):
+        """Test the execution of a custom state measurement."""
+
+        class MyMeasurement(StateMeasurement):
+            def process_state(self, state, wire_order):
+                return qml.math.sum(state)
+
+        dev = qml.device("default.qubit.legacy", wires=2)
+
+        @qml.qnode(dev)
+        def circuit():
+            return MyMeasurement()
+
+        assert circuit() == 1
+
+    def test_sample_measurement_with_shots(self):
+        """Test that executing a state measurement with shots raises a warning."""
+
+        class MyMeasurement(StateMeasurement):
+            def process_state(self, state, wire_order):
+                return qml.math.sum(state)
+
+        dev = qml.device("default.qubit.legacy", wires=2, shots=1000)
+
+        @qml.qnode(dev)
+        def circuit():
+            return MyMeasurement()
+
+        with pytest.warns(
+            UserWarning,
+            match="Requested measurement MyMeasurement with finite shots",
+        ):
+            circuit()
+
+    def test_method_overriden_by_device(self):
+        """Test that the device can override a measurement process."""
+
+        dev = qml.device("default.qubit.legacy", wires=2)
+
+        @qml.qnode(dev, interface="autograd")
+        def circuit():
+            return qml.state()
+
+        circuit.device.measurement_map[StateMP] = "test_method"
+        circuit.device.test_method = lambda obs, shot_range=None, bin_size=None: 2
+
+        assert circuit() == 2
+
+
+class TestMeasurementTransform:
+    """Tests for the MeasurementTransform class."""
+
+    def test_custom_measurement(self):
+        """Test the execution of a custom measurement."""
+
+        class MyMeasurement(MeasurementTransform):
+            def process(self, tape, device):
+                return {device.shots: len(tape)}
+
+        dev = qml.device("default.qubit.legacy", wires=2, shots=1000)
+
+        @qml.qnode(dev)
+        def circuit():
+            return MyMeasurement()
+
+        assert circuit() == {dev.shots: len(circuit.tape)}
+
+    def test_method_overriden_by_device(self):
+        """Test that the device can override a measurement process."""
+
+        dev = qml.device("default.qubit.legacy", wires=2, shots=1000)
+
+        @qml.qnode(dev)
+        def circuit():
+            return qml.classical_shadow(wires=0)
+
+        circuit.device.measurement_map[ClassicalShadowMP] = "test_method"
+        circuit.device.test_method = lambda tape: 2
+
+        assert circuit() == 2
diff --git a/tests/measurements/legacy/test_mutual_info_legacy.py b/tests/measurements/legacy/test_mutual_info_legacy.py
new file mode 100644
index 00000000000..cdab46590f6
--- /dev/null
+++ b/tests/measurements/legacy/test_mutual_info_legacy.py
@@ -0,0 +1,433 @@
+# Copyright 2018-2020 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 mutual_info module"""
+import numpy as np
+import pytest
+
+import pennylane as qml
+from pennylane.interfaces import INTERFACE_MAP
+from pennylane.measurements import Shots
+
+
+@pytest.mark.parametrize("shots, shape", [(None, ()), (10, ()), ([1, 10], ((), ()))])
+def test_shape(shots, shape):
+    """Test that the shape is correct."""
+    dev = qml.device("default.qubit.legacy", wires=3, shots=shots)
+    res = qml.mutual_info(wires0=[0], wires1=[1])
+    assert res.shape(dev, Shots(shots)) == shape
+
+
+class TestIntegration:
+    """Tests for the mutual information functions"""
+
+    @pytest.mark.all_interfaces
+    @pytest.mark.parametrize("interface", ["autograd", "jax", "tf", "torch"])
+    @pytest.mark.parametrize(
+        "state, expected",
+        [
+            ([1.0, 0.0, 0.0, 0.0], 0),
+            ([qml.math.sqrt(2) / 2, 0.0, qml.math.sqrt(2) / 2, 0.0], 0),
+            ([qml.math.sqrt(2) / 2, 0.0, 0.0, qml.math.sqrt(2) / 2], 2 * qml.math.log(2)),
+            (qml.math.ones(4) * 0.5, 0.0),
+        ],
+    )
+    def test_mutual_info_output(self, interface, state, expected):
+        """Test the output of qml.mutual_info"""
+        dev = qml.device("default.qubit.legacy", wires=4)
+
+        @qml.qnode(dev, interface=interface)
+        def circuit():
+            qml.StatePrep(state, wires=[0, 1])
+            return qml.mutual_info(wires0=[0, 2], wires1=[1, 3])
+
+        res = circuit()
+        new_res = qml.mutual_info(wires0=[0, 2], wires1=[1, 3]).process_state(
+            state=circuit.device.state, wire_order=circuit.device.wires
+        )
+        assert np.allclose(res, expected, atol=1e-6)
+        assert np.allclose(new_res, expected, atol=1e-6)
+        assert INTERFACE_MAP.get(qml.math.get_interface(new_res)) == interface
+        assert res.dtype == new_res.dtype  # pylint: disable=no-member
+
+    def test_shot_vec_error(self):
+        """Test an error is raised when using shot vectors with mutual_info."""
+        dev = qml.device("default.qubit.legacy", wires=2, shots=[1, 10, 10, 1000])
+
+        @qml.qnode(device=dev)
+        def circuit(x):
+            qml.Hadamard(wires=[0])
+            qml.CRX(x, wires=[0, 1])
+            return qml.mutual_info(wires0=[0], wires1=[1])
+
+        with pytest.raises(
+            NotImplementedError, match="mutual information is not supported with shot vectors"
+        ):
+            circuit(0.5)
+
+    diff_methods = ["backprop", "finite-diff"]
+
+    @pytest.mark.all_interfaces
+    @pytest.mark.parametrize("device", ["default.qubit.legacy", "default.mixed", "lightning.qubit"])
+    @pytest.mark.parametrize("interface", ["autograd", "jax", "tensorflow", "torch"])
+    @pytest.mark.parametrize("params", np.linspace(0, 2 * np.pi, 8))
+    def test_qnode_state(self, device, interface, params):
+        """Test that the mutual information transform works for QNodes by comparing
+        against analytic values"""
+        dev = qml.device(device, wires=2)
+
+        params = qml.math.asarray(params, like=interface)
+
+        @qml.qnode(dev, interface=interface)
+        def circuit(params):
+            qml.RY(params, wires=0)
+            qml.CNOT(wires=[0, 1])
+            return qml.state()
+
+        actual = qml.qinfo.mutual_info(circuit, wires0=[0], wires1=[1])(params)
+
+        # compare transform results with analytic values
+        expected = -2 * np.cos(params / 2) ** 2 * np.log(
+            np.cos(params / 2) ** 2 + 1e-10
+        ) - 2 * np.sin(params / 2) ** 2 * np.log(np.sin(params / 2) ** 2 + 1e-10)
+
+        assert np.allclose(actual, expected)
+
+    @pytest.mark.all_interfaces
+    @pytest.mark.parametrize("device", ["default.qubit.legacy", "default.mixed", "lightning.qubit"])
+    @pytest.mark.parametrize("interface", ["autograd", "jax", "tensorflow", "torch"])
+    @pytest.mark.parametrize("params", zip(np.linspace(0, np.pi, 8), np.linspace(0, 2 * np.pi, 8)))
+    def test_qnode_mutual_info(self, device, interface, params):
+        """Test that the measurement process for mutual information works for QNodes
+        by comparing against the mutual information transform"""
+        dev = qml.device(device, wires=2)
+
+        params = qml.math.asarray(np.array(params), like=interface)
+
+        @qml.qnode(dev, interface=interface)
+        def circuit_mutual_info(params):
+            qml.RY(params[0], wires=0)
+            qml.RY(params[1], wires=1)
+            qml.CNOT(wires=[0, 1])
+            return qml.mutual_info(wires0=[0], wires1=[1])
+
+        @qml.qnode(dev, interface=interface)
+        def circuit_state(params):
+            qml.RY(params[0], wires=0)
+            qml.RY(params[1], wires=1)
+            qml.CNOT(wires=[0, 1])
+            return qml.state()
+
+        actual = circuit_mutual_info(params)
+
+        # compare measurement results with transform results
+        expected = qml.qinfo.mutual_info(circuit_state, wires0=[0], wires1=[1])(params)
+
+        assert np.allclose(actual, expected)
+
+    @pytest.mark.parametrize("device", ["default.qubit.legacy", "default.mixed", "lightning.qubit"])
+    def test_mutual_info_wire_labels(self, device):
+        """Test that mutual_info is correct with custom wire labels"""
+        param = np.array([0.678, 1.234])
+        wires = ["a", 8]
+        dev = qml.device(device, wires=wires)
+
+        @qml.qnode(dev)
+        def circuit(param):
+            qml.RY(param, wires=wires[0])
+            qml.CNOT(wires=wires)
+            return qml.state()
+
+        actual = qml.qinfo.mutual_info(circuit, wires0=[wires[0]], wires1=[wires[1]])(param)
+
+        # compare transform results with analytic values
+        expected = -2 * np.cos(param / 2) ** 2 * np.log(np.cos(param / 2) ** 2) - 2 * np.sin(
+            param / 2
+        ) ** 2 * np.log(np.sin(param / 2) ** 2)
+
+        assert np.allclose(actual, expected)
+
+    @pytest.mark.jax
+    @pytest.mark.parametrize("params", np.linspace(0, 2 * np.pi, 8))
+    def test_qnode_state_jax_jit(self, params):
+        """Test that the mutual information transform works for QNodes by comparing
+        against analytic values, for the JAX-jit interface"""
+        import jax
+        import jax.numpy as jnp
+
+        dev = qml.device("default.qubit.legacy", wires=2)
+
+        params = jnp.array(params)
+
+        @qml.qnode(dev, interface="jax-jit")
+        def circuit(params):
+            qml.RY(params, wires=0)
+            qml.CNOT(wires=[0, 1])
+            return qml.state()
+
+        actual = jax.jit(qml.qinfo.mutual_info(circuit, wires0=[0], wires1=[1]))(params)
+
+        # compare transform results with analytic values
+        expected = -2 * jnp.cos(params / 2) ** 2 * jnp.log(
+            jnp.cos(params / 2) ** 2 + 1e-10
+        ) - 2 * jnp.sin(params / 2) ** 2 * jnp.log(jnp.sin(params / 2) ** 2 + 1e-10)
+
+        assert np.allclose(actual, expected)
+
+    @pytest.mark.jax
+    @pytest.mark.parametrize("params", zip(np.linspace(0, np.pi, 8), np.linspace(0, 2 * np.pi, 8)))
+    @pytest.mark.parametrize("interface", ["jax-jit"])
+    def test_qnode_mutual_info_jax_jit(self, params, interface):
+        """Test that the measurement process for mutual information works for QNodes
+        by comparing against the mutual information transform, for the JAX-jit interface"""
+        import jax
+        import jax.numpy as jnp
+
+        dev = qml.device("default.qubit.legacy", wires=2)
+
+        params = jnp.array(params)
+
+        @qml.qnode(dev, interface=interface)
+        def circuit_mutual_info(params):
+            qml.RY(params[0], wires=0)
+            qml.RY(params[1], wires=1)
+            qml.CNOT(wires=[0, 1])
+            return qml.mutual_info(wires0=[0], wires1=[1])
+
+        @qml.qnode(dev, interface="jax-jit")
+        def circuit_state(params):
+            qml.RY(params[0], wires=0)
+            qml.RY(params[1], wires=1)
+            qml.CNOT(wires=[0, 1])
+            return qml.state()
+
+        actual = jax.jit(circuit_mutual_info)(params)
+
+        # compare measurement results with transform results
+        expected = jax.jit(qml.qinfo.mutual_info(circuit_state, wires0=[0], wires1=[1]))(params)
+
+        assert np.allclose(actual, expected)
+
+    @pytest.mark.autograd
+    @pytest.mark.parametrize("param", np.linspace(0, 2 * np.pi, 16))
+    @pytest.mark.parametrize("diff_method", diff_methods)
+    @pytest.mark.parametrize("interface", ["auto", "autograd"])
+    def test_qnode_grad(self, param, diff_method, interface):
+        """Test that the gradient of mutual information works for QNodes
+        with the autograd interface"""
+        dev = qml.device("default.qubit.legacy", wires=2)
+
+        @qml.qnode(dev, interface=interface, diff_method=diff_method)
+        def circuit(param):
+            qml.RY(param, wires=0)
+            qml.CNOT(wires=[0, 1])
+            return qml.mutual_info(wires0=[0], wires1=[1])
+
+        if param == 0:
+            # we don't allow gradients to flow through the discontinuity at 0
+            expected = 0
+        else:
+            expected = np.sin(param) * (
+                np.log(np.cos(param / 2) ** 2) - np.log(np.sin(param / 2) ** 2)
+            )
+
+        # higher tolerance for finite-diff method
+        tol = 1e-8 if diff_method == "backprop" else 1e-5
+
+        actual = qml.grad(circuit)(param)
+        assert np.allclose(actual, expected, atol=tol)
+
+    @pytest.mark.jax
+    @pytest.mark.parametrize("param", np.linspace(0, 2 * np.pi, 16))
+    @pytest.mark.parametrize("diff_method", diff_methods)
+    @pytest.mark.parametrize("interface", ["jax"])
+    def test_qnode_grad_jax(self, param, diff_method, interface):
+        """Test that the gradient of mutual information works for QNodes
+        with the JAX interface"""
+        import jax
+        import jax.numpy as jnp
+
+        dev = qml.device("default.qubit.legacy", wires=2)
+
+        param = jnp.array(param)
+
+        @qml.qnode(dev, interface=interface, diff_method=diff_method)
+        def circuit(param):
+            qml.RY(param, wires=0)
+            qml.CNOT(wires=[0, 1])
+            return qml.mutual_info(wires0=[0], wires1=[1])
+
+        if param == 0:
+            # we don't allow gradients to flow through the discontinuity at 0
+            expected = 0
+        else:
+            expected = jnp.sin(param) * (
+                jnp.log(jnp.cos(param / 2) ** 2) - jnp.log(jnp.sin(param / 2) ** 2)
+            )
+
+        # higher tolerance for finite-diff method
+        tol = 1e-8 if diff_method == "backprop" else 1e-5
+
+        actual = jax.grad(circuit)(param)
+        assert np.allclose(actual, expected, atol=tol)
+
+    @pytest.mark.jax
+    @pytest.mark.parametrize("param", np.linspace(0, 2 * np.pi, 16))
+    @pytest.mark.parametrize("diff_method", diff_methods)
+    @pytest.mark.parametrize("interface", ["jax-jit"])
+    def test_qnode_grad_jax_jit(self, param, diff_method, interface):
+        """Test that the gradient of mutual information works for QNodes
+        with the JAX-jit interface"""
+        import jax
+        import jax.numpy as jnp
+
+        dev = qml.device("default.qubit.legacy", wires=2)
+
+        param = jnp.array(param)
+
+        @qml.qnode(dev, interface=interface, diff_method=diff_method)
+        def circuit(param):
+            qml.RY(param, wires=0)
+            qml.CNOT(wires=[0, 1])
+            return qml.mutual_info(wires0=[0], wires1=[1])
+
+        if param == 0:
+            # we don't allow gradients to flow through the discontinuity at 0
+            expected = 0
+        else:
+            expected = jnp.sin(param) * (
+                jnp.log(jnp.cos(param / 2) ** 2) - jnp.log(jnp.sin(param / 2) ** 2)
+            )
+
+        # higher tolerance for finite-diff method
+        tol = 1e-8 if diff_method == "backprop" else 1e-5
+
+        actual = jax.jit(jax.grad(circuit))(param)
+        assert np.allclose(actual, expected, atol=tol)
+
+    @pytest.mark.tf
+    @pytest.mark.parametrize("param", np.linspace(0, 2 * np.pi, 16))
+    @pytest.mark.parametrize("diff_method", diff_methods)
+    @pytest.mark.parametrize("interface", ["tf"])
+    def test_qnode_grad_tf(self, param, diff_method, interface):
+        """Test that the gradient of mutual information works for QNodes
+        with the tensorflow interface"""
+        import tensorflow as tf
+
+        dev = qml.device("default.qubit.legacy", wires=2)
+
+        param = tf.Variable(param)
+
+        @qml.qnode(dev, interface=interface, diff_method=diff_method)
+        def circuit(param):
+            qml.RY(param, wires=0)
+            qml.CNOT(wires=[0, 1])
+            return qml.mutual_info(wires0=[0], wires1=[1])
+
+        if param == 0:
+            # we don't allow gradients to flow through the discontinuity at 0
+            expected = 0
+        else:
+            expected = np.sin(param) * (
+                np.log(np.cos(param / 2) ** 2) - np.log(np.sin(param / 2) ** 2)
+            )
+
+        with tf.GradientTape() as tape:
+            out = circuit(param)
+
+        # higher tolerance for finite-diff method
+        tol = 1e-8 if diff_method == "backprop" else 1e-5
+
+        actual = tape.gradient(out, param)
+        assert np.allclose(actual, expected, atol=tol)
+
+    @pytest.mark.torch
+    @pytest.mark.parametrize("param", np.linspace(0, 2 * np.pi, 16))
+    @pytest.mark.parametrize("diff_method", diff_methods)
+    @pytest.mark.parametrize("interface", ["torch"])
+    def test_qnode_grad_torch(self, param, diff_method, interface):
+        """Test that the gradient of mutual information works for QNodes
+        with the torch interface"""
+        import torch
+
+        dev = qml.device("default.qubit.legacy", wires=2)
+
+        @qml.qnode(dev, interface=interface, diff_method=diff_method)
+        def circuit(param):
+            qml.RY(param, wires=0)
+            qml.CNOT(wires=[0, 1])
+            return qml.mutual_info(wires0=[0], wires1=[1])
+
+        if param == 0:
+            # we don't allow gradients to flow through the discontinuity at 0
+            expected = 0
+        else:
+            expected = np.sin(param) * (
+                np.log(np.cos(param / 2) ** 2) - np.log(np.sin(param / 2) ** 2)
+            )
+
+        param = torch.tensor(param, requires_grad=True)
+        out = circuit(param)
+        out.backward()  # pylint: disable=no-member
+
+        # higher tolerance for finite-diff method
+        tol = 1e-8 if diff_method == "backprop" else 1e-5
+
+        actual = param.grad
+        assert np.allclose(actual, expected, atol=tol)
+
+    @pytest.mark.all_interfaces
+    @pytest.mark.parametrize("interface", ["autograd", "jax", "tensorflow", "torch"])
+    @pytest.mark.parametrize(
+        "params", [np.array([0.0, 0.0]), np.array([0.3, 0.4]), np.array([0.6, 0.8])]
+    )
+    def test_subsystem_overlap_error(self, interface, params):
+        """Test that an error is raised when the subsystems overlap"""
+        dev = qml.device("default.qubit.legacy", wires=3)
+
+        params = qml.math.asarray(params, like=interface)
+
+        @qml.qnode(dev, interface=interface)
+        def circuit(params):
+            qml.RY(params[0], wires=0)
+            qml.RY(params[1], wires=1)
+            qml.CNOT(wires=[0, 1])
+            qml.CNOT(wires=[0, 2])
+            return qml.mutual_info(wires0=[0, 1], wires1=[1, 2])
+
+        msg = "Subsystems for computing mutual information must not overlap"
+        with pytest.raises(qml.QuantumFunctionError, match=msg):
+            circuit(params)
+
+    @pytest.mark.all_interfaces
+    @pytest.mark.parametrize("interface", ["autograd", "jax", "tensorflow", "torch"])
+    @pytest.mark.parametrize(
+        "params", [np.array([0.0, 0.0]), np.array([0.3, 0.4]), np.array([0.6, 0.8])]
+    )
+    def test_custom_wire_labels_error(self, interface, params):
+        """Tests that an error is raised when mutual information is measured
+        with custom wire labels"""
+        dev = qml.device("default.qubit.legacy", wires=["a", "b"])
+
+        params = qml.math.asarray(params, like=interface)
+
+        @qml.qnode(dev, interface=interface)
+        def circuit(params):
+            qml.RY(params[0], wires="a")
+            qml.RY(params[1], wires="b")
+            qml.CNOT(wires=["a", "b"])
+            return qml.mutual_info(wires0=["a"], wires1=["b"])
+
+        msg = "Returning the mutual information is not supported when using custom wire labels"
+        with pytest.raises(qml.QuantumFunctionError, match=msg):
+            circuit(params)
diff --git a/tests/measurements/legacy/test_probs_legacy.py b/tests/measurements/legacy/test_probs_legacy.py
new file mode 100644
index 00000000000..2e7605bc151
--- /dev/null
+++ b/tests/measurements/legacy/test_probs_legacy.py
@@ -0,0 +1,634 @@
+# Copyright 2018-2020 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 probs module"""
+import numpy as np
+import pytest
+
+import pennylane as qml
+from pennylane import numpy as pnp
+from pennylane.measurements import ProbabilityMP, Shots
+
+
+# TODO: Remove this when new CustomMP are the default
+def custom_measurement_process(device, spy):
+    assert len(spy.call_args_list) > 0  # make sure method is mocked properly
+
+    samples = device._samples  # pylint: disable=protected-access
+    state = device._state  # pylint: disable=protected-access
+    call_args_list = list(spy.call_args_list)
+    for call_args in call_args_list:
+        wires, shot_range, bin_size = (
+            call_args.kwargs["wires"],
+            call_args.kwargs["shot_range"],
+            call_args.kwargs["bin_size"],
+        )
+        # no need to use op, because the observable has already been applied to ``dev._state``
+        meas = qml.probs(wires=wires)
+        old_res = device.probability(wires=wires, shot_range=shot_range, bin_size=bin_size)
+        if device.shots is None:
+            new_res = meas.process_state(state=state, wire_order=device.wires)
+        else:
+            new_res = meas.process_samples(
+                samples=samples,
+                wire_order=device.wires,
+                shot_range=shot_range,
+                bin_size=bin_size,
+            )
+        assert qml.math.allequal(old_res, new_res)
+
+
+# make the test deterministic
+np.random.seed(42)
+
+
+@pytest.fixture(name="init_state")
+def fixture_init_state():
+    """Fixture that creates an initial state"""
+
+    def _init_state(n):
+        """An initial state over n wires"""
+        state = np.random.random([2**n]) + np.random.random([2**n]) * 1j
+        state /= np.linalg.norm(state)
+        return state
+
+    return _init_state
+
+
+class TestProbs:
+    """Tests for the probs function"""
+
+    # pylint:disable=too-many-public-methods
+
+    def test_queue(self):
+        """Test that the right measurement class is queued."""
+        dev = qml.device("default.qubit.legacy", wires=2)
+
+        @qml.qnode(dev)
+        def circuit():
+            return qml.probs(wires=0)
+
+        circuit()
+
+        assert isinstance(circuit.tape[0], ProbabilityMP)
+
+    @pytest.mark.parametrize("wires", [[0], [2, 1], ["a", "c", 3]])
+    @pytest.mark.parametrize("shots", [None, 10])
+    def test_shape(self, wires, shots):
+        """Test that the shape is correct."""
+        dev = qml.device("default.qubit.legacy", wires=3, shots=shots)
+        res = qml.probs(wires=wires)
+        assert res.shape(dev, Shots(shots)) == (2 ** len(wires),)
+
+    @pytest.mark.parametrize("wires", [[0], [2, 1], ["a", "c", 3]])
+    def test_shape_shot_vector(self, wires):
+        """Test that the shape is correct with the shot vector too."""
+        res = qml.probs(wires=wires)
+        shot_vector = (1, 2, 3)
+        dev = qml.device("default.qubit.legacy", wires=3, shots=shot_vector)
+        assert res.shape(dev, Shots(shot_vector)) == (
+            (2 ** len(wires),),
+            (2 ** len(wires),),
+            (2 ** len(wires),),
+        )
+
+    @pytest.mark.parametrize("shots", [None, 100])
+    def test_probs_no_arguments(self, shots):
+        """Test that using ``qml.probs`` with no arguments returns the probabilities of all wires."""
+        dev = qml.device("default.qubit.legacy", wires=3, shots=shots)
+
+        @qml.qnode(dev)
+        def circuit():
+            return qml.probs()
+
+        res = circuit()
+
+        assert qml.math.allequal(res, [1, 0, 0, 0, 0, 0, 0, 0])
+
+    def test_full_prob(self, init_state, tol, mocker):
+        """Test that the correct probability is returned."""
+        dev = qml.device("default.qubit.legacy", wires=4)
+        spy = mocker.spy(qml.QubitDevice, "probability")
+
+        state = init_state(4)
+
+        @qml.qnode(dev)
+        def circuit():
+            qml.StatePrep(state, wires=list(range(4)))
+            return qml.probs(wires=range(4))
+
+        res = circuit()
+        expected = np.abs(state) ** 2
+        assert np.allclose(res, expected, atol=tol, rtol=0)
+
+        custom_measurement_process(dev, spy)
+
+    def test_marginal_prob(self, init_state, tol, mocker):
+        """Test that the correct marginal probability is returned."""
+        dev = qml.device("default.qubit.legacy", wires=4)
+        spy = mocker.spy(qml.QubitDevice, "probability")
+
+        state = init_state(4)
+
+        @qml.qnode(dev)
+        def circuit():
+            qml.StatePrep(state, wires=list(range(4)))
+            return qml.probs(wires=[1, 3])
+
+        res = circuit()
+        expected = np.reshape(np.abs(state) ** 2, [2] * 4)
+        expected = np.einsum("ijkl->jl", expected).flatten()
+        assert np.allclose(res, expected, atol=tol, rtol=0)
+
+        custom_measurement_process(dev, spy)
+
+    def test_marginal_prob_more_wires(self, init_state, mocker, tol):
+        """Test that the correct marginal probability is returned, when the
+        states_to_binary method is used for probability computations."""
+        dev = qml.device("default.qubit.legacy", wires=4)
+        spy_probs = mocker.spy(qml.QubitDevice, "probability")
+        state = init_state(4)
+
+        @qml.qnode(dev)
+        def circuit():
+            qml.StatePrep(state, wires=list(range(4)))
+            return qml.probs(wires=[1, 0, 3])
+
+        res = circuit()
+
+        expected = np.reshape(np.abs(state) ** 2, [2] * 4)
+        expected = np.einsum("ijkl->jil", expected).flatten()
+        assert np.allclose(res, expected, atol=tol, rtol=0)
+
+        custom_measurement_process(dev, spy_probs)
+
+    def test_integration(self, tol, mocker):
+        """Test the probability is correct for a known state preparation."""
+        dev = qml.device("default.qubit.legacy", wires=2)
+        spy = mocker.spy(qml.QubitDevice, "probability")
+
+        @qml.qnode(dev)
+        def circuit():
+            qml.Hadamard(wires=1)
+            qml.CNOT(wires=[0, 1])
+            return qml.probs(wires=[0, 1])
+
+        # expected probability, using [00, 01, 10, 11]
+        # ordering, is [0.5, 0.5, 0, 0]
+
+        res = circuit()
+        expected = np.array([0.5, 0.5, 0, 0])
+        assert np.allclose(res, expected, atol=tol, rtol=0)
+
+        custom_measurement_process(dev, spy)
+
+    @pytest.mark.parametrize("shots", [100, [1, 10, 100]])
+    def test_integration_analytic_false(self, tol, mocker, shots):
+        """Test the probability is correct for a known state preparation when the
+        analytic attribute is set to False."""
+        dev = qml.device("default.qubit.legacy", wires=3, shots=shots)
+        spy = mocker.spy(qml.QubitDevice, "probability")
+
+        @qml.qnode(dev)
+        def circuit():
+            qml.PauliX(0)
+            return qml.probs(wires=dev.wires)
+
+        res = circuit()
+        expected = np.array([0, 0, 0, 0, 1, 0, 0, 0])
+        assert np.allclose(res, expected, atol=tol, rtol=0)
+
+        custom_measurement_process(dev, spy)
+
+    @pytest.mark.parametrize("shots", [None, 10000, [10000, 10000]])
+    @pytest.mark.parametrize("phi", np.arange(0, 2 * np.pi, np.pi / 3))
+    def test_observable_is_measurement_value(
+        self, shots, phi, mocker, tol, tol_stochastic
+    ):  # pylint: disable=too-many-arguments
+        """Test that expectation values for mid-circuit measurement values
+        are correct for a single measurement value."""
+        dev = qml.device("default.qubit", wires=2, shots=shots)
+
+        @qml.qnode(dev)
+        def circuit(phi):
+            qml.RX(phi, 0)
+            m0 = qml.measure(0)
+            return qml.probs(op=m0)
+
+        new_dev = circuit.device
+        spy = mocker.spy(qml.QubitDevice, "probability")
+
+        res = circuit(phi)
+
+        atol = tol if shots is None else tol_stochastic
+        expected = np.array([np.cos(phi / 2) ** 2, np.sin(phi / 2) ** 2])
+
+        if not isinstance(shots, list):
+            assert np.allclose(np.array(res), expected, atol=atol, rtol=0)
+        else:
+            for r in res:  # pylint: disable=not-an-iterable
+                assert np.allclose(r, expected, atol=atol, rtol=0)
+
+        custom_measurement_process(new_dev, spy)
+
+    @pytest.mark.parametrize("shots", [None, 100])
+    def test_batch_size(self, mocker, shots):
+        """Test the probability is correct for a batched input."""
+        dev = qml.device("default.qubit.legacy", wires=1, shots=shots)
+        spy = mocker.spy(qml.QubitDevice, "probability")
+
+        @qml.qnode(dev)
+        def circuit(x):
+            qml.RX(x, 0)
+            return qml.probs(wires=dev.wires)  # TODO: Use ``qml.probs()`` when supported
+
+        x = np.array([0, np.pi / 2])
+        res = circuit(x)
+        expected = [[1.0, 0.0], [0.5, 0.5]]
+        assert np.allclose(res, expected, atol=0.1, rtol=0.1)
+
+        custom_measurement_process(dev, spy)
+
+    @pytest.mark.autograd
+    def test_numerical_analytic_diff_agree(self, tol, mocker):
+        """Test that the finite difference and parameter shift rule
+        provide the same Jacobian."""
+        w = 4
+        dev = qml.device("default.qubit.legacy", wires=w)
+        spy = mocker.spy(qml.QubitDevice, "probability")
+
+        def circuit(x, y, z):
+            for i in range(w):
+                qml.RX(x, wires=i)
+                qml.PhaseShift(z, wires=i)
+                qml.RY(y, wires=i)
+
+            qml.CNOT(wires=[0, 1])
+            qml.CNOT(wires=[1, 2])
+            qml.CNOT(wires=[2, 3])
+
+            return qml.probs(wires=[1, 3])
+
+        params = pnp.array([0.543, -0.765, -0.3], requires_grad=True)
+
+        circuit_F = qml.QNode(circuit, dev, diff_method="finite-diff")
+        circuit_A = qml.QNode(circuit, dev, diff_method="parameter-shift")
+        res_F = qml.jacobian(circuit_F)(*params)
+        res_A = qml.jacobian(circuit_A)(*params)
+
+        # Both jacobians should be of shape (2**prob.wires, num_params)
+        assert isinstance(res_F, tuple) and len(res_F) == 3
+        assert all(_r.shape == (2**2,) for _r in res_F)
+        assert isinstance(res_A, tuple) and len(res_A) == 3
+        assert all(_r.shape == (2**2,) for _r in res_A)
+
+        # Check that they agree up to numeric tolerance
+        assert all(np.allclose(_rF, _rA, atol=tol, rtol=0) for _rF, _rA in zip(res_F, res_A))
+
+        custom_measurement_process(dev, spy)
+
+    @pytest.mark.parametrize("hermitian", [1 / np.sqrt(2) * np.array([[1, 1], [1, -1]])])
+    def test_prob_generalize_param_one_qubit(self, hermitian, tol, mocker):
+        """Test that the correct probability is returned."""
+        dev = qml.device("default.qubit.legacy", wires=1)
+        spy = mocker.spy(qml.QubitDevice, "probability")
+
+        @qml.qnode(dev)
+        def circuit(x):
+            qml.RZ(x, wires=0)
+            return qml.probs(op=qml.Hermitian(hermitian, wires=0))
+
+        res = circuit(0.56)
+
+        def circuit_rotated(x):
+            qml.RZ(x, wires=0)
+            qml.Hermitian(hermitian, wires=0).diagonalizing_gates()
+
+        state = np.array([1, 0])
+        matrix = qml.matrix(circuit_rotated)(0.56)
+        state = np.dot(matrix, state)
+        expected = np.reshape(np.abs(state) ** 2, [2] * 1)
+        expected = expected.flatten()
+
+        assert np.allclose(res, expected, atol=tol, rtol=0)
+
+        custom_measurement_process(dev, spy)
+
+    @pytest.mark.parametrize("hermitian", [1 / np.sqrt(2) * np.array([[1, 1], [1, -1]])])
+    def test_prob_generalize_param(self, hermitian, tol, mocker):
+        """Test that the correct probability is returned."""
+        dev = qml.device("default.qubit.legacy", wires=3)
+        spy = mocker.spy(qml.QubitDevice, "probability")
+
+        @qml.qnode(dev)
+        def circuit(x, y):
+            qml.RZ(x, wires=0)
+            qml.CNOT(wires=[0, 1])
+            qml.RY(y, wires=1)
+            qml.CNOT(wires=[0, 2])
+            return qml.probs(op=qml.Hermitian(hermitian, wires=0))
+
+        res = circuit(0.56, 0.1)
+
+        def circuit_rotated(x, y):
+            qml.RZ(x, wires=0)
+            qml.CNOT(wires=[0, 1])
+            qml.RY(y, wires=1)
+            qml.CNOT(wires=[0, 2])
+            qml.Hermitian(hermitian, wires=0).diagonalizing_gates()
+
+        state = np.array([1, 0, 0, 0, 0, 0, 0, 0])
+        matrix = qml.matrix(circuit_rotated)(0.56, 0.1)
+        state = np.dot(matrix, state)
+        expected = np.reshape(np.abs(state) ** 2, [2] * 3)
+        expected = np.einsum("ijk->i", expected).flatten()
+        assert np.allclose(res, expected, atol=tol, rtol=0)
+
+        custom_measurement_process(dev, spy)
+
+    @pytest.mark.parametrize("hermitian", [1 / np.sqrt(2) * np.array([[1, 1], [1, -1]])])
+    def test_prob_generalize_param_multiple(self, hermitian, tol, mocker):
+        """Test that the correct probability is returned."""
+        dev = qml.device("default.qubit.legacy", wires=3)
+        spy = mocker.spy(qml.QubitDevice, "probability")
+
+        @qml.qnode(dev)
+        def circuit(x, y):
+            qml.RZ(x, wires=0)
+            qml.CNOT(wires=[0, 1])
+            qml.RY(y, wires=1)
+            qml.CNOT(wires=[0, 2])
+            return (
+                qml.probs(op=qml.Hermitian(hermitian, wires=0)),
+                qml.probs(wires=[1]),
+                qml.probs(wires=[2]),
+            )
+
+        res = circuit(0.56, 0.1)
+        res = np.reshape(res, (3, 2))
+
+        def circuit_rotated(x, y):
+            qml.RZ(x, wires=0)
+            qml.CNOT(wires=[0, 1])
+            qml.RY(y, wires=1)
+            qml.CNOT(wires=[0, 2])
+            qml.Hermitian(hermitian, wires=0).diagonalizing_gates()
+
+        state = np.array([1, 0, 0, 0, 0, 0, 0, 0])
+        matrix = qml.matrix(circuit_rotated)(0.56, 0.1)
+        state = np.dot(matrix, state)
+
+        expected = np.reshape(np.abs(state) ** 2, [2] * 3)
+        expected_0 = np.einsum("ijk->i", expected).flatten()
+        expected_1 = np.einsum("ijk->j", expected).flatten()
+        expected_2 = np.einsum("ijk->k", expected).flatten()
+
+        assert np.allclose(res[0], expected_0, atol=tol, rtol=0)
+        assert np.allclose(res[1], expected_1, atol=tol, rtol=0)
+        assert np.allclose(res[2], expected_2, atol=tol, rtol=0)
+
+        custom_measurement_process(dev, spy)
+
+    @pytest.mark.parametrize("hermitian", [1 / np.sqrt(2) * np.array([[1, 1], [1, -1]])])
+    @pytest.mark.parametrize("wire", [0, 1, 2, 3])
+    def test_prob_generalize_initial_state(self, hermitian, wire, init_state, tol, mocker):
+        """Test that the correct probability is returned."""
+        # pylint:disable=too-many-arguments
+        dev = qml.device("default.qubit.legacy", wires=4)
+        spy = mocker.spy(qml.QubitDevice, "probability")
+        state = init_state(4)
+
+        @qml.qnode(dev)
+        def circuit():
+            qml.StatePrep(state, wires=list(range(4)))
+            qml.PauliX(wires=0)
+            qml.PauliX(wires=1)
+            qml.PauliX(wires=2)
+            qml.PauliX(wires=3)
+            return qml.probs(op=qml.Hermitian(hermitian, wires=wire))
+
+        res = circuit()
+
+        def circuit_rotated():
+            qml.PauliX(wires=0)
+            qml.PauliX(wires=1)
+            qml.PauliX(wires=2)
+            qml.PauliX(wires=3)
+            qml.Hermitian(hermitian, wires=wire).diagonalizing_gates()
+
+        matrix = qml.matrix(circuit_rotated)()
+        state = np.dot(matrix, state)
+        expected = np.reshape(np.abs(state) ** 2, [2] * 4)
+
+        if wire == 0:
+            expected = np.einsum("ijkl->i", expected).flatten()
+        elif wire == 1:
+            expected = np.einsum("ijkl->j", expected).flatten()
+        elif wire == 2:
+            expected = np.einsum("ijkl->k", expected).flatten()
+        elif wire == 3:
+            expected = np.einsum("ijkl->l", expected).flatten()
+
+        assert np.allclose(res, expected, atol=tol, rtol=0)
+
+        custom_measurement_process(dev, spy)
+
+    @pytest.mark.parametrize("operation", [qml.PauliX, qml.PauliY, qml.Hadamard])
+    @pytest.mark.parametrize("wire", [0, 1, 2, 3])
+    def test_operation_prob(self, operation, wire, init_state, tol, mocker):
+        "Test the rotated probability with different wires and rotating operations."
+        # pylint:disable=too-many-arguments
+        dev = qml.device("default.qubit.legacy", wires=4)
+        spy = mocker.spy(qml.QubitDevice, "probability")
+        state = init_state(4)
+
+        @qml.qnode(dev)
+        def circuit():
+            qml.StatePrep(state, wires=list(range(4)))
+            qml.PauliX(wires=0)
+            qml.PauliZ(wires=1)
+            qml.PauliY(wires=2)
+            qml.PauliZ(wires=3)
+            return qml.probs(op=operation(wires=wire))
+
+        res = circuit()
+
+        def circuit_rotated():
+            qml.PauliX(wires=0)
+            qml.PauliZ(wires=1)
+            qml.PauliY(wires=2)
+            qml.PauliZ(wires=3)
+            operation(wires=wire).diagonalizing_gates()
+
+        matrix = qml.matrix(circuit_rotated)()
+        state = np.dot(matrix, state)
+        expected = np.reshape(np.abs(state) ** 2, [2] * 4)
+
+        if wire == 0:
+            expected = np.einsum("ijkl->i", expected).flatten()
+        elif wire == 1:
+            expected = np.einsum("ijkl->j", expected).flatten()
+        elif wire == 2:
+            expected = np.einsum("ijkl->k", expected).flatten()
+        elif wire == 3:
+            expected = np.einsum("ijkl->l", expected).flatten()
+
+        assert np.allclose(res, expected, atol=tol, rtol=0)
+
+        custom_measurement_process(dev, spy)
+
+    @pytest.mark.parametrize("observable", [(qml.PauliX, qml.PauliY)])
+    def test_observable_tensor_prob(self, observable, init_state, tol, mocker):
+        "Test the rotated probability with a tensor observable."
+        dev = qml.device("default.qubit.legacy", wires=4)
+        spy = mocker.spy(qml.QubitDevice, "probability")
+        state = init_state(4)
+
+        @qml.qnode(dev)
+        def circuit():
+            qml.StatePrep(state, wires=list(range(4)))
+            qml.PauliX(wires=0)
+            qml.PauliZ(wires=1)
+            qml.PauliY(wires=2)
+            qml.PauliZ(wires=3)
+            return qml.probs(op=observable[0](wires=0) @ observable[1](wires=1))
+
+        res = circuit()
+
+        def circuit_rotated():
+            qml.PauliX(wires=0)
+            qml.PauliZ(wires=1)
+            qml.PauliY(wires=2)
+            qml.PauliZ(wires=3)
+            observable[0](wires=0).diagonalizing_gates()
+            observable[1](wires=1).diagonalizing_gates()
+
+        matrix = qml.matrix(circuit_rotated)()
+        state = np.dot(matrix, state)
+        expected = np.reshape(np.abs(state) ** 2, [2] * 4)
+
+        expected = np.einsum("ijkl->ij", expected).flatten()
+
+        assert np.allclose(res, expected, atol=tol, rtol=0)
+
+        custom_measurement_process(dev, spy)
+
+    @pytest.mark.parametrize("coeffs, obs", [([1, 1], [qml.PauliX(wires=0), qml.PauliX(wires=1)])])
+    def test_hamiltonian_error(self, coeffs, obs, init_state):
+        "Test that an error is returned for hamiltonians."
+        H = qml.Hamiltonian(coeffs, obs)
+
+        dev = qml.device("default.qubit.legacy", wires=4)
+        state = init_state(4)
+
+        @qml.qnode(dev)
+        def circuit():
+            qml.StatePrep(state, wires=list(range(4)))
+            qml.PauliX(wires=0)
+            qml.PauliZ(wires=1)
+            qml.PauliY(wires=2)
+            qml.PauliZ(wires=3)
+            return qml.probs(op=H)
+
+        with pytest.raises(
+            qml.QuantumFunctionError,
+            match="Hamiltonians are not supported for rotating probabilities.",
+        ):
+            circuit()
+
+    @pytest.mark.parametrize(
+        "operation", [qml.SingleExcitation, qml.SingleExcitationPlus, qml.SingleExcitationMinus]
+    )
+    def test_generalize_prob_not_hermitian(self, operation):
+        """Test that Operators that do not have a diagonalizing_gates representation cannot
+        be used in probability measurements."""
+
+        dev = qml.device("default.qubit.legacy", wires=2)
+
+        @qml.qnode(dev)
+        def circuit():
+            qml.PauliX(wires=0)
+            qml.PauliZ(wires=1)
+            return qml.probs(op=operation(0.56, wires=[0, 1]))
+
+        with pytest.raises(
+            qml.QuantumFunctionError,
+            match="does not define diagonalizing gates : cannot be used to rotate the probability",
+        ):
+            circuit()
+
+    @pytest.mark.parametrize("hermitian", [1 / np.sqrt(2) * np.array([[1, 1], [1, -1]])])
+    def test_prob_wires_and_hermitian(self, hermitian):
+        """Test that we can cannot give simultaneously wires and a hermitian."""
+
+        dev = qml.device("default.qubit.legacy", wires=2)
+
+        @qml.qnode(dev)
+        def circuit():
+            qml.PauliX(wires=0)
+            return qml.probs(op=qml.Hermitian(hermitian, wires=0), wires=1)
+
+        with pytest.raises(
+            qml.QuantumFunctionError,
+            match="Cannot specify the wires to probs if an observable is "
+            "provided. The wires for probs will be determined directly from the observable.",
+        ):
+            circuit()
+
+    def test_non_commuting_probs_raises_error(self):
+        dev = qml.device("default.qubit.legacy", wires=5)
+
+        @qml.qnode(dev)
+        def circuit(x, y):
+            qml.RX(x, wires=[0])
+            qml.RY(y, wires=[1])
+            qml.CNOT(wires=[0, 1])
+            return qml.expval(qml.PauliX(0)), qml.probs(wires=[0, 1])
+
+        with pytest.raises(
+            qml.QuantumFunctionError, match="Only observables that are qubit-wise commuting"
+        ):
+            circuit(1, 2)
+
+    def test_commuting_probs_in_computational_basis(self):
+        """Test that `qml.probs` can be used in the computational basis with other commuting observables."""
+        dev = qml.device("default.qubit.legacy", wires=5)
+
+        @qml.qnode(dev)
+        def circuit(x, y):
+            qml.RX(x, wires=[0])
+            qml.RY(y, wires=[1])
+            qml.CNOT(wires=[0, 1])
+            return qml.expval(qml.PauliZ(0)), qml.probs(wires=[0, 1])
+
+        res = circuit(1, 2)
+
+        @qml.qnode(dev)
+        def circuit2(x, y):
+            qml.RX(x, wires=[0])
+            qml.RY(y, wires=[1])
+            qml.CNOT(wires=[0, 1])
+            return qml.expval(qml.PauliZ(0))
+
+        @qml.qnode(dev)
+        def circuit3(x, y):
+            qml.RX(x, wires=[0])
+            qml.RY(y, wires=[1])
+            qml.CNOT(wires=[0, 1])
+            return qml.probs(wires=[0, 1])
+
+        res2 = circuit2(1, 2)
+        res3 = circuit3(1, 2)
+
+        assert res[0] == res2
+        assert qml.math.allequal(res[1:], res3)
diff --git a/tests/measurements/legacy/test_purity_measurement_legacy.py b/tests/measurements/legacy/test_purity_measurement_legacy.py
new file mode 100644
index 00000000000..7da12fdbf0c
--- /dev/null
+++ b/tests/measurements/legacy/test_purity_measurement_legacy.py
@@ -0,0 +1,303 @@
+# Copyright 2018-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 the purity measurement process"""
+
+import pytest
+
+import numpy as np
+import pennylane as qml
+
+from pennylane.measurements import Shots
+
+# pylint: disable=too-many-arguments
+
+
+def expected_purity_ising_xx(param):
+    """Returns the analytical purity for subsystems of the IsingXX"""
+
+    eig_1 = (1 + np.sqrt(1 - 4 * np.cos(param / 2) ** 2 * np.sin(param / 2) ** 2)) / 2
+    eig_2 = (1 - np.sqrt(1 - 4 * np.cos(param / 2) ** 2 * np.sin(param / 2) ** 2)) / 2
+    return eig_1**2 + eig_2**2
+
+
+def expected_purity_grad_ising_xx(param):
+    """The analytic gradient purity for the IsingXX"""
+
+    eig_1 = (1 + np.sqrt(1 - 4 * np.cos(param / 2) ** 2 * np.sin(param / 2) ** 2)) / 2
+    eig_2 = (1 - np.sqrt(1 - 4 * np.cos(param / 2) ** 2 * np.sin(param / 2) ** 2)) / 2
+    grad_expected_purity = (
+        2
+        * eig_1
+        * (np.sin(param / 2) ** 3 * np.cos(param / 2) - np.sin(param / 2) * np.cos(param / 2) ** 3)
+        / np.sqrt(1 - 4 * np.sin(param / 2) ** 2 * np.cos(param / 2) ** 2)
+    ) + (
+        2
+        * eig_2
+        * (
+            np.sin(param / 2)
+            * np.cos(param / 2)
+            * (np.cos(param / 2) ** 2 - np.sin(param / 2) ** 2)
+        )
+        / np.sqrt(1 - 4 * np.sin(param / 2) ** 2 * np.cos(param / 2) ** 2)
+    )
+    return grad_expected_purity
+
+
+@pytest.mark.parametrize("shots, shape", [(None, ()), (10, ()), ((1, 10), ((), ()))])
+def test_shape_new(shots, shape):
+    """Test the ``shape_new`` method."""
+    meas = qml.purity(wires=0)
+    dev = qml.device("default.qubit.legacy", wires=1, shots=shots)
+    assert meas.shape(dev, Shots(shots)) == shape
+
+
+class TestPurityIntegration:
+    """Test the purity meausrement with qnodes and devices."""
+
+    diff_methods = ["backprop", "finite-diff"]
+
+    parameters = np.linspace(0, 2 * np.pi, 3)
+    probs = np.array([0.001, 0.01, 0.1, 0.2])
+
+    wires_list = [([0], True), ([1], True), ([0, 1], False)]
+
+    @pytest.mark.parametrize("param", parameters)
+    @pytest.mark.parametrize("wires,is_partial", wires_list)
+    def test_IsingXX_qnode_purity(self, param, wires, is_partial):
+        """Tests purity for a qnode"""
+
+        dev = qml.device("default.qubit.legacy", wires=2)
+
+        @qml.qnode(dev)
+        def circuit(x):
+            qml.IsingXX(x, wires=[0, 1])
+            return qml.purity(wires=wires)
+
+        purity = circuit(param)
+        expected_purity = expected_purity_ising_xx(param) if is_partial else 1
+        assert qml.math.allclose(purity, expected_purity)
+
+    @pytest.mark.parametrize("param", parameters)
+    @pytest.mark.parametrize("wires,is_partial", wires_list)
+    @pytest.mark.parametrize("diff_method", diff_methods)
+    def test_IsingXX_qnode_purity_grad(self, param, wires, is_partial, diff_method):
+        """Tests purity for a qnode"""
+
+        dev = qml.device("default.qubit.legacy", wires=2)
+
+        @qml.qnode(dev, diff_method=diff_method)
+        def circuit(x):
+            qml.IsingXX(x, wires=[0, 1])
+            return qml.purity(wires=wires)
+
+        grad_purity = qml.grad(circuit)(param)
+        expected_grad = expected_purity_grad_ising_xx(param) if is_partial else 0
+        assert qml.math.allclose(grad_purity, expected_grad, rtol=1e-04, atol=1e-05)
+
+    @pytest.mark.jax
+    @pytest.mark.parametrize("param", parameters)
+    @pytest.mark.parametrize("wires,is_partial", wires_list)
+    @pytest.mark.parametrize("interface", ["jax"])
+    def test_IsingXX_qnode_purity_jax(self, param, wires, is_partial, interface):
+        """Test purity for a QNode with jax interface."""
+
+        import jax.numpy as jnp
+
+        dev = qml.device("default.qubit.legacy", wires=2)
+
+        @qml.qnode(dev, interface=interface)
+        def circuit(x):
+            qml.IsingXX(x, wires=[0, 1])
+            return qml.purity(wires=wires)
+
+        purity = circuit(jnp.array(param))
+        expected_purity = expected_purity_ising_xx(param) if is_partial else 1
+        assert qml.math.allclose(purity, expected_purity)
+
+    @pytest.mark.jax
+    @pytest.mark.parametrize("param", parameters)
+    @pytest.mark.parametrize("wires,is_partial", wires_list)
+    @pytest.mark.parametrize("diff_method", diff_methods)
+    @pytest.mark.parametrize("interface", ["jax"])
+    def test_IsingXX_qnode_purity_grad_jax(self, param, wires, is_partial, diff_method, interface):
+        """Test purity for a QNode gradient with Jax."""
+
+        import jax
+
+        dev = qml.device("default.qubit.legacy", wires=2)
+
+        @qml.qnode(dev, interface=interface, diff_method=diff_method)
+        def circuit(x):
+            qml.IsingXX(x, wires=[0, 1])
+            return qml.purity(wires=wires)
+
+        grad_purity = jax.grad(circuit)(jax.numpy.array(param))
+        grad_expected_purity = expected_purity_grad_ising_xx(param) if is_partial else 0
+
+        assert qml.math.allclose(grad_purity, grad_expected_purity, rtol=1e-04, atol=1e-05)
+
+    @pytest.mark.jax
+    @pytest.mark.parametrize("param", parameters)
+    @pytest.mark.parametrize("wires,is_partial", wires_list)
+    @pytest.mark.parametrize("interface", ["jax-jit"])
+    def test_IsingXX_qnode_purity_jax_jit(self, param, wires, is_partial, interface):
+        """Test purity for a QNode with jax interface."""
+
+        import jax
+        import jax.numpy as jnp
+
+        dev = qml.device("default.qubit.legacy", wires=2)
+
+        @qml.qnode(dev, interface=interface)
+        def circuit(x):
+            qml.IsingXX(x, wires=[0, 1])
+            return qml.purity(wires=wires)
+
+        purity = jax.jit(circuit)(jnp.array(param))
+        expected_purity = expected_purity_ising_xx(param) if is_partial else 1
+        assert qml.math.allclose(purity, expected_purity)
+
+    @pytest.mark.jax
+    @pytest.mark.parametrize("param", parameters)
+    @pytest.mark.parametrize("wires,is_partial", wires_list)
+    @pytest.mark.parametrize("diff_method", diff_methods)
+    @pytest.mark.parametrize("interface", ["jax-jit"])
+    def test_IsingXX_qnode_purity_grad_jax_jit(
+        self, param, wires, is_partial, diff_method, interface
+    ):
+        """Test purity for a QNode gradient with Jax."""
+
+        import jax
+
+        dev = qml.device("default.qubit.legacy", wires=2)
+
+        @qml.qnode(dev, interface=interface, diff_method=diff_method)
+        def circuit(x):
+            qml.IsingXX(x, wires=[0, 1])
+            return qml.purity(wires=wires)
+
+        grad_purity = jax.jit(jax.grad(circuit))(jax.numpy.array(param))
+        grad_expected_purity = expected_purity_grad_ising_xx(param) if is_partial else 0
+
+        assert qml.math.allclose(grad_purity, grad_expected_purity, rtol=1e-04, atol=1e-05)
+
+    @pytest.mark.torch
+    @pytest.mark.parametrize("param", parameters)
+    @pytest.mark.parametrize("wires,is_partial", wires_list)
+    @pytest.mark.parametrize("interface", ["torch"])
+    def test_IsingXX_qnode_purity_torch(self, param, wires, is_partial, interface):
+        """Tests purity for a qnode"""
+
+        import torch
+
+        dev = qml.device("default.qubit.legacy", wires=2)
+
+        @qml.qnode(dev, interface=interface)
+        def circuit(x):
+            qml.IsingXX(x, wires=[0, 1])
+            return qml.purity(wires=wires)
+
+        purity = circuit(torch.tensor(param))
+        expected_purity = expected_purity_ising_xx(param) if is_partial else 1
+        assert qml.math.allclose(purity, expected_purity)
+
+    @pytest.mark.torch
+    @pytest.mark.parametrize("param", parameters)
+    @pytest.mark.parametrize("wires,is_partial", wires_list)
+    @pytest.mark.parametrize("diff_method", diff_methods)
+    @pytest.mark.parametrize("interface", ["torch"])
+    def test_IsingXX_qnode_purity_grad_torch(
+        self, param, wires, is_partial, diff_method, interface
+    ):
+        """Test purity for a QNode gradient with torch."""
+
+        import torch
+
+        dev = qml.device("default.qubit.legacy", wires=2)
+
+        @qml.qnode(dev, interface=interface, diff_method=diff_method)
+        def circuit(x):
+            qml.IsingXX(x, wires=[0, 1])
+            return qml.purity(wires=wires)
+
+        expected_grad = expected_purity_grad_ising_xx(param) if is_partial else 0
+
+        param = torch.tensor(param, dtype=torch.float64, requires_grad=True)
+        purity = circuit(param)
+        purity.backward()  # pylint: disable=no-member
+        grad_purity = param.grad
+
+        assert qml.math.allclose(grad_purity, expected_grad, rtol=1e-04, atol=1e-05)
+
+    @pytest.mark.tf
+    @pytest.mark.parametrize("param", parameters)
+    @pytest.mark.parametrize("wires,is_partial", wires_list)
+    @pytest.mark.parametrize("interface", ["tf"])
+    def test_IsingXX_qnode_purity_tf(self, param, wires, is_partial, interface):
+        """Tests purity for a qnode"""
+
+        import tensorflow as tf
+
+        dev = qml.device("default.qubit.legacy", wires=2)
+
+        @qml.qnode(dev, interface=interface)
+        def circuit(x):
+            qml.IsingXX(x, wires=[0, 1])
+            return qml.purity(wires=wires)
+
+        purity = circuit(tf.Variable(param))
+        expected_purity = expected_purity_ising_xx(param) if is_partial else 1
+        assert qml.math.allclose(purity, expected_purity)
+
+    @pytest.mark.tf
+    @pytest.mark.parametrize("param", parameters)
+    @pytest.mark.parametrize("wires,is_partial", wires_list)
+    @pytest.mark.parametrize("diff_method", diff_methods)
+    @pytest.mark.parametrize("interface", ["tf"])
+    def test_IsingXX_qnode_purity_grad_tf(self, param, wires, is_partial, diff_method, interface):
+        """Test purity for a QNode gradient with tf."""
+
+        import tensorflow as tf
+
+        dev = qml.device("default.qubit.legacy", wires=2)
+
+        @qml.qnode(dev, interface=interface, diff_method=diff_method)
+        def circuit(x):
+            qml.IsingXX(x, wires=[0, 1])
+            return qml.purity(wires=wires)
+
+        grad_expected_purity = expected_purity_grad_ising_xx(param) if is_partial else 0
+
+        param = tf.Variable(param)
+        with tf.GradientTape() as tape:
+            purity = circuit(param)
+
+        grad_purity = tape.gradient(purity, param)
+
+        assert qml.math.allclose(grad_purity, grad_expected_purity, rtol=1e-04, atol=1e-05)
+
+    @pytest.mark.parametrize("param", parameters)
+    def test_qnode_entropy_custom_wires(self, param):
+        """Test that purity can be returned with custom wires."""
+
+        dev = qml.device("default.qubit.legacy", wires=["a", 1])
+
+        @qml.qnode(dev)
+        def circuit(x):
+            qml.IsingXX(x, wires=["a", 1])
+            return qml.purity(wires=["a"])
+
+        purity = circuit(param)
+        expected_purity = expected_purity_ising_xx(param)
+        assert qml.math.allclose(purity, expected_purity)
diff --git a/tests/measurements/legacy/test_sample_legacy.py b/tests/measurements/legacy/test_sample_legacy.py
new file mode 100644
index 00000000000..46fbf381a8a
--- /dev/null
+++ b/tests/measurements/legacy/test_sample_legacy.py
@@ -0,0 +1,418 @@
+# Copyright 2018-2020 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 sample module"""
+import numpy as np
+import pytest
+
+import pennylane as qml
+from pennylane.measurements import MeasurementShapeError, Sample, Shots
+from pennylane.operation import Operator
+
+# pylint: disable=protected-access, no-member
+
+
+# TODO: Remove this when new CustomMP are the default
+def custom_measurement_process(device, spy):
+    assert len(spy.call_args_list) > 0  # make sure method is mocked properly
+
+    samples = device._samples
+    call_args_list = list(spy.call_args_list)
+    for call_args in call_args_list:
+        meas = call_args.args[1]
+        shot_range, bin_size = (call_args.kwargs["shot_range"], call_args.kwargs["bin_size"])
+        if isinstance(meas, Operator):
+            meas = qml.sample(op=meas)
+        assert qml.math.allequal(
+            device.sample(call_args.args[1], **call_args.kwargs),
+            meas.process_samples(
+                samples=samples,
+                wire_order=device.wires,
+                shot_range=shot_range,
+                bin_size=bin_size,
+            ),
+        )
+
+
+class TestSample:
+    """Tests for the sample function"""
+
+    @pytest.mark.parametrize("n_sample", (1, 10))
+    def test_sample_dimension(self, mocker, n_sample):
+        """Test that the sample function outputs samples of the right size"""
+
+        dev = qml.device("default.qubit.legacy", wires=2, shots=n_sample)
+        spy = mocker.spy(qml.QubitDevice, "sample")
+
+        @qml.qnode(dev)
+        def circuit():
+            qml.RX(0.54, wires=0)
+            return qml.sample(qml.PauliZ(0)), qml.sample(qml.PauliX(1))
+
+        output = circuit()
+
+        assert len(output) == 2
+        assert circuit._qfunc_output[0].shape(dev, Shots(n_sample)) == (
+            (n_sample,) if not n_sample == 1 else ()
+        )
+        assert circuit._qfunc_output[1].shape(dev, Shots(n_sample)) == (
+            (n_sample,) if not n_sample == 1 else ()
+        )
+
+        custom_measurement_process(dev, spy)
+
+    @pytest.mark.filterwarnings("ignore:Creating an ndarray from ragged nested sequences")
+    def test_sample_combination(self, mocker):
+        """Test the output of combining expval, var and sample"""
+        n_sample = 10
+
+        dev = qml.device("default.qubit.legacy", wires=3, shots=n_sample)
+        spy = mocker.spy(qml.QubitDevice, "sample")
+
+        @qml.qnode(dev, diff_method="parameter-shift")
+        def circuit():
+            qml.RX(0.54, wires=0)
+
+            return qml.sample(qml.PauliZ(0)), qml.expval(qml.PauliX(1)), qml.var(qml.PauliY(2))
+
+        result = circuit()
+
+        assert len(result) == 3
+        assert np.array_equal(result[0].shape, (n_sample,))
+        assert circuit._qfunc_output[0].shape(dev, Shots(n_sample)) == (n_sample,)
+        assert isinstance(result[1], np.ndarray)
+        assert isinstance(result[2], np.ndarray)
+
+        custom_measurement_process(dev, spy)
+
+    def test_single_wire_sample(self, mocker):
+        """Test the return type and shape of sampling a single wire"""
+        n_sample = 10
+
+        dev = qml.device("default.qubit.legacy", wires=1, shots=n_sample)
+        spy = mocker.spy(qml.QubitDevice, "sample")
+
+        @qml.qnode(dev)
+        def circuit():
+            qml.RX(0.54, wires=0)
+            return qml.sample(qml.PauliZ(0))
+
+        result = circuit()
+
+        assert isinstance(result, np.ndarray)
+        assert np.array_equal(result.shape, (n_sample,))
+        assert circuit._qfunc_output.shape(dev, Shots(n_sample)) == (n_sample,)
+
+        custom_measurement_process(dev, spy)
+
+    def test_multi_wire_sample_regular_shape(self, mocker):
+        """Test the return type and shape of sampling multiple wires
+        where a rectangular array is expected"""
+        n_sample = 10
+
+        dev = qml.device("default.qubit.legacy", wires=3, shots=n_sample)
+        spy = mocker.spy(qml.QubitDevice, "sample")
+
+        @qml.qnode(dev)
+        def circuit():
+            return qml.sample(qml.PauliZ(0)), qml.sample(qml.PauliZ(1)), qml.sample(qml.PauliZ(2))
+
+        result = circuit()
+
+        assert circuit._qfunc_output[0].shape(dev, Shots(n_sample)) == (n_sample,)
+        assert circuit._qfunc_output[1].shape(dev, Shots(n_sample)) == (n_sample,)
+        assert circuit._qfunc_output[2].shape(dev, Shots(n_sample)) == (n_sample,)
+
+        # If all the dimensions are equal the result will end up to be a proper rectangular array
+        assert isinstance(result, tuple)
+        assert len(result) == 3
+        assert result[0].dtype == np.dtype("int")
+
+        custom_measurement_process(dev, spy)
+
+    @pytest.mark.filterwarnings("ignore:Creating an ndarray from ragged nested sequences")
+    def test_sample_output_type_in_combination(self, mocker):
+        """Test the return type and shape of sampling multiple works
+        in combination with expvals and vars"""
+        n_sample = 10
+
+        dev = qml.device("default.qubit.legacy", wires=3, shots=n_sample)
+        spy = mocker.spy(qml.QubitDevice, "sample")
+
+        @qml.qnode(dev, diff_method="parameter-shift")
+        def circuit():
+            return qml.expval(qml.PauliZ(0)), qml.var(qml.PauliZ(1)), qml.sample(qml.PauliZ(2))
+
+        result = circuit()
+
+        # If all the dimensions are equal the result will end up to be a proper rectangular array
+        assert len(result) == 3
+        assert isinstance(result[0], np.ndarray)
+        assert isinstance(result[1], np.ndarray)
+        assert result[2].dtype == np.dtype("int")
+        assert np.array_equal(result[2].shape, (n_sample,))
+
+        custom_measurement_process(dev, spy)
+
+    def test_not_an_observable(self, mocker):
+        """Test that a UserWarning is raised if the provided
+        argument might not be hermitian."""
+        dev = qml.device("default.qubit.legacy", wires=2, shots=10)
+        spy = mocker.spy(qml.QubitDevice, "sample")
+
+        @qml.qnode(dev)
+        def circuit():
+            qml.RX(0.52, wires=0)
+            return qml.sample(qml.prod(qml.PauliX(0), qml.PauliZ(0)))
+
+        with pytest.warns(UserWarning, match="Prod might not be hermitian."):
+            _ = circuit()
+
+        custom_measurement_process(dev, spy)
+
+    def test_observable_return_type_is_sample(self, mocker):
+        """Test that the return type of the observable is :attr:`ObservableReturnTypes.Sample`"""
+        n_shots = 10
+        dev = qml.device("default.qubit.legacy", wires=1, shots=n_shots)
+        spy = mocker.spy(qml.QubitDevice, "sample")
+
+        @qml.qnode(dev)
+        def circuit():
+            res = qml.sample(qml.PauliZ(0))
+            assert res.return_type is Sample
+            return res
+
+        circuit()
+
+        custom_measurement_process(dev, spy)
+
+    def test_providing_observable_and_wires(self):
+        """Test that a ValueError is raised if both an observable is provided and wires are specified"""
+        dev = qml.device("default.qubit.legacy", wires=2)
+
+        @qml.qnode(dev)
+        def circuit():
+            qml.Hadamard(wires=0)
+            return qml.sample(qml.PauliZ(0), wires=[0, 1])
+
+        with pytest.raises(
+            ValueError,
+            match="Cannot specify the wires to sample if an observable is provided."
+            " The wires to sample will be determined directly from the observable.",
+        ):
+            _ = circuit()
+
+    @pytest.mark.parametrize("shots", [5, [5, 5]])
+    @pytest.mark.parametrize("phi", np.arange(0, 2 * np.pi, np.pi / 2))
+    def test_observable_is_measurement_value(self, shots, phi, mocker):
+        """Test that expectation values for mid-circuit measurement values
+        are correct for a single measurement value."""
+        dev = qml.device("default.qubit", wires=2, shots=shots)
+
+        @qml.qnode(dev)
+        def circuit(phi):
+            qml.RX(phi, 0)
+            m0 = qml.measure(0)
+            return qml.sample(m0)
+
+        new_dev = circuit.device
+        spy = mocker.spy(qml.QubitDevice, "sample")
+
+        res = circuit(phi)
+
+        if isinstance(shots, list):
+            assert len(res) == len(shots)
+            assert all(r.shape == (s,) for r, s in zip(res, shots))
+        else:
+            assert res.shape == (shots,)
+        custom_measurement_process(new_dev, spy)
+
+    def test_providing_no_observable_and_no_wires(self, mocker):
+        """Test that we can provide no observable and no wires to sample function"""
+        dev = qml.device("default.qubit.legacy", wires=2, shots=1000)
+        spy = mocker.spy(qml.QubitDevice, "sample")
+
+        @qml.qnode(dev)
+        def circuit():
+            qml.Hadamard(wires=0)
+            res = qml.sample()
+            assert res.obs is None
+            assert res.wires == qml.wires.Wires([])
+            return res
+
+        circuit()
+
+        custom_measurement_process(dev, spy)
+
+    def test_providing_no_observable_and_no_wires_shot_vector(self, mocker):
+        """Test that we can provide no observable and no wires to sample
+        function when using a shot vector"""
+        num_wires = 2
+
+        shots1 = 1
+        shots2 = 10
+        shots3 = 1000
+        dev = qml.device("default.qubit.legacy", wires=num_wires, shots=[shots1, shots2, shots3])
+        spy = mocker.spy(qml.QubitDevice, "sample")
+
+        @qml.qnode(dev)
+        def circuit():
+            qml.Hadamard(wires=0)
+            qml.CNOT(wires=[0, 1])
+            return qml.sample()
+
+        res = circuit()
+
+        assert isinstance(res, tuple)
+
+        expected_shapes = [(num_wires,), (shots2, num_wires), (shots3, num_wires)]
+        assert len(res) == len(expected_shapes)
+        assert all(r.shape == exp_shape for r, exp_shape in zip(res, expected_shapes))
+
+        # assert first wire is always the same as second
+        # pylint: disable=unsubscriptable-object
+        assert np.all(res[0][0] == res[0][1])
+        assert np.all(res[1][:, 0] == res[1][:, 1])
+        assert np.all(res[2][:, 0] == res[2][:, 1])
+
+        custom_measurement_process(dev, spy)
+
+    def test_providing_no_observable_and_wires(self, mocker):
+        """Test that we can provide no observable but specify wires to the sample function"""
+        wires = [0, 2]
+        wires_obj = qml.wires.Wires(wires)
+        dev = qml.device("default.qubit.legacy", wires=3, shots=1000)
+        spy = mocker.spy(qml.QubitDevice, "sample")
+
+        @qml.qnode(dev)
+        def circuit():
+            qml.Hadamard(wires=0)
+            res = qml.sample(wires=wires)
+
+            assert res.obs is None
+            assert res.wires == wires_obj
+            return res
+
+        circuit()
+
+        custom_measurement_process(dev, spy)
+
+    def test_shape_no_shots_error(self):
+        """Test that the appropriate error is raised with no shots are specified"""
+        dev = qml.device("default.qubit.legacy", wires=2, shots=None)
+        shots = Shots(None)
+        mp = qml.sample()
+
+        with pytest.raises(
+            MeasurementShapeError, match="Shots are required to obtain the shape of the measurement"
+        ):
+            _ = mp.shape(dev, shots)
+
+    @pytest.mark.parametrize(
+        "obs",
+        [
+            None,
+            qml.PauliZ(0),
+            qml.Hermitian(np.diag([1, 2]), 0),
+            qml.Hermitian(np.diag([1.0, 2.0]), 0),
+        ],
+    )
+    def test_shape(self, obs):
+        """Test that the shape is correct."""
+        shots = 10
+        dev = qml.device("default.qubit.legacy", wires=3, shots=shots)
+        res = qml.sample(obs) if obs is not None else qml.sample()
+        expected = (shots,) if obs is not None else (shots, 3)
+        assert res.shape(dev, Shots(shots)) == expected
+
+    @pytest.mark.parametrize("n_samples", (1, 10))
+    def test_shape_wires(self, n_samples):
+        """Test that the shape is correct when wires are provided."""
+        dev = qml.device("default.qubit.legacy", wires=3, shots=n_samples)
+        mp = qml.sample(wires=(0, 1))
+        assert mp.shape(dev, Shots(n_samples)) == (n_samples, 2) if n_samples != 1 else (2,)
+
+    @pytest.mark.parametrize(
+        "obs",
+        [
+            None,
+            qml.PauliZ(0),
+            qml.Hermitian(np.diag([1, 2]), 0),
+            qml.Hermitian(np.diag([1.0, 2.0]), 0),
+        ],
+    )
+    def test_shape_shot_vector(self, obs):
+        """Test that the shape is correct with the shot vector too."""
+        shot_vector = (1, 2, 3)
+        dev = qml.device("default.qubit.legacy", wires=3, shots=shot_vector)
+        res = qml.sample(obs) if obs is not None else qml.sample()
+        expected = ((), (2,), (3,)) if obs is not None else ((3,), (2, 3), (3, 3))
+        assert res.shape(dev, Shots(shot_vector)) == expected
+
+    def test_shape_shot_vector_obs(self):
+        """Test that the shape is correct with the shot vector and a observable too."""
+        shot_vec = (2, 2)
+        dev = qml.device("default.qubit.legacy", wires=3, shots=shot_vec)
+
+        @qml.qnode(dev)
+        def circuit():
+            qml.Hadamard(wires=0)
+            qml.PauliZ(0)
+            return qml.sample(qml.PauliZ(0))
+
+        binned_samples = circuit()
+
+        assert isinstance(binned_samples, tuple)
+        assert len(binned_samples) == len(shot_vec)
+        # pylint: disable=unsubscriptable-object
+        assert binned_samples[0].shape == (shot_vec[0],)
+
+    @pytest.mark.parametrize("shots", [2, 100])
+    def test_sample_no_arguments(self, shots):
+        """Test that using ``qml.sample`` with no arguments returns the samples of all wires."""
+        dev = qml.device("default.qubit.legacy", wires=3, shots=shots)
+
+        @qml.qnode(dev)
+        def circuit():
+            return qml.sample()
+
+        res = circuit()
+
+        # pylint: disable=comparison-with-callable
+        assert res.shape == (shots, 3)
+
+
+@pytest.mark.jax
+@pytest.mark.parametrize("samples", (1, 10))
+def test_jitting_with_sampling_on_subset_of_wires(samples):
+    """Test case covering bug in Issue #3904.  Sampling should be jit-able
+    when sampling occurs on a subset of wires. The bug was occuring due an improperly
+    set shape method."""
+    import jax
+
+    jax.config.update("jax_enable_x64", True)
+
+    dev = qml.device("default.qubit.legacy", wires=3, shots=samples)
+
+    @qml.qnode(dev, interface="jax")
+    def circuit(x):
+        qml.RX(x, wires=0)
+        return qml.sample(wires=(0, 1))
+
+    results = jax.jit(circuit)(jax.numpy.array(0.123, dtype=jax.numpy.float64))
+
+    expected = (2,) if samples == 1 else (samples, 2)
+    assert results.shape == expected
+    assert (
+        circuit._qfunc_output.shape(dev, Shots(samples)) == (samples, 2) if samples != 1 else (2,)
+    )
diff --git a/tests/measurements/legacy/test_state_legacy.py b/tests/measurements/legacy/test_state_legacy.py
new file mode 100644
index 00000000000..a21cacfc9dd
--- /dev/null
+++ b/tests/measurements/legacy/test_state_legacy.py
@@ -0,0 +1,809 @@
+# Copyright 2018-2020 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 state module"""
+import numpy as np
+import pytest
+
+import pennylane as qml
+from pennylane import numpy as pnp
+from pennylane.devices import DefaultQubit
+from pennylane.measurements import (
+    State,
+    Shots,
+    density_matrix,
+    expval,
+    state,
+)
+
+
+class TestState:
+    """Tests for the state function"""
+
+    @pytest.mark.parametrize("wires", range(2, 5))
+    def test_state_shape_and_dtype(self, wires):
+        """Test that the state is of correct size and dtype for a trivial circuit"""
+
+        dev = qml.device("default.qubit.legacy", wires=wires)
+
+        @qml.qnode(dev)
+        def func():
+            return state()
+
+        state_val = func()
+        assert state_val.shape == (2**wires,)
+        assert state_val.dtype == np.complex128
+
+    def test_return_type_is_state(self):
+        """Test that the return type of the observable is State"""
+
+        dev = qml.device("default.qubit.legacy", wires=1)
+
+        @qml.qnode(dev)
+        def func():
+            qml.Hadamard(0)
+            return state()
+
+        func()
+        obs = func.qtape.observables
+        assert len(obs) == 1
+        assert obs[0].return_type is State
+
+    @pytest.mark.parametrize("wires", range(2, 5))
+    def test_state_correct_ghz(self, wires):
+        """Test that the correct state is returned when the circuit prepares a GHZ state"""
+
+        dev = qml.device("default.qubit.legacy", wires=wires)
+
+        @qml.qnode(dev)
+        def func():
+            qml.Hadamard(wires=0)
+            for i in range(wires - 1):
+                qml.CNOT(wires=[i, i + 1])
+            return state()
+
+        state_val = func()
+        assert np.allclose(np.sum(np.abs(state_val) ** 2), 1)
+        # pylint: disable=unsubscriptable-object
+        assert np.allclose(state_val[0], 1 / np.sqrt(2))
+        assert np.allclose(state_val[-1], 1 / np.sqrt(2))
+
+        assert np.allclose(state().process_state(state=dev.state, wire_order=dev.wires), state_val)
+
+    def test_return_with_other_types(self):
+        """Test that an exception is raised when a state is returned along with another return
+        type"""
+
+        dev = qml.device("default.qubit.legacy", wires=2)
+
+        @qml.qnode(dev)
+        def func():
+            qml.Hadamard(wires=0)
+            return state(), expval(qml.PauliZ(1))
+
+        with pytest.raises(
+            qml.QuantumFunctionError,
+            match="The state or density matrix cannot be returned in combination with other return types",
+        ):
+            func()
+
+    @pytest.mark.parametrize("wires", range(2, 5))
+    def test_state_equal_to_dev_state(self, wires):
+        """Test that the returned state is equal to the one stored in dev.state for a template
+        circuit"""
+
+        dev = qml.device("default.qubit.legacy", wires=wires)
+
+        weights = np.random.random(
+            qml.templates.StronglyEntanglingLayers.shape(n_layers=3, n_wires=wires)
+        )
+
+        @qml.qnode(dev)
+        def func():
+            qml.templates.StronglyEntanglingLayers(weights, wires=range(wires))
+            return state()
+
+        state_val = func()
+        assert np.allclose(state_val, func.device.state)
+
+    @pytest.mark.tf
+    def test_interface_tf(self):
+        """Test that the state correctly outputs in the tensorflow interface"""
+        import tensorflow as tf
+
+        dev = qml.device("default.qubit.legacy", wires=4)
+
+        @qml.qnode(dev, interface="tf")
+        def func():
+            for i in range(4):
+                qml.Hadamard(i)
+            return state()
+
+        state_expected = 0.25 * tf.ones(16)
+        state_val = func()
+
+        assert isinstance(state_val, tf.Tensor)
+        assert state_val.dtype == tf.complex128
+        assert np.allclose(state_expected, state_val.numpy())
+        assert state_val.shape == (16,)
+
+    @pytest.mark.torch
+    def test_interface_torch(self):
+        """Test that the state correctly outputs in the torch interface"""
+        import torch
+
+        dev = qml.device("default.qubit.legacy", wires=4)
+
+        @qml.qnode(dev, interface="torch")
+        def func():
+            for i in range(4):
+                qml.Hadamard(i)
+            return state()
+
+        state_expected = 0.25 * torch.ones(16, dtype=torch.complex128)
+        state_val = func()
+
+        assert isinstance(state_val, torch.Tensor)
+        assert state_val.dtype == torch.complex128
+        assert torch.allclose(state_expected, state_val)
+        assert state_val.shape == (16,)
+
+    @pytest.mark.autograd
+    def test_jacobian_not_supported(self):
+        """Test if an error is raised if the jacobian method is called via qml.grad"""
+        dev = qml.device("default.qubit.legacy", wires=4)
+
+        @qml.qnode(dev, diff_method="parameter-shift")
+        def func(x):
+            for i in range(4):
+                qml.RX(x, wires=i)
+            return state()
+
+        d_func = qml.jacobian(func)
+
+        with pytest.raises(
+            ValueError,
+            match=(
+                "Computing the gradient of circuits that return the state with the "
+                "parameter-shift rule gradient transform is not supported"
+            ),
+        ):
+            d_func(pnp.array(0.1, requires_grad=True))
+
+    def test_no_state_capability(self, monkeypatch):
+        """Test if an error is raised for devices that are not capable of returning the state.
+        This is tested by changing the capability of default.qubit"""
+        dev = qml.device("default.qubit.legacy", wires=1)
+        capabilities = dev.capabilities().copy()
+        capabilities["returns_state"] = False
+
+        @qml.qnode(dev)
+        def func():
+            return state()
+
+        with monkeypatch.context() as m:
+            m.setattr(DefaultQubit, "capabilities", lambda *args, **kwargs: capabilities)
+            with pytest.raises(qml.QuantumFunctionError, match="The current device is not capable"):
+                func()
+
+    def test_state_not_supported(self):
+        """Test if an error is raised for devices inheriting from the base Device class,
+        which do not currently support returning the state"""
+        dev = qml.device("default.gaussian", wires=1)
+
+        @qml.qnode(dev)
+        def func():
+            return state()
+
+        with pytest.raises(qml.QuantumFunctionError, match="Returning the state is not supported"):
+            func()
+
+    @pytest.mark.parametrize("diff_method", ["best", "finite-diff", "parameter-shift"])
+    def test_default_qubit(self, diff_method):
+        """Test that the returned state is equal to the expected returned state for all of
+        PennyLane's built in statevector devices"""
+
+        dev = qml.device("default.qubit.legacy", wires=4)
+
+        @qml.qnode(dev, diff_method=diff_method)
+        def func():
+            for i in range(4):
+                qml.Hadamard(i)
+            return state()
+
+        state_val = func()
+        state_expected = 0.25 * np.ones(16)
+
+        assert np.allclose(state_val, state_expected)
+        assert np.allclose(state_val, dev.state)
+
+    @pytest.mark.tf
+    @pytest.mark.parametrize("diff_method", ["best", "finite-diff", "parameter-shift"])
+    def test_default_qubit_tf(self, diff_method):
+        """Test that the returned state is equal to the expected returned state for all of
+        PennyLane's built in statevector devices"""
+
+        dev = qml.device("default.qubit.tf", wires=4)
+
+        @qml.qnode(dev, diff_method=diff_method)
+        def func():
+            for i in range(4):
+                qml.Hadamard(i)
+            return state()
+
+        state_val = func()
+        state_expected = 0.25 * np.ones(16)
+
+        assert np.allclose(state_val, state_expected)
+        assert np.allclose(state_val, dev.state)
+
+    @pytest.mark.autograd
+    @pytest.mark.parametrize("diff_method", ["best", "finite-diff", "parameter-shift"])
+    def test_default_qubit_autograd(self, diff_method):
+        """Test that the returned state is equal to the expected returned state for all of
+        PennyLane's built in statevector devices"""
+
+        dev = qml.device("default.qubit.autograd", wires=4)
+
+        @qml.qnode(dev, diff_method=diff_method)
+        def func():
+            for i in range(4):
+                qml.Hadamard(i)
+            return state()
+
+        state_val = func()
+        state_expected = 0.25 * np.ones(16)
+
+        assert np.allclose(state_val, state_expected)
+        assert np.allclose(state_val, dev.state)
+
+    @pytest.mark.tf
+    def test_gradient_with_passthru_tf(self):
+        """Test that the gradient of the state is accessible when using default.qubit.tf with the
+        backprop diff_method."""
+        import tensorflow as tf
+
+        dev = qml.device("default.qubit.tf", wires=1)
+
+        @qml.qnode(dev, interface="tf", diff_method="backprop")
+        def func(x):
+            qml.RY(x, wires=0)
+            return state()
+
+        x = tf.Variable(0.1, dtype=tf.float64)
+
+        with tf.GradientTape() as tape:
+            result = func(x)
+
+        grad = tape.jacobian(result, x)
+        expected = tf.stack([-0.5 * tf.sin(x / 2), 0.5 * tf.cos(x / 2)])
+        assert np.allclose(grad, expected)
+
+    @pytest.mark.autograd
+    def test_gradient_with_passthru_autograd(self):
+        """Test that the gradient of the state is accessible when using default.qubit.autograd
+        with the backprop diff_method."""
+
+        dev = qml.device("default.qubit.autograd", wires=1)
+
+        @qml.qnode(dev, interface="autograd", diff_method="backprop")
+        def func(x):
+            qml.RY(x, wires=0)
+            return state()
+
+        x = pnp.array(0.1, requires_grad=True)
+
+        def loss_fn(x):
+            res = func(x)
+            return pnp.real(res)  # This errors without the real. Likely an issue with complex
+            # numbers in autograd
+
+        d_loss_fn = qml.jacobian(loss_fn)
+
+        grad = d_loss_fn(x)
+        expected = np.array([-0.5 * np.sin(x / 2), 0.5 * np.cos(x / 2)])
+        assert np.allclose(grad, expected)
+
+    @pytest.mark.parametrize("wires", [[0, 2, 3, 1], ["a", -1, "b", 1000]])
+    def test_custom_wire_labels(self, wires):
+        """Test the state when custom wire labels are used"""
+        dev = qml.device("default.qubit.legacy", wires=wires)
+
+        @qml.qnode(dev, diff_method="parameter-shift")
+        def func():
+            for i in range(4):
+                qml.Hadamard(wires[i])
+            return state()
+
+        state_expected = 0.25 * np.ones(16)
+        state_val = func()
+
+        assert np.allclose(state_expected, state_val)
+
+    @pytest.mark.parametrize("shots", [None, 1, 10])
+    def test_shape(self, shots):
+        """Test that the shape is correct for qml.state."""
+        dev = qml.device("default.qubit.legacy", wires=3, shots=shots)
+        res = qml.state()
+        assert res.shape(dev, Shots(shots)) == (2**3,)
+
+    @pytest.mark.parametrize("s_vec", [(3, 2, 1), (1, 5, 10), (3, 1, 20)])
+    def test_shape_shot_vector(self, s_vec):
+        """Test that the shape is correct for qml.state with the shot vector too."""
+        dev = qml.device("default.qubit.legacy", wires=3, shots=s_vec)
+        res = qml.state()
+        assert res.shape(dev, Shots(s_vec)) == ((2**3,), (2**3,), (2**3,))
+
+
+class TestDensityMatrix:
+    """Tests for the density matrix function"""
+
+    # pylint: disable=too-many-public-methods
+
+    @pytest.mark.parametrize("wires", range(2, 5))
+    @pytest.mark.parametrize("dev_name", ["default.qubit.legacy", "default.mixed"])
+    def test_density_matrix_shape_and_dtype(self, dev_name, wires):
+        """Test that the density matrix is of correct size and dtype for a
+        trivial circuit"""
+
+        dev = qml.device(dev_name, wires=wires)
+
+        @qml.qnode(dev)
+        def circuit():
+            return density_matrix([0])
+
+        state_val = circuit()
+
+        assert state_val.shape == (2, 2)
+        assert state_val.dtype == np.complex128
+
+    @pytest.mark.parametrize("dev_name", ["default.qubit.legacy", "default.mixed"])
+    def test_return_type_is_state(self, dev_name):
+        """Test that the return type of the observable is State"""
+
+        dev = qml.device(dev_name, wires=2)
+
+        @qml.qnode(dev)
+        def func():
+            qml.Hadamard(0)
+            return density_matrix(0)
+
+        func()
+        obs = func.qtape.observables
+        assert len(obs) == 1
+        assert obs[0].return_type is State
+
+    @pytest.mark.torch
+    @pytest.mark.parametrize("dev_name", ["default.qubit.legacy", "default.mixed"])
+    @pytest.mark.parametrize("diff_method", [None, "backprop"])
+    def test_correct_density_matrix_torch(self, dev_name, diff_method):
+        """Test that the correct density matrix is returned using torch interface."""
+        dev = qml.device(dev_name, wires=2)
+
+        @qml.qnode(dev, interface="torch", diff_method=diff_method)
+        def func():
+            qml.Hadamard(wires=0)
+            return qml.density_matrix(wires=0)
+
+        density_mat = func()
+        expected = np.array([[0.5 + 0.0j, 0.5 + 0.0j], [0.5 + 0.0j, 0.5 + 0.0j]])
+        assert np.allclose(expected, density_mat)
+
+        dev = func.device
+
+        if dev_name != "default.mixed":
+            assert np.allclose(
+                expected,
+                qml.density_matrix(wires=0).process_state(state=dev.state, wire_order=dev.wires),
+            )
+
+    @pytest.mark.jax
+    @pytest.mark.parametrize("dev_name", ["default.qubit.legacy", "default.mixed"])
+    @pytest.mark.parametrize("diff_method", [None, "backprop"])
+    def test_correct_density_matrix_jax(self, dev_name, diff_method):
+        """Test that the correct density matrix is returned using JAX interface."""
+        dev = qml.device(dev_name, wires=2)
+
+        @qml.qnode(dev, interface="jax", diff_method=diff_method)
+        def func():
+            qml.Hadamard(wires=0)
+            return qml.density_matrix(wires=0)
+
+        density_mat = func()
+        expected = np.array([[0.5 + 0.0j, 0.5 + 0.0j], [0.5 + 0.0j, 0.5 + 0.0j]])
+
+        assert np.allclose(expected, density_mat)
+
+        if dev_name != "default.mixed":
+            assert np.allclose(
+                expected,
+                qml.density_matrix(wires=0).process_state(state=dev.state, wire_order=dev.wires),
+            )
+
+    @pytest.mark.tf
+    @pytest.mark.parametrize("dev_name", ["default.qubit.legacy", "default.mixed"])
+    @pytest.mark.parametrize("diff_method", [None, "backprop"])
+    def test_correct_density_matrix_tf(self, dev_name, diff_method):
+        """Test that the correct density matrix is returned using the TensorFlow interface."""
+        dev = qml.device(dev_name, wires=2)
+
+        @qml.qnode(dev, interface="tf", diff_method=diff_method)
+        def func():
+            qml.Hadamard(wires=0)
+            return qml.density_matrix(wires=0)
+
+        density_mat = func()
+        expected = np.array([[0.5 + 0.0j, 0.5 + 0.0j], [0.5 + 0.0j, 0.5 + 0.0j]])
+
+        assert np.allclose(expected, density_mat)
+
+        if dev_name != "default.mixed":
+            assert np.allclose(
+                expected,
+                qml.density_matrix(wires=0).process_state(state=dev.state, wire_order=dev.wires),
+            )
+
+    @pytest.mark.parametrize("dev_name", ["default.qubit.legacy", "default.mixed"])
+    def test_correct_density_matrix_product_state_first(self, dev_name):
+        """Test that the correct density matrix is returned when
+        tracing out a product state"""
+
+        dev = qml.device(dev_name, wires=2)
+
+        @qml.qnode(dev)
+        def func():
+            qml.Hadamard(wires=1)
+            qml.PauliY(wires=0)
+            return density_matrix(0)
+
+        density_first = func()
+        expected = np.array([[0.0 + 0.0j, 0.0 + 0.0j], [0.0 + 0.0j, 1.0 + 0.0j]])
+
+        assert np.allclose(expected, density_first)
+
+        if dev_name != "default.mixed":
+            assert np.allclose(
+                expected,
+                qml.density_matrix(wires=0).process_state(state=dev.state, wire_order=dev.wires),
+            )
+
+    @pytest.mark.parametrize("dev_name", ["default.qubit.legacy", "default.mixed"])
+    def test_correct_density_matrix_product_state_second(self, dev_name):
+        """Test that the correct density matrix is returned when
+        tracing out a product state"""
+
+        dev = qml.device(dev_name, wires=2)
+
+        @qml.qnode(dev)
+        def func():
+            qml.Hadamard(wires=1)
+            qml.PauliY(wires=0)
+            return density_matrix(1)
+
+        density_second = func()
+        expected = np.array([[0.5 + 0.0j, 0.5 + 0.0j], [0.5 + 0.0j, 0.5 + 0.0j]])
+        assert np.allclose(expected, density_second)
+
+        if dev_name != "default.mixed":
+            assert np.allclose(
+                expected,
+                qml.density_matrix(wires=1).process_state(state=dev.state, wire_order=dev.wires),
+            )
+
+    @pytest.mark.parametrize("dev_name", ["default.qubit.legacy", "default.mixed"])
+    @pytest.mark.parametrize("return_wire_order", ([0, 1], [1, 0]))
+    def test_correct_density_matrix_product_state_both(self, dev_name, return_wire_order):
+        """Test that the correct density matrix is returned
+        for a full product state on two wires."""
+
+        dev = qml.device(dev_name, wires=2)
+
+        @qml.qnode(dev)
+        def func():
+            qml.Hadamard(wires=1)
+            qml.PauliY(wires=0)
+            return density_matrix(return_wire_order)
+
+        density_both = func()
+        single_statevectors = [[0, 1j], [1 / np.sqrt(2), 1 / np.sqrt(2)]]
+        expected_statevector = np.kron(*[single_statevectors[w] for w in return_wire_order])
+        expected = np.outer(expected_statevector.conj(), expected_statevector)
+
+        assert np.allclose(expected, density_both)
+
+        if dev_name != "default.mixed":
+            assert np.allclose(
+                expected,
+                qml.density_matrix(wires=return_wire_order).process_state(
+                    state=dev.state, wire_order=dev.wires
+                ),
+            )
+
+    @pytest.mark.parametrize("dev_name", ["default.qubit.legacy", "default.mixed"])
+    def test_correct_density_matrix_three_wires_first_two(self, dev_name):
+        """Test that the correct density matrix is returned for an example with three wires,
+        and tracing out the third wire."""
+
+        dev = qml.device(dev_name, wires=3)
+
+        @qml.qnode(dev)
+        def func():
+            qml.Hadamard(wires=1)
+            qml.PauliY(wires=0)
+            return density_matrix([0, 1])
+
+        density_full = func()
+        expected = np.array(
+            [
+                [0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j],
+                [0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j],
+                [0.0 + 0.0j, 0.0 + 0.0j, 0.5 + 0.0j, 0.5 + 0.0j],
+                [0.0 + 0.0j, 0.0 + 0.0j, 0.5 + 0.0j, 0.5 + 0.0j],
+            ]
+        )
+        assert np.allclose(expected, density_full)
+
+        if dev_name != "default.mixed":
+            assert np.allclose(
+                expected,
+                qml.density_matrix(wires=[0, 1]).process_state(
+                    state=dev.state, wire_order=dev.wires
+                ),
+            )
+
+    @pytest.mark.parametrize("dev_name", ["default.qubit.legacy", "default.mixed"])
+    def test_correct_density_matrix_three_wires_last_two(self, dev_name):
+        """Test that the correct density matrix is returned for an example with three wires,
+        and tracing out the first wire."""
+
+        dev = qml.device(dev_name, wires=3)
+
+        @qml.qnode(dev)
+        def func():
+            qml.Hadamard(0)
+            qml.Hadamard(1)
+            qml.CNOT(wires=[1, 2])
+            return qml.density_matrix(wires=[1, 2])
+
+        density = func()
+        expected = np.array(
+            [
+                [
+                    [0.5 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j, 0.5 + 0.0j],
+                    [0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j],
+                    [0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j],
+                    [0.5 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j, 0.5 + 0.0j],
+                ]
+            ]
+        )
+
+        assert np.allclose(expected, density)
+
+        if dev_name != "default.mixed":
+            assert np.allclose(
+                expected,
+                qml.density_matrix(wires=[1, 2]).process_state(
+                    state=dev.state, wire_order=dev.wires
+                ),
+            )
+
+    @pytest.mark.parametrize("dev_name", ["default.qubit.legacy", "default.mixed"])
+    @pytest.mark.parametrize(
+        "return_wire_order", ([0], [1], [2], [0, 1], [1, 0], [0, 2], [2, 0], [1, 2, 0], [2, 1, 0])
+    )
+    def test_correct_density_matrix_three_wires_product(self, dev_name, return_wire_order):
+        """Test that the correct density matrix is returned for an example with
+        three wires and a product state, tracing out various combinations."""
+
+        dev = qml.device(dev_name, wires=3)
+
+        @qml.qnode(dev)
+        def func():
+            qml.Hadamard(0)
+            qml.PauliX(1)
+            qml.PauliZ(2)
+            return density_matrix(return_wire_order)
+
+        density_full = func()
+
+        single_states = [[1 / np.sqrt(2), 1 / np.sqrt(2)], [0, 1], [1, 0]]
+        if len(return_wire_order) == 1:
+            exp_statevector = np.array(single_states[return_wire_order[0]])
+        elif len(return_wire_order) == 2:
+            i, j = return_wire_order
+            exp_statevector = np.kron(single_states[i], single_states[j])
+        elif len(return_wire_order) == 3:
+            i, j, k = return_wire_order
+            exp_statevector = np.kron(np.kron(single_states[i], single_states[j]), single_states[k])
+
+        expected = np.outer(exp_statevector.conj(), exp_statevector)
+        assert np.allclose(expected, density_full)
+
+        if dev_name != "default.mixed":
+            assert np.allclose(
+                expected,
+                qml.density_matrix(wires=return_wire_order).process_state(
+                    state=dev.state, wire_order=dev.wires
+                ),
+            )
+
+    @pytest.mark.parametrize("dev_name", ["default.qubit.legacy", "default.mixed"])
+    def test_correct_density_matrix_mixed_state(self, dev_name):
+        """Test that the correct density matrix for an example with a mixed state"""
+
+        dev = qml.device(dev_name, wires=2)
+
+        @qml.qnode(dev)
+        def func():
+            qml.Hadamard(0)
+            qml.CNOT(wires=[0, 1])
+            return qml.density_matrix(wires=[1])
+
+        density = func()
+
+        assert np.allclose(np.array([[0.5 + 0.0j, 0.0 + 0.0j], [0.0 + 0.0j, 0.5 + 0.0j]]), density)
+
+    @pytest.mark.parametrize("dev_name", ["default.qubit.legacy", "default.mixed"])
+    def test_correct_density_matrix_all_wires(self, dev_name):
+        """Test that the correct density matrix is returned when all wires are given"""
+
+        dev = qml.device(dev_name, wires=2)
+
+        @qml.qnode(dev)
+        def func():
+            qml.Hadamard(0)
+            qml.CNOT(wires=[0, 1])
+            return qml.density_matrix(wires=[0, 1])
+
+        density = func()
+        expected = np.array(
+            [
+                [0.5 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j, 0.5 + 0.0j],
+                [0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j],
+                [0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j],
+                [0.5 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j, 0.5 + 0.0j],
+            ]
+        )
+
+        assert np.allclose(expected, density)
+
+        if dev_name != "default.mixed":
+            assert np.allclose(
+                expected,
+                qml.density_matrix(wires=[0, 1]).process_state(
+                    state=dev.state, wire_order=dev.wires
+                ),
+            )
+
+    @pytest.mark.parametrize("dev_name", ["default.qubit.legacy", "default.mixed"])
+    def test_return_with_other_types(self, dev_name):
+        """Test that an exception is raised when a state is returned along with another return
+        type"""
+
+        dev = qml.device(dev_name, wires=2)
+
+        @qml.qnode(dev)
+        def func():
+            qml.Hadamard(wires=0)
+            return density_matrix(0), expval(qml.PauliZ(1))
+
+        with pytest.raises(
+            qml.QuantumFunctionError,
+            match="The state or density matrix"
+            " cannot be returned in combination"
+            " with other return types",
+        ):
+            func()
+
+    def test_no_state_capability(self, monkeypatch):
+        """Test if an error is raised for devices that are not capable of returning
+        the density matrix. This is tested by changing the capability of default.qubit"""
+        dev = qml.device("default.qubit.legacy", wires=2)
+        capabilities = dev.capabilities().copy()
+        capabilities["returns_state"] = False
+
+        @qml.qnode(dev)
+        def func():
+            return density_matrix(0)
+
+        with monkeypatch.context() as m:
+            m.setattr(DefaultQubit, "capabilities", lambda *args, **kwargs: capabilities)
+            with pytest.raises(
+                qml.QuantumFunctionError,
+                match="The current device is not capable" " of returning the state",
+            ):
+                func()
+
+    def test_density_matrix_not_supported(self):
+        """Test if an error is raised for devices inheriting from the base Device class,
+        which do not currently support returning the state"""
+        dev = qml.device("default.gaussian", wires=2)
+
+        @qml.qnode(dev)
+        def func():
+            return density_matrix(0)
+
+        with pytest.raises(qml.QuantumFunctionError, match="Returning the state is not supported"):
+            func()
+
+    @pytest.mark.parametrize("wires", [[0, 2], ["a", -1]])
+    @pytest.mark.parametrize("dev_name", ["default.qubit.legacy", "default.mixed"])
+    def test_custom_wire_labels(self, wires, dev_name):
+        """Test that the correct density matrix for an example with a mixed
+        state when using custom wires"""
+
+        dev = qml.device(dev_name, wires=wires)
+
+        @qml.qnode(dev)
+        def func():
+            qml.Hadamard(wires[0])
+            qml.CNOT(wires=[wires[0], wires[1]])
+            return qml.density_matrix(wires=wires[1])
+
+        density = func()
+        expected = np.array([[0.5 + 0.0j, 0.0 + 0.0j], [0.0 + 0.0j, 0.5 + 0.0j]])
+
+        assert np.allclose(expected, density)
+
+        if dev_name != "default.mixed":
+            assert np.allclose(
+                expected,
+                qml.density_matrix(wires=wires[1]).process_state(
+                    state=dev.state, wire_order=dev.wires
+                ),
+            )
+
+    @pytest.mark.parametrize("wires", [[3, 1], ["b", 1000]])
+    @pytest.mark.parametrize("dev_name", ["default.qubit.legacy", "default.mixed"])
+    def test_custom_wire_labels_all_wires(self, wires, dev_name):
+        """Test that the correct density matrix for an example with a mixed
+        state when using custom wires"""
+        dev = qml.device(dev_name, wires=wires)
+
+        @qml.qnode(dev)
+        def func():
+            qml.Hadamard(wires[0])
+            qml.CNOT(wires=[wires[0], wires[1]])
+            return qml.density_matrix(wires=[wires[0], wires[1]])
+
+        density = func()
+
+        assert np.allclose(
+            np.array(
+                [
+                    [0.5 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j, 0.5 + 0.0j],
+                    [0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j],
+                    [0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j],
+                    [0.5 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j, 0.5 + 0.0j],
+                ]
+            ),
+            density,
+        )
+
+    @pytest.mark.parametrize("shots", [None, 1, 10])
+    def test_shape(self, shots):
+        """Test that the shape is correct for qml.density_matrix."""
+        dev = qml.device("default.qubit.legacy", wires=3, shots=shots)
+        res = qml.density_matrix(wires=[0, 1])
+        assert res.shape(dev, Shots(shots)) == (2**2, 2**2)
+
+    @pytest.mark.parametrize("s_vec", [(3, 2, 1), (1, 5, 10), (3, 1, 20)])
+    def test_shape_shot_vector(self, s_vec):
+        """Test that the shape is correct for qml.density_matrix with the shot vector too."""
+        dev = qml.device("default.qubit.legacy", wires=3, shots=s_vec)
+        res = qml.density_matrix(wires=[0, 1])
+        assert res.shape(dev, Shots(s_vec)) == (
+            (2**2, 2**2),
+            (2**2, 2**2),
+            (2**2, 2**2),
+        )
diff --git a/tests/measurements/legacy/test_var_legacy.py b/tests/measurements/legacy/test_var_legacy.py
new file mode 100644
index 00000000000..6580c3f538d
--- /dev/null
+++ b/tests/measurements/legacy/test_var_legacy.py
@@ -0,0 +1,200 @@
+# Copyright 2018-2020 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 var module"""
+import numpy as np
+import pytest
+
+import pennylane as qml
+from pennylane.measurements import Variance, Shots
+
+
+# TODO: Remove this when new CustomMP are the default
+def custom_measurement_process(device, spy):
+    assert len(spy.call_args_list) > 0  # make sure method is mocked properly
+
+    # pylint: disable=protected-access
+    samples = device._samples
+    state = device._state
+    call_args_list = list(spy.call_args_list)
+    for call_args in call_args_list:
+        obs = call_args.args[1]
+        shot_range, bin_size = (
+            call_args.kwargs["shot_range"],
+            call_args.kwargs["bin_size"],
+        )
+        meas = qml.var(op=obs)
+        old_res = device.var(obs, shot_range, bin_size)
+        if samples is not None:
+            new_res = meas.process_samples(
+                samples=samples, wire_order=device.wires, shot_range=shot_range, bin_size=bin_size
+            )
+        else:
+            new_res = meas.process_state(state=state, wire_order=device.wires)
+        assert qml.math.allclose(old_res, new_res)
+
+
+class TestVar:
+    """Tests for the var function"""
+
+    @pytest.mark.parametrize("shots", [None, 10000, [10000, 10000]])
+    @pytest.mark.parametrize("r_dtype", [np.float32, np.float64])
+    def test_value(self, tol, r_dtype, mocker, shots):
+        """Test that the var function works"""
+        dev = qml.device("default.qubit.legacy", wires=2, shots=shots)
+        spy = mocker.spy(qml.QubitDevice, "var")
+        dev.R_DTYPE = r_dtype
+
+        @qml.qnode(dev, diff_method="parameter-shift")
+        def circuit(x):
+            qml.RX(x, wires=0)
+            return qml.var(qml.PauliZ(0))
+
+        x = 0.54
+        res = circuit(x)
+        expected = [np.sin(x) ** 2, np.sin(x) ** 2] if isinstance(shots, list) else np.sin(x) ** 2
+        atol = tol if shots is None else 0.05
+        rtol = 0 if shots is None else 0.05
+
+        assert np.allclose(res, expected, atol=atol, rtol=rtol)
+        # pylint: disable=no-member, unsubscriptable-object
+        if isinstance(res, tuple):
+            assert res[0].dtype == r_dtype
+            assert res[1].dtype == r_dtype
+        else:
+            assert res.dtype == r_dtype
+
+        custom_measurement_process(dev, spy)
+
+    def test_not_an_observable(self, mocker):
+        """Test that a UserWarning is raised if the provided
+        argument might not be hermitian."""
+        dev = qml.device("default.qubit.legacy", wires=2)
+        spy = mocker.spy(qml.QubitDevice, "var")
+
+        @qml.qnode(dev)
+        def circuit():
+            qml.RX(0.52, wires=0)
+            return qml.var(qml.prod(qml.PauliX(0), qml.PauliZ(0)))
+
+        with pytest.warns(UserWarning, match="Prod might not be hermitian."):
+            _ = circuit()
+
+        custom_measurement_process(dev, spy)
+
+    def test_observable_return_type_is_variance(self, mocker):
+        """Test that the return type of the observable is :attr:`ObservableReturnTypes.Variance`"""
+        dev = qml.device("default.qubit.legacy", wires=2)
+        spy = mocker.spy(qml.QubitDevice, "var")
+
+        @qml.qnode(dev)
+        def circuit():
+            res = qml.var(qml.PauliZ(0))
+            assert res.return_type is Variance
+            return res
+
+        circuit()
+
+        custom_measurement_process(dev, spy)
+
+    @pytest.mark.parametrize("shots", [None, 10000, [10000, 10000]])
+    @pytest.mark.parametrize("phi", np.arange(0, 2 * np.pi, np.pi / 3))
+    def test_observable_is_measurement_value(
+        self, shots, phi, mocker, tol, tol_stochastic
+    ):  # pylint: disable=too-many-arguments
+        """Test that variances for mid-circuit measurement values
+        are correct for a single measurement value."""
+        dev = qml.device("default.qubit", wires=2, shots=shots)
+
+        @qml.qnode(dev)
+        def circuit(phi):
+            qml.RX(phi, 0)
+            m0 = qml.measure(0)
+            return qml.var(m0)
+
+        new_dev = circuit.device
+        spy = mocker.spy(qml.QubitDevice, "var")
+
+        res = circuit(phi)
+
+        atol = tol if shots is None else tol_stochastic
+        expected = np.sin(phi / 2) ** 2 - np.sin(phi / 2) ** 4
+        assert np.allclose(np.array(res), expected, atol=atol, rtol=0)
+        custom_measurement_process(new_dev, spy)
+
+    @pytest.mark.parametrize(
+        "obs",
+        [qml.PauliZ(0), qml.Hermitian(np.diag([1, 2]), 0), qml.Hermitian(np.diag([1.0, 2.0]), 0)],
+    )
+    def test_shape(self, obs):
+        """Test that the shape is correct."""
+        dev = qml.device("default.qubit.legacy", wires=1)
+        res = qml.var(obs)
+        # pylint: disable=use-implicit-booleaness-not-comparison
+        assert res.shape(dev, Shots(None)) == ()
+        assert res.shape(dev, Shots(100)) == ()
+
+    @pytest.mark.parametrize(
+        "obs",
+        [qml.PauliZ(0), qml.Hermitian(np.diag([1, 2]), 0), qml.Hermitian(np.diag([1.0, 2.0]), 0)],
+    )
+    def test_shape_shot_vector(self, obs):
+        """Test that the shape is correct with the shot vector too."""
+        res = qml.var(obs)
+        shot_vector = (1, 2, 3)
+        dev = qml.device("default.qubit.legacy", wires=3, shots=shot_vector)
+        assert res.shape(dev, Shots(shot_vector)) == ((), (), ())
+
+    @pytest.mark.parametrize("state", [np.array([0, 0, 0]), np.array([1, 0, 0, 0, 0, 0, 0, 0])])
+    @pytest.mark.parametrize("shots", [None, 1000, [1000, 10000]])
+    def test_projector_var(self, state, shots, mocker):
+        """Tests that the variance of a ``Projector`` object is computed correctly."""
+        dev = qml.device("default.qubit.legacy", wires=3, shots=shots)
+        spy = mocker.spy(qml.QubitDevice, "var")
+
+        @qml.qnode(dev)
+        def circuit():
+            qml.Hadamard(0)
+            return qml.var(qml.Projector(state, wires=range(3)))
+
+        res = circuit()
+        expected = [0.25, 0.25] if isinstance(shots, list) else 0.25
+
+        assert np.allclose(res, expected, atol=0.02, rtol=0.02)
+
+        custom_measurement_process(dev, spy)
+
+    def test_permuted_wires(self, mocker):
+        """Test that the variance of an operator with permuted wires is the same."""
+        obs = qml.prod(qml.PauliZ(8), qml.s_prod(2, qml.PauliZ(10)), qml.s_prod(3, qml.PauliZ("h")))
+        obs_2 = qml.prod(
+            qml.s_prod(3, qml.PauliZ("h")), qml.PauliZ(8), qml.s_prod(2, qml.PauliZ(10))
+        )
+
+        dev = qml.device("default.qubit.legacy", wires=["h", 8, 10])
+        spy = mocker.spy(qml.QubitDevice, "var")
+
+        @qml.qnode(dev)
+        def circuit():
+            qml.RX(1.23, wires=["h"])
+            qml.RY(2.34, wires=[8])
+            return qml.var(obs)
+
+        @qml.qnode(dev)
+        def circuit2():
+            qml.RX(1.23, wires=["h"])
+            qml.RY(2.34, wires=[8])
+            return qml.var(obs_2)
+
+        assert circuit() == circuit2()
+        custom_measurement_process(dev, spy)
diff --git a/tests/measurements/legacy/test_vn_entropy_legacy.py b/tests/measurements/legacy/test_vn_entropy_legacy.py
new file mode 100644
index 00000000000..3c38417e216
--- /dev/null
+++ b/tests/measurements/legacy/test_vn_entropy_legacy.py
@@ -0,0 +1,394 @@
+# Copyright 2018-2020 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 vn_entropy module"""
+import numpy as np
+import pytest
+
+import pennylane as qml
+from pennylane.interfaces import INTERFACE_MAP
+from pennylane.measurements import Shots
+from pennylane.measurements.vn_entropy import VnEntropyMP
+
+# pylint: disable=too-many-arguments
+
+
+def expected_entropy_ising_xx(param):
+    """
+    Return the analytical entropy for the IsingXX.
+    """
+    eig_1 = (1 + np.sqrt(1 - 4 * np.cos(param / 2) ** 2 * np.sin(param / 2) ** 2)) / 2
+    eig_2 = (1 - np.sqrt(1 - 4 * np.cos(param / 2) ** 2 * np.sin(param / 2) ** 2)) / 2
+    eigs = [eig_1, eig_2]
+    eigs = [eig for eig in eigs if eig > 0]
+
+    expected_entropy = eigs * np.log(eigs)
+
+    expected_entropy = -np.sum(expected_entropy)
+    return expected_entropy
+
+
+def expected_entropy_grad_ising_xx(param):
+    """
+    Return the analytical gradient entropy for the IsingXX.
+    """
+    eig_1 = (1 + np.sqrt(1 - 4 * np.cos(param / 2) ** 2 * np.sin(param / 2) ** 2)) / 2
+    eig_2 = (1 - np.sqrt(1 - 4 * np.cos(param / 2) ** 2 * np.sin(param / 2) ** 2)) / 2
+    eigs = [eig_1, eig_2]
+    eigs = np.maximum(eigs, 1e-08)
+
+    return -(
+        (np.log(eigs[0]) + 1)
+        * (np.sin(param / 2) ** 3 * np.cos(param / 2) - np.sin(param / 2) * np.cos(param / 2) ** 3)
+        / np.sqrt(1 - 4 * np.cos(param / 2) ** 2 * np.sin(param / 2) ** 2)
+    ) - (
+        (np.log(eigs[1]) + 1)
+        * (
+            np.sin(param / 2)
+            * np.cos(param / 2)
+            * (np.cos(param / 2) ** 2 - np.sin(param / 2) ** 2)
+        )
+        / np.sqrt(1 - 4 * np.cos(param / 2) ** 2 * np.sin(param / 2) ** 2)
+    )
+
+
+class TestInitialization:
+    """Unit tests for the ``qml.vn_entropy`` function."""
+
+    @pytest.mark.all_interfaces
+    @pytest.mark.parametrize(
+        "state_vector,expected",
+        [([1.0, 0.0, 0.0, 1.0] / qml.math.sqrt(2), qml.math.log(2)), ([1.0, 0.0, 0.0, 0.0], 0)],
+    )
+    @pytest.mark.parametrize("interface", ["autograd", "jax", "tf", "torch"])
+    def test_vn_entropy(self, interface, state_vector, expected):
+        """Tests the output of qml.vn_entropy"""
+        dev = qml.device("default.qubit.legacy", wires=2)
+
+        @qml.qnode(dev, interface=interface)
+        def circuit():
+            qml.StatePrep(state_vector, wires=[0, 1])
+            return qml.vn_entropy(wires=0)
+
+        res = circuit()
+        new_res = qml.vn_entropy(wires=0).process_state(
+            state=circuit.device.state, wire_order=circuit.device.wires
+        )
+        assert qml.math.allclose(res, expected)
+        assert qml.math.allclose(new_res, expected)
+        assert INTERFACE_MAP.get(qml.math.get_interface(new_res)) == interface
+        assert res.dtype == new_res.dtype
+
+    def test_queue(self):
+        """Test that the right measurement class is queued."""
+        dev = qml.device("default.qubit.legacy", wires=2)
+
+        @qml.qnode(dev)
+        def circuit():
+            return qml.vn_entropy(wires=0, log_base=2)
+
+        circuit()
+
+        assert isinstance(circuit.tape[0], VnEntropyMP)
+
+    @pytest.mark.parametrize("shots, shape", [(None, ()), (10, ()), ((1, 10), ((), ()))])
+    def test_shape(self, shots, shape):
+        """Test the ``shape`` method."""
+        meas = qml.vn_entropy(wires=0)
+        dev = qml.device("default.qubit.legacy", wires=1, shots=shots)
+
+        assert meas.shape(dev, Shots(shots)) == shape
+
+
+class TestIntegration:
+    """Integration tests for the vn_entropy measurement function."""
+
+    parameters = np.linspace(0, 2 * np.pi, 10)
+
+    single_wires_list = [
+        [0],
+        [1],
+    ]
+
+    base = [2, np.exp(1), 10]
+
+    check_state = [True, False]
+
+    diff_methods = ["backprop", "finite-diff"]
+
+    def test_shot_vec_error(self):
+        """Test an error is raised when using shot vectors with vn_entropy."""
+        dev = qml.device("default.qubit.legacy", wires=2, shots=[1, 10, 10, 1000])
+
+        @qml.qnode(device=dev)
+        def circuit(x):
+            qml.Hadamard(wires=[0])
+            qml.CRX(x, wires=[0, 1])
+            return qml.vn_entropy(wires=[0])
+
+        with pytest.raises(
+            NotImplementedError, match="Von Neumann entropy is not supported with shot vectors"
+        ):
+            circuit(0.5)
+
+    @pytest.mark.parametrize("wires", single_wires_list)
+    @pytest.mark.parametrize("param", parameters)
+    @pytest.mark.parametrize("base", base)
+    def test_IsingXX_qnode_entropy(self, param, wires, base):
+        """Test entropy for a QNode numpy."""
+
+        dev = qml.device("default.qubit.legacy", wires=2)
+
+        @qml.qnode(dev)
+        def circuit_entropy(x):
+            qml.IsingXX(x, wires=[0, 1])
+            return qml.vn_entropy(wires=wires, log_base=base)
+
+        entropy = circuit_entropy(param)
+
+        expected_entropy = expected_entropy_ising_xx(param) / np.log(base)
+        assert qml.math.allclose(entropy, expected_entropy)
+
+    @pytest.mark.autograd
+    @pytest.mark.parametrize("wires", single_wires_list)
+    @pytest.mark.parametrize("param", parameters)
+    @pytest.mark.parametrize("base", base)
+    @pytest.mark.parametrize("diff_method", diff_methods)
+    def test_IsingXX_qnode_entropy_grad(self, param, wires, base, diff_method):
+        """Test entropy for a QNode gradient with autograd."""
+
+        dev = qml.device("default.qubit.legacy", wires=2)
+
+        @qml.qnode(dev, diff_method=diff_method)
+        def circuit_entropy(x):
+            qml.IsingXX(x, wires=[0, 1])
+            return qml.vn_entropy(wires=wires, log_base=base)
+
+        grad_entropy = qml.grad(circuit_entropy)(param)
+
+        # higher tolerance for finite-diff method
+        tol = 1e-8 if diff_method == "backprop" else 1e-5
+
+        grad_expected_entropy = expected_entropy_grad_ising_xx(param) / np.log(base)
+        assert qml.math.allclose(grad_entropy, grad_expected_entropy, atol=tol)
+
+    @pytest.mark.torch
+    @pytest.mark.parametrize("wires", single_wires_list)
+    @pytest.mark.parametrize("param", parameters)
+    @pytest.mark.parametrize("base", base)
+    @pytest.mark.parametrize("interface", ["torch"])
+    def test_IsingXX_qnode_torch_entropy(self, param, wires, base, interface):
+        """Test entropy for a QNode with torch interface."""
+        import torch
+
+        dev = qml.device("default.qubit.legacy", wires=2)
+
+        @qml.qnode(dev, interface=interface)
+        def circuit_entropy(x):
+            qml.IsingXX(x, wires=[0, 1])
+            return qml.vn_entropy(wires=wires, log_base=base)
+
+        entropy = circuit_entropy(torch.tensor(param))
+
+        expected_entropy = expected_entropy_ising_xx(param) / np.log(base)
+
+        assert qml.math.allclose(entropy, expected_entropy)
+
+    @pytest.mark.torch
+    @pytest.mark.parametrize("wires", single_wires_list)
+    @pytest.mark.parametrize("param", parameters)
+    @pytest.mark.parametrize("base", base)
+    @pytest.mark.parametrize("diff_method", diff_methods)
+    def test_IsingXX_qnode_entropy_grad_torch(self, param, wires, base, diff_method):
+        """Test entropy for a QNode gradient with torch."""
+        import torch
+
+        dev = qml.device("default.qubit.legacy", wires=2)
+
+        @qml.qnode(dev, interface="torch", diff_method=diff_method)
+        def circuit_entropy(x):
+            qml.IsingXX(x, wires=[0, 1])
+            return qml.vn_entropy(wires=wires, log_base=base)
+
+        grad_expected_entropy = expected_entropy_grad_ising_xx(param) / np.log(base)
+
+        param = torch.tensor(param, dtype=torch.float64, requires_grad=True)
+        entropy = circuit_entropy(param)
+        entropy.backward()
+        grad_entropy = param.grad
+
+        # higher tolerance for finite-diff method
+        tol = 1e-8 if diff_method == "backprop" else 1e-5
+
+        assert qml.math.allclose(grad_entropy, grad_expected_entropy, atol=tol)
+
+    @pytest.mark.tf
+    @pytest.mark.parametrize("wires", single_wires_list)
+    @pytest.mark.parametrize("param", parameters)
+    @pytest.mark.parametrize("base", base)
+    @pytest.mark.parametrize("interface", ["tf"])
+    def test_IsingXX_qnode_tf_entropy(self, param, wires, base, interface):
+        """Test entropy for a QNode with tf interface."""
+        import tensorflow as tf
+
+        dev = qml.device("default.qubit.legacy", wires=2)
+
+        @qml.qnode(dev, interface=interface)
+        def circuit_entropy(x):
+            qml.IsingXX(x, wires=[0, 1])
+            return qml.vn_entropy(wires=wires, log_base=base)
+
+        entropy = circuit_entropy(tf.Variable(param))
+
+        expected_entropy = expected_entropy_ising_xx(param) / np.log(base)
+
+        assert qml.math.allclose(entropy, expected_entropy)
+
+    @pytest.mark.tf
+    @pytest.mark.parametrize("wires", single_wires_list)
+    @pytest.mark.parametrize("param", parameters)
+    @pytest.mark.parametrize("base", base)
+    @pytest.mark.parametrize("diff_method", diff_methods)
+    @pytest.mark.parametrize("interface", ["tf"])
+    def test_IsingXX_qnode_entropy_grad_tf(self, param, wires, base, diff_method, interface):
+        """Test entropy for a QNode gradient with tf."""
+        import tensorflow as tf
+
+        dev = qml.device("default.qubit.legacy", wires=2)
+
+        @qml.qnode(dev, interface=interface, diff_method=diff_method)
+        def circuit_entropy(x):
+            qml.IsingXX(x, wires=[0, 1])
+            return qml.vn_entropy(wires=wires, log_base=base)
+
+        param = tf.Variable(param)
+        with tf.GradientTape() as tape:
+            entropy = circuit_entropy(param)
+
+        grad_entropy = tape.gradient(entropy, param)
+
+        grad_expected_entropy = expected_entropy_grad_ising_xx(param) / np.log(base)
+
+        # higher tolerance for finite-diff method
+        tol = 1e-8 if diff_method == "backprop" else 1e-5
+
+        assert qml.math.allclose(grad_entropy, grad_expected_entropy, atol=tol)
+
+    @pytest.mark.jax
+    @pytest.mark.parametrize("wires", single_wires_list)
+    @pytest.mark.parametrize("param", parameters)
+    @pytest.mark.parametrize("base", base)
+    @pytest.mark.parametrize("interface", ["jax"])
+    def test_IsingXX_qnode_jax_entropy(self, param, wires, base, interface):
+        """Test entropy for a QNode with jax interface."""
+        import jax.numpy as jnp
+
+        dev = qml.device("default.qubit.legacy", wires=2)
+
+        @qml.qnode(dev, interface=interface)
+        def circuit_entropy(x):
+            qml.IsingXX(x, wires=[0, 1])
+            return qml.vn_entropy(wires=wires, log_base=base)
+
+        entropy = circuit_entropy(jnp.array(param))
+
+        expected_entropy = expected_entropy_ising_xx(param) / np.log(base)
+
+        assert qml.math.allclose(entropy, expected_entropy)
+
+    @pytest.mark.jax
+    @pytest.mark.parametrize("wires", single_wires_list)
+    @pytest.mark.parametrize("param", parameters)
+    @pytest.mark.parametrize("base", base)
+    @pytest.mark.parametrize("diff_method", diff_methods)
+    @pytest.mark.parametrize("interface", ["jax"])
+    def test_IsingXX_qnode_entropy_grad_jax(self, param, wires, base, diff_method, interface):
+        """Test entropy for a QNode gradient with Jax."""
+        import jax
+
+        dev = qml.device("default.qubit.legacy", wires=2)
+
+        @qml.qnode(dev, interface=interface, diff_method=diff_method)
+        def circuit_entropy(x):
+            qml.IsingXX(x, wires=[0, 1])
+            return qml.vn_entropy(wires=wires, log_base=base)
+
+        grad_entropy = jax.grad(circuit_entropy)(jax.numpy.array(param))
+        grad_expected_entropy = expected_entropy_grad_ising_xx(param) / np.log(base)
+
+        # higher tolerance for finite-diff method
+        tol = 1e-8 if diff_method == "backprop" else 1e-5
+
+        assert qml.math.allclose(grad_entropy, grad_expected_entropy, rtol=1e-04, atol=tol)
+
+    @pytest.mark.jax
+    @pytest.mark.parametrize("wires", single_wires_list)
+    @pytest.mark.parametrize("param", parameters)
+    @pytest.mark.parametrize("base", base)
+    @pytest.mark.parametrize("interface", ["jax"])
+    def test_IsingXX_qnode_jax_jit_entropy(self, param, wires, base, interface):
+        """Test entropy for a QNode with jax-jit interface."""
+        import jax
+        import jax.numpy as jnp
+
+        dev = qml.device("default.qubit.legacy", wires=2)
+
+        @qml.qnode(dev, interface=interface)
+        def circuit_entropy(x):
+            qml.IsingXX(x, wires=[0, 1])
+            return qml.vn_entropy(wires=wires, log_base=base)
+
+        entropy = jax.jit(circuit_entropy)(jnp.array(param))
+
+        expected_entropy = expected_entropy_ising_xx(param) / np.log(base)
+
+        assert qml.math.allclose(entropy, expected_entropy)
+
+    @pytest.mark.jax
+    @pytest.mark.parametrize("wires", single_wires_list)
+    @pytest.mark.parametrize("param", parameters)
+    @pytest.mark.parametrize("base", base)
+    @pytest.mark.parametrize("diff_method", diff_methods)
+    @pytest.mark.parametrize("interface", ["jax-jit"])
+    def test_IsingXX_qnode_entropy_grad_jax_jit(self, param, wires, base, diff_method, interface):
+        """Test entropy for a QNode gradient with Jax-jit."""
+        import jax
+
+        dev = qml.device("default.qubit.legacy", wires=2)
+
+        @qml.qnode(dev, interface=interface, diff_method=diff_method)
+        def circuit_entropy(x):
+            qml.IsingXX(x, wires=[0, 1])
+            return qml.vn_entropy(wires=wires, log_base=base)
+
+        grad_entropy = jax.jit(jax.grad(circuit_entropy))(jax.numpy.array(param))
+
+        grad_expected_entropy = expected_entropy_grad_ising_xx(param) / np.log(base)
+
+        assert qml.math.allclose(grad_entropy, grad_expected_entropy, rtol=1e-04, atol=1e-05)
+
+    def test_qnode_entropy_no_custom_wires(self):
+        """Test that entropy cannot be returned with custom wires."""
+
+        dev = qml.device("default.qubit.legacy", wires=["a", 1])
+
+        @qml.qnode(dev)
+        def circuit_entropy(x):
+            qml.IsingXX(x, wires=["a", 1])
+            return qml.vn_entropy(wires=["a"])
+
+        with pytest.raises(
+            qml.QuantumFunctionError,
+            match="Returning the Von Neumann entropy is not supported when using custom wire labels",
+        ):
+            circuit_entropy(0.1)
diff --git a/tests/measurements/test_classical_shadow.py b/tests/measurements/test_classical_shadow.py
index 06f0b3a7ca2..d67967c2ff9 100644
--- a/tests/measurements/test_classical_shadow.py
+++ b/tests/measurements/test_classical_shadow.py
@@ -23,21 +23,15 @@
 from pennylane.measurements import ClassicalShadowMP, Shots
 from pennylane.measurements.classical_shadow import ShadowExpvalMP
 
-# pylint: disable=dangerous-default-value, too-many-arguments, comparison-with-callable, no-member
+# pylint: disable=dangerous-default-value, too-many-arguments
 
 
-def get_circuit(wires, shots, seed_recipes, interface="autograd", device="default.qubit.legacy"):
+def get_circuit(wires, shots, seed_recipes, interface="autograd", device="default.qubit"):
     """
     Return a QNode that prepares the state (|00...0> + |11...1>) / sqrt(2)
         and performs the classical shadow measurement
     """
-    if device is not None:
-        dev = qml.device(device, wires=wires, shots=shots)
-    else:
-        dev = qml.device("default.qubit.legacy", wires=wires, shots=shots)
-
-        # make the device call the superclass method to switch between the general qubit device and device specific implementations (i.e. for default qubit)
-        dev.classical_shadow = super(type(dev), dev).classical_shadow
+    dev = qml.device(device, wires=wires, shots=shots)
 
     @qml.qnode(dev, interface=interface)
     def circuit():
@@ -51,17 +45,11 @@ def circuit():
     return circuit
 
 
-def get_x_basis_circuit(wires, shots, interface="autograd", device="default.qubit.legacy"):
+def get_x_basis_circuit(wires, shots, interface="autograd"):
     """
     Return a QNode that prepares the |++..+> state and performs a classical shadow measurement
     """
-    if device is not None:
-        dev = qml.device(device, wires=wires, shots=shots)
-    else:
-        dev = qml.device("default.qubit.legacy", wires=wires, shots=shots)
-
-        # make the device call the superclass method to switch between the general qubit device and device specific implementations (i.e. for default qubit)
-        dev.classical_shadow = super(type(dev), dev).classical_shadow
+    dev = qml.device("default.qubit", wires=wires, shots=shots)
 
     @qml.qnode(dev, interface=interface)
     def circuit():
@@ -72,17 +60,11 @@ def circuit():
     return circuit
 
 
-def get_y_basis_circuit(wires, shots, interface="autograd", device="default.qubit.legacy"):
+def get_y_basis_circuit(wires, shots, interface="autograd"):
     """
     Return a QNode that prepares the |+i>|+i>...|+i> state and performs a classical shadow measurement
     """
-    if device is not None:
-        dev = qml.device(device, wires=wires, shots=shots)
-    else:
-        dev = qml.device("default.qubit.legacy", wires=wires, shots=shots)
-
-        # make the device call the superclass method to switch between the general qubit device and device specific implementations (i.e. for default qubit)
-        dev.classical_shadow = super(type(dev), dev).classical_shadow
+    dev = qml.device("default.qubit", wires=wires, shots=shots)
 
     @qml.qnode(dev, interface=interface)
     def circuit():
@@ -94,17 +76,11 @@ def circuit():
     return circuit
 
 
-def get_z_basis_circuit(wires, shots, interface="autograd", device="default.qubit.legacy"):
+def get_z_basis_circuit(wires, shots, interface="autograd"):
     """
     Return a QNode that prepares the |00..0> state and performs a classical shadow measurement
     """
-    if device is not None:
-        dev = qml.device(device, wires=wires, shots=shots)
-    else:
-        dev = qml.device("default.qubit.legacy", wires=wires, shots=shots)
-
-        # make the device call the superclass method to switch between the general qubit device and device specific implementations (i.e. for default qubit)
-        dev.classical_shadow = super(type(dev), dev).classical_shadow
+    dev = qml.device("default.qubit", wires=wires, shots=shots)
 
     @qml.qnode(dev, interface=interface)
     def circuit():
@@ -263,7 +239,7 @@ def test_measurement_process_numeric_type(self, wires, seed):
     @pytest.mark.parametrize("seed", seed_recipes_list)
     def test_measurement_process_shape(self, wires, shots, seed):
         """Test that the shape of the MeasurementProcess instance is correct"""
-        dev = qml.device("default.qubit.legacy", wires=wires, shots=shots)
+        dev = qml.device("default.qubit", wires=wires, shots=shots)
         shots_obj = Shots(shots)
         res = qml.classical_shadow(wires=range(wires), seed=seed)
         assert res.shape(dev, shots_obj) == (2, shots, wires)
@@ -304,7 +280,7 @@ def test_measurement_process_copy(self, wires, seed):
     @pytest.mark.parametrize("shots", shots_list)
     @pytest.mark.parametrize("seed", seed_recipes_list)
     @pytest.mark.parametrize("interface", ["autograd", "jax", "tf", "torch"])
-    @pytest.mark.parametrize("device", ["default.qubit.legacy", "default.mixed", None])
+    @pytest.mark.parametrize("device", ["default.qubit", "default.mixed"])
     def test_format(self, wires, shots, seed, interface, device):
         """Test that the format of the returned classical shadow
         measurement is correct"""
@@ -334,63 +310,53 @@ def test_format(self, wires, shots, seed, interface, device):
 
     @pytest.mark.all_interfaces
     @pytest.mark.parametrize("interface", ["autograd", "jax", "tf", "torch"])
-    @pytest.mark.parametrize("device", ["default.qubit.legacy", None])
     @pytest.mark.parametrize(
         "circuit_fn, basis_recipe",
         [(get_x_basis_circuit, 0), (get_y_basis_circuit, 1), (get_z_basis_circuit, 2)],
     )
-    def test_return_distribution(self, wires, interface, device, circuit_fn, basis_recipe):
+    def test_return_distribution(self, wires, interface, circuit_fn, basis_recipe):
         """Test that the distribution of the bits and recipes are correct for a circuit
         that prepares all qubits in a Pauli basis"""
         # high number of shots to prevent true negatives
         np.random.seed(42)
         shots = 1000
 
-        circuit = circuit_fn(wires, shots=shots, interface=interface, device=device)
+        circuit = circuit_fn(wires, shots=shots, interface=interface)
         bits, recipes = circuit()
-        new_bits, new_recipes = circuit.tape.measurements[0].process(circuit.tape, circuit.device)
 
         # test that the recipes follow a rough uniform distribution
         ratios = np.unique(recipes, return_counts=True)[1] / (wires * shots)
         assert np.allclose(ratios, 1 / 3, atol=1e-1)
-        new_ratios = np.unique(new_recipes, return_counts=True)[1] / (wires * shots)
-        assert np.allclose(new_ratios, 1 / 3, atol=1e-1)
 
         # test that the bit is 0 for all X measurements
         assert qml.math.allequal(bits[recipes == basis_recipe], 0)
-        assert qml.math.allequal(new_bits[new_recipes == basis_recipe], 0)
 
         # test that the bits are uniformly distributed for all Y and Z measurements
         bits1 = bits[recipes == (basis_recipe + 1) % 3]
         ratios1 = np.unique(bits1, return_counts=True)[1] / bits1.shape[0]
         assert np.allclose(ratios1, 1 / 2, atol=1e-1)
-        new_bits1 = new_bits[new_recipes == (basis_recipe + 1) % 3]
-        new_ratios1 = np.unique(new_bits1, return_counts=True)[1] / new_bits1.shape[0]
-        assert np.allclose(new_ratios1, 1 / 2, atol=1e-1)
 
         bits2 = bits[recipes == (basis_recipe + 2) % 3]
         ratios2 = np.unique(bits2, return_counts=True)[1] / bits2.shape[0]
         assert np.allclose(ratios2, 1 / 2, atol=1e-1)
 
-        new_bits2 = new_bits[new_recipes == (basis_recipe + 2) % 3]
-        new_ratios2 = np.unique(new_bits2, return_counts=True)[1] / new_bits2.shape[0]
-        assert np.allclose(new_ratios2, 1 / 2, atol=1e-1)
-
+    @pytest.mark.xfail(reason="until DQ2 port")
     @pytest.mark.parametrize("seed", seed_recipes_list)
     def test_shots_none_error(self, wires, seed):
         """Test that an error is raised when a device with shots=None is used
         to obtain classical shadows"""
         circuit = get_circuit(wires, None, seed)
 
-        msg = "The number of shots has to be explicitly set on the device when using sample-based measurements"
-        with pytest.raises(qml.QuantumFunctionError, match=msg):
+        msg = "Analytic circuits must only contain StateMeasurements"
+        with pytest.raises(qml.DeviceError, match=msg):
             circuit()
 
+    @pytest.mark.xfail(reason="until DQ2 port")
     @pytest.mark.parametrize("shots", shots_list)
     def test_multi_measurement_error(self, wires, shots):
         """Test that an error is raised when classical shadows is returned
         with other measurement processes"""
-        dev = qml.device("default.qubit.legacy", wires=wires, shots=shots)
+        dev = qml.device("default.qubit", wires=wires, shots=shots)
 
         @qml.qnode(dev)
         def circuit():
@@ -401,13 +367,14 @@ def circuit():
 
             return qml.classical_shadow(wires=range(wires)), qml.expval(qml.PauliZ(0))
 
-        msg = "Classical shadows cannot be returned in combination with other return types"
-        with pytest.raises(qml.QuantumFunctionError, match=msg):
-            circuit()
+        res = circuit()
+        assert isinstance(res, tuple) and len(res) == 2
+        assert qml.math.shape(res[0]) == (2, shots, wires)
+        assert qml.math.shape(res[1]) == ()
 
 
 def hadamard_circuit(wires, shots=10000, interface="autograd"):
-    dev = qml.device("default.qubit.legacy", wires=wires, shots=shots)
+    dev = qml.device("default.qubit", wires=wires, shots=shots)
 
     @qml.qnode(dev, interface=interface)
     def circuit(obs, k=1):
@@ -419,7 +386,7 @@ def circuit(obs, k=1):
 
 
 def max_entangled_circuit(wires, shots=10000, interface="autograd"):
-    dev = qml.device("default.qubit.legacy", wires=wires, shots=shots)
+    dev = qml.device("default.qubit", wires=wires, shots=shots)
 
     @qml.qnode(dev, interface=interface)
     def circuit(obs, k=1):
@@ -432,7 +399,7 @@ def circuit(obs, k=1):
 
 
 def qft_circuit(wires, shots=10000, interface="autograd"):
-    dev = qml.device("default.qubit.legacy", wires=wires, shots=shots)
+    dev = qml.device("default.qubit", wires=wires, shots=shots)
 
     one_state = np.zeros(wires)
     one_state[-1] = 1
@@ -458,7 +425,7 @@ def test_measurement_process_numeric_type(self):
     @pytest.mark.parametrize("shots", [1, 10])
     def test_measurement_process_shape(self, wires, shots):
         """Test that the shape of the MeasurementProcess instance is correct"""
-        dev = qml.device("default.qubit.legacy", wires=wires, shots=shots)
+        dev = qml.device("default.qubit", wires=wires, shots=shots)
         H = qml.PauliZ(0)
         res = qml.shadow_expval(H)
         assert len(res.shape(dev, Shots(shots))) == 0
@@ -491,20 +458,22 @@ def test_measurement_process_copy(self):
         assert copied_res.k == res.k
         assert copied_res.seed == res.seed
 
+    @pytest.mark.xfail(reason="until DQ2 port")
     def test_shots_none_error(self):
         """Test that an error is raised when a device with shots=None is used
         to obtain classical shadows"""
         circuit = hadamard_circuit(2, None)
         H = qml.PauliZ(0)
 
-        msg = "The number of shots has to be explicitly set on the device when using sample-based measurements"
-        with pytest.raises(qml.QuantumFunctionError, match=msg):
+        msg = "Analytic circuits must only contain StateMeasurements"
+        with pytest.raises(qml.DeviceError, match=msg):
             _ = circuit(H, k=10)
 
-    def test_multi_measurement_error(self):
-        """Test that an error is raised when classical shadows is returned
+    @pytest.mark.xfail(reason="until DQ2 port")
+    def test_multi_measurement_allowed(self):
+        """Test that no error is raised when classical shadows is returned
         with other measurement processes"""
-        dev = qml.device("default.qubit.legacy", wires=2, shots=100)
+        dev = qml.device("default.qubit", wires=2, shots=10000)
 
         @qml.qnode(dev)
         def circuit():
@@ -512,9 +481,9 @@ def circuit():
             qml.CNOT(wires=[0, 1])
             return qml.shadow_expval(qml.PauliZ(0)), qml.expval(qml.PauliZ(0))
 
-        msg = "Classical shadows cannot be returned in combination with other return types"
-        with pytest.raises(qml.QuantumFunctionError, match=msg):
-            _ = circuit()
+        res = circuit()
+        assert isinstance(res, tuple)
+        assert qml.math.allclose(res, 0, atol=0.05)
 
     def test_obs_not_queued(self):
         """Test that the observable passed to qml.shadow_expval is not queued"""
@@ -587,12 +556,10 @@ def test_hadamard_expval(self, k=1, obs=obs_hadamard, expected=expected_hadamard
         superposition of qubits"""
         circuit = hadamard_circuit(3, shots=100000)
         actual = circuit(obs, k=k)
-        new_actual = circuit.tape.measurements[0].process(circuit.tape, circuit.device)
 
         assert actual.shape == (len(obs_hadamard),)
         assert actual.dtype == np.float64
         assert qml.math.allclose(actual, expected, atol=1e-1)
-        assert qml.math.allclose(new_actual, expected, atol=1e-1)
 
     def test_max_entangled_expval(
         self, k=1, obs=obs_max_entangled, expected=expected_max_entangled
@@ -601,12 +568,10 @@ def test_max_entangled_expval(
         entangled state"""
         circuit = max_entangled_circuit(3, shots=100000)
         actual = circuit(obs, k=k)
-        new_actual = circuit.tape.measurements[0].process(circuit.tape, circuit.device)
 
         assert actual.shape == (len(obs_max_entangled),)
         assert actual.dtype == np.float64
         assert qml.math.allclose(actual, expected, atol=1e-1)
-        assert qml.math.allclose(new_actual, expected, atol=1e-1)
 
     def test_non_pauli_error(self):
         """Test that an error is raised when a non-Pauli observable is passed"""
@@ -627,12 +592,10 @@ def test_qft_expval(self, interface, k=1, obs=obs_qft, expected=expected_qft):
 
         circuit = qft_circuit(3, shots=100000, interface=interface)
         actual = circuit(obs, k=k)
-        new_actual = circuit.tape.measurements[0].process(circuit.tape, circuit.device)
 
         assert actual.shape == (len(obs_qft),)
         assert actual.dtype == torch.float64 if interface == "torch" else np.float64
         assert qml.math.allclose(actual, expected, atol=1e-1)
-        assert qml.math.allclose(new_actual, expected, atol=1e-1)
 
 
 obs_strongly_entangled = [
@@ -647,7 +610,7 @@ def test_qft_expval(self, interface, k=1, obs=obs_qft, expected=expected_qft):
 
 
 def strongly_entangling_circuit(wires, shots=10000, interface="autograd"):
-    dev = qml.device("default.qubit.legacy", wires=wires, shots=shots)
+    dev = qml.device("default.qubit", wires=wires, shots=shots)
 
     @qml.qnode(dev, interface=interface)
     def circuit(x, obs, k):
@@ -658,7 +621,7 @@ def circuit(x, obs, k):
 
 
 def strongly_entangling_circuit_exact(wires, interface="autograd"):
-    dev = qml.device("default.qubit.legacy", wires=wires)
+    dev = qml.device("default.qubit", wires=wires)
 
     @qml.qnode(dev, interface=interface)
     def circuit(x, obs):
diff --git a/tests/measurements/test_counts.py b/tests/measurements/test_counts.py
index f936561995f..ac5db079eb0 100644
--- a/tests/measurements/test_counts.py
+++ b/tests/measurements/test_counts.py
@@ -18,36 +18,9 @@
 
 import pennylane as qml
 from pennylane.measurements import AllCounts, Counts, CountsMP
-from pennylane.operation import Operator
 from pennylane.wires import Wires
 
 
-# TODO: Remove this when new CustomMP are the default
-def custom_measurement_process(device, spy):
-    assert len(spy.call_args_list) > 0  # make sure method is mocked properly
-
-    samples = device._samples  # pylint:disable=protected-access
-    call_args_list = list(spy.call_args_list)
-    for call_args in call_args_list:
-        if not call_args.kwargs.get("counts", False):
-            continue
-        meas = call_args.args[1]
-        shot_range, bin_size = (call_args.kwargs["shot_range"], call_args.kwargs["bin_size"])
-        if isinstance(meas, Operator):
-            all_outcomes = meas.return_type is AllCounts
-            meas = qml.counts(op=meas, all_outcomes=all_outcomes)
-        old_res = device.sample(call_args.args[1], **call_args.kwargs)
-        new_res = meas.process_samples(
-            samples=samples, wire_order=device.wires, shot_range=shot_range, bin_size=bin_size
-        )
-        if isinstance(old_res, dict):
-            old_res = [old_res]
-            new_res = [new_res]
-        for old, new in zip(old_res, new_res):
-            assert old.keys() == new.keys()
-            assert qml.math.allequal(list(old.values()), list(new.values()))
-
-
 class TestCounts:
     """Tests for the counts function"""
 
@@ -245,12 +218,11 @@ def test_counts_all_outcomes_measurement_value(self):
 class TestCountsIntegration:
     # pylint:disable=too-many-public-methods,not-an-iterable
 
-    def test_counts_dimension(self, mocker):
+    def test_counts_dimension(self):
         """Test that the counts function outputs counts of the right size"""
         n_sample = 10
 
-        dev = qml.device("default.qubit.legacy", wires=2, shots=n_sample)
-        spy = mocker.spy(qml.QubitDevice, "sample")
+        dev = qml.device("default.qubit", wires=2, shots=n_sample)
 
         @qml.qnode(dev)
         def circuit():
@@ -262,14 +234,11 @@ def circuit():
         assert len(sample) == 2
         assert np.all([sum(s.values()) == n_sample for s in sample])
 
-        custom_measurement_process(dev, spy)
-
-    def test_batched_counts_dimension(self, mocker):
+    def test_batched_counts_dimension(self):
         """Test that the counts function outputs counts of the right size with batching"""
         n_sample = 10
 
-        dev = qml.device("default.qubit.legacy", wires=2, shots=n_sample)
-        spy = mocker.spy(qml.QubitDevice, "sample")
+        dev = qml.device("default.qubit", wires=2, shots=n_sample)
 
         @qml.qnode(dev)
         def circuit():
@@ -282,14 +251,11 @@ def circuit():
         assert len(sample) == 2
         assert np.all([sum(s.values()) == n_sample for batch in sample for s in batch])
 
-        custom_measurement_process(dev, spy)
-
-    def test_counts_combination(self, mocker):
+    def test_counts_combination(self):
         """Test the output of combining expval, var and counts"""
         n_sample = 10
 
-        dev = qml.device("default.qubit.legacy", wires=3, shots=n_sample)
-        spy = mocker.spy(qml.QubitDevice, "sample")
+        dev = qml.device("default.qubit", wires=3, shots=n_sample)
 
         @qml.qnode(dev, diff_method="parameter-shift")
         def circuit():
@@ -306,14 +272,11 @@ def circuit():
         assert len(result) == 3
         assert sum(result[0].values()) == n_sample
 
-        custom_measurement_process(dev, spy)
-
-    def test_single_wire_counts(self, mocker):
+    def test_single_wire_counts(self):
         """Test the return type and shape of sampling counts from a single wire"""
         n_sample = 10
 
-        dev = qml.device("default.qubit.legacy", wires=1, shots=n_sample)
-        spy = mocker.spy(qml.QubitDevice, "sample")
+        dev = qml.device("default.qubit", wires=1, shots=n_sample)
 
         @qml.qnode(dev)
         def circuit():
@@ -326,15 +289,12 @@ def circuit():
         assert isinstance(result, dict)
         assert sum(result.values()) == n_sample
 
-        custom_measurement_process(dev, spy)
-
-    def test_multi_wire_counts_regular_shape(self, mocker):
+    def test_multi_wire_counts_regular_shape(self):
         """Test the return type and shape of sampling multiple wires
         where a rectangular array is expected"""
         n_sample = 10
 
-        dev = qml.device("default.qubit.legacy", wires=3, shots=n_sample)
-        spy = mocker.spy(qml.QubitDevice, "sample")
+        dev = qml.device("default.qubit", wires=3, shots=n_sample)
 
         @qml.qnode(dev)
         def circuit():
@@ -352,12 +312,10 @@ def circuit():
         assert all(sum(r.values()) == n_sample for r in result)
         assert all(all(v.dtype == np.dtype("int") for v in r.values()) for r in result)
 
-        custom_measurement_process(dev, spy)
-
     def test_observable_return_type_is_counts(self):
         """Test that the return type of the observable is :attr:`ObservableReturnTypes.Counts`"""
         n_shots = 10
-        dev = qml.device("default.qubit.legacy", wires=1, shots=n_shots)
+        dev = qml.device("default.qubit", wires=1, shots=n_shots)
 
         @qml.qnode(dev)
         def circuit():
@@ -367,10 +325,9 @@ def circuit():
         circuit()
         assert circuit._qfunc_output.return_type is Counts  # pylint: disable=protected-access
 
-    def test_providing_no_observable_and_no_wires_counts(self, mocker):
+    def test_providing_no_observable_and_no_wires_counts(self):
         """Test that we can provide no observable and no wires to sample function"""
-        dev = qml.device("default.qubit.legacy", wires=2, shots=1000)
-        spy = mocker.spy(qml.QubitDevice, "sample")
+        dev = qml.device("default.qubit", wires=2, shots=1000)
 
         @qml.qnode(dev)
         def circuit():
@@ -382,14 +339,11 @@ def circuit():
 
         circuit()
 
-        custom_measurement_process(dev, spy)
-
-    def test_providing_no_observable_and_wires_counts(self, mocker):
+    def test_providing_no_observable_and_wires_counts(self):
         """Test that we can provide no observable but specify wires to the sample function"""
         wires = [0, 2]
         wires_obj = qml.wires.Wires(wires)
-        dev = qml.device("default.qubit.legacy", wires=3, shots=1000)
-        spy = mocker.spy(qml.QubitDevice, "sample")
+        dev = qml.device("default.qubit", wires=3, shots=1000)
 
         @qml.qnode(dev)
         def circuit():
@@ -402,13 +356,10 @@ def circuit():
 
         circuit()
 
-        custom_measurement_process(dev, spy)
-
-    def test_batched_counts_work_individually(self, mocker):
+    def test_batched_counts_work_individually(self):
         """Test that each counts call operates independently"""
         n_shots = 10
-        dev = qml.device("default.qubit.legacy", wires=1, shots=n_shots)
-        spy = mocker.spy(qml.QubitDevice, "sample")
+        dev = qml.device("default.qubit", wires=1, shots=n_shots)
 
         @qml.qnode(dev)
         def circuit():
@@ -416,20 +367,19 @@ def circuit():
             return qml.counts()
 
         assert circuit() == [{"1": 10}, {"0": 10}]
-        custom_measurement_process(dev, spy)
 
     @pytest.mark.all_interfaces
     @pytest.mark.parametrize("wires, basis_state", [(None, "010"), ([2, 1], "01")])
     @pytest.mark.parametrize("interface", ["autograd", "jax", "tensorflow", "torch"])
-    def test_counts_no_op_finite_shots(self, interface, wires, basis_state, mocker):
+    def test_counts_no_op_finite_shots(self, interface, wires, basis_state):
         """Check all interfaces with computational basis state counts and
         finite shot"""
         n_shots = 10
-        dev = qml.device("default.qubit.legacy", wires=3, shots=n_shots)
-        spy = mocker.spy(qml.QubitDevice, "sample")
+        dev = qml.device("default.qubit", wires=3, shots=n_shots)
 
         @qml.qnode(dev, interface=interface)
         def circuit():
+            qml.Identity(wires=[0, 1, 2])
             qml.PauliX(1)
             return qml.counts(wires=wires)
 
@@ -437,16 +387,13 @@ def circuit():
         assert res == {basis_state: n_shots}
         assert qml.math.get_interface(res[basis_state]) == interface
 
-        custom_measurement_process(dev, spy)
-
     @pytest.mark.all_interfaces
     @pytest.mark.parametrize("interface", ["autograd", "jax", "tensorflow", "torch"])
-    def test_counts_operator_finite_shots(self, interface, mocker):
+    def test_counts_operator_finite_shots(self, interface):
         """Check all interfaces with observable measurement counts and finite
         shot"""
         n_shots = 10
-        dev = qml.device("default.qubit.legacy", wires=3, shots=n_shots)
-        spy = mocker.spy(qml.QubitDevice, "sample")
+        dev = qml.device("default.qubit", wires=3, shots=n_shots)
 
         @qml.qnode(dev, interface=interface)
         def circuit():
@@ -455,22 +402,20 @@ def circuit():
         res = circuit()
         assert res == {1: n_shots}
 
-        custom_measurement_process(dev, spy)
-
     @pytest.mark.all_interfaces
     @pytest.mark.parametrize("shot_vec", [(1, 10, 10), (1, 10, 1000)])
     @pytest.mark.parametrize("wires, basis_state", [(None, "010"), ([2, 1], "01")])
     @pytest.mark.parametrize("interface", ["autograd", "jax", "tensorflow", "torch"])
     def test_counts_binned(
-        self, shot_vec, interface, wires, basis_state, mocker
+        self, shot_vec, interface, wires, basis_state
     ):  # pylint:disable=too-many-arguments
         """Check all interfaces with computational basis state counts and
         different shot vectors"""
-        dev = qml.device("default.qubit.legacy", wires=3, shots=shot_vec)
-        spy = mocker.spy(qml.QubitDevice, "sample")
+        dev = qml.device("default.qubit", wires=3, shots=shot_vec)
 
         @qml.qnode(dev, interface=interface)
         def circuit():
+            qml.Identity(wires=[0, 1, 2])
             qml.PauliX(1)
             return qml.counts(wires=wires)
 
@@ -483,16 +428,13 @@ def circuit():
         assert len(res) == len(shot_vec)
         assert sum(sum(res_bin.values()) for res_bin in res) == sum(shot_vec)
 
-        custom_measurement_process(dev, spy)
-
     @pytest.mark.all_interfaces
     @pytest.mark.parametrize("shot_vec", [(1, 10, 10), (1, 10, 1000)])
     @pytest.mark.parametrize("interface", ["autograd", "jax", "tensorflow", "torch"])
-    def test_counts_operator_binned(self, shot_vec, interface, mocker):
+    def test_counts_operator_binned(self, shot_vec, interface):
         """Check all interfaces with observable measurement counts and different
         shot vectors"""
-        dev = qml.device("default.qubit.legacy", wires=3, shots=shot_vec)
-        spy = mocker.spy(qml.QubitDevice, "sample")
+        dev = qml.device("default.qubit", wires=3, shots=shot_vec)
 
         @qml.qnode(dev, interface=interface)
         def circuit():
@@ -507,17 +449,15 @@ def circuit():
         assert len(res) == len(shot_vec)
         assert sum(sum(res_bin.values()) for res_bin in res) == sum(shot_vec)
 
-        custom_measurement_process(dev, spy)
-
     @pytest.mark.parametrize("shot_vec", [(1, 10, 10), (1, 10, 1000)])
-    def test_counts_binned_4_wires(self, shot_vec, mocker):
+    def test_counts_binned_4_wires(self, shot_vec):
         """Check the autograd interface with computational basis state counts and
         different shot vectors on a device with 4 wires"""
-        dev = qml.device("default.qubit.legacy", wires=4, shots=shot_vec)
-        spy = mocker.spy(qml.QubitDevice, "sample")
+        dev = qml.device("default.qubit", wires=4, shots=shot_vec)
 
         @qml.qnode(dev, interface="autograd")
         def circuit():
+            qml.Identity(0)
             qml.PauliX(1)
             qml.PauliX(2)
             qml.PauliX(3)
@@ -533,14 +473,11 @@ def circuit():
         assert len(res) == len(shot_vec)
         assert sum(sum(res_bin.values()) for res_bin in res) == sum(shot_vec)
 
-        custom_measurement_process(dev, spy)
-
     @pytest.mark.parametrize("shot_vec", [(1, 10, 10), (1, 10, 1000)])
-    def test_counts_operator_binned_4_wires(self, shot_vec, mocker):
+    def test_counts_operator_binned_4_wires(self, shot_vec):
         """Check the autograd interface with observable samples to obtain
         counts from and different shot vectors on a device with 4 wires"""
-        dev = qml.device("default.qubit.legacy", wires=4, shots=shot_vec)
-        spy = mocker.spy(qml.QubitDevice, "sample")
+        dev = qml.device("default.qubit", wires=4, shots=shot_vec)
 
         @qml.qnode(dev, interface="autograd")
         def circuit():
@@ -559,8 +496,6 @@ def circuit():
         assert len(res) == len(shot_vec)
         assert sum(sum(res_bin.values()) for res_bin in res) == sum(shot_vec)
 
-        custom_measurement_process(dev, spy)
-
     meas2 = [
         qml.expval(qml.PauliZ(0)),
         qml.var(qml.PauliZ(0)),
@@ -572,11 +507,10 @@ def circuit():
     @pytest.mark.parametrize("interface", ["autograd", "jax", "tensorflow", "torch"])
     @pytest.mark.parametrize("meas2", meas2)
     @pytest.mark.parametrize("shots", [1000, (1, 10)])
-    def test_counts_observable_finite_shots(self, interface, meas2, shots, mocker):
+    def test_counts_observable_finite_shots(self, interface, meas2, shots):
         """Check all interfaces with observable measurement counts and finite
         shot"""
-        dev = qml.device("default.qubit.legacy", wires=3, shots=shots)
-        spy = mocker.spy(qml.QubitDevice, "sample")
+        dev = qml.device("default.qubit", wires=3, shots=shots)
 
         if isinstance(shots, tuple) and interface == "torch":
             pytest.skip("Torch needs to be updated for shot vectors.")
@@ -603,15 +537,12 @@ def circuit():
             assert isinstance(res[1], tuple)
             assert isinstance(res[1][0], dict)
 
-        custom_measurement_process(dev, spy)
-
-    def test_all_outcomes_kwarg_providing_observable(self, mocker):
+    def test_all_outcomes_kwarg_providing_observable(self):
         """Test that the dictionary keys *all* eigenvalues of the observable,
         including 0 count values, if observable is given and all_outcomes=True"""
 
         n_shots = 10
-        dev = qml.device("default.qubit.legacy", wires=1, shots=n_shots)
-        spy = mocker.spy(qml.QubitDevice, "sample")
+        dev = qml.device("default.qubit", wires=1, shots=n_shots)
 
         @qml.qnode(dev)
         def circuit():
@@ -622,35 +553,30 @@ def circuit():
 
         assert res == {1: n_shots, -1: 0}
 
-        custom_measurement_process(dev, spy)
-
-    def test_all_outcomes_kwarg_no_observable_no_wires(self, mocker):
+    def test_all_outcomes_kwarg_no_observable_no_wires(self):
         """Test that the dictionary keys are *all* the possible combinations
         of basis states for the device, including 0 count values, if no wire
         count and no observable are given and all_outcomes=True"""
 
         n_shots = 10
-        dev = qml.device("default.qubit.legacy", wires=2, shots=n_shots)
-        spy = mocker.spy(qml.QubitDevice, "sample")
+        dev = qml.device("default.qubit", wires=2, shots=n_shots)
 
         @qml.qnode(dev)
         def circuit():
+            qml.Identity(wires=[0, 1])
             return qml.counts(all_outcomes=True)
 
         res = circuit()
 
         assert res == {"00": n_shots, "01": 0, "10": 0, "11": 0}
 
-        custom_measurement_process(dev, spy)
-
-    def test_all_outcomes_kwarg_providing_wires_and_no_observable(self, mocker):
+    def test_all_outcomes_kwarg_providing_wires_and_no_observable(self):
         """Test that the dictionary keys are *all* possible combinations
         of basis states for the specified wires, including 0 count values,
         if wire count is given and all_outcomes=True"""
 
         n_shots = 10
-        dev = qml.device("default.qubit.legacy", wires=4, shots=n_shots)
-        spy = mocker.spy(qml.QubitDevice, "sample")
+        dev = qml.device("default.qubit", wires=4, shots=n_shots)
 
         @qml.qnode(dev)
         def circuit():
@@ -660,15 +586,12 @@ def circuit():
 
         assert res == {"00": n_shots, "01": 0, "10": 0, "11": 0}
 
-        custom_measurement_process(dev, spy)
-
-    def test_all_outcomes_hermitian(self, mocker):
+    def test_all_outcomes_hermitian(self):
         """Tests that the all_outcomes=True option for counts works with the
         qml.Hermitian observable"""
 
         n_shots = 10
-        dev = qml.device("default.qubit.legacy", wires=2, shots=n_shots)
-        spy = mocker.spy(qml.QubitDevice, "sample")
+        dev = qml.device("default.qubit", wires=2, shots=n_shots)
 
         A = np.array([[1, 0], [0, -1]])
 
@@ -680,17 +603,15 @@ def circuit(x):
 
         assert res == {-1.0: 0, 1.0: n_shots}
 
-        custom_measurement_process(dev, spy)
-
-    def test_all_outcomes_multiple_measurements(self, mocker):
+    def test_all_outcomes_multiple_measurements(self):
         """Tests that the all_outcomes=True option for counts works when
         multiple measurements are performed"""
 
-        dev = qml.device("default.qubit.legacy", wires=2, shots=10)
-        spy = mocker.spy(qml.QubitDevice, "sample")
+        dev = qml.device("default.qubit", wires=2, shots=10)
 
         @qml.qnode(dev)
         def circuit():
+            qml.Identity(wires=[0, 1])
             return qml.sample(qml.PauliZ(0)), qml.counts(), qml.counts(all_outcomes=True)
 
         res = circuit()
@@ -698,13 +619,11 @@ def circuit():
         assert len(res[0]) == 10
         assert res[1] == {"00": 10}
         assert res[2] == {"00": 10, "01": 0, "10": 0, "11": 0}
-        custom_measurement_process(dev, spy)
 
-    def test_batched_all_outcomes(self, mocker):
+    def test_batched_all_outcomes(self):
         """Tests that all_outcomes=True works with broadcasting."""
         n_shots = 10
-        dev = qml.device("default.qubit.legacy", wires=1, shots=n_shots)
-        spy = mocker.spy(qml.QubitDevice, "sample")
+        dev = qml.device("default.qubit", wires=1, shots=n_shots)
 
         @qml.qnode(dev)
         def circuit():
@@ -713,8 +632,6 @@ def circuit():
 
         assert circuit() == [{1: 0, -1: n_shots}, {1: n_shots, -1: 0}]
 
-        custom_measurement_process(dev, spy)
-
     def test_counts_empty_wires(self):
         """Test that using ``qml.counts`` with an empty wire list raises an error."""
         with pytest.raises(ValueError, match="Cannot set an empty list of wires."):
@@ -723,10 +640,11 @@ def test_counts_empty_wires(self):
     @pytest.mark.parametrize("shots", [1, 100])
     def test_counts_no_arguments(self, shots):
         """Test that using ``qml.counts`` with no arguments returns the counts of all wires."""
-        dev = qml.device("default.qubit.legacy", wires=3, shots=shots)
+        dev = qml.device("default.qubit", wires=3, shots=shots)
 
         @qml.qnode(dev)
         def circuit():
+            qml.Identity(wires=[0, 1, 2])
             return qml.counts()
 
         res = circuit()
@@ -737,14 +655,14 @@ def circuit():
 @pytest.mark.all_interfaces
 @pytest.mark.parametrize("wires, basis_state", [(None, "010"), ([2, 1], "01")])
 @pytest.mark.parametrize("interface", ["autograd", "jax", "tensorflow", "torch"])
-def test_counts_no_op_finite_shots(interface, wires, basis_state, mocker):
+def test_counts_no_op_finite_shots(interface, wires, basis_state):
     """Check all interfaces with computational basis state counts and finite shot"""
     n_shots = 10
-    dev = qml.device("default.qubit.legacy", wires=3, shots=n_shots)
-    spy = mocker.spy(qml.QubitDevice, "sample")
+    dev = qml.device("default.qubit", wires=3, shots=n_shots)
 
     @qml.qnode(dev, interface=interface)
     def circuit():
+        qml.Identity(wires=[0, 1, 2])
         qml.PauliX(1)
         return qml.counts(wires=wires)
 
@@ -752,59 +670,55 @@ def circuit():
     assert res == {basis_state: n_shots}
     assert qml.math.get_interface(res[basis_state]) == interface
 
-    custom_measurement_process(dev, spy)
-
 
 @pytest.mark.all_interfaces
 @pytest.mark.parametrize("wires, basis_states", [(None, ("010", "000")), ([2, 1], ("01", "00"))])
 @pytest.mark.parametrize("interface", ["autograd", "jax", "tensorflow", "torch"])
-def test_batched_counts_no_op_finite_shots(interface, wires, basis_states, mocker):
+def test_batched_counts_no_op_finite_shots(interface, wires, basis_states):
     """Check all interfaces with computational basis state counts and
     finite shot"""
     n_shots = 10
-    dev = qml.device("default.qubit.legacy", wires=3, shots=n_shots)
-    spy = mocker.spy(qml.QubitDevice, "sample")
+    dev = qml.device("default.qubit", wires=3, shots=n_shots)
 
     @qml.qnode(dev, interface=interface)
     def circuit():
+        qml.Identity(wires=[0, 1, 2])
         qml.pow(qml.PauliX(1), z=[1, 2])
         return qml.counts(wires=wires)
 
     assert circuit() == [{basis_state: n_shots} for basis_state in basis_states]
 
-    custom_measurement_process(dev, spy)
-
 
 @pytest.mark.all_interfaces
 @pytest.mark.parametrize("wires, basis_states", [(None, ("010", "000")), ([2, 1], ("01", "00"))])
 @pytest.mark.parametrize("interface", ["autograd", "jax", "tensorflow", "torch"])
-def test_batched_counts_and_expval_no_op_finite_shots(interface, wires, basis_states, mocker):
+def test_batched_counts_and_expval_no_op_finite_shots(interface, wires, basis_states):
     """Check all interfaces with computational basis state counts and
     finite shot"""
     n_shots = 10
-    dev = qml.device("default.qubit.legacy", wires=3, shots=n_shots)
-    spy = mocker.spy(qml.QubitDevice, "sample")
+    dev = qml.device("default.qubit", wires=3, shots=n_shots)
 
     @qml.qnode(dev, interface=interface)
     def circuit():
+        qml.Identity(2)
         qml.pow(qml.PauliX(1), z=[1, 2])
         return qml.counts(wires=wires), qml.expval(qml.PauliZ(0))
 
     res = circuit()
     assert isinstance(res, tuple) and len(res) == 2
-    assert res[0] == [{basis_state: n_shots} for basis_state in basis_states]
+    for i, basis_state in enumerate(basis_states):
+        assert list(res[0][i].keys()) == [basis_state]
+        assert qml.math.allequal(list(res[0][i].values()), n_shots)
+    # assert res[0] == [{basis_state: expected_n_shots} for basis_state in basis_states]
     assert len(res[1]) == 2 and qml.math.allequal(res[1], 1)
 
-    custom_measurement_process(dev, spy)
-
 
 @pytest.mark.all_interfaces
 @pytest.mark.parametrize("interface", ["autograd", "jax", "tensorflow", "torch"])
-def test_batched_counts_operator_finite_shots(interface, mocker):
+def test_batched_counts_operator_finite_shots(interface):
     """Check all interfaces with observable measurement counts, batching and finite shots"""
     n_shots = 10
-    dev = qml.device("default.qubit.legacy", wires=3, shots=n_shots)
-    spy = mocker.spy(qml.QubitDevice, "sample")
+    dev = qml.device("default.qubit", wires=3, shots=n_shots)
 
     @qml.qnode(dev, interface=interface)
     def circuit():
@@ -813,16 +727,13 @@ def circuit():
 
     assert circuit() == [{-1: n_shots}, {1: n_shots}]
 
-    custom_measurement_process(dev, spy)
-
 
 @pytest.mark.all_interfaces
 @pytest.mark.parametrize("interface", ["autograd", "jax", "tensorflow", "torch"])
-def test_batched_counts_and_expval_operator_finite_shots(interface, mocker):
+def test_batched_counts_and_expval_operator_finite_shots(interface):
     """Check all interfaces with observable measurement counts, batching and finite shots"""
     n_shots = 10
-    dev = qml.device("default.qubit.legacy", wires=3, shots=n_shots)
-    spy = mocker.spy(qml.QubitDevice, "sample")
+    dev = qml.device("default.qubit", wires=3, shots=n_shots)
 
     @qml.qnode(dev, interface=interface)
     def circuit():
@@ -833,5 +744,3 @@ def circuit():
     assert isinstance(res, tuple) and len(res) == 2
     assert res[0] == [{-1: n_shots}, {1: n_shots}]
     assert len(res[1]) == 2 and qml.math.allequal(res[1], [-1, 1])
-
-    custom_measurement_process(dev, spy)
diff --git a/tests/measurements/test_expval.py b/tests/measurements/test_expval.py
index c5a6850d4ae..c50a87b503d 100644
--- a/tests/measurements/test_expval.py
+++ b/tests/measurements/test_expval.py
@@ -21,49 +21,19 @@
 from pennylane.measurements.expval import ExpectationMP
 
 
-# TODO: Remove this when new CustomMP are the default
-def custom_measurement_process(device, spy):
-    assert len(spy.call_args_list) > 0  # make sure method is mocked properly
-
-    samples = device._samples  # pylint: disable=protected-access
-    state = device._state  # pylint: disable=protected-access
-    call_args_list = list(spy.call_args_list)
-    for call_args in call_args_list:
-        obs = call_args.args[1]
-        shot_range, bin_size = (
-            call_args.kwargs["shot_range"],
-            call_args.kwargs["bin_size"],
-        )
-        # no need to use op, because the observable has already been applied to ``self.dev._state``
-        meas = qml.expval(op=obs)
-        old_res = device.expval(obs, shot_range=shot_range, bin_size=bin_size)
-        if device.shots is None:
-            new_res = meas.process_state(state=state, wire_order=device.wires)
-        else:
-            new_res = meas.process_samples(
-                samples=samples, wire_order=device.wires, shot_range=shot_range, bin_size=bin_size
-            )
-        assert qml.math.allclose(old_res, new_res)
-
-
 class TestExpval:
     """Tests for the expval function"""
 
     @pytest.mark.parametrize("shots", [None, 10000, [10000, 10000]])
-    @pytest.mark.parametrize("r_dtype", [np.float32, np.float64])
-    def test_value(self, tol, r_dtype, mocker, shots):
+    def test_value(self, tol, shots):
         """Test that the expval interface works"""
-        dev = qml.device("default.qubit.legacy", wires=2, shots=shots)
-        dev.R_DTYPE = r_dtype
+        dev = qml.device("default.qubit", wires=2, shots=shots)
 
         @qml.qnode(dev, diff_method="parameter-shift")
         def circuit(x):
             qml.RX(x, wires=0)
             return qml.expval(qml.PauliY(0))
 
-        new_dev = circuit.device
-        spy = mocker.spy(qml.QubitDevice, "expval")
-
         x = 0.54
         res = circuit(x)
         expected = -np.sin(x)
@@ -72,36 +42,29 @@ def circuit(x):
         rtol = 0 if shots is None else 0.05
         assert np.allclose(res, expected, atol=atol, rtol=rtol)
 
-        # pylint: disable=no-member, unsubscriptable-object
+        r_dtype = np.float64
         if isinstance(res, tuple):
             assert res[0].dtype == r_dtype
             assert res[1].dtype == r_dtype
         else:
             assert res.dtype == r_dtype
 
-        custom_measurement_process(new_dev, spy)
-
-    def test_not_an_observable(self, mocker):
+    def test_not_an_observable(self):
         """Test that a warning is raised if the provided
         argument might not be hermitian."""
-        dev = qml.device("default.qubit.legacy", wires=2)
+        dev = qml.device("default.qubit", wires=2)
 
         @qml.qnode(dev)
         def circuit():
             qml.RX(0.52, wires=0)
             return qml.expval(qml.prod(qml.PauliX(0), qml.PauliZ(0)))
 
-        new_dev = circuit.device
-        spy = mocker.spy(qml.QubitDevice, "expval")
-
         with pytest.warns(UserWarning, match="Prod might not be hermitian."):
             _ = circuit()
 
-        custom_measurement_process(new_dev, spy)
-
-    def test_observable_return_type_is_expectation(self, mocker):
+    def test_observable_return_type_is_expectation(self):
         """Test that the return type of the observable is :attr:`ObservableReturnTypes.Expectation`"""
-        dev = qml.device("default.qubit.legacy", wires=2)
+        dev = qml.device("default.qubit", wires=2)
 
         @qml.qnode(dev)
         def circuit():
@@ -109,13 +72,8 @@ def circuit():
             assert res.return_type is Expectation
             return res
 
-        new_dev = circuit.device
-        spy = mocker.spy(qml.QubitDevice, "expval")
-
         circuit()
 
-        custom_measurement_process(new_dev, spy)
-
     @pytest.mark.parametrize("shots", [None, 10000, [10000, 10000]])
     @pytest.mark.parametrize("phi", np.arange(0, 2 * np.pi, np.pi / 3))
     def test_observable_is_measurement_value(
@@ -155,7 +113,7 @@ def test_numeric_type(self, obs):
     )
     def test_shape(self, obs):
         """Test that the shape is correct."""
-        dev = qml.device("default.qubit.legacy", wires=1)
+        dev = qml.device("default.qubit", wires=1)
 
         res = qml.expval(obs)
         # pylint: disable=use-implicit-booleaness-not-comparison
@@ -170,15 +128,15 @@ def test_shape_shot_vector(self, obs):
         """Test that the shape is correct with the shot vector too."""
         res = qml.expval(obs)
         shot_vector = (1, 2, 3)
-        dev = qml.device("default.qubit.legacy", wires=3, shots=shot_vector)
+        dev = qml.device("default.qubit", wires=3, shots=shot_vector)
         assert res.shape(dev, Shots(shot_vector)) == ((), (), ())
 
     @pytest.mark.parametrize("state", [np.array([0, 0, 0]), np.array([1, 0, 0, 0, 0, 0, 0, 0])])
     @pytest.mark.parametrize("shots", [None, 1000, [1000, 10000]])
-    def test_projector_expval(self, state, shots, mocker):
+    def test_projector_expval(self, state, shots):
         """Tests that the expectation of a ``Projector`` object is computed correctly for both of
         its subclasses."""
-        dev = qml.device("default.qubit.legacy", wires=3, shots=shots)
+        dev = qml.device("default.qubit", wires=3, shots=shots)
         np.random.seed(42)
 
         @qml.qnode(dev)
@@ -186,24 +144,18 @@ def circuit():
             qml.Hadamard(0)
             return qml.expval(qml.Projector(state, wires=range(3)))
 
-        new_dev = circuit.device
-        spy = mocker.spy(qml.QubitDevice, "expval")
-
         res = circuit()
         expected = [0.5, 0.5] if isinstance(shots, list) else 0.5
         assert np.allclose(res, expected, atol=0.02, rtol=0.02)
 
-        custom_measurement_process(new_dev, spy)
-
-    def test_permuted_wires(self, mocker):
+    def test_permuted_wires(self):
         """Test that the expectation value of an operator with permuted wires is the same."""
         obs = qml.prod(qml.PauliZ(8), qml.s_prod(2, qml.PauliZ(10)), qml.s_prod(3, qml.PauliZ("h")))
         obs_2 = qml.prod(
             qml.s_prod(3, qml.PauliZ("h")), qml.PauliZ(8), qml.s_prod(2, qml.PauliZ(10))
         )
 
-        dev = qml.device("default.qubit.legacy", wires=["h", 8, 10])
-        spy = mocker.spy(qml.QubitDevice, "expval")
+        dev = qml.device("default.qubit", wires=["h", 8, 10])
 
         @qml.qnode(dev)
         def circuit():
@@ -218,7 +170,6 @@ def circuit2():
             return qml.expval(obs_2)
 
         assert circuit() == circuit2()
-        custom_measurement_process(dev, spy)
 
     def test_copy_observable(self):
         """Test that the observable is copied if present."""
diff --git a/tests/measurements/test_measurements.py b/tests/measurements/test_measurements.py
index a88b72f6391..ad06fdd5346 100644
--- a/tests/measurements/test_measurements.py
+++ b/tests/measurements/test_measurements.py
@@ -78,7 +78,7 @@ def test_ObservableReturnTypes(return_type, value):
 def test_no_measure():
     """Test that failing to specify a measurement
     raises an exception"""
-    dev = qml.device("default.qubit.legacy", wires=2)
+    dev = qml.device("default.qubit", wires=2)
 
     @qml.qnode(dev)
     def circuit(x):
@@ -104,7 +104,7 @@ def test_numeric_type_unrecognized_error():
 def test_shape_unrecognized_error():
     """Test that querying the shape of a measurement process with an
     unrecognized return type raises an error."""
-    dev = qml.device("default.qubit.legacy", wires=2)
+    dev = qml.device("default.qubit", wires=2)
     mp = NotValidMeasurement()
     with pytest.raises(
         qml.QuantumFunctionError,
@@ -244,7 +244,7 @@ def test_not_an_observable(self, stat_func, return_type):  # pylint: disable=unu
         if stat_func is sample:
             pytest.skip("Sampling is not yet supported with symbolic operators.")
 
-        dev = qml.device("default.qubit.legacy", wires=2)
+        dev = qml.device("default.qubit", wires=2)
 
         @qml.qnode(dev)
         def circuit():
@@ -489,10 +489,10 @@ def test_custom_sample_measurement(self):
 
         class MyMeasurement(SampleMeasurement):
             # pylint: disable=signature-differs
-            def process_samples(self, samples, wire_order, shot_range, bin_size):
+            def process_samples(self, samples, wire_order, shot_range=None, bin_size=None):
                 return qml.math.sum(samples[..., self.wires])
 
-        dev = qml.device("default.qubit.legacy", wires=2, shots=1000)
+        dev = qml.device("default.qubit", wires=2, shots=1000)
 
         @qml.qnode(dev)
         def circuit():
@@ -501,6 +501,7 @@ def circuit():
 
         assert qml.math.allequal(circuit(), [1000, 0])
 
+    @pytest.mark.xfail(reason="until DQ2 port")
     def test_sample_measurement_without_shots(self):
         """Test that executing a sampled measurement with ``shots=None`` raises an error."""
 
@@ -509,7 +510,11 @@ class MyMeasurement(SampleMeasurement):
             def process_samples(self, samples, wire_order, shot_range, bin_size):
                 return qml.math.sum(samples[..., self.wires])
 
-        dev = qml.device("default.qubit.legacy", wires=2)
+            @property
+            def return_type(self):
+                return Sample
+
+        dev = qml.device("default.qubit", wires=2)
 
         @qml.qnode(dev)
         def circuit():
@@ -517,25 +522,11 @@ def circuit():
             return MyMeasurement(wires=[0]), MyMeasurement(wires=[1])
 
         with pytest.raises(
-            ValueError, match="Shots must be specified in the device to compute the measurement "
+            qml.DeviceError,
+            match="Analytic circuits must only contain StateMeasurements; got sample",
         ):
             circuit()
 
-    def test_method_overriden_by_device(self):
-        """Test that the device can override a measurement process."""
-
-        dev = qml.device("default.qubit.legacy", wires=2, shots=1000)
-
-        @qml.qnode(dev)
-        def circuit():
-            qml.PauliX(0)
-            return qml.sample(wires=[0]), qml.sample(wires=[1])
-
-        circuit.device.measurement_map[SampleMP] = "test_method"
-        circuit.device.test_method = lambda obs, shot_range=None, bin_size=None: 2
-
-        assert qml.math.allequal(circuit(), [2, 2])
-
 
 class TestStateMeasurement:
     """Tests for the SampleMeasurement class."""
@@ -547,7 +538,7 @@ class MyMeasurement(StateMeasurement):
             def process_state(self, state, wire_order):
                 return qml.math.sum(state)
 
-        dev = qml.device("default.qubit.legacy", wires=2)
+        dev = qml.device("default.qubit", wires=2)
 
         @qml.qnode(dev)
         def circuit():
@@ -555,68 +546,47 @@ def circuit():
 
         assert circuit() == 1
 
-    def test_sample_measurement_with_shots(self):
-        """Test that executing a state measurement with shots raises a warning."""
+    @pytest.mark.xfail(reason="until DQ2 port")
+    def test_state_measurement_with_shots(self):
+        """Test that executing a state measurement with shots raises an error."""
 
         class MyMeasurement(StateMeasurement):
             def process_state(self, state, wire_order):
                 return qml.math.sum(state)
 
-        dev = qml.device("default.qubit.legacy", wires=2, shots=1000)
+            @property
+            def return_type(self):
+                return State
+
+        dev = qml.device("default.qubit", wires=2, shots=1000)
 
         @qml.qnode(dev)
         def circuit():
             return MyMeasurement()
 
-        with pytest.warns(
-            UserWarning,
-            match="Requested measurement MyMeasurement with finite shots",
-        ):
+        with pytest.raises(qml.DeviceError, match="Circuits with finite shots must only contain"):
             circuit()
 
-    def test_method_overriden_by_device(self):
-        """Test that the device can override a measurement process."""
-
-        dev = qml.device("default.qubit.legacy", wires=2)
-
-        @qml.qnode(dev, interface="autograd")
-        def circuit():
-            return qml.state()
-
-        circuit.device.measurement_map[StateMP] = "test_method"
-        circuit.device.test_method = lambda obs, shot_range=None, bin_size=None: 2
-
-        assert circuit() == 2
-
 
 class TestMeasurementTransform:
     """Tests for the MeasurementTransform class."""
 
+    @pytest.mark.xfail(reason="until DQ2 port")
     def test_custom_measurement(self):
         """Test the execution of a custom measurement."""
 
-        class MyMeasurement(MeasurementTransform):
+        class CountTapesMP(MeasurementTransform, SampleMeasurement):
             def process(self, tape, device):
-                return {device.shots: len(tape)}
-
-        dev = qml.device("default.qubit.legacy", wires=2, shots=1000)
+                tapes, _, _ = device.preprocess(tape)
+                return len(tapes)
 
-        @qml.qnode(dev)
-        def circuit():
-            return MyMeasurement()
-
-        assert circuit() == {dev.shots: len(circuit.tape)}
+            def process_samples(self, samples, wire_order, shot_range=None, bin_size=None):
+                return [True]
 
-    def test_method_overriden_by_device(self):
-        """Test that the device can override a measurement process."""
-
-        dev = qml.device("default.qubit.legacy", wires=2, shots=1000)
+        dev = qml.device("default.qubit", wires=2, shots=1000)
 
         @qml.qnode(dev)
         def circuit():
-            return qml.classical_shadow(wires=0)
-
-        circuit.device.measurement_map[ClassicalShadowMP] = "test_method"
-        circuit.device.test_method = lambda tape: 2
+            return CountTapesMP(wires=[0])
 
-        assert circuit() == 2
+        assert circuit() == 1
diff --git a/tests/measurements/test_mutual_info.py b/tests/measurements/test_mutual_info.py
index d75ce6de280..0473a71f7f0 100644
--- a/tests/measurements/test_mutual_info.py
+++ b/tests/measurements/test_mutual_info.py
@@ -18,7 +18,6 @@
 import pytest
 
 import pennylane as qml
-from pennylane.interfaces import INTERFACE_MAP
 from pennylane.measurements import MutualInfo, Shots
 from pennylane.measurements.mutual_info import MutualInfoMP
 from pennylane.wires import Wires
@@ -39,7 +38,7 @@ def test_queue(self):
     @pytest.mark.parametrize("shots, shape", [(None, ()), (10, ()), ([1, 10], ((), ()))])
     def test_shape(self, shots, shape):
         """Test that the shape is correct."""
-        dev = qml.device("default.qubit.legacy", wires=3, shots=shots)
+        dev = qml.device("default.qubit", wires=3, shots=shots)
         res = qml.mutual_info(wires0=[0], wires1=[1])
         assert res.shape(dev, Shots(shots)) == shape
 
@@ -71,6 +70,18 @@ def test_hash(self):
         m4 = MutualInfoMP(wires=[Wires((0)), Wires((1, 2))])
         assert m3.hash != m4.hash
 
+    def test_map_wires(self):
+        """Test that map_wires works as expected."""
+        mapped1 = MutualInfoMP(wires=[Wires([0]), Wires([1])]).map_wires({0: 1, 1: 0})
+        assert mapped1.raw_wires == [Wires([1]), Wires([0])]
+        assert qml.equal(mapped1, MutualInfoMP(wires=[Wires([1]), Wires([0])]))
+
+        mapped2 = MutualInfoMP(wires=[Wires(["a", "b"]), Wires(["c"])]).map_wires(
+            {"a": 0, "b": 1, "c": 2}
+        )
+        assert mapped2.raw_wires == [Wires([0, 1]), Wires([2])]
+        assert qml.equal(mapped2, MutualInfoMP(wires=[Wires([0, 1]), Wires([2])]))
+
 
 class TestIntegration:
     """Tests for the mutual information functions"""
@@ -88,7 +99,7 @@ class TestIntegration:
     )
     def test_mutual_info_output(self, interface, state, expected):
         """Test the output of qml.mutual_info"""
-        dev = qml.device("default.qubit.legacy", wires=4)
+        dev = qml.device("default.qubit", wires=4)
 
         @qml.qnode(dev, interface=interface)
         def circuit():
@@ -96,17 +107,13 @@ def circuit():
             return qml.mutual_info(wires0=[0, 2], wires1=[1, 3])
 
         res = circuit()
-        new_res = qml.mutual_info(wires0=[0, 2], wires1=[1, 3]).process_state(
-            state=circuit.device.state, wire_order=circuit.device.wires
-        )
         assert np.allclose(res, expected, atol=1e-6)
-        assert np.allclose(new_res, expected, atol=1e-6)
-        assert INTERFACE_MAP.get(qml.math.get_interface(new_res)) == interface
-        assert res.dtype == new_res.dtype  # pylint: disable=no-member
 
-    def test_shot_vec_error(self):
+    @pytest.mark.xfail(reason="until DQ2 port")
+    @pytest.mark.parametrize("shots", [1000, [1, 10, 10, 1000]])
+    def test_finite_shots_error(self, shots):
         """Test an error is raised when using shot vectors with mutual_info."""
-        dev = qml.device("default.qubit.legacy", wires=2, shots=[1, 10, 10, 1000])
+        dev = qml.device("default.qubit", wires=2, shots=shots)
 
         @qml.qnode(device=dev)
         def circuit(x):
@@ -114,15 +121,13 @@ def circuit(x):
             qml.CRX(x, wires=[0, 1])
             return qml.mutual_info(wires0=[0], wires1=[1])
 
-        with pytest.raises(
-            NotImplementedError, match="mutual information is not supported with shot vectors"
-        ):
+        with pytest.raises(qml.DeviceError, match="Circuits with finite shots must only contain"):
             circuit(0.5)
 
     diff_methods = ["backprop", "finite-diff"]
 
     @pytest.mark.all_interfaces
-    @pytest.mark.parametrize("device", ["default.qubit.legacy", "default.mixed", "lightning.qubit"])
+    @pytest.mark.parametrize("device", ["default.qubit", "default.mixed", "lightning.qubit"])
     @pytest.mark.parametrize("interface", ["autograd", "jax", "tensorflow", "torch"])
     @pytest.mark.parametrize("params", np.linspace(0, 2 * np.pi, 8))
     def test_qnode_state(self, device, interface, params):
@@ -148,7 +153,7 @@ def circuit(params):
         assert np.allclose(actual, expected)
 
     @pytest.mark.all_interfaces
-    @pytest.mark.parametrize("device", ["default.qubit.legacy", "default.mixed", "lightning.qubit"])
+    @pytest.mark.parametrize("device", ["default.qubit", "default.mixed", "lightning.qubit"])
     @pytest.mark.parametrize("interface", ["autograd", "jax", "tensorflow", "torch"])
     @pytest.mark.parametrize("params", zip(np.linspace(0, np.pi, 8), np.linspace(0, 2 * np.pi, 8)))
     def test_qnode_mutual_info(self, device, interface, params):
@@ -179,7 +184,7 @@ def circuit_state(params):
 
         assert np.allclose(actual, expected)
 
-    @pytest.mark.parametrize("device", ["default.qubit.legacy", "default.mixed", "lightning.qubit"])
+    @pytest.mark.parametrize("device", ["default.qubit", "default.mixed", "lightning.qubit"])
     def test_mutual_info_wire_labels(self, device):
         """Test that mutual_info is correct with custom wire labels"""
         param = np.array([0.678, 1.234])
@@ -209,7 +214,7 @@ def test_qnode_state_jax_jit(self, params):
         import jax
         import jax.numpy as jnp
 
-        dev = qml.device("default.qubit.legacy", wires=2)
+        dev = qml.device("default.qubit", wires=2)
 
         params = jnp.array(params)
 
@@ -237,7 +242,7 @@ def test_qnode_mutual_info_jax_jit(self, params, interface):
         import jax
         import jax.numpy as jnp
 
-        dev = qml.device("default.qubit.legacy", wires=2)
+        dev = qml.device("default.qubit", wires=2)
 
         params = jnp.array(params)
 
@@ -269,7 +274,7 @@ def circuit_state(params):
     def test_qnode_grad(self, param, diff_method, interface):
         """Test that the gradient of mutual information works for QNodes
         with the autograd interface"""
-        dev = qml.device("default.qubit.legacy", wires=2)
+        dev = qml.device("default.qubit", wires=2)
 
         @qml.qnode(dev, interface=interface, diff_method=diff_method)
         def circuit(param):
@@ -301,7 +306,7 @@ def test_qnode_grad_jax(self, param, diff_method, interface):
         import jax
         import jax.numpy as jnp
 
-        dev = qml.device("default.qubit.legacy", wires=2)
+        dev = qml.device("default.qubit", wires=2)
 
         param = jnp.array(param)
 
@@ -335,7 +340,7 @@ def test_qnode_grad_jax_jit(self, param, diff_method, interface):
         import jax
         import jax.numpy as jnp
 
-        dev = qml.device("default.qubit.legacy", wires=2)
+        dev = qml.device("default.qubit", wires=2)
 
         param = jnp.array(param)
 
@@ -368,7 +373,7 @@ def test_qnode_grad_tf(self, param, diff_method, interface):
         with the tensorflow interface"""
         import tensorflow as tf
 
-        dev = qml.device("default.qubit.legacy", wires=2)
+        dev = qml.device("default.qubit", wires=2)
 
         param = tf.Variable(param)
 
@@ -404,7 +409,7 @@ def test_qnode_grad_torch(self, param, diff_method, interface):
         with the torch interface"""
         import torch
 
-        dev = qml.device("default.qubit.legacy", wires=2)
+        dev = qml.device("default.qubit", wires=2)
 
         @qml.qnode(dev, interface=interface, diff_method=diff_method)
         def circuit(param):
@@ -431,7 +436,7 @@ def circuit(param):
         assert np.allclose(actual, expected, atol=tol)
 
     @pytest.mark.all_interfaces
-    @pytest.mark.parametrize("device", ["default.qubit.legacy", "default.mixed", "lightning.qubit"])
+    @pytest.mark.parametrize("device", ["default.qubit", "default.mixed", "lightning.qubit"])
     @pytest.mark.parametrize("interface", ["autograd", "jax", "tensorflow", "torch"])
     @pytest.mark.parametrize(
         "params", [np.array([0.0, 0.0]), np.array([0.3, 0.4]), np.array([0.6, 0.8])]
@@ -454,16 +459,14 @@ def circuit(params):
         with pytest.raises(qml.QuantumFunctionError, match=msg):
             circuit(params)
 
+    @pytest.mark.xfail(reason="until DQ2 port")
     @pytest.mark.all_interfaces
-    @pytest.mark.parametrize("device", ["default.qubit.legacy", "default.mixed", "lightning.qubit"])
     @pytest.mark.parametrize("interface", ["autograd", "jax", "tensorflow", "torch"])
-    @pytest.mark.parametrize(
-        "params", [np.array([0.0, 0.0]), np.array([0.3, 0.4]), np.array([0.6, 0.8])]
-    )
-    def test_custom_wire_labels_error(self, device, interface, params):
-        """Tests that an error is raised when mutual information is measured
+    @pytest.mark.parametrize("params", [np.array([0.0, 0.0]), np.array([0.3, 0.4])])
+    def test_custom_wire_labels_works(self, interface, params):
+        """Tests that no error is raised when mutual information is measured
         with custom wire labels"""
-        dev = qml.device(device, wires=["a", "b"])
+        dev = qml.device("default.qubit", wires=["a", "b"])
 
         params = qml.math.asarray(params, like=interface)
 
@@ -474,6 +477,13 @@ def circuit(params):
             qml.CNOT(wires=["a", "b"])
             return qml.mutual_info(wires0=["a"], wires1=["b"])
 
-        msg = "Returning the mutual information is not supported when using custom wire labels"
-        with pytest.raises(qml.QuantumFunctionError, match=msg):
-            circuit(params)
+        @qml.qnode(dev, interface=interface)
+        def circuit_expected(params):
+            qml.RY(params[0], wires="a")
+            qml.RY(params[1], wires="b")
+            qml.CNOT(wires=["a", "b"])
+            return qml.state()
+
+        actual = circuit(params)
+        expected = qml.qinfo.mutual_info(circuit_expected, wires0=["a"], wires1=["b"])(params)
+        assert np.allclose(actual, expected)
diff --git a/tests/measurements/test_probs.py b/tests/measurements/test_probs.py
index fd3fc698920..4ec43054914 100644
--- a/tests/measurements/test_probs.py
+++ b/tests/measurements/test_probs.py
@@ -26,34 +26,6 @@
 from pennylane.queuing import AnnotatedQueue
 
 
-# TODO: Remove this when new CustomMP are the default
-def custom_measurement_process(device, spy):
-    assert len(spy.call_args_list) > 0  # make sure method is mocked properly
-
-    samples = device._samples  # pylint: disable=protected-access
-    state = device._state  # pylint: disable=protected-access
-    call_args_list = list(spy.call_args_list)
-    for call_args in call_args_list:
-        wires, shot_range, bin_size = (
-            call_args.kwargs["wires"],
-            call_args.kwargs["shot_range"],
-            call_args.kwargs["bin_size"],
-        )
-        # no need to use op, because the observable has already been applied to ``dev._state``
-        meas = qml.probs(wires=wires)
-        old_res = device.probability(wires=wires, shot_range=shot_range, bin_size=bin_size)
-        if device.shots is None:
-            new_res = meas.process_state(state=state, wire_order=device.wires)
-        else:
-            new_res = meas.process_samples(
-                samples=samples,
-                wire_order=device.wires,
-                shot_range=shot_range,
-                bin_size=bin_size,
-            )
-        assert qml.math.allequal(old_res, new_res)
-
-
 # make the test deterministic
 np.random.seed(42)
 
@@ -78,7 +50,7 @@ class TestProbs:
 
     def test_queue(self):
         """Test that the right measurement class is queued."""
-        dev = qml.device("default.qubit.legacy", wires=2)
+        dev = qml.device("default.qubit", wires=2)
 
         @qml.qnode(dev)
         def circuit():
@@ -97,7 +69,7 @@ def test_numeric_type(self):
     @pytest.mark.parametrize("shots", [None, 10])
     def test_shape(self, wires, shots):
         """Test that the shape is correct."""
-        dev = qml.device("default.qubit.legacy", wires=3, shots=shots)
+        dev = qml.device("default.qubit", wires=3, shots=shots)
         res = qml.probs(wires=wires)
         assert res.shape(dev, Shots(shots)) == (2 ** len(wires),)
 
@@ -106,7 +78,7 @@ def test_shape_shot_vector(self, wires):
         """Test that the shape is correct with the shot vector too."""
         res = qml.probs(wires=wires)
         shot_vector = (1, 2, 3)
-        dev = qml.device("default.qubit.legacy", wires=3, shots=shot_vector)
+        dev = qml.device("default.qubit", wires=3, shots=shot_vector)
         assert res.shape(dev, Shots(shot_vector)) == (
             (2 ** len(wires),),
             (2 ** len(wires),),
@@ -133,20 +105,20 @@ def test_probs_empty_wires(self):
     @pytest.mark.parametrize("shots", [None, 100])
     def test_probs_no_arguments(self, shots):
         """Test that using ``qml.probs`` with no arguments returns the probabilities of all wires."""
-        dev = qml.device("default.qubit.legacy", wires=3, shots=shots)
+        dev = qml.device("default.qubit", wires=3, shots=shots)
 
         @qml.qnode(dev)
         def circuit():
+            qml.Identity(wires=[0, 1, 2])
             return qml.probs()
 
         res = circuit()
 
         assert qml.math.allequal(res, [1, 0, 0, 0, 0, 0, 0, 0])
 
-    def test_full_prob(self, init_state, tol, mocker):
+    def test_full_prob(self, init_state, tol):
         """Test that the correct probability is returned."""
-        dev = qml.device("default.qubit.legacy", wires=4)
-        spy = mocker.spy(qml.QubitDevice, "probability")
+        dev = qml.device("default.qubit", wires=4)
 
         state = init_state(4)
 
@@ -159,12 +131,9 @@ def circuit():
         expected = np.abs(state) ** 2
         assert np.allclose(res, expected, atol=tol, rtol=0)
 
-        custom_measurement_process(dev, spy)
-
-    def test_marginal_prob(self, init_state, tol, mocker):
+    def test_marginal_prob(self, init_state, tol):
         """Test that the correct marginal probability is returned."""
-        dev = qml.device("default.qubit.legacy", wires=4)
-        spy = mocker.spy(qml.QubitDevice, "probability")
+        dev = qml.device("default.qubit", wires=4)
 
         state = init_state(4)
 
@@ -178,13 +147,10 @@ def circuit():
         expected = np.einsum("ijkl->jl", expected).flatten()
         assert np.allclose(res, expected, atol=tol, rtol=0)
 
-        custom_measurement_process(dev, spy)
-
-    def test_marginal_prob_more_wires(self, init_state, mocker, tol):
+    def test_marginal_prob_more_wires(self, init_state, tol):
         """Test that the correct marginal probability is returned, when the
         states_to_binary method is used for probability computations."""
-        dev = qml.device("default.qubit.legacy", wires=4)
-        spy_probs = mocker.spy(qml.QubitDevice, "probability")
+        dev = qml.device("default.qubit", wires=4)
         state = init_state(4)
 
         @qml.qnode(dev)
@@ -198,8 +164,6 @@ def circuit():
         expected = np.einsum("ijkl->jil", expected).flatten()
         assert np.allclose(res, expected, atol=tol, rtol=0)
 
-        custom_measurement_process(dev, spy_probs)
-
     @pytest.mark.all_interfaces
     @pytest.mark.parametrize("interface", ["numpy", "jax", "torch", "tensorflow"])
     @pytest.mark.parametrize(
@@ -242,10 +206,9 @@ def test_process_state_batched(self, interface, subset_wires, expected):
         assert subset_probs.shape == qml.math.shape(expected)
         assert qml.math.allclose(subset_probs, expected)
 
-    def test_integration(self, tol, mocker):
+    def test_integration(self, tol):
         """Test the probability is correct for a known state preparation."""
-        dev = qml.device("default.qubit.legacy", wires=2)
-        spy = mocker.spy(qml.QubitDevice, "probability")
+        dev = qml.device("default.qubit", wires=2)
 
         @qml.qnode(dev)
         def circuit():
@@ -260,14 +223,11 @@ def circuit():
         expected = np.array([0.5, 0.5, 0, 0])
         assert np.allclose(res, expected, atol=tol, rtol=0)
 
-        custom_measurement_process(dev, spy)
-
     @pytest.mark.parametrize("shots", [100, [1, 10, 100]])
-    def test_integration_analytic_false(self, tol, mocker, shots):
+    def test_integration_analytic_false(self, tol, shots):
         """Test the probability is correct for a known state preparation when the
         analytic attribute is set to False."""
-        dev = qml.device("default.qubit.legacy", wires=3, shots=shots)
-        spy = mocker.spy(qml.QubitDevice, "probability")
+        dev = qml.device("default.qubit", wires=3, shots=shots)
 
         @qml.qnode(dev)
         def circuit():
@@ -278,12 +238,10 @@ def circuit():
         expected = np.array([0, 0, 0, 0, 1, 0, 0, 0])
         assert np.allclose(res, expected, atol=tol, rtol=0)
 
-        custom_measurement_process(dev, spy)
-
     @pytest.mark.parametrize("shots", [None, 10000, [10000, 10000]])
     @pytest.mark.parametrize("phi", np.arange(0, 2 * np.pi, np.pi / 3))
     def test_observable_is_measurement_value(
-        self, shots, phi, mocker, tol, tol_stochastic
+        self, shots, phi, tol, tol_stochastic
     ):  # pylint: disable=too-many-arguments
         """Test that expectation values for mid-circuit measurement values
         are correct for a single measurement value."""
@@ -295,9 +253,6 @@ def circuit(phi):
             m0 = qml.measure(0)
             return qml.probs(op=m0)
 
-        new_dev = circuit.device
-        spy = mocker.spy(qml.QubitDevice, "probability")
-
         res = circuit(phi)
 
         atol = tol if shots is None else tol_stochastic
@@ -309,13 +264,10 @@ def circuit(phi):
             for r in res:  # pylint: disable=not-an-iterable
                 assert np.allclose(r, expected, atol=atol, rtol=0)
 
-        custom_measurement_process(new_dev, spy)
-
     @pytest.mark.parametrize("shots", [None, 100])
-    def test_batch_size(self, mocker, shots):
+    def test_batch_size(self, shots):
         """Test the probability is correct for a batched input."""
-        dev = qml.device("default.qubit.legacy", wires=1, shots=shots)
-        spy = mocker.spy(qml.QubitDevice, "probability")
+        dev = qml.device("default.qubit", wires=1, shots=shots)
 
         @qml.qnode(dev)
         def circuit(x):
@@ -327,14 +279,12 @@ def circuit(x):
         expected = [[1.0, 0.0], [0.5, 0.5]]
         assert np.allclose(res, expected, atol=0.1, rtol=0.1)
 
-        custom_measurement_process(dev, spy)
-
     @pytest.mark.tf
     @pytest.mark.parametrize(
         "prep_op, expected",
         [
             (None, [1, 0]),
-            (qml.QubitStateVector([[0, 1], [1, 0]], 0), [[0, 1], [1, 0]]),
+            (qml.StatePrep([[0, 1], [1, 0]], 0), [[0, 1], [1, 0]]),
         ],
     )
     def test_process_state_tf_autograph(self, prep_op, expected):
@@ -351,12 +301,11 @@ def probs_from_state(state):
         assert np.allclose(probs_from_state(state), expected)
 
     @pytest.mark.autograd
-    def test_numerical_analytic_diff_agree(self, tol, mocker):
+    def test_numerical_analytic_diff_agree(self, tol):
         """Test that the finite difference and parameter shift rule
         provide the same Jacobian."""
         w = 4
-        dev = qml.device("default.qubit.legacy", wires=w)
-        spy = mocker.spy(qml.QubitDevice, "probability")
+        dev = qml.device("default.qubit", wires=w)
 
         def circuit(x, y, z):
             for i in range(w):
@@ -386,13 +335,10 @@ def circuit(x, y, z):
         # Check that they agree up to numeric tolerance
         assert all(np.allclose(_rF, _rA, atol=tol, rtol=0) for _rF, _rA in zip(res_F, res_A))
 
-        custom_measurement_process(dev, spy)
-
     @pytest.mark.parametrize("hermitian", [1 / np.sqrt(2) * np.array([[1, 1], [1, -1]])])
-    def test_prob_generalize_param_one_qubit(self, hermitian, tol, mocker):
+    def test_prob_generalize_param_one_qubit(self, hermitian, tol):
         """Test that the correct probability is returned."""
-        dev = qml.device("default.qubit.legacy", wires=1)
-        spy = mocker.spy(qml.QubitDevice, "probability")
+        dev = qml.device("default.qubit", wires=1)
 
         @qml.qnode(dev)
         def circuit(x):
@@ -413,13 +359,10 @@ def circuit_rotated(x):
 
         assert np.allclose(res, expected, atol=tol, rtol=0)
 
-        custom_measurement_process(dev, spy)
-
     @pytest.mark.parametrize("hermitian", [1 / np.sqrt(2) * np.array([[1, 1], [1, -1]])])
-    def test_prob_generalize_param(self, hermitian, tol, mocker):
+    def test_prob_generalize_param(self, hermitian, tol):
         """Test that the correct probability is returned."""
-        dev = qml.device("default.qubit.legacy", wires=3)
-        spy = mocker.spy(qml.QubitDevice, "probability")
+        dev = qml.device("default.qubit", wires=3)
 
         @qml.qnode(dev)
         def circuit(x, y):
@@ -445,13 +388,10 @@ def circuit_rotated(x, y):
         expected = np.einsum("ijk->i", expected).flatten()
         assert np.allclose(res, expected, atol=tol, rtol=0)
 
-        custom_measurement_process(dev, spy)
-
     @pytest.mark.parametrize("hermitian", [1 / np.sqrt(2) * np.array([[1, 1], [1, -1]])])
-    def test_prob_generalize_param_multiple(self, hermitian, tol, mocker):
+    def test_prob_generalize_param_multiple(self, hermitian, tol):
         """Test that the correct probability is returned."""
-        dev = qml.device("default.qubit.legacy", wires=3)
-        spy = mocker.spy(qml.QubitDevice, "probability")
+        dev = qml.device("default.qubit", wires=3)
 
         @qml.qnode(dev)
         def circuit(x, y):
@@ -488,15 +428,13 @@ def circuit_rotated(x, y):
         assert np.allclose(res[1], expected_1, atol=tol, rtol=0)
         assert np.allclose(res[2], expected_2, atol=tol, rtol=0)
 
-        custom_measurement_process(dev, spy)
-
     @pytest.mark.parametrize("hermitian", [1 / np.sqrt(2) * np.array([[1, 1], [1, -1]])])
     @pytest.mark.parametrize("wire", [0, 1, 2, 3])
-    def test_prob_generalize_initial_state(self, hermitian, wire, init_state, tol, mocker):
+    def test_prob_generalize_initial_state(self, hermitian, wire, init_state, tol):
         """Test that the correct probability is returned."""
         # pylint:disable=too-many-arguments
-        dev = qml.device("default.qubit.legacy", wires=4)
-        spy = mocker.spy(qml.QubitDevice, "probability")
+        dev = qml.device("default.qubit", wires=4)
+
         state = init_state(4)
 
         @qml.qnode(dev)
@@ -532,15 +470,13 @@ def circuit_rotated():
 
         assert np.allclose(res, expected, atol=tol, rtol=0)
 
-        custom_measurement_process(dev, spy)
-
     @pytest.mark.parametrize("operation", [qml.PauliX, qml.PauliY, qml.Hadamard])
     @pytest.mark.parametrize("wire", [0, 1, 2, 3])
-    def test_operation_prob(self, operation, wire, init_state, tol, mocker):
+    def test_operation_prob(self, operation, wire, init_state, tol):
         "Test the rotated probability with different wires and rotating operations."
         # pylint:disable=too-many-arguments
-        dev = qml.device("default.qubit.legacy", wires=4)
-        spy = mocker.spy(qml.QubitDevice, "probability")
+        dev = qml.device("default.qubit", wires=4)
+
         state = init_state(4)
 
         @qml.qnode(dev)
@@ -576,13 +512,11 @@ def circuit_rotated():
 
         assert np.allclose(res, expected, atol=tol, rtol=0)
 
-        custom_measurement_process(dev, spy)
-
     @pytest.mark.parametrize("observable", [(qml.PauliX, qml.PauliY)])
-    def test_observable_tensor_prob(self, observable, init_state, tol, mocker):
+    def test_observable_tensor_prob(self, observable, init_state, tol):
         "Test the rotated probability with a tensor observable."
-        dev = qml.device("default.qubit.legacy", wires=4)
-        spy = mocker.spy(qml.QubitDevice, "probability")
+        dev = qml.device("default.qubit", wires=4)
+
         state = init_state(4)
 
         @qml.qnode(dev)
@@ -612,14 +546,12 @@ def circuit_rotated():
 
         assert np.allclose(res, expected, atol=tol, rtol=0)
 
-        custom_measurement_process(dev, spy)
-
     @pytest.mark.parametrize("coeffs, obs", [([1, 1], [qml.PauliX(wires=0), qml.PauliX(wires=1)])])
     def test_hamiltonian_error(self, coeffs, obs, init_state):
         "Test that an error is returned for hamiltonians."
         H = qml.Hamiltonian(coeffs, obs)
 
-        dev = qml.device("default.qubit.legacy", wires=4)
+        dev = qml.device("default.qubit", wires=4)
         state = init_state(4)
 
         @qml.qnode(dev)
@@ -644,7 +576,7 @@ def test_generalize_prob_not_hermitian(self, operation):
         """Test that Operators that do not have a diagonalizing_gates representation cannot
         be used in probability measurements."""
 
-        dev = qml.device("default.qubit.legacy", wires=2)
+        dev = qml.device("default.qubit", wires=2)
 
         @qml.qnode(dev)
         def circuit():
@@ -662,7 +594,7 @@ def circuit():
     def test_prob_wires_and_hermitian(self, hermitian):
         """Test that we can cannot give simultaneously wires and a hermitian."""
 
-        dev = qml.device("default.qubit.legacy", wires=2)
+        dev = qml.device("default.qubit", wires=2)
 
         @qml.qnode(dev)
         def circuit():
@@ -713,24 +645,24 @@ def test_estimate_probability_with_binsize_with_broadcasting(self, wires, expect
 
         assert np.allclose(res, expected)
 
-    def test_non_commuting_probs_raises_error(self):
-        dev = qml.device("default.qubit.legacy", wires=5)
+    @pytest.mark.xfail(reason="until DQ2 port")
+    def test_non_commuting_probs_does_not_raises_error(self):
+        """Tests that non-commuting probs with expval does not raise an error."""
+        dev = qml.device("default.qubit", wires=5)
 
         @qml.qnode(dev)
         def circuit(x, y):
             qml.RX(x, wires=[0])
             qml.RY(y, wires=[1])
             qml.CNOT(wires=[0, 1])
-            return qml.expval(qml.PauliX(0)), qml.probs(wires=[0, 1])
+            return qml.expval(qml.PauliX(1)), qml.probs(wires=[0, 1])
 
-        with pytest.raises(
-            qml.QuantumFunctionError, match="Only observables that are qubit-wise commuting"
-        ):
-            circuit(1, 2)
+        res = circuit(1, 2)
+        assert isinstance(res, tuple) and len(res) == 2
 
     def test_commuting_probs_in_computational_basis(self):
         """Test that `qml.probs` can be used in the computational basis with other commuting observables."""
-        dev = qml.device("default.qubit.legacy", wires=5)
+        dev = qml.device("default.qubit", wires=5)
 
         @qml.qnode(dev)
         def circuit(x, y):
diff --git a/tests/measurements/test_purity_measurement.py b/tests/measurements/test_purity_measurement.py
index 969f9a16c14..ccaab4fb5bc 100644
--- a/tests/measurements/test_purity_measurement.py
+++ b/tests/measurements/test_purity_measurement.py
@@ -71,15 +71,15 @@ def test_numeric_type(self):
     def test_shape_new(self, shots, shape):
         """Test the ``shape_new`` method."""
         meas = qml.purity(wires=0)
-        dev = qml.device("default.qubit.legacy", wires=1, shots=shots)
+        dev = qml.device("default.qubit", wires=1, shots=shots)
         assert meas.shape(dev, Shots(shots)) == shape
 
 
 class TestPurityIntegration:
     """Test the purity meausrement with qnodes and devices."""
 
-    devices = ["default.qubit.legacy", "lightning.qubit", "default.mixed"]
-    grad_supported_devices = ["default.qubit.legacy", "default.mixed"]
+    devices = ["default.qubit", "lightning.qubit", "default.mixed"]
+    grad_supported_devices = ["default.qubit", "default.mixed"]
     mix_supported_devices = ["default.mixed"]
 
     diff_methods = ["backprop", "finite-diff"]
diff --git a/tests/measurements/test_sample.py b/tests/measurements/test_sample.py
index 3dfcf34280b..b5ba8e8e83b 100644
--- a/tests/measurements/test_sample.py
+++ b/tests/measurements/test_sample.py
@@ -22,37 +22,14 @@
 # pylint: disable=protected-access, no-member
 
 
-# TODO: Remove this when new CustomMP are the default
-def custom_measurement_process(device, spy):
-    assert len(spy.call_args_list) > 0  # make sure method is mocked properly
-
-    samples = device._samples
-    call_args_list = list(spy.call_args_list)
-    for call_args in call_args_list:
-        meas = call_args.args[1]
-        shot_range, bin_size = (call_args.kwargs["shot_range"], call_args.kwargs["bin_size"])
-        if isinstance(meas, (Operator, MeasurementValue)):
-            meas = qml.sample(op=meas)
-        assert qml.math.allequal(
-            device.sample(call_args.args[1], **call_args.kwargs),
-            meas.process_samples(
-                samples=samples,
-                wire_order=device.wires,
-                shot_range=shot_range,
-                bin_size=bin_size,
-            ),
-        )
-
-
-class TestSample:  # pylint: disable=too-many-public-methods
+class TestSample:
     """Tests for the sample function"""
 
     @pytest.mark.parametrize("n_sample", (1, 10))
-    def test_sample_dimension(self, mocker, n_sample):
+    def test_sample_dimension(self, n_sample):
         """Test that the sample function outputs samples of the right size"""
 
-        dev = qml.device("default.qubit.legacy", wires=2, shots=n_sample)
-        spy = mocker.spy(qml.QubitDevice, "sample")
+        dev = qml.device("default.qubit", wires=2, shots=n_sample)
 
         @qml.qnode(dev)
         def circuit():
@@ -69,15 +46,12 @@ def circuit():
             (n_sample,) if not n_sample == 1 else ()
         )
 
-        custom_measurement_process(dev, spy)
-
-    @pytest.mark.filterwarnings("ignore:Creating an ndarray from ragged nested sequences")
-    def test_sample_combination(self, mocker):
+    @pytest.mark.xfail(reason="until DQ2 port")
+    def test_sample_combination(self):
         """Test the output of combining expval, var and sample"""
         n_sample = 10
 
-        dev = qml.device("default.qubit.legacy", wires=3, shots=n_sample)
-        spy = mocker.spy(qml.QubitDevice, "sample")
+        dev = qml.device("default.qubit", wires=3, shots=n_sample)
 
         @qml.qnode(dev, diff_method="parameter-shift")
         def circuit():
@@ -90,17 +64,14 @@ def circuit():
         assert len(result) == 3
         assert np.array_equal(result[0].shape, (n_sample,))
         assert circuit._qfunc_output[0].shape(dev, Shots(n_sample)) == (n_sample,)
-        assert isinstance(result[1], np.ndarray)
-        assert isinstance(result[2], np.ndarray)
+        assert isinstance(result[1], np.float64)
+        assert isinstance(result[2], np.float64)
 
-        custom_measurement_process(dev, spy)
-
-    def test_single_wire_sample(self, mocker):
+    def test_single_wire_sample(self):
         """Test the return type and shape of sampling a single wire"""
         n_sample = 10
 
-        dev = qml.device("default.qubit.legacy", wires=1, shots=n_sample)
-        spy = mocker.spy(qml.QubitDevice, "sample")
+        dev = qml.device("default.qubit", wires=1, shots=n_sample)
 
         @qml.qnode(dev)
         def circuit():
@@ -113,15 +84,12 @@ def circuit():
         assert np.array_equal(result.shape, (n_sample,))
         assert circuit._qfunc_output.shape(dev, Shots(n_sample)) == (n_sample,)
 
-        custom_measurement_process(dev, spy)
-
-    def test_multi_wire_sample_regular_shape(self, mocker):
+    def test_multi_wire_sample_regular_shape(self):
         """Test the return type and shape of sampling multiple wires
         where a rectangular array is expected"""
         n_sample = 10
 
-        dev = qml.device("default.qubit.legacy", wires=3, shots=n_sample)
-        spy = mocker.spy(qml.QubitDevice, "sample")
+        dev = qml.device("default.qubit", wires=3, shots=n_sample)
 
         @qml.qnode(dev)
         def circuit():
@@ -138,16 +106,13 @@ def circuit():
         assert len(result) == 3
         assert result[0].dtype == np.dtype("int")
 
-        custom_measurement_process(dev, spy)
-
     @pytest.mark.filterwarnings("ignore:Creating an ndarray from ragged nested sequences")
-    def test_sample_output_type_in_combination(self, mocker):
+    def test_sample_output_type_in_combination(self):
         """Test the return type and shape of sampling multiple works
         in combination with expvals and vars"""
         n_sample = 10
 
-        dev = qml.device("default.qubit.legacy", wires=3, shots=n_sample)
-        spy = mocker.spy(qml.QubitDevice, "sample")
+        dev = qml.device("default.qubit", wires=3, shots=n_sample)
 
         @qml.qnode(dev, diff_method="parameter-shift")
         def circuit():
@@ -162,13 +127,10 @@ def circuit():
         assert result[2].dtype == np.dtype("int")
         assert np.array_equal(result[2].shape, (n_sample,))
 
-        custom_measurement_process(dev, spy)
-
-    def test_not_an_observable(self, mocker):
+    def test_not_an_observable(self):
         """Test that a UserWarning is raised if the provided
         argument might not be hermitian."""
-        dev = qml.device("default.qubit.legacy", wires=2, shots=10)
-        spy = mocker.spy(qml.QubitDevice, "sample")
+        dev = qml.device("default.qubit", wires=2, shots=10)
 
         @qml.qnode(dev)
         def circuit():
@@ -178,13 +140,10 @@ def circuit():
         with pytest.warns(UserWarning, match="Prod might not be hermitian."):
             _ = circuit()
 
-        custom_measurement_process(dev, spy)
-
-    def test_observable_return_type_is_sample(self, mocker):
+    def test_observable_return_type_is_sample(self):
         """Test that the return type of the observable is :attr:`ObservableReturnTypes.Sample`"""
         n_shots = 10
-        dev = qml.device("default.qubit.legacy", wires=1, shots=n_shots)
-        spy = mocker.spy(qml.QubitDevice, "sample")
+        dev = qml.device("default.qubit", wires=1, shots=n_shots)
 
         @qml.qnode(dev)
         def circuit():
@@ -194,11 +153,9 @@ def circuit():
 
         circuit()
 
-        custom_measurement_process(dev, spy)
-
     @pytest.mark.parametrize("shots", [5, [5, 5]])
     @pytest.mark.parametrize("phi", np.arange(0, 2 * np.pi, np.pi / 2))
-    def test_observable_is_measurement_value(self, shots, phi, mocker):
+    def test_observable_is_measurement_value(self, shots, phi):
         """Test that expectation values for mid-circuit measurement values
         are correct for a single measurement value."""
         dev = qml.device("default.qubit", wires=2, shots=shots)
@@ -209,9 +166,6 @@ def circuit(phi):
             m0 = qml.measure(0)
             return qml.sample(m0)
 
-        new_dev = circuit.device
-        spy = mocker.spy(qml.QubitDevice, "sample")
-
         res = circuit(phi)
 
         if isinstance(shots, list):
@@ -219,11 +173,10 @@ def circuit(phi):
             assert all(r.shape == (s,) for r, s in zip(res, shots))
         else:
             assert res.shape == (shots,)
-        custom_measurement_process(new_dev, spy)
 
     def test_providing_observable_and_wires(self):
         """Test that a ValueError is raised if both an observable is provided and wires are specified"""
-        dev = qml.device("default.qubit.legacy", wires=2)
+        dev = qml.device("default.qubit", wires=2)
 
         @qml.qnode(dev)
         def circuit():
@@ -237,10 +190,9 @@ def circuit():
         ):
             _ = circuit()
 
-    def test_providing_no_observable_and_no_wires(self, mocker):
+    def test_providing_no_observable_and_no_wires(self):
         """Test that we can provide no observable and no wires to sample function"""
-        dev = qml.device("default.qubit.legacy", wires=2, shots=1000)
-        spy = mocker.spy(qml.QubitDevice, "sample")
+        dev = qml.device("default.qubit", wires=2, shots=1000)
 
         @qml.qnode(dev)
         def circuit():
@@ -252,9 +204,7 @@ def circuit():
 
         circuit()
 
-        custom_measurement_process(dev, spy)
-
-    def test_providing_no_observable_and_no_wires_shot_vector(self, mocker):
+    def test_providing_no_observable_and_no_wires_shot_vector(self):
         """Test that we can provide no observable and no wires to sample
         function when using a shot vector"""
         num_wires = 2
@@ -262,8 +212,7 @@ def test_providing_no_observable_and_no_wires_shot_vector(self, mocker):
         shots1 = 1
         shots2 = 10
         shots3 = 1000
-        dev = qml.device("default.qubit.legacy", wires=num_wires, shots=[shots1, shots2, shots3])
-        spy = mocker.spy(qml.QubitDevice, "sample")
+        dev = qml.device("default.qubit", wires=num_wires, shots=[shots1, shots2, shots3])
 
         @qml.qnode(dev)
         def circuit():
@@ -285,14 +234,11 @@ def circuit():
         assert np.all(res[1][:, 0] == res[1][:, 1])
         assert np.all(res[2][:, 0] == res[2][:, 1])
 
-        custom_measurement_process(dev, spy)
-
-    def test_providing_no_observable_and_wires(self, mocker):
+    def test_providing_no_observable_and_wires(self):
         """Test that we can provide no observable but specify wires to the sample function"""
         wires = [0, 2]
         wires_obj = qml.wires.Wires(wires)
-        dev = qml.device("default.qubit.legacy", wires=3, shots=1000)
-        spy = mocker.spy(qml.QubitDevice, "sample")
+        dev = qml.device("default.qubit", wires=3, shots=1000)
 
         @qml.qnode(dev)
         def circuit():
@@ -305,8 +251,6 @@ def circuit():
 
         circuit()
 
-        custom_measurement_process(dev, spy)
-
     @pytest.mark.parametrize(
         "obs,exp",
         [
@@ -343,7 +287,7 @@ def test_numeric_type(self, obs, exp):
 
     def test_shape_no_shots_error(self):
         """Test that the appropriate error is raised with no shots are specified"""
-        dev = qml.device("default.qubit.legacy", wires=2, shots=None)
+        dev = qml.device("default.qubit", wires=2, shots=None)
         shots = Shots(None)
         mp = qml.sample()
 
@@ -364,7 +308,7 @@ def test_shape_no_shots_error(self):
     def test_shape(self, obs):
         """Test that the shape is correct."""
         shots = 10
-        dev = qml.device("default.qubit.legacy", wires=3, shots=shots)
+        dev = qml.device("default.qubit", wires=3, shots=shots)
         res = qml.sample(obs) if obs is not None else qml.sample()
         expected = (shots,) if obs is not None else (shots, 3)
         assert res.shape(dev, Shots(shots)) == expected
@@ -372,7 +316,7 @@ def test_shape(self, obs):
     @pytest.mark.parametrize("n_samples", (1, 10))
     def test_shape_wires(self, n_samples):
         """Test that the shape is correct when wires are provided."""
-        dev = qml.device("default.qubit.legacy", wires=3, shots=n_samples)
+        dev = qml.device("default.qubit", wires=3, shots=n_samples)
         mp = qml.sample(wires=(0, 1))
         assert mp.shape(dev, Shots(n_samples)) == (n_samples, 2) if n_samples != 1 else (2,)
 
@@ -388,7 +332,7 @@ def test_shape_wires(self, n_samples):
     def test_shape_shot_vector(self, obs):
         """Test that the shape is correct with the shot vector too."""
         shot_vector = (1, 2, 3)
-        dev = qml.device("default.qubit.legacy", wires=3, shots=shot_vector)
+        dev = qml.device("default.qubit", wires=3, shots=shot_vector)
         res = qml.sample(obs) if obs is not None else qml.sample()
         expected = ((), (2,), (3,)) if obs is not None else ((3,), (2, 3), (3, 3))
         assert res.shape(dev, Shots(shot_vector)) == expected
@@ -396,7 +340,7 @@ def test_shape_shot_vector(self, obs):
     def test_shape_shot_vector_obs(self):
         """Test that the shape is correct with the shot vector and a observable too."""
         shot_vec = (2, 2)
-        dev = qml.device("default.qubit.legacy", wires=3, shots=shot_vec)
+        dev = qml.device("default.qubit", wires=3, shots=shot_vec)
 
         @qml.qnode(dev)
         def circuit():
@@ -419,10 +363,11 @@ def test_sample_empty_wires(self):
     @pytest.mark.parametrize("shots", [2, 100])
     def test_sample_no_arguments(self, shots):
         """Test that using ``qml.sample`` with no arguments returns the samples of all wires."""
-        dev = qml.device("default.qubit.legacy", wires=3, shots=shots)
+        dev = qml.device("default.qubit", wires=3, shots=shots)
 
         @qml.qnode(dev)
         def circuit():
+            qml.Identity(wires=list(range(3)))
             return qml.sample()
 
         res = circuit()
@@ -450,7 +395,7 @@ def test_jitting_with_sampling_on_subset_of_wires(samples):
 
     jax.config.update("jax_enable_x64", True)
 
-    dev = qml.device("default.qubit.legacy", wires=3, shots=samples)
+    dev = qml.device("default.qubit", wires=3, shots=samples)
 
     @qml.qnode(dev, interface="jax")
     def circuit(x):
diff --git a/tests/measurements/test_state.py b/tests/measurements/test_state.py
index 909d1276d59..c843f0feb1b 100644
--- a/tests/measurements/test_state.py
+++ b/tests/measurements/test_state.py
@@ -226,24 +226,28 @@ def test_process_state_matrix_from_matrix(self, mat, wires):
 class TestState:
     """Tests for the state function"""
 
+    @pytest.mark.xfail(reason="until DQ2 port")
     @pytest.mark.parametrize("wires", range(2, 5))
-    def test_state_shape_and_dtype(self, wires):
+    @pytest.mark.parametrize("op,dtype", [(qml.PauliX, np.float64), (qml.PauliY, np.complex128)])
+    def test_state_shape_and_dtype(self, op, dtype, wires):
         """Test that the state is of correct size and dtype for a trivial circuit"""
 
-        dev = qml.device("default.qubit.legacy", wires=wires)
+        dev = qml.device("default.qubit", wires=wires)
 
         @qml.qnode(dev)
         def func():
+            qml.Identity(wires=list(range(wires)))
+            op(0)
             return state()
 
         state_val = func()
         assert state_val.shape == (2**wires,)
-        assert state_val.dtype == np.complex128
+        assert state_val.dtype == dtype
 
     def test_return_type_is_state(self):
         """Test that the return type of the observable is State"""
 
-        dev = qml.device("default.qubit.legacy", wires=1)
+        dev = qml.device("default.qubit", wires=1)
 
         @qml.qnode(dev)
         def func():
@@ -259,7 +263,7 @@ def func():
     def test_state_correct_ghz(self, wires):
         """Test that the correct state is returned when the circuit prepares a GHZ state"""
 
-        dev = qml.device("default.qubit.legacy", wires=wires)
+        dev = qml.device("default.qubit", wires=wires)
 
         @qml.qnode(dev)
         def func():
@@ -270,35 +274,32 @@ def func():
 
         state_val = func()
         assert np.allclose(np.sum(np.abs(state_val) ** 2), 1)
-        # pylint: disable=unsubscriptable-object
         assert np.allclose(state_val[0], 1 / np.sqrt(2))
         assert np.allclose(state_val[-1], 1 / np.sqrt(2))
 
-        assert np.allclose(state().process_state(state=dev.state, wire_order=dev.wires), state_val)
-
-    def test_return_with_other_types(self):
-        """Test that an exception is raised when a state is returned along with another return
+    @pytest.mark.xfail(reason="until DQ2 port")
+    def test_return_with_other_types_works(self):
+        """Test that no exception is raised when a state is returned along with another return
         type"""
 
-        dev = qml.device("default.qubit.legacy", wires=2)
+        dev = qml.device("default.qubit", wires=2)
 
         @qml.qnode(dev)
         def func():
             qml.Hadamard(wires=0)
             return state(), expval(qml.PauliZ(1))
 
-        with pytest.raises(
-            qml.QuantumFunctionError,
-            match="The state or density matrix cannot be returned in combination with other return types",
-        ):
-            func()
+        res = func()
+        assert isinstance(res, tuple)
+        assert np.allclose(res[0], np.array([1, 0, 1, 0]) / np.sqrt(2))
+        assert np.isclose(res[1], 1)
 
+    @pytest.mark.xfail(reason="until DQ2 port")
     @pytest.mark.parametrize("wires", range(2, 5))
-    def test_state_equal_to_dev_state(self, wires):
-        """Test that the returned state is equal to the one stored in dev.state for a template
-        circuit"""
+    def test_state_equal_to_expected_state(self, wires):
+        """Test that the returned state is equal to the expected state for a template circuit"""
 
-        dev = qml.device("default.qubit.legacy", wires=wires)
+        dev = qml.device("default.qubit", wires=wires)
 
         weights = np.random.random(
             qml.templates.StronglyEntanglingLayers.shape(n_layers=3, n_wires=wires)
@@ -310,17 +311,23 @@ def func():
             return state()
 
         state_val = func()
-        assert np.allclose(state_val, func.device.state)
+        scripts, _, _ = dev.preprocess(func.tape)
+        assert len(scripts) == 1
+        expected_state, _ = qml.devices.qubit.get_final_state(scripts[0])
+        assert np.allclose(state_val, expected_state.flatten())
 
     @pytest.mark.tf
-    def test_interface_tf(self):
+    @pytest.mark.parametrize("op", [qml.PauliX, qml.PauliY])
+    def test_interface_tf(self, op):
         """Test that the state correctly outputs in the tensorflow interface"""
         import tensorflow as tf
 
-        dev = qml.device("default.qubit.legacy", wires=4)
+        dev = qml.device("default.qubit", wires=4)
 
         @qml.qnode(dev, interface="tf")
         def func():
+            op(0)
+            op(0)
             for i in range(4):
                 qml.Hadamard(i)
             return state()
@@ -329,35 +336,40 @@ def func():
         state_val = func()
 
         assert isinstance(state_val, tf.Tensor)
-        assert state_val.dtype == tf.complex128
+        assert state_val.dtype == tf.complex128 if op is qml.PauliY else tf.float64
         assert np.allclose(state_expected, state_val.numpy())
         assert state_val.shape == (16,)
 
+    @pytest.mark.xfail(reason="until DQ2 port")
     @pytest.mark.torch
-    def test_interface_torch(self):
+    @pytest.mark.parametrize("op", [qml.PauliX, qml.PauliY])
+    def test_interface_torch(self, op):
         """Test that the state correctly outputs in the torch interface"""
         import torch
 
-        dev = qml.device("default.qubit.legacy", wires=4)
+        dev = qml.device("default.qubit", wires=4)
 
         @qml.qnode(dev, interface="torch")
         def func():
+            op(0)
+            op(0)
             for i in range(4):
                 qml.Hadamard(i)
             return state()
 
-        state_expected = 0.25 * torch.ones(16, dtype=torch.complex128)
+        dtype = torch.complex128 if op is qml.PauliY else torch.float64
+        state_expected = 0.25 * torch.ones(16, dtype=dtype)
         state_val = func()
 
         assert isinstance(state_val, torch.Tensor)
-        assert state_val.dtype == torch.complex128
+        assert state_val.dtype == dtype
         assert torch.allclose(state_expected, state_val)
         assert state_val.shape == (16,)
 
     @pytest.mark.autograd
     def test_jacobian_not_supported(self):
         """Test if an error is raised if the jacobian method is called via qml.grad"""
-        dev = qml.device("default.qubit.legacy", wires=4)
+        dev = qml.device("default.qubit", wires=4)
 
         @qml.qnode(dev, diff_method="parameter-shift")
         def func(x):
@@ -409,7 +421,7 @@ def test_default_qubit(self, diff_method):
         """Test that the returned state is equal to the expected returned state for all of
         PennyLane's built in statevector devices"""
 
-        dev = qml.device("default.qubit.legacy", wires=4)
+        dev = qml.device("default.qubit", wires=4)
 
         @qml.qnode(dev, diff_method=diff_method)
         def func():
@@ -421,7 +433,6 @@ def func():
         state_expected = 0.25 * np.ones(16)
 
         assert np.allclose(state_val, state_expected)
-        assert np.allclose(state_val, dev.state)
 
     @pytest.mark.tf
     @pytest.mark.parametrize("diff_method", ["best", "finite-diff", "parameter-shift"])
@@ -441,7 +452,6 @@ def func():
         state_expected = 0.25 * np.ones(16)
 
         assert np.allclose(state_val, state_expected)
-        assert np.allclose(state_val, dev.state)
 
     @pytest.mark.autograd
     @pytest.mark.parametrize("diff_method", ["best", "finite-diff", "parameter-shift"])
@@ -461,7 +471,6 @@ def func():
         state_expected = 0.25 * np.ones(16)
 
         assert np.allclose(state_val, state_expected)
-        assert np.allclose(state_val, dev.state)
 
     @pytest.mark.tf
     def test_gradient_with_passthru_tf(self):
@@ -513,7 +522,7 @@ def loss_fn(x):
     @pytest.mark.parametrize("wires", [[0, 2, 3, 1], ["a", -1, "b", 1000]])
     def test_custom_wire_labels(self, wires):
         """Test the state when custom wire labels are used"""
-        dev = qml.device("default.qubit.legacy", wires=wires)
+        dev = qml.device("default.qubit", wires=wires)
 
         @qml.qnode(dev, diff_method="parameter-shift")
         def func():
@@ -529,14 +538,14 @@ def func():
     @pytest.mark.parametrize("shots", [None, 1, 10])
     def test_shape(self, shots):
         """Test that the shape is correct for qml.state."""
-        dev = qml.device("default.qubit.legacy", wires=3, shots=shots)
+        dev = qml.device("default.qubit", wires=3, shots=shots)
         res = qml.state()
         assert res.shape(dev, Shots(shots)) == (2**3,)
 
     @pytest.mark.parametrize("s_vec", [(3, 2, 1), (1, 5, 10), (3, 1, 20)])
     def test_shape_shot_vector(self, s_vec):
         """Test that the shape is correct for qml.state with the shot vector too."""
-        dev = qml.device("default.qubit.legacy", wires=3, shots=s_vec)
+        dev = qml.device("default.qubit", wires=3, shots=s_vec)
         res = qml.state()
         assert res.shape(dev, Shots(s_vec)) == ((2**3,), (2**3,), (2**3,))
 
@@ -551,9 +560,11 @@ class TestDensityMatrix:
 
     # pylint: disable=too-many-public-methods
 
+    @pytest.mark.xfail(reason="until DQ2 port")
     @pytest.mark.parametrize("wires", range(2, 5))
-    @pytest.mark.parametrize("dev_name", ["default.qubit.legacy", "default.mixed"])
-    def test_density_matrix_shape_and_dtype(self, dev_name, wires):
+    @pytest.mark.parametrize("dev_name", ["default.qubit", "default.mixed"])
+    @pytest.mark.parametrize("op,dtype", [(qml.PauliX, np.float64), (qml.PauliY, np.complex128)])
+    def test_density_matrix_shape_and_dtype(self, dev_name, op, dtype, wires):
         """Test that the density matrix is of correct size and dtype for a
         trivial circuit"""
 
@@ -561,14 +572,15 @@ def test_density_matrix_shape_and_dtype(self, dev_name, wires):
 
         @qml.qnode(dev)
         def circuit():
+            op(0)
             return density_matrix([0])
 
         state_val = circuit()
 
         assert state_val.shape == (2, 2)
-        assert state_val.dtype == np.complex128
+        assert state_val.dtype == dtype if dev_name == "default.qubit" else np.complex128
 
-    @pytest.mark.parametrize("dev_name", ["default.qubit.legacy", "default.mixed"])
+    @pytest.mark.parametrize("dev_name", ["default.qubit", "default.mixed"])
     def test_return_type_is_state(self, dev_name):
         """Test that the return type of the observable is State"""
 
@@ -585,7 +597,7 @@ def func():
         assert obs[0].return_type is State
 
     @pytest.mark.torch
-    @pytest.mark.parametrize("dev_name", ["default.qubit.legacy", "default.mixed"])
+    @pytest.mark.parametrize("dev_name", ["default.qubit", "default.mixed"])
     @pytest.mark.parametrize("diff_method", [None, "backprop"])
     def test_correct_density_matrix_torch(self, dev_name, diff_method):
         """Test that the correct density matrix is returned using torch interface."""
@@ -600,16 +612,8 @@ def func():
         expected = np.array([[0.5 + 0.0j, 0.5 + 0.0j], [0.5 + 0.0j, 0.5 + 0.0j]])
         assert np.allclose(expected, density_mat)
 
-        dev = func.device
-
-        if dev_name != "default.mixed":
-            assert np.allclose(
-                expected,
-                qml.density_matrix(wires=0).process_state(state=dev.state, wire_order=dev.wires),
-            )
-
     @pytest.mark.jax
-    @pytest.mark.parametrize("dev_name", ["default.qubit.legacy", "default.mixed"])
+    @pytest.mark.parametrize("dev_name", ["default.qubit", "default.mixed"])
     @pytest.mark.parametrize("diff_method", [None, "backprop"])
     def test_correct_density_matrix_jax(self, dev_name, diff_method):
         """Test that the correct density matrix is returned using JAX interface."""
@@ -625,18 +629,11 @@ def func():
 
         assert np.allclose(expected, density_mat)
 
-        if dev_name != "default.mixed":
-            assert np.allclose(
-                expected,
-                qml.density_matrix(wires=0).process_state(state=dev.state, wire_order=dev.wires),
-            )
-
     @pytest.mark.tf
-    @pytest.mark.parametrize("dev_name", ["default.qubit.legacy", "default.mixed"])
     @pytest.mark.parametrize("diff_method", [None, "backprop"])
-    def test_correct_density_matrix_tf(self, dev_name, diff_method):
+    def test_correct_density_matrix_tf_default_mixed(self, diff_method):
         """Test that the correct density matrix is returned using the TensorFlow interface."""
-        dev = qml.device(dev_name, wires=2)
+        dev = qml.device("default.mixed", wires=2)
 
         @qml.qnode(dev, interface="tf", diff_method=diff_method)
         def func():
@@ -648,18 +645,33 @@ def func():
 
         assert np.allclose(expected, density_mat)
 
-        if dev_name != "default.mixed":
-            assert np.allclose(
-                expected,
-                qml.density_matrix(wires=0).process_state(state=dev.state, wire_order=dev.wires),
-            )
+    @pytest.mark.xfail(reason="until DQ2 port")
+    @pytest.mark.tf
+    @pytest.mark.parametrize("diff_method", [None, "backprop"])
+    def test_correct_density_matrix_tf_default_qubit(self, diff_method):
+        """Test that the correct density matrix is returned using the TensorFlow interface."""
+        dev = qml.device("default.qubit", wires=2)
+
+        @qml.qnode(dev, interface="tf", diff_method=diff_method)
+        def func():
+            qml.Hadamard(wires=0)
+            return qml.density_matrix(wires=0), state()
+
+        density_mat, dev_state = func()
+        expected = np.array([[0.5 + 0.0j, 0.5 + 0.0j], [0.5 + 0.0j, 0.5 + 0.0j]])
+
+        assert np.allclose(expected, density_mat)
 
-    @pytest.mark.parametrize("dev_name", ["default.qubit.legacy", "default.mixed"])
-    def test_correct_density_matrix_product_state_first(self, dev_name):
+        assert np.allclose(
+            expected,
+            qml.density_matrix(wires=0).process_state(state=dev_state, wire_order=dev.wires),
+        )
+
+    def test_correct_density_matrix_product_state_first_default_mixed(self):
         """Test that the correct density matrix is returned when
         tracing out a product state"""
 
-        dev = qml.device(dev_name, wires=2)
+        dev = qml.device("default.mixed", wires=2)
 
         @qml.qnode(dev)
         def func():
@@ -672,18 +684,33 @@ def func():
 
         assert np.allclose(expected, density_first)
 
-        if dev_name != "default.mixed":
-            assert np.allclose(
-                expected,
-                qml.density_matrix(wires=0).process_state(state=dev.state, wire_order=dev.wires),
-            )
+    @pytest.mark.xfail(reason="until DQ2 port")
+    def test_correct_density_matrix_product_state_first_default_qubit(self):
+        """Test that the correct density matrix is returned when
+        tracing out a product state"""
+
+        dev = qml.device("default.qubit", wires=2)
+
+        @qml.qnode(dev)
+        def func():
+            qml.Hadamard(wires=1)
+            qml.PauliY(wires=0)
+            return density_matrix(0), state()
+
+        density_first, dev_state = func()
+        expected = np.array([[0.0 + 0.0j, 0.0 + 0.0j], [0.0 + 0.0j, 1.0 + 0.0j]])
+
+        assert np.allclose(expected, density_first)
+        assert np.allclose(
+            expected,
+            qml.density_matrix(wires=0).process_state(state=dev_state, wire_order=dev.wires),
+        )
 
-    @pytest.mark.parametrize("dev_name", ["default.qubit.legacy", "default.mixed"])
-    def test_correct_density_matrix_product_state_second(self, dev_name):
+    def test_correct_density_matrix_product_state_second_default_mixed(self):
         """Test that the correct density matrix is returned when
         tracing out a product state"""
 
-        dev = qml.device(dev_name, wires=2)
+        dev = qml.device("default.mixed", wires=2)
 
         @qml.qnode(dev)
         def func():
@@ -695,19 +722,34 @@ def func():
         expected = np.array([[0.5 + 0.0j, 0.5 + 0.0j], [0.5 + 0.0j, 0.5 + 0.0j]])
         assert np.allclose(expected, density_second)
 
-        if dev_name != "default.mixed":
-            assert np.allclose(
-                expected,
-                qml.density_matrix(wires=1).process_state(state=dev.state, wire_order=dev.wires),
-            )
+    @pytest.mark.xfail(reason="until DQ2 port")
+    def test_correct_density_matrix_product_state_second_default_qubit(self):
+        """Test that the correct density matrix is returned when
+        tracing out a product state"""
+
+        dev = qml.device("default.qubit", wires=2)
+
+        @qml.qnode(dev)
+        def func():
+            qml.Hadamard(wires=1)
+            qml.PauliY(wires=0)
+            return density_matrix(1), state()
+
+        density_second, dev_state = func()
+        expected = np.array([[0.5 + 0.0j, 0.5 + 0.0j], [0.5 + 0.0j, 0.5 + 0.0j]])
+        assert np.allclose(expected, density_second)
+
+        assert np.allclose(
+            expected,
+            qml.density_matrix(wires=1).process_state(state=dev_state, wire_order=dev.wires),
+        )
 
-    @pytest.mark.parametrize("dev_name", ["default.qubit.legacy", "default.mixed"])
     @pytest.mark.parametrize("return_wire_order", ([0, 1], [1, 0]))
-    def test_correct_density_matrix_product_state_both(self, dev_name, return_wire_order):
+    def test_correct_density_matrix_product_state_both_default_mixed(self, return_wire_order):
         """Test that the correct density matrix is returned
         for a full product state on two wires."""
 
-        dev = qml.device(dev_name, wires=2)
+        dev = qml.device("default.mixed", wires=2)
 
         @qml.qnode(dev)
         def func():
@@ -722,20 +764,39 @@ def func():
 
         assert np.allclose(expected, density_both)
 
-        if dev_name != "default.mixed":
-            assert np.allclose(
-                expected,
-                qml.density_matrix(wires=return_wire_order).process_state(
-                    state=dev.state, wire_order=dev.wires
-                ),
-            )
+    @pytest.mark.xfail(reason="until DQ2 port")
+    @pytest.mark.parametrize("return_wire_order", ([0, 1], [1, 0]))
+    def test_correct_density_matrix_product_state_both_default_qubit(self, return_wire_order):
+        """Test that the correct density matrix is returned
+        for a full product state on two wires."""
+
+        dev = qml.device("default.qubit", wires=2)
+
+        @qml.qnode(dev)
+        def func():
+            qml.Hadamard(wires=1)
+            qml.PauliY(wires=0)
+            return density_matrix(return_wire_order), state()
+
+        density_both, dev_state = func()
+        single_statevectors = [[0, 1j], [1 / np.sqrt(2), 1 / np.sqrt(2)]]
+        expected_statevector = np.kron(*[single_statevectors[w] for w in return_wire_order])
+        expected = np.outer(expected_statevector.conj(), expected_statevector)
+
+        assert np.allclose(expected, density_both)
+
+        assert np.allclose(
+            expected,
+            qml.density_matrix(wires=return_wire_order).process_state(
+                state=dev_state, wire_order=dev.wires
+            ),
+        )
 
-    @pytest.mark.parametrize("dev_name", ["default.qubit.legacy", "default.mixed"])
-    def test_correct_density_matrix_three_wires_first_two(self, dev_name):
+    def test_correct_density_matrix_three_wires_first_two_default_mixed(self):
         """Test that the correct density matrix is returned for an example with three wires,
         and tracing out the third wire."""
 
-        dev = qml.device(dev_name, wires=3)
+        dev = qml.device("default.mixed", wires=3)
 
         @qml.qnode(dev)
         def func():
@@ -754,15 +815,35 @@ def func():
         )
         assert np.allclose(expected, density_full)
 
-        if dev_name != "default.mixed":
-            assert np.allclose(
-                expected,
-                qml.density_matrix(wires=[0, 1]).process_state(
-                    state=dev.state, wire_order=dev.wires
-                ),
-            )
+    @pytest.mark.xfail(reason="until DQ2 port")
+    def test_correct_density_matrix_three_wires_first_two_default_qubit(self):
+        """Test that the correct density matrix is returned for an example with three wires,
+        and tracing out the third wire."""
+
+        dev = qml.device("default.qubit", wires=3)
+
+        @qml.qnode(dev)
+        def func():
+            qml.Hadamard(wires=1)
+            qml.PauliY(wires=0)
+            return density_matrix([0, 1]), state()
+
+        density_full, dev_state = func()
+        expected = np.array(
+            [
+                [0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j],
+                [0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j],
+                [0.0 + 0.0j, 0.0 + 0.0j, 0.5 + 0.0j, 0.5 + 0.0j],
+                [0.0 + 0.0j, 0.0 + 0.0j, 0.5 + 0.0j, 0.5 + 0.0j],
+            ]
+        )
+        assert np.allclose(expected, density_full)
+        assert np.allclose(
+            expected,
+            qml.density_matrix(wires=[0, 1]).process_state(state=dev_state, wire_order=dev.wires),
+        )
 
-    @pytest.mark.parametrize("dev_name", ["default.qubit.legacy", "default.mixed"])
+    @pytest.mark.parametrize("dev_name", ["default.qubit", "default.mixed"])
     def test_correct_density_matrix_three_wires_last_two(self, dev_name):
         """Test that the correct density matrix is returned for an example with three wires,
         and tracing out the first wire."""
@@ -790,23 +871,14 @@ def func():
 
         assert np.allclose(expected, density)
 
-        if dev_name != "default.mixed":
-            assert np.allclose(
-                expected,
-                qml.density_matrix(wires=[1, 2]).process_state(
-                    state=dev.state, wire_order=dev.wires
-                ),
-            )
-
-    @pytest.mark.parametrize("dev_name", ["default.qubit.legacy", "default.mixed"])
     @pytest.mark.parametrize(
         "return_wire_order", ([0], [1], [2], [0, 1], [1, 0], [0, 2], [2, 0], [1, 2, 0], [2, 1, 0])
     )
-    def test_correct_density_matrix_three_wires_product(self, dev_name, return_wire_order):
+    def test_correct_density_matrix_three_wires_product_default_mixed(self, return_wire_order):
         """Test that the correct density matrix is returned for an example with
         three wires and a product state, tracing out various combinations."""
 
-        dev = qml.device(dev_name, wires=3)
+        dev = qml.device("default.mixed", wires=3)
 
         @qml.qnode(dev)
         def func():
@@ -830,15 +902,46 @@ def func():
         expected = np.outer(exp_statevector.conj(), exp_statevector)
         assert np.allclose(expected, density_full)
 
-        if dev_name != "default.mixed":
-            assert np.allclose(
-                expected,
-                qml.density_matrix(wires=return_wire_order).process_state(
-                    state=dev.state, wire_order=dev.wires
-                ),
-            )
+    @pytest.mark.xfail(reason="until DQ2 port")
+    @pytest.mark.parametrize(
+        "return_wire_order", ([0], [1], [2], [0, 1], [1, 0], [0, 2], [2, 0], [1, 2, 0], [2, 1, 0])
+    )
+    def test_correct_density_matrix_three_wires_product_default_qubit(self, return_wire_order):
+        """Test that the correct density matrix is returned for an example with
+        three wires and a product state, tracing out various combinations."""
 
-    @pytest.mark.parametrize("dev_name", ["default.qubit.legacy", "default.mixed"])
+        dev = qml.device("default.qubit", wires=3)
+
+        @qml.qnode(dev)
+        def func():
+            qml.Hadamard(0)
+            qml.PauliX(1)
+            qml.PauliZ(2)
+            return density_matrix(return_wire_order), state()
+
+        density_full, dev_state = func()
+
+        single_states = [[1 / np.sqrt(2), 1 / np.sqrt(2)], [0, 1], [1, 0]]
+        if len(return_wire_order) == 1:
+            exp_statevector = np.array(single_states[return_wire_order[0]])
+        elif len(return_wire_order) == 2:
+            i, j = return_wire_order
+            exp_statevector = np.kron(single_states[i], single_states[j])
+        elif len(return_wire_order) == 3:
+            i, j, k = return_wire_order
+            exp_statevector = np.kron(np.kron(single_states[i], single_states[j]), single_states[k])
+
+        expected = np.outer(exp_statevector.conj(), exp_statevector)
+        assert np.allclose(expected, density_full)
+
+        assert np.allclose(
+            expected,
+            qml.density_matrix(wires=return_wire_order).process_state(
+                state=dev_state, wire_order=dev.wires
+            ),
+        )
+
+    @pytest.mark.parametrize("dev_name", ["default.qubit", "default.mixed"])
     def test_correct_density_matrix_mixed_state(self, dev_name):
         """Test that the correct density matrix for an example with a mixed state"""
 
@@ -854,11 +957,10 @@ def func():
 
         assert np.allclose(np.array([[0.5 + 0.0j, 0.0 + 0.0j], [0.0 + 0.0j, 0.5 + 0.0j]]), density)
 
-    @pytest.mark.parametrize("dev_name", ["default.qubit.legacy", "default.mixed"])
-    def test_correct_density_matrix_all_wires(self, dev_name):
+    def test_correct_density_matrix_all_wires_default_mixed(self):
         """Test that the correct density matrix is returned when all wires are given"""
 
-        dev = qml.device(dev_name, wires=2)
+        dev = qml.device("default.mixed", wires=2)
 
         @qml.qnode(dev)
         def func():
@@ -878,32 +980,63 @@ def func():
 
         assert np.allclose(expected, density)
 
-        if dev_name != "default.mixed":
-            assert np.allclose(
-                expected,
-                qml.density_matrix(wires=[0, 1]).process_state(
-                    state=dev.state, wire_order=dev.wires
-                ),
-            )
-
-    @pytest.mark.parametrize("dev_name", ["default.qubit.legacy", "default.mixed"])
-    def test_return_with_other_types(self, dev_name):
-        """Test that an exception is raised when a state is returned along with another return
+    @pytest.mark.xfail(reason="until DQ2 port")
+    def test_correct_density_matrix_all_wires_default_qubit(self):
+        """Test that the correct density matrix is returned when all wires are given"""
+
+        dev = qml.device("default.qubit", wires=2)
+
+        @qml.qnode(dev)
+        def func():
+            qml.Hadamard(0)
+            qml.CNOT(wires=[0, 1])
+            return qml.density_matrix(wires=[0, 1]), state()
+
+        density, dev_state = func()
+        expected = np.array(
+            [
+                [0.5 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j, 0.5 + 0.0j],
+                [0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j],
+                [0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j],
+                [0.5 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j, 0.5 + 0.0j],
+            ]
+        )
+
+        assert np.allclose(expected, density)
+        assert np.allclose(
+            expected,
+            qml.density_matrix(wires=[0, 1]).process_state(state=dev_state, wire_order=dev.wires),
+        )
+
+    @pytest.mark.xfail(reason="until DQ2 port")
+    def test_return_with_other_types_works(self):
+        """Test that no exception is raised when a state is returned along with another return
         type"""
 
-        dev = qml.device(dev_name, wires=2)
+        dev = qml.device("default.qubit", wires=2)
 
         @qml.qnode(dev)
         def func():
             qml.Hadamard(wires=0)
             return density_matrix(0), expval(qml.PauliZ(1))
 
-        with pytest.raises(
-            qml.QuantumFunctionError,
-            match="The state or density matrix"
-            " cannot be returned in combination"
-            " with other return types",
-        ):
+        res = func()
+        assert isinstance(res, tuple)
+        assert np.allclose(res[0], np.ones((2, 2)) / 2)
+        assert np.isclose(res[1], 1)
+
+    def test_return_with_other_types_fails(self):
+        """Test that no exception is raised when a state is returned along with another return
+        type"""
+
+        dev = qml.device("default.mixed", wires=2)
+
+        @qml.qnode(dev)
+        def func():
+            qml.Hadamard(wires=0)
+            return density_matrix(0), expval(qml.PauliZ(1))
+
+        with pytest.raises(qml.QuantumFunctionError, match="cannot be returned in combination"):
             func()
 
     def test_no_state_capability(self, monkeypatch):
@@ -938,7 +1071,7 @@ def func():
             func()
 
     @pytest.mark.parametrize("wires", [[0, 2], ["a", -1]])
-    @pytest.mark.parametrize("dev_name", ["default.qubit.legacy", "default.mixed"])
+    @pytest.mark.parametrize("dev_name", ["default.qubit", "default.mixed"])
     def test_custom_wire_labels(self, wires, dev_name):
         """Test that the correct density matrix for an example with a mixed
         state when using custom wires"""
@@ -956,16 +1089,8 @@ def func():
 
         assert np.allclose(expected, density)
 
-        if dev_name != "default.mixed":
-            assert np.allclose(
-                expected,
-                qml.density_matrix(wires=wires[1]).process_state(
-                    state=dev.state, wire_order=dev.wires
-                ),
-            )
-
     @pytest.mark.parametrize("wires", [[3, 1], ["b", 1000]])
-    @pytest.mark.parametrize("dev_name", ["default.qubit.legacy", "default.mixed"])
+    @pytest.mark.parametrize("dev_name", ["default.qubit", "default.mixed"])
     def test_custom_wire_labels_all_wires(self, wires, dev_name):
         """Test that the correct density matrix for an example with a mixed
         state when using custom wires"""
@@ -994,14 +1119,14 @@ def func():
     @pytest.mark.parametrize("shots", [None, 1, 10])
     def test_shape(self, shots):
         """Test that the shape is correct for qml.density_matrix."""
-        dev = qml.device("default.qubit.legacy", wires=3, shots=shots)
+        dev = qml.device("default.qubit", wires=3, shots=shots)
         res = qml.density_matrix(wires=[0, 1])
         assert res.shape(dev, Shots(shots)) == (2**2, 2**2)
 
     @pytest.mark.parametrize("s_vec", [(3, 2, 1), (1, 5, 10), (3, 1, 20)])
     def test_shape_shot_vector(self, s_vec):
         """Test that the shape is correct for qml.density_matrix with the shot vector too."""
-        dev = qml.device("default.qubit.legacy", wires=3, shots=s_vec)
+        dev = qml.device("default.qubit", wires=3, shots=s_vec)
         res = qml.density_matrix(wires=[0, 1])
         assert res.shape(dev, Shots(s_vec)) == (
             (2**2, 2**2),
diff --git a/tests/measurements/test_var.py b/tests/measurements/test_var.py
index 331defbae60..ef19d1ba49b 100644
--- a/tests/measurements/test_var.py
+++ b/tests/measurements/test_var.py
@@ -19,41 +19,13 @@
 from pennylane.measurements import Variance, Shots
 
 
-# TODO: Remove this when new CustomMP are the default
-def custom_measurement_process(device, spy):
-    assert len(spy.call_args_list) > 0  # make sure method is mocked properly
-
-    # pylint: disable=protected-access
-    samples = device._samples
-    state = device._state
-    call_args_list = list(spy.call_args_list)
-    for call_args in call_args_list:
-        obs = call_args.args[1]
-        shot_range, bin_size = (
-            call_args.kwargs["shot_range"],
-            call_args.kwargs["bin_size"],
-        )
-        meas = qml.var(op=obs)
-        old_res = device.var(obs, shot_range=shot_range, bin_size=bin_size)
-        if samples is not None:
-            new_res = meas.process_samples(
-                samples=samples, wire_order=device.wires, shot_range=shot_range, bin_size=bin_size
-            )
-        else:
-            new_res = meas.process_state(state=state, wire_order=device.wires)
-        assert qml.math.allclose(old_res, new_res)
-
-
 class TestVar:
     """Tests for the var function"""
 
     @pytest.mark.parametrize("shots", [None, 10000, [10000, 10000]])
-    @pytest.mark.parametrize("r_dtype", [np.float32, np.float64])
-    def test_value(self, tol, r_dtype, mocker, shots):
+    def test_value(self, tol, shots):
         """Test that the var function works"""
-        dev = qml.device("default.qubit.legacy", wires=2, shots=shots)
-        spy = mocker.spy(qml.QubitDevice, "var")
-        dev.R_DTYPE = r_dtype
+        dev = qml.device("default.qubit", wires=2, shots=shots)
 
         @qml.qnode(dev, diff_method="parameter-shift")
         def circuit(x):
@@ -67,20 +39,17 @@ def circuit(x):
         rtol = 0 if shots is None else 0.05
 
         assert np.allclose(res, expected, atol=atol, rtol=rtol)
-        # pylint: disable=no-member, unsubscriptable-object
+        r_dtype = np.float64
         if isinstance(res, tuple):
             assert res[0].dtype == r_dtype
             assert res[1].dtype == r_dtype
         else:
             assert res.dtype == r_dtype
 
-        custom_measurement_process(dev, spy)
-
-    def test_not_an_observable(self, mocker):
+    def test_not_an_observable(self):
         """Test that a UserWarning is raised if the provided
         argument might not be hermitian."""
-        dev = qml.device("default.qubit.legacy", wires=2)
-        spy = mocker.spy(qml.QubitDevice, "var")
+        dev = qml.device("default.qubit", wires=2)
 
         @qml.qnode(dev)
         def circuit():
@@ -90,12 +59,9 @@ def circuit():
         with pytest.warns(UserWarning, match="Prod might not be hermitian."):
             _ = circuit()
 
-        custom_measurement_process(dev, spy)
-
-    def test_observable_return_type_is_variance(self, mocker):
+    def test_observable_return_type_is_variance(self):
         """Test that the return type of the observable is :attr:`ObservableReturnTypes.Variance`"""
-        dev = qml.device("default.qubit.legacy", wires=2)
-        spy = mocker.spy(qml.QubitDevice, "var")
+        dev = qml.device("default.qubit", wires=2)
 
         @qml.qnode(dev)
         def circuit():
@@ -105,12 +71,10 @@ def circuit():
 
         circuit()
 
-        custom_measurement_process(dev, spy)
-
     @pytest.mark.parametrize("shots", [None, 10000, [10000, 10000]])
     @pytest.mark.parametrize("phi", np.arange(0, 2 * np.pi, np.pi / 3))
     def test_observable_is_measurement_value(
-        self, shots, phi, mocker, tol, tol_stochastic
+        self, shots, phi, tol, tol_stochastic
     ):  # pylint: disable=too-many-arguments
         """Test that variances for mid-circuit measurement values
         are correct for a single measurement value."""
@@ -122,15 +86,11 @@ def circuit(phi):
             m0 = qml.measure(0)
             return qml.var(m0)
 
-        new_dev = circuit.device
-        spy = mocker.spy(qml.QubitDevice, "var")
-
         res = circuit(phi)
 
         atol = tol if shots is None else tol_stochastic
         expected = np.sin(phi / 2) ** 2 - np.sin(phi / 2) ** 4
         assert np.allclose(np.array(res), expected, atol=atol, rtol=0)
-        custom_measurement_process(new_dev, spy)
 
     @pytest.mark.parametrize(
         "obs",
@@ -147,7 +107,7 @@ def test_numeric_type(self, obs):
     )
     def test_shape(self, obs):
         """Test that the shape is correct."""
-        dev = qml.device("default.qubit.legacy", wires=1)
+        dev = qml.device("default.qubit", wires=1)
         res = qml.var(obs)
         # pylint: disable=use-implicit-booleaness-not-comparison
         assert res.shape(dev, Shots(None)) == ()
@@ -161,15 +121,14 @@ def test_shape_shot_vector(self, obs):
         """Test that the shape is correct with the shot vector too."""
         res = qml.var(obs)
         shot_vector = (1, 2, 3)
-        dev = qml.device("default.qubit.legacy", wires=3, shots=shot_vector)
+        dev = qml.device("default.qubit", wires=3, shots=shot_vector)
         assert res.shape(dev, Shots(shot_vector)) == ((), (), ())
 
     @pytest.mark.parametrize("state", [np.array([0, 0, 0]), np.array([1, 0, 0, 0, 0, 0, 0, 0])])
     @pytest.mark.parametrize("shots", [None, 1000, [1000, 10000]])
-    def test_projector_var(self, state, shots, mocker):
+    def test_projector_var(self, state, shots):
         """Tests that the variance of a ``Projector`` object is computed correctly."""
-        dev = qml.device("default.qubit.legacy", wires=3, shots=shots)
-        spy = mocker.spy(qml.QubitDevice, "var")
+        dev = qml.device("default.qubit", wires=3, shots=shots)
 
         @qml.qnode(dev)
         def circuit():
@@ -181,17 +140,14 @@ def circuit():
 
         assert np.allclose(res, expected, atol=0.02, rtol=0.02)
 
-        custom_measurement_process(dev, spy)
-
-    def test_permuted_wires(self, mocker):
+    def test_permuted_wires(self):
         """Test that the variance of an operator with permuted wires is the same."""
         obs = qml.prod(qml.PauliZ(8), qml.s_prod(2, qml.PauliZ(10)), qml.s_prod(3, qml.PauliZ("h")))
         obs_2 = qml.prod(
             qml.s_prod(3, qml.PauliZ("h")), qml.PauliZ(8), qml.s_prod(2, qml.PauliZ(10))
         )
 
-        dev = qml.device("default.qubit.legacy", wires=["h", 8, 10])
-        spy = mocker.spy(qml.QubitDevice, "var")
+        dev = qml.device("default.qubit", wires=["h", 8, 10])
 
         @qml.qnode(dev)
         def circuit():
@@ -206,4 +162,3 @@ def circuit2():
             return qml.var(obs_2)
 
         assert circuit() == circuit2()
-        custom_measurement_process(dev, spy)
diff --git a/tests/measurements/test_vn_entropy.py b/tests/measurements/test_vn_entropy.py
index e69e88e5256..432447227fe 100644
--- a/tests/measurements/test_vn_entropy.py
+++ b/tests/measurements/test_vn_entropy.py
@@ -18,7 +18,6 @@
 import pytest
 
 import pennylane as qml
-from pennylane.interfaces import INTERFACE_MAP
 from pennylane.measurements import VnEntropy, Shots
 from pennylane.measurements.vn_entropy import VnEntropyMP
 from pennylane.wires import Wires
@@ -76,25 +75,18 @@ class TestInitialization:
     @pytest.mark.parametrize("interface", ["autograd", "jax", "tf", "torch"])
     def test_vn_entropy(self, interface, state_vector, expected):
         """Tests the output of qml.vn_entropy"""
-        dev = qml.device("default.qubit.legacy", wires=2)
+        dev = qml.device("default.qubit", wires=2)
 
         @qml.qnode(dev, interface=interface)
         def circuit():
             qml.StatePrep(state_vector, wires=[0, 1])
             return qml.vn_entropy(wires=0)
 
-        res = circuit()
-        new_res = qml.vn_entropy(wires=0).process_state(
-            state=circuit.device.state, wire_order=circuit.device.wires
-        )
-        assert qml.math.allclose(res, expected)
-        assert qml.math.allclose(new_res, expected)
-        assert INTERFACE_MAP.get(qml.math.get_interface(new_res)) == interface
-        assert res.dtype == new_res.dtype
+        assert qml.math.allclose(circuit(), expected)
 
     def test_queue(self):
         """Test that the right measurement class is queued."""
-        dev = qml.device("default.qubit.legacy", wires=2)
+        dev = qml.device("default.qubit", wires=2)
 
         @qml.qnode(dev)
         def circuit():
@@ -121,7 +113,7 @@ def test_properties(self):
     def test_shape(self, shots, shape):
         """Test the ``shape`` method."""
         meas = qml.vn_entropy(wires=0)
-        dev = qml.device("default.qubit.legacy", wires=1, shots=shots)
+        dev = qml.device("default.qubit", wires=1, shots=shots)
 
         assert meas.shape(dev, Shots(shots)) == shape
 
@@ -131,7 +123,7 @@ class TestIntegration:
 
     parameters = np.linspace(0, 2 * np.pi, 10)
 
-    devices = ["default.qubit.legacy", "default.mixed", "lightning.qubit"]
+    devices = ["default.qubit", "default.mixed", "lightning.qubit"]
 
     single_wires_list = [
         [0],
@@ -142,12 +134,13 @@ class TestIntegration:
 
     check_state = [True, False]
 
-    devices = ["default.qubit.legacy", "default.mixed"]
     diff_methods = ["backprop", "finite-diff"]
 
-    def test_shot_vec_error(self):
+    @pytest.mark.xfail(reason="until DQ2 port")
+    @pytest.mark.parametrize("shots", [1000, [1, 10, 10, 1000]])
+    def test_finite_shots_error(self, shots):
         """Test an error is raised when using shot vectors with vn_entropy."""
-        dev = qml.device("default.qubit.legacy", wires=2, shots=[1, 10, 10, 1000])
+        dev = qml.device("default.qubit", wires=2, shots=shots)
 
         @qml.qnode(device=dev)
         def circuit(x):
@@ -155,9 +148,7 @@ def circuit(x):
             qml.CRX(x, wires=[0, 1])
             return qml.vn_entropy(wires=[0])
 
-        with pytest.raises(
-            NotImplementedError, match="Von Neumann entropy is not supported with shot vectors"
-        ):
+        with pytest.raises(qml.DeviceError, match="Circuits with finite shots must only contain"):
             circuit(0.5)
 
     @pytest.mark.parametrize("wires", single_wires_list)
@@ -187,7 +178,7 @@ def circuit_entropy(x):
     def test_IsingXX_qnode_entropy_grad(self, param, wires, base, diff_method):
         """Test entropy for a QNode gradient with autograd."""
 
-        dev = qml.device("default.qubit.legacy", wires=2)
+        dev = qml.device("default.qubit", wires=2)
 
         @qml.qnode(dev, diff_method=diff_method)
         def circuit_entropy(x):
@@ -234,7 +225,7 @@ def test_IsingXX_qnode_entropy_grad_torch(self, param, wires, base, diff_method)
         """Test entropy for a QNode gradient with torch."""
         import torch
 
-        dev = qml.device("default.qubit.legacy", wires=2)
+        dev = qml.device("default.qubit", wires=2)
 
         @qml.qnode(dev, interface="torch", diff_method=diff_method)
         def circuit_entropy(x):
@@ -286,7 +277,7 @@ def test_IsingXX_qnode_entropy_grad_tf(self, param, wires, base, diff_method, in
         """Test entropy for a QNode gradient with tf."""
         import tensorflow as tf
 
-        dev = qml.device("default.qubit.legacy", wires=2)
+        dev = qml.device("default.qubit", wires=2)
 
         @qml.qnode(dev, interface=interface, diff_method=diff_method)
         def circuit_entropy(x):
@@ -339,7 +330,7 @@ def test_IsingXX_qnode_entropy_grad_jax(self, param, wires, base, diff_method, i
         """Test entropy for a QNode gradient with Jax."""
         import jax
 
-        dev = qml.device("default.qubit.legacy", wires=2)
+        dev = qml.device("default.qubit", wires=2)
 
         @qml.qnode(dev, interface=interface, diff_method=diff_method)
         def circuit_entropy(x):
@@ -388,7 +379,7 @@ def test_IsingXX_qnode_entropy_grad_jax_jit(self, param, wires, base, diff_metho
         """Test entropy for a QNode gradient with Jax-jit."""
         import jax
 
-        dev = qml.device("default.qubit.legacy", wires=2)
+        dev = qml.device("default.qubit", wires=2)
 
         @qml.qnode(dev, interface=interface, diff_method=diff_method)
         def circuit_entropy(x):
@@ -401,19 +392,15 @@ def circuit_entropy(x):
 
         assert qml.math.allclose(grad_entropy, grad_expected_entropy, rtol=1e-04, atol=1e-05)
 
-    @pytest.mark.parametrize("device", devices)
-    def test_qnode_entropy_no_custom_wires(self, device):
-        """Test that entropy cannot be returned with custom wires."""
+    @pytest.mark.xfail(reason="until DQ2 port")
+    def test_qnode_entropy_custom_wires(self):
+        """Test that entropy can be returned with custom wires."""
 
-        dev = qml.device(device, wires=["a", 1])
+        dev = qml.device("default.qubit", wires=["a", 1])
 
         @qml.qnode(dev)
         def circuit_entropy(x):
             qml.IsingXX(x, wires=["a", 1])
             return qml.vn_entropy(wires=["a"])
 
-        with pytest.raises(
-            qml.QuantumFunctionError,
-            match="Returning the Von Neumann entropy is not supported when using custom wire labels",
-        ):
-            circuit_entropy(0.1)
+        assert np.isclose(circuit_entropy(0.1), expected_entropy_ising_xx(0.1))

From bcd09c47e8757c41a8a3fc5fd43929158c1d03c4 Mon Sep 17 00:00:00 2001
From: soranjh <40344468+soranjh@users.noreply.github.com>
Date: Wed, 6 Sep 2023 17:47:47 -0400
Subject: [PATCH 035/131] Use fermi sentence in qchem dipole and number
 functions (#4546)

---
 doc/releases/changelog-dev.md    |   3 +
 pennylane/qchem/dipole.py        |  28 +++++---
 pennylane/qchem/number.py        |  15 ++---
 tests/qchem/test_dipole.py       | 106 +++++++++----------------------
 tests/qchem/test_hamiltonians.py |  14 ++--
 5 files changed, 64 insertions(+), 102 deletions(-)

diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 59afddf3409..9da39f449fb 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -31,6 +31,8 @@
   [(4440)](https://github.com/PennyLaneAI/pennylane/pull/4440)
 
 <h3>Improvements 🛠</h3>
+* The qchem ``fermionic_dipole`` and ``particle_number`` functions are updated to use a ``FermiSentence``
+  [(#4546)](https://github.com/PennyLaneAI/pennylane/pull/4546)
 
 * Add the method ``add_transform`` and ``insert_front_transform`` transform in the ``TransformProgram``.
   [(#4559)](https://github.com/PennyLaneAI/pennylane/pull/4559)
@@ -186,6 +188,7 @@
 
 This release contains contributions from (in alphabetical order):
 
+Soran Jahangiri,
 Lillian M. A. Frederiksen,
 Romain Moyard,
 Mudit Pandey,
diff --git a/pennylane/qchem/dipole.py b/pennylane/qchem/dipole.py
index 876049cfa72..2ee1e2b32b7 100644
--- a/pennylane/qchem/dipole.py
+++ b/pennylane/qchem/dipole.py
@@ -15,6 +15,7 @@
 This module contains the functions needed for computing the dipole moment.
 """
 import pennylane as qml
+from pennylane.fermi import FermiSentence, FermiWord
 
 from .basis_data import atomic_numbers
 from .hartree_fock import scf
@@ -188,9 +189,16 @@ def fermionic_dipole(mol, cutoff=1.0e-18, core=None, active=None):
     >>>                   [3.42525091, 0.62391373, 0.1688554]], requires_grad=True)
     >>> mol = qml.qchem.Molecule(symbols, geometry, alpha=alpha)
     >>> args = [alpha]
-    >>> coeffs, ops = fermionic_dipole(mol)(*args)[2]
-    >>> ops
-    [[], [0, 0], [0, 2], [1, 1], [1, 3], [2, 0], [2, 2], [3, 1], [3, 3]]
+    >>> fermionic_dipole(mol)(*args)[2]
+    -0.4999999988651487 * a⁺(0) a(0)
+    + 0.82709948984052 * a⁺(0) a(2)
+    + -0.4999999988651487 * a⁺(1) a(1)
+    + 0.82709948984052 * a⁺(1) a(3)
+    + 0.82709948984052 * a⁺(2) a(0)
+    + -0.4999999899792451 * a⁺(2) a(2)
+    + 0.82709948984052 * a⁺(3) a(1)
+    + -0.4999999899792451 * a⁺(3) a(3)
+    + 1.0 * I
     """
 
     def _fermionic_dipole(*args):
@@ -200,8 +208,7 @@ def _fermionic_dipole(*args):
             *args (array[array[float]]): initial values of the differentiable parameters
 
         Returns:
-            tuple(array[float], list[list[int]]): the dipole moment coefficients and the indices of
-            the spin orbitals the creation and annihilation operators act on
+            FermiSentence: fermionic dipole moment
         """
         constants, integrals = dipole_integrals(mol, core, active)(*args)
 
@@ -211,12 +218,12 @@ def _fermionic_dipole(*args):
             nd[1] = nd[1] + atomic_numbers[s] * mol.coordinates[i][1]
             nd[2] = nd[2] + atomic_numbers[s] * mol.coordinates[i][2]
 
-        f = []
+        d_ferm = []
         for i in range(3):
-            coeffs, ops = fermionic_observable(constants[i], integrals[i], cutoff=cutoff)
-            f.append((qml.math.concatenate((nd[i], coeffs * (-1))), [[]] + ops))
+            f = fermionic_observable(constants[i], integrals[i], cutoff=cutoff, fs=True)
+            d_ferm.append(FermiSentence({FermiWord({}): nd[i][0]}) - f)
 
-        return f
+        return d_ferm
 
     return _fermionic_dipole
 
@@ -283,7 +290,8 @@ def dipole_moment(mol, cutoff=1.0e-18, core=None, active=None):
     >>> mol = qml.qchem.Molecule(symbols, geometry, alpha=alpha)
     >>> args = [alpha]
     >>> dipole_moment(mol)(*args)[2].ops
-    [PauliZ(wires=[0]),
+    [Identity(wires=[0]),
+     PauliZ(wires=[0]),
      PauliY(wires=[0]) @ PauliZ(wires=[1]) @ PauliY(wires=[2]),
      PauliX(wires=[0]) @ PauliZ(wires=[1]) @ PauliX(wires=[2]),
      PauliZ(wires=[1]),
diff --git a/pennylane/qchem/number.py b/pennylane/qchem/number.py
index 517f5f61826..d31a1d51429 100644
--- a/pennylane/qchem/number.py
+++ b/pennylane/qchem/number.py
@@ -14,7 +14,7 @@
 """
 This module contains the functions needed for computing the particle number observable.
 """
-from pennylane import numpy as np
+from pennylane.fermi import FermiSentence, FermiWord
 
 from .observable_hf import qubit_observable
 
@@ -57,14 +57,7 @@ def particle_number(orbitals):
     if orbitals <= 0:
         raise ValueError(f"'orbitals' must be greater than 0; got for 'orbitals' {orbitals}")
 
-    r = np.arange(orbitals)
-    table = np.vstack([r, r, np.ones([orbitals])]).T
+    sentence = FermiSentence({FermiWord({(0, i): "+", (1, i): "-"}): 1.0 for i in range(orbitals)})
+    sentence.simplify()
 
-    coeffs = np.array([])
-    ops = []
-
-    for i in table:
-        coeffs = np.concatenate((coeffs, np.array([i[2]])))
-        ops.append([int(i[0]), int(i[1])])
-
-    return qubit_observable((coeffs, ops))
+    return qubit_observable(sentence)
diff --git a/tests/qchem/test_dipole.py b/tests/qchem/test_dipole.py
index 42e249b89a6..1a11338f93a 100644
--- a/tests/qchem/test_dipole.py
+++ b/tests/qchem/test_dipole.py
@@ -21,7 +21,8 @@
 from pennylane import PauliX, PauliY, PauliZ
 from pennylane import numpy as np
 from pennylane import qchem
-from pennylane.operation import enable_new_opmath, disable_new_opmath
+from pennylane.fermi import from_string
+from pennylane.operation import disable_new_opmath, enable_new_opmath
 
 
 @pytest.mark.parametrize(
@@ -104,52 +105,25 @@ def test_dipole_integrals(symbols, geometry, charge, core, active, core_ref, int
             None,
             # x component of fermionic dipole computed with PL-QChem dipole (format is modified:
             # the signs of the coefficients, except that from the nuclear contribution, is flipped.
-            (
-                np.array(
-                    [
-                        2.869,
-                        -0.956224634652776,
-                        -0.782727697897828,
-                        0.532222940905614,
-                        -0.956224634652776,
-                        -0.782727697897828,
-                        0.532222940905614,
-                        -0.782727697897828,
-                        -1.42895581236226,
-                        -0.234699175620383,
-                        -0.782727697897828,
-                        -1.42895581236226,
-                        -0.234699175620383,
-                        0.532222940905614,
-                        -0.234699175620383,
-                        -0.483819552892797,
-                        0.532222940905614,
-                        -0.234699175620383,
-                        -0.483819552892797,
-                    ]
-                ),
-                [
-                    [],
-                    [0, 0],
-                    [0, 2],
-                    [0, 4],
-                    [1, 1],
-                    [1, 3],
-                    [1, 5],
-                    [2, 0],
-                    [2, 2],
-                    [2, 4],
-                    [3, 1],
-                    [3, 3],
-                    [3, 5],
-                    [4, 0],
-                    [4, 2],
-                    [4, 4],
-                    [5, 1],
-                    [5, 3],
-                    [5, 5],
-                ],
-            ),
+            2.869 * from_string("")
+            + -0.956224634652776 * from_string("0+ 0-")
+            + -0.782727697897828 * from_string("0+ 2-")
+            + 0.532222940905614 * from_string("0+ 4-")
+            + -0.956224634652776 * from_string("1+ 1-")
+            + -0.782727697897828 * from_string("1+ 3-")
+            + 0.532222940905614 * from_string("1+ 5-")
+            + -0.782727697897828 * from_string("2+ 0-")
+            + -1.42895581236226 * from_string("2+ 2-")
+            + -0.234699175620383 * from_string("2+ 4-")
+            + -0.782727697897828 * from_string("3+ 1-")
+            + -1.42895581236226 * from_string("3+ 3-")
+            + -0.234699175620383 * from_string("3+ 5-")
+            + 0.532222940905614 * from_string("4+ 0-")
+            + -0.234699175620383 * from_string("4+ 2-")
+            + -0.483819552892797 * from_string("4+ 4-")
+            + 0.532222940905614 * from_string("5+ 1-")
+            + -0.234699175620383 * from_string("5+ 3-")
+            + -0.483819552892797 * from_string("5+ 5-"),
         ),
         (
             ["H", "H", "H"],
@@ -161,34 +135,16 @@ def test_dipole_integrals(symbols, geometry, charge, core, active, core_ref, int
             [1, 2],
             # x component of fermionic dipole computed with PL-QChem dipole (format is modified:
             # the signs of the coefficients, except that from the nuclear contribution, is flipped.
-            (
-                np.array(
-                    [
-                        2.869,
-                        -1.912449269305551,
-                        -1.4289558123627388,
-                        -0.2346991756194219,
-                        -1.4289558123627388,
-                        -0.2346991756194219,
-                        -0.2346991756194219,
-                        -0.48381955289231976,
-                        -0.2346991756194219,
-                        -0.48381955289231976,
-                    ]
-                ),
-                [
-                    [],
-                    [],
-                    [0, 0],
-                    [0, 2],
-                    [1, 1],
-                    [1, 3],
-                    [2, 0],
-                    [2, 2],
-                    [3, 1],
-                    [3, 3],
-                ],
-            ),
+            2.869 * from_string("")
+            - 1.912449269305551 * from_string("")
+            + -1.4289558123627388 * from_string("0+ 0-")
+            + -0.2346991756194219 * from_string("0+ 2-")
+            + -1.4289558123627388 * from_string("1+ 1-")
+            + -0.2346991756194219 * from_string("1+ 3-")
+            + -0.2346991756194219 * from_string("2+ 0-")
+            + -0.48381955289231976 * from_string("2+ 2-")
+            + -0.2346991756194219 * from_string("3+ 1-")
+            + -0.48381955289231976 * from_string("3+ 3-"),
         ),
     ],
 )
diff --git a/tests/qchem/test_hamiltonians.py b/tests/qchem/test_hamiltonians.py
index d608630df2d..837807c3ec5 100644
--- a/tests/qchem/test_hamiltonians.py
+++ b/tests/qchem/test_hamiltonians.py
@@ -463,7 +463,9 @@ def circuit(*args):
                 qml.PauliX(0)
                 qml.PauliX(1)
                 qml.DoubleExcitation(0.22350048111151138, wires=[0, 1, 2, 3])
-                h_qubit = qchem.diff_hamiltonian(mol, fs=False)(*args)
+                enable_new_opmath()
+                h_qubit = qchem.diff_hamiltonian(mol, fs=True)(*args)
+                disable_new_opmath()
                 return qml.expval(h_qubit)
 
             return circuit
@@ -471,16 +473,16 @@ def circuit(*args):
         grad_jax = jax.grad(energy(mol), argnums=0)(*args)
 
         alpha_1 = np.array(
-            [[3.42515091, 0.62391373, 0.1688554], [3.42525091, 0.62391373, 0.1688554]],
-        )  # alpha[0][0] -= 0.0001
+            [[3.42425091, 0.62391373, 0.1688554], [3.42525091, 0.62391373, 0.1688554]],
+        )  # alpha[0][0] -= 0.001
 
         alpha_2 = np.array(
-            [[3.42535091, 0.62391373, 0.1688554], [3.42525091, 0.62391373, 0.1688554]],
-        )  # alpha[0][0] += 0.0001
+            [[3.42625091, 0.62391373, 0.1688554], [3.42525091, 0.62391373, 0.1688554]],
+        )  # alpha[0][0] += 0.001
 
         e_1 = energy(mol)(*[alpha_1])
         e_2 = energy(mol)(*[alpha_2])
 
-        grad_finitediff = (e_2 - e_1) / 0.0002
+        grad_finitediff = (e_2 - e_1) / 0.002
 
         assert np.allclose(grad_jax[0][0], grad_finitediff, rtol=1e-02)

From 33ec1e61e834ec1fa0f5f519239466e844fd442c Mon Sep 17 00:00:00 2001
From: Matthew Silverman <matthews@xanadu.ai>
Date: Thu, 7 Sep 2023 11:10:00 -0400
Subject: [PATCH 036/131] fix np.random.seed usage (#4581)

No clue why regular CI doesn't catch this, but [docker CI
does](https://github.com/PennyLaneAI/pennylane/actions/runs/6102750018/job/16561865816#step:4:4086).
A lot.
---
 tests/gradients/finite_diff/test_spsa_gradient.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/gradients/finite_diff/test_spsa_gradient.py b/tests/gradients/finite_diff/test_spsa_gradient.py
index ddfde3c3533..1721961d479 100644
--- a/tests/gradients/finite_diff/test_spsa_gradient.py
+++ b/tests/gradients/finite_diff/test_spsa_gradient.py
@@ -81,7 +81,7 @@ def test_same_seeds(self):
         num = 5
         rng = np.random.default_rng(42)
         first_direction = _rademacher_sampler(ids, num, rng=rng)
-        np.random.seed = 0  # Setting the global seed should have no effect.
+        np.random.seed(0)  # Setting the global seed should have no effect.
         rng = np.random.default_rng(42)
         second_direction = _rademacher_sampler(ids, num, rng=rng)
         assert np.allclose(first_direction, second_direction)

From 58fb45f9dcb683a238ee6cff18072752864fb9ca Mon Sep 17 00:00:00 2001
From: soranjh <40344468+soranjh@users.noreply.github.com>
Date: Thu, 7 Sep 2023 17:23:04 -0400
Subject: [PATCH 037/131] Remove qchem deprecated functionality for fermionic
 observables (#4556)

---
 doc/development/deprecations.rst  |  28 +-
 doc/introduction/chemistry.rst    |   2 +-
 doc/releases/changelog-dev.md     |   6 +-
 pennylane/fermi/conversion.py     | 109 +------
 pennylane/qchem/__init__.py       |   2 +-
 pennylane/qchem/dipole.py         |   2 +-
 pennylane/qchem/hamiltonian.py    |  22 +-
 pennylane/qchem/observable_hf.py  | 171 +++--------
 tests/fermi/test_fermi_mapping.py |   6 -
 tests/qchem/test_hamiltonians.py  | 131 +--------
 tests/qchem/test_observable_hf.py | 464 +-----------------------------
 11 files changed, 74 insertions(+), 869 deletions(-)

diff --git a/doc/development/deprecations.rst b/doc/development/deprecations.rst
index 03eb3009f1e..b60d381599c 100644
--- a/doc/development/deprecations.rst
+++ b/doc/development/deprecations.rst
@@ -62,17 +62,27 @@ Pending deprecations
   `discussion forum <https://discuss.pennylane.ai/>`_.
 
   - Deprecated in v0.31
+  
+* The ``prep`` keyword argument in ``QuantumScript`` is deprecated and will be removed from ``QuantumScript``.
+  ``StatePrepBase`` operations should be placed at the beginning of the `ops` list instead.
+
+  - Deprecated in v0.33
+  - Will be removed in v0.34
 
-* ``qml.qchem.jordan_wigner`` is deprecated, and usage will now raise a warning.
+
+Completed deprecation cycles
+----------------------------
+
+* ``qml.qchem.jordan_wigner`` had been removed.
   Use ``qml.jordan_wigner`` instead. List input to define the fermionic operator
-  is also deprecated; the fermionic operators ``qml.FermiA``, ``qml.FermiC``,
+  is no longer accepted; the fermionic operators ``qml.FermiA``, ``qml.FermiC``,
   ``qml.FermiWord`` and ``qml.FermiSentence`` should be used instead. See the
   :mod:`pennylane.fermi` module documentation and the
   `Fermionic Operator <https://pennylane.ai/qml/demos/tutorial_fermionic_operators>`_
   tutorial for more details.
 
   - Deprecated in v0.32
-  - Will be removed in v0.33
+  - Removed in v0.33
 
 * The ``tuple`` input type in ``qubit_observable`` has been deprecated. Please use a fermionic
   operator object. The ``tuple`` return type in ``fermionic_hamiltonian`` and
@@ -80,17 +90,7 @@ Pending deprecations
   by default.
 
   - Deprecated in v0.32
-  - Will be removed in v0.33
-  
-* The ``prep`` keyword argument in ``QuantumScript`` is deprecated and will be removed from `QuantumScript`.
-  ``StatePrepBase`` operations should be placed at the beginning of the `ops` list instead.
-
-  - Deprecated in v0.33
-  - Will be removed in v0.34
-
-
-Completed deprecation cycles
-----------------------------
+  - Removed in v0.33
 
 * The ``sampler_seed`` argument of ``qml.gradients.spsa_grad`` has been removed.
   Instead, the ``sampler_rng`` argument should be set, either to an integer value, which will be used
diff --git a/doc/introduction/chemistry.rst b/doc/introduction/chemistry.rst
index cf056f68ec3..71a5de0f0b9 100644
--- a/doc/introduction/chemistry.rst
+++ b/doc/introduction/chemistry.rst
@@ -207,7 +207,7 @@ Differentiable observables
     ~pennylane.qchem.fermionic_dipole
     ~pennylane.qchem.fermionic_hamiltonian
     ~pennylane.qchem.fermionic_observable
-    ~pennylane.qchem.jordan_wigner
+    ~pennylane.jordan_wigner
     ~pennylane.qchem.molecular_hamiltonian
     ~pennylane.qchem.qubit_observable
 
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 9da39f449fb..b463e4eee6e 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -31,8 +31,12 @@
   [(4440)](https://github.com/PennyLaneAI/pennylane/pull/4440)
 
 <h3>Improvements 🛠</h3>
-* The qchem ``fermionic_dipole`` and ``particle_number`` functions are updated to use a ``FermiSentence``
+
+* The qchem ``fermionic_dipole`` and ``particle_number`` functions are updated to use a
+  ``FermiSentence``. The deprecated features for using tuples to represent fermionic operations are
+  removed.
   [(#4546)](https://github.com/PennyLaneAI/pennylane/pull/4546)
+  [(#4556)](https://github.com/PennyLaneAI/pennylane/pull/4556)
 
 * Add the method ``add_transform`` and ``insert_front_transform`` transform in the ``TransformProgram``.
   [(#4559)](https://github.com/PennyLaneAI/pennylane/pull/4559)
diff --git a/pennylane/fermi/conversion.py b/pennylane/fermi/conversion.py
index 1fc4b42b345..45b109fd28a 100644
--- a/pennylane/fermi/conversion.py
+++ b/pennylane/fermi/conversion.py
@@ -13,21 +13,16 @@
 # limitations under the License.
 """Functions to convert a fermionic operator to the qubit basis."""
 
-import warnings
 from functools import singledispatch
 from typing import Union
-import numpy as np
-
-import pennylane as qml
-from pennylane.operation import Operator, Tensor
+from pennylane.operation import Operator
 from pennylane.pauli import PauliWord, PauliSentence
-from pennylane.pauli.utils import _get_pauli_map, _pauli_mult
 from .fermionic import FermiWord, FermiSentence
 
 
 # pylint: disable=unexpected-keyword-arg
 def jordan_wigner(
-    fermi_operator: (Union[FermiWord, FermiSentence]), **kwargs
+    fermi_operator: (Union[FermiWord, FermiSentence]), ps: bool = False, wire_map: dict = None
 ) -> Union[Operator, PauliSentence]:
     r"""Convert a fermionic operator to a qubit operator using the Jordan-Wigner mapping.
 
@@ -77,13 +72,12 @@ def jordan_wigner(
     + (0.25+0j) * X(2) @ X(3)
     + 0.25j * X(2) @ Y(3)
     """
-    return _jordan_wigner_dispatch(fermi_operator, **kwargs)
+    return _jordan_wigner_dispatch(fermi_operator, ps, wire_map)
 
 
 @singledispatch
-def _jordan_wigner_dispatch(fermi_operator, **kwargs):
-    """Dispatches to appropriate function if fermi_operator is a FermiWord,
-    FermiSentence or list"""
+def _jordan_wigner_dispatch(fermi_operator, ps, wire_map):
+    """Dispatches to appropriate function if fermi_operator is a FermiWord or FermiSentence."""
     raise ValueError(f"fermi_operator must be a FermiWord or FermiSentence, got: {fermi_operator}")
 
 
@@ -140,96 +134,3 @@ def _(fermi_operator: FermiSentence, ps=False, wire_map=None):
         return qubit_operator.map_wires(wire_map)
 
     return qubit_operator
-
-
-@_jordan_wigner_dispatch.register
-def _jordan_wigner_legacy(op: list, notation="physicist"):  # pylint:disable=too-many-branches
-    r"""Convert a fermionic operator to a qubit operator using the Jordan-Wigner mapping.
-
-    For instance, the one-body fermionic operator :math:`a_2^\dagger a_0` should be constructed as
-    [2, 0]. The two-body operator :math:`a_4^\dagger a_3^\dagger a_2 a_1` should be constructed
-    as [4, 3, 2, 1] with ``notation='physicist'``. If ``notation`` is set to ``'chemist'``, the
-    two-body operator [4, 3, 2, 1] is constructed as :math:`a_4^\dagger a_3 a_2^\dagger a_1`.
-
-    Args:
-        op (list[int]): the fermionic operator
-        notation (str): notation specifying the order of the two-body fermionic operators
-
-    Returns:
-        tuple(list[complex], list[Operation]): list of coefficients and qubit operators
-
-    **Example**
-
-    >>> f  = [0, 0]
-    >>> q = jordan_wigner(f)
-    >>> q # corresponds to :math:`\frac{1}{2}(I_0 - Z_0)`
-    ([(0.5+0j), (-0.5+0j)], [Identity(wires=[0]), PauliZ(wires=[0])])
-    """
-
-    warnings.warn(
-        "List input for the jordan_wigner function is deprecated; please use the fermionic operators format. For "
-        "details, see the Fermionic Operators tutorial: https://pennylane.ai/qml/demos/tutorial_fermionic_operators"
-    )
-
-    if len(op) == 1:
-        op = [(op[0], 1)]
-
-    if len(op) == 2:
-        op = [(op[0], 1), (op[1], 0)]
-
-    if len(op) == 4:
-        if notation == "physicist":
-            if op[0] == op[1] or op[2] == op[3]:
-                return [0], [qml.Identity(wires=[min(op)])]
-            op = [(op[0], 1), (op[1], 1), (op[2], 0), (op[3], 0)]
-        elif notation == "chemist":
-            if (op[0] == op[2] or op[1] == op[3]) and op[1] != op[2]:
-                return [0], [qml.Identity(wires=[min(op)])]
-            op = [(op[0], 1), (op[1], 0), (op[2], 1), (op[3], 0)]
-        else:
-            raise ValueError(
-                f"Currently, the only supported notations for the two-body terms are 'physicist'"
-                f" and 'chemist', got notation = '{notation}'."
-            )
-
-    q = [[(0, "I"), 1.0]]
-    for l in op:
-        z = [(index, "Z") for index in range(l[0])]
-        x = z + [(l[0], "X"), 0.5]
-        if l[1]:
-            y = z + [(l[0], "Y"), -0.5j]
-        else:
-            y = z + [(l[0], "Y"), 0.5j]
-
-        m = []
-        for t1 in q:
-            for t2 in [x, y]:
-                q1, c1 = _pauli_mult(t1[:-1], t2[:-1])
-                m.append(q1 + [c1 * t1[-1] * t2[-1]])
-        q = m
-
-    c = [p[-1] for p in q]
-    o = [p[:-1] for p in q]
-
-    for item in o:
-        k = [i for i, x in enumerate(o) if x == item]
-        if len(k) >= 2:
-            for j in k[::-1][:-1]:
-                del o[j]
-                c[k[0]] = c[k[0]] + c[j]
-                del c[j]
-
-    # Pauli gates objects pregenerated for speed
-    pauli_map = _get_pauli_map(np.max(op))
-    for i, term in enumerate(o):
-        if len(term) == 0:
-            # moved function from qchem.observable_hf without any changes to tests
-            # codecov complained this line is not covered
-            # function to be deprecated next release
-            # not going to write a test to cover this line
-            o[i] = qml.Identity(0)  # pragma: no cover
-        else:
-            k = [pauli_map[t[0]][t[1]] for t in term]
-            o[i] = Tensor(*k)
-
-    return c, o
diff --git a/pennylane/qchem/__init__.py b/pennylane/qchem/__init__.py
index 9af19dedd25..6c91b561243 100644
--- a/pennylane/qchem/__init__.py
+++ b/pennylane/qchem/__init__.py
@@ -56,7 +56,7 @@
     core_matrix,
 )
 from .molecule import Molecule
-from .observable_hf import fermionic_observable, qubit_observable, jordan_wigner
+from .observable_hf import fermionic_observable, qubit_observable
 from .number import particle_number
 from .spin import spin2, spinz
 from .structure import (
diff --git a/pennylane/qchem/dipole.py b/pennylane/qchem/dipole.py
index 2ee1e2b32b7..909869b8f8d 100644
--- a/pennylane/qchem/dipole.py
+++ b/pennylane/qchem/dipole.py
@@ -220,7 +220,7 @@ def _fermionic_dipole(*args):
 
         d_ferm = []
         for i in range(3):
-            f = fermionic_observable(constants[i], integrals[i], cutoff=cutoff, fs=True)
+            f = fermionic_observable(constants[i], integrals[i], cutoff=cutoff)
             d_ferm.append(FermiSentence({FermiWord({}): nd[i][0]}) - f)
 
         return d_ferm
diff --git a/pennylane/qchem/hamiltonian.py b/pennylane/qchem/hamiltonian.py
index cb3bf6e7dd5..629bc229ae5 100644
--- a/pennylane/qchem/hamiltonian.py
+++ b/pennylane/qchem/hamiltonian.py
@@ -15,9 +15,6 @@
 This module contains the functions needed for computing the molecular Hamiltonian.
 """
 # pylint: disable= too-many-branches, too-many-arguments, too-many-locals, too-many-nested-blocks
-
-import warnings
-
 import pennylane as qml
 
 from .hartree_fock import nuclear_energy, scf
@@ -144,7 +141,7 @@ def _electron_integrals(*args):
     return _electron_integrals
 
 
-def fermionic_hamiltonian(mol, cutoff=1.0e-12, core=None, active=None, fs=False):
+def fermionic_hamiltonian(mol, cutoff=1.0e-12, core=None, active=None):
     r"""Return a function that computes the fermionic Hamiltonian.
 
     Args:
@@ -152,7 +149,6 @@ def fermionic_hamiltonian(mol, cutoff=1.0e-12, core=None, active=None, fs=False)
         cutoff (float): cutoff value for discarding the negligible electronic integrals
         core (list[int]): indices of the core orbitals
         active (list[int]): indices of the active orbitals
-        fs (bool): if True, a fermi sentence will be returned
 
     Returns:
         function: function that computes the fermionic hamiltonian
@@ -175,24 +171,17 @@ def _fermionic_hamiltonian(*args):
             *args (array[array[float]]): initial values of the differentiable parameters
 
         Returns:
-            Union[FermiSentence, tuple(array[float], list[list[int]])]: fermionic Hamiltonian
+            FermiSentence: fermionic Hamiltonian
         """
 
         core_constant, one, two = electron_integrals(mol, core, active)(*args)
 
-        if not fs:
-            warnings.warn(
-                "This function will return a fermionic operator by default in the next release. For "
-                "details, see the Fermionic Operators tutorial: "
-                "https://pennylane.ai/qml/demos/tutorial_fermionic_operators. "
-                "Currently, a fermionic operator can be returned by setting the `fs` kwarg to `True`."
-            )
-        return fermionic_observable(core_constant, one, two, cutoff, fs)
+        return fermionic_observable(core_constant, one, two, cutoff)
 
     return _fermionic_hamiltonian
 
 
-def diff_hamiltonian(mol, cutoff=1.0e-12, core=None, active=None, fs=True):
+def diff_hamiltonian(mol, cutoff=1.0e-12, core=None, active=None):
     r"""Return a function that computes the qubit Hamiltonian.
 
     Args:
@@ -200,7 +189,6 @@ def diff_hamiltonian(mol, cutoff=1.0e-12, core=None, active=None, fs=True):
         cutoff (float): cutoff value for discarding the negligible electronic integrals
         core (list[int]): indices of the core orbitals
         active (list[int]): indices of the active orbitals
-        fs (bool): if True, a fermi sentence is constructed and used internally
 
     Returns:
         function: function that computes the qubit hamiltonian
@@ -232,7 +220,7 @@ def _molecular_hamiltonian(*args):
             Hamiltonian: the qubit Hamiltonian
         """
 
-        h_ferm = fermionic_hamiltonian(mol, cutoff, core, active, fs=fs)(*args)
+        h_ferm = fermionic_hamiltonian(mol, cutoff, core, active)(*args)
 
         return qubit_observable(h_ferm)
 
diff --git a/pennylane/qchem/observable_hf.py b/pennylane/qchem/observable_hf.py
index 0819318cac4..7005928a069 100644
--- a/pennylane/qchem/observable_hf.py
+++ b/pennylane/qchem/observable_hf.py
@@ -14,10 +14,7 @@
 """
 This module contains the functions needed for creating fermionic and qubit observables.
 """
-import warnings
-
 # pylint: disable= too-many-branches, too-many-return-statements
-
 import pennylane as qml
 from pennylane import numpy as np
 from pennylane.fermi import FermiSentence, FermiWord
@@ -25,7 +22,7 @@
 from pennylane.pauli.utils import simplify
 
 
-def fermionic_observable(constant, one=None, two=None, cutoff=1.0e-12, fs=False):
+def fermionic_observable(constant, one=None, two=None, cutoff=1.0e-12):
     r"""Create a fermionic observable from molecular orbital integrals.
 
     Args:
@@ -33,24 +30,15 @@ def fermionic_observable(constant, one=None, two=None, cutoff=1.0e-12, fs=False)
         one (array[float]): the one-particle molecular orbital integrals
         two (array[float]): the two-particle molecular orbital integrals
         cutoff (float): cutoff value for discarding the negligible integrals
-        fs (bool): if True, a fermi sentence will be returned
 
     Returns:
-        Union[FermiSentence, tuple(array[float], list[list[int]])]: fermionic Hamiltonian
+        FermiSentence: fermionic observable
 
     **Example**
 
     >>> constant = np.array([1.0])
     >>> integral = np.array([[0.5, -0.8270995], [-0.8270995, 0.5]])
-    >>> coeffs, ops = fermionic_observable(constant, integral)
-    >>> ops
-    [[], [0, 0], [0, 2], [1, 1], [1, 3], [2, 0], [2, 2], [3, 1], [3, 3]]
-
-    If the `fs` kwarg is `True`, a fermionic operator is returned.
-
-    >>> constant = np.array([1.0])
-    >>> integral = np.array([[0.5, -0.8270995], [-0.8270995, 0.5]])
-    >>> fermionic_observable(constant, integral, fs=True)
+    >>> fermionic_observable(constant, integral)
     1.0 * I
     + 0.5 * a⁺(0) a(0)
     + -0.8270995 * a⁺(0) a(2)
@@ -96,54 +84,31 @@ def fermionic_observable(constant, one=None, two=None, cutoff=1.0e-12, fs=False)
     if indices_sort:
         indices_sort = qml.math.array(indices_sort)
 
-    if fs:
-        sentence = FermiSentence({FermiWord({}): constant[0]})
-        for c, o in zip(coeffs[indices_sort], sorted(operators)):
-            if len(o) == 2:
-                sentence.update({FermiWord({(0, o[0]): "+", (1, o[1]): "-"}): c})
-            if len(o) == 4:
-                sentence.update(
-                    {FermiWord({(0, o[0]): "+", (1, o[1]): "+", (2, o[2]): "-", (3, o[3]): "-"}): c}
-                )
-        sentence.simplify()
-
-        return sentence
-
-    warnings.warn(
-        "This function will return a fermionic operator by default in the next release. For details, "
-        "see the Fermionic Operators tutorial: "
-        "https://pennylane.ai/qml/demos/tutorial_fermionic_operators. "
-        "Currently, a fermionic operator can be returned by setting the `fs` kwarg to `True`."
-    )
+    sentence = FermiSentence({FermiWord({}): constant[0]})
+    for c, o in zip(coeffs[indices_sort], sorted(operators)):
+        if len(o) == 2:
+            sentence.update({FermiWord({(0, o[0]): "+", (1, o[1]): "-"}): c})
+        if len(o) == 4:
+            sentence.update(
+                {FermiWord({(0, o[0]): "+", (1, o[1]): "+", (2, o[2]): "-", (3, o[3]): "-"}): c}
+            )
+    sentence.simplify()
 
-    return coeffs[indices_sort], sorted(operators)
+    return sentence
 
 
 def qubit_observable(o_ferm, cutoff=1.0e-12):
     r"""Convert a fermionic observable to a PennyLane qubit observable.
 
-    The fermionic operator is a tuple containing the fermionic coefficients and operators. For
-    instance, the one-body fermionic operator :math:`a_2^\dagger a_0` is specified as [2, 0] and the
-    two-body operator :math:`a_4^\dagger a_3^\dagger a_2 a_1` is specified as [4, 3, 2, 1].
-
     Args:
-        o_ferm Union[FermiSentence, tuple(array[float], list[list[int]])]: fermionic operator
+        o_ferm (Union[FermiWord, FermiSentence]): fermionic operator
         cutoff (float): cutoff value for discarding the negligible terms
 
     Returns:
-        (Operator): Simplified PennyLane Hamiltonian
+        Operator: Simplified PennyLane Hamiltonian
 
     **Example**
 
-    >>> coeffs = np.array([1.0, 1.0])
-    >>> ops = [[0, 0], [0, 0]]
-    >>> f = (coeffs, ops)
-    >>> print(qubit_observable(f))
-    ((-1+0j)) [Z0]
-    + ((1+0j)) [I0]
-
-    The input can also be a fermionic operator.
-
     >>> w1 = qml.fermi.FermiWord({(0, 0) : '+', (1, 1) : '-'})
     >>> w2 = qml.fermi.FermiWord({(0, 1) : '+', (1, 2) : '-'})
     >>> s = qml.fermi.FermiSentence({w1 : 1.2, w2: 3.1})
@@ -160,97 +125,31 @@ def qubit_observable(o_ferm, cutoff=1.0e-12):
     If the new op-math is active, an arithmetic operator is returned.
 
     >>> qml.operation.enable_new_opmath()
-    >>> coeffs = np.array([1.0, 1.0])
-    >>> ops = [[0, 0], [0, 0]]
-    >>> f = (coeffs, ops)
-    >>> print(qubit_observable(f))
-    Identity(wires=[0]) + ((-1+0j)*(PauliZ(wires=[0])))
+    >>> w1 = qml.fermi.FermiWord({(0, 0) : '+', (1, 1) : '-'})
+    >>> w2 = qml.fermi.FermiWord({(0, 0) : '+', (1, 1) : '-'})
+    >>> s = qml.fermi.FermiSentence({w1 : 1.2, w2: 3.1})
+    >>> print(qubit_observable(s))
+    (-0.775j*(PauliY(wires=[0]) @ PauliX(wires=[1])))
+    + ((0.775+0j)*(PauliY(wires=[0]) @ PauliY(wires=[1])))
+    + ((0.775+0j)*(PauliX(wires=[0]) @ PauliX(wires=[1])))
+    + (0.775j*(PauliX(wires=[0]) @ PauliY(wires=[1])))
     """
-    if isinstance(o_ferm, (FermiWord, FermiSentence)):
-        h = qml.jordan_wigner(o_ferm, ps=True)
-        h.simplify(tol=cutoff)
-
-        if active_new_opmath():
-            if not h.wires:
-                return h.operation(wire_order=[0])
-            return h.operation()
+    h = qml.jordan_wigner(o_ferm, ps=True)
+    h.simplify(tol=cutoff)
 
+    if active_new_opmath():
         if not h.wires:
-            h = h.hamiltonian(wire_order=[0])
-            return qml.Hamiltonian(
-                h.coeffs, [qml.Identity(0) if o.name == "Identity" else o for o in h.ops]
-            )
+            return h.operation(wire_order=[0])
+        return h.operation()
 
-        h = h.hamiltonian()
-
-        return simplify(
-            qml.Hamiltonian(
-                h.coeffs, [qml.Identity(0) if o.name == "Identity" else o for o in h.ops]
-            )
+    if not h.wires:
+        h = h.hamiltonian(wire_order=[0])
+        return qml.Hamiltonian(
+            h.coeffs, [qml.Identity(0) if o.name == "Identity" else o for o in h.ops]
         )
 
-    warnings.warn(
-        "Tuple input for the qubit_observable function is deprecated; please use the fermionic "
-        "operators format. For details, see the Fermionic Operators tutorial: "
-        "https://pennylane.ai/qml/demos/tutorial_fermionic_operators"
-    )
-
-    ops = []
-    coeffs = qml.math.array([])
-
-    for n, t in enumerate(o_ferm[1]):
-        if len(t) == 0:
-            ops = ops + [qml.Identity(0)]
-            coeffs = qml.math.array([0.0])
-            coeffs = coeffs + o_ferm[0][n]
-        else:
-            op = jordan_wigner(t)
-            if op != 0:
-                ops = ops + op[1]
-                coeffs = qml.math.concatenate([coeffs, qml.math.array(op[0]) * o_ferm[0][n]])
-
-    if active_new_opmath():
-        ps = qml.dot(coeffs, ops, pauli=True)
-        ps.simplify(tol=cutoff)
-
-        if len(ps) == 0:
-            return qml.s_prod(
-                0, qml.Identity(ops[0].wires[0])
-            )  # use any op and any wire to represent the null op
-        if (len(ps) == 1) and ((identity := qml.pauli.PauliWord({})) in ps):
-            return qml.s_prod(
-                ps[identity], qml.Identity(ops[0].wires[0])
-            )  # use any op and any wire to represent the null op
-        return ps.operation()
-
-    return simplify(qml.Hamiltonian(coeffs, ops), cutoff=cutoff)
-
-
-def jordan_wigner(op: list, notation="physicist"):  # pylint:disable=too-many-branches
-    r"""Convert a fermionic operator to a qubit operator using the Jordan-Wigner mapping.
+    h = h.hamiltonian()
 
-    For instance, the one-body fermionic operator :math:`a_2^\dagger a_0` should be constructed as
-    [2, 0]. The two-body operator :math:`a_4^\dagger a_3^\dagger a_2 a_1` should be constructed
-    as [4, 3, 2, 1] with ``notation='physicist'``. If ``notation`` is set to ``'chemist'``, the
-    two-body operator [4, 3, 2, 1] is constructed as :math:`a_4^\dagger a_3 a_2^\dagger a_1`.
-
-    Args:
-        op (list[int]): the fermionic operator
-        notation (str): notation specifying the order of the two-body fermionic operators
-
-    Returns:
-        tuple(list[complex], list[Operation]): list of coefficients and qubit operators
-
-    **Example**
-
-    >>> f  = [0, 0]
-    >>> q = jordan_wigner(f)
-    >>> q # corresponds to :math:`\frac{1}{2}(I_0 - Z_0)`
-    ([(0.5+0j), (-0.5+0j)], [Identity(wires=[0]), PauliZ(wires=[0])])
-    """
-
-    warnings.warn(
-        "Use of qml.qchem.jordan_wigner is deprecated; please use qml.jordan_wigner instead."
+    return simplify(
+        qml.Hamiltonian(h.coeffs, [qml.Identity(0) if o.name == "Identity" else o for o in h.ops])
     )
-
-    return qml.fermi.jordan_wigner(op, notation=notation)
diff --git a/tests/fermi/test_fermi_mapping.py b/tests/fermi/test_fermi_mapping.py
index caa6a744dd2..21f218eb405 100644
--- a/tests/fermi/test_fermi_mapping.py
+++ b/tests/fermi/test_fermi_mapping.py
@@ -654,9 +654,3 @@ def test_providing_wire_map_fermi_word_to_ps(wire_map, ops):
     op.simplify()
 
     assert ps == op
-
-
-def test_list_input_raises_warning():
-    """Test that using the old input format (list) raises a deprecation warning"""
-    with pytest.warns(UserWarning, match="List input for the jordan_wigner function is deprecated"):
-        _ = qml.jordan_wigner([3, 3, 3, 1], notation="chemist")
diff --git a/tests/qchem/test_hamiltonians.py b/tests/qchem/test_hamiltonians.py
index 837807c3ec5..c6c9f3c091a 100644
--- a/tests/qchem/test_hamiltonians.py
+++ b/tests/qchem/test_hamiltonians.py
@@ -107,123 +107,6 @@ def test_electron_integrals(symbols, geometry, core, active, e_core, one_ref, tw
     assert np.allclose(two, two_ref)
 
 
-@pytest.mark.parametrize(
-    ("symbols", "geometry", "alpha", "coeffs_h_ref", "ops_h_ref"),
-    [
-        (
-            ["H", "H"],
-            np.array([[0.0, 0.0, 0.0], [0.0, 0.0, 1.0]], requires_grad=False),
-            np.array(
-                [[3.42525091, 0.62391373, 0.1688554], [3.42525091, 0.62391373, 0.1688554]],
-                requires_grad=True,
-            ),
-            # Hamiltonian coefficients and operators computed with OpenFermion using
-            # openfermion.transforms.get_fermion_operator(molecule.get_molecular_hamiltonian())
-            # The "^" symbols in the operators are removed and "," is added for consistency
-            np.array(
-                [
-                    1.0000000000321256,
-                    -1.3902192706002598,
-                    0.35721953951840535,
-                    0.08512072192002007,
-                    0.35721953951840535,
-                    0.08512072192002007,
-                    0.08512072192002007,
-                    0.35092657803574406,
-                    0.08512072192002007,
-                    0.35092657803574406,
-                    0.35721953951840535,
-                    0.08512072192002007,
-                    -1.3902192706002598,
-                    0.35721953951840535,
-                    0.08512072192002007,
-                    0.08512072192002007,
-                    0.35092657803574406,
-                    0.08512072192002007,
-                    0.35092657803574406,
-                    0.35092657803574495,
-                    0.08512072192002007,
-                    0.35092657803574495,
-                    0.08512072192002007,
-                    -0.2916533049477536,
-                    0.08512072192002007,
-                    0.3694183466586136,
-                    0.08512072192002007,
-                    0.3694183466586136,
-                    0.35092657803574495,
-                    0.08512072192002007,
-                    0.35092657803574495,
-                    0.08512072192002007,
-                    0.08512072192002007,
-                    0.3694183466586136,
-                    -0.2916533049477536,
-                    0.08512072192002007,
-                    0.3694183466586136,
-                ]
-            ),
-            [
-                [],
-                [0, 0],
-                [0, 0, 0, 0],
-                [0, 0, 2, 2],
-                [0, 1, 1, 0],
-                [0, 1, 3, 2],
-                [0, 2, 0, 2],
-                [0, 2, 2, 0],
-                [0, 3, 1, 2],
-                [0, 3, 3, 0],
-                [1, 0, 0, 1],
-                [1, 0, 2, 3],
-                [1, 1],
-                [1, 1, 1, 1],
-                [1, 1, 3, 3],
-                [1, 2, 0, 3],
-                [1, 2, 2, 1],
-                [1, 3, 1, 3],
-                [1, 3, 3, 1],
-                [2, 0, 0, 2],
-                [2, 0, 2, 0],
-                [2, 1, 1, 2],
-                [2, 1, 3, 0],
-                [2, 2],
-                [2, 2, 0, 0],
-                [2, 2, 2, 2],
-                [2, 3, 1, 0],
-                [2, 3, 3, 2],
-                [3, 0, 0, 3],
-                [3, 0, 2, 1],
-                [3, 1, 1, 3],
-                [3, 1, 3, 1],
-                [3, 2, 0, 1],
-                [3, 2, 2, 3],
-                [3, 3],
-                [3, 3, 1, 1],
-                [3, 3, 3, 3],
-            ],
-        )
-    ],
-)
-def test_fermionic_hamiltonian_fs_False(symbols, geometry, alpha, coeffs_h_ref, ops_h_ref):
-    r"""Test that fermionic_hamiltonian returns the correct Hamiltonian."""
-    # TODO: remove this test when supporting tuple output by fermionic_hamiltonian is deprecated.
-    mol = qchem.Molecule(symbols, geometry, alpha=alpha)
-    args = [alpha]
-    h = qchem.fermionic_hamiltonian(mol, fs=False)(*args)
-
-    assert np.allclose(h[0], coeffs_h_ref)
-    assert h[1] == ops_h_ref
-
-
-def test_fermionic_hamiltonian_warning():
-    r"""Test that a warning is raised if `fs = False` in fermionic_hamiltonian."""
-    # TODO: remove this test when supporting tuple output by fermionic_hamiltonian is deprecated.
-    symbols = ["H", "H"]
-    geometry = np.array([[0.0, 0.0, 0.0], [0.0, 0.0, 1.0]], requires_grad=False)
-    mol = qml.qchem.Molecule(symbols, geometry)
-    with pytest.warns(UserWarning, match="This function will return a fermionic operator"):
-        qchem.fermionic_hamiltonian(mol, fs=False)()
-
-
 @pytest.mark.parametrize(
     ("symbols", "geometry", "alpha", "h_ref"),
     [
@@ -283,7 +166,7 @@ def test_fermionic_hamiltonian(symbols, geometry, alpha, h_ref):
     r"""Test that fermionic_hamiltonian returns the correct Hamiltonian."""
     mol = qchem.Molecule(symbols, geometry, alpha=alpha)
     args = [alpha]
-    h = qchem.fermionic_hamiltonian(mol, fs=True)(*args)
+    h = qchem.fermionic_hamiltonian(mol)(*args)
 
     assert np.allclose(list(h.values()), list(h_ref.values()))
     assert h.keys() == h_ref.keys()
@@ -341,13 +224,12 @@ def test_fermionic_hamiltonian(symbols, geometry, alpha, h_ref):
         )
     ],
 )
-@pytest.mark.parametrize("fs", [True, False])
-def test_diff_hamiltonian(symbols, geometry, h_ref_data, fs):
+def test_diff_hamiltonian(symbols, geometry, h_ref_data):
     r"""Test that diff_hamiltonian returns the correct Hamiltonian."""
 
     mol = qchem.Molecule(symbols, geometry)
     args = []
-    h = qchem.diff_hamiltonian(mol, fs=fs)(*args)
+    h = qchem.diff_hamiltonian(mol)(*args)
     h_ref = qml.Hamiltonian(h_ref_data[0], h_ref_data[1])
 
     assert np.allclose(np.sort(h.terms()[0]), np.sort(h_ref.terms()[0]))
@@ -387,8 +269,7 @@ def test_diff_hamiltonian_active_space():
     assert not isinstance(h_op, qml.Hamiltonian)
 
 
-@pytest.mark.parametrize("fs", [True, False])
-def test_gradient_expvalH(fs):
+def test_gradient_expvalH():
     r"""Test that the gradient of expval(H) computed with ``qml.grad`` is equal to the value
     obtained with the finite difference method."""
     symbols = ["H", "H"]
@@ -411,7 +292,7 @@ def circuit(*args):
             qml.PauliX(0)
             qml.PauliX(1)
             qml.DoubleExcitation(0.22350048111151138, wires=[0, 1, 2, 3])
-            h_qubit = qchem.diff_hamiltonian(mol, fs=fs)(*args)
+            h_qubit = qchem.diff_hamiltonian(mol)(*args)
             return qml.expval(h_qubit)
 
         return circuit
@@ -464,7 +345,7 @@ def circuit(*args):
                 qml.PauliX(1)
                 qml.DoubleExcitation(0.22350048111151138, wires=[0, 1, 2, 3])
                 enable_new_opmath()
-                h_qubit = qchem.diff_hamiltonian(mol, fs=True)(*args)
+                h_qubit = qchem.diff_hamiltonian(mol)(*args)
                 disable_new_opmath()
                 return qml.expval(h_qubit)
 
diff --git a/tests/qchem/test_observable_hf.py b/tests/qchem/test_observable_hf.py
index 2c0f125c324..41eb15a644a 100644
--- a/tests/qchem/test_observable_hf.py
+++ b/tests/qchem/test_observable_hf.py
@@ -24,189 +24,6 @@
 from pennylane.fermi import from_string
 
 
-@pytest.mark.parametrize(
-    ("core_constant", "integral_one", "integral_two", "f_ref"),
-    [
-        (  # computed with openfermion for H2 (format is modified):
-            # H2 bond length: 1 Angstrom, basis = 'sto-3g', multiplicity = 1, charge = 0
-            # molecule = openfermion.MolecularData(geometry, basis, multiplicity, charge)
-            # run_pyscf(molecule).get_integrals()
-            np.array([0.529177210903]),  # nuclear repulsion 1 / 1.88973 Bohr
-            np.array([[-1.11084418e00, 1.01781501e-16], [7.32122533e-17, -5.89121004e-01]]),
-            np.array(
-                [
-                    [
-                        [[6.26402500e-01, -1.84129592e-16], [-2.14279171e-16, 1.96790583e-01]],
-                        [[-2.14279171e-16, 1.96790583e-01], [6.21706763e-01, -1.84062159e-17]],
-                    ],
-                    [
-                        [[-1.84129592e-16, 6.21706763e-01], [1.96790583e-01, -5.28427412e-17]],
-                        [[1.96790583e-01, -5.28427412e-17], [-1.84062159e-17, 6.53070747e-01]],
-                    ],
-                ]
-            ),
-            # computed with openfermion for H2 (format is modified):
-            # get_fermion_operator(run_pyscf(molecule).get_molecular_hamiltonian())
-            (
-                np.array(
-                    [
-                        0.52917721092,
-                        -1.1108441798837276,
-                        0.31320124976475916,
-                        0.09839529174273519,
-                        0.31320124976475916,
-                        0.09839529174273519,
-                        0.09839529174273519,
-                        0.3108533815598568,
-                        0.09839529174273519,
-                        0.3108533815598568,
-                        0.31320124976475916,
-                        0.09839529174273519,
-                        -1.1108441798837276,
-                        0.31320124976475916,
-                        0.09839529174273519,
-                        0.09839529174273519,
-                        0.3108533815598568,
-                        0.09839529174273519,
-                        0.3108533815598568,
-                        0.3108533815598569,
-                        0.09839529174273519,
-                        0.3108533815598569,
-                        0.09839529174273519,
-                        -0.5891210037060831,
-                        0.09839529174273519,
-                        0.32653537347128725,
-                        0.09839529174273519,
-                        0.32653537347128725,
-                        0.3108533815598569,
-                        0.09839529174273519,
-                        0.3108533815598569,
-                        0.09839529174273519,
-                        0.09839529174273519,
-                        0.32653537347128725,
-                        -0.5891210037060831,
-                        0.09839529174273519,
-                        0.32653537347128725,
-                    ]
-                ),
-                [
-                    [],
-                    [0, 0],
-                    [0, 0, 0, 0],
-                    [0, 0, 2, 2],
-                    [0, 1, 1, 0],
-                    [0, 1, 3, 2],
-                    [0, 2, 0, 2],
-                    [0, 2, 2, 0],
-                    [0, 3, 1, 2],
-                    [0, 3, 3, 0],
-                    [1, 0, 0, 1],
-                    [1, 0, 2, 3],
-                    [1, 1],
-                    [1, 1, 1, 1],
-                    [1, 1, 3, 3],
-                    [1, 2, 0, 3],
-                    [1, 2, 2, 1],
-                    [1, 3, 1, 3],
-                    [1, 3, 3, 1],
-                    [2, 0, 0, 2],
-                    [2, 0, 2, 0],
-                    [2, 1, 1, 2],
-                    [2, 1, 3, 0],
-                    [2, 2],
-                    [2, 2, 0, 0],
-                    [2, 2, 2, 2],
-                    [2, 3, 1, 0],
-                    [2, 3, 3, 2],
-                    [3, 0, 0, 3],
-                    [3, 0, 2, 1],
-                    [3, 1, 1, 3],
-                    [3, 1, 3, 1],
-                    [3, 2, 0, 1],
-                    [3, 2, 2, 3],
-                    [3, 3],
-                    [3, 3, 1, 1],
-                    [3, 3, 3, 3],
-                ],
-            ),
-        ),
-        (
-            np.array([2.869]),
-            np.array(
-                [
-                    [0.95622463, 0.7827277, -0.53222294],
-                    [0.7827277, 1.42895581, 0.23469918],
-                    [-0.53222294, 0.23469918, 0.48381955],
-                ]
-            ),
-            None,
-            # computed with PL-QChem dipole (format is modified)
-            (
-                np.array(
-                    [
-                        2.869,
-                        0.956224634652776,
-                        0.782727697897828,
-                        -0.532222940905614,
-                        0.956224634652776,
-                        0.782727697897828,
-                        -0.532222940905614,
-                        0.782727697897828,
-                        1.42895581236226,
-                        0.234699175620383,
-                        0.782727697897828,
-                        1.42895581236226,
-                        0.234699175620383,
-                        -0.532222940905614,
-                        0.234699175620383,
-                        0.483819552892797,
-                        -0.532222940905614,
-                        0.234699175620383,
-                        0.483819552892797,
-                    ]
-                ),
-                [
-                    [],
-                    [0, 0],
-                    [0, 2],
-                    [0, 4],
-                    [1, 1],
-                    [1, 3],
-                    [1, 5],
-                    [2, 0],
-                    [2, 2],
-                    [2, 4],
-                    [3, 1],
-                    [3, 3],
-                    [3, 5],
-                    [4, 0],
-                    [4, 2],
-                    [4, 4],
-                    [5, 1],
-                    [5, 3],
-                    [5, 5],
-                ],
-            ),
-        ),
-    ],
-)
-def test_fermionic_observable_fs_False(core_constant, integral_one, integral_two, f_ref):
-    r"""Test that fermionic_observable returns the correct fermionic observable."""
-    # TODO: remove this test when supporting tuple output by fermionic_observable is deprecated.
-    f = qchem.fermionic_observable(core_constant, integral_one, integral_two)
-    assert np.allclose(f[0], f_ref[0])  # fermionic coefficients
-    assert f[1] == f_ref[1]  # fermionic operators
-
-
-def test_fermionic_observable_warning():
-    r"""Test that a warning is raised if `fs = False` in fermionic_observable."""
-    # TODO: remove this test when supporting tuple output by fermionic_observable is deprecated.
-    constant = np.array([1.0])
-    integral = np.array([[0.5, -0.8270995], [-0.8270995, 0.5]])
-    with pytest.warns(UserWarning, match="This function will return a fermionic operator"):
-        qchem.fermionic_observable(constant, integral, fs=False)
-
-
 @pytest.mark.parametrize(
     ("core_constant", "integral_one", "integral_two", "f_ref"),
     [
@@ -303,82 +120,12 @@ def test_fermionic_observable_warning():
 )
 def test_fermionic_observable(core_constant, integral_one, integral_two, f_ref):
     r"""Test that fermionic_observable returns the correct fermionic observable."""
-    f = qchem.fermionic_observable(core_constant, integral_one, integral_two, fs=True)
+    f = qchem.fermionic_observable(core_constant, integral_one, integral_two)
 
     assert np.allclose(list(f.values()), list(f_ref.values()))
     assert f.keys() == f_ref.keys()
 
 
-@pytest.mark.parametrize(
-    ("f_observable", "q_observable"),
-    [
-        (
-            (np.array([1.0]), [[0, 0]]),
-            # obtained with openfermion: jordan_wigner(FermionOperator('0^ 0', 1)) and reformatted
-            [[0.5 + 0j, -0.5 + 0j], [qml.Identity(0), qml.PauliZ(0)]],
-        ),
-        (
-            (np.array([1.0, 1.0]), [[0, 0], [0, 0]]),
-            # obtained with openfermion: jordan_wigner(FermionOperator('0^ 0', 1)) and reformatted
-            [[1.0 + 0j, -1.0 + 0j], [qml.Identity(0), qml.PauliZ(0)]],
-        ),
-        (
-            (np.array([1.0]), [[2, 0, 2, 0]]),
-            # obtained with openfermion: jordan_wigner(FermionOperator('0^ 0', 1)) and reformatted
-            [
-                [-0.25 + 0j, 0.25 + 0j, -0.25 + 0j, 0.25 + 0j],
-                [qml.Identity(0), qml.PauliZ(0), qml.PauliZ(0) @ qml.PauliZ(2), qml.PauliZ(2)],
-            ],
-        ),
-        (
-            (np.array([1.0, 1.0]), [[2, 0, 2, 0], [2, 0]]),
-            # obtained with openfermion: jordan_wigner(FermionOperator('0^ 0', 1)) and reformatted
-            [
-                [-0.25 + 0j, 0.25 + 0j, -0.25j, 0.25j, 0.25 + 0j, 0.25 + 0j, -0.25 + 0j, 0.25 + 0j],
-                [
-                    qml.Identity(0),
-                    qml.PauliX(0) @ qml.PauliZ(1) @ qml.PauliX(2),
-                    qml.PauliX(0) @ qml.PauliZ(1) @ qml.PauliY(2),
-                    qml.PauliY(0) @ qml.PauliZ(1) @ qml.PauliX(2),
-                    qml.PauliY(0) @ qml.PauliZ(1) @ qml.PauliY(2),
-                    qml.PauliZ(0),
-                    qml.PauliZ(0) @ qml.PauliZ(2),
-                    qml.PauliZ(2),
-                ],
-            ],
-        ),
-        ((np.array([1.23]), [[]]), [[1.23], [qml.Identity(0)]]),
-    ],
-)
-def test_qubit_observable_tuple_input(f_observable, q_observable):
-    r"""Test that qubit_observable returns the correct operator."""
-    # TODO: remove this test when supporting tuple input by qubit_observable is deprecated.
-    h_as_hamiltonian = qchem.qubit_observable(f_observable)
-    h_ref = qml.Hamiltonian(q_observable[0], q_observable[1])
-
-    enable_new_opmath()
-
-    h_as_op = qchem.qubit_observable(f_observable)
-    h_ref_as_op = pauli_sentence(h_ref).operation(
-        h_ref.wires
-    )  # easy conversion from ham to operation
-
-    disable_new_opmath()
-
-    assert h_as_hamiltonian.compare(h_ref)
-    assert np.allclose(
-        qml.matrix(h_as_op, wire_order=[0, 1, 2]), qml.matrix(h_ref_as_op, wire_order=[0, 1, 2])
-    )
-
-
-def test_qubit_observable_warning():
-    r"""Test that a warning is raised if input is not a fermionic operator in qubit_observable."""
-    # TODO: remove this test when supporting tuple input by qubit_observable is deprecated.
-    f = (np.array([1.0, 1.0]), [[2, 0, 2, 0], [2, 0]])
-    with pytest.warns(UserWarning, match="Tuple input for the qubit_observable function is"):
-        qchem.qubit_observable(f)
-
-
 @pytest.mark.parametrize(
     ("f_observable", "q_observable"),
     [
@@ -441,36 +188,6 @@ def test_qubit_observable(f_observable, q_observable):
     )
 
 
-@pytest.mark.parametrize(
-    ("f_observable", "cut_off"),
-    [
-        (
-            (np.array([0.01]), [[0, 0]]),
-            0.1,
-        ),
-        (
-            (np.array([1.0]), [[0, 0, 1, 1]]),  # should produce the 0 operator
-            0.1,
-        ),
-    ],
-)
-def test_qubit_observable_cutoff_tuple_input(f_observable, cut_off):
-    """Test that qubit_observable returns the correct operator when a cutoff is provided."""
-    # TODO: remove this test when supporting tuple input by qubit_observable is deprecated.
-    h_ref, h_ref_op = (qml.Hamiltonian([], []), qml.s_prod(0, qml.Identity(0)))
-    h_as_hamiltonian = qchem.qubit_observable(f_observable, cutoff=cut_off)
-
-    enable_new_opmath()
-    h_as_op = qchem.qubit_observable(f_observable, cutoff=cut_off)
-    disable_new_opmath()
-
-    assert h_as_hamiltonian.compare(h_ref)
-    assert qml.equal(h_as_op, h_ref_op)
-    assert np.allclose(
-        qml.matrix(h_as_op, wire_order=[0, 1, 2]), qml.matrix(h_ref_op, wire_order=[0, 1, 2])
-    )
-
-
 @pytest.mark.parametrize(
     ("f_observable", "cut_off"),
     [
@@ -498,182 +215,3 @@ def test_qubit_observable_cutoff(f_observable, cut_off):
     assert np.allclose(
         qml.matrix(h_as_op, wire_order=[0, 1, 2]), qml.matrix(h_ref_op, wire_order=[0, 1, 2])
     )
-
-
-@pytest.mark.parametrize(
-    ("f_obs", "q_obs", "notation"),
-    [
-        (
-            [0],
-            # trivial case of a creation operator, 0^ -> (X_0 - iY_0) / 2
-            # reformatted the original openfermion output: (0.5+0j) [] + (-0.5+0j) [Z0]
-            ([(0.5 + 0j), (0.0 - 0.5j)], [qml.PauliX(0), qml.PauliY(0)]),
-            None,
-        ),
-        (
-            [0, 0],
-            # obtained with openfermion using: jordan_wigner(FermionOperator('0^ 0', 1))
-            # reformatted the original openfermion output: (0.5+0j) [] + (-0.5+0j) [Z0]
-            ([(0.5 + 0j), (-0.5 + 0j)], [qml.Identity(0), qml.PauliZ(0)]),
-            None,
-        ),
-        (
-            [3, 0],
-            # obtained with openfermion using: jordan_wigner(FermionOperator('3^ 0', 1))
-            # reformatted the original openfermion output
-            (
-                [(0.25 + 0j), -0.25j, 0.25j, (0.25 + 0j)],
-                [
-                    qml.PauliX(0) @ qml.PauliZ(1) @ qml.PauliZ(2) @ qml.PauliX(3),
-                    qml.PauliX(0) @ qml.PauliZ(1) @ qml.PauliZ(2) @ qml.PauliY(3),
-                    qml.PauliY(0) @ qml.PauliZ(1) @ qml.PauliZ(2) @ qml.PauliX(3),
-                    qml.PauliY(0) @ qml.PauliZ(1) @ qml.PauliZ(2) @ qml.PauliY(3),
-                ],
-            ),
-            None,
-        ),
-        (
-            [1, 4],
-            # obtained with openfermion using: jordan_wigner(FermionOperator('1^ 4', 1))
-            # reformatted the original openfermion output
-            (
-                [(0.25 + 0j), 0.25j, -0.25j, (0.25 + 0j)],
-                [
-                    qml.PauliX(1) @ qml.PauliZ(2) @ qml.PauliZ(3) @ qml.PauliX(4),
-                    qml.PauliX(1) @ qml.PauliZ(2) @ qml.PauliZ(3) @ qml.PauliY(4),
-                    qml.PauliY(1) @ qml.PauliZ(2) @ qml.PauliZ(3) @ qml.PauliX(4),
-                    qml.PauliY(1) @ qml.PauliZ(2) @ qml.PauliZ(3) @ qml.PauliY(4),
-                ],
-            ),
-            None,
-        ),
-        (
-            [1, 1, 1, 1],
-            # obtained with openfermion using: jordan_wigner(FermionOperator('1^ 1^ 1 1', 1))
-            ([0], [qml.Identity(1)]),
-            "physicist",
-        ),
-        (
-            [3, 1, 3, 1],
-            # obtained with openfermion using: jordan_wigner(FermionOperator('3^ 1^ 3 1', 1))
-            # reformatted the original openfermion output
-            (
-                [(-0.25 + 0j), (0.25 + 0j), (-0.25 + 0j), (0.25 + 0j)],
-                [qml.Identity(0), qml.PauliZ(1), qml.PauliZ(1) @ qml.PauliZ(3), qml.PauliZ(3)],
-            ),
-            "physicist",
-        ),
-        (
-            [3, 1, 3, 1],
-            # obtained with openfermion using: jordan_wigner(FermionOperator('3^ 1 3^ 1', 1))
-            ([0], [qml.Identity(1)]),
-            "chemist",
-        ),
-        (
-            [1, 0, 1, 1],
-            # obtained with openfermion using: jordan_wigner(FermionOperator('1^ 0 1^ 1', 1))
-            ([0], [qml.Identity(0)]),
-            "chemist",
-        ),
-        (
-            [1, 1, 0, 0],
-            # obtained with openfermion using: jordan_wigner(FermionOperator('1^ 1 0^ 0', 1))
-            (
-                [(0.25 + 0j), (-0.25 + 0j), (0.25 + 0j), (-0.25 + 0j)],
-                [qml.Identity(0), qml.PauliZ(0), qml.PauliZ(0) @ qml.PauliZ(1), qml.PauliZ(1)],
-            ),
-            "chemist",
-        ),
-        (
-            [5, 5, 5, 5],
-            # obtained with openfermion using: jordan_wigner(FermionOperator('5^ 5 5^ 5', 1))
-            (
-                [(0.5 + 0j), (-0.5 + 0j)],
-                [qml.Identity(0), qml.PauliZ(5)],
-            ),
-            "chemist",
-        ),
-        (
-            [3, 3, 3, 1],
-            # obtained with openfermion using: jordan_wigner(FermionOperator('3^ 3 3^ 1', 1))
-            (
-                [(0.25 + 0j), (-0.25j), (0.25j), (0.25 + 0j)],
-                [
-                    qml.PauliX(1) @ qml.PauliZ(2) @ qml.PauliX(3),
-                    qml.PauliX(1) @ qml.PauliZ(2) @ qml.PauliY(3),
-                    qml.PauliY(1) @ qml.PauliZ(2) @ qml.PauliX(3),
-                    qml.PauliY(1) @ qml.PauliZ(2) @ qml.PauliY(3),
-                ],
-            ),
-            "chemist",
-        ),
-        (
-            [3, 0, 2, 1],
-            # obtained with openfermion using: jordan_wigner(FermionOperator('3^ 0 2^ 1', 1))
-            (
-                [
-                    (-0.0625 + 0j),
-                    0.0625j,
-                    0.0625j,
-                    (0.0625 + 0j),
-                    -0.0625j,
-                    (-0.0625 + 0j),
-                    (-0.0625 + 0j),
-                    0.0625j,
-                    -0.0625j,
-                    (-0.0625 + 0j),
-                    (-0.0625 + 0j),
-                    0.0625j,
-                    (0.0625 + 0j),
-                    -0.0625j,
-                    -0.0625j,
-                    (-0.0625 + 0j),
-                ],
-                [
-                    qml.PauliX(0) @ qml.PauliX(1) @ qml.PauliX(2) @ qml.PauliX(3),
-                    qml.PauliX(0) @ qml.PauliX(1) @ qml.PauliX(2) @ qml.PauliY(3),
-                    qml.PauliX(0) @ qml.PauliX(1) @ qml.PauliY(2) @ qml.PauliX(3),
-                    qml.PauliX(0) @ qml.PauliX(1) @ qml.PauliY(2) @ qml.PauliY(3),
-                    qml.PauliX(0) @ qml.PauliY(1) @ qml.PauliX(2) @ qml.PauliX(3),
-                    qml.PauliX(0) @ qml.PauliY(1) @ qml.PauliX(2) @ qml.PauliY(3),
-                    qml.PauliX(0) @ qml.PauliY(1) @ qml.PauliY(2) @ qml.PauliX(3),
-                    qml.PauliX(0) @ qml.PauliY(1) @ qml.PauliY(2) @ qml.PauliY(3),
-                    qml.PauliY(0) @ qml.PauliX(1) @ qml.PauliX(2) @ qml.PauliX(3),
-                    qml.PauliY(0) @ qml.PauliX(1) @ qml.PauliX(2) @ qml.PauliY(3),
-                    qml.PauliY(0) @ qml.PauliX(1) @ qml.PauliY(2) @ qml.PauliX(3),
-                    qml.PauliY(0) @ qml.PauliX(1) @ qml.PauliY(2) @ qml.PauliY(3),
-                    qml.PauliY(0) @ qml.PauliY(1) @ qml.PauliX(2) @ qml.PauliX(3),
-                    qml.PauliY(0) @ qml.PauliY(1) @ qml.PauliX(2) @ qml.PauliY(3),
-                    qml.PauliY(0) @ qml.PauliY(1) @ qml.PauliY(2) @ qml.PauliX(3),
-                    qml.PauliY(0) @ qml.PauliY(1) @ qml.PauliY(2) @ qml.PauliY(3),
-                ],
-            ),
-            "chemist",
-        ),
-    ],
-)
-def test_jordan_wigner(f_obs, q_obs, notation):
-    r"""Test that jordan_wigner returns the correct operator."""
-    res = qchem.jordan_wigner(f_obs, notation=notation)
-    assert qml.Hamiltonian(res[0], res[1]).compare(qml.Hamiltonian(q_obs[0], q_obs[1]))
-
-
-def test_jordan_wigner_raises_warning():
-    r"""Test that jordan_wigner raisese the expected deprecation warning"""
-    with pytest.warns(UserWarning, match="Use of qml.qchem.jordan_wigner is deprecated"):
-        _ = qchem.jordan_wigner([3, 3, 3, 1], notation="chemist")
-
-
-@pytest.mark.parametrize(
-    ("f_obs", "notation"),
-    [
-        (
-            [1, 1, 1, 1],
-            "random_notation",
-        ),
-    ],
-)
-def test_jordan_wigner_error(f_obs, notation):
-    r"""Test that an error is raised if a wrong/not-supported notation is used."""
-    with pytest.raises(ValueError, match="the only supported notations for the two-body terms are"):
-        qchem.jordan_wigner(f_obs, notation=notation)

From 603074fa0a91961cf672fd287e2136941f053534 Mon Sep 17 00:00:00 2001
From: Jay Soni <jbsoni@uwaterloo.ca>
Date: Fri, 8 Sep 2023 14:03:36 -0400
Subject: [PATCH 038/131] Update `DecompositionUndefinedError` for `Exp` class
 (#4571)

**Context:**
When calling the decomposition for the `Exp` class with `num_steps`
specified, an error was raised suggesting to set `num_steps`. The real
source of the error was caused because the underlying operator was not
unitary. The error message was misleading.

**Description of the Change:**
Update the decomposition function to perform added validation before
raising error.

**Benefits:**
More informative error messages.

**Possible Drawbacks:**
None

**Related GitHub Issues:**
Fixes https://github.com/PennyLaneAI/pennylane/issues/4537
---
 doc/releases/changelog-dev.md |  4 ++++
 pennylane/ops/op_math/exp.py  | 19 ++++++++++++++-----
 tests/ops/op_math/test_exp.py | 28 +++++++++++++++++++++++++---
 3 files changed, 43 insertions(+), 8 deletions(-)

diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index b463e4eee6e..74d7977bc59 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -58,6 +58,9 @@
   process, `DensityMatrixMP`.
   [(#4558)](https://github.com/PennyLaneAI/pennylane/pull/4558)
 
+* `qml.exp` returns a more informative error message when decomposition is unavailable for non-unitary operator.
+  [(#4571)](https://github.com/PennyLaneAI/pennylane/pull/4571)
+
 * The `StateMP` measurement now accepts a wire order (eg. a device wire order). The `process_state`
   method will re-order the given state to go from the inputted wire-order to the process's wire-order.
   If the process's wire-order contains extra wires, it will assume those are in the zero-state.
@@ -197,3 +200,4 @@ Lillian M. A. Frederiksen,
 Romain Moyard,
 Mudit Pandey,
 Matthew Silverman,
+Jay Soni,
diff --git a/pennylane/ops/op_math/exp.py b/pennylane/ops/op_math/exp.py
index ac8ec144747..71156153cfd 100644
--- a/pennylane/ops/op_math/exp.py
+++ b/pennylane/ops/op_math/exp.py
@@ -313,11 +313,20 @@ def _recursive_decomposition(self, base: Operator, coeff: complex):
             # Check if the exponential can be decomposed into a PauliRot gate
             return self._pauli_rot_decomposition(base, coeff)
 
-        raise DecompositionUndefinedError(
-            f"The decomposition of the {self} operator is not defined. "
-            "Please set a value to ``num_steps`` when instantiating the ``Exp`` operator "
-            "if a Suzuki-Trotter decomposition is required."
-        )
+        error_msg = f"The decomposition of the {self} operator is not defined. "
+
+        if not self.num_steps:  # if num_steps was not set
+            error_msg += (
+                "Please set a value to ``num_steps`` when instantiating the ``Exp`` operator "
+                "if a Suzuki-Trotter decomposition is required. "
+            )
+
+        if math.real(self.coeff) != 0 and self.base.is_hermitian:
+            error_msg += (
+                "Decomposition is not defined for real coefficients of hermitian operators."
+            )
+
+        raise DecompositionUndefinedError(error_msg)
 
     @staticmethod
     def _pauli_rot_decomposition(base: Operator, coeff: complex):
diff --git a/tests/ops/op_math/test_exp.py b/tests/ops/op_math/test_exp.py
index 981ab8221c3..2e2848b8987 100644
--- a/tests/ops/op_math/test_exp.py
+++ b/tests/ops/op_math/test_exp.py
@@ -13,6 +13,7 @@
 # limitations under the License.
 """Unit tests for the ``Exp`` class"""
 import copy
+import re
 
 import pytest
 
@@ -400,19 +401,28 @@ def test_non_imag_no_decomposition(self, coeff):
         """Tests that the decomposition doesn't exist if the coefficient has a real component."""
         op = Exp(qml.PauliX(0), coeff)
         assert not op.has_decomposition
-        with pytest.raises(DecompositionUndefinedError):
+        with pytest.raises(
+            DecompositionUndefinedError,
+            match="Decomposition is not defined for real coefficients of hermitian operators.",
+        ):
             op.decomposition()
 
     def test_non_pauli_word_base_no_decomposition(self):
         """Tests that the decomposition doesn't exist if the base is not a pauli word."""
         op = Exp(qml.S(0), -0.5j, num_steps=100)
         assert not op.has_decomposition
-        with pytest.raises(DecompositionUndefinedError):
+        with pytest.raises(
+            DecompositionUndefinedError,
+            match=re.escape(f"The decomposition of the {op} operator is not defined. "),
+        ):
             op.decomposition()
 
         op = Exp(2 * qml.S(0) + qml.PauliZ(1), -0.5j, num_steps=100)
         assert not op.has_decomposition
-        with pytest.raises(DecompositionUndefinedError):
+        with pytest.raises(
+            DecompositionUndefinedError,
+            match=re.escape(f"The decomposition of the {op} operator is not defined. "),
+        ):
             op.decomposition()
 
     def test_nontensor_tensor_no_decomposition(self):
@@ -568,6 +578,18 @@ def test_trotter_decomposition_raises_error(self):
         ):
             op.decomposition()
 
+    def test_real_coeff_and_none_num_steps_error(self):
+        """Test that the decomposition raises an error if ``num_steps`` is None and
+        the coefficient has non-zero real part"""
+        op = qml.exp(qml.sum(qml.PauliX(0), qml.PauliY(1)), 1.23 + 0.5j)
+        msg = (
+            "Please set a value to ``num_steps`` when instantiating the ``Exp`` operator "
+            "if a Suzuki-Trotter decomposition is required. "
+            "Decomposition is not defined for real coefficients of hermitian operators."
+        )
+        with pytest.raises(DecompositionUndefinedError, match=msg):
+            op.decomposition()
+
     @pytest.mark.parametrize(
         "coeff, hamiltonian",
         [

From 8816391c4b6f0330e2ab92a895f07a299f8775bc Mon Sep 17 00:00:00 2001
From: Vincent Michaud-Rioux <vincentm@nanoacademic.com>
Date: Fri, 8 Sep 2023 15:03:14 -0400
Subject: [PATCH 039/131] Fix expand_tape_state_prep. (#4564)

### 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:**

[43708](https://app.shortcut.com/xanaduai/story/43708/lightning-ecosystem-devices-allow-stateprepbase-ops-to-be-used-mid-circuit)
demands mid-circuit `StatePrep` support in Lightning. This PR fixes a
small bug in `expand_tape_state_prep`.

**Description of the Change:**
`first_op`-bugfix

**Benefits:**

**Possible Drawbacks:**

**Related GitHub Issues:**

---------

Co-authored-by: Matthew Silverman <matthews@xanadu.ai>
Co-authored-by: Jay Soni <jbsoni@uwaterloo.ca>
---
 doc/releases/changelog-dev.md | 4 ++++
 pennylane/tape/tape.py        | 2 +-
 tests/tape/test_tape.py       | 5 ++++-
 3 files changed, 9 insertions(+), 2 deletions(-)

diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 74d7977bc59..ccff4f8b45d 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -185,6 +185,9 @@
 
 <h3>Bug fixes 🐛</h3>
 
+* Fix `skip_first` option in `expand_tape_state_prep`.
+  [(#4564)](https://github.com/PennyLaneAI/pennylane/pull/4564)
+
 * `convert_to_numpy_parameters` now uses `qml.ops.functions.bind_new_parameters`. This reinitializes the operation and
   makes sure everything references the new numpy parameters.
 
@@ -197,6 +200,7 @@ This release contains contributions from (in alphabetical order):
 
 Soran Jahangiri,
 Lillian M. A. Frederiksen,
+Vincent Michaud-Rioux,
 Romain Moyard,
 Mudit Pandey,
 Matthew Silverman,
diff --git a/pennylane/tape/tape.py b/pennylane/tape/tape.py
index 226776bf4e0..d8e9cd4fbd8 100644
--- a/pennylane/tape/tape.py
+++ b/pennylane/tape/tape.py
@@ -263,7 +263,7 @@ def expand_tape_state_prep(tape, skip_first=True):
     first_op = tape.operations[0]
     new_ops = (
         [first_op]
-        if isinstance(first_op, StatePrepBase) and skip_first
+        if not isinstance(first_op, StatePrepBase) or skip_first
         else first_op.decomposition()
     )
 
diff --git a/tests/tape/test_tape.py b/tests/tape/test_tape.py
index ba535958095..bf873447dd6 100644
--- a/tests/tape/test_tape.py
+++ b/tests/tape/test_tape.py
@@ -1014,10 +1014,12 @@ def test_depth_expansion(self):
             [
                 qml.BasisState([1, 0], wires=[0, 1]),
                 qml.StatePrep([0, 1, 0, 0], wires=[0, 1]),
+                qml.PauliZ(0),
             ],
             [
                 qml.BasisStatePreparation([1, 0], wires=[0, 1]),
                 qml.MottonenStatePreparation([0, 1, 0, 0], wires=[0, 1]),
+                qml.PauliZ(0),
             ],
         ),
     )
@@ -1026,12 +1028,13 @@ def test_expansion_state_prep(self, skip_first, op, decomp):
         expanding other operations in the tape.
         """
         ops = [
+            op,
             qml.PauliZ(wires=0),
             qml.Rot(0.1, 0.2, 0.3, wires=0),
             qml.BasisState([0], wires=1),
             qml.StatePrep([0, 1], wires=0),
         ]
-        tape = QuantumTape(ops=ops, measurements=[], prep=[op])
+        tape = QuantumTape(ops=ops, measurements=[])
         new_tape = expand_tape_state_prep(tape, skip_first=skip_first)
 
         true_decomposition = []

From 2f48b93103bef83025b069d35154cfcc2adb1fe8 Mon Sep 17 00:00:00 2001
From: Mudit Pandey <mudit.pandey@xanadu.ai>
Date: Fri, 8 Sep 2023 16:35:08 -0400
Subject: [PATCH 040/131] Added counts legacy tests

---
 .../measurements/legacy/test_counts_legacy.py | 34 +++++++++++++++++++
 1 file changed, 34 insertions(+)

diff --git a/tests/measurements/legacy/test_counts_legacy.py b/tests/measurements/legacy/test_counts_legacy.py
index 58b0a95966d..22af6b292f3 100644
--- a/tests/measurements/legacy/test_counts_legacy.py
+++ b/tests/measurements/legacy/test_counts_legacy.py
@@ -428,6 +428,40 @@ def circuit():
 
         custom_measurement_process(dev, spy)
 
+    def test_counts_shape_single_measurement_value(self):
+        """Test that the counts output is correct for single mid-circuit measurement
+        values."""
+        shots = 1000
+        samples = np.random.choice([0, 1], size=(shots, 2)).astype(np.int64)
+        mv = qml.measure(0)
+
+        result = qml.counts(mv).process_samples(samples, wire_order=[0])
+
+        assert len(result) == 2
+        assert set(result.keys()) == {"0", "1"}
+        assert result["0"] == np.count_nonzero(samples[:, 0] == 0)
+        assert result["1"] == np.count_nonzero(samples[:, 0] == 1)
+
+    def test_counts_all_outcomes_measurement_value(self):
+        """Test that the counts output is correct when all_outcomes is passed
+        for mid-circuit measurement values."""
+        shots = 1000
+        samples = np.zeros((shots, 2)).astype(np.int64)
+        mv = qml.measure(0)
+
+        result1 = qml.counts(mv, all_outcomes=False).process_samples(samples, wire_order=[0])
+
+        assert len(result1) == 1
+        assert set(result1.keys()) == {"0"}
+        assert result1["0"] == shots
+
+        result2 = qml.counts(mv, all_outcomes=True).process_samples(samples, wire_order=[0])
+
+        assert len(result2) == 2
+        assert set(result2.keys()) == {"0", "1"}
+        assert result2["0"] == shots
+        assert result2["1"] == 0
+
     def test_all_outcomes_kwarg_no_observable_no_wires(self, mocker):
         """Test that the dictionary keys are *all* the possible combinations
         of basis states for the device, including 0 count values, if no wire

From 929b463a51ffed85f930fe382bed20a722cc9084 Mon Sep 17 00:00:00 2001
From: Mudit Pandey <mudit.pandey@xanadu.ai>
Date: Fri, 8 Sep 2023 17:07:16 -0400
Subject: [PATCH 041/131] Fixing sample legacy tests

---
 tests/measurements/legacy/test_sample_legacy.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tests/measurements/legacy/test_sample_legacy.py b/tests/measurements/legacy/test_sample_legacy.py
index 46fbf381a8a..8c92be31930 100644
--- a/tests/measurements/legacy/test_sample_legacy.py
+++ b/tests/measurements/legacy/test_sample_legacy.py
@@ -16,7 +16,7 @@
 import pytest
 
 import pennylane as qml
-from pennylane.measurements import MeasurementShapeError, Sample, Shots
+from pennylane.measurements import MeasurementShapeError, Sample, Shots, MeasurementValue
 from pennylane.operation import Operator
 
 # pylint: disable=protected-access, no-member
@@ -31,7 +31,7 @@ def custom_measurement_process(device, spy):
     for call_args in call_args_list:
         meas = call_args.args[1]
         shot_range, bin_size = (call_args.kwargs["shot_range"], call_args.kwargs["bin_size"])
-        if isinstance(meas, Operator):
+        if isinstance(meas, (Operator, MeasurementValue)):
             meas = qml.sample(op=meas)
         assert qml.math.allequal(
             device.sample(call_args.args[1], **call_args.kwargs),

From a6ce9aacefd7b1ac09626af7e9f49afae7b85cf7 Mon Sep 17 00:00:00 2001
From: Mudit Pandey <mudit.pandey@xanadu.ai>
Date: Mon, 11 Sep 2023 11:22:20 -0400
Subject: [PATCH 042/131] Reverted Hamiltonian.compare; linting

---
 pennylane/ops/qubit/hamiltonian.py | 13 +++++++------
 tests/measurements/test_sample.py  |  4 ++--
 2 files changed, 9 insertions(+), 8 deletions(-)

diff --git a/pennylane/ops/qubit/hamiltonian.py b/pennylane/ops/qubit/hamiltonian.py
index 8991b31497c..1a2ff7ac6b2 100644
--- a/pennylane/ops/qubit/hamiltonian.py
+++ b/pennylane/ops/qubit/hamiltonian.py
@@ -613,15 +613,16 @@ def compare(self, other):
         >>> ob1.compare(ob2)
         False
         """
-        # pylint: disable=protected-access
         if isinstance(other, Hamiltonian):
-            H1 = copy(self).simplify()
-            H2 = copy(other).simplify()
-            return H1._obs_data() == H2._obs_data()
+            self.simplify()
+            other.simplify()
+            return self._obs_data() == other._obs_data()  # pylint: disable=protected-access
 
         if isinstance(other, (Tensor, Observable)):
-            H1 = copy(self).simplify()
-            return H1._obs_data() == {(1, frozenset(other._obs_data()))}
+            self.simplify()
+            return self._obs_data() == {
+                (1, frozenset(other._obs_data()))  # pylint: disable=protected-access
+            }
 
         raise ValueError("Can only compare a Hamiltonian, and a Hamiltonian/Observable/Tensor.")
 
diff --git a/tests/measurements/test_sample.py b/tests/measurements/test_sample.py
index b5ba8e8e83b..cdb6b7d4d77 100644
--- a/tests/measurements/test_sample.py
+++ b/tests/measurements/test_sample.py
@@ -16,10 +16,10 @@
 import pytest
 
 import pennylane as qml
-from pennylane.measurements import MeasurementShapeError, Sample, Shots, MeasurementValue
+from pennylane.measurements import MeasurementShapeError, Sample, Shots
 from pennylane.operation import EigvalsUndefinedError, Operator
 
-# pylint: disable=protected-access, no-member
+# pylint: disable=protected-access, no-member, too-many-public-methods
 
 
 class TestSample:

From de5f3b13c17e2ee3bc3ef5eb68f5ddbb4b528fcf Mon Sep 17 00:00:00 2001
From: Mudit Pandey <mudit.pandey@xanadu.ai>
Date: Mon, 11 Sep 2023 11:30:16 -0400
Subject: [PATCH 043/131] Removed legacy tests

---
 .../measurements/legacy/test_counts_legacy.py | 34 -------------------
 .../measurements/legacy/test_expval_legacy.py | 24 -------------
 .../measurements/legacy/test_probs_legacy.py  | 31 -----------------
 .../measurements/legacy/test_sample_legacy.py | 25 --------------
 tests/measurements/legacy/test_var_legacy.py  | 25 --------------
 5 files changed, 139 deletions(-)

diff --git a/tests/measurements/legacy/test_counts_legacy.py b/tests/measurements/legacy/test_counts_legacy.py
index 22af6b292f3..58b0a95966d 100644
--- a/tests/measurements/legacy/test_counts_legacy.py
+++ b/tests/measurements/legacy/test_counts_legacy.py
@@ -428,40 +428,6 @@ def circuit():
 
         custom_measurement_process(dev, spy)
 
-    def test_counts_shape_single_measurement_value(self):
-        """Test that the counts output is correct for single mid-circuit measurement
-        values."""
-        shots = 1000
-        samples = np.random.choice([0, 1], size=(shots, 2)).astype(np.int64)
-        mv = qml.measure(0)
-
-        result = qml.counts(mv).process_samples(samples, wire_order=[0])
-
-        assert len(result) == 2
-        assert set(result.keys()) == {"0", "1"}
-        assert result["0"] == np.count_nonzero(samples[:, 0] == 0)
-        assert result["1"] == np.count_nonzero(samples[:, 0] == 1)
-
-    def test_counts_all_outcomes_measurement_value(self):
-        """Test that the counts output is correct when all_outcomes is passed
-        for mid-circuit measurement values."""
-        shots = 1000
-        samples = np.zeros((shots, 2)).astype(np.int64)
-        mv = qml.measure(0)
-
-        result1 = qml.counts(mv, all_outcomes=False).process_samples(samples, wire_order=[0])
-
-        assert len(result1) == 1
-        assert set(result1.keys()) == {"0"}
-        assert result1["0"] == shots
-
-        result2 = qml.counts(mv, all_outcomes=True).process_samples(samples, wire_order=[0])
-
-        assert len(result2) == 2
-        assert set(result2.keys()) == {"0", "1"}
-        assert result2["0"] == shots
-        assert result2["1"] == 0
-
     def test_all_outcomes_kwarg_no_observable_no_wires(self, mocker):
         """Test that the dictionary keys are *all* the possible combinations
         of basis states for the device, including 0 count values, if no wire
diff --git a/tests/measurements/legacy/test_expval_legacy.py b/tests/measurements/legacy/test_expval_legacy.py
index 436b6b2e121..68fcecc8368 100644
--- a/tests/measurements/legacy/test_expval_legacy.py
+++ b/tests/measurements/legacy/test_expval_legacy.py
@@ -114,30 +114,6 @@ def circuit():
 
         custom_measurement_process(new_dev, spy)
 
-    @pytest.mark.parametrize("shots", [None, 10000, [10000, 10000]])
-    @pytest.mark.parametrize("phi", np.arange(0, 2 * np.pi, np.pi / 3))
-    def test_observable_is_measurement_value(
-        self, shots, phi, mocker, tol, tol_stochastic
-    ):  # pylint: disable=too-many-arguments
-        """Test that expectation values for mid-circuit measurement values
-        are correct for a single measurement value."""
-        dev = qml.device("default.qubit", wires=2, shots=shots)
-
-        @qml.qnode(dev)
-        def circuit(phi):
-            qml.RX(phi, 0)
-            m0 = qml.measure(0)
-            return qml.expval(m0)
-
-        new_dev = circuit.device
-        spy = mocker.spy(qml.QubitDevice, "expval")
-
-        res = circuit(phi)
-
-        atol = tol if shots is None else tol_stochastic
-        assert np.allclose(np.array(res), np.sin(phi / 2) ** 2, atol=atol, rtol=0)
-        custom_measurement_process(new_dev, spy)
-
     @pytest.mark.parametrize(
         "obs",
         [qml.PauliZ(0), qml.Hermitian(np.diag([1, 2]), 0), qml.Hermitian(np.diag([1.0, 2.0]), 0)],
diff --git a/tests/measurements/legacy/test_probs_legacy.py b/tests/measurements/legacy/test_probs_legacy.py
index 2e7605bc151..ad08007fcc2 100644
--- a/tests/measurements/legacy/test_probs_legacy.py
+++ b/tests/measurements/legacy/test_probs_legacy.py
@@ -210,37 +210,6 @@ def circuit():
 
         custom_measurement_process(dev, spy)
 
-    @pytest.mark.parametrize("shots", [None, 10000, [10000, 10000]])
-    @pytest.mark.parametrize("phi", np.arange(0, 2 * np.pi, np.pi / 3))
-    def test_observable_is_measurement_value(
-        self, shots, phi, mocker, tol, tol_stochastic
-    ):  # pylint: disable=too-many-arguments
-        """Test that expectation values for mid-circuit measurement values
-        are correct for a single measurement value."""
-        dev = qml.device("default.qubit", wires=2, shots=shots)
-
-        @qml.qnode(dev)
-        def circuit(phi):
-            qml.RX(phi, 0)
-            m0 = qml.measure(0)
-            return qml.probs(op=m0)
-
-        new_dev = circuit.device
-        spy = mocker.spy(qml.QubitDevice, "probability")
-
-        res = circuit(phi)
-
-        atol = tol if shots is None else tol_stochastic
-        expected = np.array([np.cos(phi / 2) ** 2, np.sin(phi / 2) ** 2])
-
-        if not isinstance(shots, list):
-            assert np.allclose(np.array(res), expected, atol=atol, rtol=0)
-        else:
-            for r in res:  # pylint: disable=not-an-iterable
-                assert np.allclose(r, expected, atol=atol, rtol=0)
-
-        custom_measurement_process(new_dev, spy)
-
     @pytest.mark.parametrize("shots", [None, 100])
     def test_batch_size(self, mocker, shots):
         """Test the probability is correct for a batched input."""
diff --git a/tests/measurements/legacy/test_sample_legacy.py b/tests/measurements/legacy/test_sample_legacy.py
index 8c92be31930..d03a697d7b2 100644
--- a/tests/measurements/legacy/test_sample_legacy.py
+++ b/tests/measurements/legacy/test_sample_legacy.py
@@ -212,31 +212,6 @@ def circuit():
         ):
             _ = circuit()
 
-    @pytest.mark.parametrize("shots", [5, [5, 5]])
-    @pytest.mark.parametrize("phi", np.arange(0, 2 * np.pi, np.pi / 2))
-    def test_observable_is_measurement_value(self, shots, phi, mocker):
-        """Test that expectation values for mid-circuit measurement values
-        are correct for a single measurement value."""
-        dev = qml.device("default.qubit", wires=2, shots=shots)
-
-        @qml.qnode(dev)
-        def circuit(phi):
-            qml.RX(phi, 0)
-            m0 = qml.measure(0)
-            return qml.sample(m0)
-
-        new_dev = circuit.device
-        spy = mocker.spy(qml.QubitDevice, "sample")
-
-        res = circuit(phi)
-
-        if isinstance(shots, list):
-            assert len(res) == len(shots)
-            assert all(r.shape == (s,) for r, s in zip(res, shots))
-        else:
-            assert res.shape == (shots,)
-        custom_measurement_process(new_dev, spy)
-
     def test_providing_no_observable_and_no_wires(self, mocker):
         """Test that we can provide no observable and no wires to sample function"""
         dev = qml.device("default.qubit.legacy", wires=2, shots=1000)
diff --git a/tests/measurements/legacy/test_var_legacy.py b/tests/measurements/legacy/test_var_legacy.py
index 6580c3f538d..5d5ed0ac523 100644
--- a/tests/measurements/legacy/test_var_legacy.py
+++ b/tests/measurements/legacy/test_var_legacy.py
@@ -107,31 +107,6 @@ def circuit():
 
         custom_measurement_process(dev, spy)
 
-    @pytest.mark.parametrize("shots", [None, 10000, [10000, 10000]])
-    @pytest.mark.parametrize("phi", np.arange(0, 2 * np.pi, np.pi / 3))
-    def test_observable_is_measurement_value(
-        self, shots, phi, mocker, tol, tol_stochastic
-    ):  # pylint: disable=too-many-arguments
-        """Test that variances for mid-circuit measurement values
-        are correct for a single measurement value."""
-        dev = qml.device("default.qubit", wires=2, shots=shots)
-
-        @qml.qnode(dev)
-        def circuit(phi):
-            qml.RX(phi, 0)
-            m0 = qml.measure(0)
-            return qml.var(m0)
-
-        new_dev = circuit.device
-        spy = mocker.spy(qml.QubitDevice, "var")
-
-        res = circuit(phi)
-
-        atol = tol if shots is None else tol_stochastic
-        expected = np.sin(phi / 2) ** 2 - np.sin(phi / 2) ** 4
-        assert np.allclose(np.array(res), expected, atol=atol, rtol=0)
-        custom_measurement_process(new_dev, spy)
-
     @pytest.mark.parametrize(
         "obs",
         [qml.PauliZ(0), qml.Hermitian(np.diag([1, 2]), 0), qml.Hermitian(np.diag([1.0, 2.0]), 0)],

From ffca3180d2e114ff1f1a49737bd43ba224fa8e59 Mon Sep 17 00:00:00 2001
From: Mudit Pandey <mudit.pandey@xanadu.ai>
Date: Mon, 11 Sep 2023 11:34:00 -0400
Subject: [PATCH 044/131] Reverted sample legacy tests

---
 tests/measurements/legacy/test_sample_legacy.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tests/measurements/legacy/test_sample_legacy.py b/tests/measurements/legacy/test_sample_legacy.py
index d03a697d7b2..ee8bb3bca01 100644
--- a/tests/measurements/legacy/test_sample_legacy.py
+++ b/tests/measurements/legacy/test_sample_legacy.py
@@ -16,7 +16,7 @@
 import pytest
 
 import pennylane as qml
-from pennylane.measurements import MeasurementShapeError, Sample, Shots, MeasurementValue
+from pennylane.measurements import MeasurementShapeError, Sample, Shots
 from pennylane.operation import Operator
 
 # pylint: disable=protected-access, no-member
@@ -31,7 +31,7 @@ def custom_measurement_process(device, spy):
     for call_args in call_args_list:
         meas = call_args.args[1]
         shot_range, bin_size = (call_args.kwargs["shot_range"], call_args.kwargs["bin_size"])
-        if isinstance(meas, (Operator, MeasurementValue)):
+        if isinstance(meas, Operator):
             meas = qml.sample(op=meas)
         assert qml.math.allequal(
             device.sample(call_args.args[1], **call_args.kwargs),

From 8b22ca0a230c8c8dd3cb618d201f40e4e268dd74 Mon Sep 17 00:00:00 2001
From: Mudit Pandey <mudit.pandey@xanadu.ai>
Date: Mon, 11 Sep 2023 11:37:46 -0400
Subject: [PATCH 045/131] Update test to measure 3 wires

---
 tests/transforms/test_defer_measurements.py | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/tests/transforms/test_defer_measurements.py b/tests/transforms/test_defer_measurements.py
index efd4d3f17aa..7e95f564adc 100644
--- a/tests/transforms/test_defer_measurements.py
+++ b/tests/transforms/test_defer_measurements.py
@@ -1057,7 +1057,7 @@ class TestQubitReuseAndReset:
     def test_new_wire_for_multiple_measurements(self):
         """Test that a new wire is added if there are multiple mid-circuit measurements
         on the same wire."""
-        dev = qml.device("default.qubit", wires=3)
+        dev = qml.device("default.qubit", wires=4)
 
         @qml.qnode(dev)
         def circ(x, y):
@@ -1065,6 +1065,8 @@ def circ(x, y):
             qml.measure(0)
             qml.RY(y, 1)
             qml.measure(0)
+            qml.RZ(x + y, 1)
+            qml.measure(0)
             return qml.expval(qml.PauliZ(1))
 
         _ = circ(1.0, 2.0)
@@ -1073,9 +1075,11 @@ def circ(x, y):
             qml.RX(1.0, 0),
             qml.CNOT([0, 2]),
             qml.RY(2.0, 1),
+            qml.CNOT([0, 3]),
+            qml.RZ(3.0, 1),
         ]
 
-        assert len(circ.qtape.operations) == 3
+        assert len(circ.qtape.operations) == 5
         for op, exp in zip(circ.qtape.operations, expected):
             assert qml.equal(op, exp)
 

From 5bf50617b1bc6d56069b4d12ebbcc92ca4157af6 Mon Sep 17 00:00:00 2001
From: Mudit Pandey <mudit.pandey@xanadu.ai>
Date: Tue, 12 Sep 2023 16:19:13 -0400
Subject: [PATCH 046/131] Updated MP eigvals

---
 pennylane/measurements/counts.py      | 11 +++--------
 pennylane/measurements/expval.py      |  6 +-----
 pennylane/measurements/mid_measure.py |  2 ++
 pennylane/measurements/sample.py      | 23 ++++++++---------------
 pennylane/measurements/var.py         |  6 +-----
 5 files changed, 15 insertions(+), 33 deletions(-)

diff --git a/pennylane/measurements/counts.py b/pennylane/measurements/counts.py
index 83e17a43b4b..e317785fed7 100644
--- a/pennylane/measurements/counts.py
+++ b/pennylane/measurements/counts.py
@@ -214,14 +214,9 @@ def process_samples(
         shot_range: Tuple[int] = None,
         bin_size: int = None,
     ):
-        if not isinstance(self.obs, MeasurementValue):
-            samples = qml.sample(op=self.obs, wires=self._wires).process_samples(
-                samples, wire_order, shot_range, bin_size
-            )
-        else:
-            samples = qml.sample(wires=self.obs.wires).process_samples(
-                samples, wire_order, shot_range, bin_size
-            )
+        samples = qml.sample(op=self.obs, wires=self._wires).process_samples(
+            samples, wire_order, shot_range, bin_size
+        )
 
         if bin_size is None:
             return self._samples_to_counts(samples)
diff --git a/pennylane/measurements/expval.py b/pennylane/measurements/expval.py
index 015b8c46c2d..b9c53195bda 100644
--- a/pennylane/measurements/expval.py
+++ b/pennylane/measurements/expval.py
@@ -130,11 +130,7 @@ def process_state(self, state: Sequence[complex], wire_order: Wires):
             idx = int("".join(str(i) for i in self.obs.parameters[0]), 2)
             probs = qml.probs(wires=self.wires).process_state(state=state, wire_order=wire_order)
             return probs[idx]
-        eigvals = (
-            qml.math.asarray(self.obs.eigvals(), dtype="float64")
-            if not isinstance(self.obs, MeasurementValue)
-            else qml.math.asarray(qml.math.arange(0, 2 ** len(self.wires)), dtype="float64")
-        )
+        eigvals = qml.math.asarray(self.eigvals(), dtype="float64")
         # we use ``self.wires`` instead of ``self.obs`` because the observable was
         # already applied to the state
         prob = qml.probs(wires=self.wires).process_state(state=state, wire_order=wire_order)
diff --git a/pennylane/measurements/mid_measure.py b/pennylane/measurements/mid_measure.py
index 9602a1501a1..6bbb90a14b8 100644
--- a/pennylane/measurements/mid_measure.py
+++ b/pennylane/measurements/mid_measure.py
@@ -164,6 +164,8 @@ class MeasurementValue(Generic[T]):
         processing_fn (callable): A lazily transformation applied to the measurement values.
     """
 
+    name = "MeasurementValue"
+
     def __init__(self, measurements, processing_fn):
         self.measurements = measurements
         self.processing_fn = processing_fn
diff --git a/pennylane/measurements/sample.py b/pennylane/measurements/sample.py
index 28b00626249..e0db9ed186d 100644
--- a/pennylane/measurements/sample.py
+++ b/pennylane/measurements/sample.py
@@ -193,11 +193,7 @@ def process_samples(
     ):
         wire_map = dict(zip(wire_order, range(len(wire_order))))
         mapped_wires = [wire_map[w] for w in self.wires]
-        name = (
-            self.obs.name
-            if self.obs is not None and not isinstance(self.obs, MeasurementValue)
-            else None
-        )
+        name = self.obs.name if self.obs is not None else None
         # Select the samples from samples that correspond to ``shot_range`` if provided
         if shot_range is not None:
             # Indexing corresponds to: (potential broadcasting, shots, wires). Note that the last
@@ -226,15 +222,12 @@ def process_samples(
             powers_of_two = 2 ** qml.math.arange(num_wires)[::-1]
             indices = samples @ powers_of_two
             indices = qml.math.array(indices)  # Add np.array here for Jax support.
-            if isinstance(self.obs, MeasurementValue):
-                samples = qml.math.arange(0, 2 ** len(self.obs.wires), 1)[indices]
-            else:
-                try:
-                    samples = self.obs.eigvals()[indices]
-                except qml.operation.EigvalsUndefinedError as e:
-                    # if observable has no info on eigenvalues, we cannot return this measurement
-                    raise qml.operation.EigvalsUndefinedError(
-                        f"Cannot compute samples of {self.obs.name}."
-                    ) from e
+            try:
+                samples = self.eigvals()[indices]
+            except qml.operation.EigvalsUndefinedError as e:
+                # if observable has no info on eigenvalues, we cannot return this measurement
+                raise qml.operation.EigvalsUndefinedError(
+                    f"Cannot compute samples of {self.obs.name}."
+                ) from e
 
         return samples if bin_size is None else samples.reshape((bin_size, -1))
diff --git a/pennylane/measurements/var.py b/pennylane/measurements/var.py
index 818f241a91d..fa526dceb0f 100644
--- a/pennylane/measurements/var.py
+++ b/pennylane/measurements/var.py
@@ -135,11 +135,7 @@ def process_state(self, state: Sequence[complex], wire_order: Wires):
             probs = qml.probs(wires=self.wires).process_state(state=state, wire_order=wire_order)
             return probs[idx] - probs[idx] ** 2
 
-        eigvals = (
-            qml.math.asarray(self.obs.eigvals(), dtype="float64")
-            if not isinstance(self.obs, MeasurementValue)
-            else qml.math.asarray(qml.math.arange(0, 2 ** len(self.wires)), dtype="float64")
-        )
+        eigvals = qml.math.asarray(self.eigvals(), dtype="float64")
 
         # we use ``wires`` instead of ``op`` because the observable was
         # already applied to the state

From a7d0ebbf4ed60048b4022a5299e6b75e454ae54d Mon Sep 17 00:00:00 2001
From: Mudit Pandey <mudit.pandey@xanadu.ai>
Date: Tue, 12 Sep 2023 16:21:59 -0400
Subject: [PATCH 047/131] reformatting

---
 pennylane/measurements/expval.py | 5 +----
 pennylane/measurements/var.py    | 5 +----
 2 files changed, 2 insertions(+), 8 deletions(-)

diff --git a/pennylane/measurements/expval.py b/pennylane/measurements/expval.py
index b9c53195bda..d9e60d3c232 100644
--- a/pennylane/measurements/expval.py
+++ b/pennylane/measurements/expval.py
@@ -112,10 +112,7 @@ def process_samples(
 
         # estimate the ev
         samples = qml.sample(op=self.obs).process_samples(
-            samples=samples,
-            wire_order=wire_order,
-            shot_range=shot_range,
-            bin_size=bin_size,
+            samples=samples, wire_order=wire_order, shot_range=shot_range, bin_size=bin_size
         )
 
         # With broadcasting, we want to take the mean over axis 1, which is the -1st/-2nd with/
diff --git a/pennylane/measurements/var.py b/pennylane/measurements/var.py
index fa526dceb0f..f3ab27cd4f4 100644
--- a/pennylane/measurements/var.py
+++ b/pennylane/measurements/var.py
@@ -114,10 +114,7 @@ def process_samples(
 
         # estimate the variance
         samples = qml.sample(op=self.obs).process_samples(
-            samples=samples,
-            wire_order=wire_order,
-            shot_range=shot_range,
-            bin_size=bin_size,
+            samples=samples, wire_order=wire_order, shot_range=shot_range, bin_size=bin_size
         )
 
         # With broadcasting, we want to take the variance over axis 1, which is the -1st/-2nd with/

From 031b2d2f5211850bc43cefc62e9ac7c726b65903 Mon Sep 17 00:00:00 2001
From: Mudit Pandey <mudit.pandey@xanadu.ai>
Date: Tue, 12 Sep 2023 17:07:05 -0400
Subject: [PATCH 048/131] Reverted counts changes

---
 pennylane/measurements/counts.py | 11 ++++++++---
 1 file changed, 8 insertions(+), 3 deletions(-)

diff --git a/pennylane/measurements/counts.py b/pennylane/measurements/counts.py
index e317785fed7..83e17a43b4b 100644
--- a/pennylane/measurements/counts.py
+++ b/pennylane/measurements/counts.py
@@ -214,9 +214,14 @@ def process_samples(
         shot_range: Tuple[int] = None,
         bin_size: int = None,
     ):
-        samples = qml.sample(op=self.obs, wires=self._wires).process_samples(
-            samples, wire_order, shot_range, bin_size
-        )
+        if not isinstance(self.obs, MeasurementValue):
+            samples = qml.sample(op=self.obs, wires=self._wires).process_samples(
+                samples, wire_order, shot_range, bin_size
+            )
+        else:
+            samples = qml.sample(wires=self.obs.wires).process_samples(
+                samples, wire_order, shot_range, bin_size
+            )
 
         if bin_size is None:
             return self._samples_to_counts(samples)

From 74e896e4035f02a355ab04403c3fc292c836b3bd Mon Sep 17 00:00:00 2001
From: Mudit Pandey <mudit.pandey@xanadu.ai>
Date: Wed, 13 Sep 2023 09:46:14 -0400
Subject: [PATCH 049/131] Removed error supression

---
 pennylane/measurements/measurements.py | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/pennylane/measurements/measurements.py b/pennylane/measurements/measurements.py
index 018a9d91138..5096bee9739 100644
--- a/pennylane/measurements/measurements.py
+++ b/pennylane/measurements/measurements.py
@@ -16,7 +16,6 @@
 outcomes from quantum observables - expectation values, variances of expectations,
 and measurement samples using AnnotatedQueues.
 """
-import contextlib
 import copy
 import functools
 from abc import ABC, abstractmethod
@@ -335,8 +334,7 @@ def eigvals(self):
             return qml.math.arange(0, 2 ** len(self.wires), 1)
 
         if self.obs is not None:
-            with contextlib.suppress(qml.operation.EigvalsUndefinedError):
-                return self.obs.eigvals()
+            return self.obs.eigvals()
         return self._eigvals
 
     @property

From 150304db7946540875ad49b12e71beb1e4d00a79 Mon Sep 17 00:00:00 2001
From: Mudit Pandey <mudit.pandey@xanadu.ai>
Date: Wed, 13 Sep 2023 09:57:43 -0400
Subject: [PATCH 050/131] Merging

---
 pennylane/transforms/defer_measurements.py | 61 ++++++++++++++++------
 1 file changed, 46 insertions(+), 15 deletions(-)

diff --git a/pennylane/transforms/defer_measurements.py b/pennylane/transforms/defer_measurements.py
index 066bfd88c85..ec1b97a40d2 100644
--- a/pennylane/transforms/defer_measurements.py
+++ b/pennylane/transforms/defer_measurements.py
@@ -12,19 +12,22 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 """Code for the tape transform implementing the deferred measurement principle."""
+from typing import Sequence, Callable
 import pennylane as qml
 from pennylane.measurements import MidMeasureMP, MeasurementValue
 from pennylane.ops.op_math import ctrl
-from pennylane.queuing import apply
+
 from pennylane.tape import QuantumTape
-from pennylane.transforms import qfunc_transform
+from pennylane.transforms.core import transform
+
 from pennylane.wires import Wires
+from pennylane.queuing import QueuingManager
 
 # pylint: disable=too-many-branches
 
 
-@qfunc_transform
-def defer_measurements(tape: QuantumTape):
+@transform
+def defer_measurements(tape: QuantumTape) -> (Sequence[QuantumTape], Callable):
     """Quantum function transform that substitutes operations conditioned on
     measurement outcomes to controlled operations.
 
@@ -67,6 +70,13 @@ def defer_measurements(tape: QuantumTape):
     Args:
         tape (.QuantumTape): a quantum tape
 
+    Returns:
+        pennylane.QNode or qfunc or tuple[List[.QuantumTape], function]: If a QNode is passed,
+        it returns a QNode with the transform added to its transform program.
+        If a tape is passed, returns a tuple containing a list of
+        quantum tapes to be evaluated, and a function to be applied to these
+        tape executions.
+
     **Example**
 
     Suppose we have a quantum function with mid-circuit measurements and
@@ -127,6 +137,8 @@ def func(x, y):
     if ops_cv or obs_cv:
         raise ValueError("Continuous variable operations and observables are not supported.")
 
+    new_operations = []
+
     # Find wires that are reused after measurement
     measured_wires = []
     reused_measurement_wires = set()
@@ -160,35 +172,54 @@ def func(x, y):
             # Store measurement outcome in new wire if wire gets reused
             if op.wires[0] in reused_measurement_wires or op.wires[0] in measured_wires:
                 control_wires[op.id] = cur_wire
-
-                qml.CNOT([op.wires[0], cur_wire])
+                with QueuingManager.stop_recording():
+                    new_operations.append(qml.CNOT([op.wires[0], cur_wire]))
                 if op.reset:
-                    qml.CNOT([cur_wire, op.wires[0]])
+                    with QueuingManager.stop_recording():
+                        new_operations.append(qml.CNOT([cur_wire, op.wires[0]]))
 
                 cur_wire += 1
             else:
                 control_wires[op.id] = op.wires[0]
 
         elif op.__class__.__name__ == "Conditional":
-            _add_control_gate(op, control_wires)
+            with QueuingManager.stop_recording():
+                new_operations.extend(_add_control_gate(op, control_wires))
         else:
-            apply(op)
+            new_operations.append(op)
+
+    new_measurements = []
 
     for mp in tape.measurements:
         if isinstance(mp.obs, MeasurementValue):
             mp.obs._wires = Wires([control_wires[mp.obs.measurements[0].id]])
-        apply(mp)
+        new_measurements.append(mp)
+
+    new_tape = QuantumTape(new_operations, new_measurements, shots=tape.shots)
+    new_tape._qfunc_output = tape._qfunc_output  # pylint: disable=protected-access
 
-    return tape._qfunc_output  # pylint: disable=protected-access
+    def null_postprocessing(results):
+        """A postprocesing function returned by a transform that only converts the batch of results
+        into a result for a single ``QuantumTape``.
+        """
+        return results[0]
+
+    return [new_tape], null_postprocessing
 
 
 def _add_control_gate(op, control_wires):
     """Helper function to add control gates"""
     control = [control_wires[m.id] for m in op.meas_val.measurements]
+    new_ops = []
+
     for branch, value in op.meas_val._items():  # pylint: disable=protected-access
         if value:
-            ctrl(
-                lambda: apply(op.then_op),  # pylint: disable=cell-var-from-loop
-                control=Wires(control),
-                control_values=branch,
+            qscript = qml.tape.make_qscript(
+                ctrl(
+                    lambda: qml.apply(op.then_op),  # pylint: disable=cell-var-from-loop
+                    control=Wires(control),
+                    control_values=branch,
+                )
             )()
+            new_ops.extend(qscript.circuit)
+    return new_ops

From 55c3a3872dd2ddef150682c7e92b2f2da974839d Mon Sep 17 00:00:00 2001
From: Mudit Pandey <mudit.pandey@xanadu.ai>
Date: Wed, 13 Sep 2023 16:52:01 -0400
Subject: [PATCH 051/131] Updated measurements eigvals test

---
 tests/measurements/test_measurements.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/tests/measurements/test_measurements.py b/tests/measurements/test_measurements.py
index ad06fdd5346..04ec964338e 100644
--- a/tests/measurements/test_measurements.py
+++ b/tests/measurements/test_measurements.py
@@ -308,7 +308,8 @@ def test_observable_with_no_eigvals(self):
         the eigvals method to return a NotImplementedError"""
         obs = qml.NumberOperator(wires=0)
         m = qml.expval(op=obs)
-        assert m.eigvals() is None
+        with pytest.raises(qml.operation.EigvalsUndefinedError):
+            _ = m.eigvals()
 
     def test_repr(self):
         """Test the string representation of a MeasurementProcess."""

From bd64ea9baaea6612b96c039fdf194a3ce4330ee7 Mon Sep 17 00:00:00 2001
From: Matthew Silverman <matthews@xanadu.ai>
Date: Wed, 13 Sep 2023 12:42:49 -0400
Subject: [PATCH 052/131] AmplitudeEmbedding inherits from StatePrep (#4583)

**Context:**
While porting DQ2, the merge_amplitude_embedding tests were failing
because `AmplitudeEmbedding` was being decomposed. I figured that since
it behaves like a `StatePrep` op, we should make it inherit from it and
gain the benefit!

**Description of the Change:**
`AmplitudeEmbedding` inherits from `StatePrep`, gaining its
`state_vector` implementation and various properties.

The only real change now that it inherits from `StatePrep` is that it
will return True for `isinstance(op, StatePrep)` checks but I think this
is not a problem.

**Benefits:**
`AmplitudeEmbedding` will behave the same on DQ2 as it did on DQL (it
would decompose once to `StatePrep`, then be "supported" by DQL, and
then it would have the custom apply function - on DQ2, it was
decomposing all the way, and then I hit a bug where it thought it could
support batching but the full decomposition disagreed). I confirmed
locally that all changed tests, `test_amplitude.py`, and
`test_merge_amplitude_embeddings.py` all passed with DQ2 as well.

**Possible Drawbacks:**
If someone somehow calls `expand_tape_state_prep` with
`AmplitudeEmbedding` and `skip_first=False`, it will decompose once to a
`StatePrep` and call it a day. This is reflected in tests, but I figured
it shouldn't be too much of a concern... users don't even have a way to
set `skip_first` to False (...yet?)
---
 doc/releases/changelog-dev.md                |  4 +++
 pennylane/templates/embeddings/amplitude.py  | 21 +++++----------
 tests/devices/qubit/test_initialize_state.py | 28 +++++++++-----------
 tests/devices/qubit/test_preprocess.py       | 15 ++++++++---
 tests/tape/test_tape.py                      |  2 ++
 5 files changed, 36 insertions(+), 34 deletions(-)

diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index d19e0ba7063..1c1e1bfcfb2 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -87,6 +87,10 @@
 * `DefaultQubit2` now works as expected with measurement processes that don't specify wires.
   [(#4580)](https://github.com/PennyLaneAI/pennylane/pull/4580)
 
+* `AmplitudeEmbedding` now inherits from `StatePrep`, allowing for it to not be decomposed
+  when at the beginning of a circuit, thus behaving like `StatePrep`.
+  [(#4583)](https://github.com/PennyLaneAI/pennylane/pull/4583)
+
 <h3>Breaking changes 💔</h3>
 
 * The `__eq__` and `__hash__` methods of `Operator` and `MeasurementProcess` no longer rely on the
diff --git a/pennylane/templates/embeddings/amplitude.py b/pennylane/templates/embeddings/amplitude.py
index 67cc0e2f62e..96009db8d2a 100644
--- a/pennylane/templates/embeddings/amplitude.py
+++ b/pennylane/templates/embeddings/amplitude.py
@@ -18,7 +18,6 @@
 import numpy as np
 
 import pennylane as qml
-from pennylane.operation import Operation, AnyWires
 from pennylane.ops import StatePrep
 from pennylane.wires import Wires
 
@@ -26,7 +25,7 @@
 TOLERANCE = 1e-10
 
 
-class AmplitudeEmbedding(Operation):
+class AmplitudeEmbedding(StatePrep):
     r"""Encodes :math:`2^n` features into the amplitude vector of :math:`n` qubits.
 
     By setting ``pad_with`` to a real or complex number, ``features`` is automatically padded to dimension
@@ -120,26 +119,18 @@ def circuit(f=None):
 
     """
 
-    num_wires = AnyWires
-    grad_method = None
-
     def __init__(self, features, wires, pad_with=None, normalize=False, id=None):
+        # pylint:disable=bad-super-call
         wires = Wires(wires)
         self.pad_with = pad_with
         self.normalize = normalize
         features = self._preprocess(features, wires, pad_with, normalize)
-        super().__init__(features, wires=wires, id=id)
-
-    @property
-    def num_params(self):
-        return 1
-
-    @property
-    def ndim_params(self):
-        return (1,)
+        super(StatePrep, self).__init__(features, wires=wires, id=id)
 
     @staticmethod
-    def compute_decomposition(features, wires):  # pylint: disable=arguments-differ
+    def compute_decomposition(
+        features, wires
+    ):  # pylint: disable=arguments-differ,arguments-renamed
         r"""Representation of the operator as a product of other operators.
 
         .. math:: O = O_1 O_2 \dots O_n.
diff --git a/tests/devices/qubit/test_initialize_state.py b/tests/devices/qubit/test_initialize_state.py
index 9f8e15ff0d0..5b008b78115 100644
--- a/tests/devices/qubit/test_initialize_state.py
+++ b/tests/devices/qubit/test_initialize_state.py
@@ -58,28 +58,24 @@ def test_create_initial_state_with_BasisState(self):
         state[0, 1, 0] = 0  # set to zero to make test below simple
         assert qml.math.allequal(state, np.zeros((2, 2, 2)))
 
-    def test_create_initial_state_with_StatePrep(self):
+    @pytest.mark.parametrize("prep_op_cls", [qml.StatePrep, qml.AmplitudeEmbedding])
+    def test_create_initial_state_with_StatePrep(self, prep_op_cls):
         """Tests that create_initial_state works with the StatePrep operator."""
-        prep_op = qml.StatePrep(np.array([0, 1, 0, 0, 0, 0, 0, 1]) / np.sqrt(2), wires=[0, 1, 2])
+        prep_op = prep_op_cls(np.array([0, 1, 0, 0, 0, 0, 0, 1]) / np.sqrt(2), wires=[0, 1, 2])
         state = create_initial_state([0, 1, 2], prep_operation=prep_op)
-        assert state[0, 0, 1] == 1 / np.sqrt(2)
-        assert state[1, 1, 1] == 1 / np.sqrt(2)
-        state[0, 0, 1] = 0
-        state[1, 1, 1] = 0  # set to zero to make test below simple
-        assert qml.math.allequal(state, np.zeros((2, 2, 2)))
+        expected = np.zeros((2, 2, 2))
+        expected[0, 0, 1] = expected[1, 1, 1] = 1 / np.sqrt(2)
+        assert np.array_equal(state, expected)
 
-    def test_create_initial_state_with_StatePrep_broadcasted(self):
+    @pytest.mark.parametrize("prep_op_cls", [qml.StatePrep, qml.AmplitudeEmbedding])
+    def test_create_initial_state_with_StatePrep_broadcasted(self, prep_op_cls):
         """Tests that create_initial_state works with a broadcasted StatePrep
         operator."""
-        prep_op = qml.StatePrep(np.array([[0, 1, 0, 0], [0, 0, 0, 1], [0, 0, 1, 0]]), wires=[0, 1])
+        prep_op = prep_op_cls(np.array([[0, 1, 0, 0], [0, 0, 0, 1], [0, 0, 1, 0]]), wires=[0, 1])
         state = create_initial_state([0, 1], prep_operation=prep_op)
-        assert state[0, 0, 1] == 1
-        assert state[1, 1, 1] == 1
-        assert state[2, 1, 0] == 1
-        state[0, 0, 1] = 0
-        state[1, 1, 1] = 0
-        state[2, 1, 0] = 0  # set to zero to make test below simple
-        assert qml.math.allequal(state, np.zeros((3, 2, 2)))
+        expected = np.zeros((3, 2, 2))
+        expected[0, 0, 1] = expected[1, 1, 1] = expected[2, 1, 0] = 1
+        assert np.array_equal(state, expected)
 
     @pytest.mark.torch
     def test_create_initial_state_casts_to_like_with_prep_op(self):
diff --git a/tests/devices/qubit/test_preprocess.py b/tests/devices/qubit/test_preprocess.py
index 07dfb4a243e..01c7f528554 100644
--- a/tests/devices/qubit/test_preprocess.py
+++ b/tests/devices/qubit/test_preprocess.py
@@ -260,7 +260,12 @@ def test_expand_fn_non_commuting_measurements(self):
         assert new_qs.measurements == qs.measurements
 
     @pytest.mark.parametrize(
-        "prep_op", (qml.BasisState([1], wires=0), qml.StatePrep([0, 1], wires=1))
+        "prep_op",
+        (
+            qml.BasisState([1], wires=0),
+            qml.StatePrep([0, 1], wires=1),
+            qml.AmplitudeEmbedding([0, 1], wires=1),
+        ),
     )
     def test_expand_fn_state_prep(self, prep_op):
         """Test that the expand_fn only expands mid-circuit instances of StatePrepBase"""
@@ -270,6 +275,7 @@ def test_expand_fn_state_prep(self, prep_op):
             qml.StatePrep([0, 1], wires=1),
             qml.BasisState([1], wires=0),
             qml.RZ(0.123, wires=1),
+            qml.AmplitudeEmbedding([0, 1, 0, 0], wires=[0, 1]),
         ]
         measurements = [qml.expval(qml.PauliZ(0)), qml.probs()]
         tape = QuantumScript(ops=ops, measurements=measurements)
@@ -282,10 +288,13 @@ def test_expand_fn_state_prep(self, prep_op):
             qml.RY(3.14159265, wires=1),  # decomposition of StatePrep
             qml.PauliX(wires=0),  # decomposition of BasisState
             qml.RZ(0.123, wires=1),
+            qml.RY(1.57079633, wires=[1]),  # decomposition of AmplitudeEmbedding
+            qml.CNOT(wires=[0, 1]),
+            qml.RY(1.57079633, wires=[1]),
+            qml.CNOT(wires=[0, 1]),
         ]
 
-        for op, exp in zip(expanded_tape.circuit, expected + measurements):
-            assert qml.equal(op, exp)
+        assert expanded_tape.circuit == expected + measurements
 
 
 class TestValidateMeasurements:
diff --git a/tests/tape/test_tape.py b/tests/tape/test_tape.py
index bf873447dd6..a8c33a9d73f 100644
--- a/tests/tape/test_tape.py
+++ b/tests/tape/test_tape.py
@@ -1014,11 +1014,13 @@ def test_depth_expansion(self):
             [
                 qml.BasisState([1, 0], wires=[0, 1]),
                 qml.StatePrep([0, 1, 0, 0], wires=[0, 1]),
+                qml.AmplitudeEmbedding([0, 1, 0, 0], wires=[0, 1]),
                 qml.PauliZ(0),
             ],
             [
                 qml.BasisStatePreparation([1, 0], wires=[0, 1]),
                 qml.MottonenStatePreparation([0, 1, 0, 0], wires=[0, 1]),
+                qml.StatePrep([0, 1, 0, 0], wires=[0, 1]),  # still a StatePrepBase :/
                 qml.PauliZ(0),
             ],
         ),

From ba35f2abc261e4eca009ee50ffb5592299faf6da Mon Sep 17 00:00:00 2001
From: Utkarsh <utkarshazad98@gmail.com>
Date: Wed, 13 Sep 2023 15:43:03 -0400
Subject: [PATCH 053/131] Add support for offset in `qml.MPS` template (#4531)

**Context:** `qml.MPS` allows users to create MPS circuits according to
the paper [Towards Quantum Machine Learning with Tensor
Networks](https://arxiv.org/abs/1803.11537). But users would like to
have more flexibility in controlling the tensor-network connectivity.

**Description of the Change:** Adds a new keyword argument `offset` that
allows users to change position of subsequent block from the standard
`n_block_wires/2`

**Benefits:** Makes the template more flexible

**Possible Drawbacks:** N/A

**Related GitHub Issues:** [Forum
Discussion](https://discuss.pennylane.ai/t/qml-mps-offset-control/2890)
---
 doc/releases/changelog-dev.md                 |   3 +
 pennylane/templates/tensornetworks/mps.py     | 161 ++++++++++-----
 .../templates/test_tensornetworks/test_MPS.py | 186 ++++++++++++++----
 3 files changed, 262 insertions(+), 88 deletions(-)

diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 1c1e1bfcfb2..a6a72b3bccd 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -38,6 +38,9 @@
 
 <h3>Improvements 🛠</h3>
 
+* Tensor-network template `qml.MPS` now supports changing `offset` between subsequent blocks for more flexibility.
+ [(#4531)](https://github.com/PennyLaneAI/pennylane/pull/4531)
+
 * The qchem ``fermionic_dipole`` and ``particle_number`` functions are updated to use a
   ``FermiSentence``. The deprecated features for using tuples to represent fermionic operations are
   removed.
diff --git a/pennylane/templates/tensornetworks/mps.py b/pennylane/templates/tensornetworks/mps.py
index db87064061f..2c0888677b7 100644
--- a/pennylane/templates/tensornetworks/mps.py
+++ b/pennylane/templates/tensornetworks/mps.py
@@ -17,28 +17,28 @@
 # pylint: disable-msg=too-many-branches,too-many-arguments,protected-access
 import warnings
 import pennylane as qml
-import pennylane.numpy as np
 from pennylane.operation import Operation, AnyWires
 
 
-def compute_indices_MPS(wires, n_block_wires):
-    """Generate a list containing the wires for each block.
+def compute_indices_MPS(wires, n_block_wires, offset=None):
+    r"""Generate a list containing the wires for each block.
 
     Args:
         wires (Iterable): wires that the template acts on
-        n_block_wires (int): number of wires per block
+        n_block_wires (int): number of wires per block_gen
+        offset (int): offset value for positioning the subsequent blocks relative to each other.
+            If ``None``, it defaults to :math:`\text{offset} = \lfloor \text{n_block_wires}/2  \rfloor`,
+            otherwise :math:`\text{offset} \in [1, \text{n_block_wires} - 1]`.
+
     Returns:
         layers (Tuple[Tuple]]): array of wire indices or wire labels for each block
     """
 
     n_wires = len(wires)
 
-    if n_block_wires % 2 != 0:
-        raise ValueError(f"n_block_wires must be an even integer; got {n_block_wires}")
-
     if n_block_wires < 2:
         raise ValueError(
-            f"number of wires in each block must be larger than or equal to 2; got n_block_wires = {n_block_wires}"
+            f"The number of wires in each block must be larger than or equal to 2; got n_block_wires = {n_block_wires}"
         )
 
     if n_block_wires > n_wires:
@@ -46,27 +46,35 @@ def compute_indices_MPS(wires, n_block_wires):
             f"n_block_wires must be smaller than or equal to the number of wires; got n_block_wires = {n_block_wires} and number of wires = {n_wires}"
         )
 
-    if n_wires % (n_block_wires / 2) > 0:
-        warnings.warn(
-            f"The number of wires should be a multiple of {int(n_block_wires/2)}; got {n_wires}"
+    if offset is None:
+        offset = n_block_wires // 2
+
+    if offset < 1 or offset > n_block_wires - 1:
+        raise ValueError(
+            f"Provided offset is outside the expected range; the expected range for n_block_wires = {n_block_wires} is range{1, n_block_wires - 1}"
         )
 
+    n_step = offset
+    n_layers = len(wires) - int(len(wires) % (n_block_wires // 2)) - n_step
+
     return tuple(
         tuple(wires[idx] for idx in range(j, j + n_block_wires))
         for j in range(
             0,
-            len(wires) - int(len(wires) % (n_block_wires // 2)) - n_block_wires // 2,
-            n_block_wires // 2,
+            n_layers,
+            n_step,
         )
+        if not j + n_block_wires > len(wires)
     )
 
 
 class MPS(Operation):
-    """The MPS template broadcasts an input circuit across many wires following the architecture of a Matrix Product State tensor network.
+    r"""The MPS template broadcasts an input circuit across many wires following the architecture of a Matrix Product State tensor network.
     The result is similar to the architecture in `arXiv:1803.11537 <https://arxiv.org/abs/1803.11537>`_.
 
-    The argument ``block`` is a user-defined quantum circuit.``block`` should have two arguments: ``weights`` and ``wires``.
-    For clarity, it is recommended to use a one-dimensional list or array for the block weights.
+    The keyword argument ``block`` is a user-defined quantum circuit that should accept two arguments: ``wires`` and ``weights``.
+    The latter argument is optional in case the implementation of ``block`` doesn't require any weights. Any additional arguments
+    should be provided using the ``kwargs``.
 
     Args:
         wires (Iterable): wires that the template acts on
@@ -74,11 +82,17 @@ class MPS(Operation):
         block (Callable): quantum circuit that defines a block
         n_params_block (int): the number of parameters in a block; equal to the length of the ``weights`` argument in ``block``
         template_weights (Sequence): list containing the weights for all blocks
+        offset (int): offset value for positioning the subsequent blocks relative to each other.
+            If ``None``, it defaults to :math:`\text{offset} = \lfloor \text{n_block_wires}/2  \rfloor`,
+            otherwise :math:`\text{offset} \in [1, \text{n_block_wires} - 1]`
+        **kwargs: additional keyword arguments for implementing the ``block``
 
     .. note::
 
-        The expected number of blocks can be obtained from ``qml.MPS.get_n_blocks(wires, n_block_wires)``.
-        The length of ``template_weights`` argument should match the number of blocks.
+        The expected number of blocks can be obtained from ``qml.MPS.get_n_blocks(wires, n_block_wires, offset=0)``, and
+        the length of ``template_weights`` argument should match the number of blocks. Whenever either ``n_block_wires``
+        is odd or ``offset`` is not :math:`\lfloor \text{n_block_wires}/2  \rfloor`, the template deviates from the maximally
+        unbalanced tree architecture described in `arXiv:1803.11537 <https://arxiv.org/abs/1803.11537>`_.
 
     .. details::
         :title: Usage Details
@@ -99,7 +113,7 @@ def block(weights, wires):
             n_block_wires = 2
             n_params_block = 2
             n_blocks = qml.MPS.get_n_blocks(range(n_wires),n_block_wires)
-            template_weights = [[0.1,-0.3]]*n_blocks
+            template_weights = [[0.1, -0.3]] * n_blocks
 
             dev= qml.device('default.qubit',wires=range(n_wires))
             @qml.qnode(dev)
@@ -113,10 +127,37 @@ def circuit(template_weights):
         2: ───────────────╰X──RY(-0.30)─╭●──RY(0.10)──┤
         3: ─────────────────────────────╰X──RY(-0.30)─┤  <Z>
 
-    """
+        MPS can also be used with an ``offset`` argument that shifts the positioning the subsequent blocks from the default ``n_block_wires/2``.
+
+        .. code-block:: python
 
-    num_params = 1
-    """int: Number of trainable parameters that the operator depends on."""
+            import pennylane as qml
+            import numpy as np
+
+            def block(wires):
+                qml.MultiControlledX(wires=[wires[i] for i in range(len(wires))])
+
+            n_wires = 8
+            n_block_wires = 4
+            n_params_block = 2
+
+            dev= qml.device('default.qubit',wires=n_wires)
+            @qml.qnode(dev)
+            def circuit():
+                qml.MPS(range(n_wires),n_block_wires, block, n_params_block, offset = 1)
+                return qml.state()
+
+        >>> print(qml.draw(circuit, expansion_strategy='device')())
+        0: ─╭●─────────────┤  State
+        1: ─├●─╭●──────────┤  State
+        2: ─├●─├●─╭●───────┤  State
+        3: ─╰X─├●─├●─╭●────┤  State
+        4: ────╰X─├●─├●─╭●─┤  State
+        5: ───────╰X─├●─├●─┤  State
+        6: ──────────╰X─├●─┤  State
+        7: ─────────────╰X─┤  State
+
+    """
 
     num_wires = AnyWires
     par_domain = "A"
@@ -124,8 +165,9 @@ def circuit(template_weights):
     @classmethod
     def _unflatten(cls, data, metadata):
         new_op = cls.__new__(cls)
-        new_op._hyperparameters = dict(metadata[1])
-        Operation.__init__(new_op, data, wires=metadata[0])
+        setattr(new_op, "_hyperparameters", dict(metadata[1]))
+        setattr(new_op, "_weights", data[0] if len(data) else None)
+        Operation.__init__(new_op, *data, wires=metadata[0])
         return new_op
 
     def __init__(
@@ -133,19 +175,17 @@ def __init__(
         wires,
         n_block_wires,
         block,
-        n_params_block,
+        n_params_block=0,
         template_weights=None,
+        offset=None,
         id=None,
+        **kwargs,
     ):
-        ind_gates = compute_indices_MPS(wires, n_block_wires)
-        n_wires = len(wires)
-        n_blocks = int(n_wires / (n_block_wires / 2) - 1)
-
-        if template_weights is None:
-            template_weights = np.random.rand(n_params_block, int(n_blocks))
+        ind_gates = compute_indices_MPS(wires, n_block_wires, offset)
+        n_blocks = self.get_n_blocks(wires, n_block_wires, offset)
 
-        else:
-            shape = qml.math.shape(template_weights)[-4:]  # (n_params_block, n_blocks)
+        if template_weights is not None:
+            shape = qml.math.shape(template_weights)  # (n_blocks, n_params_block)
             if shape[0] != n_blocks:
                 raise ValueError(
                     f"Weights tensor must have first dimension of length {n_blocks}; got {shape[0]}"
@@ -155,19 +195,27 @@ def __init__(
                     f"Weights tensor must have last dimension of length {n_params_block}; got {shape[-1]}"
                 )
 
-        self._hyperparameters = {"ind_gates": ind_gates, "block": block}
-        super().__init__(template_weights, wires=wires, id=id)
+        self._weights = template_weights
+        self._hyperparameters = {"ind_gates": ind_gates, "block": block, **kwargs}
+
+        if self._weights is None:
+            super().__init__(wires=wires, id=id)
+        else:
+            super().__init__(self._weights, wires=wires, id=id)
+
+    @property
+    def num_params(self):
+        """int: Number of trainable parameters that the operator depends on."""
+        return 0 if self._weights is None else 1
 
     @staticmethod
     def compute_decomposition(
-        weights, wires, ind_gates, block
+        weights=None, wires=None, ind_gates=None, block=None, **kwargs
     ):  # pylint: disable=arguments-differ,unused-argument
         r"""Representation of the operator as a product of other operators.
 
         .. math:: O = O_1 O_2 \dots O_n.
 
-
-
         .. seealso:: :meth:`~.MPS.decomposition`.
 
         Args:
@@ -175,27 +223,40 @@ def compute_decomposition(
             wires (Iterable): wires that the template acts on
             block (Callable): quantum circuit that defines a block
             ind_gates (array): array of wire indices
+            **kwargs: additional keyword arguments for implementing the ``block``
 
         Returns:
             list[.Operator]: decomposition of the operator
         """
         decomp = []
+        itrweights = iter([]) if weights is None else iter(weights)
         block_gen = qml.tape.make_qscript(block)
-        for idx, w in enumerate(ind_gates):
-            decomp += block_gen(weights=weights[idx][:], wires=w)
+        for w in ind_gates:
+            weight = next(itrweights, None)
+            decomp += (
+                block_gen(wires=w, **kwargs)
+                if weight is None
+                else block_gen(weights=weight, wires=w, **kwargs)
+            )
         return [qml.apply(op) for op in decomp] if qml.QueuingManager.recording() else decomp
 
     @staticmethod
-    def get_n_blocks(wires, n_block_wires):
-        """Returns the expected number of blocks for a set of wires and number of wires per block.
+    def get_n_blocks(wires, n_block_wires, offset=None):
+        r"""Returns the expected number of blocks for a set of wires and number of wires per block.
+
         Args:
             wires (Sequence): number of wires the template acts on
             n_block_wires (int): number of wires per block
+            offset (int): offset value for positioning the subsequent blocks relative to each other.
+                If ``None``, it defaults to :math:`\text{offset} = \lfloor \text{n_block_wires}/2  \rfloor`,
+                otherwise :math:`\text{offset} \in [1, \text{n_block_wires} - 1]`.
+
         Returns:
             n_blocks (int): number of blocks; expected length of the template_weights argument
         """
         n_wires = len(wires)
-        if n_wires % (n_block_wires / 2) > 0:
+
+        if offset is None and not n_block_wires % 2 and n_wires % (n_block_wires // 2) > 0:
             warnings.warn(
                 f"The number of wires should be a multiple of {int(n_block_wires/2)}; got {n_wires}"
             )
@@ -205,5 +266,15 @@ def get_n_blocks(wires, n_block_wires):
                 f"n_block_wires must be smaller than or equal to the number of wires; got n_block_wires = {n_block_wires} and number of wires = {n_wires}"
             )
 
-        n_blocks = int(n_wires / (n_block_wires / 2) - 1)
-        return n_blocks
+        if offset is None:
+            offset = n_block_wires // 2
+
+        if offset < 1 or offset > n_block_wires - 1:
+            raise ValueError(
+                f"Provided offset is outside the expected range; the expected range for n_block_wires = {n_block_wires} is range{1, n_block_wires - 1}"
+            )
+
+        n_step = offset
+        n_layers = n_wires - int(n_wires % (n_block_wires // 2)) - n_step
+
+        return len([idx for idx in range(0, n_layers, n_step) if not idx + n_block_wires > n_wires])
diff --git a/tests/templates/test_tensornetworks/test_MPS.py b/tests/templates/test_tensornetworks/test_MPS.py
index 9d03571f362..e7eb3beb614 100644
--- a/tests/templates/test_tensornetworks/test_MPS.py
+++ b/tests/templates/test_tensornetworks/test_MPS.py
@@ -25,8 +25,9 @@
 def test_flatten_unflatten():
     """Test the flatten and unflatten methods."""
 
-    def block(weights, wires):
-        qml.CNOT(wires=[wires[0], wires[1]])
+    def block(weights, wires, use_CNOT=True):
+        if use_CNOT:
+            qml.CNOT(wires=[wires[0], wires[1]])
         qml.RY(weights[0], wires=wires[0])
         qml.RY(weights[1], wires=wires[1])
 
@@ -38,7 +39,7 @@ def block(weights, wires):
 
     wires = qml.wires.Wires((0, 1, 2, 3))
 
-    op = qml.MPS(wires, n_block_wires, block, n_params_block, template_weights)
+    op = qml.MPS(wires, n_block_wires, block, n_params_block, template_weights, use_CNOT=True)
 
     data, metadata = op._flatten()
     assert len(data) == 1
@@ -59,22 +60,6 @@ def block(weights, wires):
 class TestIndicesMPS:
     """Test function that computes MPS indices"""
 
-    @pytest.mark.parametrize(
-        ("n_wires", "n_block_wires"),
-        [
-            (5, 3),
-            (9, 5),
-            (11, 7),
-        ],
-    )
-    def test_exception_n_block_wires_uneven(self, n_wires, n_block_wires):
-        """Verifies that an exception is raised if n_block_wires is not even."""
-
-        with pytest.raises(
-            ValueError, match=f"n_block_wires must be an even integer; got {n_block_wires}"
-        ):
-            compute_indices_MPS(range(n_wires), n_block_wires)
-
     @pytest.mark.parametrize(
         ("n_wires", "n_block_wires"),
         [
@@ -106,33 +91,46 @@ def test_exception_n_block_wires_small(self):
             compute_indices_MPS(range(n_wires), n_block_wires)
 
     @pytest.mark.parametrize(
-        ("n_wires", "n_block_wires"),
+        ("n_wires", "n_block_wires", "offset"),
         [
-            (5, 4),
-            (9, 4),
-            (7, 6),
+            (18, 6, 6),
+            (12, 4, 0),
+            (10, 4, 4),
         ],
     )
-    def test_warning_many_wires(self, n_wires, n_block_wires):
-        """Verifies that a warning is raised if n_wires doesn't correspond to n_block_wires."""
+    def test_exception_offset(self, n_wires, n_block_wires, offset):
+        """Verifies that an exception is raised when abs(offset) is more than n_block_wires / 2."""
 
-        with pytest.warns(
-            Warning,
-            match=f"The number of wires should be a multiple of {int(n_block_wires/2)}; "
-            f"got {n_wires}",
+        with pytest.raises(
+            ValueError,
+            match="Provided offset is outside the expected range; "
+            f"the expected range for n_block_wires = {n_block_wires}",
         ):
-            compute_indices_MPS(range(n_wires), n_block_wires)
+            compute_indices_MPS(range(n_wires), n_block_wires, offset)
 
     @pytest.mark.parametrize(
-        ("wires", "n_block_wires", "expected_indices"),
+        ("wires", "n_block_wires", "offset", "expected_indices"),
         [
-            ([1, 2, 3, 4], 2, ((1, 2), (2, 3), (3, 4))),
-            (["a", "b", "c", "d"], 2, (("a", "b"), ("b", "c"), ("c", "d"))),
+            ([1, 2, 3, 4], 2, 1, ((1, 2), (2, 3), (3, 4))),
+            (["a", "b", "c", "d"], 2, 1, (("a", "b"), ("b", "c"), ("c", "d"))),
+            ([1, 2, 3, 4, 5, 6, 7, 8], 4, 3, ((1, 2, 3, 4), (4, 5, 6, 7))),
+            (
+                ["a", "b", "c", "d", "e", "f", "g", "h"],
+                4,
+                1,
+                (
+                    ("a", "b", "c", "d"),
+                    ("b", "c", "d", "e"),
+                    ("c", "d", "e", "f"),
+                    ("d", "e", "f", "g"),
+                    ("e", "f", "g", "h"),
+                ),
+            ),
         ],
     )
-    def test_indices_output(self, wires, n_block_wires, expected_indices):
+    def test_indices_output(self, wires, n_block_wires, offset, expected_indices):
         """Verifies the indices are correct for both integer and string wire labels."""
-        indices = compute_indices_MPS(wires, n_block_wires)
+        indices = compute_indices_MPS(wires, n_block_wires, offset)
         assert indices == expected_indices
 
 
@@ -140,14 +138,14 @@ class TestTemplateInputs:
     """Test template inputs and pre-processing (ensure the correct exceptions are thrown for the inputs)"""
 
     @pytest.mark.parametrize(
-        ("block", "n_params_block", "wires", "n_block_wires", "msg_match"),
+        ("block", "n_params_block", "wires", "n_block_wires", "offset", "msg_match"),
         [
-            (None, None, [1, 2, 3, 4], 7, "n_block_wires must be an even integer; got 7"),
             (
                 None,
                 None,
                 [1, 2, 3, 4],
                 6,
+                None,
                 "n_block_wires must be smaller than or equal to the number of wires; "
                 "got n_block_wires = 6 and number of wires = 4",
             ),
@@ -156,15 +154,26 @@ class TestTemplateInputs:
                 None,
                 [1, 2, 3, 4],
                 0,
-                "number of wires in each block must be larger than or equal to 2; "
+                None,
+                "The number of wires in each block must be larger than or equal to 2; "
                 "got n_block_wires = 0",
             ),
+            (
+                None,
+                None,
+                [1, 2, 3, 4, 5, 6, 7, 8],
+                4,
+                4,
+                "Provided offset is outside the expected range; ",
+            ),
         ],
     )
-    def test_exception_wrong_input(self, block, n_params_block, wires, n_block_wires, msg_match):
+    def test_exception_wrong_input(
+        self, block, n_params_block, wires, n_block_wires, offset, msg_match
+    ):
         """Verifies that an exception is raised if the number of wires or n_block_wires is incorrect."""
         with pytest.raises(ValueError, match=msg_match):
-            MPS(wires, n_block_wires, block, n_params_block)
+            MPS(wires, n_block_wires, block, n_params_block, offset=offset)
 
     def test_warning_many_wires(self):
         """Verifies that a warning is raised if n_wires doesn't correspond to n_block_wires."""
@@ -256,6 +265,36 @@ def test_get_n_blocks_error(self, wires, n_block_wires):
         ):
             qml.MPS.get_n_blocks(wires, n_block_wires)
 
+    @pytest.mark.filterwarnings("ignore")
+    @pytest.mark.parametrize(
+        ("wires", "n_block_wires", "offset", "expected_n_blocks"),
+        [
+            (range(14), 4, 1, 11),
+            (range(15), 4, 3, 4),
+            (range(18), 6, 1, 13),
+            (range(20), 6, 5, 3),
+        ],
+    )
+    def test_get_n_blocks_with_offset(self, wires, n_block_wires, offset, expected_n_blocks):
+        """Test that the number of blocks attribute returns the correct number of blocks with offset."""
+
+        assert qml.MPS.get_n_blocks(wires, n_block_wires, offset) == expected_n_blocks
+
+    @pytest.mark.filterwarnings("ignore")
+    @pytest.mark.parametrize(
+        ("wires", "n_block_wires", "offset"),
+        [(range(12), 6, 6), (range(9), 4, 0)],
+    )
+    def test_get_n_blocks_error_with_offset(self, wires, offset, n_block_wires):
+        """Test that the number of blocks attribute raises an error when offset is out of bounds."""
+
+        with pytest.raises(
+            ValueError,
+            match=r"Provided offset is outside the expected range; "
+            f"the expected range for n_block_wires = {n_block_wires}",
+        ):
+            qml.MPS.get_n_blocks(wires, n_block_wires, offset)
+
 
 class TestTemplateOutputs:
     """Test the output of the MPS template."""
@@ -302,6 +341,25 @@ def circuit2_MPS(weights, wires):
         qml.StronglyEntanglingLayers(SELWeights1, wires=wires[0:2])
         qml.StronglyEntanglingLayers(SELWeights2, wires=wires[1:3])
 
+    @staticmethod
+    def circuit3_block(wires, k=None):
+        qml.MultiControlledX(wires=[wires[i] for i in range(len(wires))])
+        assert k == 2
+
+    @staticmethod
+    def circuit3_MPS(wires, **kwargs):  # pylint: disable=unused-argument
+        qml.MultiControlledX(wires=[wires[0], wires[1], wires[2], wires[3]])
+        qml.MultiControlledX(wires=[wires[1], wires[2], wires[3], wires[4]])
+        qml.MultiControlledX(wires=[wires[2], wires[3], wires[4], wires[5]])
+        qml.MultiControlledX(wires=[wires[3], wires[4], wires[5], wires[6]])
+        qml.MultiControlledX(wires=[wires[4], wires[5], wires[6], wires[7]])
+
+    @staticmethod
+    def circuit4_MPS(wires, **kwargs):  # pylint: disable=unused-argument
+        qml.MultiControlledX(wires=[wires[0], wires[1], wires[2], wires[3]])
+        qml.MultiControlledX(wires=[wires[2], wires[3], wires[4], wires[5]])
+        qml.MultiControlledX(wires=[wires[4], wires[5], wires[6], wires[7]])
+
     @pytest.mark.parametrize(
         (
             "block",
@@ -309,6 +367,8 @@ def circuit2_MPS(weights, wires):
             "wires",
             "n_block_wires",
             "template_weights",
+            "offset",
+            "kwargs",
             "expected_circuit",
         ),
         [
@@ -318,6 +378,8 @@ def circuit2_MPS(weights, wires):
                 [1, 2, 3, 4],
                 2,
                 [[0.1, 0.2], [-0.2, 0.3], [0.3, 0.4]],
+                None,
+                {},
                 "circuit1_MPS",
             ),
             (
@@ -326,12 +388,42 @@ def circuit2_MPS(weights, wires):
                 [1, 2, 3],
                 2,
                 [[0.1, 0.2, 0.3], [0.2, 0.3, -0.4]],
+                None,
+                {},
                 "circuit2_MPS",
             ),
+            (
+                "circuit3_block",
+                0,
+                [1, 2, 3, 4, 5, 6, 7, 8],
+                4,
+                None,
+                1,
+                {"k": 2},
+                "circuit3_MPS",
+            ),
+            (
+                "circuit3_block",
+                0,
+                [1, 2, 3, 4, 5, 6, 7, 8, 9],
+                5,
+                None,
+                2,
+                {"k": 2},
+                "circuit4_MPS",
+            ),
         ],
     )
     def test_output(
-        self, block, n_params_block, wires, n_block_wires, template_weights, expected_circuit
+        self,
+        block,
+        n_params_block,
+        wires,
+        n_block_wires,
+        template_weights,
+        offset,
+        kwargs,
+        expected_circuit,
     ):
         """Verifies that the output of the circuits is correct."""
         dev = qml.device("default.qubit", wires=wires)
@@ -340,14 +432,22 @@ def test_output(
 
         @qml.qnode(dev)
         def circuit_template():
-            qml.MPS(wires, n_block_wires, block, n_params_block, template_weights)
+            qml.MPS(
+                wires,
+                n_block_wires,
+                block,
+                n_params_block,
+                template_weights,
+                offset=offset,
+                **kwargs,
+            )
             return qml.expval(qml.PauliZ(wires=wires[-1]))
 
         template_result = circuit_template()
 
         @qml.qnode(dev)
         def circuit_manual():
-            expected_circuit(template_weights, wires)
+            expected_circuit(weights=template_weights, wires=wires)
             return qml.expval(qml.PauliZ(wires=wires[-1]))
 
         manual_result = circuit_manual()

From 4c14d705e1db4aff5243d58bf195db08ccba3b5e Mon Sep 17 00:00:00 2001
From: Diego <67476785+DSGuala@users.noreply.github.com>
Date: Wed, 13 Sep 2023 17:20:30 -0400
Subject: [PATCH 054/131] Fix copying the Select template (#4551)

**Context:**
`qml.Select` could not be copied with `copy.copy(qml.Select(...))`,
resulting in errors with `qml.qsvt`.

**Description of the Change:**
Overrode the `__copy__` function in `qml.Select`

**Benefits:**
Allows `qml.Select` to be used with `qml.qsvt` more easily and allows it
to be copied.

**Possible Drawbacks:**
Changes to `__copy__` in the base `Operator` class will no longer affect
`qml.Select`

---------

Co-authored-by: Jay Soni <jbsoni@uwaterloo.ca>
---
 doc/releases/changelog-dev.md                   |  6 +++++-
 pennylane/templates/subroutines/select.py       | 16 ++++++++++++++++
 tests/templates/test_subroutines/test_select.py |  9 +++++++++
 3 files changed, 30 insertions(+), 1 deletion(-)

diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index a6a72b3bccd..1bf1d2337df 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -208,6 +208,9 @@
 
 <h3>Bug fixes 🐛</h3>
 
+* Fixed issue where `__copy__` method of the `qml.Select()` operator attempted to access un-initialized data.
+[(#4551)](https://github.com/PennyLaneAI/pennylane/pull/4551)
+
 * Fix `skip_first` option in `expand_tape_state_prep`.
   [(#4564)](https://github.com/PennyLaneAI/pennylane/pull/4564)
 
@@ -229,7 +232,8 @@
 
 This release contains contributions from (in alphabetical order):
 
-Utkarsh Azad
+Utkarsh Azad,
+Diego Guala,
 Soran Jahangiri,
 Lillian M. A. Frederiksen,
 Vincent Michaud-Rioux,
diff --git a/pennylane/templates/subroutines/select.py b/pennylane/templates/subroutines/select.py
index 373517030c3..48ece88e59d 100644
--- a/pennylane/templates/subroutines/select.py
+++ b/pennylane/templates/subroutines/select.py
@@ -16,6 +16,7 @@
 """
 # pylint: disable=too-many-arguments
 
+import copy
 import itertools
 import pennylane as qml
 from pennylane.operation import Operation
@@ -98,6 +99,21 @@ def __init__(self, ops, control, id=None):
         all_wires = target_wires + control
         super().__init__(*self.data, wires=all_wires, id=id)
 
+    def __copy__(self):
+        """Copy this op"""
+        cls = self.__class__
+        copied_op = cls.__new__(cls)
+
+        new_data = copy.copy(self.data)
+
+        for attr, value in vars(self).items():
+            if attr != "data":
+                setattr(copied_op, attr, value)
+
+        copied_op.data = new_data
+
+        return copied_op
+
     @property
     def data(self):
         """Create data property"""
diff --git a/tests/templates/test_subroutines/test_select.py b/tests/templates/test_subroutines/test_select.py
index 596376d9efd..db3c7b850c7 100644
--- a/tests/templates/test_subroutines/test_select.py
+++ b/tests/templates/test_subroutines/test_select.py
@@ -15,6 +15,7 @@
 Tests for the Select template.
 """
 # pylint: disable=protected-access,too-many-arguments,import-outside-toplevel, no-self-use
+import copy
 import pytest
 import numpy as np
 from pennylane import numpy as pnp
@@ -206,6 +207,14 @@ def test_flatten_unflatten(self):
         assert op.target_wires == new_op.target_wires
         assert op is not new_op
 
+    def test_copy(self):
+        """Test that the copy function of Select works correctly."""
+        ops = [qml.PauliX(wires=2), qml.RX(0.2, wires=3), qml.PauliY(wires=2), qml.SWAP([2, 3])]
+        op = qml.Select(ops, control=[0, 1])
+        op_copy = copy.copy(op)
+
+        assert qml.equal(op, op_copy)
+
 
 class TestErrorMessages:
     """Test that the correct errors are raised"""

From 788fa67d20edb160e077b4306c4cd1a18b05bb5d Mon Sep 17 00:00:00 2001
From: Matthew Silverman <matthews@xanadu.ai>
Date: Thu, 14 Sep 2023 17:30:54 -0400
Subject: [PATCH 055/131] Various changes for DQ2 to work (#4534)

**Context:**
In #4436, I'm working on a mega-fix PR with far too many things to be
reviewed at once. That said, CI is now passing there with most tests
using the new device. I've copied over all changes to `pennylane/` to
this fresh PR, and if CI passes here then we should be well on our way
to swap out the device.

**Description of the Change:**
A lot of little things, but I'll try to sum them up here:
- if a device doesn't provide a `pennylane_requires`, assume it is
always supported
- use observable data (as well as the state) to determine if we should
use "backprop mode". If you have a circuit with non-trainable ops, but
the observable has trainable data, the state will still be a
non-trainable dtype at that point, and we might wrongly do CSR
dot-product instead of sum of terms.
- don't support trainable Pow ops because they don't work in backprop
(same as DQL)
- Squeeze all `CountsMP` results, not just dict results, to support
batched CountsMP results
- Use `device.preprocess` instead of changing the expansion strategy for
DQ2 in the `draw` module
- When sampling, only squeeze the last dimension with Pauli words to
properly support `1xN` ShotCopies
- don't validate jitted DiagonalQubitUnitary
- `metric_tensor` has better interface handling so it can be multiplied
by vanilla numpy arrays without breaking (needed for DQ2)
- QCut supports the new device interface

**Benefits:**
This PR is less than 200 lines and contains a significant portion of the
little things needed to make DQ2 work as the new default.qubit

**Possible Drawbacks:**
Lots of unrelated changes, might be tough to review.
---
 pennylane/__init__.py                         |  4 +-
 pennylane/devices/qubit/measure.py            |  2 +-
 pennylane/devices/qubit/preprocess.py         |  7 +-
 pennylane/drawer/draw.py                      | 56 +++++++----
 pennylane/gradients/spsa_gradient.py          |  2 +-
 pennylane/measurements/sample.py              |  4 +-
 pennylane/ops/qubit/matrix_ops.py             |  4 +-
 pennylane/ops/qubit/state_preparation.py      |  2 +-
 pennylane/qinfo/transforms.py                 | 12 ++-
 pennylane/qnode.py                            |  4 +-
 pennylane/tape/qscript.py                     |  6 +-
 pennylane/transforms/metric_tensor.py         | 17 ++--
 pennylane/transforms/qcut/cutstrategy.py      |  6 +-
 pennylane/transforms/qcut/montecarlo.py       |  6 +-
 pennylane/transforms/specs.py                 |  8 +-
 tests/devices/qubit/test_preprocess.py        |  2 +
 tests/drawer/test_draw.py                     |  8 +-
 tests/drawer/test_draw_mpl.py                 |  8 +-
 .../test_jax_jit_qnode_default_qubit_2.py     |  4 +-
 tests/ops/op_math/test_sum.py                 |  7 +-
 tests/transforms/test_qcut.py                 | 93 ++++++++++---------
 21 files changed, 156 insertions(+), 106 deletions(-)

diff --git a/pennylane/__init__.py b/pennylane/__init__.py
index 41398739e42..4b33c0cb011 100644
--- a/pennylane/__init__.py
+++ b/pennylane/__init__.py
@@ -334,7 +334,9 @@ def run_cnot():
         # loads the device class
         plugin_device_class = plugin_devices[name].load()
 
-        if Version(version()) not in SimpleSpec(plugin_device_class.pennylane_requires):
+        if hasattr(plugin_device_class, "pennylane_requires") and Version(
+            version()
+        ) not in SimpleSpec(plugin_device_class.pennylane_requires):
             raise DeviceError(
                 f"The {name} plugin requires PennyLane versions {plugin_device_class.pennylane_requires}, "
                 f"however PennyLane version {__version__} is installed."
diff --git a/pennylane/devices/qubit/measure.py b/pennylane/devices/qubit/measure.py
index fef0ac5fd4c..65aed196dd8 100644
--- a/pennylane/devices/qubit/measure.py
+++ b/pennylane/devices/qubit/measure.py
@@ -142,7 +142,7 @@ def get_measurement_function(
             if measurementprocess.obs.name == "SparseHamiltonian":
                 return csr_dot_products
 
-            backprop_mode = math.get_interface(state) != "numpy"
+            backprop_mode = math.get_interface(state, *measurementprocess.obs.data) != "numpy"
             if isinstance(measurementprocess.obs, Hamiltonian):
                 # need to work out thresholds for when its faster to use "backprop mode" measurements
                 return sum_of_terms_method if backprop_mode else csr_dot_products
diff --git a/pennylane/devices/qubit/preprocess.py b/pennylane/devices/qubit/preprocess.py
index 1fdd4404ef2..c1be5a802d3 100644
--- a/pennylane/devices/qubit/preprocess.py
+++ b/pennylane/devices/qubit/preprocess.py
@@ -70,7 +70,12 @@ def _accepted_operator(op: qml.operation.Operator) -> bool:
         return False
     if op.name == "GroverOperator" and len(op.wires) >= 13:
         return False
-    return op.name == "Snapshot" or op.has_matrix
+    if op.name == "Snapshot":
+        return True
+    if op.__class__.__name__ == "Pow" and qml.operation.is_trainable(op):
+        return False
+
+    return op.has_matrix
 
 
 def _accepted_adjoint_operator(op: qml.operation.Operator) -> bool:
diff --git a/pennylane/drawer/draw.py b/pennylane/drawer/draw.py
index 13ce6c2b544..560b6e4df5b 100644
--- a/pennylane/drawer/draw.py
+++ b/pennylane/drawer/draw.py
@@ -240,15 +240,23 @@ def _draw_qnode(
 ):
     @wraps(qnode)
     def wrapper(*args, **kwargs):
-        original_expansion_strategy = getattr(qnode, "expansion_strategy", None)
-
-        try:
-            qnode.expansion_strategy = expansion_strategy or original_expansion_strategy
-            tapes = qnode.construct(args, kwargs)
-        finally:
-            qnode.expansion_strategy = original_expansion_strategy
-
-        _wire_order = wire_order or getattr(qnode.device, "wires", None)
+        if expansion_strategy == "device" and isinstance(
+            qnode.device, qml.devices.experimental.Device
+        ):
+            qnode.construct(args, kwargs)
+            program, _ = qnode.device.preprocess()
+            tapes = program([qnode.tape])
+            _wire_order = wire_order or qnode.tape.wires
+        else:
+            original_expansion_strategy = getattr(qnode, "expansion_strategy", None)
+
+            try:
+                qnode.expansion_strategy = expansion_strategy or original_expansion_strategy
+                tapes = qnode.construct(args, kwargs)
+            finally:
+                qnode.expansion_strategy = original_expansion_strategy
+
+            _wire_order = wire_order or qnode.device.wires
 
         if tapes is not None:
             cache = {"tape_offset": 0, "matrices": []}
@@ -535,18 +543,28 @@ def _draw_mpl_qnode(
 ):
     @wraps(qnode)
     def wrapper(*args, **kwargs_qnode):
-        original_expansion_strategy = getattr(qnode, "expansion_strategy", None)
-
-        try:
-            qnode.expansion_strategy = expansion_strategy or original_expansion_strategy
-            qnode.construct(args, kwargs_qnode)
-        finally:
-            qnode.expansion_strategy = original_expansion_strategy
-
-        _wire_order = wire_order or getattr(qnode.device, "wires", None)
+        if expansion_strategy == "device" and isinstance(
+            qnode.device, qml.devices.experimental.Device
+        ):
+            qnode.construct(args, kwargs)
+            program, _ = qnode.device.preprocess()
+            tapes, _ = program([qnode.tape])
+            tape = tapes[0]
+            _wire_order = wire_order or qnode.tape.wires
+        else:
+            original_expansion_strategy = getattr(qnode, "expansion_strategy", None)
+
+            try:
+                qnode.expansion_strategy = expansion_strategy or original_expansion_strategy
+                qnode.construct(args, kwargs_qnode)
+            finally:
+                qnode.expansion_strategy = original_expansion_strategy
+
+            tape = qnode.tape
+            _wire_order = wire_order or qnode.device.wires
 
         return tape_mpl(
-            qnode.qtape,
+            tape,
             wire_order=_wire_order,
             show_all_wires=show_all_wires,
             decimals=decimals,
diff --git a/pennylane/gradients/spsa_gradient.py b/pennylane/gradients/spsa_gradient.py
index 9aca9706b6a..70239178d1c 100644
--- a/pennylane/gradients/spsa_gradient.py
+++ b/pennylane/gradients/spsa_gradient.py
@@ -314,7 +314,7 @@ def _single_shot_batch_result(results):
         if num_measurements == 1:
             grads = 0
             for rep, _coeffs in enumerate(all_coeffs):
-                res = results[rep * tapes_per_grad : (rep + 1) * tapes_per_grad]
+                res = list(results[rep * tapes_per_grad : (rep + 1) * tapes_per_grad])
                 if r0 is not None:
                     res.insert(0, r0)
                 res = qml.math.stack(res)
diff --git a/pennylane/measurements/sample.py b/pennylane/measurements/sample.py
index e0db9ed186d..dfaeb3ed805 100644
--- a/pennylane/measurements/sample.py
+++ b/pennylane/measurements/sample.py
@@ -213,9 +213,7 @@ def process_samples(
 
         if str(name) in {"PauliX", "PauliY", "PauliZ", "Hadamard"}:
             # Process samples for observables with eigenvalues {1, -1}
-            samples = 1 - 2 * qml.math.squeeze(samples)
-            if samples.shape == ():
-                samples = samples.reshape((1,))
+            samples = 1 - 2 * qml.math.squeeze(samples, axis=-1)
         else:
             # Replace the basis state in the computational basis with the correct eigenvalue.
             # Extract only the columns of the basis samples required based on ``wires``.
diff --git a/pennylane/ops/qubit/matrix_ops.py b/pennylane/ops/qubit/matrix_ops.py
index 878f5d98958..eca006e0529 100644
--- a/pennylane/ops/qubit/matrix_ops.py
+++ b/pennylane/ops/qubit/matrix_ops.py
@@ -293,7 +293,9 @@ def compute_matrix(D):  # pylint: disable=arguments-differ
         """
         D = qml.math.asarray(D)
 
-        if not qml.math.allclose(D * qml.math.conj(D), qml.math.ones_like(D)):
+        if not qml.math.is_abstract(D) and not qml.math.allclose(
+            D * qml.math.conj(D), qml.math.ones_like(D)
+        ):
             raise ValueError("Operator must be unitary.")
 
         # The diagonal is supposed to have one-dimension. If it is broadcasted, it has two
diff --git a/pennylane/ops/qubit/state_preparation.py b/pennylane/ops/qubit/state_preparation.py
index 1b6839796ee..8e904f98960 100644
--- a/pennylane/ops/qubit/state_preparation.py
+++ b/pennylane/ops/qubit/state_preparation.py
@@ -214,7 +214,7 @@ def state_vector(self, wire_order=None):
 
         wire_order = Wires(wire_order)
         if not wire_order.contains_wires(self.wires):
-            raise WireError("Custom wire_order must contain all StatePrep wires")
+            raise WireError(f"Custom wire_order must contain all {self.name} wires")
 
         num_total_wires = len(wire_order)
         indices = tuple(
diff --git a/pennylane/qinfo/transforms.py b/pennylane/qinfo/transforms.py
index 3c7adc577c8..02cb95df2a7 100644
--- a/pennylane/qinfo/transforms.py
+++ b/pennylane/qinfo/transforms.py
@@ -71,7 +71,7 @@ def circuit(x):
     def processing_fn(res):
         # device is provided by the custom QNode transform
         device = kwargs.get("device", None)
-        c_dtype = device.C_DTYPE if device else "complex128"
+        c_dtype = getattr(device, "C_DTYPE", "complex128")
 
         # determine the density matrix
         dm_func = (
@@ -169,7 +169,7 @@ def circuit(x):
     def processing_fn(res):
         # device is provided by the custom QNode transform
         device = kwargs.get("device", None)
-        c_dtype = device.C_DTYPE if device else "complex128"
+        c_dtype = getattr(device, "C_DTYPE", "complex128")
 
         # determine the density matrix
         density_matrix = (
@@ -256,7 +256,7 @@ def circuit(x):
     def processing_fn(res):
         # device is provided by the custom QNode transform
         device = kwargs.get("device", None)
-        c_dtype = device.C_DTYPE if device else "complex128"
+        c_dtype = getattr(device, "C_DTYPE", "complex128")
 
         # determine if the measurement is a state vector or a density matrix
         if not isinstance(measurements[0], DensityMatrixMP) and not isinstance(
@@ -362,7 +362,7 @@ def circuit(x):
     def processing_fn(res):
         # device is provided by the custom QNode transform
         device = kwargs.get("device", None)
-        c_dtype = device.C_DTYPE if device else "complex128"
+        c_dtype = getattr(device, "C_DTYPE", "complex128")
 
         density_matrix = (
             res[0]
@@ -731,7 +731,9 @@ def circ(params):
 
     """
 
-    if qnode.device.shots is not None and isinstance(qnode.device, DefaultQubit):
+    if qnode.device.shots and isinstance(
+        qnode.device, (DefaultQubit, qml.devices.experimental.DefaultQubit2)
+    ):
 
         def wrapper(*args0, **kwargs0):
             return 4 * metric_tensor(qnode, *args, **kwargs)(*args0, **kwargs0)
diff --git a/pennylane/qnode.py b/pennylane/qnode.py
index ef88013e822..6793cc2f8f5 100644
--- a/pennylane/qnode.py
+++ b/pennylane/qnode.py
@@ -907,7 +907,9 @@ def construct(self, args, kwargs):  # pylint: disable=too-many-branches
             tapes, _ = qml.defer_measurements(self._tape)
             self._tape = tapes[0]
 
-        if self.expansion_strategy == "device":
+        if self.expansion_strategy == "device" and not isinstance(
+            self.device, qml.devices.experimental.Device
+        ):
             self._tape = self.device.expand_fn(self.tape, max_expansion=self.max_expansion)
 
         # If the gradient function is a transform, expand the tape so that
diff --git a/pennylane/tape/qscript.py b/pennylane/tape/qscript.py
index fe4cca4f163..d4515c752ba 100644
--- a/pennylane/tape/qscript.py
+++ b/pennylane/tape/qscript.py
@@ -799,11 +799,9 @@ def shape(self, device):
         if isinstance(device, qml.devices.experimental.Device):
             # MP.shape (called below) takes 2 arguments: `device` and `shots`.
             # With the new device interface, shots are stored on tapes rather than the device
-            # As well, MP.shape needs the device largely to see the device wires, and this is
-            # also stored on tapes in the new device interface. TODO: refactor MP.shape to accept
-            # `wires` instead of device (not currently done because probs.shape uses device.cutoff)
+            # TODO: refactor MP.shape to accept `wires` instead of device (not currently done
+            # because probs.shape uses device.cutoff)
             shots = self.shots
-            device = self
         else:
             shots = (
                 Shots(device._raw_shot_sequence)
diff --git a/pennylane/transforms/metric_tensor.py b/pennylane/transforms/metric_tensor.py
index c1b92d4ea55..639f28c94dc 100644
--- a/pennylane/transforms/metric_tensor.py
+++ b/pennylane/transforms/metric_tensor.py
@@ -17,9 +17,9 @@
 """
 import functools
 import warnings
+import numpy as np
 
 import pennylane as qml
-from pennylane import numpy as np
 from pennylane.circuit_graph import LayerData
 from pennylane.queuing import WrappedObj
 
@@ -843,14 +843,15 @@ def processing_fn(results):
         # Rescale first and second term
         coeffs_gen = (c for c in qml.math.hstack(coeffs_list))
         # flattened coeffs_list but also with 0s where parameters are not in argnum
-        extended_coeffs_list = [
-            next(coeffs_gen) if param_in_argnum else 0.0
-            for param_in_argnum in qml.math.hstack(in_argnum_list)
-        ]
-
-        scale = qml.math.convert_like(
-            qml.math.tensordot(extended_coeffs_list, extended_coeffs_list, axes=0), results[0]
+        interface = qml.math.get_interface(*results)
+        extended_coeffs_list = qml.math.asarray(
+            [
+                next(coeffs_gen) if param_in_argnum else 0.0
+                for param_in_argnum in qml.math.hstack(in_argnum_list)
+            ],
+            like=interface,
         )
+        scale = qml.math.tensordot(extended_coeffs_list, extended_coeffs_list, axes=0)
         off_diag_mt = scale * off_diag_mt
 
         # Combine block-diagonal and off block-diagonal
diff --git a/pennylane/transforms/qcut/cutstrategy.py b/pennylane/transforms/qcut/cutstrategy.py
index fc4dd5c1e83..553defa079f 100644
--- a/pennylane/transforms/qcut/cutstrategy.py
+++ b/pennylane/transforms/qcut/cutstrategy.py
@@ -125,16 +125,16 @@ def __post_init__(
         if devices is None and self.max_free_wires is None:
             raise ValueError("One of arguments `devices` and max_free_wires` must be provided.")
 
-        if isinstance(devices, qml.Device):
+        if isinstance(devices, (qml.Device, qml.devices.experimental.Device)):
             devices = (devices,)
 
         if devices is not None:
             if not isinstance(devices, SequenceType) or any(
-                (not isinstance(d, qml.Device) for d in devices)
+                (not isinstance(d, (qml.Device, qml.devices.experimental.Device)) for d in devices)
             ):
                 raise ValueError(
                     "Argument `devices` must be a list or tuple containing elements of type "
-                    "`qml.Device`"
+                    "`qml.Device` or `qml.devices.experimental.Device`"
                 )
 
             device_wire_sizes = [len(d.wires) for d in devices]
diff --git a/pennylane/transforms/qcut/montecarlo.py b/pennylane/transforms/qcut/montecarlo.py
index 7ea1a76978e..013216339bb 100644
--- a/pennylane/transforms/qcut/montecarlo.py
+++ b/pennylane/transforms/qcut/montecarlo.py
@@ -495,11 +495,13 @@ def __call__(self, *args, **kwargs):
         shots = kwargs.pop("shots", False)
         shots = shots or self.device.shots
 
-        if shots is None:
+        if not shots:
             raise ValueError(
                 "A shots value must be provided in the device "
                 "or when calling the QNode to be cut"
             )
+        if isinstance(shots, qml.measurements.Shots):
+            shots = shots.total_shots
 
         # find the qcut transform inside the transform program and set the shots argument
         qcut_tc = [
@@ -696,7 +698,7 @@ def expand_fragment_tapes_mc(
                         meas_w = op.wires[0]
                         MC_MEASUREMENTS[meas_settings[op.id][shot]](meas_w)
 
-            frag_config.append(QuantumScript.from_queue(q))
+            frag_config.append(QuantumScript.from_queue(q, shots=1))
 
         all_configs.append(frag_config)
 
diff --git a/pennylane/transforms/specs.py b/pennylane/transforms/specs.py
index e508b808a46..da03b023827 100644
--- a/pennylane/transforms/specs.py
+++ b/pennylane/transforms/specs.py
@@ -123,8 +123,12 @@ def specs_qnode(*args, **kwargs):
 
         info = qnode.qtape.specs.copy()
 
-        info["num_device_wires"] = qnode.device.num_wires
-        info["device_name"] = qnode.device.short_name
+        info["num_device_wires"] = (
+            len(qnode.tape.wires)
+            if isinstance(qnode.device, qml.devices.experimental.Device)
+            else len(qnode.device.wires)
+        )
+        info["device_name"] = getattr(qnode.device, "short_name", qnode.device.name)
         info["expansion_strategy"] = qnode.expansion_strategy
         info["gradient_options"] = qnode.gradient_kwargs
         info["interface"] = qnode.interface
diff --git a/tests/devices/qubit/test_preprocess.py b/tests/devices/qubit/test_preprocess.py
index 01c7f528554..405657967fd 100644
--- a/tests/devices/qubit/test_preprocess.py
+++ b/tests/devices/qubit/test_preprocess.py
@@ -73,6 +73,8 @@ class TestPrivateHelpers:
             (qml.QFT(wires=range(10)), False),
             (qml.GroverOperator(wires=range(10)), True),
             (qml.GroverOperator(wires=range(14)), False),
+            (qml.pow(qml.RX(1.1, 0), 3), True),
+            (qml.pow(qml.RX(pnp.array(1.1), 0), 3), False),
         ],
     )
     def test_accepted_operator(self, op, expected):
diff --git a/tests/drawer/test_draw.py b/tests/drawer/test_draw.py
index 8529c816cba..998a907a61d 100644
--- a/tests/drawer/test_draw.py
+++ b/tests/drawer/test_draw.py
@@ -329,12 +329,16 @@ def circ():
     assert draw(circ)() == expected
 
 
-def test_expansion_strategy():
+@pytest.mark.parametrize(
+    "device",
+    [qml.device("default.qubit.legacy", wires=2), qml.devices.experimental.DefaultQubit2(wires=2)],
+)
+def test_expansion_strategy(device):
     """Test expansion strategy keyword modifies tape expansion."""
 
     H = qml.PauliX(0) + qml.PauliZ(1) + 0.5 * qml.PauliX(0) @ qml.PauliX(1)
 
-    @qml.qnode(qml.device("default.qubit", wires=2))
+    @qml.qnode(device)
     def circ(t):
         qml.ApproxTimeEvolution(H, t, 2)
         return qml.probs(wires=0)
diff --git a/tests/drawer/test_draw_mpl.py b/tests/drawer/test_draw_mpl.py
index 8168f0103c1..96ebfaf858c 100644
--- a/tests/drawer/test_draw_mpl.py
+++ b/tests/drawer/test_draw_mpl.py
@@ -70,13 +70,17 @@ def test_standard_use():
     plt.close()
 
 
+@pytest.mark.parametrize(
+    "device",
+    [qml.device("default.qubit.legacy", wires=3), qml.devices.experimental.DefaultQubit2(wires=3)],
+)
 @pytest.mark.parametrize(
     "strategy, initial_strategy, n_lines", [("gradient", "device", 3), ("device", "gradient", 13)]
 )
-def test_expansion_strategy(strategy, initial_strategy, n_lines):
+def test_expansion_strategy(device, strategy, initial_strategy, n_lines):
     """Test that the expansion strategy keyword controls what operations are drawn."""
 
-    @qml.qnode(qml.device("default.qubit", wires=3), expansion_strategy=initial_strategy)
+    @qml.qnode(device, expansion_strategy=initial_strategy)
     def circuit():
         qml.Permute([2, 0, 1], wires=(0, 1, 2))
         return qml.expval(qml.PauliZ(0))
diff --git a/tests/interfaces/default_qubit_2_integration/test_jax_jit_qnode_default_qubit_2.py b/tests/interfaces/default_qubit_2_integration/test_jax_jit_qnode_default_qubit_2.py
index 3e337200de7..e15e14bc19c 100644
--- a/tests/interfaces/default_qubit_2_integration/test_jax_jit_qnode_default_qubit_2.py
+++ b/tests/interfaces/default_qubit_2_integration/test_jax_jit_qnode_default_qubit_2.py
@@ -1223,6 +1223,8 @@ def test_state(self, dev, diff_method, grad_on_execution, interface, tol):
 
         x = jax.numpy.array(0.543)
         y = jax.numpy.array(-0.654)
+        if not dev.wires:
+            dev._wires = qml.wires.Wires([0, 1])  # pylint:disable=protected-access
 
         @qnode(
             dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution
@@ -1235,7 +1237,7 @@ def circuit(x, y):
 
         def cost_fn(x, y):
             res = circuit(x, y)
-            assert res.dtype is np.dtype("complex128")  # pylint:disable=no-member
+            assert res.dtype is np.dtype("complex128")
             probs = jax.numpy.abs(res) ** 2
             return probs[0] + probs[2]
 
diff --git a/tests/ops/op_math/test_sum.py b/tests/ops/op_math/test_sum.py
index 102021e1cd5..719bedd38e3 100644
--- a/tests/ops/op_math/test_sum.py
+++ b/tests/ops/op_math/test_sum.py
@@ -1081,10 +1081,11 @@ def test_params_can_be_considered_trainable(self):
 
         @qml.qnode(dev, interface=None)
         def circuit():
-            return qml.expval(Sum(qml.RX(1.1, 0), qml.RY(qnp.array(2.2), 0)))
+            return qml.expval(
+                Sum(qml.s_prod(1.1, qml.PauliX(0)), qml.s_prod(qnp.array(2.2), qml.PauliY(1)))
+            )
 
-        with pytest.warns(UserWarning):
-            circuit()
+        circuit()
         assert circuit.tape.trainable_params == [1]
 
 
diff --git a/tests/transforms/test_qcut.py b/tests/transforms/test_qcut.py
index 61b1ab37eb4..092924907d3 100644
--- a/tests/transforms/test_qcut.py
+++ b/tests/transforms/test_qcut.py
@@ -2457,19 +2457,22 @@ def func(x):
             )
 
 
+@pytest.mark.parametrize(
+    "dev_fn", [qml.devices.DefaultQubit, qml.devices.experimental.DefaultQubit2]
+)
 class TestCutCircuitMCTransform:
     """
     Tests that the `cut_circuit_mc` transform gives the correct results.
     """
 
     @flaky(max_runs=3)
-    def test_cut_circuit_mc_expval(self):
+    def test_cut_circuit_mc_expval(self, dev_fn):
         """
         Tests that a circuit containing sampling measurements can be cut and
         recombined to give the correct expectation value
         """
 
-        dev_sim = qml.device("default.qubit", wires=3)
+        dev_sim = dev_fn(wires=3)
 
         @qml.qnode(dev_sim)
         def target_circuit(v):
@@ -2486,7 +2489,7 @@ def target_circuit(v):
             qml.RX(2.3, wires=2)
             return qml.expval(qml.PauliZ(wires=0) @ qml.PauliZ(wires=2))
 
-        dev = qml.device("default.qubit", wires=2, shots=10000)
+        dev = dev_fn(wires=2, shots=10000)
 
         @partial(qml.cut_circuit_mc, classical_processing_fn=fn)
         @qml.qnode(dev)
@@ -2510,13 +2513,13 @@ def circuit(v):
         target = target_circuit(v)
         assert np.isclose(cut_res_mc, target, atol=0.1)  # not guaranteed to pass each time
 
-    def test_cut_circuit_mc_sample(self):
+    def test_cut_circuit_mc_sample(self, dev_fn):
         """
         Tests that a circuit containing sampling measurements can be cut and
         postprocessed to return bitstrings of the original circuit size.
         """
 
-        dev = qml.device("default.qubit", wires=3, shots=100)
+        dev = dev_fn(wires=3, shots=100)
 
         @qml.qnode(dev)
         def circuit(x):
@@ -2542,13 +2545,13 @@ def circuit(x):
         assert cut_res_bs.shape == target.shape
         assert isinstance(cut_res_bs, type(target))
 
-    def test_override_samples(self):
+    def test_override_samples(self, dev_fn):
         """
         Tests that the number of shots used on a device can be temporarily
         altered when executing the QNode
         """
         shots = 100
-        dev = qml.device("default.qubit", wires=2, shots=shots)
+        dev = dev_fn(wires=2, shots=shots)
 
         @qml.cut_circuit_mc
         @qml.qnode(dev)
@@ -2577,13 +2580,13 @@ def cut_circuit(x):
 
         assert cut_res_original.shape == (shots, 2)
 
-    def test_no_shots(self):
+    def test_no_shots(self, dev_fn):
         """
         Tests that the correct error message is given if a device is provided
         without shots
         """
 
-        dev = qml.device("default.qubit", wires=2)
+        dev = dev_fn(wires=2)
 
         @qml.cut_circuit_mc
         @qml.qnode(dev)
@@ -2605,13 +2608,13 @@ def cut_circuit(x):
         with pytest.raises(ValueError, match="A shots value must be provided in the device "):
             cut_circuit(v)
 
-    def test_sample_obs_error(self):
+    def test_sample_obs_error(self, dev_fn):
         """
         Tests that a circuit with sample measurements containing observables
         gives the correct error
         """
         shots = 100
-        dev = qml.device("default.qubit", wires=2, shots=shots)
+        dev = dev_fn(wires=2, shots=shots)
 
         @qml.cut_circuit_mc
         @qml.qnode(dev)
@@ -2633,13 +2636,13 @@ def cut_circuit(x):
         with pytest.raises(ValueError, match="The Monte Carlo circuit cutting workflow only "):
             cut_circuit(v)
 
-    def test_transform_shots_error(self):
+    def test_transform_shots_error(self, dev_fn):
         """
         Tests that the correct error is given when a `shots` argument is passed
         when transforming a qnode
         """
 
-        dev = qml.device("default.qubit", wires=2)
+        dev = dev_fn(wires=2)
 
         with pytest.raises(
             ValueError, match="Cannot provide a 'shots' value directly to the cut_circuit_mc "
@@ -2661,12 +2664,12 @@ def cut_circuit(x):  # pylint: disable=unused-variable,unused-argument
                 qml.RX(2.3, wires=2)
                 return qml.sample(wires=[0, 2])
 
-    def test_multiple_meas_error(self):
+    def test_multiple_meas_error(self, dev_fn):
         """
         Tests that attempting to cut a circuit with multiple sample measurements
         using `cut_circuit_mc` gives the correct error
         """
-        dev = qml.device("default.qubit", wires=3, shots=100)
+        dev = dev_fn(wires=3, shots=100)
 
         @qml.cut_circuit_mc
         @qml.qnode(dev)
@@ -2690,12 +2693,12 @@ def cut_circuit(x):
         ):
             cut_circuit(v)
 
-    def test_non_sample_meas_error(self):
+    def test_non_sample_meas_error(self, dev_fn):
         """
         Tests that attempting to cut a circuit with non-sample measurements
         using `cut_circuit_mc` gives the correct error
         """
-        dev = qml.device("default.qubit", wires=2, shots=100)
+        dev = dev_fn(wires=2, shots=100)
 
         @qml.cut_circuit_mc
         @qml.qnode(dev)
@@ -2719,13 +2722,13 @@ def cut_circuit(x):
         ):
             cut_circuit(v)
 
-    def test_qnode_shots_arg_error(self):
+    def test_qnode_shots_arg_error(self, dev_fn):
         """
         Tests that if a shots argument is passed directly to the qnode when using
         `cut_circuit_mc` the correct error is given
         """
         shots = 100
-        dev = qml.device("default.qubit", wires=2, shots=shots)
+        dev = dev_fn(wires=2, shots=shots)
 
         with pytest.raises(
             ValueError,
@@ -2737,13 +2740,13 @@ def test_qnode_shots_arg_error(self):
             def cut_circuit(x, shots=shots):  # pylint: disable=unused-variable,unused-argument
                 return qml.sample(wires=[0, 2])
 
-    def test_no_interface(self):
+    def test_no_interface(self, dev_fn):
         """
         Tests that if no interface is provided when using `cut_circuit_mc` the
         correct output is given
         """
         shots = 100
-        dev = qml.device("default.qubit", wires=2, shots=shots)
+        dev = dev_fn(wires=2, shots=shots)
 
         @qml.cut_circuit_mc
         @qml.qnode(dev, interface=None)
@@ -2767,7 +2770,7 @@ def cut_circuit(x):
         assert isinstance(res, np.ndarray)
 
     @pytest.mark.autograd
-    def test_samples_autograd(self):
+    def test_samples_autograd(self, dev_fn):
         """
         Tests that `cut_circuit_mc` returns the correct type of sample
         output value in autograd
@@ -2775,7 +2778,7 @@ def test_samples_autograd(self):
         import autograd
 
         shots = 10
-        dev = qml.device("default.qubit", wires=2, shots=shots)
+        dev = dev_fn(wires=2, shots=shots)
 
         @qml.cut_circuit_mc
         @qml.qnode(dev)
@@ -2802,7 +2805,7 @@ def cut_circuit(x):
         assert isinstance(res, type(convert_input))
 
     @pytest.mark.tf
-    def test_samples_tf(self):
+    def test_samples_tf(self, dev_fn):
         """
         Tests that `cut_circuit_mc` returns the correct type of sample
         output value in tf
@@ -2810,7 +2813,7 @@ def test_samples_tf(self):
         import tensorflow as tf
 
         shots = 10
-        dev = qml.device("default.qubit", wires=2, shots=shots)
+        dev = dev_fn(wires=2, shots=shots)
 
         @qml.cut_circuit_mc
         @qml.qnode(dev)
@@ -2837,7 +2840,7 @@ def cut_circuit(x):
         assert isinstance(res, type(convert_input))
 
     @pytest.mark.torch
-    def test_samples_torch(self):
+    def test_samples_torch(self, dev_fn):
         """
         Tests that `cut_circuit_mc` returns the correct type of sample
         output value in torch
@@ -2845,7 +2848,7 @@ def test_samples_torch(self):
         import torch
 
         shots = 10
-        dev = qml.device("default.qubit", wires=2, shots=shots)
+        dev = dev_fn(wires=2, shots=shots)
 
         @qml.cut_circuit_mc
         @qml.qnode(dev)
@@ -2872,7 +2875,7 @@ def cut_circuit(x):
         assert isinstance(res, type(convert_input))
 
     @pytest.mark.jax
-    def test_samples_jax(self):
+    def test_samples_jax(self, dev_fn):
         """
         Tests that `cut_circuit_mc` returns the correct type of sample
         output value in Jax
@@ -2880,7 +2883,7 @@ def test_samples_jax(self):
         import jax
 
         shots = 10
-        dev = qml.device("default.qubit", wires=2, shots=shots)
+        dev = dev_fn(wires=2, shots=shots)
 
         @qml.cut_circuit_mc
         @qml.qnode(dev)
@@ -2907,7 +2910,7 @@ def cut_circuit(x):
         assert isinstance(res, type(convert_input))
 
     @pytest.mark.autograd
-    def test_mc_autograd(self):
+    def test_mc_autograd(self, dev_fn):
         """
         Tests that `cut_circuit_mc` returns the correct type of expectation
         value output in autograd
@@ -2915,7 +2918,7 @@ def test_mc_autograd(self):
         import autograd
 
         shots = 10
-        dev = qml.device("default.qubit", wires=2, shots=shots)
+        dev = dev_fn(wires=2, shots=shots)
 
         @partial(qml.cut_circuit_mc, classical_processing_fn=fn)
         @qml.qnode(dev)
@@ -2940,7 +2943,7 @@ def cut_circuit(x):
         assert isinstance(res, type(convert_input))
 
     @pytest.mark.tf
-    def test_mc_tf(self):
+    def test_mc_tf(self, dev_fn):
         """
         Tests that `cut_circuit_mc` returns the correct type of expectation
         value output in tf
@@ -2948,7 +2951,7 @@ def test_mc_tf(self):
         import tensorflow as tf
 
         shots = 10
-        dev = qml.device("default.qubit", wires=2, shots=shots)
+        dev = dev_fn(wires=2, shots=shots)
 
         @partial(qml.cut_circuit_mc, classical_processing_fn=fn)
         @qml.qnode(dev)
@@ -2973,7 +2976,7 @@ def cut_circuit(x):
         assert isinstance(res, type(convert_input))
 
     @pytest.mark.torch
-    def test_mc_torch(self):
+    def test_mc_torch(self, dev_fn):
         """
         Tests that `cut_circuit_mc` returns the correct type of expectation
         value output in torch
@@ -2981,7 +2984,7 @@ def test_mc_torch(self):
         import torch
 
         shots = 10
-        dev = qml.device("default.qubit", wires=2, shots=shots)
+        dev = dev_fn(wires=2, shots=shots)
 
         @partial(qml.cut_circuit_mc, classical_processing_fn=fn)
         @qml.qnode(dev)
@@ -3006,7 +3009,7 @@ def cut_circuit(x):
         assert isinstance(res, type(convert_input))
 
     @pytest.mark.jax
-    def test_mc_jax(self):
+    def test_mc_jax(self, dev_fn):
         """
         Tests that `cut_circuit_mc` returns the correct type of expectation
         value output in Jax
@@ -3014,7 +3017,7 @@ def test_mc_jax(self):
         import jax
 
         shots = 10
-        dev = qml.device("default.qubit", wires=2, shots=shots)
+        dev = dev_fn(wires=2, shots=shots)
 
         @partial(qml.cut_circuit_mc, classical_processing_fn=fn)
         @qml.qnode(dev)
@@ -3038,14 +3041,14 @@ def cut_circuit(x):
 
         assert isinstance(res, type(convert_input))
 
-    def test_mc_with_mid_circuit_measurement(self, mocker):
+    def test_mc_with_mid_circuit_measurement(self, dev_fn, mocker):
         """Tests the full sample-based circuit cutting pipeline successfully returns a
         single value for a circuit that contains mid-circuit
         measurements and terminal sample measurements using the `cut_circuit_mc`
         transform."""
 
         shots = 10
-        dev = qml.device("default.qubit", wires=3, shots=shots)
+        dev = dev_fn(wires=3, shots=shots)
 
         @partial(qml.cut_circuit_mc, classical_processing_fn=fn)
         @qml.qnode(dev)
@@ -3066,12 +3069,12 @@ def circuit(x):
         spy.assert_called_once()
         assert res.size == 1
 
-    def test_mc_circuit_with_disconnected_components(self):
+    def test_mc_circuit_with_disconnected_components(self, dev_fn):
         """Tests if a sample-based circuit that is fragmented into subcircuits such
         that some of the subcircuits are disconnected from the final terminal sample
         measurements is executed successfully"""
         shots = 10
-        dev = qml.device("default.qubit", wires=3, shots=shots)
+        dev = dev_fn(wires=3, shots=shots)
 
         @partial(qml.cut_circuit_mc, classical_processing_fn=fn)
         @qml.qnode(dev)
@@ -3087,11 +3090,11 @@ def circuit(x):
         res = circuit(x)
         assert res.size == 1
 
-    def test_mc_circuit_with_trivial_wire_cut(self, mocker):
+    def test_mc_circuit_with_trivial_wire_cut(self, dev_fn, mocker):
         """Tests that a sample-based circuit with a trivial wire cut (not
         separating the circuit into fragments) is executed successfully"""
         shots = 10
-        dev = qml.device("default.qubit", wires=2, shots=shots)
+        dev = dev_fn(wires=2, shots=shots)
 
         @partial(qml.cut_circuit_mc, classical_processing_fn=fn)
         @qml.qnode(dev)
@@ -3110,7 +3113,7 @@ def circuit(x):
         assert len(spy.call_args[0][0]) == shots
         assert len(spy.call_args[0]) == 1
 
-    def test_mc_complicated_circuit(self, mocker):
+    def test_mc_complicated_circuit(self, dev_fn, mocker):
         """
         Tests that the full sample-based circuit cutting pipeline successfully returns a
         value for a complex circuit with multiple wire cut scenarios. The circuit is cut into
@@ -3125,7 +3128,7 @@ def test_mc_complicated_circuit(self, mocker):
 
         # We need a 4-qubit device to account for mid-circuit measurements
         shots = 10
-        dev = qml.device("default.qubit", wires=4, shots=shots)
+        dev = dev_fn(wires=4, shots=shots)
 
         def two_qubit_unitary(param, wires):
             qml.Hadamard(wires=[wires[0]])

From 4ef39139732750dd2ef7f76025b1ae00a0da0d2b Mon Sep 17 00:00:00 2001
From: Mudit Pandey <mudit.pandey@xanadu.ai>
Date: Thu, 14 Sep 2023 18:36:59 -0400
Subject: [PATCH 056/131] Update `op_transforms` (#4573)

**Context:**
Updating `op_transform`s to use new `transform` decorator.

**Description of the Change:**
* Added `last_transform` kwarg to `TransformDispatcher` and
`TransformContainer` to differentiate between informative transforms
(which do not require execution), and any other transforms that should
be at the end of the transform program but require execution.
* `TransformDispatcher` now returns the processed results for
informative transforms rather than the transformed tapes and processing
function.

TODO:

Update op transforms:
* [x] eigvals
* [x] matrix
* [x] to_zx
* [x] generator

**Benefits:**

**Possible Drawbacks:**

**Related GitHub Issues:**

---------

Co-authored-by: rmoyard <rmoyard@gmail.com>
Co-authored-by: Edward Jiang <34989448+eddddddy@users.noreply.github.com>
Co-authored-by: Edward Jiang <eddy3.14159@gmail.com>
Co-authored-by: soranjh <40344468+soranjh@users.noreply.github.com>
Co-authored-by: Matthew Silverman <matthews@xanadu.ai>
Co-authored-by: Utkarsh <utkarshazad98@gmail.com>
Co-authored-by: lillian542 <Lillian.frederiksen@xanadu.ai>
Co-authored-by: lillian542 <38584660+lillian542@users.noreply.github.com>
---
 doc/releases/changelog-dev.md                 |   4 +
 pennylane/interfaces/execution.py             |   3 +
 pennylane/ops/functions/eigvals.py            |  92 +++++++------
 pennylane/ops/functions/generator.py          |  75 ++++++-----
 pennylane/ops/functions/matrix.py             |  83 ++++++++----
 pennylane/qinfo/transforms.py                 |  14 +-
 pennylane/transforms/core/transform.py        |  18 ++-
 .../transforms/core/transform_dispatcher.py   |  53 ++++++--
 .../transforms/core/transform_program.py      |  35 +++--
 pennylane/transforms/transpile.py             |   3 +-
 pennylane/transforms/zx/converter.py          | 127 ++++++++++--------
 tests/ops/functions/test_eigvals.py           |  11 +-
 tests/ops/functions/test_matrix.py            |  23 ++--
 .../test_transform_dispatcher.py              |  64 ++++++++-
 .../test_transform_program.py                 | 120 +++++++++++++++--
 tests/transforms/test_zx.py                   |  12 +-
 16 files changed, 520 insertions(+), 217 deletions(-)

diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 1bf1d2337df..177052d23ce 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -27,6 +27,10 @@
   (array(0.6), array([1, 1, 1, 0, 1]))
   ```
 
+* Operator transforms `qml.matrix`, `qml.eigvals`, `qml.generator`, and `qml.transforms.to_zx` are updated
+  to the new transform program system.
+  [(#4573)](https://github.com/PennyLaneAI/pennylane/pull/4573)
+
 * All quantum functions transforms are update to the new transform program system.
  [(#4439)](https://github.com/PennyLaneAI/pennylane/pull/4439)
 
diff --git a/pennylane/interfaces/execution.py b/pennylane/interfaces/execution.py
index d9b4a58f50a..d42cf584d09 100644
--- a/pennylane/interfaces/execution.py
+++ b/pennylane/interfaces/execution.py
@@ -617,6 +617,9 @@ def inner_execute_with_empty_jac(tapes, **_):
         def post_processing(results):
             return program_post_processing(program_pre_processing(results))
 
+    if transform_program.is_informative:
+        return post_processing(tapes)
+
     # Exiting early if we do not need to deal with an interface boundary
     if no_interface_boundary_required:
         results = inner_execute(tapes)
diff --git a/pennylane/ops/functions/eigvals.py b/pennylane/ops/functions/eigvals.py
index fbd002af9a0..4201e38d29e 100644
--- a/pennylane/ops/functions/eigvals.py
+++ b/pennylane/ops/functions/eigvals.py
@@ -14,18 +14,20 @@
 """
 This module contains the qml.eigvals function.
 """
+from typing import Sequence, Callable
 import warnings
 
 # pylint: disable=protected-access
-from functools import reduce
-
+from functools import reduce, partial
 import scipy
 
 import pennylane as qml
+from pennylane.transforms.op_transforms import OperationTransformError
+from pennylane.transforms.core import transform
+from pennylane.typing import TensorLike
 
 
-@qml.op_transform
-def eigvals(op, k=1, which="SA"):
+def eigvals(op: qml.operation.Operator, k=1, which="SA") -> TensorLike:
     r"""The eigenvalues of one or more operations.
 
     .. note::
@@ -39,8 +41,7 @@ def eigvals(op, k=1, which="SA"):
         `documentation <https://docs.scipy.org/doc/scipy/reference/generated/scipy.sparse.linalg.eigsh.html#scipy.sparse.linalg.eigsh>`_.
 
     Args:
-        op (.Operator, pennylane.QNode, .QuantumTape, or Callable): An operator, quantum node, tape,
-            or function that applies quantum operations.
+        op (.Operator or .QuantumTape): A quantum operator or tape.
         k (int): The number of eigenvalues to be returned for a :class:`~.SparseHamiltonian`.
         which (str): Method for computing the eigenvalues of a :class:`~.SparseHamiltonian`. The
             possible methods are ``'LM'`` (largest in magnitude), ``'SM'`` (smallest in magnitude),
@@ -48,11 +49,10 @@ def eigvals(op, k=1, which="SA"):
             from each end of the spectrum).
 
     Returns:
-        tensor_like or function: If an operator is provided as input, the eigenvalues are returned
-        directly. If a QNode or quantum function is provided as input, a function which accepts the
-        same arguments as the QNode or quantum function is returned. When called, this function will
-        return the unitary matrix in the appropriate autodiff framework (Autograd, TensorFlow,
-        PyTorch, JAX) given its parameters.
+        TensorLike or (Sequence[.QuantumTape], Callable): If an operator is provided as input, the eigenvalues
+        are returned directly. If a quantum tape is provided as input, a list of transformed tapes and a post-processing
+        function are returned. When called, this function will return the unitary matrix in the appropriate autodiff
+        framework (Autograd, TensorFlow, PyTorch, JAX) given its parameters.
 
     **Example**
 
@@ -104,6 +104,13 @@ def circuit(theta):
         array([ 0.92387953+0.38268343j,  0.92387953-0.38268343j,
                -0.92387953+0.38268343j, -0.92387953-0.38268343j])
     """
+    if not isinstance(op, qml.operation.Operator):
+        if not isinstance(op, (qml.tape.QuantumScript, qml.QNode)) and not callable(op):
+            raise OperationTransformError(
+                "Input is not an Operator, tape, QNode, or quantum function"
+            )
+        return _eigvals_tranform(op, k=k, which=which)
+
     if isinstance(op, qml.Hamiltonian):
         warnings.warn(
             "For Hamiltonians, the eigenvalues will be computed numerically. "
@@ -120,30 +127,39 @@ def circuit(theta):
         return qml.math.linalg.eigvalsh(sparse_matrix.toarray())
 
     # TODO: make `eigvals` take a `wire_order` argument to mimic `matrix`
-    return op.eigvals()
-
-
-@eigvals.tape_transform
-def _eigvals(tape):
-    op_wires = [op.wires for op in tape.operations]
-    all_wires = qml.wires.Wires.all_wires(op_wires).tolist()
-    unique_wires = qml.wires.Wires.unique_wires(op_wires).tolist()
-
-    if len(all_wires) != len(unique_wires):
-        warnings.warn(
-            "For multiple operations, the eigenvalues will be computed numerically. "
-            "This may be computationally intensive for a large number of wires.",
-            UserWarning,
-        )
-        return qml.math.linalg.eigvals(qml.matrix(tape))
-
-    # TODO: take into account wire ordering, by reordering eigenvalues
-    # as per operator wires/wire ordering, and by inserting implicit identity
-    # matrices (eigenvalues [1, 1]) at missing locations.
-
-    ev = [eigvals(op) for op in tape.operations]
-
-    if len(ev) == 1:
-        return ev[0]
-
-    return reduce(qml.math.kron, ev)
+    try:
+        return op.eigvals()
+    except qml.operation.EigvalsUndefinedError:
+        return eigvals(op.expand(), k=k, which=which)
+
+
+@partial(transform, is_informative=True)
+def _eigvals_tranform(
+    tape: qml.tape.QuantumTape, k=1, which="SA"
+) -> (Sequence[qml.tape.QuantumTape], Callable):
+    def processing_fn(res):
+        op_wires = [op.wires for op in res[0].operations]
+        all_wires = qml.wires.Wires.all_wires(op_wires).tolist()
+        unique_wires = qml.wires.Wires.unique_wires(op_wires).tolist()
+
+        if len(all_wires) != len(unique_wires):
+            warnings.warn(
+                "For multiple operations, the eigenvalues will be computed numerically. "
+                "This may be computationally intensive for a large number of wires.",
+                UserWarning,
+            )
+            matrix = qml.matrix(res[0])
+            return qml.math.linalg.eigvals(matrix)
+
+        # TODO: take into account wire ordering, by reordering eigenvalues
+        # as per operator wires/wire ordering, and by inserting implicit identity
+        # matrices (eigenvalues [1, 1]) at missing locations.
+
+        ev = [eigvals(op, k=k, which=which) for op in res[0].operations]
+
+        if len(ev) == 1:
+            return ev[0]
+
+        return reduce(qml.math.kron, ev)
+
+    return [tape], processing_fn
diff --git a/pennylane/ops/functions/generator.py b/pennylane/ops/functions/generator.py
index 5b3a4cb1a9e..f98569ffd7d 100644
--- a/pennylane/ops/functions/generator.py
+++ b/pennylane/ops/functions/generator.py
@@ -101,8 +101,7 @@ def _generator_backcompatibility(op):
     raise qml.operation.GeneratorUndefinedError
 
 
-@qml.op_transform
-def generator(op, format="prefactor"):
+def generator(op: qml.operation.Operator, format="prefactor"):
     r"""Returns the generator of an operation.
 
     Args:
@@ -167,34 +166,48 @@ def generator(op, format="prefactor"):
     >>> qml.generator(op, format="arithmetic")  # output is an instance of `SProd`
     -0.5*(PauliX(wires=[0]))
     """
-    if op.num_params != 1:
-        raise ValueError(f"Operation {op.name} is not written in terms of a single parameter")
-
-    try:
-        gen = op.generator()
-    except TypeError:
-        # For backwards compatibility with PennyLane
-        # versions <=0.22, assume op.generator is a property
-        gen = _generator_backcompatibility(op)
-
-    if not gen.is_hermitian:
-        raise qml.QuantumFunctionError(
-            f"Generator {gen.name} of operation {op.name} is not hermitian"
-        )
-
-    if format == "prefactor":
-        return _generator_prefactor(gen)
-
-    if format == "hamiltonian":
-        return _generator_hamiltonian(gen, op)
 
-    if format == "arithmetic":
-        h = _generator_hamiltonian(gen, op)
-        return qml.dot(h.coeffs, h.ops)
-
-    if format == "observable":
-        return gen
+    def processing_fn(*args, **kwargs):
+        if callable(op):
+            with qml.queuing.QueuingManager.stop_recording():
+                gen_op = op(*args, **kwargs)
+        else:
+            gen_op = op
+
+        if gen_op.num_params != 1:
+            raise ValueError(
+                f"Operation {gen_op.name} is not written in terms of a single parameter"
+            )
+
+        try:
+            gen = gen_op.generator()
+        except TypeError:
+            # For backwards compatibility with PennyLane
+            # versions <=0.22, assume gen_op.generator is a property
+            gen = _generator_backcompatibility(gen_op)
+
+        if not gen.is_hermitian:
+            raise qml.QuantumFunctionError(
+                f"Generator {gen.name} of operation {gen_op.name} is not hermitian"
+            )
+
+        if format == "prefactor":
+            return _generator_prefactor(gen)
+
+        if format == "hamiltonian":
+            return _generator_hamiltonian(gen, gen_op)
+
+        if format == "arithmetic":
+            h = _generator_hamiltonian(gen, gen_op)
+            return qml.dot(h.coeffs, h.ops)
+
+        if format == "observable":
+            return gen
+
+        raise ValueError(
+            "format must be one of ('prefactor', 'hamiltonian', 'observable', 'arithmetic')"
+        )
 
-    raise ValueError(
-        "format must be one of ('prefactor', 'hamiltonian', 'observable', 'arithmetic')"
-    )
+    if callable(op):
+        return processing_fn
+    return processing_fn()
diff --git a/pennylane/ops/functions/matrix.py b/pennylane/ops/functions/matrix.py
index cf03a59399a..28c5f3e2a39 100644
--- a/pennylane/ops/functions/matrix.py
+++ b/pennylane/ops/functions/matrix.py
@@ -15,22 +15,27 @@
 This module contains the qml.matrix function.
 """
 # pylint: disable=protected-access
+from typing import Sequence, Callable
+from functools import partial
+
 import pennylane as qml
+from pennylane.transforms.op_transforms import OperationTransformError
+from pennylane.transforms.core import transform
+from pennylane.typing import TensorLike
 
 
-@qml.op_transform
-def matrix(op, *, wire_order=None):
+def matrix(op: qml.operation.Operator, wire_order=None) -> TensorLike:
     r"""The matrix representation of an operation or quantum circuit.
 
     Args:
-        op (.Operator, pennylane.QNode, .QuantumTape, or Callable): An operator, quantum node, tape,
-            or function that applies quantum operations.
+        op (.Operator or .QuantumTape): A quantum operator or tape.
         wire_order (Sequence[Any], optional): Order of the wires in the quantum circuit.
             Defaults to the order in which the wires appear in the quantum function.
 
     Returns:
-        tensor_like or function: Function which accepts the same arguments as the QNode or quantum
-        function. When called, this function will return the unitary matrix in the appropriate
+        TensorLike or (Sequence[.QuantumTape], Callable): If an operator is provided as input, the matrix
+        is returned directly. If a quantum tape is provided, a list of transformed tapes and a post-processing
+        function are returned. When called, this function will return unitary matrix in the appropriate
         autodiff framework (Autograd, TensorFlow, PyTorch, JAX) given its parameters.
 
     **Example**
@@ -118,35 +123,65 @@ def cost(theta):
         >>> qml.grad(cost)(theta)
         -0.14943813247359922
     """
+    if not isinstance(op, qml.operation.Operator):
+        if not isinstance(op, (qml.tape.QuantumScript, qml.QNode)) and not callable(op):
+            raise OperationTransformError(
+                "Input is not an Operator, tape, QNode, or quantum function"
+            )
+        return _matrix_transform(op, wire_order=wire_order)
+
     if isinstance(op, qml.operation.Tensor) and wire_order is not None:
         op = 1.0 * op  # convert to a Hamiltonian
 
     if isinstance(op, qml.Hamiltonian):
         return op.sparse_matrix(wire_order=wire_order).toarray()
 
-    return op.matrix(wire_order=wire_order)
+    try:
+        return op.matrix(wire_order=wire_order)
+    except:  # pylint: disable=bare-except
+        return matrix(op.expand(), wire_order=wire_order)
 
 
-@matrix.tape_transform
-def _matrix(tape, wire_order=None):
-    """Defines how matrix works if applied to a tape containing multiple operations."""
+@partial(transform, is_informative=True)
+def _matrix_transform(
+    tape: qml.tape.QuantumTape, wire_order=None, **kwargs
+) -> (Sequence[qml.tape.QuantumTape], Callable):
     if not tape.wires:
         raise qml.operation.MatrixUndefinedError
-    params = tape.get_parameters(trainable_only=False)
-    interface = qml.math.get_interface(*params)
 
-    wire_order = wire_order or tape.wires
+    if wire_order and not set(tape.wires).issubset(wire_order):
+        raise OperationTransformError(
+            f"Wires in circuit {list(tape.wires)} are inconsistent with "
+            f"those in wire_order {list(wire_order)}"
+        )
+
+    wires = kwargs.get("device_wires", None) or tape.wires
+    wire_order = wire_order or wires
+
+    def processing_fn(res):
+        """Defines how matrix works if applied to a tape containing multiple operations."""
+
+        params = res[0].get_parameters(trainable_only=False)
+        interface = qml.math.get_interface(*params)
+
+        # initialize the unitary matrix
+        if len(res[0].operations) == 0:
+            result = qml.math.eye(2 ** len(wire_order), like=interface)
+        else:
+            result = matrix(res[0].operations[0], wire_order=wire_order)
+
+        for op in res[0].operations[1:]:
+            U = matrix(op, wire_order=wire_order)
+            # Coerce the matrices U and result and use matrix multiplication. Broadcasted axes
+            # are handled correctly automatically by ``matmul`` (See e.g. NumPy documentation)
+            result = qml.math.matmul(*qml.math.coerce([U, result], like=interface), like=interface)
+
+        return result
 
-    # initialize the unitary matrix
-    if len(tape.operations) == 0:
-        result = qml.math.eye(2 ** len(wire_order), like=interface)
-    else:
-        result = matrix(tape.operations[0], wire_order=wire_order)
+    return [tape], processing_fn
 
-    for op in tape.operations[1:]:
-        U = matrix(op, wire_order=wire_order)
-        # Coerce the matrices U and result and use matrix multiplication. Broadcasted axes
-        # are handled correctly automatically by ``matmul`` (See e.g. NumPy documentation)
-        result = qml.math.matmul(*qml.math.coerce([U, result], like=interface), like=interface)
 
-    return result
+@_matrix_transform.custom_qnode_transform
+def _matrix_transform_qnode(self, qnode, targs, tkwargs):
+    tkwargs.setdefault("device_wires", qnode.device.wires)
+    return self.default_qnode_transform(qnode, targs, tkwargs)
diff --git a/pennylane/qinfo/transforms.py b/pennylane/qinfo/transforms.py
index 02cb95df2a7..20665b98003 100644
--- a/pennylane/qinfo/transforms.py
+++ b/pennylane/qinfo/transforms.py
@@ -13,8 +13,8 @@
 # limitations under the License.
 """QNode transforms for the quantum information quantities."""
 # pylint: disable=import-outside-toplevel, not-callable
-import functools
-from typing import Sequence, Callable
+from functools import partial
+from typing import Callable, Sequence
 
 import pennylane as qml
 from pennylane.tape import QuantumTape
@@ -24,7 +24,7 @@
 from pennylane.transforms.core import transform
 
 
-@transform
+@partial(transform, final_transform=True)
 def reduced_dm(tape: QuantumTape, wires, **kwargs) -> (Sequence[QuantumTape], Callable):
     """Compute the reduced density matrix from a :class:`~.QNode` returning
     :func:`~pennylane.state`.
@@ -104,7 +104,7 @@ def _reduced_dm_qnode(self, qnode, targs, tkwargs):
     return self.default_qnode_transform(qnode, targs, tkwargs)
 
 
-@transform
+@partial(transform, final_transform=True)
 def purity(tape: QuantumTape, wires, **kwargs) -> (Sequence[QuantumTape], Callable):
     r"""Compute the purity of a :class:`~.QuantumTape` returning :func:`~pennylane.state`.
 
@@ -201,7 +201,7 @@ def _purity_qnode(self, qnode, targs, tkwargs):
     return self.default_qnode_transform(qnode, targs, tkwargs)
 
 
-@transform
+@partial(transform, final_transform=True)
 def vn_entropy(
     tape: QuantumTape, wires: Sequence[int], base: float = None, **kwargs
 ) -> (Sequence[QuantumTape], Callable):
@@ -295,7 +295,7 @@ def _vn_entropy_qnode(self, qnode, targs, tkwargs):
     return self.default_qnode_transform(qnode, targs, tkwargs)
 
 
-@transform
+@partial(transform, final_transform=True)
 def mutual_info(
     tape: QuantumTape, wires0: Sequence[int], wires1: Sequence[int], base: float = None, **kwargs
 ) -> (Sequence[QuantumTape], Callable):
@@ -401,7 +401,7 @@ def _torch_jac(circ):
     import torch
 
     def wrapper(*args, **kwargs):
-        loss = functools.partial(circ, **kwargs)
+        loss = partial(circ, **kwargs)
         if len(args) > 1:
             return torch.autograd.functional.jacobian(loss, args, create_graph=True)
         return torch.autograd.functional.jacobian(loss, *args, create_graph=True)
diff --git a/pennylane/transforms/core/transform.py b/pennylane/transforms/core/transform.py
index 3fd69ce8e0e..aabe860d5a4 100644
--- a/pennylane/transforms/core/transform.py
+++ b/pennylane/transforms/core/transform.py
@@ -20,13 +20,17 @@
 
 
 def transform(
-    quantum_transform, expand_transform=None, classical_cotransform=None, is_informative=None
+    quantum_transform,
+    expand_transform=None,
+    classical_cotransform=None,
+    is_informative=None,
+    final_transform=False,
 ):
     """The transform function is to be used to validate and dispatch a quantum transform on PennyLane objects (tape,
     qfunc and Qnode). It can be used directly as a decorator on qfunc and qnodes.
 
     Args:
-        quantum_transform(callable): A quantum transform is defined as a function that has the following requirements:
+        quantum_transform (callable): A quantum transform is defined as a function that has the following requirements:
 
             * A quantum transform is a function that takes a quantum tape as first input and returns a sequence of tapes
               and a processing function.
@@ -34,7 +38,7 @@ def transform(
             * The transform must have type hinting of the following form: my_quantum_transform(tape:
               qml.tape.QuantumTape, ...) -> ( Sequence[qml.tape.QuantumTape], callable)
 
-        expand_transform(callable): An expand transform is defined as a function that has the following requirements:
+        expand_transform (callable): An expand transform is defined as a function that has the following requirements:
 
             * An expand transform is a function that is applied before applying the defined quantum transform. It
               takes the same arguments as the transform and returns a single tape in a sequence with a dummy processing
@@ -42,8 +46,11 @@ def transform(
 
             * The expand transform must have the same type hinting as a quantum transform.
 
-        classical_cotransform(callable): A classical co-transform. NOT YET SUPPORTED.
-        is_informative(bool): If true the execution is skipped, because the transform is informative.
+        classical_cotransform (callable): A classical co-transform. NOT YET SUPPORTED.
+        is_informative (bool): Whether or not a transform is informative. If true the transform is queued at the end
+            of the transform program and the tapes or qnode aren't executed.
+        final_transform (bool): Whether or not the transform is terminal. If true the transform is queued at the end
+            of the transform program. ``is_informative`` supersedes ``final_transform``.
 
     **Example**
 
@@ -133,6 +140,7 @@ def qnode_circuit(a):
         expand_transform=expand_transform,
         classical_cotransform=classical_cotransform,
         is_informative=is_informative,
+        final_transform=final_transform,
     )
 
 
diff --git a/pennylane/transforms/core/transform_dispatcher.py b/pennylane/transforms/core/transform_dispatcher.py
index 5e360dfc548..b46aee18b41 100644
--- a/pennylane/transforms/core/transform_dispatcher.py
+++ b/pennylane/transforms/core/transform_dispatcher.py
@@ -39,17 +39,21 @@ class TransformDispatcher:
     .. seealso:: :func:`~.pennylane.transforms.core.transform`
     """
 
+    # pylint: disable=too-many-arguments
     def __init__(
         self,
         transform,
         expand_transform=None,
         classical_cotransform=None,
         is_informative=False,
+        final_transform=False,
     ):  # pylint:disable=redefined-outer-name
         self._transform = transform
         self._expand_transform = expand_transform
         self._classical_cotransform = classical_cotransform
         self._is_informative = is_informative
+        # is_informative supersedes final_transform
+        self._final_transform = is_informative or final_transform
 
         self._qnode_transform = self.default_qnode_transform
 
@@ -62,13 +66,14 @@ def __call__(self, *targs, **tkwargs):
             obj, *targs = targs
 
         if isinstance(obj, qml.tape.QuantumScript):
-            return self._transform(obj, *targs, **tkwargs)
+            transformed_tapes, processing_fn = self._transform(obj, *targs, **tkwargs)
+
+            if self.is_informative:
+                return processing_fn(transformed_tapes)
+            return transformed_tapes, processing_fn
+
         if isinstance(obj, qml.QNode):
-            return self._qnode_transform(
-                obj,
-                targs,
-                tkwargs,
-            )
+            return self._qnode_transform(obj, targs, tkwargs)
         if callable(obj):
             return self._qfunc_transform(obj, targs, tkwargs)
 
@@ -114,9 +119,14 @@ def classical_cotransform(self):
 
     @property
     def is_informative(self):
-        """Return True is the transform does not need to be executed."""
+        """Return True is the transform is informative."""
         return self._is_informative
 
+    @property
+    def final_transform(self):
+        """Return True if the transformed tapes must be executed."""
+        return self._final_transform
+
     def custom_qnode_transform(self, fn):
         """Register a custom QNode execution wrapper function
         for the batch transform.
@@ -168,7 +178,12 @@ def default_qnode_transform(self, qnode, targs, tkwargs):
             qnode.add_transform(TransformContainer(self._expand_transform, targs, tkwargs))
         qnode.add_transform(
             TransformContainer(
-                self._transform, targs, tkwargs, self._classical_cotransform, self._is_informative
+                self._transform,
+                targs,
+                tkwargs,
+                self._classical_cotransform,
+                self._is_informative,
+                self._final_transform,
             )
         )
         return qnode
@@ -178,7 +193,7 @@ def _qfunc_transform(self, qfunc, targs, tkwargs):
 
         def qfunc_transformed(*args, **kwargs):
             tape = qml.tape.make_qscript(qfunc)(*args, **kwargs)
-            transformed_tapes, _ = self._transform(tape, *targs, **tkwargs)
+            transformed_tapes, processing_fn = self._transform(tape, *targs, **tkwargs)
 
             if len(transformed_tapes) != 1:
                 raise TransformError(
@@ -188,6 +203,9 @@ def qfunc_transformed(*args, **kwargs):
 
             transformed_tape = transformed_tapes[0]
 
+            if self.is_informative:
+                return processing_fn(transformed_tapes)
+
             for op in transformed_tape.circuit:
                 qml.apply(op)
 
@@ -208,13 +226,20 @@ class TransformContainer:
     """
 
     def __init__(
-        self, transform, args=None, kwargs=None, classical_cotransform=None, is_informative=False
+        self,
+        transform,
+        args=None,
+        kwargs=None,
+        classical_cotransform=None,
+        is_informative=False,
+        final_transform=False,
     ):  # pylint:disable=redefined-outer-name,too-many-arguments
         self._transform = transform
         self._args = args or []
         self._kwargs = kwargs or {}
         self._classical_cotransform = classical_cotransform
         self._is_informative = is_informative
+        self._final_transform = is_informative or final_transform
 
     def __iter__(self):
         return iter(
@@ -224,6 +249,7 @@ def __iter__(self):
                 self._kwargs,
                 self._classical_cotransform,
                 self._is_informative,
+                self.final_transform,
             )
         )
 
@@ -249,5 +275,10 @@ def classical_cotransform(self):
 
     @property
     def is_informative(self):
-        """Return True is the transform does not need to be executed."""
+        """Return True is the transform is informative."""
         return self._is_informative
+
+    @property
+    def final_transform(self):
+        """Return True if the transform needs to be executed"""
+        return self._final_transform
diff --git a/pennylane/transforms/core/transform_program.py b/pennylane/transforms/core/transform_program.py
index ef181cd9573..73d716243de 100644
--- a/pennylane/transforms/core/transform_program.py
+++ b/pennylane/transforms/core/transform_program.py
@@ -131,10 +131,14 @@ def __bool__(self):
         return bool(self._transform_program)
 
     def __add__(self, other):
-        program = TransformProgram(self._transform_program)
-        for container in other:
-            program.push_back(container)
-        return program
+        if self.has_final_transform and other.has_final_transform:
+            raise TransformError("The transform program already has a terminal transform.")
+
+        transforms = self._transform_program + other._transform_program
+        if self.has_final_transform:
+            transforms.append(transforms.pop(len(self) - 1))
+
+        return TransformProgram(transforms)
 
     def __repr__(self):
         """The string representation of the transform program class."""
@@ -151,8 +155,8 @@ def push_back(self, transform_container: TransformContainer):
             raise TransformError("Only transform container can be added to the transform program.")
 
         # Program can only contain one informative transform and at the end of the program
-        if self.is_informative:
-            raise TransformError("The transform program already has an informative transform.")
+        if self.has_final_transform:
+            raise TransformError("The transform program already has a terminal transform.")
         self._transform_program.append(transform_container)
 
     def insert_front(self, transform_container: TransformContainer):
@@ -161,7 +165,7 @@ def insert_front(self, transform_container: TransformContainer):
         Args:
             transform_container(TransformContainer): A transform represented by its container.
         """
-        if transform_container.is_informative and not self.is_empty():
+        if (transform_container.final_transform) and not self.is_empty():
             raise TransformError(
                 "Informative transforms can only be added at the end of the program."
             )
@@ -193,6 +197,7 @@ def add_transform(self, transform: TransformDispatcher, *targs, **tkwargs):
                 tkwargs,
                 transform.classical_cotransform,
                 transform.is_informative,
+                transform.final_transform,
             )
         )
 
@@ -207,7 +212,7 @@ def insert_front_transform(self, transform: TransformDispatcher, *targs, **tkwar
             **tkwargs: Any additional keyword arguments that are passed to the transform.
 
         """
-        if transform.is_informative and not self.is_empty():
+        if transform.final_transform and not self.is_empty():
             raise TransformError(
                 "Informative transforms can only be added at the end of the program."
             )
@@ -219,6 +224,7 @@ def insert_front_transform(self, transform: TransformDispatcher, *targs, **tkwar
                 tkwargs,
                 transform.classical_cotransform,
                 transform.is_informative,
+                transform.final_transform,
             )
         )
 
@@ -265,16 +271,20 @@ def is_informative(self) -> bool:
         """
         return self[-1].is_informative if self else False
 
-    def __call__(self, tapes: Tuple[QuantumTape]) -> Tuple[ResultBatch, BatchPostProcessingFn]:
-        if self.is_informative:
-            raise NotImplementedError("Informative transforms are not yet supported.")
+    @property
+    def has_final_transform(self) -> bool:
+        """Check if the transform program has a terminal transform or not."""
+        return self[-1].final_transform if self else False
 
+    def __call__(self, tapes: Tuple[QuantumTape]) -> Tuple[ResultBatch, BatchPostProcessingFn]:
         if not self:
             return tapes, null_postprocessing
+
         processing_fns_stack = []
 
         for transform_container in self:
-            transform, args, kwargs, cotransform, _ = transform_container
+            transform, args, kwargs, cotransform, _, _ = transform_container
+
             if cotransform:
                 raise NotImplementedError(
                     "cotransforms are not yet integrated with TransformProgram"
@@ -305,6 +315,7 @@ def __call__(self, tapes: Tuple[QuantumTape]) -> Tuple[ResultBatch, BatchPostPro
             _apply_postprocessing_stack,
             postprocessing_stack=processing_fns_stack,
         )
+
         postprocessing_fn.__doc__ = _apply_postprocessing_stack.__doc__
 
         return tuple(tapes), postprocessing_fn
diff --git a/pennylane/transforms/transpile.py b/pennylane/transforms/transpile.py
index c8814efdb6c..e4e10ed52f8 100644
--- a/pennylane/transforms/transpile.py
+++ b/pennylane/transforms/transpile.py
@@ -5,7 +5,6 @@
 
 import networkx as nx
 
-import pennylane as qml
 from pennylane.transforms.core import transform
 from pennylane import Hamiltonian
 from pennylane.operation import Tensor
@@ -194,7 +193,7 @@ def _adjust_mmt_indices(_m, _map_wires):
 
     # change wires of observable
     if _m.obs is None:
-        return type(_m)(eigvals=qml.eigvals(_m), wires=_new_wires)
+        return type(_m)(eigvals=_m.eigvals, wires=_new_wires)
 
     _new_obs = type(_m.obs)(wires=_new_wires, id=_m.obs.id)
     return type(_m)(obs=_new_obs)
diff --git a/pennylane/transforms/zx/converter.py b/pennylane/transforms/zx/converter.py
index 2ba98532740..5115855f46b 100644
--- a/pennylane/transforms/zx/converter.py
+++ b/pennylane/transforms/zx/converter.py
@@ -14,11 +14,14 @@
 """Transforms for interacting with PyZX, framework for ZX calculus."""
 # pylint: disable=too-many-statements, too-many-branches, too-many-return-statements, too-many-arguments
 
+from functools import partial
+from typing import Sequence, Callable
 from collections import OrderedDict
 import numpy as np
 import pennylane as qml
-from pennylane.transforms.op_transforms import op_transform
-from pennylane.tape import QuantumScript
+from pennylane.tape import QuantumScript, QuantumTape
+from pennylane.transforms.op_transforms import OperationTransformError
+from pennylane.transforms.core import transform
 from pennylane.wires import Wires
 
 
@@ -46,8 +49,7 @@ class EdgeType:  # pylint: disable=too-few-public-methods
     HADAMARD = 2
 
 
-@op_transform
-def to_zx(tape, expand_measurement=False):  # pylint: disable=unused-argument
+def to_zx(tape, expand_measurements=False):  # pylint: disable=unused-argument
     """This transform converts a PennyLane quantum tape to a ZX-Graph in the `PyZX framework <https://pyzx.readthedocs.io/en/latest/>`_.
     The graph can be optimized and transformed by well-known ZX-calculus reductions.
 
@@ -244,11 +246,20 @@ def mod_5_4():
         Copyright (C) 2018 - Aleks Kissinger and John van de Wetering
     """
     # If it is a simple operation just transform it to a tape
-    return _to_zx(QuantumScript([tape]))
+    if not isinstance(tape, qml.operation.Operator):
+        if not isinstance(tape, (qml.tape.QuantumScript, qml.QNode)) and not callable(tape):
+            raise OperationTransformError(
+                "Input is not an Operator, tape, QNode, or quantum function"
+            )
+        return _to_zx_transform(tape, expand_measurements=expand_measurements)
+
+    return to_zx(QuantumScript([tape]))
 
 
-@to_zx.tape_transform
-def _to_zx(tape, expand_measurements=False):
+@partial(transform, is_informative=True)
+def _to_zx_transform(
+    tape: QuantumTape, expand_measurements=False
+) -> (Sequence[QuantumTape], Callable):
     """Private function to convert a PennyLane tape to a `PyZX graph <https://pyzx.readthedocs.io/en/latest/>`_ ."""
     # Avoid to make PyZX a requirement for PennyLane.
     try:
@@ -280,68 +291,72 @@ def _to_zx(tape, expand_measurements=False):
         "CCZ": pyzx.circuit.gates.CCZ,
         "Toffoli": pyzx.circuit.gates.Tofolli,
     }
-    # Create the graph, a qubit mapper, the classical mapper stays empty as PennyLane does not support classical bits.
-    graph = Graph(None)
-    q_mapper = TargetMapper()
-    c_mapper = TargetMapper()
 
-    # Map the wires to consecutive wires
+    def processing_fn(res):
+        # Create the graph, a qubit mapper, the classical mapper stays empty as PennyLane does not support classical bits.
+        graph = Graph(None)
+        q_mapper = TargetMapper()
+        c_mapper = TargetMapper()
 
-    consecutive_wires = Wires(range(len(tape.wires)))
-    consecutive_wires_map = OrderedDict(zip(tape.wires, consecutive_wires))
-    tape = qml.map_wires(input=tape, wire_map=consecutive_wires_map)
+        # Map the wires to consecutive wires
 
-    inputs = []
+        consecutive_wires = Wires(range(len(res[0].wires)))
+        consecutive_wires_map = OrderedDict(zip(res[0].wires, consecutive_wires))
+        mapped_tape = qml.map_wires(input=res[0], wire_map=consecutive_wires_map)
 
-    # Create the qubits in the graph and the qubit mapper
-    for i in range(len(tape.wires)):
-        vertex = graph.add_vertex(VertexType.BOUNDARY, i, 0)
-        inputs.append(vertex)
-        q_mapper.set_prev_vertex(i, vertex)
-        q_mapper.set_next_row(i, 1)
-        q_mapper.set_qubit(i, i)
+        inputs = []
 
-    # Expand the tape to be compatible with PyZX and add rotations first for measurements
-    stop_crit = qml.BooleanFn(lambda obj: obj.name in gate_types)
-    tape = qml.tape.tape.expand_tape(
-        tape, depth=10, stop_at=stop_crit, expand_measurements=expand_measurements
-    )
+        # Create the qubits in the graph and the qubit mapper
+        for i in range(len(mapped_tape.wires)):
+            vertex = graph.add_vertex(VertexType.BOUNDARY, i, 0)
+            inputs.append(vertex)
+            q_mapper.set_prev_vertex(i, vertex)
+            q_mapper.set_next_row(i, 1)
+            q_mapper.set_qubit(i, i)
 
-    expanded_operations = []
+        # Expand the tape to be compatible with PyZX and add rotations first for measurements
+        stop_crit = qml.BooleanFn(lambda obj: obj.name in gate_types)
+        mapped_tape = qml.tape.tape.expand_tape(
+            mapped_tape, depth=10, stop_at=stop_crit, expand_measurements=expand_measurements
+        )
 
-    # Define specific decompositions
-    for op in tape.operations:
-        if op.name == "RY":
-            theta = op.data[0]
-            decomp = [
-                qml.RX(np.pi / 2, wires=op.wires),
-                qml.RZ(theta + np.pi, wires=op.wires),
-                qml.RX(np.pi / 2, wires=op.wires),
-                qml.RZ(3 * np.pi, wires=op.wires),
-            ]
-            expanded_operations.extend(decomp)
-        else:
-            expanded_operations.append(op)
+        expanded_operations = []
+
+        # Define specific decompositions
+        for op in mapped_tape.operations:
+            if op.name == "RY":
+                theta = op.data[0]
+                decomp = [
+                    qml.RX(np.pi / 2, wires=op.wires),
+                    qml.RZ(theta + np.pi, wires=op.wires),
+                    qml.RX(np.pi / 2, wires=op.wires),
+                    qml.RZ(3 * np.pi, wires=op.wires),
+                ]
+                expanded_operations.extend(decomp)
+            else:
+                expanded_operations.append(op)
+
+        expanded_tape = QuantumScript(expanded_operations, mapped_tape.measurements, [])
 
-    expanded_tape = QuantumScript(expanded_operations, tape.measurements, [])
+        _add_operations_to_graph(expanded_tape, graph, gate_types, q_mapper, c_mapper)
 
-    _add_operations_to_graph(expanded_tape, graph, gate_types, q_mapper, c_mapper)
+        row = max(q_mapper.max_row(), c_mapper.max_row())
 
-    row = max(q_mapper.max_row(), c_mapper.max_row())
+        outputs = []
+        for mapper in (q_mapper, c_mapper):
+            for label in mapper.labels():
+                qubit = mapper.to_qubit(label)
+                vertex = graph.add_vertex(VertexType.BOUNDARY, qubit, row)
+                outputs.append(vertex)
+                pre_vertex = mapper.prev_vertex(label)
+                graph.add_edge(graph.edge(pre_vertex, vertex))
 
-    outputs = []
-    for mapper in (q_mapper, c_mapper):
-        for label in mapper.labels():
-            qubit = mapper.to_qubit(label)
-            vertex = graph.add_vertex(VertexType.BOUNDARY, qubit, row)
-            outputs.append(vertex)
-            pre_vertex = mapper.prev_vertex(label)
-            graph.add_edge(graph.edge(pre_vertex, vertex))
+        graph.set_inputs(tuple(inputs))
+        graph.set_outputs(tuple(outputs))
 
-    graph.set_inputs(tuple(inputs))
-    graph.set_outputs(tuple(outputs))
+        return graph
 
-    return graph
+    return [tape], processing_fn
 
 
 def _add_operations_to_graph(tape, graph, gate_types, q_mapper, c_mapper):
diff --git a/tests/ops/functions/test_eigvals.py b/tests/ops/functions/test_eigvals.py
index f17cc8a5a31..e9cda9579b4 100644
--- a/tests/ops/functions/test_eigvals.py
+++ b/tests/ops/functions/test_eigvals.py
@@ -16,13 +16,13 @@
 """
 # pylint: disable=too-few-public-methods
 from functools import reduce
-
 import pytest
 import scipy
 
 from gate_data import CNOT, H, I, S, X, Y, Z
 import pennylane as qml
 from pennylane import numpy as np
+from pennylane.transforms.op_transforms import OperationTransformError
 
 one_qubit_no_parameter = [
     qml.PauliX,
@@ -37,6 +37,15 @@
 one_qubit_one_parameter = [qml.RX, qml.RY, qml.RZ, qml.PhaseShift]
 
 
+def test_invalid_argument():
+    """Assert error raised when input is neither a tape, QNode, nor quantum function"""
+    with pytest.raises(
+        OperationTransformError,
+        match="Input is not an Operator, tape, QNode, or quantum function",
+    ):
+        _ = qml.eigvals(None)
+
+
 class TestSingleOperation:
     @pytest.mark.parametrize("op_class", one_qubit_no_parameter)
     def test_non_parametric_instantiated(self, op_class):
diff --git a/tests/ops/functions/test_matrix.py b/tests/ops/functions/test_matrix.py
index 49767fdf4d8..cfeacfbe698 100644
--- a/tests/ops/functions/test_matrix.py
+++ b/tests/ops/functions/test_matrix.py
@@ -15,7 +15,7 @@
 Unit tests for the get_unitary_matrix transform
 """
 # pylint: disable=too-few-public-methods,too-many-function-args
-from functools import reduce
+from functools import reduce, partial
 
 import pytest
 
@@ -369,7 +369,7 @@ def test_qnode_wireorder(self):
         """Test changing the wire order when using a QNode"""
         dev = qml.device("default.qubit", wires=[1, 0, 2, 3])
 
-        @qml.matrix()
+        @qml.matrix
         @qml.qnode(dev)
         def testcircuit1(x):
             qml.PauliX(wires=0)
@@ -383,7 +383,7 @@ def testcircuit1(x):
         expected_matrix = np.kron(RY(x), np.kron(X, np.kron(Z, I)))
         assert np.allclose(testcircuit1(x), expected_matrix)
 
-        @qml.matrix(wire_order=[1, 0, 2])
+        @partial(qml.matrix, wire_order=[1, 0, 2])
         @qml.qnode(dev)
         def testcircuit2(x):
             qml.PauliX(wires=0)
@@ -493,16 +493,7 @@ def test_invalid_argument(self):
             OperationTransformError,
             match="Input is not an Operator, tape, QNode, or quantum function",
         ):
-            qml.matrix(None)(0.5)
-
-    def test_wrong_function(self):
-        """Assert error raised when input function is not a quantum function"""
-
-        def testfunction(x):
-            return x
-
-        with pytest.raises(OperationTransformError, match="function contains no quantum operation"):
-            qml.matrix(testfunction)(0)
+            _ = qml.matrix(None)
 
     def test_inconsistent_wires(self):
         """Assert error raised when wire labels in wire_order and circuit are inconsistent"""
@@ -519,6 +510,12 @@ def circuit():
         ):
             qml.matrix(circuit, wire_order=wires)()
 
+        with pytest.raises(
+            OperationTransformError,
+            match=r"Wires in circuit \[0\] are inconsistent with those in wire_order \[1\]",
+        ):
+            qml.matrix(qml.PauliX(0), wire_order=[1])
+
 
 class TestInterfaces:
     @pytest.mark.tf
diff --git a/tests/transforms/test_experimental/test_transform_dispatcher.py b/tests/transforms/test_experimental/test_transform_dispatcher.py
index a37403d2eb0..17d9da15691 100644
--- a/tests/transforms/test_experimental/test_transform_dispatcher.py
+++ b/tests/transforms/test_experimental/test_transform_dispatcher.py
@@ -139,7 +139,18 @@ def non_valid_expand_transform(
     return [tape], lambda x: x
 
 
-class TestTransformDispatcher:
+##########################################
+# Valid informative transform
+def informative_transform(tape: qml.tape.QuantumTape) -> (Sequence[qml.tape.QuantumTape], Callable):
+    """A valid informative transform"""
+
+    def fn(results):
+        return len(results[0].operations)
+
+    return [tape], fn
+
+
+class TestTransformDispatcher:  # pylint: disable=too-many-public-methods
     """Test the transform function (validate and dispatch)."""
 
     @pytest.mark.parametrize("valid_transform", valid_transforms)
@@ -179,6 +190,42 @@ def qnode_circuit(a):
         )
         assert not dispatched_transform.is_informative
 
+    def test_integration_dispatcher_with_informative_transform(self):
+        """Test that no error is raised with the transform function and that the transform dispatcher returns
+        the right object when an informative transform is applied."""
+
+        dispatched_transform = transform(informative_transform, is_informative=True)
+
+        # Applied on a tape (return processed results)
+        expected = len(tape_circuit.operations)
+        num_ops = dispatched_transform(tape_circuit)
+        assert num_ops == expected
+
+        # Applied on a qfunc (return a qfunc)
+        qfunc = dispatched_transform(qfunc_circuit)
+        assert callable(qfunc)
+
+        # Applied on a qnode (return a qnode with populated the program)
+        @qml.qnode(device=dev)
+        def qnode_circuit(a):
+            """QNode circuit."""
+            qml.Hadamard(wires=0)
+            qml.CNOT(wires=[0, 1])
+            qml.PauliX(wires=0)
+            qml.RZ(a, wires=1)
+            return qml.expval(qml.PauliZ(wires=0))
+
+        qnode_transformed = dispatched_transform(qnode_circuit)
+        assert not qnode_circuit.transform_program
+
+        assert qnode_transformed(0.1) == 4
+        assert isinstance(qnode_transformed, qml.QNode)
+        assert isinstance(qnode_transformed.transform_program, qml.transforms.core.TransformProgram)
+        assert isinstance(
+            qnode_transformed.transform_program.pop_front(), qml.transforms.core.TransformContainer
+        )
+        assert dispatched_transform.is_informative
+
     @pytest.mark.parametrize("valid_transform", valid_transforms)
     def test_integration_dispatcher_with_valid_transform_decorator_partial(self, valid_transform):
         """Test that no error is raised with the transform function and that the transform dispatcher returns
@@ -380,6 +427,17 @@ def test_qfunc_transform_multiple_tapes(self):
         ):
             dispatched_transform(qfunc_circuit, 0)(0.42)
 
+    def test_informative_transform_tape_return(self):
+        """Test that disaptched informative transforms return processed results instead of
+        a list of tapes and processing function."""
+        tape = qml.tape.QuantumScript(
+            [qml.PauliX(0), qml.CNOT([0, 1]), qml.RX(0.234, 1), qml.Hadamard(1)]
+        )
+        dispatched_transform = transform(informative_transform, is_informative=True)
+
+        num_ops = dispatched_transform(tape)
+        assert num_ops == 4
+
     def test_dispatched_transform_attribute(self):
         """Test the dispatcher attributes."""
         dispatched_transform = transform(first_valid_transform)
@@ -394,19 +452,21 @@ def test_the_transform_container_attributes(self):
             first_valid_transform, args=[0], kwargs={}, classical_cotransform=None
         )
 
-        q_transform, args, kwargs, cotransform, is_informative = container
+        q_transform, args, kwargs, cotransform, is_informative, final_transform = container
 
         assert q_transform is first_valid_transform
         assert args == [0]
         assert kwargs == {}
         assert cotransform is None
         assert not is_informative
+        assert not final_transform
 
         assert container.transform is first_valid_transform
         assert container.args == [0]
         assert not container.kwargs
         assert container.classical_cotransform is None
         assert not container.is_informative
+        assert not container.final_transform
 
     @pytest.mark.parametrize("valid_transform", valid_transforms)
     def test_custom_qnode_transform(self, valid_transform):
diff --git a/tests/transforms/test_experimental/test_transform_program.py b/tests/transforms/test_experimental/test_transform_program.py
index 00c3153867c..4926411d0f7 100644
--- a/tests/transforms/test_experimental/test_transform_program.py
+++ b/tests/transforms/test_experimental/test_transform_program.py
@@ -62,6 +62,15 @@ def fn(results):
     return [tape1, tape2], fn
 
 
+def informative_transform(tape: qml.tape.QuantumTape) -> (Sequence[qml.tape.QuantumTape], Callable):
+    """A valid informative transform"""
+
+    def fn(results):
+        return len(results[0].operations)
+
+    return [tape], fn
+
+
 class TestUtilityHelpers:
     """test the private functions used in post processing."""
 
@@ -187,6 +196,62 @@ def test_add_two_programs(self):
         assert isinstance(transform_program[4], TransformContainer)
         assert transform_program[4].transform is second_valid_transform
 
+    def test_add_both_final_transform_programs(self):
+        """Test that an error is raised if two programs are added when both have
+        terminal transforms"""
+        transform1 = TransformContainer(transform=first_valid_transform)
+        transform2 = TransformContainer(transform=second_valid_transform, final_transform=True)
+
+        transform_program1 = TransformProgram()
+        transform_program1.push_back(transform1)
+        transform_program1.push_back(transform2)
+
+        transform_program2 = TransformProgram()
+        transform_program2.push_back(transform1)
+        transform_program2.push_back(transform2)
+
+        with pytest.raises(
+            TransformError, match="The transform program already has a terminal transform"
+        ):
+            _ = transform_program1 + transform_program2
+
+    def test_add_programs_with_one_final_transform(self):
+        """Test that transform programs are added correctly when one of them has a terminal
+        transform."""
+        transform1 = TransformContainer(transform=first_valid_transform)
+        transform2 = TransformContainer(transform=second_valid_transform, final_transform=True)
+
+        transform_program1 = TransformProgram()
+        transform_program1.push_back(transform1)
+
+        transform_program2 = TransformProgram()
+        transform_program2.push_back(transform1)
+        transform_program2.push_back(transform2)
+
+        merged_program1 = transform_program1 + transform_program2
+        assert len(merged_program1) == 3
+
+        assert isinstance(merged_program1[0], TransformContainer)
+        assert merged_program1[0].transform is first_valid_transform
+
+        assert isinstance(merged_program1[1], TransformContainer)
+        assert merged_program1[1].transform is first_valid_transform
+
+        assert isinstance(merged_program1[2], TransformContainer)
+        assert merged_program1[2].transform is second_valid_transform
+
+        merged_program2 = transform_program2 + transform_program1
+        assert len(merged_program2) == 3
+
+        assert isinstance(merged_program2[0], TransformContainer)
+        assert merged_program2[0].transform is first_valid_transform
+
+        assert isinstance(merged_program2[1], TransformContainer)
+        assert merged_program2[1].transform is first_valid_transform
+
+        assert isinstance(merged_program2[2], TransformContainer)
+        assert merged_program2[2].transform is second_valid_transform
+
     def test_repr_program(self):
         """Test the string representation of a program."""
         transform_program = TransformProgram()
@@ -406,14 +471,14 @@ def test_valid_transforms(self):
         transform_program.push_back(transform1)
 
         with pytest.raises(
-            TransformError, match="The transform program already has an informative transform."
+            TransformError, match="The transform program already has a terminal transform."
         ):
             transform_program.push_back(transform1)
 
         transform2 = TransformContainer(transform=second_valid_transform, is_informative=False)
 
         with pytest.raises(
-            TransformError, match="The transform program already has an informative transform."
+            TransformError, match="The transform program already has a terminal transform."
         ):
             transform_program.push_back(transform2)
 
@@ -435,17 +500,6 @@ def test_call_on_empty_program(self):
         obj = [1, 2, 3, "b"]
         assert null_postprocessing(obj) is obj
 
-    def test_informative_transforms_not_supported(self):
-        """Test that a program with an informative raises a `NotImplementedError` on call."""
-        my_transform = TransformContainer(first_valid_transform, is_informative=True)
-        prog = TransformProgram((my_transform,))
-        batch = (qml.tape.QuantumScript([], [qml.state()]),)
-
-        with pytest.raises(
-            NotImplementedError, match="Informative transforms are not yet supported."
-        ):
-            prog(batch)
-
     def test_cotransform_support_notimplemented(self):
         """Test that a transform with a cotransform raises a not implemented error."""
 
@@ -553,7 +607,7 @@ def transform_mul(tape: qml.tape.QuantumTape):
         assert postprocessing1.keywords["slices"] == [slice(0, 1)]
 
         results = (1.0,)
-        expected = (3.0,)  # 2.0 *1.0 + 1.0
+        expected = (3.0,)  # 2.0 * 1.0 + 1.0
         assert fn(results) == expected
 
         # Test reverse direction
@@ -676,6 +730,44 @@ def qnode_circuit(a):
         assert program[0].transform is first_valid_transform
         assert program[1].transform is first_valid_transform
 
+    def test_qnode_integration_informative_transform(self):
+        """Test the integration with QNode with two transforms, one of which is
+        informative."""
+        dispatched_transform_1 = transform(first_valid_transform)
+        dispatched_transform_2 = transform(informative_transform, is_informative=True)
+
+        dev = qml.device("default.qubit", wires=2)
+
+        @qml.qnode(device=dev)
+        def qnode_circuit(a):
+            """QNode circuit."""
+            qml.Hadamard(wires=0)
+            qml.CNOT(wires=[0, 1])
+            qml.PauliX(wires=0)
+            qml.RZ(a, wires=1)
+            return qml.expval(qml.PauliZ(wires=0))
+
+        new_qnode = dispatched_transform_2(dispatched_transform_1(qnode_circuit, 0))
+
+        program = new_qnode.transform_program
+        transformed_qnode_rep = repr(program)
+        assert (
+            transformed_qnode_rep
+            == "TransformProgram("
+            + str(first_valid_transform.__name__)
+            + ", "
+            + str(informative_transform.__name__)
+            + ")"
+        )
+
+        assert not program.is_empty()
+        assert len(program) == 2
+        assert program[0].transform is first_valid_transform
+        assert program[1].transform is informative_transform
+
+        result = new_qnode(0.1)
+        assert result == (3,)
+
     def test_qnode_integration_different_transforms(self):
         """Test the integration with QNode with two different transforms."""
 
diff --git a/tests/transforms/test_zx.py b/tests/transforms/test_zx.py
index 847ba14236f..34232d4267e 100644
--- a/tests/transforms/test_zx.py
+++ b/tests/transforms/test_zx.py
@@ -15,11 +15,13 @@
 Unit tests for the `pennylane.transforms.zx` folder.
 """
 import sys
+from functools import partial
 
 import numpy as np
 import pytest
 import pennylane as qml
 from pennylane.tape import QuantumScript
+from pennylane.transforms.op_transforms import OperationTransformError
 
 pyzx = pytest.importorskip("pyzx")
 
@@ -67,6 +69,14 @@ def test_import_pyzx(monkeypatch):
 class TestConvertersZX:
     """Test converters to_zx and from_zx."""
 
+    def test_invalid_argument(self):
+        """Assert error raised when input is neither a tape, QNode, nor quantum function"""
+        with pytest.raises(
+            OperationTransformError,
+            match="Input is not an Operator, tape, QNode, or quantum function",
+        ):
+            _ = qml.transforms.to_zx(None)
+
     @pytest.mark.parametrize("script", qscript)
     @pytest.mark.parametrize("operation", supported_operations)
     def test_supported_operation_no_params(self, operation, script):
@@ -636,7 +646,7 @@ def test_qnode_decorator(self):
         """Test the QNode decorator."""
         dev = qml.device("default.qubit", wires=2)
 
-        @qml.transforms.to_zx(expand_measurements=True)
+        @partial(qml.transforms.to_zx, expand_measurements=True)
         @qml.qnode(device=dev)
         def circuit(p):
             qml.RZ(p[0], wires=1)

From ac42ff9b91d2d077d9d043886133da96a28afd2c Mon Sep 17 00:00:00 2001
From: Matthew Silverman <matthews@xanadu.ai>
Date: Fri, 15 Sep 2023 10:15:14 -0400
Subject: [PATCH 057/131] Make the device API not-experimental (#4594)

**Context:**
DefaultQubit2 is the new DefaultQubit! Let us name it accordingly.

**Description of the Change:**
1. Rename the existing `DefaultQubit` to `DefaultQubitLegacy`
2. Move the contents of `pennylane/devices/experimental` into
`pennylane/devices`
3. Rename `DefaultQubit2` to `DefaultQubit`

Notes/Follow-up:
- The short names for DQL and DQ2 are still `default.qubit` and
`default.qubit.2` respectively, and `default.qubit` still results in
DQL. I will make that switch in another PR.
- Some test files still need moving/renaming, but I'll do that in
another PR. All files for DQ2 should be at the higher level, and all DQL
files should be in a `legacy` folder with a `legacy` suffix, so that's
next.

To be extra clear about what's happened:
- the old device API is still `qml.Device`
- the new device API is now `qml.devices.Device`
- If called directly, `qml.devices.DefaultQubit` return DQ2.
`qml.devices.DefaultQubitLegacy` returns DQL

**Benefits:**
The Device API is no longer experimental! Will make the final switch
much easier!

**Possible Drawbacks:**
I'm probably going to screw up merging this into mega-PR #4436

---------

Co-authored-by: Mudit Pandey <mudit.pandey@xanadu.ai>
---
 doc/development/adding_operators.rst          |    6 +-
 doc/development/deprecations.rst              |    6 +-
 doc/development/plugins.rst                   |    6 +-
 doc/releases/changelog-dev.md                 |    4 +
 pennylane/_qubit_device.py                    |    4 +-
 pennylane/debugging.py                        |    3 +-
 pennylane/devices/__init__.py                 |   21 +-
 pennylane/devices/default_mixed.py            |    6 +-
 pennylane/devices/default_qubit.py            | 1416 +++++------------
 pennylane/devices/default_qubit_autograd.py   |    8 +-
 pennylane/devices/default_qubit_jax.py        |   10 +-
 pennylane/devices/default_qubit_legacy.py     | 1085 +++++++++++++
 pennylane/devices/default_qubit_tf.py         |   14 +-
 pennylane/devices/default_qubit_torch.py      |    8 +-
 pennylane/devices/default_qutrit.py           |    8 +-
 .../devices/{experimental => }/device_api.py  |    2 +-
 .../{experimental => }/execution_config.py    |    0
 pennylane/devices/experimental/__init__.py    |   20 -
 .../devices/experimental/default_qubit_2.py   |  529 ------
 pennylane/devices/null_qubit.py               |    6 +-
 pennylane/devices/qubit/preprocess.py         |   14 +-
 pennylane/drawer/draw.py                      |    8 +-
 pennylane/interfaces/execution.py             |   18 +-
 pennylane/interfaces/jax_jit.py               |    6 +-
 pennylane/interfaces/set_shots.py             |    2 +-
 pennylane/qinfo/transforms.py                 |    6 +-
 pennylane/qnode.py                            |   28 +-
 pennylane/tape/qscript.py                     |    2 +-
 pennylane/transforms/qcut/cutstrategy.py      |    6 +-
 pennylane/transforms/specs.py                 |    2 +-
 setup.py                                      |    4 +-
 .../experimental/test_default_qubit_2.py      |  232 +--
 tests/devices/experimental/test_device_api.py |    4 +-
 tests/devices/qubit/test_adjoint_jacobian.py  |    2 +-
 tests/devices/qubit/test_preprocess.py        |   14 +-
 tests/devices/test_default_qubit_legacy.py    |   16 +-
 .../test_default_qubit_legacy_broadcasting.py |   20 +-
 tests/drawer/test_draw.py                     |    2 +-
 tests/drawer/test_draw_mpl.py                 |    2 +-
 .../finite_diff/test_finite_difference.py     |    6 +-
 .../test_finite_difference_shot_vec.py        |    6 +-
 .../finite_diff/test_spsa_gradient.py         |    6 +-
 .../test_spsa_gradient_shot_vec.py            |    6 +-
 .../parameter_shift/test_parameter_shift.py   |    6 +-
 .../test_parameter_shift_shot_vec.py          |    6 +-
 .../test_autograd_default_qubit_2.py          |   20 +-
 .../test_autograd_qnode_default_qubit_2.py    |   62 +-
 .../test_execute_default_qubit_2.py           |   16 +-
 .../test_jax_default_qubit_2.py               |   16 +-
 .../test_jax_jit_qnode_default_qubit_2.py     |   26 +-
 .../test_jax_qnode_default_qubit_2.py         |   26 +-
 ...graph_qnode_shot_vector_default_qubit_2.py |    8 +-
 .../test_tensorflow_default_qubit_2.py        |   24 +-
 .../test_tensorflow_qnode_default_qubit_2.py  |   54 +-
 ...rflow_qnode_shot_vector_default_qubit_2.py |    8 +-
 .../test_torch_default_qubit_2.py             |   18 +-
 .../test_torch_qnode_default_qubit_2.py       |   50 +-
 tests/interfaces/test_autograd.py             |   14 +-
 tests/interfaces/test_jax.py                  |    4 +-
 tests/interfaces/test_jax_jit.py              |    4 +-
 tests/interfaces/test_set_shots.py            |    6 +-
 tests/interfaces/test_tensorflow.py           |    4 +-
 tests/interfaces/test_torch.py                |    4 +-
 .../test_transform_program_integration.py     |    8 +-
 tests/logging/test_logging_autograd.py        |    4 +-
 .../measurements/legacy/test_state_legacy.py  |    6 +-
 tests/measurements/test_state.py              |    6 +-
 tests/test_debugging.py                       |    6 +-
 tests/test_qnode.py                           |   25 +-
 tests/test_return_types_qnode.py              |    8 +-
 .../transforms/test_adjoint_metric_tensor.py  |    2 +-
 .../test_optimization/test_cancel_inverses.py |    2 +-
 .../test_commute_controlled.py                |    2 +-
 .../test_optimization/test_merge_rotations.py |    2 +-
 tests/transforms/test_qcut.py                 |    4 +-
 75 files changed, 1996 insertions(+), 2029 deletions(-)
 create mode 100644 pennylane/devices/default_qubit_legacy.py
 rename pennylane/devices/{experimental => }/device_api.py (99%)
 rename pennylane/devices/{experimental => }/execution_config.py (100%)
 delete mode 100644 pennylane/devices/experimental/__init__.py
 delete mode 100644 pennylane/devices/experimental/default_qubit_2.py

diff --git a/doc/development/adding_operators.rst b/doc/development/adding_operators.rst
index 4343e2183ee..1be3872a709 100644
--- a/doc/development/adding_operators.rst
+++ b/doc/development/adding_operators.rst
@@ -212,11 +212,11 @@ FlipAndRotate(-0.1, wires=['q3', 'q1'])
 The new gate can be used with PennyLane devices. Device support for an operation can be checked via
 ``dev.stopping_condition(op)``.  If ``True``, then the device supports the operation.
 
-``DefaultQubit`` first checks if the operator has a matrix using the :attr:`~.Operator.has_matrix` property.
+``DefaultQubitLegacy`` first checks if the operator has a matrix using the :attr:`~.Operator.has_matrix` property.
 If the Operator doesn't have a matrix, the device then checks if the name of the Operator is explicitly specified in 
-:attr:`~DefaultQubit.operations` or :attr:`~DefaultQubit.observables`.
+:attr:`~DefaultQubitLegacy.operations` or :attr:`~DefaultQubitLegacy.observables`.
 
-Other devices that do not inherit from ``DefaultQubit`` only check if the name is explicitly specified in the ``operations``
+Other devices that do not inherit from ``DefaultQubitLegacy`` only check if the name is explicitly specified in the ``operations``
 property.
 
 - If the device registers support for an operation with the same name,
diff --git a/doc/development/deprecations.rst b/doc/development/deprecations.rst
index b60d381599c..347c5641adf 100644
--- a/doc/development/deprecations.rst
+++ b/doc/development/deprecations.rst
@@ -44,7 +44,7 @@ Pending deprecations
   - Behaviour will change in v0.33
 
 * The public methods of ``DefaultQubit`` are pending changes to
-  follow the new device API, as used in ``DefaultQubit2``.
+  follow the new device API.
 
   We will be switching to the new device interface in a coming release.
   In this new interface, simulation implementation details
@@ -53,8 +53,8 @@ Pending deprecations
   in a workflow will remain the same.
   
   If you directly interact with device methods, please consult
-  :class:`pennylane.devices.experimental.Device` and
-  :class:`pennylane.devices.experimental.DefaultQubit2`
+  :class:`pennylane.devices.Device` and
+  :class:`pennylane.devices.DefaultQubit`
   for more information on what the new interface will look like and be prepared
   to make updates in a coming release. If you have any feedback on these
   changes, please create an
diff --git a/doc/development/plugins.rst b/doc/development/plugins.rst
index 6c6988ca2fe..3f2d415dd5c 100644
--- a/doc/development/plugins.rst
+++ b/doc/development/plugins.rst
@@ -61,7 +61,7 @@ and subclassing it:
 .. warning::
 
     The API of PennyLane devices will be updated soon to follow a new interface defined by
-    the experimental :class:`pennylane.devices.experimental.Device` class. This guide describes
+    the :class:`pennylane.devices.Device` class. This guide describes
     how to create a device with the current :class:`pennylane.Device` and
     :class:`pennylane.QubitDevice` base classes, and will be updated as we switch to the new API.
     In the meantime, please reach out to the PennyLane team if you would like help with building
@@ -108,9 +108,9 @@ as well as potential further capabilities, by providing the following class attr
               return obj.name in {'CNOT', 'PauliX', 'PauliY', 'PauliZ'}
           return qml.BooleanFn(accepts_obj)
 
-  If the device does *not* inherit from :class:`~.DefaultQubit`, then supported operations can be determined
+  If the device does *not* inherit from :class:`~.DefaultQubitLegacy`, then supported operations can be determined
   by the :attr:`pennylane.Device.operations` property.  This property is a list of string names for supported operations.
-  :class:`~.DefaultQubit` supports any operation with a matrix, even if it's name isn't specifically enumerated
+  :class:`~.DefaultQubitLegacy` supports any operation with a matrix, even if it's name isn't specifically enumerated
   in :attr:`pennylane.Device.operations`.
 
   .. code-block:: python
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 177052d23ce..5d5500d5de2 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -40,6 +40,10 @@
 * Quantum information transforms are updated to the new transform program system.
   [(#4569)](https://github.com/PennyLaneAI/pennylane/pull/4569)
 
+* `qml.devices.DefaultQubit` now implements the new device API. The old version of `default.qubit`
+  is still accessible via `qml.devices.DefaultQubitLegacy`, or via short name `default.qubit.legacy`.
+  [(#4594)](https://github.com/PennyLaneAI/pennylane/pull/4594)
+
 <h3>Improvements 🛠</h3>
 
 * Tensor-network template `qml.MPS` now supports changing `offset` between subsequent blocks for more flexibility.
diff --git a/pennylane/_qubit_device.py b/pennylane/_qubit_device.py
index fa9d86fd587..cb6713cdac4 100644
--- a/pennylane/_qubit_device.py
+++ b/pennylane/_qubit_device.py
@@ -178,13 +178,13 @@ def _const_mul(constant, array):
 
     **Example:**
 
-    Let's create device that inherits from :class:`~pennylane.devices.DefaultQubit` and overrides the
+    Let's create a device that inherits from :class:`~pennylane.devices.DefaultQubitLegacy` and overrides the
     logic of the `qml.sample` measurement. To do so we will need to update the ``measurement_map``
     dictionary:
 
     .. code-block:: python
 
-        class NewDevice(DefaultQubit):
+        class NewDevice(DefaultQubitLegacy):
             def __init__(self, wires, shots):
                 super().__init__(wires=wires, shots=shots)
                 self.measurement_map[SampleMP] = "sample_measurement"
diff --git a/pennylane/debugging.py b/pennylane/debugging.py
index 76e77b79db0..eded61f7e04 100644
--- a/pennylane/debugging.py
+++ b/pennylane/debugging.py
@@ -16,7 +16,6 @@
 """
 import pennylane as qml
 from pennylane import DeviceError
-from pennylane.devices import experimental
 
 
 class _Debugger:
@@ -35,7 +34,7 @@ def __init__(self, dev):
             raise DeviceError("Device does not support snapshots.")
 
         # new device API: check if it's the simulator device
-        if isinstance(dev, experimental.Device) and not isinstance(dev, experimental.DefaultQubit2):
+        if isinstance(dev, qml.devices.Device) and not isinstance(dev, qml.devices.DefaultQubit):
             raise DeviceError("Device does not support snapshots.")
 
         self.snapshots = {}
diff --git a/pennylane/devices/__init__.py b/pennylane/devices/__init__.py
index 29626abb916..870c9a5eb4f 100644
--- a/pennylane/devices/__init__.py
+++ b/pennylane/devices/__init__.py
@@ -25,7 +25,7 @@
     :toctree: api
 
 
-    default_qubit
+    default_qubit_legacy
     default_qubit_jax
     default_qubit_torch
     default_qubit_tf
@@ -38,19 +38,19 @@
 Next generation devices
 -----------------------
 
-:class:`pennylane.devices.experimental.Device` in an experimental interface for the next generation of devices that
-will eventually replace :class:`pennylane.Device` and :class:`pennylane.QubitDevice`.
+:class:`pennylane.devices.Device` is the latest interface for the next generation of devices that
+replaces :class:`pennylane.Device` and :class:`pennylane.QubitDevice`.
 
-While the current interface :class:`pennylane.Device` is imported top level, the new :class:`pennylane.devices.experimental.Device` is
-accessible from the ``pennylane.devices.experimental`` submodule.
+While the previous interface :class:`pennylane.Device` is imported top level, the new :class:`pennylane.devices.Device` is
+accessible from the ``pennylane.devices`` submodule.
 
-.. currentmodule:: pennylane.devices.experimental
+.. currentmodule:: pennylane.devices
 .. autosummary::
     :toctree: api
 
     ExecutionConfig
     Device
-    DefaultQubit2
+    DefaultQubit
 
 Qubit Simulation Tools
 ----------------------
@@ -60,13 +60,14 @@
 
 """
 
-from . import experimental
-from .experimental import ExecutionConfig
+from .execution_config import ExecutionConfig, DefaultExecutionConfig
+from .device_api import Device
+from .default_qubit import DefaultQubit
 
 # DefaultQubitTF and DefaultQubitAutograd not imported here since this
 # would lead to an automatic import of tensorflow and autograd, which are
 # not PennyLane core dependencies
-from .default_qubit import DefaultQubit
+from .default_qubit_legacy import DefaultQubitLegacy
 from .default_gaussian import DefaultGaussian
 from .default_mixed import DefaultMixed
 from .null_qubit import NullQubit
diff --git a/pennylane/devices/default_mixed.py b/pennylane/devices/default_mixed.py
index 63eac84d4b8..7581f4fee57 100644
--- a/pennylane/devices/default_mixed.py
+++ b/pennylane/devices/default_mixed.py
@@ -52,12 +52,12 @@ class DefaultMixed(QubitDevice):
     .. warning::
 
         The API of ``DefaultMixed`` will be updated soon to follow a new device interface described
-        in :class:`pennylane.devices.experimental.Device`.
+        in :class:`pennylane.devices.Device`.
 
         This change will not alter device behaviour for most workflows, but may have implications for
         plugin developers and users who directly interact with device methods. Please consult
-        :class:`pennylane.devices.experimental.Device` and the implementation in
-        :class:`pennylane.devices.experimental.DefaultQubit2` for more information on what the new
+        :class:`pennylane.devices.Device` and the implementation in
+        :class:`pennylane.devices.DefaultQubit` for more information on what the new
         interface will look like and be prepared to make updates in a coming release. If you have any
         feedback on these changes, please create an
         `issue <https://github.com/PennyLaneAI/pennylane/issues>`_ or post in our
diff --git a/pennylane/devices/default_qubit.py b/pennylane/devices/default_qubit.py
index 1995b598d66..e908996d879 100644
--- a/pennylane/devices/default_qubit.py
+++ b/pennylane/devices/default_qubit.py
@@ -1,4 +1,4 @@
-# Copyright 2018-2021 Xanadu Quantum Technologies Inc.
+# Copyright 2018-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.
@@ -11,1077 +11,519 @@
 # 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.
-r"""
-The default.qubit device is PennyLane's standard qubit-based device.
-
-It implements the necessary :class:`~pennylane._device.Device` methods as well as some built-in
-:mod:`qubit operations <pennylane.ops.qubit>`, and provides a very simple pure state
-simulation of a qubit-based quantum circuit architecture.
 """
-import functools
-import itertools
-from string import ascii_letters as ABC
-from typing import List
+This module contains the next generation successor to default qubit
+"""
 
+from functools import partial
+from numbers import Number
+from typing import Union, Callable, Tuple, Optional, Sequence
+import concurrent.futures
 import numpy as np
-from scipy.sparse import csr_matrix
-
-import pennylane as qml
-from pennylane import BasisState, DeviceError, QubitDevice, StatePrep, Snapshot
-from pennylane.devices.qubit import measure
-from pennylane.operation import Operation
-from pennylane.ops import Sum
-from pennylane.ops.qubit.attributes import diagonal_in_z_basis
-from pennylane.pulse import ParametrizedEvolution
-from pennylane.measurements import ExpectationMP, MeasurementValue
-from pennylane.typing import TensorLike
-from pennylane.wires import WireError
 
-from .._version import __version__
+from pennylane.tape import QuantumTape, QuantumScript
+from pennylane.typing import Result, ResultBatch
+from pennylane.transforms import convert_to_numpy_parameters
+from pennylane.transforms.core import TransformProgram
 
-ABC_ARRAY = np.array(list(ABC))
+from . import Device
+from .execution_config import ExecutionConfig, DefaultExecutionConfig
+from .qubit.simulate import simulate, get_final_state, measure_final_state
+from .qubit.preprocess import (
+    preprocess,
+    validate_and_expand_adjoint,
+    validate_multiprocessing_workers,
+    validate_device_wires,
+)
+from .qubit.adjoint_jacobian import adjoint_jacobian, adjoint_vjp, adjoint_jvp
 
-# tolerance for numerical errors
-tolerance = 1e-10
-SQRT2INV = 1 / np.sqrt(2)
-TPHASE = np.exp(1j * np.pi / 4)
+Result_or_ResultBatch = Union[Result, ResultBatch]
+QuantumTapeBatch = Sequence[QuantumTape]
+QuantumTape_or_Batch = Union[QuantumTape, QuantumTapeBatch]
+# always a function from a resultbatch to either a result or a result batch
+PostprocessingFn = Callable[[ResultBatch], Result_or_ResultBatch]
 
 
-def _get_slice(index, axis, num_axes):
-    """Allows slicing along an arbitrary axis of an array or tensor.
+class DefaultQubit(Device):
+    """A PennyLane device written in Python and capable of backpropagation derivatives.
 
     Args:
-        index (int): the index to access
-        axis (int): the axis to slice into
-        num_axes (int): total number of axes
-
-    Returns:
-        tuple[slice or int]: a tuple that can be used to slice into an array or tensor
+        shots (int, Sequence[int], Sequence[Union[int, Sequence[int]]]): The default number of shots to use in executions involving
+            this device.
+        seed (Union[str, None, int, array_like[int], SeedSequence, BitGenerator, Generator]): A
+            seed-like parameter matching that of ``seed`` for ``numpy.random.default_rng`` or
+            a request to seed from numpy's global random number generator.
+            The default, ``seed="global"`` pulls a seed from NumPy's global generator. ``seed=None``
+            will pull a seed from the OS entropy.
+        max_workers (int): A ``ProcessPoolExecutor`` executes tapes asynchronously
+            using a pool of at most ``max_workers`` processes. If ``max_workers`` is ``None``,
+            only the current process executes tapes. If you experience any
+            issue, say using JAX, TensorFlow, Torch, try setting ``max_workers`` to ``None``.
 
     **Example:**
 
-    Accessing the 2 index along axis 1 of a 3-axis array:
-
-    >>> sl = _get_slice(2, 1, 3)
-    >>> sl
-    (slice(None, None, None), 2, slice(None, None, None))
-    >>> a = np.arange(27).reshape((3, 3, 3))
-    >>> a[sl]
-    array([[ 6,  7,  8],
-           [15, 16, 17],
-           [24, 25, 26]])
-    """
-    idx = [slice(None)] * num_axes
-    idx[axis] = index
-    return tuple(idx)
-
-
-# pylint: disable=unused-argument
-class DefaultQubit(QubitDevice):
-    """Default qubit device for PennyLane.
-
-    .. warning::
-
-        The API of ``DefaultQubit`` will be updated soon to follow a new device interface described
-        in :class:`pennylane.devices.experimental.Device`.
-
-        This change will not alter device behaviour for most workflows, but may have implications for
-        plugin developers and users who directly interact with device methods. Please consult
-        :class:`pennylane.devices.experimental.Device` and the implementation in
-        :class:`pennylane.devices.experimental.DefaultQubit2` for more information on what the new
-        interface will look like and be prepared to make updates in a coming release. If you have any
-        feedback on these changes, please create an
-        `issue <https://github.com/PennyLaneAI/pennylane/issues>`_ or post in our
-        `discussion forum <https://discuss.pennylane.ai/>`_.
+    .. code-block:: python
+
+        n_layers = 5
+        n_wires = 10
+        num_qscripts = 5
+
+        shape = qml.StronglyEntanglingLayers.shape(n_layers=n_layers, n_wires=n_wires)
+        rng = qml.numpy.random.default_rng(seed=42)
+
+        qscripts = []
+        for i in range(num_qscripts):
+            params = rng.random(shape)
+            op = qml.StronglyEntanglingLayers(params, wires=range(n_wires))
+            qs = qml.tape.QuantumScript([op], [qml.expval(qml.PauliZ(0))])
+            qscripts.append(qs)
+
+    >>> dev = DefaultQubit()
+    >>> program, execution_config = dev.preprocess()
+    >>> new_batch, post_processing_fn = program(qscripts)
+    >>> results = dev.execute(new_batch, execution_config=execution_config)
+    >>> post_processing_fn(results)
+    [-0.0006888975950537501,
+    0.025576307134457577,
+    -0.0038567269892757494,
+    0.1339705146860149,
+    -0.03780669772690448]
+
+    Suppose one has a processor with 5 cores or more, these scripts can be executed in
+    parallel as follows
+
+    >>> dev = DefaultQubit(max_workers=5)
+    >>> program, execution_config = dev.preprocess()
+    >>> new_batch, post_processing_fn = program(qscripts)
+    >>> results = dev.execute(new_batch, execution_config=execution_config)
+    >>> post_processing_fn(results)
+
+    If you monitor your CPU usage, you should see 5 new Python processes pop up to
+    crunch through those ``QuantumScript``'s. Beware not oversubscribing your machine.
+    This may happen if a single device already uses many cores, if NumPy uses a multi-
+    threaded BLAS library like MKL or OpenBLAS for example. The number of threads per
+    process times the number of processes should not exceed the number of cores on your
+    machine. You can control the number of threads per process with the environment
+    variables:
+
+    * OMP_NUM_THREADS
+    * MKL_NUM_THREADS
+    * OPENBLAS_NUM_THREADS
+
+    where the last two are specific to the MKL and OpenBLAS libraries specifically.
+
+    This device currently supports backpropagation derivatives:
+
+    >>> from pennylane.devices import ExecutionConfig
+    >>> dev.supports_derivatives(ExecutionConfig(gradient_method="backprop"))
+    True
+
+    For example, we can use jax to jit computing the derivative:
+
+    .. code-block:: python
+
+        import jax
+
+        @jax.jit
+        def f(x):
+            qs = qml.tape.QuantumScript([qml.RX(x, 0)], [qml.expval(qml.PauliZ(0))])
+            program, execution_config = dev.preprocess()
+            new_batch, post_processing_fn = program([qs])
+            results = dev.execute(new_batch, execution_config=execution_config)
+            return post_processing_fn(results)
+
+    >>> f(jax.numpy.array(1.2))
+    DeviceArray(0.36235774, dtype=float32)
+    >>> jax.grad(f)(jax.numpy.array(1.2))
+    DeviceArray(-0.93203914, dtype=float32, weak_type=True)
 
-    Args:
-        wires (int, Iterable[Number, str]): Number of subsystems represented by the device,
-            or iterable that contains unique labels for the subsystems as numbers (i.e., ``[-1, 0, 2]``)
-            or strings (``['ancilla', 'q1', 'q2']``). Default 1 if not specified.
-        shots (None, int): How many times the circuit should be evaluated (or sampled) to estimate
-            the expectation values. Defaults to ``None`` if not specified, which means that the device
-            returns analytical results.
     """
 
-    name = "Default qubit PennyLane plugin"
-    short_name = "default.qubit"
-    pennylane_requires = __version__
-    version = __version__
-    author = "Xanadu Inc."
-
-    operations = {
-        "GlobalPhase",
-        "Identity",
-        "Snapshot",
-        "BasisState",
-        "StatePrep",
-        "QubitStateVector",
-        "QubitUnitary",
-        "ControlledQubitUnitary",
-        "BlockEncode",
-        "MultiControlledX",
-        "IntegerComparator",
-        "DiagonalQubitUnitary",
-        "PauliX",
-        "PauliY",
-        "PauliZ",
-        "MultiRZ",
-        "Hadamard",
-        "S",
-        "Adjoint(S)",
-        "T",
-        "Adjoint(T)",
-        "SX",
-        "Adjoint(SX)",
-        "CNOT",
-        "SWAP",
-        "ISWAP",
-        "PSWAP",
-        "Adjoint(ISWAP)",
-        "SISWAP",
-        "Adjoint(SISWAP)",
-        "SQISW",
-        "CSWAP",
-        "Toffoli",
-        "CCZ",
-        "CY",
-        "CZ",
-        "CH",
-        "PhaseShift",
-        "PCPhase",
-        "ControlledPhaseShift",
-        "CPhaseShift00",
-        "CPhaseShift01",
-        "CPhaseShift10",
-        "CPhase",
-        "RX",
-        "RY",
-        "RZ",
-        "Rot",
-        "CRX",
-        "CRY",
-        "CRZ",
-        "CRot",
-        "IsingXX",
-        "IsingYY",
-        "IsingZZ",
-        "IsingXY",
-        "SingleExcitation",
-        "SingleExcitationPlus",
-        "SingleExcitationMinus",
-        "DoubleExcitation",
-        "DoubleExcitationPlus",
-        "DoubleExcitationMinus",
-        "QubitCarry",
-        "QubitSum",
-        "OrbitalRotation",
-        "QFT",
-        "ECR",
-    }
-
-    observables = {
-        "PauliX",
-        "PauliY",
-        "PauliZ",
-        "Hadamard",
-        "Hermitian",
-        "Identity",
-        "Projector",
-        "SparseHamiltonian",
-        "Hamiltonian",
-        "Sum",
-        "SProd",
-        "Prod",
-        "Exp",
-        "Evolution",
-    }
-
-    def __init__(
-        self, wires, *, r_dtype=np.float64, c_dtype=np.complex128, shots=None, analytic=None
-    ):
-        super().__init__(wires, shots, r_dtype=r_dtype, c_dtype=c_dtype, analytic=analytic)
-        self._debugger = None
-
-        # Create the initial state. Internally, we store the
-        # state as an array of dimension [2]*wires.
-        self._state = self._create_basis_state(0)
-        self._pre_rotated_state = self._state
-
-        self._apply_ops = {
-            "PauliX": self._apply_x,
-            "PauliY": self._apply_y,
-            "PauliZ": self._apply_z,
-            "Hadamard": self._apply_hadamard,
-            "S": self._apply_s,
-            "T": self._apply_t,
-            "SX": self._apply_sx,
-            "CNOT": self._apply_cnot,
-            "SWAP": self._apply_swap,
-            "CZ": self._apply_cz,
-            "Toffoli": self._apply_toffoli,
-        }
-
     @property
-    def stopping_condition(self):
-        def accepts_obj(obj):
-            if obj.name == "QFT" and len(obj.wires) >= 6:
-                return False
-            if obj.name == "GroverOperator" and len(obj.wires) >= 13:
-                return False
-            if getattr(obj, "has_matrix", False):
-                # pow operations dont work with backprop or adjoint without decomposition
-                # use class name string so we don't need to use isinstance check
-                return not (obj.__class__.__name__ == "Pow" and qml.operation.is_trainable(obj))
-            return obj.name in self.observables.union(self.operations)
-
-        return qml.BooleanFn(accepts_obj)
-
-    @functools.lru_cache()
-    def map_wires(self, wires):
-        # temporarily overwrite this method to bypass
-        # wire map that produces Wires objects
-        try:
-            mapped_wires = [self.wire_map[w] for w in wires]
-        except KeyError as e:
-            raise WireError(
-                f"Did not find some of the wires {wires.labels} on device with wires {self.wires.labels}."
-            ) from e
-
-        return mapped_wires
-
-    def define_wire_map(self, wires):
-        # temporarily overwrite this method to bypass
-        # wire map that produces Wires objects
-        consecutive_wires = range(self.num_wires)
-        wire_map = zip(wires, consecutive_wires)
-        return dict(wire_map)
-
-    # pylint: disable=arguments-differ
-    def _get_batch_size(self, tensor, expected_shape, expected_size):
-        """Determine whether a tensor has an additional batch dimension for broadcasting,
-        compared to an expected_shape."""
-        size = self._size(tensor)
-        if self._ndim(tensor) > len(expected_shape) or size > expected_size:
-            return size // expected_size
-
-        return None
-
-    # pylint: disable=arguments-differ
-    def apply(self, operations, rotations=None, **kwargs):
-        rotations = rotations or []
-
-        # apply the circuit operations
-        for i, operation in enumerate(operations):
-            if i > 0 and isinstance(operation, (StatePrep, BasisState)):
-                raise DeviceError(
-                    f"Operation {operation.name} cannot be used after other Operations have already been applied "
-                    f"on a {self.short_name} device."
-                )
-
-            if isinstance(operation, StatePrep):
-                self._apply_state_vector(operation.parameters[0], operation.wires)
-            elif isinstance(operation, BasisState):
-                self._apply_basis_state(operation.parameters[0], operation.wires)
-            elif isinstance(operation, Snapshot):
-                if self._debugger and self._debugger.active:
-                    state_vector = np.array(self._flatten(self._state))
-                    if operation.tag:
-                        self._debugger.snapshots[operation.tag] = state_vector
-                    else:
-                        self._debugger.snapshots[len(self._debugger.snapshots)] = state_vector
-            elif isinstance(operation, ParametrizedEvolution):
-                self._state = self._apply_parametrized_evolution(self._state, operation)
-            else:
-                self._state = self._apply_operation(self._state, operation)
-
-        # store the pre-rotated state
-        self._pre_rotated_state = self._state
-
-        # apply the circuit rotations
-        for operation in rotations:
-            self._state = self._apply_operation(self._state, operation)
-
-    def _apply_parametrized_evolution(self, state: TensorLike, operation: ParametrizedEvolution):
-        """Applies a parametrized evolution to the input state.
-
-        Args:
-            state (array[complex]): input state
-            operation (ParametrizedEvolution): operation to apply on the state
-        """
-        raise NotImplementedError(
-            f"The device {self.short_name} cannot execute a ParametrizedEvolution operation. "
-            "Please use the jax interface."
-        )
-
-    def _apply_operation(self, state, operation):
-        """Applies operations to the input state.
-
-        Args:
-            state (array[complex]): input state
-            operation (~.Operation): operation to apply on the device
-
-        Returns:
-            array[complex]: output state
-        """
-        if operation.__class__.__name__ == "Identity":
-            return state
-        if operation.name == "GlobalPhase":
-            return self._apply_global_phase(state, operation)
-        wires = operation.wires
-
-        if str(operation.name) in self._apply_ops:  # cast to string because of Tensor
-            shift = int(self._ndim(state) > self.num_wires)
-            axes = [ax + shift for ax in self.wires.indices(wires)]
-            return self._apply_ops[operation.name](state, axes)
-
-        matrix = self._asarray(self._get_unitary_matrix(operation), dtype=self.C_DTYPE)
-
-        if operation in diagonal_in_z_basis:
-            return self._apply_diagonal_unitary(state, matrix, wires)
-        if len(wires) <= 2:
-            # Einsum is faster for small gates
-            return self._apply_unitary_einsum(state, matrix, wires)
-
-        return self._apply_unitary(state, matrix, wires)
-
-    def _apply_global_phase(self, state, operation: qml.GlobalPhase):  # pylint: disable=no-self-use
-        """Applies a :class:`~.GlobalPhase` operation to the state."""
-        return qml.math.exp(-1j * operation.data[0]) * state
+    def name(self):
+        """The name of the device."""
+        return "default.qubit.2"
+
+    def __init__(self, wires=None, shots=None, seed="global", max_workers=None) -> None:
+        super().__init__(wires=wires, shots=shots)
+        self._max_workers = max_workers
+        seed = np.random.randint(0, high=10000000) if seed == "global" else seed
+        self._rng = np.random.default_rng(seed)
+        self._debugger = None
 
-    def _apply_x(self, state, axes, **kwargs):
-        """Applies a PauliX gate by rolling 1 unit along the axis specified in ``axes``.
+    def supports_derivatives(
+        self,
+        execution_config: Optional[ExecutionConfig] = None,
+        circuit: Optional[QuantumTape] = None,
+    ) -> bool:
+        """Check whether or not derivatives are available for a given configuration and circuit.
 
-        Rolling by 1 unit along the axis means that the :math:`|0 \rangle` state with index ``0`` is
-        shifted to the :math:`|1 \rangle` state with index ``1``. Likewise, since rolling beyond
-        the last index loops back to the first, :math:`|1 \rangle` is transformed to
-        :math:`|0\rangle`.
+        ``DefaultQubit`` supports backpropagation derivatives with analytic results, as well as
+        adjoint differentiation.
 
         Args:
-            state (array[complex]): input state
-            axes (List[int]): target axes to apply transformation
+            execution_config (ExecutionConfig): The configuration of the desired derivative calculation
+            circuit (QuantumTape): An optional circuit to check derivatives support for.
 
         Returns:
-            array[complex]: output state
-        """
-        return self._roll(state, 1, axes[0])
-
-    def _apply_y(self, state, axes, **kwargs):
-        """Applies a PauliY gate by adding a negative sign to the 1 index along the axis specified
-        in ``axes``, rolling one unit along the same axis, and multiplying the result by 1j.
-
-        Args:
-            state (array[complex]): input state
-            axes (List[int]): target axes to apply transformation
+            Bool: Whether or not a derivative can be calculated provided the given information
 
-        Returns:
-            array[complex]: output state
         """
-        return 1j * self._apply_x(self._apply_z(state, axes), axes)
-
-    def _apply_z(self, state, axes, **kwargs):
-        """Applies a PauliZ gate by adding a negative sign to the 1 index along the axis specified
-        in ``axes``.
-
-        Args:
-            state (array[complex]): input state
-            axes (List[int]): target axes to apply transformation
-
-        Returns:
-            array[complex]: output state
-        """
-        return self._apply_phase(state, axes, -1)
+        if execution_config is None:
+            return True
+        # backpropagation currently supported for all supported circuits
+        # will later need to add logic if backprop requested with finite shots
+        # do once device accepts finite shots
+        if (
+            execution_config.gradient_method == "backprop"
+            and execution_config.device_options.get("max_workers", self._max_workers) is None
+        ):
+            return True
 
-    def _apply_hadamard(self, state, axes, **kwargs):
-        """Apply the Hadamard gate by combining the results of applying the PauliX and PauliZ gates.
+        if execution_config.gradient_method == "adjoint" and execution_config.use_device_gradient:
+            if circuit is None:
+                return True
 
-        Args:
-            state (array[complex]): input state
-            axes (List[int]): target axes to apply transformation
+            return isinstance(validate_and_expand_adjoint(circuit)[0][0], QuantumScript)
 
-        Returns:
-            array[complex]: output state
-        """
-        state_x = self._apply_x(state, axes)
-        state_z = self._apply_z(state, axes)
-        return self._const_mul(SQRT2INV, state_x + state_z)
+        return False
 
-    def _apply_s(self, state, axes, inverse=False):
-        return self._apply_phase(state, axes, 1j, inverse)
-
-    def _apply_t(self, state, axes, inverse=False):
-        return self._apply_phase(state, axes, TPHASE, inverse)
-
-    def _apply_sx(self, state, axes, inverse=False):
-        """Apply the Square Root X gate.
+    def preprocess(
+        self,
+        execution_config: ExecutionConfig = DefaultExecutionConfig,
+    ) -> Tuple[QuantumTapeBatch, PostprocessingFn, ExecutionConfig]:
+        """This function defines the device transform program to be applied and an updated device configuration.
 
         Args:
-            state (array[complex]): input state
-            axes (List[int]): target axes to apply transformation
+            execution_config (Union[ExecutionConfig, Sequence[ExecutionConfig]]): A data structure describing the
+                parameters needed to fully describe the execution.
 
         Returns:
-            array[complex]: output state
-        """
-        if inverse:
-            return 0.5 * ((1 - 1j) * state + (1 + 1j) * self._apply_x(state, axes))
+            TransformProgram, ExecutionConfig: A transform program that when called returns QuantumTapes that the device
+            can natively execute as well as a postprocessing function to be called after execution, and a configuration with
+            unset specifications filled in.
 
-        return 0.5 * ((1 + 1j) * state + (1 - 1j) * self._apply_x(state, axes))
+        This device:
 
-    def _apply_cnot(self, state, axes, **kwargs):
-        """Applies a CNOT gate by slicing along the first axis specified in ``axes`` and then
-        applying an X transformation along the second axis.
+        * Supports any qubit operations that provide a matrix
+        * Currently does not support finite shots
+        * Currently does not intrinsically support parameter broadcasting
 
-        By slicing along the first axis, we are able to select all of the amplitudes with a
-        corresponding :math:`|1\rangle` for the control qubit. This means we then just need to apply
-        a :class:`~.PauliX` (NOT) gate to the result.
-
-        Args:
-            state (array[complex]): input state
-            axes (List[int]): target axes to apply transformation
-
-        Returns:
-            array[complex]: output state
         """
-        ndim = self._ndim(state)
-        sl_0 = _get_slice(0, axes[0], ndim)
-        sl_1 = _get_slice(1, axes[0], ndim)
-
-        # We will be slicing into the state according to state[sl_1], giving us all of the
-        # amplitudes with a |1> for the control qubit. The resulting array has lost an axis
-        # relative to state and we need to be careful about the axis we apply the PauliX rotation
-        # to. If axes[1] is larger than axes[0], then we need to shift the target axis down by
-        # one, otherwise we can leave as-is. For example: a state has [0, 1, 2, 3], control=1,
-        # target=3. Then, state[sl_1] has 3 axes and target=3 now corresponds to the second axis.
-        if axes[1] > axes[0]:
-            target_axes = [axes[1] - 1]
+        transform_program = TransformProgram()
+        # Validate device wires
+        transform_program.add_transform(validate_device_wires, self)
+
+        # Validate multi processing
+        max_workers = execution_config.device_options.get("max_workers", self._max_workers)
+        transform_program.add_transform(validate_multiprocessing_workers, max_workers, self)
+
+        # General preprocessing (Validate measurement, expand, adjoint expand, broadcast expand)
+        transform_program_preprocess, config = preprocess(execution_config=execution_config)
+        transform_program = transform_program + transform_program_preprocess
+        return transform_program, config
+
+    def execute(
+        self,
+        circuits: QuantumTape_or_Batch,
+        execution_config: ExecutionConfig = DefaultExecutionConfig,
+    ) -> Result_or_ResultBatch:
+        is_single_circuit = False
+        if isinstance(circuits, QuantumScript):
+            is_single_circuit = True
+            circuits = [circuits]
+
+        if self.tracker.active:
+            for c in circuits:
+                self.tracker.update(resources=c.specs["resources"])
+            self.tracker.update(batches=1, executions=len(circuits))
+            self.tracker.record()
+
+        max_workers = execution_config.device_options.get("max_workers", self._max_workers)
+        interface = (
+            execution_config.interface
+            if execution_config.gradient_method in {"backprop", None}
+            else None
+        )
+        if max_workers is None:
+            results = tuple(
+                simulate(c, rng=self._rng, debugger=self._debugger, interface=interface)
+                for c in circuits
+            )
         else:
-            target_axes = [axes[1]]
-
-        state_x = self._apply_x(state[sl_1], axes=target_axes)
-        return self._stack([state[sl_0], state_x], axis=axes[0])
+            vanilla_circuits = [convert_to_numpy_parameters(c) for c in circuits]
+            seeds = self._rng.integers(2**31 - 1, size=len(vanilla_circuits))
+            _wrap_simulate = partial(simulate, debugger=None, interface=interface)
+            with concurrent.futures.ProcessPoolExecutor(max_workers=max_workers) as executor:
+                exec_map = executor.map(_wrap_simulate, vanilla_circuits, seeds)
+                results = tuple(exec_map)
+
+            # reset _rng to mimic serial behavior
+            self._rng = np.random.default_rng(self._rng.integers(2**31 - 1))
+
+        return results[0] if is_single_circuit else results
+
+    def compute_derivatives(
+        self,
+        circuits: QuantumTape_or_Batch,
+        execution_config: ExecutionConfig = DefaultExecutionConfig,
+    ):
+        is_single_circuit = False
+        if isinstance(circuits, QuantumScript):
+            is_single_circuit = True
+            circuits = [circuits]
+
+        if self.tracker.active:
+            self.tracker.update(derivative_batches=1, derivatives=len(circuits))
+            self.tracker.record()
+
+        max_workers = execution_config.device_options.get("max_workers", self._max_workers)
+        if max_workers is None:
+            res = tuple(adjoint_jacobian(circuit) for circuit in circuits)
+        else:
+            vanilla_circuits = [convert_to_numpy_parameters(c) for c in circuits]
+            with concurrent.futures.ProcessPoolExecutor(max_workers=max_workers) as executor:
+                exec_map = executor.map(adjoint_jacobian, vanilla_circuits)
+                res = tuple(exec_map)
 
-    def _apply_toffoli(self, state, axes, **kwargs):
-        """Applies a Toffoli gate by slicing along the axis of the greater control qubit, slicing
-        each of the resulting sub-arrays along the axis of the smaller control qubit, and then applying
-        an X transformation along the axis of the target qubit of the fourth sub-sub-array.
+            # reset _rng to mimic serial behavior
+            self._rng = np.random.default_rng(self._rng.integers(2**31 - 1))
 
-        By performing two consecutive slices in this way, we are able to select all of the amplitudes with
-        a corresponding :math:`|11\rangle` for the two control qubits. This means we then just need to apply
-        a :class:`~.PauliX` (NOT) gate to the result.
+        return res[0] if is_single_circuit else res
 
-        Args:
-            state (array[complex]): input state
-            axes (List[int]): target axes to apply transformation
+    def execute_and_compute_derivatives(
+        self,
+        circuits: QuantumTape_or_Batch,
+        execution_config: ExecutionConfig = DefaultExecutionConfig,
+    ):
+        is_single_circuit = False
+        if isinstance(circuits, QuantumScript):
+            is_single_circuit = True
+            circuits = [circuits]
+
+        if self.tracker.active:
+            for c in circuits:
+                self.tracker.update(resources=c.specs["resources"])
+            self.tracker.update(
+                execute_and_derivative_batches=1,
+                executions=len(circuits),
+                derivatives=len(circuits),
+            )
+            self.tracker.record()
 
-        Returns:
-            array[complex]: output state
-        """
-        cntrl_max = np.argmax(axes[:2])
-        cntrl_min = cntrl_max ^ 1
-        ndim = self._ndim(state)
-        sl_a0 = _get_slice(0, axes[cntrl_max], ndim)
-        sl_a1 = _get_slice(1, axes[cntrl_max], ndim)
-        sl_b0 = _get_slice(0, axes[cntrl_min], ndim - 1)
-        sl_b1 = _get_slice(1, axes[cntrl_min], ndim - 1)
-
-        # If both controls are smaller than the target, shift the target axis down by two. If one
-        # control is greater and one control is smaller than the target, shift the target axis
-        # down by one. If both controls are greater than the target, leave the target axis as-is.
-        if axes[cntrl_min] > axes[2]:
-            target_axes = [axes[2]]
-        elif axes[cntrl_max] > axes[2]:
-            target_axes = [axes[2] - 1]
+        max_workers = execution_config.device_options.get("max_workers", self._max_workers)
+        if max_workers is None:
+            results = tuple(
+                _adjoint_jac_wrapper(c, rng=self._rng, debugger=self._debugger) for c in circuits
+            )
         else:
-            target_axes = [axes[2] - 2]
+            vanilla_circuits = [convert_to_numpy_parameters(c) for c in circuits]
+            seeds = self._rng.integers(2**31 - 1, size=len(vanilla_circuits))
 
-        # state[sl_a1][sl_b1] gives us all of the amplitudes with a |11> for the two control qubits.
-        state_x = self._apply_x(state[sl_a1][sl_b1], axes=target_axes)
-        state_stacked_a1 = self._stack([state[sl_a1][sl_b0], state_x], axis=axes[cntrl_min])
-        return self._stack([state[sl_a0], state_stacked_a1], axis=axes[cntrl_max])
+            with concurrent.futures.ProcessPoolExecutor(max_workers=max_workers) as executor:
+                results = tuple(executor.map(_adjoint_jac_wrapper, vanilla_circuits, seeds))
 
-    def _apply_swap(self, state, axes, **kwargs):
-        """Applies a SWAP gate by performing a partial transposition along the specified axes.
+            # reset _rng to mimic serial behavior
+            self._rng = np.random.default_rng(self._rng.integers(2**31 - 1))
 
-        Args:
-            state (array[complex]): input state
-            axes (List[int]): target axes to apply transformation
+        results, jacs = tuple(zip(*results))
+        return (results[0], jacs[0]) if is_single_circuit else (results, jacs)
 
-        Returns:
-            array[complex]: output state
-        """
-        all_axes = list(range(len(state.shape)))
-        all_axes[axes[0]] = axes[1]
-        all_axes[axes[1]] = axes[0]
-        return self._transpose(state, all_axes)
+    def supports_jvp(
+        self,
+        execution_config: Optional[ExecutionConfig] = None,
+        circuit: Optional[QuantumTape] = None,
+    ) -> bool:
+        """Whether or not this device defines a custom jacobian vector product.
 
-    def _apply_cz(self, state, axes, **kwargs):
-        """Applies a CZ gate by slicing along the first axis specified in ``axes`` and then
-        applying a Z transformation along the second axis.
+        ``DefaultQubit`` supports backpropagation derivatives with analytic results, as well as
+        adjoint differentiation.
 
         Args:
-            state (array[complex]): input state
-            axes (List[int]): target axes to apply transformation
+            execution_config (ExecutionConfig): The configuration of the desired derivative calculation
+            circuit (QuantumTape): An optional circuit to check derivatives support for.
 
         Returns:
-            array[complex]: output state
+            bool: Whether or not a derivative can be calculated provided the given information
         """
-        ndim = self._ndim(state)
-        sl_0 = _get_slice(0, axes[0], ndim)
-        sl_1 = _get_slice(1, axes[0], ndim)
+        return self.supports_derivatives(execution_config, circuit)
 
-        if axes[1] > axes[0]:
-            target_axes = [axes[1] - 1]
+    def compute_jvp(
+        self,
+        circuits: QuantumTape_or_Batch,
+        tangents: Tuple[Number],
+        execution_config: ExecutionConfig = DefaultExecutionConfig,
+    ):
+        is_single_circuit = False
+        if isinstance(circuits, QuantumScript):
+            is_single_circuit = True
+            circuits = [circuits]
+            tangents = [tangents]
+
+        if self.tracker.active:
+            self.tracker.update(jvp_batches=1, jvps=len(circuits))
+            self.tracker.record()
+
+        max_workers = execution_config.device_options.get("max_workers", self._max_workers)
+        if max_workers is None:
+            res = tuple(adjoint_jvp(circuit, tans) for circuit, tans in zip(circuits, tangents))
         else:
-            target_axes = [axes[1]]
-
-        state_z = self._apply_z(state[sl_1], axes=target_axes)
-        return self._stack([state[sl_0], state_z], axis=axes[0])
+            vanilla_circuits = [convert_to_numpy_parameters(c) for c in circuits]
+            with concurrent.futures.ProcessPoolExecutor(max_workers=max_workers) as executor:
+                res = tuple(executor.map(adjoint_jvp, vanilla_circuits, tangents))
 
-    def _apply_phase(self, state, axes, parameters, inverse=False):
-        """Applies a phase onto the 1 index along the axis specified in ``axes``.
+            # reset _rng to mimic serial behavior
+            self._rng = np.random.default_rng(self._rng.integers(2**31 - 1))
 
-        Args:
-            state (array[complex]): input state
-            axes (List[int]): target axes to apply transformation
-            parameters (float): phase to apply
-            inverse (bool): whether to apply the inverse phase
+        return res[0] if is_single_circuit else res
 
-        Returns:
-            array[complex]: output state
-        """
-        ndim = self._ndim(state)
-        sl_0 = _get_slice(0, axes[0], ndim)
-        sl_1 = _get_slice(1, axes[0], ndim)
-
-        phase = self._conj(parameters) if inverse else parameters
-        return self._stack([state[sl_0], self._const_mul(phase, state[sl_1])], axis=axes[0])
-
-    def expval(self, observable, shot_range=None, bin_size=None):
-        """Returns the expectation value of a Hamiltonian observable. When the observable is a
-        ``Hamiltonian`` or ``SparseHamiltonian`` object, the expectation value is computed directly
-        from the sparse matrix representation, which leads to faster execution.
-
-        Args:
-            observable (~.Observable): a PennyLane observable
-            shot_range (tuple[int]): 2-tuple of integers specifying the range of samples
-                to use. If not specified, all samples are used.
-            bin_size (int): Divides the shot range into bins of size ``bin_size``, and
-                returns the measurement statistic separately over each bin. If not
-                provided, the entire shot range is treated as a single bin.
-
-        Returns:
-            float: returns the expectation value of the observable
-
-        .. warning::
-
-            This function does not support broadcasting if ``observable`` is a
-            :class:``~.Hamiltonian`` and the device interface or the interface of the
-            Hamiltonian is not NumPy or Autograd
-
-        """
-        is_state_batched = self._ndim(self.state) == 2
-        # intercept Sums
-        if isinstance(observable, Sum) and not self.shots:
-            return measure(
-                ExpectationMP(observable.map_wires(self.wire_map)),
-                self._pre_rotated_state,
-                is_state_batched,
+    def execute_and_compute_jvp(
+        self,
+        circuits: QuantumTape_or_Batch,
+        tangents: Tuple[Number],
+        execution_config: ExecutionConfig = DefaultExecutionConfig,
+    ):
+        is_single_circuit = False
+        if isinstance(circuits, QuantumScript):
+            is_single_circuit = True
+            circuits = [circuits]
+            tangents = [tangents]
+
+        if self.tracker.active:
+            for c in circuits:
+                self.tracker.update(resources=c.specs["resources"])
+            self.tracker.update(
+                execute_and_jvp_batches=1, executions=len(circuits), jvps=len(circuits)
             )
+            self.tracker.record()
 
-        # intercept other Hamiltonians
-        # TODO: Ideally, this logic should not live in the Device, but be moved
-        # to a component that can be re-used by devices as needed.
-        if isinstance(observable, MeasurementValue) or observable.name not in (
-            "Hamiltonian",
-            "SparseHamiltonian",
-        ):
-            return super().expval(observable, shot_range=shot_range, bin_size=bin_size)
-
-        assert self.shots is None, f"{observable.name} must be used with shots=None"
-
-        self.map_wires(observable.wires)
-        backprop_mode = (
-            not isinstance(self.state, np.ndarray)
-            or any(not isinstance(d, (float, np.ndarray)) for d in observable.data)
-        ) and observable.name == "Hamiltonian"
-
-        if backprop_mode:
-            # TODO[dwierichs]: This branch is not adapted to broadcasting yet
-            if is_state_batched:
-                raise NotImplementedError(
-                    "Expectation values of Hamiltonians for interface!=None are "
-                    "not supported together with parameter broadcasting yet"
-                )
-            # We must compute the expectation value assuming that the Hamiltonian
-            # coefficients *and* the quantum states are tensor objects.
-
-            # Compute  <psi| H |psi> via sum_i coeff_i * <psi| PauliWord |psi> using a sparse
-            # representation of the Pauliword
-            res = qml.math.cast(qml.math.convert_like(0.0, observable.data), dtype=complex)
-            interface = qml.math.get_interface(self.state)
-
-            # Note: it is important that we use the Hamiltonian's data and not the coeffs
-            # attribute. This is because the .data attribute may be 'unwrapped' as required by
-            # the interfaces, whereas the .coeff attribute will always be the same input dtype
-            # that the user provided.
-            for op, coeff in zip(observable.ops, observable.data):
-                # extract a scipy.sparse.coo_matrix representation of this Pauli word
-                coo = qml.operation.Tensor(op).sparse_matrix(wire_order=self.wires, format="coo")
-                Hmat = qml.math.cast(qml.math.convert_like(coo.data, self.state), self.C_DTYPE)
-
-                product = (
-                    self._gather(self._conj(self.state), coo.row)
-                    * Hmat
-                    * self._gather(self.state, coo.col)
-                )
-                c = qml.math.convert_like(coeff, product)
-
-                if interface == "tensorflow":
-                    c = qml.math.cast(c, "complex128")
-
-                res = qml.math.convert_like(res, product) + qml.math.sum(c * product)
-
+        max_workers = execution_config.device_options.get("max_workers", self._max_workers)
+        if max_workers is None:
+            results = tuple(
+                _adjoint_jvp_wrapper(c, t, rng=self._rng, debugger=self._debugger)
+                for c, t in zip(circuits, tangents)
+            )
         else:
-            # Coefficients and the state are not trainable, we can be more
-            # efficient in how we compute the Hamiltonian sparse matrix.
-            Hmat = observable.sparse_matrix(wire_order=self.wires)
-
-            state = qml.math.toarray(self.state)
-            if is_state_batched:
-                res = qml.math.array(
-                    [
-                        csr_matrix.dot(
-                            csr_matrix(self._conj(_state)),
-                            csr_matrix.dot(Hmat, csr_matrix(_state[..., None])),
-                        ).toarray()[0]
-                        for _state in state
-                    ]
+            vanilla_circuits = [convert_to_numpy_parameters(c) for c in circuits]
+            seeds = self._rng.integers(2**31 - 1, size=len(vanilla_circuits))
+
+            with concurrent.futures.ProcessPoolExecutor(max_workers=max_workers) as executor:
+                results = tuple(
+                    executor.map(_adjoint_jvp_wrapper, vanilla_circuits, tangents, seeds)
                 )
-            else:
-                res = csr_matrix.dot(
-                    csr_matrix(self._conj(state)),
-                    csr_matrix.dot(Hmat, csr_matrix(state[..., None])),
-                ).toarray()[0]
 
-        if observable.name == "Hamiltonian":
-            res = qml.math.squeeze(res)
+            # reset _rng to mimic serial behavior
+            self._rng = np.random.default_rng(self._rng.integers(2**31 - 1))
 
-        return self._real(res)
+        results, jvps = tuple(zip(*results))
+        return (results[0], jvps[0]) if is_single_circuit else (results, jvps)
 
-    def _get_unitary_matrix(self, unitary):  # pylint: disable=no-self-use
-        """Return the matrix representing a unitary operation.
+    def supports_vjp(
+        self,
+        execution_config: Optional[ExecutionConfig] = None,
+        circuit: Optional[QuantumTape] = None,
+    ) -> bool:
+        """Whether or not this device defines a custom vector jacobian product.
 
-        Args:
-            unitary (~.Operation): a PennyLane unitary operation
-
-        Returns:
-            array[complex]: Returns a 2D matrix representation of
-            the unitary in the computational basis, or, in the case of a diagonal unitary,
-            a 1D array representing the matrix diagonal.
-        """
-        if unitary in diagonal_in_z_basis:
-            return unitary.eigvals()
-
-        return unitary.matrix()
-
-    @classmethod
-    def capabilities(cls):
-        capabilities = super().capabilities().copy()
-        capabilities.update(
-            model="qubit",
-            supports_inverse_operations=True,
-            supports_analytic_computation=True,
-            supports_broadcasting=True,
-            returns_state=True,
-            passthru_devices={
-                "tf": "default.qubit.tf",
-                "torch": "default.qubit.torch",
-                "autograd": "default.qubit.autograd",
-                "jax": "default.qubit.jax",
-            },
-        )
-        return capabilities
-
-    def _create_basis_state(self, index):
-        """Return a computational basis state over all wires.
+        ``DefaultQubit`` supports backpropagation derivatives with analytic results, as well as
+        adjoint differentiation.
 
         Args:
-            index (int): integer representing the computational basis state
+            execution_config (ExecutionConfig): A description of the hyperparameters for the desired computation.
+            circuit (None, QuantumTape): A specific circuit to check differentation for.
 
         Returns:
-            array[complex]: complex array of shape ``[2]*self.num_wires``
-            representing the statevector of the basis state
-
-        Note: This function does not support broadcasted inputs yet.
-        """
-        state = np.zeros(2**self.num_wires, dtype=np.complex128)
-        state[index] = 1
-        state = self._asarray(state, dtype=self.C_DTYPE)
-        return self._reshape(state, [2] * self.num_wires)
-
-    @property
-    def state(self):
-        dim = 2**self.num_wires
-        batch_size = self._get_batch_size(self._pre_rotated_state, (2,) * self.num_wires, dim)
-        # Do not flatten the state completely but leave the broadcasting dimension if there is one
-        shape = (batch_size, dim) if batch_size is not None else (dim,)
-        return self._reshape(self._pre_rotated_state, shape)
-
-    def _apply_state_vector(self, state, device_wires):
-        """Initialize the internal state vector in a specified state.
-
-        Args:
-            state (array[complex]): normalized input state of length ``2**len(wires)``
-                or broadcasted state of shape ``(batch_size, 2**len(wires))``
-            device_wires (Wires): wires that get initialized in the state
+            bool: Whether or not a derivative can be calculated provided the given information
         """
+        return self.supports_derivatives(execution_config, circuit)
 
-        # translate to wire labels used by device
-        device_wires = self.map_wires(device_wires)
-        dim = 2 ** len(device_wires)
-
-        state = self._asarray(state, dtype=self.C_DTYPE)
-        batch_size = self._get_batch_size(state, (dim,), dim)
-        output_shape = [2] * self.num_wires
-        if batch_size is not None:
-            output_shape.insert(0, batch_size)
-
-        if len(device_wires) == self.num_wires and sorted(device_wires) == device_wires:
-            # Initialize the entire device state with the input state
-            self._state = self._reshape(state, output_shape)
-            return
+    def compute_vjp(
+        self,
+        circuits: QuantumTape_or_Batch,
+        cotangents: Tuple[Number],
+        execution_config: ExecutionConfig = DefaultExecutionConfig,
+    ):
+        is_single_circuit = False
+        if isinstance(circuits, QuantumScript):
+            is_single_circuit = True
+            circuits = [circuits]
+            cotangents = [cotangents]
+
+        if self.tracker.active:
+            self.tracker.update(vjp_batches=1, vjps=len(circuits))
+            self.tracker.record()
+
+        max_workers = execution_config.device_options.get("max_workers", self._max_workers)
+        if max_workers is None:
+            res = tuple(adjoint_vjp(circuit, cots) for circuit, cots in zip(circuits, cotangents))
+        else:
+            vanilla_circuits = [convert_to_numpy_parameters(c) for c in circuits]
+            with concurrent.futures.ProcessPoolExecutor(max_workers=max_workers) as executor:
+                res = tuple(executor.map(adjoint_vjp, vanilla_circuits, cotangents))
 
-        # generate basis states on subset of qubits via the cartesian product
-        basis_states = np.array(list(itertools.product([0, 1], repeat=len(device_wires))))
+            # reset _rng to mimic serial behavior
+            self._rng = np.random.default_rng(self._rng.integers(2**31 - 1))
 
-        # get basis states to alter on full set of qubits
-        unravelled_indices = np.zeros((2 ** len(device_wires), self.num_wires), dtype=int)
-        unravelled_indices[:, device_wires] = basis_states
+        return res[0] if is_single_circuit else res
 
-        # get indices for which the state is changed to input state vector elements
-        ravelled_indices = np.ravel_multi_index(unravelled_indices.T, [2] * self.num_wires)
+    def execute_and_compute_vjp(
+        self,
+        circuits: QuantumTape_or_Batch,
+        cotangents: Tuple[Number],
+        execution_config: ExecutionConfig = DefaultExecutionConfig,
+    ):
+        is_single_circuit = False
+        if isinstance(circuits, QuantumScript):
+            is_single_circuit = True
+            circuits = [circuits]
+            cotangents = [cotangents]
+
+        if self.tracker.active:
+            for c in circuits:
+                self.tracker.update(resources=c.specs["resources"])
+            self.tracker.update(
+                execute_and_vjp_batches=1, executions=len(circuits), vjps=len(circuits)
+            )
+            self.tracker.record()
 
-        if batch_size is not None:
-            state = self._scatter(
-                (slice(None), ravelled_indices), state, [batch_size, 2**self.num_wires]
+        max_workers = execution_config.device_options.get("max_workers", self._max_workers)
+        if max_workers is None:
+            results = tuple(
+                _adjoint_vjp_wrapper(c, t, rng=self._rng, debugger=self._debugger)
+                for c, t in zip(circuits, cotangents)
             )
         else:
-            state = self._scatter(ravelled_indices, state, [2**self.num_wires])
-        state = self._reshape(state, output_shape)
-        self._state = self._asarray(state, dtype=self.C_DTYPE)
-
-    def _apply_basis_state(self, state, wires):
-        """Initialize the state vector in a specified computational basis state.
-
-        Args:
-            state (array[int]): computational basis state of shape ``(wires,)``
-                consisting of 0s and 1s.
-            wires (Wires): wires that the provided computational state should be initialized on
-
-        Note: This function does not support broadcasted inputs yet.
-        """
-        # translate to wire labels used by device
-        device_wires = self.map_wires(wires)
-
-        # length of basis state parameter
-        n_basis_state = len(state)
-
-        if not set(state.tolist()).issubset({0, 1}):
-            raise ValueError("BasisState parameter must consist of 0 or 1 integers.")
-
-        if n_basis_state != len(device_wires):
-            raise ValueError("BasisState parameter and wires must be of equal length.")
-
-        # get computational basis state number
-        basis_states = 2 ** (self.num_wires - 1 - np.array(device_wires))
-        basis_states = qml.math.convert_like(basis_states, state)
-        num = int(qml.math.dot(state, basis_states))
-
-        self._state = self._create_basis_state(num)
-
-    def _apply_unitary(self, state, mat, wires):
-        r"""Apply multiplication of a matrix to subsystems of the quantum state.
-
-        Args:
-            state (array[complex]): input state
-            mat (array): matrix to multiply
-            wires (Wires): target wires
-
-        Returns:
-            array[complex]: output state
+            vanilla_circuits = [convert_to_numpy_parameters(c) for c in circuits]
+            seeds = self._rng.integers(2**31 - 1, size=len(vanilla_circuits))
 
-        Note: This function does not support simultaneously broadcasted states and matrices yet.
-        """
-        # translate to wire labels used by device
-        device_wires = self.map_wires(wires)
-
-        dim = 2 ** len(device_wires)
-        mat_batch_size = self._get_batch_size(mat, (dim, dim), dim**2)
-        state_batch_size = self._get_batch_size(state, (2,) * self.num_wires, 2**self.num_wires)
-
-        shape = [2] * (len(device_wires) * 2)
-        state_axes = device_wires
-        # If the matrix is broadcasted, it is reshaped to have leading axis of size mat_batch_size
-        if mat_batch_size:
-            shape.insert(0, mat_batch_size)
-            if state_batch_size:
-                raise NotImplementedError(
-                    "Applying a broadcasted unitary to an already broadcasted state via "
-                    "_apply_unitary is not supported. Broadcasting sizes are "
-                    f"({mat_batch_size}, {state_batch_size})."
+            with concurrent.futures.ProcessPoolExecutor(max_workers=max_workers) as executor:
+                results = tuple(
+                    executor.map(_adjoint_vjp_wrapper, vanilla_circuits, cotangents, seeds)
                 )
-        # If the state is broadcasted, the affected state axes need to be shifted by 1.
-        if state_batch_size:
-            state_axes = [ax + 1 for ax in state_axes]
-        mat = self._cast(self._reshape(mat, shape), dtype=self.C_DTYPE)
-        axes = (np.arange(-len(device_wires), 0), state_axes)
-        tdot = self._tensordot(mat, state, axes=axes)
-
-        # tensordot causes the axes given in `wires` to end up in the first positions
-        # of the resulting tensor. This corresponds to a (partial) transpose of
-        # the correct output state
-        # We'll need to invert this permutation to put the indices in the correct place
-        unused_idxs = [idx for idx in range(self.num_wires) if idx not in device_wires]
-        perm = list(device_wires) + unused_idxs
-        # If the matrix is broadcasted, all but the first dimension are shifted by 1
-        if mat_batch_size:
-            perm = [idx + 1 for idx in perm]
-            perm.insert(0, 0)
-        if state_batch_size:
-            # As the state broadcasting dimension always is the first in the state, it always
-            # ends up in position `len(device_wires)` after the tensordot. The -1 causes it
-            # being permuted to the leading dimension after transposition
-            perm.insert(len(device_wires), -1)
-
-        inv_perm = np.argsort(perm)  # argsort gives inverse permutation
-        return self._transpose(tdot, inv_perm)
-
-    def _apply_unitary_einsum(self, state, mat, wires):
-        r"""Apply multiplication of a matrix to subsystems of the quantum state.
-
-        This function uses einsum instead of tensordot. This approach is only
-        faster for single- and two-qubit gates.
-
-        Args:
-            state (array[complex]): input state
-            mat (array): matrix to multiply
-            wires (Wires): target wires
-
-        Returns:
-            array[complex]: output state
-        """
-        # translate to wire labels used by device
-        device_wires = self.map_wires(wires)
-
-        dim = 2 ** len(device_wires)
-        batch_size = self._get_batch_size(mat, (dim, dim), dim**2)
-
-        # If the matrix is broadcasted, it is reshaped to have leading axis of size mat_batch_size
-        shape = [2] * (len(device_wires) * 2)
-        if batch_size is not None:
-            shape.insert(0, batch_size)
-        mat = self._cast(self._reshape(mat, shape), dtype=self.C_DTYPE)
-
-        # Tensor indices of the quantum state
-        state_indices = ABC[: self.num_wires]
-
-        # Indices of the quantum state affected by this operation
-        affected_indices = "".join(ABC_ARRAY[list(device_wires)].tolist())
-
-        # All affected indices will be summed over, so we need the same number of new indices
-        new_indices = ABC[self.num_wires : self.num_wires + len(device_wires)]
-
-        # The new indices of the state are given by the old ones with the affected indices
-        # replaced by the new_indices
-        new_state_indices = functools.reduce(
-            lambda old_string, idx_pair: old_string.replace(idx_pair[0], idx_pair[1]),
-            zip(affected_indices, new_indices),
-            state_indices,
-        )
-
-        # We now put together the indices in the notation numpy's einsum requires
-        # This notation allows for the state, the matrix, or both to be broadcasted
-        einsum_indices = (
-            f"...{new_indices}{affected_indices},...{state_indices}->...{new_state_indices}"
-        )
-
-        return self._einsum(einsum_indices, mat, state)
-
-    def _apply_diagonal_unitary(self, state, phases, wires):
-        r"""Apply multiplication of a phase vector to subsystems of the quantum state.
-
-        This represents the multiplication with diagonal gates in a more efficient manner.
-
-        Args:
-            state (array[complex]): input state
-            phases (array): vector to multiply
-            wires (Wires): target wires
 
-        Returns:
-            array[complex]: output state
-        """
-        # translate to wire labels used by device
-        device_wires = self.map_wires(wires)
-        dim = 2 ** len(device_wires)
-        batch_size = self._get_batch_size(phases, (dim,), dim)
-
-        # reshape vectors
-        shape = [2] * len(device_wires)
-        if batch_size is not None:
-            shape.insert(0, batch_size)
-        phases = self._cast(self._reshape(phases, shape), dtype=self.C_DTYPE)
-
-        state_indices = ABC[: self.num_wires]
-        affected_indices = "".join(ABC_ARRAY[list(device_wires)].tolist())
-
-        einsum_indices = f"...{affected_indices},...{state_indices}->...{state_indices}"
-        return self._einsum(einsum_indices, phases, state)
-
-    def reset(self):
-        """Reset the device"""
-        super().reset()
-
-        # init the state vector to |00..0>
-        self._state = self._create_basis_state(0)
-        self._pre_rotated_state = self._state
-
-    def analytic_probability(self, wires=None):
-        if self._state is None:
-            return None
-
-        dim = 2**self.num_wires
-        batch_size = self._get_batch_size(self._state, [2] * self.num_wires, dim)
-        flat_state = self._reshape(
-            self._state, (batch_size, dim) if batch_size is not None else (dim,)
-        )
-        real_state = self._real(flat_state)
-        imag_state = self._imag(flat_state)
-        return self.marginal_prob(real_state**2 + imag_state**2, wires)
-
-    def classical_shadow(self, obs, circuit):
-        """
-        Returns the measured bits and recipes in the classical shadow protocol.
-
-        The protocol is described in detail in the `classical shadows paper <https://arxiv.org/abs/2002.08953>`_.
-        This measurement process returns the randomized Pauli measurements (the ``recipes``)
-        that are performed for each qubit and snapshot as an integer:
+            # reset _rng to mimic serial behavior
+            self._rng = np.random.default_rng(self._rng.integers(2**31 - 1))
 
-        - 0 for Pauli X,
-        - 1 for Pauli Y, and
-        - 2 for Pauli Z.
+        results, vjps = tuple(zip(*results))
+        return (results[0], vjps[0]) if is_single_circuit else (results, vjps)
 
-        It also returns the measurement results (the ``bits``); 0 if the 1 eigenvalue
-        is sampled, and 1 if the -1 eigenvalue is sampled.
 
-        The device shots are used to specify the number of snapshots. If ``T`` is the number
-        of shots and ``n`` is the number of qubits, then both the measured bits and the
-        Pauli measurements have shape ``(T, n)``.
+def _adjoint_jac_wrapper(c, rng=None, debugger=None):
+    state, is_state_batched = get_final_state(c, debugger=debugger)
+    jac = adjoint_jacobian(c, state=state)
+    res = measure_final_state(c, state, is_state_batched, rng=rng)
+    return res, jac
 
-        This implementation leverages vectorization and offers a significant speed-up over
-        the generic implementation.
-
-        .. Note::
-
-            This method internally calls ``np.einsum`` which supports at most 52 indices,
-            thus the classical shadow measurement for this device supports at most 52
-            qubits.
-
-        .. seealso:: :func:`~pennylane.classical_shadow`
-
-        Args:
-            obs (~.pennylane.measurements.ClassicalShadowMP): The classical shadow measurement process
-            circuit (~.tape.QuantumTape): The quantum tape that is being executed
 
-        Returns:
-            tensor_like[int]: A tensor with shape ``(2, T, n)``, where the first row represents
-            the measured bits and the second represents the recipes used.
-        """
-        wires = obs.wires
-        seed = obs.seed
-
-        n_qubits = len(wires)
-        n_snapshots = self.shots
-        device_qubits = len(self.wires)
-        mapped_wires = np.array(self.map_wires(wires))
-
-        # seed the random measurement generation so that recipes
-        # are the same for different executions with the same seed
-        rng = np.random.RandomState(seed)
-        recipes = rng.randint(0, 3, size=(n_snapshots, n_qubits))
-
-        obs_list = self._stack(
-            [
-                qml.PauliX.compute_matrix(),
-                qml.PauliY.compute_matrix(),
-                qml.PauliZ.compute_matrix(),
-            ]
-        )
-        uni_list = self._stack(
-            [
-                qml.Hadamard.compute_matrix(),
-                qml.Hadamard.compute_matrix() @ qml.RZ.compute_matrix(-np.pi / 2),
-                qml.Identity.compute_matrix(),
-            ]
-        )
-        obs = obs_list[recipes]
-        uni = uni_list[recipes]
-
-        # There's a significant speedup if we use the following iterative
-        # process to perform the randomized Pauli measurements:
-        #   1. Randomly generate Pauli observables for all snapshots for
-        #      a single qubit (e.g. the first qubit).
-        #   2. Compute the expectation of each Pauli observable on the first
-        #      qubit by tracing out all other qubits.
-        #   3. Sample the first qubit based on each Pauli expectation.
-        #   4. For all snapshots, determine the collapsed state of the remaining
-        #      qubits based on the sample result.
-        #   4. Repeat iteratively until no qubits are remaining.
-        #
-        # Observe that after the first iteration, the second qubit will become the
-        # "first" qubit in the process. The advantage to this approach as opposed to
-        # simulataneously computing the Pauli expectations for each qubit is that
-        # the partial traces are computed over iteratively smaller subsystems, leading
-        # to a significant speed-up.
-
-        # transpose the state so that the measured wires appear first
-        unmeasured_wires = [i for i in range(len(self.wires)) if i not in mapped_wires]
-        transposed_state = np.transpose(self._state, axes=mapped_wires.tolist() + unmeasured_wires)
-
-        outcomes = np.zeros((n_snapshots, n_qubits))
-        stacked_state = self._stack([transposed_state for _ in range(n_snapshots)])
-
-        for i in range(n_qubits):
-            # trace out every qubit except the first
-            first_qubit_state = self._einsum(
-                f"{ABC[device_qubits - i + 1]}{ABC[:device_qubits - i]},{ABC[device_qubits - i + 1]}{ABC[device_qubits - i]}{ABC[1:device_qubits - i]}"
-                f"->{ABC[device_qubits - i + 1]}a{ABC[device_qubits - i]}",
-                stacked_state,
-                self._conj(stacked_state),
-            )
-
-            # sample the observables on the first qubit
-            probs = (self._einsum("abc,acb->a", first_qubit_state, obs[:, i]) + 1) / 2
-            samples = np.random.uniform(0, 1, size=probs.shape) > probs
-            outcomes[:, i] = samples
-
-            # collapse the state of the remaining qubits; the next qubit in line
-            # becomes the first qubit for the next iteration
-            rotated_state = self._einsum("ab...,acb->ac...", stacked_state, uni[:, i])
-            stacked_state = rotated_state[np.arange(n_snapshots), self._cast(samples, np.int8)]
-
-            # re-normalize the collapsed state
-            norms = np.sqrt(
-                np.sum(
-                    np.abs(stacked_state) ** 2, tuple(range(1, device_qubits - i)), keepdims=True
-                )
-            )
-            stacked_state /= norms
+def _adjoint_jvp_wrapper(c, t, rng=None, debugger=None):
+    state, is_state_batched = get_final_state(c, debugger=debugger)
+    jvp = adjoint_jvp(c, t, state=state)
+    res = measure_final_state(c, state, is_state_batched, rng=rng)
+    return res, jvp
 
-        return self._cast(self._stack([outcomes, recipes]), dtype=np.int8)
 
-    def _get_diagonalizing_gates(self, circuit: qml.tape.QuantumTape) -> List[Operation]:
-        meas_filtered = [
-            m
-            for m in circuit.measurements
-            if m.obs is None or not isinstance(m.obs, qml.Hamiltonian)
-        ]
-        return super()._get_diagonalizing_gates(qml.tape.QuantumScript(measurements=meas_filtered))
+def _adjoint_vjp_wrapper(c, t, rng=None, debugger=None):
+    state, is_state_batched = get_final_state(c, debugger=debugger)
+    vjp = adjoint_vjp(c, t, state=state)
+    res = measure_final_state(c, state, is_state_batched, rng=rng)
+    return res, vjp
diff --git a/pennylane/devices/default_qubit_autograd.py b/pennylane/devices/default_qubit_autograd.py
index dfe038a6ed2..dcf57f0b84c 100644
--- a/pennylane/devices/default_qubit_autograd.py
+++ b/pennylane/devices/default_qubit_autograd.py
@@ -11,16 +11,16 @@
 # 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.
-"""This module contains an autograd implementation of the :class:`~.DefaultQubit`
+"""This module contains an autograd implementation of the :class:`~.DefaultQubitLegacy`
 reference plugin.
 """
 from pennylane import numpy as np
 
-from pennylane.devices import DefaultQubit
+from pennylane.devices import DefaultQubitLegacy
 
 
-class DefaultQubitAutograd(DefaultQubit):
-    """Simulator plugin based on ``"default.qubit"``, written using Autograd.
+class DefaultQubitAutograd(DefaultQubitLegacy):
+    """Simulator plugin based on ``"default.qubit.legacy"``, written using Autograd.
 
     **Short name:** ``default.qubit.autograd``
 
diff --git a/pennylane/devices/default_qubit_jax.py b/pennylane/devices/default_qubit_jax.py
index 411d77fea02..38bccaeea05 100644
--- a/pennylane/devices/default_qubit_jax.py
+++ b/pennylane/devices/default_qubit_jax.py
@@ -11,14 +11,14 @@
 # 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.
-"""This module contains a jax implementation of the :class:`~.DefaultQubit`
+"""This module contains a jax implementation of the :class:`~.DefaultQubitLegacy`
 reference plugin.
 """
 # pylint: disable=ungrouped-imports
 import numpy as np
 
 import pennylane as qml
-from pennylane.devices import DefaultQubit
+from pennylane.devices import DefaultQubitLegacy
 from pennylane.pulse import ParametrizedEvolution
 from pennylane.typing import TensorLike
 
@@ -34,8 +34,8 @@
     raise ImportError("default.qubit.jax device requires installing jax>0.3.20") from e
 
 
-class DefaultQubitJax(DefaultQubit):
-    """Simulator plugin based on ``"default.qubit"``, written using jax.
+class DefaultQubitJax(DefaultQubitLegacy):
+    """Simulator plugin based on ``"default.qubit.legacy"``, written using jax.
 
     **Short name:** ``default.qubit.jax``
 
@@ -163,7 +163,7 @@ def circuit():
     _size = staticmethod(jnp.size)
     _ndim = staticmethod(jnp.ndim)
 
-    operations = DefaultQubit.operations.union({"ParametrizedEvolution"})
+    operations = DefaultQubitLegacy.operations.union({"ParametrizedEvolution"})
 
     def __init__(self, wires, *, shots=None, prng_key=None, analytic=None):
         if jax_config.read("jax_enable_x64"):
diff --git a/pennylane/devices/default_qubit_legacy.py b/pennylane/devices/default_qubit_legacy.py
new file mode 100644
index 00000000000..dc4983933c3
--- /dev/null
+++ b/pennylane/devices/default_qubit_legacy.py
@@ -0,0 +1,1085 @@
+# Copyright 2018-2021 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.
+r"""
+The default.qubit device is PennyLane's standard qubit-based device.
+
+It implements the necessary :class:`~pennylane._device.Device` methods as well as some built-in
+:mod:`qubit operations <pennylane.ops.qubit>`, and provides a very simple pure state
+simulation of a qubit-based quantum circuit architecture.
+"""
+import functools
+import itertools
+from string import ascii_letters as ABC
+from typing import List
+
+import numpy as np
+from scipy.sparse import csr_matrix
+
+import pennylane as qml
+from pennylane import BasisState, DeviceError, QubitDevice, StatePrep, Snapshot
+from pennylane.devices.qubit import measure
+from pennylane.operation import Operation
+from pennylane.ops import Sum
+from pennylane.ops.qubit.attributes import diagonal_in_z_basis
+from pennylane.pulse import ParametrizedEvolution
+from pennylane.measurements import ExpectationMP
+from pennylane.typing import TensorLike
+from pennylane.wires import WireError
+
+from .._version import __version__
+
+ABC_ARRAY = np.array(list(ABC))
+
+# tolerance for numerical errors
+tolerance = 1e-10
+SQRT2INV = 1 / np.sqrt(2)
+TPHASE = np.exp(1j * np.pi / 4)
+
+
+def _get_slice(index, axis, num_axes):
+    """Allows slicing along an arbitrary axis of an array or tensor.
+
+    Args:
+        index (int): the index to access
+        axis (int): the axis to slice into
+        num_axes (int): total number of axes
+
+    Returns:
+        tuple[slice or int]: a tuple that can be used to slice into an array or tensor
+
+    **Example:**
+
+    Accessing the 2 index along axis 1 of a 3-axis array:
+
+    >>> sl = _get_slice(2, 1, 3)
+    >>> sl
+    (slice(None, None, None), 2, slice(None, None, None))
+    >>> a = np.arange(27).reshape((3, 3, 3))
+    >>> a[sl]
+    array([[ 6,  7,  8],
+           [15, 16, 17],
+           [24, 25, 26]])
+    """
+    idx = [slice(None)] * num_axes
+    idx[axis] = index
+    return tuple(idx)
+
+
+# pylint: disable=unused-argument
+class DefaultQubitLegacy(QubitDevice):
+    """Default qubit device for PennyLane.
+
+    .. warning::
+
+        This is the legacy implementation of DefaultQubit. It has been replaced by
+        ``qml.devices.DefaultQubit``, which can be accessed with the familiar constructor,
+        ``qml.device("default.qubit")``.
+
+        This change will not alter device behaviour for most workflows, but may have implications for
+        plugin developers and users who directly interact with device methods. Please consult
+        :class:`pennylane.devices.Device` and the implementation in
+        :class:`pennylane.devices.DefaultQubit` for more information on what the new
+        interface will look like and be prepared to make updates in a coming release. If you have any
+        feedback on these changes, please create an
+        `issue <https://github.com/PennyLaneAI/pennylane/issues>`_ or post in our
+        `discussion forum <https://discuss.pennylane.ai/>`_.
+
+    Args:
+        wires (int, Iterable[Number, str]): Number of subsystems represented by the device,
+            or iterable that contains unique labels for the subsystems as numbers (i.e., ``[-1, 0, 2]``)
+            or strings (``['ancilla', 'q1', 'q2']``). Default 1 if not specified.
+        shots (None, int): How many times the circuit should be evaluated (or sampled) to estimate
+            the expectation values. Defaults to ``None`` if not specified, which means that the device
+            returns analytical results.
+    """
+
+    name = "Default qubit PennyLane plugin"
+    short_name = "default.qubit"
+    pennylane_requires = __version__
+    version = __version__
+    author = "Xanadu Inc."
+
+    operations = {
+        "GlobalPhase",
+        "Identity",
+        "Snapshot",
+        "BasisState",
+        "StatePrep",
+        "QubitStateVector",
+        "QubitUnitary",
+        "ControlledQubitUnitary",
+        "BlockEncode",
+        "MultiControlledX",
+        "IntegerComparator",
+        "DiagonalQubitUnitary",
+        "PauliX",
+        "PauliY",
+        "PauliZ",
+        "MultiRZ",
+        "Hadamard",
+        "S",
+        "Adjoint(S)",
+        "T",
+        "Adjoint(T)",
+        "SX",
+        "Adjoint(SX)",
+        "CNOT",
+        "SWAP",
+        "ISWAP",
+        "PSWAP",
+        "Adjoint(ISWAP)",
+        "SISWAP",
+        "Adjoint(SISWAP)",
+        "SQISW",
+        "CSWAP",
+        "Toffoli",
+        "CCZ",
+        "CY",
+        "CZ",
+        "CH",
+        "PhaseShift",
+        "PCPhase",
+        "ControlledPhaseShift",
+        "CPhaseShift00",
+        "CPhaseShift01",
+        "CPhaseShift10",
+        "CPhase",
+        "RX",
+        "RY",
+        "RZ",
+        "Rot",
+        "CRX",
+        "CRY",
+        "CRZ",
+        "CRot",
+        "IsingXX",
+        "IsingYY",
+        "IsingZZ",
+        "IsingXY",
+        "SingleExcitation",
+        "SingleExcitationPlus",
+        "SingleExcitationMinus",
+        "DoubleExcitation",
+        "DoubleExcitationPlus",
+        "DoubleExcitationMinus",
+        "QubitCarry",
+        "QubitSum",
+        "OrbitalRotation",
+        "QFT",
+        "ECR",
+    }
+
+    observables = {
+        "PauliX",
+        "PauliY",
+        "PauliZ",
+        "Hadamard",
+        "Hermitian",
+        "Identity",
+        "Projector",
+        "SparseHamiltonian",
+        "Hamiltonian",
+        "Sum",
+        "SProd",
+        "Prod",
+        "Exp",
+        "Evolution",
+    }
+
+    def __init__(
+        self, wires, *, r_dtype=np.float64, c_dtype=np.complex128, shots=None, analytic=None
+    ):
+        super().__init__(wires, shots, r_dtype=r_dtype, c_dtype=c_dtype, analytic=analytic)
+        self._debugger = None
+
+        # Create the initial state. Internally, we store the
+        # state as an array of dimension [2]*wires.
+        self._state = self._create_basis_state(0)
+        self._pre_rotated_state = self._state
+
+        self._apply_ops = {
+            "PauliX": self._apply_x,
+            "PauliY": self._apply_y,
+            "PauliZ": self._apply_z,
+            "Hadamard": self._apply_hadamard,
+            "S": self._apply_s,
+            "T": self._apply_t,
+            "SX": self._apply_sx,
+            "CNOT": self._apply_cnot,
+            "SWAP": self._apply_swap,
+            "CZ": self._apply_cz,
+            "Toffoli": self._apply_toffoli,
+        }
+
+    @property
+    def stopping_condition(self):
+        def accepts_obj(obj):
+            if obj.name == "QFT" and len(obj.wires) >= 6:
+                return False
+            if obj.name == "GroverOperator" and len(obj.wires) >= 13:
+                return False
+            if getattr(obj, "has_matrix", False):
+                # pow operations dont work with backprop or adjoint without decomposition
+                # use class name string so we don't need to use isinstance check
+                return not (obj.__class__.__name__ == "Pow" and qml.operation.is_trainable(obj))
+            return obj.name in self.observables.union(self.operations)
+
+        return qml.BooleanFn(accepts_obj)
+
+    @functools.lru_cache()
+    def map_wires(self, wires):
+        # temporarily overwrite this method to bypass
+        # wire map that produces Wires objects
+        try:
+            mapped_wires = [self.wire_map[w] for w in wires]
+        except KeyError as e:
+            raise WireError(
+                f"Did not find some of the wires {wires.labels} on device with wires {self.wires.labels}."
+            ) from e
+
+        return mapped_wires
+
+    def define_wire_map(self, wires):
+        # temporarily overwrite this method to bypass
+        # wire map that produces Wires objects
+        consecutive_wires = range(self.num_wires)
+        wire_map = zip(wires, consecutive_wires)
+        return dict(wire_map)
+
+    # pylint: disable=arguments-differ
+    def _get_batch_size(self, tensor, expected_shape, expected_size):
+        """Determine whether a tensor has an additional batch dimension for broadcasting,
+        compared to an expected_shape."""
+        size = self._size(tensor)
+        if self._ndim(tensor) > len(expected_shape) or size > expected_size:
+            return size // expected_size
+
+        return None
+
+    # pylint: disable=arguments-differ
+    def apply(self, operations, rotations=None, **kwargs):
+        rotations = rotations or []
+
+        # apply the circuit operations
+        for i, operation in enumerate(operations):
+            if i > 0 and isinstance(operation, (StatePrep, BasisState)):
+                raise DeviceError(
+                    f"Operation {operation.name} cannot be used after other Operations have already been applied "
+                    f"on a {self.short_name} device."
+                )
+
+            if isinstance(operation, StatePrep):
+                self._apply_state_vector(operation.parameters[0], operation.wires)
+            elif isinstance(operation, BasisState):
+                self._apply_basis_state(operation.parameters[0], operation.wires)
+            elif isinstance(operation, Snapshot):
+                if self._debugger and self._debugger.active:
+                    state_vector = np.array(self._flatten(self._state))
+                    if operation.tag:
+                        self._debugger.snapshots[operation.tag] = state_vector
+                    else:
+                        self._debugger.snapshots[len(self._debugger.snapshots)] = state_vector
+            elif isinstance(operation, ParametrizedEvolution):
+                self._state = self._apply_parametrized_evolution(self._state, operation)
+            else:
+                self._state = self._apply_operation(self._state, operation)
+
+        # store the pre-rotated state
+        self._pre_rotated_state = self._state
+
+        # apply the circuit rotations
+        for operation in rotations:
+            self._state = self._apply_operation(self._state, operation)
+
+    def _apply_parametrized_evolution(self, state: TensorLike, operation: ParametrizedEvolution):
+        """Applies a parametrized evolution to the input state.
+
+        Args:
+            state (array[complex]): input state
+            operation (ParametrizedEvolution): operation to apply on the state
+        """
+        raise NotImplementedError(
+            f"The device {self.short_name} cannot execute a ParametrizedEvolution operation. "
+            "Please use the jax interface."
+        )
+
+    def _apply_operation(self, state, operation):
+        """Applies operations to the input state.
+
+        Args:
+            state (array[complex]): input state
+            operation (~.Operation): operation to apply on the device
+
+        Returns:
+            array[complex]: output state
+        """
+        if operation.__class__.__name__ == "Identity":
+            return state
+        if operation.name == "GlobalPhase":
+            return self._apply_global_phase(state, operation)
+        wires = operation.wires
+
+        if str(operation.name) in self._apply_ops:  # cast to string because of Tensor
+            shift = int(self._ndim(state) > self.num_wires)
+            axes = [ax + shift for ax in self.wires.indices(wires)]
+            return self._apply_ops[operation.name](state, axes)
+
+        matrix = self._asarray(self._get_unitary_matrix(operation), dtype=self.C_DTYPE)
+
+        if operation in diagonal_in_z_basis:
+            return self._apply_diagonal_unitary(state, matrix, wires)
+        if len(wires) <= 2:
+            # Einsum is faster for small gates
+            return self._apply_unitary_einsum(state, matrix, wires)
+
+        return self._apply_unitary(state, matrix, wires)
+
+    def _apply_global_phase(self, state, operation: qml.GlobalPhase):  # pylint: disable=no-self-use
+        """Applies a :class:`~.GlobalPhase` operation to the state."""
+        return qml.math.exp(-1j * operation.data[0]) * state
+
+    def _apply_x(self, state, axes, **kwargs):
+        """Applies a PauliX gate by rolling 1 unit along the axis specified in ``axes``.
+
+        Rolling by 1 unit along the axis means that the :math:`|0 \rangle` state with index ``0`` is
+        shifted to the :math:`|1 \rangle` state with index ``1``. Likewise, since rolling beyond
+        the last index loops back to the first, :math:`|1 \rangle` is transformed to
+        :math:`|0\rangle`.
+
+        Args:
+            state (array[complex]): input state
+            axes (List[int]): target axes to apply transformation
+
+        Returns:
+            array[complex]: output state
+        """
+        return self._roll(state, 1, axes[0])
+
+    def _apply_y(self, state, axes, **kwargs):
+        """Applies a PauliY gate by adding a negative sign to the 1 index along the axis specified
+        in ``axes``, rolling one unit along the same axis, and multiplying the result by 1j.
+
+        Args:
+            state (array[complex]): input state
+            axes (List[int]): target axes to apply transformation
+
+        Returns:
+            array[complex]: output state
+        """
+        return 1j * self._apply_x(self._apply_z(state, axes), axes)
+
+    def _apply_z(self, state, axes, **kwargs):
+        """Applies a PauliZ gate by adding a negative sign to the 1 index along the axis specified
+        in ``axes``.
+
+        Args:
+            state (array[complex]): input state
+            axes (List[int]): target axes to apply transformation
+
+        Returns:
+            array[complex]: output state
+        """
+        return self._apply_phase(state, axes, -1)
+
+    def _apply_hadamard(self, state, axes, **kwargs):
+        """Apply the Hadamard gate by combining the results of applying the PauliX and PauliZ gates.
+
+        Args:
+            state (array[complex]): input state
+            axes (List[int]): target axes to apply transformation
+
+        Returns:
+            array[complex]: output state
+        """
+        state_x = self._apply_x(state, axes)
+        state_z = self._apply_z(state, axes)
+        return self._const_mul(SQRT2INV, state_x + state_z)
+
+    def _apply_s(self, state, axes, inverse=False):
+        return self._apply_phase(state, axes, 1j, inverse)
+
+    def _apply_t(self, state, axes, inverse=False):
+        return self._apply_phase(state, axes, TPHASE, inverse)
+
+    def _apply_sx(self, state, axes, inverse=False):
+        """Apply the Square Root X gate.
+
+        Args:
+            state (array[complex]): input state
+            axes (List[int]): target axes to apply transformation
+
+        Returns:
+            array[complex]: output state
+        """
+        if inverse:
+            return 0.5 * ((1 - 1j) * state + (1 + 1j) * self._apply_x(state, axes))
+
+        return 0.5 * ((1 + 1j) * state + (1 - 1j) * self._apply_x(state, axes))
+
+    def _apply_cnot(self, state, axes, **kwargs):
+        """Applies a CNOT gate by slicing along the first axis specified in ``axes`` and then
+        applying an X transformation along the second axis.
+
+        By slicing along the first axis, we are able to select all of the amplitudes with a
+        corresponding :math:`|1\rangle` for the control qubit. This means we then just need to apply
+        a :class:`~.PauliX` (NOT) gate to the result.
+
+        Args:
+            state (array[complex]): input state
+            axes (List[int]): target axes to apply transformation
+
+        Returns:
+            array[complex]: output state
+        """
+        ndim = self._ndim(state)
+        sl_0 = _get_slice(0, axes[0], ndim)
+        sl_1 = _get_slice(1, axes[0], ndim)
+
+        # We will be slicing into the state according to state[sl_1], giving us all of the
+        # amplitudes with a |1> for the control qubit. The resulting array has lost an axis
+        # relative to state and we need to be careful about the axis we apply the PauliX rotation
+        # to. If axes[1] is larger than axes[0], then we need to shift the target axis down by
+        # one, otherwise we can leave as-is. For example: a state has [0, 1, 2, 3], control=1,
+        # target=3. Then, state[sl_1] has 3 axes and target=3 now corresponds to the second axis.
+        if axes[1] > axes[0]:
+            target_axes = [axes[1] - 1]
+        else:
+            target_axes = [axes[1]]
+
+        state_x = self._apply_x(state[sl_1], axes=target_axes)
+        return self._stack([state[sl_0], state_x], axis=axes[0])
+
+    def _apply_toffoli(self, state, axes, **kwargs):
+        """Applies a Toffoli gate by slicing along the axis of the greater control qubit, slicing
+        each of the resulting sub-arrays along the axis of the smaller control qubit, and then applying
+        an X transformation along the axis of the target qubit of the fourth sub-sub-array.
+
+        By performing two consecutive slices in this way, we are able to select all of the amplitudes with
+        a corresponding :math:`|11\rangle` for the two control qubits. This means we then just need to apply
+        a :class:`~.PauliX` (NOT) gate to the result.
+
+        Args:
+            state (array[complex]): input state
+            axes (List[int]): target axes to apply transformation
+
+        Returns:
+            array[complex]: output state
+        """
+        cntrl_max = np.argmax(axes[:2])
+        cntrl_min = cntrl_max ^ 1
+        ndim = self._ndim(state)
+        sl_a0 = _get_slice(0, axes[cntrl_max], ndim)
+        sl_a1 = _get_slice(1, axes[cntrl_max], ndim)
+        sl_b0 = _get_slice(0, axes[cntrl_min], ndim - 1)
+        sl_b1 = _get_slice(1, axes[cntrl_min], ndim - 1)
+
+        # If both controls are smaller than the target, shift the target axis down by two. If one
+        # control is greater and one control is smaller than the target, shift the target axis
+        # down by one. If both controls are greater than the target, leave the target axis as-is.
+        if axes[cntrl_min] > axes[2]:
+            target_axes = [axes[2]]
+        elif axes[cntrl_max] > axes[2]:
+            target_axes = [axes[2] - 1]
+        else:
+            target_axes = [axes[2] - 2]
+
+        # state[sl_a1][sl_b1] gives us all of the amplitudes with a |11> for the two control qubits.
+        state_x = self._apply_x(state[sl_a1][sl_b1], axes=target_axes)
+        state_stacked_a1 = self._stack([state[sl_a1][sl_b0], state_x], axis=axes[cntrl_min])
+        return self._stack([state[sl_a0], state_stacked_a1], axis=axes[cntrl_max])
+
+    def _apply_swap(self, state, axes, **kwargs):
+        """Applies a SWAP gate by performing a partial transposition along the specified axes.
+
+        Args:
+            state (array[complex]): input state
+            axes (List[int]): target axes to apply transformation
+
+        Returns:
+            array[complex]: output state
+        """
+        all_axes = list(range(len(state.shape)))
+        all_axes[axes[0]] = axes[1]
+        all_axes[axes[1]] = axes[0]
+        return self._transpose(state, all_axes)
+
+    def _apply_cz(self, state, axes, **kwargs):
+        """Applies a CZ gate by slicing along the first axis specified in ``axes`` and then
+        applying a Z transformation along the second axis.
+
+        Args:
+            state (array[complex]): input state
+            axes (List[int]): target axes to apply transformation
+
+        Returns:
+            array[complex]: output state
+        """
+        ndim = self._ndim(state)
+        sl_0 = _get_slice(0, axes[0], ndim)
+        sl_1 = _get_slice(1, axes[0], ndim)
+
+        if axes[1] > axes[0]:
+            target_axes = [axes[1] - 1]
+        else:
+            target_axes = [axes[1]]
+
+        state_z = self._apply_z(state[sl_1], axes=target_axes)
+        return self._stack([state[sl_0], state_z], axis=axes[0])
+
+    def _apply_phase(self, state, axes, parameters, inverse=False):
+        """Applies a phase onto the 1 index along the axis specified in ``axes``.
+
+        Args:
+            state (array[complex]): input state
+            axes (List[int]): target axes to apply transformation
+            parameters (float): phase to apply
+            inverse (bool): whether to apply the inverse phase
+
+        Returns:
+            array[complex]: output state
+        """
+        ndim = self._ndim(state)
+        sl_0 = _get_slice(0, axes[0], ndim)
+        sl_1 = _get_slice(1, axes[0], ndim)
+
+        phase = self._conj(parameters) if inverse else parameters
+        return self._stack([state[sl_0], self._const_mul(phase, state[sl_1])], axis=axes[0])
+
+    def expval(self, observable, shot_range=None, bin_size=None):
+        """Returns the expectation value of a Hamiltonian observable. When the observable is a
+        ``Hamiltonian`` or ``SparseHamiltonian`` object, the expectation value is computed directly
+        from the sparse matrix representation, which leads to faster execution.
+
+        Args:
+            observable (~.Observable): a PennyLane observable
+            shot_range (tuple[int]): 2-tuple of integers specifying the range of samples
+                to use. If not specified, all samples are used.
+            bin_size (int): Divides the shot range into bins of size ``bin_size``, and
+                returns the measurement statistic separately over each bin. If not
+                provided, the entire shot range is treated as a single bin.
+
+        Returns:
+            float: returns the expectation value of the observable
+
+        .. warning::
+
+            This function does not support broadcasting if ``observable`` is a
+            :class:``~.Hamiltonian`` and the device interface or the interface of the
+            Hamiltonian is not NumPy or Autograd
+
+        """
+        is_state_batched = self._ndim(self.state) == 2
+        # intercept Sums
+        if isinstance(observable, Sum) and not self.shots:
+            return measure(
+                ExpectationMP(observable.map_wires(self.wire_map)),
+                self._pre_rotated_state,
+                is_state_batched,
+            )
+
+        # intercept other Hamiltonians
+        # TODO: Ideally, this logic should not live in the Device, but be moved
+        # to a component that can be re-used by devices as needed.
+        if observable.name not in ("Hamiltonian", "SparseHamiltonian"):
+            return super().expval(observable, shot_range=shot_range, bin_size=bin_size)
+
+        assert self.shots is None, f"{observable.name} must be used with shots=None"
+
+        self.map_wires(observable.wires)
+        backprop_mode = (
+            not isinstance(self.state, np.ndarray)
+            or any(not isinstance(d, (float, np.ndarray)) for d in observable.data)
+        ) and observable.name == "Hamiltonian"
+
+        if backprop_mode:
+            # TODO[dwierichs]: This branch is not adapted to broadcasting yet
+            if is_state_batched:
+                raise NotImplementedError(
+                    "Expectation values of Hamiltonians for interface!=None are "
+                    "not supported together with parameter broadcasting yet"
+                )
+            # We must compute the expectation value assuming that the Hamiltonian
+            # coefficients *and* the quantum states are tensor objects.
+
+            # Compute  <psi| H |psi> via sum_i coeff_i * <psi| PauliWord |psi> using a sparse
+            # representation of the Pauliword
+            res = qml.math.cast(qml.math.convert_like(0.0, observable.data), dtype=complex)
+            interface = qml.math.get_interface(self.state)
+
+            # Note: it is important that we use the Hamiltonian's data and not the coeffs
+            # attribute. This is because the .data attribute may be 'unwrapped' as required by
+            # the interfaces, whereas the .coeff attribute will always be the same input dtype
+            # that the user provided.
+            for op, coeff in zip(observable.ops, observable.data):
+                # extract a scipy.sparse.coo_matrix representation of this Pauli word
+                coo = qml.operation.Tensor(op).sparse_matrix(wire_order=self.wires, format="coo")
+                Hmat = qml.math.cast(qml.math.convert_like(coo.data, self.state), self.C_DTYPE)
+
+                product = (
+                    self._gather(self._conj(self.state), coo.row)
+                    * Hmat
+                    * self._gather(self.state, coo.col)
+                )
+                c = qml.math.convert_like(coeff, product)
+
+                if interface == "tensorflow":
+                    c = qml.math.cast(c, "complex128")
+
+                res = qml.math.convert_like(res, product) + qml.math.sum(c * product)
+
+        else:
+            # Coefficients and the state are not trainable, we can be more
+            # efficient in how we compute the Hamiltonian sparse matrix.
+            Hmat = observable.sparse_matrix(wire_order=self.wires)
+
+            state = qml.math.toarray(self.state)
+            if is_state_batched:
+                res = qml.math.array(
+                    [
+                        csr_matrix.dot(
+                            csr_matrix(self._conj(_state)),
+                            csr_matrix.dot(Hmat, csr_matrix(_state[..., None])),
+                        ).toarray()[0]
+                        for _state in state
+                    ]
+                )
+            else:
+                res = csr_matrix.dot(
+                    csr_matrix(self._conj(state)),
+                    csr_matrix.dot(Hmat, csr_matrix(state[..., None])),
+                ).toarray()[0]
+
+        if observable.name == "Hamiltonian":
+            res = qml.math.squeeze(res)
+
+        return self._real(res)
+
+    def _get_unitary_matrix(self, unitary):  # pylint: disable=no-self-use
+        """Return the matrix representing a unitary operation.
+
+        Args:
+            unitary (~.Operation): a PennyLane unitary operation
+
+        Returns:
+            array[complex]: Returns a 2D matrix representation of
+            the unitary in the computational basis, or, in the case of a diagonal unitary,
+            a 1D array representing the matrix diagonal.
+        """
+        if unitary in diagonal_in_z_basis:
+            return unitary.eigvals()
+
+        return unitary.matrix()
+
+    @classmethod
+    def capabilities(cls):
+        capabilities = super().capabilities().copy()
+        capabilities.update(
+            model="qubit",
+            supports_inverse_operations=True,
+            supports_analytic_computation=True,
+            supports_broadcasting=True,
+            returns_state=True,
+            passthru_devices={
+                "tf": "default.qubit.tf",
+                "torch": "default.qubit.torch",
+                "autograd": "default.qubit.autograd",
+                "jax": "default.qubit.jax",
+            },
+        )
+        return capabilities
+
+    def _create_basis_state(self, index):
+        """Return a computational basis state over all wires.
+
+        Args:
+            index (int): integer representing the computational basis state
+
+        Returns:
+            array[complex]: complex array of shape ``[2]*self.num_wires``
+            representing the statevector of the basis state
+
+        Note: This function does not support broadcasted inputs yet.
+        """
+        state = np.zeros(2**self.num_wires, dtype=np.complex128)
+        state[index] = 1
+        state = self._asarray(state, dtype=self.C_DTYPE)
+        return self._reshape(state, [2] * self.num_wires)
+
+    @property
+    def state(self):
+        dim = 2**self.num_wires
+        batch_size = self._get_batch_size(self._pre_rotated_state, (2,) * self.num_wires, dim)
+        # Do not flatten the state completely but leave the broadcasting dimension if there is one
+        shape = (batch_size, dim) if batch_size is not None else (dim,)
+        return self._reshape(self._pre_rotated_state, shape)
+
+    def _apply_state_vector(self, state, device_wires):
+        """Initialize the internal state vector in a specified state.
+
+        Args:
+            state (array[complex]): normalized input state of length ``2**len(wires)``
+                or broadcasted state of shape ``(batch_size, 2**len(wires))``
+            device_wires (Wires): wires that get initialized in the state
+        """
+
+        # translate to wire labels used by device
+        device_wires = self.map_wires(device_wires)
+        dim = 2 ** len(device_wires)
+
+        state = self._asarray(state, dtype=self.C_DTYPE)
+        batch_size = self._get_batch_size(state, (dim,), dim)
+        output_shape = [2] * self.num_wires
+        if batch_size is not None:
+            output_shape.insert(0, batch_size)
+
+        if len(device_wires) == self.num_wires and sorted(device_wires) == device_wires:
+            # Initialize the entire device state with the input state
+            self._state = self._reshape(state, output_shape)
+            return
+
+        # generate basis states on subset of qubits via the cartesian product
+        basis_states = np.array(list(itertools.product([0, 1], repeat=len(device_wires))))
+
+        # get basis states to alter on full set of qubits
+        unravelled_indices = np.zeros((2 ** len(device_wires), self.num_wires), dtype=int)
+        unravelled_indices[:, device_wires] = basis_states
+
+        # get indices for which the state is changed to input state vector elements
+        ravelled_indices = np.ravel_multi_index(unravelled_indices.T, [2] * self.num_wires)
+
+        if batch_size is not None:
+            state = self._scatter(
+                (slice(None), ravelled_indices), state, [batch_size, 2**self.num_wires]
+            )
+        else:
+            state = self._scatter(ravelled_indices, state, [2**self.num_wires])
+        state = self._reshape(state, output_shape)
+        self._state = self._asarray(state, dtype=self.C_DTYPE)
+
+    def _apply_basis_state(self, state, wires):
+        """Initialize the state vector in a specified computational basis state.
+
+        Args:
+            state (array[int]): computational basis state of shape ``(wires,)``
+                consisting of 0s and 1s.
+            wires (Wires): wires that the provided computational state should be initialized on
+
+        Note: This function does not support broadcasted inputs yet.
+        """
+        # translate to wire labels used by device
+        device_wires = self.map_wires(wires)
+
+        # length of basis state parameter
+        n_basis_state = len(state)
+
+        if not set(state.tolist()).issubset({0, 1}):
+            raise ValueError("BasisState parameter must consist of 0 or 1 integers.")
+
+        if n_basis_state != len(device_wires):
+            raise ValueError("BasisState parameter and wires must be of equal length.")
+
+        # get computational basis state number
+        basis_states = 2 ** (self.num_wires - 1 - np.array(device_wires))
+        basis_states = qml.math.convert_like(basis_states, state)
+        num = int(qml.math.dot(state, basis_states))
+
+        self._state = self._create_basis_state(num)
+
+    def _apply_unitary(self, state, mat, wires):
+        r"""Apply multiplication of a matrix to subsystems of the quantum state.
+
+        Args:
+            state (array[complex]): input state
+            mat (array): matrix to multiply
+            wires (Wires): target wires
+
+        Returns:
+            array[complex]: output state
+
+        Note: This function does not support simultaneously broadcasted states and matrices yet.
+        """
+        # translate to wire labels used by device
+        device_wires = self.map_wires(wires)
+
+        dim = 2 ** len(device_wires)
+        mat_batch_size = self._get_batch_size(mat, (dim, dim), dim**2)
+        state_batch_size = self._get_batch_size(state, (2,) * self.num_wires, 2**self.num_wires)
+
+        shape = [2] * (len(device_wires) * 2)
+        state_axes = device_wires
+        # If the matrix is broadcasted, it is reshaped to have leading axis of size mat_batch_size
+        if mat_batch_size:
+            shape.insert(0, mat_batch_size)
+            if state_batch_size:
+                raise NotImplementedError(
+                    "Applying a broadcasted unitary to an already broadcasted state via "
+                    "_apply_unitary is not supported. Broadcasting sizes are "
+                    f"({mat_batch_size}, {state_batch_size})."
+                )
+        # If the state is broadcasted, the affected state axes need to be shifted by 1.
+        if state_batch_size:
+            state_axes = [ax + 1 for ax in state_axes]
+        mat = self._cast(self._reshape(mat, shape), dtype=self.C_DTYPE)
+        axes = (np.arange(-len(device_wires), 0), state_axes)
+        tdot = self._tensordot(mat, state, axes=axes)
+
+        # tensordot causes the axes given in `wires` to end up in the first positions
+        # of the resulting tensor. This corresponds to a (partial) transpose of
+        # the correct output state
+        # We'll need to invert this permutation to put the indices in the correct place
+        unused_idxs = [idx for idx in range(self.num_wires) if idx not in device_wires]
+        perm = list(device_wires) + unused_idxs
+        # If the matrix is broadcasted, all but the first dimension are shifted by 1
+        if mat_batch_size:
+            perm = [idx + 1 for idx in perm]
+            perm.insert(0, 0)
+        if state_batch_size:
+            # As the state broadcasting dimension always is the first in the state, it always
+            # ends up in position `len(device_wires)` after the tensordot. The -1 causes it
+            # being permuted to the leading dimension after transposition
+            perm.insert(len(device_wires), -1)
+
+        inv_perm = np.argsort(perm)  # argsort gives inverse permutation
+        return self._transpose(tdot, inv_perm)
+
+    def _apply_unitary_einsum(self, state, mat, wires):
+        r"""Apply multiplication of a matrix to subsystems of the quantum state.
+
+        This function uses einsum instead of tensordot. This approach is only
+        faster for single- and two-qubit gates.
+
+        Args:
+            state (array[complex]): input state
+            mat (array): matrix to multiply
+            wires (Wires): target wires
+
+        Returns:
+            array[complex]: output state
+        """
+        # translate to wire labels used by device
+        device_wires = self.map_wires(wires)
+
+        dim = 2 ** len(device_wires)
+        batch_size = self._get_batch_size(mat, (dim, dim), dim**2)
+
+        # If the matrix is broadcasted, it is reshaped to have leading axis of size mat_batch_size
+        shape = [2] * (len(device_wires) * 2)
+        if batch_size is not None:
+            shape.insert(0, batch_size)
+        mat = self._cast(self._reshape(mat, shape), dtype=self.C_DTYPE)
+
+        # Tensor indices of the quantum state
+        state_indices = ABC[: self.num_wires]
+
+        # Indices of the quantum state affected by this operation
+        affected_indices = "".join(ABC_ARRAY[list(device_wires)].tolist())
+
+        # All affected indices will be summed over, so we need the same number of new indices
+        new_indices = ABC[self.num_wires : self.num_wires + len(device_wires)]
+
+        # The new indices of the state are given by the old ones with the affected indices
+        # replaced by the new_indices
+        new_state_indices = functools.reduce(
+            lambda old_string, idx_pair: old_string.replace(idx_pair[0], idx_pair[1]),
+            zip(affected_indices, new_indices),
+            state_indices,
+        )
+
+        # We now put together the indices in the notation numpy's einsum requires
+        # This notation allows for the state, the matrix, or both to be broadcasted
+        einsum_indices = (
+            f"...{new_indices}{affected_indices},...{state_indices}->...{new_state_indices}"
+        )
+
+        return self._einsum(einsum_indices, mat, state)
+
+    def _apply_diagonal_unitary(self, state, phases, wires):
+        r"""Apply multiplication of a phase vector to subsystems of the quantum state.
+
+        This represents the multiplication with diagonal gates in a more efficient manner.
+
+        Args:
+            state (array[complex]): input state
+            phases (array): vector to multiply
+            wires (Wires): target wires
+
+        Returns:
+            array[complex]: output state
+        """
+        # translate to wire labels used by device
+        device_wires = self.map_wires(wires)
+        dim = 2 ** len(device_wires)
+        batch_size = self._get_batch_size(phases, (dim,), dim)
+
+        # reshape vectors
+        shape = [2] * len(device_wires)
+        if batch_size is not None:
+            shape.insert(0, batch_size)
+        phases = self._cast(self._reshape(phases, shape), dtype=self.C_DTYPE)
+
+        state_indices = ABC[: self.num_wires]
+        affected_indices = "".join(ABC_ARRAY[list(device_wires)].tolist())
+
+        einsum_indices = f"...{affected_indices},...{state_indices}->...{state_indices}"
+        return self._einsum(einsum_indices, phases, state)
+
+    def reset(self):
+        """Reset the device"""
+        super().reset()
+
+        # init the state vector to |00..0>
+        self._state = self._create_basis_state(0)
+        self._pre_rotated_state = self._state
+
+    def analytic_probability(self, wires=None):
+        if self._state is None:
+            return None
+
+        dim = 2**self.num_wires
+        batch_size = self._get_batch_size(self._state, [2] * self.num_wires, dim)
+        flat_state = self._reshape(
+            self._state, (batch_size, dim) if batch_size is not None else (dim,)
+        )
+        real_state = self._real(flat_state)
+        imag_state = self._imag(flat_state)
+        return self.marginal_prob(real_state**2 + imag_state**2, wires)
+
+    def classical_shadow(self, obs, circuit):
+        """
+        Returns the measured bits and recipes in the classical shadow protocol.
+
+        The protocol is described in detail in the `classical shadows paper <https://arxiv.org/abs/2002.08953>`_.
+        This measurement process returns the randomized Pauli measurements (the ``recipes``)
+        that are performed for each qubit and snapshot as an integer:
+
+        - 0 for Pauli X,
+        - 1 for Pauli Y, and
+        - 2 for Pauli Z.
+
+        It also returns the measurement results (the ``bits``); 0 if the 1 eigenvalue
+        is sampled, and 1 if the -1 eigenvalue is sampled.
+
+        The device shots are used to specify the number of snapshots. If ``T`` is the number
+        of shots and ``n`` is the number of qubits, then both the measured bits and the
+        Pauli measurements have shape ``(T, n)``.
+
+        This implementation leverages vectorization and offers a significant speed-up over
+        the generic implementation.
+
+        .. Note::
+
+            This method internally calls ``np.einsum`` which supports at most 52 indices,
+            thus the classical shadow measurement for this device supports at most 52
+            qubits.
+
+        .. seealso:: :func:`~pennylane.classical_shadow`
+
+        Args:
+            obs (~.pennylane.measurements.ClassicalShadowMP): The classical shadow measurement process
+            circuit (~.tape.QuantumTape): The quantum tape that is being executed
+
+        Returns:
+            tensor_like[int]: A tensor with shape ``(2, T, n)``, where the first row represents
+            the measured bits and the second represents the recipes used.
+        """
+        wires = obs.wires
+        seed = obs.seed
+
+        n_qubits = len(wires)
+        n_snapshots = self.shots
+        device_qubits = len(self.wires)
+        mapped_wires = np.array(self.map_wires(wires))
+
+        # seed the random measurement generation so that recipes
+        # are the same for different executions with the same seed
+        rng = np.random.RandomState(seed)
+        recipes = rng.randint(0, 3, size=(n_snapshots, n_qubits))
+
+        obs_list = self._stack(
+            [
+                qml.PauliX.compute_matrix(),
+                qml.PauliY.compute_matrix(),
+                qml.PauliZ.compute_matrix(),
+            ]
+        )
+        uni_list = self._stack(
+            [
+                qml.Hadamard.compute_matrix(),
+                qml.Hadamard.compute_matrix() @ qml.RZ.compute_matrix(-np.pi / 2),
+                qml.Identity.compute_matrix(),
+            ]
+        )
+        obs = obs_list[recipes]
+        uni = uni_list[recipes]
+
+        # There's a significant speedup if we use the following iterative
+        # process to perform the randomized Pauli measurements:
+        #   1. Randomly generate Pauli observables for all snapshots for
+        #      a single qubit (e.g. the first qubit).
+        #   2. Compute the expectation of each Pauli observable on the first
+        #      qubit by tracing out all other qubits.
+        #   3. Sample the first qubit based on each Pauli expectation.
+        #   4. For all snapshots, determine the collapsed state of the remaining
+        #      qubits based on the sample result.
+        #   4. Repeat iteratively until no qubits are remaining.
+        #
+        # Observe that after the first iteration, the second qubit will become the
+        # "first" qubit in the process. The advantage to this approach as opposed to
+        # simulataneously computing the Pauli expectations for each qubit is that
+        # the partial traces are computed over iteratively smaller subsystems, leading
+        # to a significant speed-up.
+
+        # transpose the state so that the measured wires appear first
+        unmeasured_wires = [i for i in range(len(self.wires)) if i not in mapped_wires]
+        transposed_state = np.transpose(self._state, axes=mapped_wires.tolist() + unmeasured_wires)
+
+        outcomes = np.zeros((n_snapshots, n_qubits))
+        stacked_state = self._stack([transposed_state for _ in range(n_snapshots)])
+
+        for i in range(n_qubits):
+            # trace out every qubit except the first
+            first_qubit_state = self._einsum(
+                f"{ABC[device_qubits - i + 1]}{ABC[:device_qubits - i]},{ABC[device_qubits - i + 1]}{ABC[device_qubits - i]}{ABC[1:device_qubits - i]}"
+                f"->{ABC[device_qubits - i + 1]}a{ABC[device_qubits - i]}",
+                stacked_state,
+                self._conj(stacked_state),
+            )
+
+            # sample the observables on the first qubit
+            probs = (self._einsum("abc,acb->a", first_qubit_state, obs[:, i]) + 1) / 2
+            samples = np.random.uniform(0, 1, size=probs.shape) > probs
+            outcomes[:, i] = samples
+
+            # collapse the state of the remaining qubits; the next qubit in line
+            # becomes the first qubit for the next iteration
+            rotated_state = self._einsum("ab...,acb->ac...", stacked_state, uni[:, i])
+            stacked_state = rotated_state[np.arange(n_snapshots), self._cast(samples, np.int8)]
+
+            # re-normalize the collapsed state
+            norms = np.sqrt(
+                np.sum(
+                    np.abs(stacked_state) ** 2, tuple(range(1, device_qubits - i)), keepdims=True
+                )
+            )
+            stacked_state /= norms
+
+        return self._cast(self._stack([outcomes, recipes]), dtype=np.int8)
+
+    def _get_diagonalizing_gates(self, circuit: qml.tape.QuantumTape) -> List[Operation]:
+        meas_filtered = [
+            m
+            for m in circuit.measurements
+            if m.obs is None or not isinstance(m.obs, qml.Hamiltonian)
+        ]
+        return super()._get_diagonalizing_gates(qml.tape.QuantumScript(measurements=meas_filtered))
diff --git a/pennylane/devices/default_qubit_tf.py b/pennylane/devices/default_qubit_tf.py
index 702bf34af0b..3b09adf31e7 100644
--- a/pennylane/devices/default_qubit_tf.py
+++ b/pennylane/devices/default_qubit_tf.py
@@ -11,7 +11,7 @@
 # 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.
-"""This module contains a TensorFlow implementation of the :class:`~.DefaultQubit`
+"""This module contains a TensorFlow implementation of the :class:`~.DefaultQubitLegacy`
 reference plugin.
 """
 import itertools
@@ -35,12 +35,12 @@
 
 
 from pennylane.math.single_dispatch import _ndim_tf
-from . import DefaultQubit
-from .default_qubit import tolerance
+from . import DefaultQubitLegacy
+from .default_qubit_legacy import tolerance
 
 
-class DefaultQubitTF(DefaultQubit):
-    """Simulator plugin based on ``"default.qubit"``, written using TensorFlow.
+class DefaultQubitTF(DefaultQubitLegacy):
+    """Simulator plugin based on ``"default.qubit.legacy"``, written using TensorFlow.
 
     **Short name:** ``default.qubit.tf``
 
@@ -250,12 +250,12 @@ def _apply_state_vector(self, state, device_wires):
         ravelled_indices = np.ravel_multi_index(unravelled_indices.T, [2] * self.num_wires)
 
         if batch_size:
-            # This is the only logical branch that differs from DefaultQubit
+            # This is the only logical branch that differs from DefaultQubitLegacy
             raise NotImplementedError(
                 "Parameter broadcasting is not supported together with initializing the state "
                 "vector of a subsystem of the device when using DefaultQubitTF."
             )
-        # The following line is unchanged in the "else"-clause in DefaultQubit's implementation
+        # The following line is unchanged in the "else"-clause in DefaultQubitLegacy's implementation
         state = self._scatter(ravelled_indices, state, [2**self.num_wires])
         state = self._reshape(state, output_shape)
         self._state = self._asarray(state, dtype=self.C_DTYPE)
diff --git a/pennylane/devices/default_qubit_torch.py b/pennylane/devices/default_qubit_torch.py
index 61cd6e90cc9..97bc9fbe4a7 100644
--- a/pennylane/devices/default_qubit_torch.py
+++ b/pennylane/devices/default_qubit_torch.py
@@ -11,7 +11,7 @@
 # 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.
-"""This module contains a PyTorch implementation of the :class:`~.DefaultQubit`
+"""This module contains a PyTorch implementation of the :class:`~.DefaultQubitLegacy`
 reference plugin.
 """
 import warnings
@@ -31,14 +31,14 @@
 
 import numpy as np
 from pennylane.ops.qubit.attributes import diagonal_in_z_basis
-from . import DefaultQubit
+from . import DefaultQubitLegacy
 
 logger = logging.getLogger(__name__)
 logger.addHandler(logging.NullHandler())
 
 
-class DefaultQubitTorch(DefaultQubit):
-    """Simulator plugin based on ``"default.qubit"``, written using PyTorch.
+class DefaultQubitTorch(DefaultQubitLegacy):
+    """Simulator plugin based on ``"default.qubit.legacy"``, written using PyTorch.
 
     **Short name:** ``default.qubit.torch``
 
diff --git a/pennylane/devices/default_qutrit.py b/pennylane/devices/default_qutrit.py
index 4f70b22f6a2..44bd1ed79e6 100644
--- a/pennylane/devices/default_qutrit.py
+++ b/pennylane/devices/default_qutrit.py
@@ -24,7 +24,7 @@
 import pennylane as qml  # pylint: disable=unused-import
 from pennylane import QutritDevice, QutritBasisState, DeviceError
 from pennylane.wires import WireError
-from pennylane.devices.default_qubit import _get_slice
+from pennylane.devices.default_qubit_legacy import _get_slice
 from .._version import __version__
 
 # tolerance for numerical errors
@@ -39,12 +39,12 @@ class DefaultQutrit(QutritDevice):
     .. warning::
 
         The API of ``DefaultQutrit`` will be updated soon to follow a new device interface described
-        in :class:`pennylane.devices.experimental.Device`.
+        in :class:`pennylane.devices.Device`.
 
         This change will not alter device behaviour for most workflows, but may have implications for
         plugin developers and users who directly interact with device methods. Please consult
-        :class:`pennylane.devices.experimental.Device` and the implementation in
-        :class:`pennylane.devices.experimental.DefaultQubit2` for more information on what the new
+        :class:`pennylane.devices.Device` and the implementation in
+        :class:`pennylane.devices.DefaultQubit` for more information on what the new
         interface will look like and be prepared to make updates in a coming release. If you have any
         feedback on these changes, please create an
         `issue <https://github.com/PennyLaneAI/pennylane/issues>`_ or post in our
diff --git a/pennylane/devices/experimental/device_api.py b/pennylane/devices/device_api.py
similarity index 99%
rename from pennylane/devices/experimental/device_api.py
rename to pennylane/devices/device_api.py
index 554d25c7d2e..1fb8dcb81b1 100644
--- a/pennylane/devices/experimental/device_api.py
+++ b/pennylane/devices/device_api.py
@@ -96,7 +96,7 @@ class Device(abc.ABC):
 
         >>> op = qml.Permute(["c", 3,"a",2,0], wires=[3,2,"a",0,"c"])
         >>> circuit = qml.tape.QuantumScript([op], [qml.state()])
-        >>> dev = DefaultQubit2()
+        >>> dev = DefaultQubit()
         >>> dev.execute(circuit)
         MatrixUndefinedError
         >>> circuit = qml.tape.QuantumScript([qml.Rot(1.2, 2.3, 3.4, 0)], [qml.expval(qml.PauliZ(0))])
diff --git a/pennylane/devices/experimental/execution_config.py b/pennylane/devices/execution_config.py
similarity index 100%
rename from pennylane/devices/experimental/execution_config.py
rename to pennylane/devices/execution_config.py
diff --git a/pennylane/devices/experimental/__init__.py b/pennylane/devices/experimental/__init__.py
deleted file mode 100644
index 350e2f8ae19..00000000000
--- a/pennylane/devices/experimental/__init__.py
+++ /dev/null
@@ -1,20 +0,0 @@
-# Copyright 2018-2021 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.
-"""This experimental directory contains the next generation interface
-for PennyLane devices.
-
-"""
-from .execution_config import ExecutionConfig, DefaultExecutionConfig
-from .device_api import Device
-from .default_qubit_2 import DefaultQubit2
diff --git a/pennylane/devices/experimental/default_qubit_2.py b/pennylane/devices/experimental/default_qubit_2.py
deleted file mode 100644
index 9052c01313e..00000000000
--- a/pennylane/devices/experimental/default_qubit_2.py
+++ /dev/null
@@ -1,529 +0,0 @@
-# Copyright 2018-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.
-"""
-This module contains the next generation successor to default qubit
-"""
-
-from functools import partial
-from numbers import Number
-from typing import Union, Callable, Tuple, Optional, Sequence
-import concurrent.futures
-import numpy as np
-
-from pennylane.tape import QuantumTape, QuantumScript
-from pennylane.typing import Result, ResultBatch
-from pennylane.transforms import convert_to_numpy_parameters
-from pennylane.transforms.core import TransformProgram
-
-from . import Device
-from .execution_config import ExecutionConfig, DefaultExecutionConfig
-from ..qubit.simulate import simulate, get_final_state, measure_final_state
-from ..qubit.preprocess import (
-    preprocess,
-    validate_and_expand_adjoint,
-    validate_multiprocessing_workers,
-    validate_device_wires,
-)
-from ..qubit.adjoint_jacobian import adjoint_jacobian, adjoint_vjp, adjoint_jvp
-
-Result_or_ResultBatch = Union[Result, ResultBatch]
-QuantumTapeBatch = Sequence[QuantumTape]
-QuantumTape_or_Batch = Union[QuantumTape, QuantumTapeBatch]
-# always a function from a resultbatch to either a result or a result batch
-PostprocessingFn = Callable[[ResultBatch], Result_or_ResultBatch]
-
-
-class DefaultQubit2(Device):
-    """A PennyLane device written in Python and capable of backpropagation derivatives.
-
-    Args:
-        shots (int, Sequence[int], Sequence[Union[int, Sequence[int]]]): The default number of shots to use in executions involving
-            this device.
-        seed (Union[str, None, int, array_like[int], SeedSequence, BitGenerator, Generator]): A
-            seed-like parameter matching that of ``seed`` for ``numpy.random.default_rng`` or
-            a request to seed from numpy's global random number generator.
-            The default, ``seed="global"`` pulls a seed from NumPy's global generator. ``seed=None``
-            will pull a seed from the OS entropy.
-        max_workers (int): A ``ProcessPoolExecutor`` executes tapes asynchronously
-            using a pool of at most ``max_workers`` processes. If ``max_workers`` is ``None``,
-            only the current process executes tapes. If you experience any
-            issue, say using JAX, TensorFlow, Torch, try setting ``max_workers`` to ``None``.
-
-    **Example:**
-
-    .. code-block:: python
-
-        n_layers = 5
-        n_wires = 10
-        num_qscripts = 5
-
-        shape = qml.StronglyEntanglingLayers.shape(n_layers=n_layers, n_wires=n_wires)
-        rng = qml.numpy.random.default_rng(seed=42)
-
-        qscripts = []
-        for i in range(num_qscripts):
-            params = rng.random(shape)
-            op = qml.StronglyEntanglingLayers(params, wires=range(n_wires))
-            qs = qml.tape.QuantumScript([op], [qml.expval(qml.PauliZ(0))])
-            qscripts.append(qs)
-
-    >>> dev = DefaultQubit2()
-    >>> program, execution_config = dev.preprocess()
-    >>> new_batch, post_processing_fn = program(qscripts)
-    >>> results = dev.execute(new_batch, execution_config=execution_config)
-    >>> post_processing_fn(results)
-    [-0.0006888975950537501,
-    0.025576307134457577,
-    -0.0038567269892757494,
-    0.1339705146860149,
-    -0.03780669772690448]
-
-    Suppose one has a processor with 5 cores or more, these scripts can be executed in
-    parallel as follows
-
-    >>> dev = DefaultQubit2(max_workers=5)
-    >>> program, execution_config = dev.preprocess()
-    >>> new_batch, post_processing_fn = program(qscripts)
-    >>> results = dev.execute(new_batch, execution_config=execution_config)
-    >>> post_processing_fn(results)
-
-    If you monitor your CPU usage, you should see 5 new Python processes pop up to
-    crunch through those ``QuantumScript``'s. Beware not oversubscribing your machine.
-    This may happen if a single device already uses many cores, if NumPy uses a multi-
-    threaded BLAS library like MKL or OpenBLAS for example. The number of threads per
-    process times the number of processes should not exceed the number of cores on your
-    machine. You can control the number of threads per process with the environment
-    variables:
-
-    * OMP_NUM_THREADS
-    * MKL_NUM_THREADS
-    * OPENBLAS_NUM_THREADS
-
-    where the last two are specific to the MKL and OpenBLAS libraries specifically.
-
-    This device currently supports backpropagation derivatives:
-
-    >>> from pennylane.devices import ExecutionConfig
-    >>> dev.supports_derivatives(ExecutionConfig(gradient_method="backprop"))
-    True
-
-    For example, we can use jax to jit computing the derivative:
-
-    .. code-block:: python
-
-        import jax
-
-        @jax.jit
-        def f(x):
-            qs = qml.tape.QuantumScript([qml.RX(x, 0)], [qml.expval(qml.PauliZ(0))])
-            program, execution_config = dev.preprocess()
-            new_batch, post_processing_fn = program([qs])
-            results = dev.execute(new_batch, execution_config=execution_config)
-            return post_processing_fn(results)
-
-    >>> f(jax.numpy.array(1.2))
-    DeviceArray(0.36235774, dtype=float32)
-    >>> jax.grad(f)(jax.numpy.array(1.2))
-    DeviceArray(-0.93203914, dtype=float32, weak_type=True)
-
-    """
-
-    @property
-    def name(self):
-        """The name of the device."""
-        return "default.qubit.2"
-
-    def __init__(self, wires=None, shots=None, seed="global", max_workers=None) -> None:
-        super().__init__(wires=wires, shots=shots)
-        self._max_workers = max_workers
-        seed = np.random.randint(0, high=10000000) if seed == "global" else seed
-        self._rng = np.random.default_rng(seed)
-        self._debugger = None
-
-    def supports_derivatives(
-        self,
-        execution_config: Optional[ExecutionConfig] = None,
-        circuit: Optional[QuantumTape] = None,
-    ) -> bool:
-        """Check whether or not derivatives are available for a given configuration and circuit.
-
-        ``DefaultQubit2`` supports backpropagation derivatives with analytic results, as well as
-        adjoint differentiation.
-
-        Args:
-            execution_config (ExecutionConfig): The configuration of the desired derivative calculation
-            circuit (QuantumTape): An optional circuit to check derivatives support for.
-
-        Returns:
-            Bool: Whether or not a derivative can be calculated provided the given information
-
-        """
-        if execution_config is None:
-            return True
-        # backpropagation currently supported for all supported circuits
-        # will later need to add logic if backprop requested with finite shots
-        # do once device accepts finite shots
-        if (
-            execution_config.gradient_method == "backprop"
-            and execution_config.device_options.get("max_workers", self._max_workers) is None
-        ):
-            return True
-
-        if execution_config.gradient_method == "adjoint" and execution_config.use_device_gradient:
-            if circuit is None:
-                return True
-
-            return isinstance(validate_and_expand_adjoint(circuit)[0][0], QuantumScript)
-
-        return False
-
-    def preprocess(
-        self,
-        execution_config: ExecutionConfig = DefaultExecutionConfig,
-    ) -> Tuple[QuantumTapeBatch, PostprocessingFn, ExecutionConfig]:
-        """This function defines the device transform program to be applied and an updated device configuration.
-
-        Args:
-            execution_config (Union[ExecutionConfig, Sequence[ExecutionConfig]]): A data structure describing the
-                parameters needed to fully describe the execution.
-
-        Returns:
-            TransformProgram, ExecutionConfig: A transform program that when called returns QuantumTapes that the device
-            can natively execute as well as a postprocessing function to be called after execution, and a configuration with
-            unset specifications filled in.
-
-        This device:
-
-        * Supports any qubit operations that provide a matrix
-        * Currently does not support finite shots
-        * Currently does not intrinsically support parameter broadcasting
-
-        """
-        transform_program = TransformProgram()
-        # Validate device wires
-        transform_program.add_transform(validate_device_wires, self)
-
-        # Validate multi processing
-        max_workers = execution_config.device_options.get("max_workers", self._max_workers)
-        transform_program.add_transform(validate_multiprocessing_workers, max_workers, self)
-
-        # General preprocessing (Validate measurement, expand, adjoint expand, broadcast expand)
-        transform_program_preprocess, config = preprocess(execution_config=execution_config)
-        transform_program = transform_program + transform_program_preprocess
-        return transform_program, config
-
-    def execute(
-        self,
-        circuits: QuantumTape_or_Batch,
-        execution_config: ExecutionConfig = DefaultExecutionConfig,
-    ) -> Result_or_ResultBatch:
-        is_single_circuit = False
-        if isinstance(circuits, QuantumScript):
-            is_single_circuit = True
-            circuits = [circuits]
-
-        if self.tracker.active:
-            for c in circuits:
-                self.tracker.update(resources=c.specs["resources"])
-            self.tracker.update(batches=1, executions=len(circuits))
-            self.tracker.record()
-
-        max_workers = execution_config.device_options.get("max_workers", self._max_workers)
-        interface = (
-            execution_config.interface
-            if execution_config.gradient_method in {"backprop", None}
-            else None
-        )
-        if max_workers is None:
-            results = tuple(
-                simulate(c, rng=self._rng, debugger=self._debugger, interface=interface)
-                for c in circuits
-            )
-        else:
-            vanilla_circuits = [convert_to_numpy_parameters(c) for c in circuits]
-            seeds = self._rng.integers(2**31 - 1, size=len(vanilla_circuits))
-            _wrap_simulate = partial(simulate, debugger=None, interface=interface)
-            with concurrent.futures.ProcessPoolExecutor(max_workers=max_workers) as executor:
-                exec_map = executor.map(_wrap_simulate, vanilla_circuits, seeds)
-                results = tuple(exec_map)
-
-            # reset _rng to mimic serial behavior
-            self._rng = np.random.default_rng(self._rng.integers(2**31 - 1))
-
-        return results[0] if is_single_circuit else results
-
-    def compute_derivatives(
-        self,
-        circuits: QuantumTape_or_Batch,
-        execution_config: ExecutionConfig = DefaultExecutionConfig,
-    ):
-        is_single_circuit = False
-        if isinstance(circuits, QuantumScript):
-            is_single_circuit = True
-            circuits = [circuits]
-
-        if self.tracker.active:
-            self.tracker.update(derivative_batches=1, derivatives=len(circuits))
-            self.tracker.record()
-
-        max_workers = execution_config.device_options.get("max_workers", self._max_workers)
-        if max_workers is None:
-            res = tuple(adjoint_jacobian(circuit) for circuit in circuits)
-        else:
-            vanilla_circuits = [convert_to_numpy_parameters(c) for c in circuits]
-            with concurrent.futures.ProcessPoolExecutor(max_workers=max_workers) as executor:
-                exec_map = executor.map(adjoint_jacobian, vanilla_circuits)
-                res = tuple(exec_map)
-
-            # reset _rng to mimic serial behavior
-            self._rng = np.random.default_rng(self._rng.integers(2**31 - 1))
-
-        return res[0] if is_single_circuit else res
-
-    def execute_and_compute_derivatives(
-        self,
-        circuits: QuantumTape_or_Batch,
-        execution_config: ExecutionConfig = DefaultExecutionConfig,
-    ):
-        is_single_circuit = False
-        if isinstance(circuits, QuantumScript):
-            is_single_circuit = True
-            circuits = [circuits]
-
-        if self.tracker.active:
-            for c in circuits:
-                self.tracker.update(resources=c.specs["resources"])
-            self.tracker.update(
-                execute_and_derivative_batches=1,
-                executions=len(circuits),
-                derivatives=len(circuits),
-            )
-            self.tracker.record()
-
-        max_workers = execution_config.device_options.get("max_workers", self._max_workers)
-        if max_workers is None:
-            results = tuple(
-                _adjoint_jac_wrapper(c, rng=self._rng, debugger=self._debugger) for c in circuits
-            )
-        else:
-            vanilla_circuits = [convert_to_numpy_parameters(c) for c in circuits]
-            seeds = self._rng.integers(2**31 - 1, size=len(vanilla_circuits))
-
-            with concurrent.futures.ProcessPoolExecutor(max_workers=max_workers) as executor:
-                results = tuple(executor.map(_adjoint_jac_wrapper, vanilla_circuits, seeds))
-
-            # reset _rng to mimic serial behavior
-            self._rng = np.random.default_rng(self._rng.integers(2**31 - 1))
-
-        results, jacs = tuple(zip(*results))
-        return (results[0], jacs[0]) if is_single_circuit else (results, jacs)
-
-    def supports_jvp(
-        self,
-        execution_config: Optional[ExecutionConfig] = None,
-        circuit: Optional[QuantumTape] = None,
-    ) -> bool:
-        """Whether or not this device defines a custom jacobian vector product.
-
-        ``DefaultQubit2`` supports backpropagation derivatives with analytic results, as well as
-        adjoint differentiation.
-
-        Args:
-            execution_config (ExecutionConfig): The configuration of the desired derivative calculation
-            circuit (QuantumTape): An optional circuit to check derivatives support for.
-
-        Returns:
-            bool: Whether or not a derivative can be calculated provided the given information
-        """
-        return self.supports_derivatives(execution_config, circuit)
-
-    def compute_jvp(
-        self,
-        circuits: QuantumTape_or_Batch,
-        tangents: Tuple[Number],
-        execution_config: ExecutionConfig = DefaultExecutionConfig,
-    ):
-        is_single_circuit = False
-        if isinstance(circuits, QuantumScript):
-            is_single_circuit = True
-            circuits = [circuits]
-            tangents = [tangents]
-
-        if self.tracker.active:
-            self.tracker.update(jvp_batches=1, jvps=len(circuits))
-            self.tracker.record()
-
-        max_workers = execution_config.device_options.get("max_workers", self._max_workers)
-        if max_workers is None:
-            res = tuple(adjoint_jvp(circuit, tans) for circuit, tans in zip(circuits, tangents))
-        else:
-            vanilla_circuits = [convert_to_numpy_parameters(c) for c in circuits]
-            with concurrent.futures.ProcessPoolExecutor(max_workers=max_workers) as executor:
-                res = tuple(executor.map(adjoint_jvp, vanilla_circuits, tangents))
-
-            # reset _rng to mimic serial behavior
-            self._rng = np.random.default_rng(self._rng.integers(2**31 - 1))
-
-        return res[0] if is_single_circuit else res
-
-    def execute_and_compute_jvp(
-        self,
-        circuits: QuantumTape_or_Batch,
-        tangents: Tuple[Number],
-        execution_config: ExecutionConfig = DefaultExecutionConfig,
-    ):
-        is_single_circuit = False
-        if isinstance(circuits, QuantumScript):
-            is_single_circuit = True
-            circuits = [circuits]
-            tangents = [tangents]
-
-        if self.tracker.active:
-            for c in circuits:
-                self.tracker.update(resources=c.specs["resources"])
-            self.tracker.update(
-                execute_and_jvp_batches=1, executions=len(circuits), jvps=len(circuits)
-            )
-            self.tracker.record()
-
-        max_workers = execution_config.device_options.get("max_workers", self._max_workers)
-        if max_workers is None:
-            results = tuple(
-                _adjoint_jvp_wrapper(c, t, rng=self._rng, debugger=self._debugger)
-                for c, t in zip(circuits, tangents)
-            )
-        else:
-            vanilla_circuits = [convert_to_numpy_parameters(c) for c in circuits]
-            seeds = self._rng.integers(2**31 - 1, size=len(vanilla_circuits))
-
-            with concurrent.futures.ProcessPoolExecutor(max_workers=max_workers) as executor:
-                results = tuple(
-                    executor.map(_adjoint_jvp_wrapper, vanilla_circuits, tangents, seeds)
-                )
-
-            # reset _rng to mimic serial behavior
-            self._rng = np.random.default_rng(self._rng.integers(2**31 - 1))
-
-        results, jvps = tuple(zip(*results))
-        return (results[0], jvps[0]) if is_single_circuit else (results, jvps)
-
-    def supports_vjp(
-        self,
-        execution_config: Optional[ExecutionConfig] = None,
-        circuit: Optional[QuantumTape] = None,
-    ) -> bool:
-        """Whether or not this device defines a custom vector jacobian product.
-
-        ``DefaultQubit2`` supports backpropagation derivatives with analytic results, as well as
-        adjoint differentiation.
-
-        Args:
-            execution_config (ExecutionConfig): A description of the hyperparameters for the desired computation.
-            circuit (None, QuantumTape): A specific circuit to check differentation for.
-
-        Returns:
-            bool: Whether or not a derivative can be calculated provided the given information
-        """
-        return self.supports_derivatives(execution_config, circuit)
-
-    def compute_vjp(
-        self,
-        circuits: QuantumTape_or_Batch,
-        cotangents: Tuple[Number],
-        execution_config: ExecutionConfig = DefaultExecutionConfig,
-    ):
-        is_single_circuit = False
-        if isinstance(circuits, QuantumScript):
-            is_single_circuit = True
-            circuits = [circuits]
-            cotangents = [cotangents]
-
-        if self.tracker.active:
-            self.tracker.update(vjp_batches=1, vjps=len(circuits))
-            self.tracker.record()
-
-        max_workers = execution_config.device_options.get("max_workers", self._max_workers)
-        if max_workers is None:
-            res = tuple(adjoint_vjp(circuit, cots) for circuit, cots in zip(circuits, cotangents))
-        else:
-            vanilla_circuits = [convert_to_numpy_parameters(c) for c in circuits]
-            with concurrent.futures.ProcessPoolExecutor(max_workers=max_workers) as executor:
-                res = tuple(executor.map(adjoint_vjp, vanilla_circuits, cotangents))
-
-            # reset _rng to mimic serial behavior
-            self._rng = np.random.default_rng(self._rng.integers(2**31 - 1))
-
-        return res[0] if is_single_circuit else res
-
-    def execute_and_compute_vjp(
-        self,
-        circuits: QuantumTape_or_Batch,
-        cotangents: Tuple[Number],
-        execution_config: ExecutionConfig = DefaultExecutionConfig,
-    ):
-        is_single_circuit = False
-        if isinstance(circuits, QuantumScript):
-            is_single_circuit = True
-            circuits = [circuits]
-            cotangents = [cotangents]
-
-        if self.tracker.active:
-            for c in circuits:
-                self.tracker.update(resources=c.specs["resources"])
-            self.tracker.update(
-                execute_and_vjp_batches=1, executions=len(circuits), vjps=len(circuits)
-            )
-            self.tracker.record()
-
-        max_workers = execution_config.device_options.get("max_workers", self._max_workers)
-        if max_workers is None:
-            results = tuple(
-                _adjoint_vjp_wrapper(c, t, rng=self._rng, debugger=self._debugger)
-                for c, t in zip(circuits, cotangents)
-            )
-        else:
-            vanilla_circuits = [convert_to_numpy_parameters(c) for c in circuits]
-            seeds = self._rng.integers(2**31 - 1, size=len(vanilla_circuits))
-
-            with concurrent.futures.ProcessPoolExecutor(max_workers=max_workers) as executor:
-                results = tuple(
-                    executor.map(_adjoint_vjp_wrapper, vanilla_circuits, cotangents, seeds)
-                )
-
-            # reset _rng to mimic serial behavior
-            self._rng = np.random.default_rng(self._rng.integers(2**31 - 1))
-
-        results, vjps = tuple(zip(*results))
-        return (results[0], vjps[0]) if is_single_circuit else (results, vjps)
-
-
-def _adjoint_jac_wrapper(c, rng=None, debugger=None):
-    state, is_state_batched = get_final_state(c, debugger=debugger)
-    jac = adjoint_jacobian(c, state=state)
-    res = measure_final_state(c, state, is_state_batched, rng=rng)
-    return res, jac
-
-
-def _adjoint_jvp_wrapper(c, t, rng=None, debugger=None):
-    state, is_state_batched = get_final_state(c, debugger=debugger)
-    jvp = adjoint_jvp(c, t, state=state)
-    res = measure_final_state(c, state, is_state_batched, rng=rng)
-    return res, jvp
-
-
-def _adjoint_vjp_wrapper(c, t, rng=None, debugger=None):
-    state, is_state_batched = get_final_state(c, debugger=debugger)
-    vjp = adjoint_vjp(c, t, state=state)
-    res = measure_final_state(c, state, is_state_batched, rng=rng)
-    return res, vjp
diff --git a/pennylane/devices/null_qubit.py b/pennylane/devices/null_qubit.py
index 99ec456ed91..28032a42194 100644
--- a/pennylane/devices/null_qubit.py
+++ b/pennylane/devices/null_qubit.py
@@ -33,12 +33,12 @@ class NullQubit(QubitDevice):
     .. warning::
 
         The API of ``NullQubit`` will be updated soon to follow a new device interface described
-        in :class:`pennylane.devices.experimental.Device`.
+        in :class:`pennylane.devices.Device`.
 
         This change will not alter device behaviour for most workflows, but may have implications for
         plugin developers and users who directly interact with device methods. Please consult
-        :class:`pennylane.devices.experimental.Device` and the implementation in
-        :class:`pennylane.devices.experimental.DefaultQubit2` for more information on what the new
+        :class:`pennylane.devices.Device` and the implementation in
+        :class:`pennylane.devices.DefaultQubit` for more information on what the new
         interface will look like and be prepared to make updates in a coming release. If you have any
         feedback on these changes, please create an
         `issue <https://github.com/PennyLaneAI/pennylane/issues>`_ or post in our
diff --git a/pennylane/devices/qubit/preprocess.py b/pennylane/devices/qubit/preprocess.py
index c1be5a802d3..326ba6343ce 100644
--- a/pennylane/devices/qubit/preprocess.py
+++ b/pennylane/devices/qubit/preprocess.py
@@ -37,7 +37,7 @@
 from pennylane.transforms.core import transform, TransformProgram
 from pennylane.wires import WireError
 
-from ..experimental import ExecutionConfig, DefaultExecutionConfig
+from pennylane.devices import ExecutionConfig, DefaultExecutionConfig
 
 PostprocessingFn = Callable[[ResultBatch], Union[Result, ResultBatch]]
 
@@ -86,7 +86,7 @@ def _accepted_adjoint_operator(op: qml.operation.Operator) -> bool:
 def _operator_decomposition_gen(
     op: qml.operation.Operator, acceptance_function: Callable[[qml.operation.Operator], bool]
 ) -> Generator[qml.operation.Operator, None, None]:
-    """A generator that yields the next operation that is accepted by DefaultQubit2."""
+    """A generator that yields the next operation that is accepted by DefaultQubit."""
     if acceptance_function(op):
         yield op
     else:
@@ -94,7 +94,7 @@ def _operator_decomposition_gen(
             decomp = op.decomposition()
         except qml.operation.DecompositionUndefinedError as e:
             raise DeviceError(
-                f"Operator {op} not supported on DefaultQubit2. Must provide either a matrix or a decomposition."
+                f"Operator {op} not supported on DefaultQubit. Must provide either a matrix or a decomposition."
             ) from e
 
         for sub_op in decomp:
@@ -112,7 +112,7 @@ def validate_device_wires(
 
     Args:
         tape (QuantumTape): a quantum circuit.
-        device (pennylane.devices.experimental.Device): The device to be checked.
+        device (pennylane.devices.Device): The device to be checked.
 
     Returns:
         pennylane.QNode or qfunc or Tuple[List[.QuantumTape], Callable]: If a QNode is passed,
@@ -160,7 +160,7 @@ def validate_multiprocessing_workers(
     Args:
         tape (QuantumTape): a quantum circuit.
         max_workers (int): Maximal number of multiprocessing workers
-        device (pennylane.devices.experimental.Device): The device to be checked.
+        device (pennylane.devices.Device): The device to be checked.
 
     Returns:
         pennylane.QNode or qfunc or Tuple[List[.QuantumTape], Callable]: If a QNode is passed,
@@ -380,9 +380,9 @@ def expand_fn(tape: qml.tape.QuantumTape) -> (Sequence[qml.tape.QuantumTape], Ca
     for observable in tape.observables:
         if isinstance(observable, Tensor):
             if any(o.name not in _observables for o in observable.obs):
-                raise DeviceError(f"Observable {observable} not supported on DefaultQubit2")
+                raise DeviceError(f"Observable {observable} not supported on DefaultQubit")
         elif observable.name not in _observables:
-            raise DeviceError(f"Observable {observable} not supported on DefaultQubit2")
+            raise DeviceError(f"Observable {observable} not supported on DefaultQubit")
 
     def null_postprocessing(results):
         """A postprocesing function returned by a transform that only converts the batch of results
diff --git a/pennylane/drawer/draw.py b/pennylane/drawer/draw.py
index 560b6e4df5b..7b6454dd65e 100644
--- a/pennylane/drawer/draw.py
+++ b/pennylane/drawer/draw.py
@@ -240,9 +240,7 @@ def _draw_qnode(
 ):
     @wraps(qnode)
     def wrapper(*args, **kwargs):
-        if expansion_strategy == "device" and isinstance(
-            qnode.device, qml.devices.experimental.Device
-        ):
+        if expansion_strategy == "device" and isinstance(qnode.device, qml.devices.Device):
             qnode.construct(args, kwargs)
             program, _ = qnode.device.preprocess()
             tapes = program([qnode.tape])
@@ -543,9 +541,7 @@ def _draw_mpl_qnode(
 ):
     @wraps(qnode)
     def wrapper(*args, **kwargs_qnode):
-        if expansion_strategy == "device" and isinstance(
-            qnode.device, qml.devices.experimental.Device
-        ):
+        if expansion_strategy == "device" and isinstance(qnode.device, qml.devices.Device):
             qnode.construct(args, kwargs)
             program, _ = qnode.device.preprocess()
             tapes, _ = program([qnode.tape])
diff --git a/pennylane/interfaces/execution.py b/pennylane/interfaces/execution.py
index d42cf584d09..a5d07e979ae 100644
--- a/pennylane/interfaces/execution.py
+++ b/pennylane/interfaces/execution.py
@@ -40,7 +40,7 @@
 logger = logging.getLogger(__name__)
 logger.addHandler(logging.NullHandler())
 
-device_type = Union[qml.Device, "qml.devices.experimental.Device"]
+device_type = Union[qml.Device, "qml.devices.Device"]
 
 INTERFACE_MAP = {
     None: "Numpy",
@@ -140,16 +140,16 @@ def _get_ml_boundary_execute(interface: str, grad_on_execution: bool) -> Callabl
 def _batch_transform(
     tapes: Sequence[QuantumTape],
     device: device_type,
-    config: "qml.devices.experimental.ExecutionConfig",
+    config: "qml.devices.ExecutionConfig",
     override_shots: Union[bool, int, Sequence[int]] = False,
     device_batch_transform: bool = True,
-) -> Tuple[Sequence[QuantumTape], Callable, "qml.devices.experimental.ExecutionConfig"]:
+) -> Tuple[Sequence[QuantumTape], Callable, "qml.devices.ExecutionConfig"]:
     """Apply the device batch transform unless requested not to.
 
     Args:
         tapes (Tuple[.QuantumTape]): batch of tapes to preprocess
-        device (Device, devices.experimental.Device): the device that defines the required batch transformation
-        config (qml.devices.experimental.ExecutionConfig): the config that characterizes the requested computation
+        device (Device, devices.Device): the device that defines the required batch transformation
+        config (qml.devices.ExecutionConfig): the config that characterizes the requested computation
         override_shots (int): The number of shots to use for the execution. If ``False``, then the
             number of shots on the device is used.
         device_batch_transform (bool): Whether to apply any batch transforms defined by the device
@@ -181,7 +181,7 @@ def _preprocess_expand_fn(
     Args:
         expand_fn (str, Callable): If string, then it must be "device".  Otherwise, it should be a map
             from one tape to a new tape. The final tape must be natively executable by the device.
-        device (Device, devices.experimental.Device): The device that we will be executing on.
+        device (Device, devices.Device): The device that we will be executing on.
         max_expansion (int): The number of times the internal circuit should be expanded when
             executed on a device. Expansion occurs when an operation or measurement is not
             supported, and results in a gate decomposition. If any operations in the decomposition
@@ -193,7 +193,7 @@ def _preprocess_expand_fn(
     """
     if expand_fn != "device":
         return expand_fn
-    if isinstance(device, qml.devices.experimental.Device):
+    if isinstance(device, qml.devices.Device):
 
         def blank_expansion_function(tape):  # pylint: disable=function-redefined
             """A blank expansion function since the new device handles expansion in preprocessing."""
@@ -558,7 +558,7 @@ def cost_fn(params, x):
         _gradient_method = gradient_fn
     else:
         _gradient_method = "gradient-transform"
-    config = qml.devices.experimental.ExecutionConfig(
+    config = qml.devices.ExecutionConfig(
         interface=interface,
         gradient_method=_gradient_method,
         grad_on_execution=None if grad_on_execution == "best" else grad_on_execution,
@@ -598,7 +598,7 @@ def inner_execute_with_empty_jac(tapes, **_):
 
     #### Executing the configured setup #####
 
-    if isinstance(device, qml.devices.experimental.Device):
+    if isinstance(device, qml.devices.Device):
         if not device_batch_transform:
             warnings.warn(
                 "device batch transforms cannot be turned off with the new device interface.",
diff --git a/pennylane/interfaces/jax_jit.py b/pennylane/interfaces/jax_jit.py
index b763374ecef..902128e787a 100644
--- a/pennylane/interfaces/jax_jit.py
+++ b/pennylane/interfaces/jax_jit.py
@@ -76,7 +76,7 @@ def _jac_shape_dtype_struct(tape: "qml.tape.QuantumScript", device: "qml.Device"
         device (Device): the device used to execute the tape.
 
     >>> tape = qml.tape.QuantumScript([qml.RX(1.0, wires=0)], [qml.expval(qml.PauliX(0)), qml.probs(0)])
-    >>> dev = qml.devices.experimental.DefaultQubit2()
+    >>> dev = qml.devices.DefaultQubit()
     >>> _jac_shape_dtype_struct(tape, dev)
     (ShapeDtypeStruct(shape=(), dtype=float64),
     ShapeDtypeStruct(shape=(2,), dtype=float64))
@@ -112,8 +112,8 @@ def _switched_jacobian(tape, device, original_trainable_parameters, jacs):
     Note that this adds an additional nesting dimension to ``jacs``, with one jacobian
     for each original trainable parameter. I am unsure why this works.
 
-    >>> dev = qml.devices.experimental.DefaultQubit2()
-    >>> config = qml.devices.experimental.ExecutionConfig(gradient_method="adjoint")
+    >>> dev = qml.devices.DefaultQubit()
+    >>> config = qml.devices.ExecutionConfig(gradient_method="adjoint")
     >>> tape = qml.tape.QuantumTape([qml.RY(1.0, 0), qml.RX(0.6, 0), qml.RX(0.7, 0)], [qml.expval(qml.PauliZ(0))])
     >>> tape.trainable_params = [0, 2]
     >>> jac = dev.compute_derivatives(tape, config)
diff --git a/pennylane/interfaces/set_shots.py b/pennylane/interfaces/set_shots.py
index 779cb3246c8..652a7af8966 100644
--- a/pennylane/interfaces/set_shots.py
+++ b/pennylane/interfaces/set_shots.py
@@ -43,7 +43,7 @@ def set_shots(device, shots):
     >>> set_shots(dev, shots=100)(lambda: dev.shots)()
     100
     """
-    if isinstance(device, qml.devices.experimental.Device):
+    if isinstance(device, qml.devices.Device):
         raise ValueError(
             "The new device interface is not compatible with `set_shots`. "
             "Set shots when calling the qnode or put the shots on the QuantumTape."
diff --git a/pennylane/qinfo/transforms.py b/pennylane/qinfo/transforms.py
index 20665b98003..1ab0c0a8c6a 100644
--- a/pennylane/qinfo/transforms.py
+++ b/pennylane/qinfo/transforms.py
@@ -18,7 +18,7 @@
 
 import pennylane as qml
 from pennylane.tape import QuantumTape
-from pennylane.devices import DefaultQubit, DefaultMixed
+from pennylane.devices import DefaultQubit, DefaultQubitLegacy, DefaultMixed
 from pennylane.measurements import StateMP, DensityMatrixMP
 from pennylane.transforms import adjoint_metric_tensor, metric_tensor
 from pennylane.transforms.core import transform
@@ -731,9 +731,7 @@ def circ(params):
 
     """
 
-    if qnode.device.shots and isinstance(
-        qnode.device, (DefaultQubit, qml.devices.experimental.DefaultQubit2)
-    ):
+    if qnode.device.shots and isinstance(qnode.device, (DefaultQubitLegacy, DefaultQubit)):
 
         def wrapper(*args0, **kwargs0):
             return 4 * metric_tensor(qnode, *args, **kwargs)(*args0, **kwargs0)
diff --git a/pennylane/qnode.py b/pennylane/qnode.py
index 6793cc2f8f5..cccecc36485 100644
--- a/pennylane/qnode.py
+++ b/pennylane/qnode.py
@@ -385,7 +385,7 @@ def circuit_unpacking(x):
     def __init__(
         self,
         func,
-        device: Union[Device, "qml.devices.experimental.Device"],
+        device: Union[Device, "qml.devices.Device"],
         interface="auto",
         diff_method="best",
         expansion_strategy="gradient",
@@ -420,7 +420,7 @@ def __init__(
                 f"one of {SUPPORTED_INTERFACES}."
             )
 
-        if not isinstance(device, (Device, qml.devices.experimental.Device)):
+        if not isinstance(device, (Device, qml.devices.Device)):
             raise qml.QuantumFunctionError(
                 "Invalid device. Device must be a valid PennyLane device."
             )
@@ -484,7 +484,7 @@ def __init__(
 
     def __repr__(self):
         """String representation."""
-        if isinstance(self.device, qml.devices.experimental.Device):
+        if isinstance(self.device, qml.devices.Device):
             return f"<QNode: device='{self.device}', interface='{self.interface}', diff_method='{self.diff_method}'>"
 
         detail = "<QNode: wires={}, device='{}', interface='{}', diff_method='{}'>"
@@ -701,10 +701,8 @@ def _validate_backprop_method(device, interface, shots=None):
         if shots is not None or _get_device_shots(device):
             raise qml.QuantumFunctionError("Backpropagation is only supported when shots=None.")
 
-        if isinstance(device, qml.devices.experimental.Device):
-            config = qml.devices.experimental.ExecutionConfig(
-                gradient_method="backprop", interface=interface
-            )
+        if isinstance(device, qml.devices.Device):
+            config = qml.devices.ExecutionConfig(gradient_method="backprop", interface=interface)
             if device.supports_derivatives(config):
                 return "backprop", {}, device
             raise qml.QuantumFunctionError(f"Device {device.name} does not support backprop")
@@ -766,8 +764,8 @@ def _validate_adjoint_method(device):
         # need to inspect the circuit measurements to ensure only expectation values are taken. This
         # cannot be done here since we don't yet know the composition of the circuit.
 
-        if isinstance(device, qml.devices.experimental.Device):
-            config = qml.devices.experimental.ExecutionConfig(
+        if isinstance(device, qml.devices.Device):
+            config = qml.devices.ExecutionConfig(
                 gradient_method="adjoint", use_device_gradient=True
             )
             if device.supports_derivatives(config):
@@ -798,7 +796,7 @@ def _validate_device_method(device):
                 return "device", {}, device
             name = device.short_name
         else:
-            config = qml.devices.experimental.ExecutionConfig(gradient_method="device")
+            config = qml.devices.ExecutionConfig(gradient_method="device")
             if device.supports_derivatives(config):
                 return "device", {}, device
             name = device.name
@@ -809,7 +807,7 @@ def _validate_device_method(device):
 
     @staticmethod
     def _validate_parameter_shift(device):
-        if isinstance(device, qml.devices.experimental.Device):
+        if isinstance(device, qml.devices.Device):
             return qml.gradients.param_shift, {}, device
         model = device.capabilities().get("model", None)
 
@@ -900,16 +898,14 @@ def construct(self, args, kwargs):  # pylint: disable=too-many-branches
         # Apply the deferred measurement principle if the device doesn't
         # support mid-circuit measurements natively
         expand_mid_measure = any(isinstance(op, MidMeasureMP) for op in self.tape.operations) and (
-            isinstance(self.device, qml.devices.experimental.Device)
+            isinstance(self.device, qml.devices.Device)
             or not self.device.capabilities().get("supports_mid_measure", False)
         )
         if expand_mid_measure:
             tapes, _ = qml.defer_measurements(self._tape)
             self._tape = tapes[0]
 
-        if self.expansion_strategy == "device" and not isinstance(
-            self.device, qml.devices.experimental.Device
-        ):
+        if self.expansion_strategy == "device" and not isinstance(self.device, qml.devices.Device):
             self._tape = self.device.expand_fn(self.tape, max_expansion=self.max_expansion)
 
         # If the gradient function is a transform, expand the tape so that
@@ -993,7 +989,7 @@ def __call__(self, *args, **kwargs) -> qml.typing.Result:
         if not isinstance(self._qfunc_output, (tuple, qml.measurements.MeasurementProcess)):
             has_partitioned_shots = (
                 self.tape.shots.has_partitioned_shots
-                if isinstance(self.device, qml.devices.experimental.Device)
+                if isinstance(self.device, qml.devices.Device)
                 else self.device._shot_vector
             )
             if has_partitioned_shots:
diff --git a/pennylane/tape/qscript.py b/pennylane/tape/qscript.py
index d4515c752ba..23b3f20e2b4 100644
--- a/pennylane/tape/qscript.py
+++ b/pennylane/tape/qscript.py
@@ -796,7 +796,7 @@ def shape(self, device):
             ((4,), (), (4,))
         """
 
-        if isinstance(device, qml.devices.experimental.Device):
+        if isinstance(device, qml.devices.Device):
             # MP.shape (called below) takes 2 arguments: `device` and `shots`.
             # With the new device interface, shots are stored on tapes rather than the device
             # TODO: refactor MP.shape to accept `wires` instead of device (not currently done
diff --git a/pennylane/transforms/qcut/cutstrategy.py b/pennylane/transforms/qcut/cutstrategy.py
index 553defa079f..511acb0e71f 100644
--- a/pennylane/transforms/qcut/cutstrategy.py
+++ b/pennylane/transforms/qcut/cutstrategy.py
@@ -125,16 +125,16 @@ def __post_init__(
         if devices is None and self.max_free_wires is None:
             raise ValueError("One of arguments `devices` and max_free_wires` must be provided.")
 
-        if isinstance(devices, (qml.Device, qml.devices.experimental.Device)):
+        if isinstance(devices, (qml.Device, qml.devices.Device)):
             devices = (devices,)
 
         if devices is not None:
             if not isinstance(devices, SequenceType) or any(
-                (not isinstance(d, (qml.Device, qml.devices.experimental.Device)) for d in devices)
+                (not isinstance(d, (qml.Device, qml.devices.Device)) for d in devices)
             ):
                 raise ValueError(
                     "Argument `devices` must be a list or tuple containing elements of type "
-                    "`qml.Device` or `qml.devices.experimental.Device`"
+                    "`qml.Device` or `qml.devices.Device`"
                 )
 
             device_wire_sizes = [len(d.wires) for d in devices]
diff --git a/pennylane/transforms/specs.py b/pennylane/transforms/specs.py
index da03b023827..1bda5cbd805 100644
--- a/pennylane/transforms/specs.py
+++ b/pennylane/transforms/specs.py
@@ -125,7 +125,7 @@ def specs_qnode(*args, **kwargs):
 
         info["num_device_wires"] = (
             len(qnode.tape.wires)
-            if isinstance(qnode.device, qml.devices.experimental.Device)
+            if isinstance(qnode.device, qml.devices.Device)
             else len(qnode.device.wires)
         )
         info["device_name"] = getattr(qnode.device, "short_name", qnode.device.name)
diff --git a/setup.py b/setup.py
index a759d389dff..5bd049958f4 100644
--- a/setup.py
+++ b/setup.py
@@ -47,8 +47,8 @@
         # TODO: rename entry point 'pennylane.plugins' to 'pennylane.devices'.
         # This requires a rename in the setup file of all devices, and is best done during another refactor
         "pennylane.plugins": [
-            "default.qubit = pennylane.devices:DefaultQubit",
-            "default.qubit.legacy = pennylane.devices:DefaultQubit",
+            "default.qubit = pennylane.devices:DefaultQubitLegacy",
+            "default.qubit.legacy = pennylane.devices:DefaultQubitLegacy",
             "default.gaussian = pennylane.devices:DefaultGaussian",
             "default.qubit.tf = pennylane.devices.default_qubit_tf:DefaultQubitTF",
             "default.qubit.torch = pennylane.devices.default_qubit_torch:DefaultQubitTorch",
diff --git a/tests/devices/experimental/test_default_qubit_2.py b/tests/devices/experimental/test_default_qubit_2.py
index fa22113f70d..6f07f5c9786 100644
--- a/tests/devices/experimental/test_default_qubit_2.py
+++ b/tests/devices/experimental/test_default_qubit_2.py
@@ -21,38 +21,38 @@
 import pennylane as qml
 from pennylane.measurements import SampleMP, StateMP, ProbabilityMP
 from pennylane.resource import Resources
-from pennylane.devices.experimental import DefaultQubit2, ExecutionConfig
+from pennylane.devices import DefaultQubit, ExecutionConfig
 from pennylane.devices.qubit.preprocess import validate_and_expand_adjoint
 
 
 def test_name():
-    """Tests the name of DefaultQubit2."""
-    assert DefaultQubit2().name == "default.qubit.2"
+    """Tests the name of DefaultQubit."""
+    assert DefaultQubit().name == "default.qubit.2"
 
 
 def test_shots():
-    """Test the shots property of DefaultQubit2."""
-    assert DefaultQubit2().shots == qml.measurements.Shots(None)
-    assert DefaultQubit2(shots=100).shots == qml.measurements.Shots(100)
+    """Test the shots property of DefaultQubit."""
+    assert DefaultQubit().shots == qml.measurements.Shots(None)
+    assert DefaultQubit(shots=100).shots == qml.measurements.Shots(100)
 
     with pytest.raises(AttributeError):
-        DefaultQubit2().shots = 10
+        DefaultQubit().shots = 10
 
 
 def test_wires():
     """Test that a device can be created with wires."""
-    assert DefaultQubit2().wires is None
-    assert DefaultQubit2(wires=2).wires == qml.wires.Wires([0, 1])
-    assert DefaultQubit2(wires=[0, 2]).wires == qml.wires.Wires([0, 2])
+    assert DefaultQubit().wires is None
+    assert DefaultQubit(wires=2).wires == qml.wires.Wires([0, 1])
+    assert DefaultQubit(wires=[0, 2]).wires == qml.wires.Wires([0, 2])
 
     with pytest.raises(AttributeError):
-        DefaultQubit2().wires = [0, 1]
+        DefaultQubit().wires = [0, 1]
 
 
 def test_debugger_attribute():
-    """Test that DefaultQubit2 has a debugger attribute and that it is `None`"""
+    """Test that DefaultQubit has a debugger attribute and that it is `None`"""
     # pylint: disable=protected-access
-    dev = DefaultQubit2()
+    dev = DefaultQubit()
 
     assert hasattr(dev, "_debugger")
     assert dev._debugger is None
@@ -60,8 +60,8 @@ def test_debugger_attribute():
 
 class TestSnapshotMulti:
     def test_snapshot_multiprocessing_execute(self):
-        """DefaultQubit2 cannot execute tapes with Snapshot if `max_workers` is not `None`"""
-        dev = DefaultQubit2(max_workers=2)
+        """DefaultQubit cannot execute tapes with Snapshot if `max_workers` is not `None`"""
+        dev = DefaultQubit(max_workers=2)
 
         tape = qml.tape.QuantumScript(
             [
@@ -80,8 +80,8 @@ def test_snapshot_multiprocessing_execute(self):
             program([tape])
 
     def test_snapshot_multiprocessing_qnode(self):
-        """DefaultQubit2 cannot execute tapes with Snapshot if `max_workers` is not `None`"""
-        dev = DefaultQubit2(max_workers=2)
+        """DefaultQubit cannot execute tapes with Snapshot if `max_workers` is not `None`"""
+        dev = DefaultQubit(max_workers=2)
 
         @qml.qnode(dev)
         def circuit():
@@ -99,15 +99,15 @@ def circuit():
 
 
 class TestTracking:
-    """Testing the tracking capabilities of DefaultQubit2."""
+    """Testing the tracking capabilities of DefaultQubit."""
 
     def test_tracker_set_upon_initialization(self):
         """Test that a new tracker is intialized with each device."""
-        assert DefaultQubit2().tracker is not DefaultQubit2().tracker
+        assert DefaultQubit().tracker is not DefaultQubit().tracker
 
     def test_tracker_not_updated_if_not_active(self):
         """Test that the tracker is not updated if not active."""
-        dev = DefaultQubit2()
+        dev = DefaultQubit()
         assert len(dev.tracker.totals) == 0
 
         dev.execute(qml.tape.QuantumScript())
@@ -115,11 +115,11 @@ def test_tracker_not_updated_if_not_active(self):
         assert len(dev.tracker.history) == 0
 
     def test_tracking_batch(self):
-        """Test that the experimental default qubit integrates with the tracker."""
+        """Test that the new default qubit integrates with the tracker."""
 
         qs = qml.tape.QuantumScript([], [qml.expval(qml.PauliZ(0))])
 
-        dev = DefaultQubit2()
+        dev = DefaultQubit()
         config = ExecutionConfig(gradient_method="adjoint")
         with qml.Tracker(dev) as tracker:
             dev.execute(qs)
@@ -143,10 +143,10 @@ def test_tracking_batch(self):
 
     def test_tracking_execute_and_derivatives(self):
         """Test that the execute_and_compute_* calls are being tracked for the
-        experimental default qubit device"""
+        new default qubit device"""
 
         qs = qml.tape.QuantumScript([], [qml.expval(qml.PauliZ(0))])
-        dev = DefaultQubit2()
+        dev = DefaultQubit()
         config = ExecutionConfig(gradient_method="adjoint")
 
         with qml.Tracker(dev) as tracker:
@@ -172,7 +172,7 @@ def test_tracking_execute_and_derivatives(self):
         }
 
     def test_tracking_resources(self):
-        """Test that resources are tracked for the experimental default qubit device."""
+        """Test that resources are tracked for the new default qubit device."""
         qs = qml.tape.QuantumScript(
             [
                 qml.Hadamard(0),
@@ -193,7 +193,7 @@ def test_tracking_resources(self):
             depth=3,
         )
 
-        dev = DefaultQubit2()
+        dev = DefaultQubit()
         with qml.Tracker(dev) as tracker:
             dev.execute(qs)
 
@@ -207,7 +207,7 @@ class TestPreprocessing:
 
     def test_chooses_best_gradient_method(self):
         """Test that preprocessing chooses backprop as the best gradient method."""
-        dev = DefaultQubit2()
+        dev = DefaultQubit()
 
         config = ExecutionConfig(
             gradient_method="best", use_device_gradient=None, grad_on_execution=None
@@ -221,7 +221,7 @@ def test_chooses_best_gradient_method(self):
 
     def test_config_choices_for_adjoint(self):
         """Test that preprocessing request grad on execution and says to use the device gradient if adjoint is requested."""
-        dev = DefaultQubit2()
+        dev = DefaultQubit()
 
         config = ExecutionConfig(
             gradient_method="adjoint", use_device_gradient=None, grad_on_execution=None
@@ -235,7 +235,7 @@ def test_config_choices_for_adjoint(self):
     @pytest.mark.parametrize("max_workers", [None, 1, 2, 3])
     def test_config_choices_for_threading(self, max_workers):
         """Test that preprocessing request grad on execution and says to use the device gradient if adjoint is requested."""
-        dev = DefaultQubit2()
+        dev = DefaultQubit()
 
         config = ExecutionConfig(device_options={"max_workers": max_workers})
         _, new_config = dev.preprocess(config)
@@ -244,7 +244,7 @@ def test_config_choices_for_threading(self, max_workers):
 
     def test_circuit_wire_validation(self):
         """Test that preprocessing validates wires on the circuits being executed."""
-        dev = DefaultQubit2(wires=3)
+        dev = DefaultQubit(wires=3)
         circuit_valid_0 = qml.tape.QuantumScript([qml.PauliX(0)])
         program, _ = dev.preprocess()
         circuits, _ = program([circuit_valid_0])
@@ -278,7 +278,7 @@ def test_circuit_wire_validation(self):
     )
     def test_measurement_is_swapped_out(self, mp_fn, mp_cls, shots):
         """Test that preprocessing swaps out any MP with no wires or obs"""
-        dev = DefaultQubit2(wires=3)
+        dev = DefaultQubit(wires=3)
         original_mp = mp_fn()
         exp_z = qml.expval(qml.PauliZ(0))
         qs = qml.tape.QuantumScript([qml.Hadamard(0)], [original_mp, exp_z], shots=shots)
@@ -293,11 +293,11 @@ def test_measurement_is_swapped_out(self, mp_fn, mp_cls, shots):
 
 
 class TestSupportsDerivatives:
-    """Test that DefaultQubit2 states what kind of derivatives it supports."""
+    """Test that DefaultQubit states what kind of derivatives it supports."""
 
     def test_supports_backprop(self):
-        """Test that DefaultQubit2 says that it supports backpropagation."""
-        dev = DefaultQubit2()
+        """Test that DefaultQubit says that it supports backpropagation."""
+        dev = DefaultQubit()
         assert dev.supports_derivatives() is True
         assert dev.supports_jvp() is True
         assert dev.supports_vjp() is True
@@ -318,8 +318,8 @@ def test_supports_backprop(self):
         assert dev.supports_vjp(config) is False
 
     def test_supports_adjoint(self):
-        """Test that DefaultQubit2 says that it supports adjoint differentiation."""
-        dev = DefaultQubit2()
+        """Test that DefaultQubit says that it supports adjoint differentiation."""
+        dev = DefaultQubit()
         config = ExecutionConfig(gradient_method="adjoint", use_device_gradient=True)
         assert dev.supports_derivatives(config) is True
         assert dev.supports_jvp(config) is True
@@ -340,8 +340,8 @@ def test_supports_adjoint(self):
         assert dev.supports_vjp(config, qs) is False
 
     def test_doesnt_support_adjoint_with_invalid_tape(self):
-        """Tests that DefaultQubit2 does not support adjoint differentiation with invalid circuits."""
-        dev = DefaultQubit2()
+        """Tests that DefaultQubit does not support adjoint differentiation with invalid circuits."""
+        dev = DefaultQubit()
         config = ExecutionConfig(gradient_method="adjoint")
         circuit = qml.tape.QuantumScript([], [qml.probs()])
         assert dev.supports_derivatives(config, circuit=circuit) is False
@@ -350,8 +350,8 @@ def test_doesnt_support_adjoint_with_invalid_tape(self):
 
     @pytest.mark.parametrize("gradient_method", ["parameter-shift", "finite-diff", "device"])
     def test_doesnt_support_other_gradient_methods(self, gradient_method):
-        """Test that DefaultQubit2 currently does not support other gradient methods natively."""
-        dev = DefaultQubit2()
+        """Test that DefaultQubit currently does not support other gradient methods natively."""
+        dev = DefaultQubit()
         config = ExecutionConfig(gradient_method=gradient_method)
         assert dev.supports_derivatives(config) is False
         assert dev.supports_jvp(config) is False
@@ -369,7 +369,7 @@ def test_basic_circuit_numpy(self, max_workers):
             [qml.RX(phi, wires=0)], [qml.expval(qml.PauliY(0)), qml.expval(qml.PauliZ(0))]
         )
 
-        dev = DefaultQubit2(max_workers=max_workers)
+        dev = DefaultQubit(max_workers=max_workers)
         result = dev.execute(qs)
 
         assert isinstance(result, tuple)
@@ -386,7 +386,7 @@ def test_basic_circuit_numpy_with_config(self, max_workers):
             [qml.RX(phi, wires=0)], [qml.expval(qml.PauliY(0)), qml.expval(qml.PauliZ(0))]
         )
 
-        dev = DefaultQubit2(max_workers=max_workers)
+        dev = DefaultQubit(max_workers=max_workers)
         config = ExecutionConfig(
             device_options={"max_workers": dev._max_workers}  # pylint: disable=protected-access
         )
@@ -404,7 +404,7 @@ def test_autograd_results_and_backprop(self, max_workers):
         """Tests execution and gradients with autograd"""
         phi = qml.numpy.array(-0.52)
 
-        dev = DefaultQubit2(max_workers=max_workers)
+        dev = DefaultQubit(max_workers=max_workers)
 
         def f(x):
             qs = qml.tape.QuantumScript(
@@ -432,7 +432,7 @@ def test_jax_results_and_backprop(self, use_jit, max_workers):
 
         phi = jax.numpy.array(0.678)
 
-        dev = DefaultQubit2(max_workers=max_workers)
+        dev = DefaultQubit(max_workers=max_workers)
 
         def f(x):
             qs = qml.tape.QuantumScript(
@@ -465,7 +465,7 @@ def test_torch_results_and_backprop(self, max_workers):
 
         phi = torch.tensor(-0.526, requires_grad=True)
 
-        dev = DefaultQubit2(max_workers=max_workers)
+        dev = DefaultQubit(max_workers=max_workers)
 
         def f(x):
             qs = qml.tape.QuantumScript(
@@ -493,7 +493,7 @@ def test_tf_results_and_backprop(self, max_workers):
 
         phi = tf.Variable(4.873)
 
-        dev = DefaultQubit2(max_workers=max_workers)
+        dev = DefaultQubit(max_workers=max_workers)
 
         with tf.GradientTape(persistent=True) as grad_tape:
             qs = qml.tape.QuantumScript(
@@ -517,7 +517,7 @@ def test_tf_results_and_backprop(self, max_workers):
     @pytest.mark.parametrize("op,param", [(qml.RX, np.pi), (qml.BasisState, [1])])
     def test_qnode_returns_correct_interface(self, op, param):
         """Test that even if no interface parameters are given, result is correct."""
-        dev = DefaultQubit2()
+        dev = DefaultQubit()
 
         @qml.qnode(dev, interface="tf")
         def circuit(p):
@@ -538,7 +538,7 @@ def test_single_expval(self, max_workers):
         x = np.array(0.732)
         qs = qml.tape.QuantumScript([qml.RY(x, wires=0)], [qml.expval(qml.PauliZ(0))], shots=10000)
 
-        dev = DefaultQubit2(max_workers=max_workers)
+        dev = DefaultQubit(max_workers=max_workers)
         result = dev.execute(qs)
 
         assert isinstance(result, (float, np.ndarray))
@@ -551,7 +551,7 @@ def test_single_probs(self, max_workers):
         x = np.array(0.732)
         qs = qml.tape.QuantumScript([qml.RY(x, wires=0)], [qml.probs(wires=0)], shots=10000)
 
-        dev = DefaultQubit2(max_workers=max_workers)
+        dev = DefaultQubit(max_workers=max_workers)
         result = dev.execute(qs)
 
         assert isinstance(result, (float, np.ndarray))
@@ -564,7 +564,7 @@ def test_single_sample(self, max_workers):
         x = np.array(0.732)
         qs = qml.tape.QuantumScript([qml.RY(x, wires=0)], [qml.sample(wires=range(2))], shots=10000)
 
-        dev = DefaultQubit2(max_workers=max_workers)
+        dev = DefaultQubit(max_workers=max_workers)
         result = dev.execute(qs)
 
         assert isinstance(result, (float, np.ndarray))
@@ -583,7 +583,7 @@ def test_multi_measurements(self, max_workers):
             shots=10000,
         )
 
-        dev = DefaultQubit2(max_workers=max_workers)
+        dev = DefaultQubit(max_workers=max_workers)
         result = dev.execute(qs)
 
         assert isinstance(result, tuple)
@@ -624,7 +624,7 @@ def test_expval_shot_vector(self, max_workers, shots):
         shots = qml.measurements.Shots(shots)
         qs = qml.tape.QuantumScript([qml.RY(x, wires=0)], [qml.expval(qml.PauliZ(0))], shots=shots)
 
-        dev = DefaultQubit2(max_workers=max_workers)
+        dev = DefaultQubit(max_workers=max_workers)
         result = dev.execute(qs)
 
         assert isinstance(result, tuple)
@@ -642,7 +642,7 @@ def test_probs_shot_vector(self, max_workers, shots):
         shots = qml.measurements.Shots(shots)
         qs = qml.tape.QuantumScript([qml.RY(x, wires=0)], [qml.probs(wires=0)], shots=shots)
 
-        dev = DefaultQubit2(max_workers=max_workers)
+        dev = DefaultQubit(max_workers=max_workers)
         result = dev.execute(qs)
 
         assert isinstance(result, tuple)
@@ -662,7 +662,7 @@ def test_sample_shot_vector(self, max_workers, shots):
         shots = qml.measurements.Shots(shots)
         qs = qml.tape.QuantumScript([qml.RY(x, wires=0)], [qml.sample(wires=range(2))], shots=shots)
 
-        dev = DefaultQubit2(max_workers=max_workers)
+        dev = DefaultQubit(max_workers=max_workers)
         result = dev.execute(qs)
 
         assert isinstance(result, tuple)
@@ -689,7 +689,7 @@ def test_multi_measurement_shot_vector(self, max_workers, shots):
             shots=shots,
         )
 
-        dev = DefaultQubit2(max_workers=max_workers)
+        dev = DefaultQubit(max_workers=max_workers)
         result = dev.execute(qs)
 
         assert isinstance(result, tuple)
@@ -732,7 +732,7 @@ def test_custom_wire_labels(self, max_workers):
             shots=10000,
         )
 
-        dev = DefaultQubit2(max_workers=max_workers)
+        dev = DefaultQubit(max_workers=max_workers)
         result = dev.execute(qs)
 
         assert isinstance(result, tuple)
@@ -764,7 +764,7 @@ def test_batch_tapes(self, max_workers):
         qs1 = qml.tape.QuantumScript([qml.RX(x, wires=0)], [qml.sample(wires=(0, 1))], shots=100)
         qs2 = qml.tape.QuantumScript([qml.RX(x, wires=0)], [qml.sample(wires=1)], shots=50)
 
-        dev = DefaultQubit2(max_workers=max_workers)
+        dev = DefaultQubit(max_workers=max_workers)
         results = dev.execute((qs1, qs2))
 
         assert isinstance(results, tuple)
@@ -779,7 +779,7 @@ def test_counts_wires(self, max_workers):
         x = np.array(np.pi / 2)
         qs = qml.tape.QuantumScript([qml.RY(x, wires=0)], [qml.counts(wires=[0, 1])], shots=10000)
 
-        dev = DefaultQubit2(seed=123, max_workers=max_workers)
+        dev = DefaultQubit(seed=123, max_workers=max_workers)
         result = dev.execute(qs)
 
         assert isinstance(result, dict)
@@ -800,7 +800,7 @@ def test_counts_obs(self, all_outcomes, max_workers):
             shots=10000,
         )
 
-        dev = DefaultQubit2(seed=123, max_workers=max_workers)
+        dev = DefaultQubit(seed=123, max_workers=max_workers)
         result = dev.execute(qs)
 
         assert isinstance(result, dict)
@@ -816,7 +816,7 @@ class TestExecutingBatches:
 
     @staticmethod
     def f(dev, phi):
-        """A function that executes a batch of scripts on DefaultQubit2 without preprocessing."""
+        """A function that executes a batch of scripts on DefaultQubit without preprocessing."""
         ops = [
             qml.PauliX("a"),
             qml.PauliX("b"),
@@ -837,7 +837,7 @@ def f(dev, phi):
 
     @staticmethod
     def f_hashable(phi):
-        """A function that executes a batch of scripts on DefaultQubit2 without preprocessing."""
+        """A function that executes a batch of scripts on DefaultQubit without preprocessing."""
         ops = [
             qml.PauliX("a"),
             qml.PauliX("b"),
@@ -854,7 +854,7 @@ def f_hashable(phi):
 
         ops = [qml.Hadamard(0), qml.IsingXX(phi, wires=(0, 1))]
         qs2 = qml.tape.QuantumScript(ops, [qml.probs(wires=(0, 1))])
-        return DefaultQubit2().execute((qs1, qs2))
+        return DefaultQubit().execute((qs1, qs2))
 
     @staticmethod
     def expected(phi):
@@ -878,7 +878,7 @@ def nested_compare(x1, x2):
     @pytest.mark.parametrize("max_workers", [None, 1, 2])
     def test_numpy(self, max_workers):
         """Test that results are expected when the parameter does not have a parameter."""
-        dev = DefaultQubit2(max_workers=max_workers)
+        dev = DefaultQubit(max_workers=max_workers)
 
         phi = 0.892
         results = self.f(dev, phi)
@@ -890,7 +890,7 @@ def test_numpy(self, max_workers):
     @pytest.mark.parametrize("max_workers", [None, 1, 2])
     def test_autograd(self, max_workers):
         """Test batches can be executed and have backprop derivatives in autograd."""
-        dev = DefaultQubit2(max_workers=max_workers)
+        dev = DefaultQubit(max_workers=max_workers)
 
         phi = qml.numpy.array(-0.629)
         results = self.f(dev, phi)
@@ -934,7 +934,7 @@ def test_torch(self, max_workers):
         """Test batches can be executed and have backprop derivatives in torch."""
         import torch
 
-        dev = DefaultQubit2(max_workers=max_workers)
+        dev = DefaultQubit(max_workers=max_workers)
 
         x = torch.tensor(9.6243)
 
@@ -962,7 +962,7 @@ def test_tf(self, max_workers):
 
         import tensorflow as tf
 
-        dev = DefaultQubit2(max_workers=max_workers)
+        dev = DefaultQubit(max_workers=max_workers)
 
         x = tf.Variable(5.2281)
         with tf.GradientTape(persistent=True) as tape:
@@ -1012,7 +1012,7 @@ def f_hashable(scale, n_wires=10, offset=0.1, convert_to_hamiltonian=False):
         if convert_to_hamiltonian:
             H = H._pauli_rep.hamiltonian()  # pylint: disable=protected-access
         qs = qml.tape.QuantumScript(ops, [qml.expval(H)])
-        return DefaultQubit2().execute(qs)
+        return DefaultQubit().execute(qs)
 
     @staticmethod
     def expected(scale, n_wires=10, offset=0.1, like="numpy"):
@@ -1026,7 +1026,7 @@ def expected(scale, n_wires=10, offset=0.1, like="numpy"):
     @pytest.mark.parametrize("convert_to_hamiltonian", (True, False))
     def test_autograd_backprop(self, convert_to_hamiltonian):
         """Test that backpropagation derivatives work in autograd with hamiltonians and large sums."""
-        dev = DefaultQubit2()
+        dev = DefaultQubit()
         x = qml.numpy.array(0.52)
         out = self.f(dev, x, convert_to_hamiltonian=convert_to_hamiltonian)
         expected_out = self.expected(x)
@@ -1062,7 +1062,7 @@ def test_torch_backprop(self, convert_to_hamiltonian):
         """Test that backpropagation derivatives work with torch with hamiltonians and large sums."""
         import torch
 
-        dev = DefaultQubit2()
+        dev = DefaultQubit()
 
         x = torch.tensor(-0.289, requires_grad=True)
         x2 = torch.tensor(-0.289, requires_grad=True)
@@ -1080,7 +1080,7 @@ def test_tf_backprop(self, convert_to_hamiltonian):
         """Test that backpropagation derivatives work with tensorflow with hamiltonians and large sums."""
         import tensorflow as tf
 
-        dev = DefaultQubit2()
+        dev = DefaultQubit()
 
         x = tf.Variable(0.5)
 
@@ -1098,13 +1098,13 @@ def test_tf_backprop(self, convert_to_hamiltonian):
 
 @pytest.mark.parametrize("max_workers", [None, 1, 2])
 class TestAdjointDifferentiation:
-    """Tests adjoint differentiation integration with DefaultQubit2."""
+    """Tests adjoint differentiation integration with DefaultQubit."""
 
     ec = ExecutionConfig(gradient_method="adjoint")
 
     def test_derivatives_single_circuit(self, max_workers):
         """Tests derivatives with a single circuit."""
-        dev = DefaultQubit2(max_workers=max_workers)
+        dev = DefaultQubit(max_workers=max_workers)
         x = np.array(np.pi / 7)
         qs = qml.tape.QuantumScript([qml.RX(x, 0)], [qml.expval(qml.PauliZ(0))])
         qs, _ = validate_and_expand_adjoint(qs)
@@ -1122,7 +1122,7 @@ def test_derivatives_single_circuit(self, max_workers):
 
     def test_derivatives_list_with_single_circuit(self, max_workers):
         """Tests a basic example with a batch containing a single circuit."""
-        dev = DefaultQubit2(max_workers=max_workers)
+        dev = DefaultQubit(max_workers=max_workers)
         x = np.array(np.pi / 7)
         qs = qml.tape.QuantumScript([qml.RX(x, 0)], [qml.expval(qml.PauliZ(0))])
         qs, _ = validate_and_expand_adjoint(qs)
@@ -1140,7 +1140,7 @@ def test_derivatives_list_with_single_circuit(self, max_workers):
 
     def test_derivatives_many_tapes_many_results(self, max_workers):
         """Tests a basic example with a batch of circuits of varying return shapes."""
-        dev = DefaultQubit2(max_workers=max_workers)
+        dev = DefaultQubit(max_workers=max_workers)
         x = np.array(np.pi / 7)
         single_meas = qml.tape.QuantumScript([qml.RX(x, 0)], [qml.expval(qml.PauliZ(0))])
         multi_meas = qml.tape.QuantumScript(
@@ -1154,7 +1154,7 @@ def test_derivatives_many_tapes_many_results(self, max_workers):
 
     def test_derivatives_integration(self, max_workers):
         """Tests the expected workflow done by a calling method."""
-        dev = DefaultQubit2(max_workers=max_workers)
+        dev = DefaultQubit(max_workers=max_workers)
         x = np.array(np.pi / 7)
         expected_grad = (-qml.math.sin(x), (qml.math.cos(x), -qml.math.sin(x)))
         single_meas = qml.tape.QuantumScript([qml.RX(x, 0)], [qml.expval(qml.PauliZ(0))])
@@ -1175,7 +1175,7 @@ def test_derivatives_integration(self, max_workers):
 
     def test_jvps_single_circuit(self, max_workers):
         """Tests jvps with a single circuit."""
-        dev = DefaultQubit2(max_workers=max_workers)
+        dev = DefaultQubit(max_workers=max_workers)
         x = np.array(np.pi / 7)
         tangent = (0.456,)
 
@@ -1197,7 +1197,7 @@ def test_jvps_single_circuit(self, max_workers):
 
     def test_jvps_list_with_single_circuit(self, max_workers):
         """Tests a basic example with a batch containing a single circuit."""
-        dev = DefaultQubit2(max_workers=max_workers)
+        dev = DefaultQubit(max_workers=max_workers)
         x = np.array(np.pi / 7)
         tangent = (0.456,)
 
@@ -1219,7 +1219,7 @@ def test_jvps_list_with_single_circuit(self, max_workers):
 
     def test_jvps_many_tapes_many_results(self, max_workers):
         """Tests a basic example with a batch of circuits of varying return shapes."""
-        dev = DefaultQubit2(max_workers=max_workers)
+        dev = DefaultQubit(max_workers=max_workers)
         x = np.array(np.pi / 7)
         single_meas = qml.tape.QuantumScript([qml.RX(x, 0)], [qml.expval(qml.PauliZ(0))])
         multi_meas = qml.tape.QuantumScript(
@@ -1247,7 +1247,7 @@ def test_jvps_many_tapes_many_results(self, max_workers):
 
     def test_jvps_integration(self, max_workers):
         """Tests the expected workflow done by a calling method."""
-        dev = DefaultQubit2(max_workers=max_workers)
+        dev = DefaultQubit(max_workers=max_workers)
         x = np.array(np.pi / 7)
 
         single_meas = qml.tape.QuantumScript([qml.RX(x, 0)], [qml.expval(qml.PauliZ(0))])
@@ -1273,7 +1273,7 @@ def test_jvps_integration(self, max_workers):
 
     def test_vjps_single_circuit(self, max_workers):
         """Tests vjps with a single circuit."""
-        dev = DefaultQubit2(max_workers=max_workers)
+        dev = DefaultQubit(max_workers=max_workers)
         x = np.array(np.pi / 7)
         cotangent = (0.456,)
 
@@ -1294,7 +1294,7 @@ def test_vjps_single_circuit(self, max_workers):
 
     def test_vjps_list_with_single_circuit(self, max_workers):
         """Tests a basic example with a batch containing a single circuit."""
-        dev = DefaultQubit2(max_workers=max_workers)
+        dev = DefaultQubit(max_workers=max_workers)
         x = np.array(np.pi / 7)
         cotangent = (0.456,)
 
@@ -1315,7 +1315,7 @@ def test_vjps_list_with_single_circuit(self, max_workers):
 
     def test_vjps_many_tapes_many_results(self, max_workers):
         """Tests a basic example with a batch of circuits of varying return shapes."""
-        dev = DefaultQubit2(max_workers=max_workers)
+        dev = DefaultQubit(max_workers=max_workers)
         x = np.array(np.pi / 7)
         single_meas = qml.tape.QuantumScript([qml.RX(x, 0)], [qml.expval(qml.PauliZ(0))])
         multi_meas = qml.tape.QuantumScript(
@@ -1342,7 +1342,7 @@ def test_vjps_many_tapes_many_results(self, max_workers):
 
     def test_vjps_integration(self, max_workers):
         """Tests the expected workflow done by a calling method."""
-        dev = DefaultQubit2(max_workers=max_workers)
+        dev = DefaultQubit(max_workers=max_workers)
         x = np.array(np.pi / 7)
 
         single_meas = qml.tape.QuantumScript([qml.RX(x, 0)], [qml.expval(qml.PauliZ(0))])
@@ -1395,7 +1395,7 @@ def decomposition(self):
             [qml.expval(qml.PauliY("a")), qml.expval(qml.PauliZ("a")), qml.expval(qml.PauliX("b"))],
         )
 
-        dev = DefaultQubit2(max_workers=max_workers)
+        dev = DefaultQubit(max_workers=max_workers)
         tapes = tuple([qscript])
         program, config = dev.preprocess()
         tapes, pre_processing_fn = program(tapes)
@@ -1451,7 +1451,7 @@ def decomposition(self):
 
         initial_batch = [qs1, qs2]
 
-        dev = DefaultQubit2(max_workers=max_workers)
+        dev = DefaultQubit(max_workers=max_workers)
 
         program, config = dev.preprocess()
         batch, pre_processing_fn = program(initial_batch)
@@ -1479,7 +1479,7 @@ def decomposition(self):
     def test_preprocess_defer_measurements(self, mocker):
         """Test that a QNode with mid-circuit measurements is transformed
         using defer_measurements."""
-        dev = DefaultQubit2()
+        dev = DefaultQubit()
 
         tape = qml.tape.QuantumScript(
             [qml.PauliX(0), qml.measurements.MidMeasureMP(qml.wires.Wires(0))],
@@ -1510,10 +1510,10 @@ def test_same_seed(self, measurements, max_workers):
         the same results"""
         qs = qml.tape.QuantumScript([qml.Hadamard(0)], measurements, shots=1000)
 
-        dev1 = DefaultQubit2(seed=123, max_workers=max_workers)
+        dev1 = DefaultQubit(seed=123, max_workers=max_workers)
         result1 = dev1.execute(qs)
 
-        dev2 = DefaultQubit2(seed=123, max_workers=max_workers)
+        dev2 = DefaultQubit(seed=123, max_workers=max_workers)
         result2 = dev2.execute(qs)
 
         if len(measurements) == 1:
@@ -1527,13 +1527,13 @@ def test_different_seed(self, max_workers):
         different results (with almost certainty)"""
         qs = qml.tape.QuantumScript([qml.Hadamard(0)], [qml.sample(wires=0)], shots=1000)
 
-        dev1 = DefaultQubit2(seed=None, max_workers=max_workers)
+        dev1 = DefaultQubit(seed=None, max_workers=max_workers)
         result1 = dev1.execute(qs)
 
-        dev2 = DefaultQubit2(seed=123, max_workers=max_workers)
+        dev2 = DefaultQubit(seed=123, max_workers=max_workers)
         result2 = dev2.execute(qs)
 
-        dev3 = DefaultQubit2(seed=456, max_workers=max_workers)
+        dev3 = DefaultQubit(seed=456, max_workers=max_workers)
         result3 = dev3.execute(qs)
 
         # assert results are pairwise different
@@ -1555,7 +1555,7 @@ def test_different_executions(self, measurements, max_workers):
         """Test that the same device will produce different results every execution"""
         qs = qml.tape.QuantumScript([qml.Hadamard(0)], measurements, shots=1000)
 
-        dev = DefaultQubit2(seed=123, max_workers=max_workers)
+        dev = DefaultQubit(seed=123, max_workers=max_workers)
         result1 = dev.execute(qs)
         result2 = dev.execute(qs)
 
@@ -1580,13 +1580,13 @@ def test_global_seed_and_device_seed(self, measurements, max_workers):
         qs = qml.tape.QuantumScript([qml.Hadamard(0)], measurements, shots=1000)
 
         # expected result
-        dev1 = DefaultQubit2(seed=123, max_workers=max_workers)
+        dev1 = DefaultQubit(seed=123, max_workers=max_workers)
         result1 = dev1.execute(qs)
 
         # set a global seed both before initialization of the
         # device and before execution of the tape
         np.random.seed(456)
-        dev2 = DefaultQubit2(seed=123, max_workers=max_workers)
+        dev2 = DefaultQubit(seed=123, max_workers=max_workers)
         np.random.seed(789)
         result2 = dev2.execute(qs)
 
@@ -1598,17 +1598,17 @@ def test_global_seed_and_device_seed(self, measurements, max_workers):
     def test_global_seed_no_device_seed_by_default(self):
         """Test that the global numpy seed initializes the rng if device seed is none."""
         np.random.seed(42)
-        dev = DefaultQubit2()
+        dev = DefaultQubit()
         first_num = dev._rng.random()  # pylint: disable=protected-access
 
         np.random.seed(42)
-        dev2 = DefaultQubit2()
+        dev2 = DefaultQubit()
         second_num = dev2._rng.random()  # pylint: disable=protected-access
 
         assert qml.math.allclose(first_num, second_num)
 
         np.random.seed(42)
-        dev2 = DefaultQubit2(seed="global")
+        dev2 = DefaultQubit(seed="global")
         third_num = dev2._rng.random()  # pylint: disable=protected-access
 
         assert qml.math.allclose(third_num, first_num)
@@ -1616,11 +1616,11 @@ def test_global_seed_no_device_seed_by_default(self):
     def test_None_seed_not_using_global_rng(self):
         """Test that if the seed is None, it is uncorrelated with the global rng."""
         np.random.seed(42)
-        dev = DefaultQubit2(seed=None)
+        dev = DefaultQubit(seed=None)
         first_nums = dev._rng.random(10)  # pylint: disable=protected-access
 
         np.random.seed(42)
-        dev2 = DefaultQubit2(seed=None)
+        dev2 = DefaultQubit(seed=None)
         second_nums = dev2._rng.random(10)  # pylint: disable=protected-access
 
         assert not qml.math.allclose(first_nums, second_nums)
@@ -1631,7 +1631,7 @@ def test_rng_as_seed(self):
         first_num = rng1.random()
 
         rng = np.random.default_rng(42)
-        dev = DefaultQubit2(seed=rng)
+        dev = DefaultQubit(seed=rng)
         second_num = dev._rng.random()  # pylint: disable=protected-access
 
         assert qml.math.allclose(first_num, second_num)
@@ -1650,7 +1650,7 @@ def test_hamiltonian_expval(self, max_workers):
         ops = [qml.RY(x, wires=0), qml.RZ(y, wires=0)]
         meas = [qml.expval(qml.Hamiltonian([0.8, 0.5], [qml.PauliZ(0), qml.PauliX(0)]))]
 
-        dev = DefaultQubit2(seed=100, max_workers=max_workers)
+        dev = DefaultQubit(seed=100, max_workers=max_workers)
         qs = qml.tape.QuantumScript(ops, meas, shots=10000)
         res = dev.execute(qs)
 
@@ -1664,7 +1664,7 @@ def test_sum_expval(self, max_workers):
         ops = [qml.RY(x, wires=0), qml.RZ(y, wires=0)]
         meas = [qml.expval(qml.s_prod(0.8, qml.PauliZ(0)) + qml.s_prod(0.5, qml.PauliX(0)))]
 
-        dev = DefaultQubit2(seed=100, max_workers=max_workers)
+        dev = DefaultQubit(seed=100, max_workers=max_workers)
         qs = qml.tape.QuantumScript(ops, meas, shots=10000)
         res = dev.execute(qs)
 
@@ -1684,7 +1684,7 @@ def test_multi_wires(self, max_workers):
         t2 = 6.2 * qml.prod(*(qml.PauliY(i) for i in range(n_wires)))
         H = t1 + t2
 
-        dev = DefaultQubit2(seed=100, max_workers=max_workers)
+        dev = DefaultQubit(seed=100, max_workers=max_workers)
         qs = qml.tape.QuantumScript(ops, [qml.expval(H)], shots=100000)
         res = dev.execute(qs)
 
@@ -1753,7 +1753,7 @@ def test_complex_hamiltonian(self, max_workers):
             ],
         )
 
-        dev = DefaultQubit2(seed=100, max_workers=max_workers)
+        dev = DefaultQubit(seed=100, max_workers=max_workers)
         qs = qml.tape.QuantumScript(ops, [qml.expval(H)], shots=100000)
         res = dev.execute(qs)
 
@@ -1770,7 +1770,7 @@ class TestClassicalShadows:
     @pytest.mark.parametrize("max_workers", [None, 1, 2])
     def test_shape_and_dtype(self, max_workers, n_qubits):
         """Test that the shape and dtype of the measurement is correct"""
-        dev = DefaultQubit2(max_workers=max_workers)
+        dev = DefaultQubit(max_workers=max_workers)
 
         ops = [qml.Hadamard(i) for i in range(n_qubits)]
         qs = qml.tape.QuantumScript(ops, [qml.classical_shadow(range(n_qubits))], shots=100)
@@ -1788,7 +1788,7 @@ def test_shape_and_dtype(self, max_workers, n_qubits):
     @pytest.mark.parametrize("max_workers", [None, 1, 2])
     def test_expval(self, max_workers):
         """Test that shadow expval measurements work as expected"""
-        dev = DefaultQubit2(seed=100, max_workers=max_workers)
+        dev = DefaultQubit(seed=100, max_workers=max_workers)
 
         ops = [qml.Hadamard(0), qml.Hadamard(1)]
         meas = [qml.shadow_expval(qml.PauliX(0) @ qml.PauliX(1), seed=200)]
@@ -1802,7 +1802,7 @@ def test_expval(self, max_workers):
     @pytest.mark.parametrize("max_workers", [None, 1, 2])
     def test_multiple_shadow_measurements(self, n_qubits, max_workers):
         """Test that multiple classical shadow measurements work as expected"""
-        dev = DefaultQubit2(max_workers=max_workers)
+        dev = DefaultQubit(max_workers=max_workers)
 
         ops = [qml.Hadamard(i) for i in range(n_qubits)]
         mps = [qml.classical_shadow(range(n_qubits)), qml.classical_shadow(range(n_qubits))]
@@ -1828,7 +1828,7 @@ def test_multiple_shadow_measurements(self, n_qubits, max_workers):
     @pytest.mark.parametrize("max_workers", [None, 1, 2])
     def test_reconstruct_bell_state(self, max_workers):
         """Test that a bell state can be faithfully reconstructed"""
-        dev = DefaultQubit2(seed=100, max_workers=max_workers)
+        dev = DefaultQubit(seed=100, max_workers=max_workers)
 
         ops = [qml.Hadamard(0), qml.CNOT([0, 1])]
         meas = [qml.classical_shadow(wires=[0, 1], seed=200)]
@@ -1875,7 +1875,7 @@ def test_reconstruct_bell_state(self, max_workers):
     @pytest.mark.parametrize("max_workers", [None, 1, 2])
     def test_shot_vectors(self, max_workers, n_qubits, shots):
         """Test that classical shadows works when given a shot vector"""
-        dev = DefaultQubit2(max_workers=max_workers)
+        dev = DefaultQubit(max_workers=max_workers)
         shots = qml.measurements.Shots(shots)
 
         ops = [qml.Hadamard(i) for i in range(n_qubits)]
@@ -1904,7 +1904,7 @@ class TestDynamicType:
     def test_projector(self, max_workers, n_wires):
         """Test that qml.Projector yields the expected results for both of its subclasses."""
         wires = list(range(n_wires))
-        dev = DefaultQubit2(max_workers=max_workers)
+        dev = DefaultQubit(max_workers=max_workers)
         ops = [qml.adjoint(qml.Hadamard(q)) for q in wires]
         basis_state = np.zeros((n_wires,))
         state_vector = np.zeros((2**n_wires,))
@@ -1922,7 +1922,7 @@ class TestIntegration:
     @pytest.mark.parametrize("wires,expected", [(None, [1, 0]), (3, [0, 0, 1])])
     def test_sample_uses_device_wires(self, wires, expected):
         """Test that if device wires are given, then they are used by sample."""
-        dev = DefaultQubit2(wires=wires, shots=5)
+        dev = DefaultQubit(wires=wires, shots=5)
 
         @qml.qnode(dev)
         def circuit():
@@ -1941,7 +1941,7 @@ def circuit():
     )
     def test_state_uses_device_wires(self, wires, expected):
         """Test that if device wires are given, then they are used by state."""
-        dev = DefaultQubit2(wires=wires)
+        dev = DefaultQubit(wires=wires)
 
         @qml.qnode(dev)
         def circuit():
@@ -1960,7 +1960,7 @@ def circuit():
     )
     def test_probs_uses_device_wires(self, wires, expected):
         """Test that if device wires are given, then they are used by probs."""
-        dev = DefaultQubit2(wires=wires)
+        dev = DefaultQubit(wires=wires)
 
         @qml.qnode(dev)
         def circuit():
@@ -1985,7 +1985,7 @@ def circuit():
     )
     def test_counts_uses_device_wires(self, wires, all_outcomes, expected):
         """Test that if device wires are given, then they are used by probs."""
-        dev = DefaultQubit2(wires=wires, shots=10)
+        dev = DefaultQubit(wires=wires, shots=10)
 
         @qml.qnode(dev, interface=None)
         def circuit():
@@ -1998,8 +1998,8 @@ def circuit():
 
 @pytest.mark.parametrize("max_workers", [None, 1, 2])
 def test_broadcasted_parameter(max_workers):
-    """Test that DefaultQubit2 handles broadcasted parameters as expected."""
-    dev = DefaultQubit2(max_workers=max_workers)
+    """Test that DefaultQubit handles broadcasted parameters as expected."""
+    dev = DefaultQubit(max_workers=max_workers)
     x = np.array([0.536, 0.894])
     qs = qml.tape.QuantumScript([qml.RX(x, 0)], [qml.expval(qml.PauliZ(0))])
 
diff --git a/tests/devices/experimental/test_device_api.py b/tests/devices/experimental/test_device_api.py
index 8dfe47296a3..78566e6d401 100644
--- a/tests/devices/experimental/test_device_api.py
+++ b/tests/devices/experimental/test_device_api.py
@@ -18,7 +18,7 @@
 import pytest
 
 import pennylane as qml
-from pennylane.devices.experimental import Device, ExecutionConfig, DefaultExecutionConfig
+from pennylane.devices import Device, ExecutionConfig, DefaultExecutionConfig
 from pennylane.wires import Wires
 
 
@@ -92,7 +92,7 @@ def test_preprocess_single_circuit(self):
 
         a = (1,)
         assert fn(a) == (1,)
-        assert config is qml.devices.experimental.DefaultExecutionConfig
+        assert config is qml.devices.DefaultExecutionConfig
 
     def test_preprocess_batch_circuits(self):
         """Test that preprocessing a batch doesn't do anything."""
diff --git a/tests/devices/qubit/test_adjoint_jacobian.py b/tests/devices/qubit/test_adjoint_jacobian.py
index e931dc68ef5..afac77d5854 100644
--- a/tests/devices/qubit/test_adjoint_jacobian.py
+++ b/tests/devices/qubit/test_adjoint_jacobian.py
@@ -11,7 +11,7 @@
 # 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 and integration tests for the adjoint_jacobian function for DefaultQubit2"""
+"""Unit and integration tests for the adjoint_jacobian function for DefaultQubit"""
 import pytest
 import pennylane as qml
 from pennylane.devices.qubit import adjoint_jacobian, adjoint_jvp, adjoint_vjp
diff --git a/tests/devices/qubit/test_preprocess.py b/tests/devices/qubit/test_preprocess.py
index 405657967fd..58a3bddea1e 100644
--- a/tests/devices/qubit/test_preprocess.py
+++ b/tests/devices/qubit/test_preprocess.py
@@ -29,7 +29,7 @@
     validate_measurements,
     validate_multiprocessing_workers,
 )
-from pennylane.devices.experimental import ExecutionConfig
+from pennylane.devices import ExecutionConfig
 from pennylane.measurements import MidMeasureMP, MeasurementValue
 from pennylane.tape import QuantumScript
 from pennylane import DeviceError
@@ -388,7 +388,7 @@ def test_batch_transform_no_batching(self):
         measurements = [qml.expval(qml.PauliZ(1))]
         tape = QuantumScript(ops=ops, measurements=measurements)
 
-        device = qml.devices.experimental.DefaultQubit2()
+        device = qml.devices.DefaultQubit()
 
         program, _ = device.preprocess()
         tapes, _ = program([tape])
@@ -403,7 +403,7 @@ def test_batch_transform_broadcast_not_adjoint(self):
         ops = [qml.Hadamard(0), qml.CNOT([0, 1]), qml.RX([np.pi, np.pi / 2], wires=1)]
         measurements = [qml.expval(qml.PauliZ(1))]
         tape = QuantumScript(ops=ops, measurements=measurements)
-        device = qml.devices.experimental.DefaultQubit2()
+        device = qml.devices.DefaultQubit()
 
         program, _ = device.preprocess()
         tapes, _ = program([tape])
@@ -422,7 +422,7 @@ def test_batch_transform_broadcast_adjoint(self):
         execution_config = ExecutionConfig()
         execution_config.gradient_method = "adjoint"
 
-        device = qml.devices.experimental.DefaultQubit2()
+        device = qml.devices.DefaultQubit()
 
         program, _ = device.preprocess(execution_config=execution_config)
         tapes, _ = program([tape])
@@ -798,7 +798,7 @@ def test_preprocess_invalid_tape_adjoint(self, ops, measurement, message):
         """Test that preprocessing fails if adjoint differentiation is requested and an
         invalid tape is used"""
         qs = QuantumScript(ops, measurement)
-        execution_config = qml.devices.experimental.ExecutionConfig(gradient_method="adjoint")
+        execution_config = qml.devices.ExecutionConfig(gradient_method="adjoint")
 
         with pytest.raises(DeviceError, match=message):
             program, _ = preprocess(execution_config)
@@ -810,7 +810,7 @@ def test_preprocess_tape_for_adjoint(self):
             [qml.Rot(0.1, 0.2, 0.3, wires=0), qml.CNOT([0, 1])],
             [qml.expval(qml.PauliZ(1))],
         )
-        execution_config = qml.devices.experimental.ExecutionConfig(gradient_method="adjoint")
+        execution_config = qml.devices.ExecutionConfig(gradient_method="adjoint")
 
         program, _ = preprocess(execution_config)
         expanded_tapes, _ = program([qs])
@@ -838,5 +838,5 @@ def test_validate_multiprocessing_workers_None():
         [qml.Rot(0.1, 0.2, 0.3, wires=0), qml.CNOT([0, 1])],
         [qml.expval(qml.PauliZ(1))],
     )
-    device = qml.devices.experimental.DefaultQubit2()
+    device = qml.devices.DefaultQubit()
     validate_multiprocessing_workers(qs, None, device)
diff --git a/tests/devices/test_default_qubit_legacy.py b/tests/devices/test_default_qubit_legacy.py
index a7c299d149a..ec98819a7d4 100644
--- a/tests/devices/test_default_qubit_legacy.py
+++ b/tests/devices/test_default_qubit_legacy.py
@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 """
-Unit tests for the :mod:`pennylane.plugin.DefaultQubit` device.
+Unit tests for the :mod:`pennylane.devices.DefaultQubitLegacy` device.
 """
 # pylint: disable=too-many-arguments,too-few-public-methods
 # pylint: disable=protected-access,cell-var-from-loop
@@ -26,7 +26,7 @@
 import pennylane as qml
 from pennylane import DeviceError
 from pennylane import numpy as np
-from pennylane.devices.default_qubit import DefaultQubit, _get_slice
+from pennylane.devices.default_qubit_legacy import DefaultQubitLegacy, _get_slice
 from pennylane.pulse import ParametrizedHamiltonian
 from pennylane.wires import WireError, Wires
 
@@ -988,7 +988,7 @@ def test_sample_values(self, tol):
         assert np.allclose(s1**2, 1, atol=tol, rtol=0)
 
 
-class TestDefaultQubitIntegration:
+class TestDefaultQubitLegacyIntegration:
     """Integration tests for default.qubit. This test ensures it integrates
     properly with the PennyLane interface, in particular QNode."""
 
@@ -1999,7 +1999,7 @@ def test_get_slice_1d(self):
 
 class TestApplyOps:
     """Tests for special methods listed in _apply_ops that use array manipulation tricks to apply
-    gates in DefaultQubit."""
+    gates in DefaultQubitLegacy."""
 
     state = np.arange(2**4, dtype=np.complex128).reshape((2, 2, 2, 2))
     dev = qml.device("default.qubit.legacy", wires=4)
@@ -2119,7 +2119,7 @@ class TestStateVector:
 
     def test_full_subsystem(self, mocker):
         """Test applying a state vector to the full subsystem"""
-        dev = DefaultQubit(wires=["a", "b", "c"])
+        dev = DefaultQubitLegacy(wires=["a", "b", "c"])
         state = np.array([1, 0, 0, 0, 1, 0, 1, 1]) / 2.0
         state_wires = qml.wires.Wires(["a", "b", "c"])
 
@@ -2132,7 +2132,7 @@ def test_full_subsystem(self, mocker):
     def test_partial_subsystem(self, mocker):
         """Test applying a state vector to a subset of wires of the full subsystem"""
 
-        dev = DefaultQubit(wires=["a", "b", "c"])
+        dev = DefaultQubitLegacy(wires=["a", "b", "c"])
         state = np.array([1, 0, 1, 0]) / np.sqrt(2.0)
         state_wires = qml.wires.Wires(["a", "c"])
 
@@ -2389,7 +2389,7 @@ def test_Hamiltonian_filtered_from_rotations(self, mocker):
 
 @pytest.mark.parametrize("is_state_batched", [False, True])
 class TestSumSupport:
-    """Tests for custom Sum support in DefaultQubit."""
+    """Tests for custom Sum support in DefaultQubitLegacy."""
 
     @staticmethod
     def expected_grad(is_state_batched):
@@ -2504,4 +2504,4 @@ def test_threshold(self, op, n_wires, condition):
         wires = np.linspace(0, n_wires - 1, n_wires, dtype=int)
         op = op(wires=wires)
         # pylint:disable=no-member
-        assert DefaultQubit.stopping_condition.__get__(op)(op) == condition
+        assert DefaultQubitLegacy.stopping_condition.__get__(op)(op) == condition
diff --git a/tests/devices/test_default_qubit_legacy_broadcasting.py b/tests/devices/test_default_qubit_legacy_broadcasting.py
index 594a228cd05..8040f505e4c 100644
--- a/tests/devices/test_default_qubit_legacy_broadcasting.py
+++ b/tests/devices/test_default_qubit_legacy_broadcasting.py
@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 """
-Unit tests for the :mod:`pennylane.plugin.DefaultQubit` device when using broadcasting.
+Unit tests for the :class:`pennylane.devices.DefaultQubitLegacy` device when using broadcasting.
 """
 from itertools import product
 
@@ -53,7 +53,7 @@
 
 import pennylane as qml
 from pennylane import numpy as np, DeviceError
-from pennylane.devices.default_qubit import DefaultQubit
+from pennylane.devices.default_qubit_legacy import DefaultQubitLegacy
 
 THETA = np.linspace(0.11, 1, 3)
 PHI = np.linspace(0.32, 1, 3)
@@ -766,7 +766,7 @@ def test_sample_values_broadcasted(self, tol):
 
 
 class TestDefaultQubitIntegrationBroadcasted:
-    """Integration tests for default.qubit. This test ensures it integrates
+    """Integration tests for default.qubit.legacy. This test ensures it integrates
     properly with the PennyLane interface, in particular QNode."""
 
     @pytest.mark.parametrize("r_dtype", [np.float32, np.float64])
@@ -1616,7 +1616,7 @@ def test_wires_probs_broadcasted(self, wires1, wires2, tol):
 
 class TestApplyOpsBroadcasted:
     """Tests for special methods listed in _apply_ops that use array manipulation tricks to apply
-    gates in DefaultQubit."""
+    gates in DefaultQubitLegacy."""
 
     broadcasted_state = np.arange(2**4 * 3, dtype=np.complex128).reshape((3, 2, 2, 2, 2))
     dev = qml.device("default.qubit.legacy", wires=4)
@@ -1710,7 +1710,7 @@ class TestStateVectorBroadcasted:
 
     def test_full_subsystem_broadcasted(self, mocker):
         """Test applying a state vector to the full subsystem"""
-        dev = DefaultQubit(wires=["a", "b", "c"])
+        dev = DefaultQubitLegacy(wires=["a", "b", "c"])
         state = np.array([[0, 1, 1, 0, 1, 1, 0, 0], [1, 0, 0, 0, 1, 0, 1, 1]]) / 2.0
         state_wires = qml.wires.Wires(["a", "b", "c"])
 
@@ -1723,7 +1723,7 @@ def test_full_subsystem_broadcasted(self, mocker):
     def test_partial_subsystem_broadcasted(self, mocker):
         """Test applying a state vector to a subset of wires of the full subsystem"""
 
-        dev = DefaultQubit(wires=["a", "b", "c"])
+        dev = DefaultQubitLegacy(wires=["a", "b", "c"])
         state = np.array([[0, 1, 1, 0], [1, 0, 1, 0], [1, 1, 0, 0]]) / np.sqrt(2.0)
         state_wires = qml.wires.Wires(["a", "c"])
 
@@ -1743,7 +1743,7 @@ def test_internal_apply_ops_case_broadcasted(self, monkeypatch):
         """Tests that if we provide an operation that has an internal
         implementation, then we use that specific implementation.
 
-        This test provides a new internal function that `default.qubit` uses to
+        This test provides a new internal function that `default.qubit.legacy` uses to
         apply `PauliX` (rather than redefining the gate itself).
         """
         dev = qml.device("default.qubit.legacy", wires=1)
@@ -1923,7 +1923,7 @@ def circuit():
         assert spy.call_count == 2
 
 
-original_capabilities = qml.devices.DefaultQubit.capabilities()
+original_capabilities = qml.devices.DefaultQubitLegacy.capabilities()
 
 
 @pytest.fixture(scope="function", name="mock_default_qubit")
@@ -1938,10 +1938,10 @@ def overwrite_support(*cls):
         return capabilities
 
     with monkeypatch.context() as m:
-        m.setattr(qml.devices.DefaultQubit, "capabilities", overwrite_support)
+        m.setattr(qml.devices.DefaultQubitLegacy, "capabilities", overwrite_support)
 
         def get_default_qubit(wires=1, shots=None):
-            dev = qml.devices.DefaultQubit(wires=wires, shots=shots)
+            dev = qml.devices.DefaultQubitLegacy(wires=wires, shots=shots)
             return dev
 
         yield get_default_qubit
diff --git a/tests/drawer/test_draw.py b/tests/drawer/test_draw.py
index 998a907a61d..d6686d62209 100644
--- a/tests/drawer/test_draw.py
+++ b/tests/drawer/test_draw.py
@@ -331,7 +331,7 @@ def circ():
 
 @pytest.mark.parametrize(
     "device",
-    [qml.device("default.qubit.legacy", wires=2), qml.devices.experimental.DefaultQubit2(wires=2)],
+    [qml.device("default.qubit.legacy", wires=2), qml.devices.DefaultQubit(wires=2)],
 )
 def test_expansion_strategy(device):
     """Test expansion strategy keyword modifies tape expansion."""
diff --git a/tests/drawer/test_draw_mpl.py b/tests/drawer/test_draw_mpl.py
index 96ebfaf858c..b5551a2bbfc 100644
--- a/tests/drawer/test_draw_mpl.py
+++ b/tests/drawer/test_draw_mpl.py
@@ -72,7 +72,7 @@ def test_standard_use():
 
 @pytest.mark.parametrize(
     "device",
-    [qml.device("default.qubit.legacy", wires=3), qml.devices.experimental.DefaultQubit2(wires=3)],
+    [qml.device("default.qubit.legacy", wires=3), qml.devices.DefaultQubit(wires=3)],
 )
 @pytest.mark.parametrize(
     "strategy, initial_strategy, n_lines", [("gradient", "device", 3), ("device", "gradient", 13)]
diff --git a/tests/gradients/finite_diff/test_finite_difference.py b/tests/gradients/finite_diff/test_finite_difference.py
index 405b63b9ccd..003a3cfe665 100644
--- a/tests/gradients/finite_diff/test_finite_difference.py
+++ b/tests/gradients/finite_diff/test_finite_difference.py
@@ -19,7 +19,7 @@
 
 import pennylane as qml
 from pennylane import numpy as np
-from pennylane.devices import DefaultQubit
+from pennylane.devices import DefaultQubitLegacy
 from pennylane.gradients import finite_diff, finite_diff_coeffs
 from pennylane.operation import AnyWires, Observable
 
@@ -499,12 +499,12 @@ def diagonalizing_gates(self):
                 return []
 
         # pylint: disable=too-few-public-methods
-        class DeviceSupportingSpecialObservable(DefaultQubit):
+        class DeviceSupportingSpecialObservable(DefaultQubitLegacy):
             """A device that supports the above SpecialObservable as a return type."""
 
             name = "Device supporting SpecialObservable"
             short_name = "default.qubit.specialobservable"
-            observables = DefaultQubit.observables.union({"SpecialObservable"})
+            observables = DefaultQubitLegacy.observables.union({"SpecialObservable"})
 
             # pylint: disable=unused-argument
             @staticmethod
diff --git a/tests/gradients/finite_diff/test_finite_difference_shot_vec.py b/tests/gradients/finite_diff/test_finite_difference_shot_vec.py
index 5359f93cccd..d335db4068f 100644
--- a/tests/gradients/finite_diff/test_finite_difference_shot_vec.py
+++ b/tests/gradients/finite_diff/test_finite_difference_shot_vec.py
@@ -21,7 +21,7 @@
 
 import pennylane as qml
 from pennylane.gradients import finite_diff
-from pennylane.devices import DefaultQubit
+from pennylane.devices import DefaultQubitLegacy
 from pennylane.measurements import Shots
 from pennylane.operation import Observable, AnyWires
 
@@ -451,12 +451,12 @@ def diagonalizing_gates(self):
                 return []
 
         # pylint: disable=too-few-public-methods
-        class DeviceSupportingSpecialObservable(DefaultQubit):
+        class DeviceSupportingSpecialObservable(DefaultQubitLegacy):
             """A device that supports the above SpecialObservable as a return type."""
 
             name = "Device supporting SpecialObservable"
             short_name = "default.qubit.specialobservable"
-            observables = DefaultQubit.observables.union({"SpecialObservable"})
+            observables = DefaultQubitLegacy.observables.union({"SpecialObservable"})
 
             # pylint: disable=unused-argument
             @staticmethod
diff --git a/tests/gradients/finite_diff/test_spsa_gradient.py b/tests/gradients/finite_diff/test_spsa_gradient.py
index 1721961d479..44a149041cb 100644
--- a/tests/gradients/finite_diff/test_spsa_gradient.py
+++ b/tests/gradients/finite_diff/test_spsa_gradient.py
@@ -19,7 +19,7 @@
 
 import pennylane as qml
 from pennylane import numpy as np
-from pennylane.devices import DefaultQubit
+from pennylane.devices import DefaultQubitLegacy
 from pennylane.gradients import spsa_grad
 from pennylane.gradients.spsa_gradient import _rademacher_sampler
 from pennylane.operation import AnyWires, Observable
@@ -588,13 +588,13 @@ def diagonalizing_gates(self):
                 """Diagonalizing gates"""
                 return []
 
-        class DeviceSupportingSpecialObservable(DefaultQubit):
+        class DeviceSupportingSpecialObservable(DefaultQubitLegacy):
             """A device class supporting SpecialObservable."""
 
             # pylint:disable=too-few-public-methods
             name = "Device supporting SpecialObservable"
             short_name = "default.qubit.specialobservable"
-            observables = DefaultQubit.observables.union({"SpecialObservable"})
+            observables = DefaultQubitLegacy.observables.union({"SpecialObservable"})
 
             @staticmethod
             def _asarray(arr, dtype=None):
diff --git a/tests/gradients/finite_diff/test_spsa_gradient_shot_vec.py b/tests/gradients/finite_diff/test_spsa_gradient_shot_vec.py
index 621f54558f5..3d8fa8017df 100644
--- a/tests/gradients/finite_diff/test_spsa_gradient_shot_vec.py
+++ b/tests/gradients/finite_diff/test_spsa_gradient_shot_vec.py
@@ -21,7 +21,7 @@
 
 import pennylane as qml
 from pennylane.gradients import spsa_grad
-from pennylane.devices import DefaultQubit
+from pennylane.devices import DefaultQubitLegacy
 from pennylane.measurements import Shots
 from pennylane.operation import Observable, AnyWires
 
@@ -513,13 +513,13 @@ def diagonalizing_gates(self):
                 """Diagonalizing gates"""
                 return []
 
-        class DeviceSupportingSpecialObservable(DefaultQubit):
+        class DeviceSupportingSpecialObservable(DefaultQubitLegacy):
             """A device class supporting SpecialObservable."""
 
             # pylint:disable=too-few-public-methods
             name = "Device supporting SpecialObservable"
             short_name = "default.qubit.specialobservable"
-            observables = DefaultQubit.observables.union({"SpecialObservable"})
+            observables = DefaultQubitLegacy.observables.union({"SpecialObservable"})
 
             @staticmethod
             def _asarray(arr, dtype=None):
diff --git a/tests/gradients/parameter_shift/test_parameter_shift.py b/tests/gradients/parameter_shift/test_parameter_shift.py
index 0b044daf338..a8e0aaa1858 100644
--- a/tests/gradients/parameter_shift/test_parameter_shift.py
+++ b/tests/gradients/parameter_shift/test_parameter_shift.py
@@ -22,7 +22,7 @@
     _put_zeros_in_pdA2_involutory,
     _make_zero_rep,
 )
-from pennylane.devices import DefaultQubit
+from pennylane.devices import DefaultQubitLegacy
 from pennylane.operation import Observable, AnyWires
 
 
@@ -2373,12 +2373,12 @@ def diagonalizing_gates(self):
                 return []
 
         # pylint: disable=too-few-public-methods
-        class DeviceSupporingSpecialObservable(DefaultQubit):
+        class DeviceSupporingSpecialObservable(DefaultQubitLegacy):
             """A custom device that supports the above special observable."""
 
             name = "Device supporting SpecialObservable"
             short_name = "default.qubit.specialobservable"
-            observables = DefaultQubit.observables.union({"SpecialObservable"})
+            observables = DefaultQubitLegacy.observables.union({"SpecialObservable"})
 
             # pylint: disable=unused-argument
             @staticmethod
diff --git a/tests/gradients/parameter_shift/test_parameter_shift_shot_vec.py b/tests/gradients/parameter_shift/test_parameter_shift_shot_vec.py
index be15d8efb04..16499641081 100644
--- a/tests/gradients/parameter_shift/test_parameter_shift_shot_vec.py
+++ b/tests/gradients/parameter_shift/test_parameter_shift_shot_vec.py
@@ -19,7 +19,7 @@
 import pennylane as qml
 from pennylane import numpy as np
 from pennylane.gradients import param_shift
-from pennylane.devices import DefaultQubit
+from pennylane.devices import DefaultQubitLegacy
 from pennylane.measurements import Shots
 from pennylane.operation import Observable, AnyWires
 
@@ -2014,12 +2014,12 @@ def diagonalizing_gates(self):
                 """Diagonalizing gates"""
                 return []
 
-        class DeviceSupporingSpecialObservable(DefaultQubit):
+        class DeviceSupporingSpecialObservable(DefaultQubitLegacy):
             """A custom device that supports the above SpecialObservable."""
 
             name = "Device supporting SpecialObservable"
             short_name = "default.qubit.specialobservable"
-            observables = DefaultQubit.observables.union({"SpecialObservable"})
+            observables = DefaultQubitLegacy.observables.union({"SpecialObservable"})
 
             # pylint: disable=unused-argument
             @staticmethod
diff --git a/tests/interfaces/default_qubit_2_integration/test_autograd_default_qubit_2.py b/tests/interfaces/default_qubit_2_integration/test_autograd_default_qubit_2.py
index f6377858885..0f2c822f0ba 100644
--- a/tests/interfaces/default_qubit_2_integration/test_autograd_default_qubit_2.py
+++ b/tests/interfaces/default_qubit_2_integration/test_autograd_default_qubit_2.py
@@ -17,7 +17,7 @@
 from pennylane import numpy as np
 
 import pennylane as qml
-from pennylane.devices.experimental import DefaultQubit2
+from pennylane.devices import DefaultQubit
 from pennylane.gradients import param_shift
 from pennylane.interfaces import execute
 
@@ -33,7 +33,7 @@ def test_caching_param_shift_hessian(self, num_params):
         """Test that, when using parameter-shift transform,
         caching reduces the number of evaluations to their optimum
         when computing Hessians."""
-        dev = DefaultQubit2()
+        dev = DefaultQubit()
         params = np.arange(1, num_params + 1) / 10
 
         N = len(params)
@@ -101,7 +101,7 @@ def test_single_backward_pass_batch(self):
         """Tests that the backward pass is one single batch, not a bunch of batches, when parameter shift
         is requested for multiple tapes."""
 
-        dev = DefaultQubit2()
+        dev = DefaultQubit()
 
         def f(x):
             tape1 = qml.tape.QuantumScript([qml.RX(x, 0)], [qml.probs(wires=0)])
@@ -123,11 +123,11 @@ def f(x):
 # add tests for lightning 2 when possible
 # set rng for device when possible
 test_matrix = [
-    ({"gradient_fn": param_shift}, 100000, DefaultQubit2(seed=42)),
-    ({"gradient_fn": param_shift}, None, DefaultQubit2()),
-    ({"gradient_fn": "backprop"}, None, DefaultQubit2()),
-    ({"gradient_fn": "adjoint", "grad_on_execution": True}, None, DefaultQubit2()),
-    ({"gradient_fn": "adjoint", "grad_on_execution": False}, None, DefaultQubit2()),
+    ({"gradient_fn": param_shift}, 100000, DefaultQubit(seed=42)),
+    ({"gradient_fn": param_shift}, None, DefaultQubit()),
+    ({"gradient_fn": "backprop"}, None, DefaultQubit()),
+    ({"gradient_fn": "adjoint", "grad_on_execution": True}, None, DefaultQubit()),
+    ({"gradient_fn": "adjoint", "grad_on_execution": False}, None, DefaultQubit()),
 ]
 
 
@@ -572,7 +572,7 @@ class TestHigherOrderDerivatives:
     def test_parameter_shift_hessian(self, params, tol):
         """Tests that the output of the parameter-shift transform
         can be differentiated using autograd, yielding second derivatives."""
-        dev = DefaultQubit2()
+        dev = DefaultQubit()
 
         def cost_fn(x):
             ops1 = [qml.RX(x[0], 0), qml.RY(x[1], 1), qml.CNOT((0, 1))]
@@ -606,7 +606,7 @@ def cost_fn(x):
     def test_max_diff(self, tol):
         """Test that setting the max_diff parameter blocks higher-order
         derivatives"""
-        dev = DefaultQubit2()
+        dev = DefaultQubit()
         params = np.array([0.543, -0.654], requires_grad=True)
 
         def cost_fn(x):
diff --git a/tests/interfaces/default_qubit_2_integration/test_autograd_qnode_default_qubit_2.py b/tests/interfaces/default_qubit_2_integration/test_autograd_qnode_default_qubit_2.py
index dc4b54bf66d..3463e581096 100644
--- a/tests/interfaces/default_qubit_2_integration/test_autograd_qnode_default_qubit_2.py
+++ b/tests/interfaces/default_qubit_2_integration/test_autograd_qnode_default_qubit_2.py
@@ -21,33 +21,33 @@
 import pennylane as qml
 from pennylane import numpy as np
 from pennylane import qnode
-from pennylane.devices.experimental import DefaultQubit2
+from pennylane.devices import DefaultQubit
 
 qubit_device_and_diff_method = [
-    [DefaultQubit2(), "finite-diff", False],
-    [DefaultQubit2(), "parameter-shift", False],
-    [DefaultQubit2(), "backprop", True],
-    [DefaultQubit2(), "adjoint", True],
-    [DefaultQubit2(), "adjoint", False],
-    [DefaultQubit2(), "spsa", False],
-    [DefaultQubit2(), "hadamard", False],
+    [DefaultQubit(), "finite-diff", False],
+    [DefaultQubit(), "parameter-shift", False],
+    [DefaultQubit(), "backprop", True],
+    [DefaultQubit(), "adjoint", True],
+    [DefaultQubit(), "adjoint", False],
+    [DefaultQubit(), "spsa", False],
+    [DefaultQubit(), "hadamard", False],
 ]
 
 interface_qubit_device_and_diff_method = [
-    ["autograd", DefaultQubit2(), "finite-diff", False],
-    ["autograd", DefaultQubit2(), "parameter-shift", False],
-    ["autograd", DefaultQubit2(), "backprop", True],
-    ["autograd", DefaultQubit2(), "adjoint", True],
-    ["autograd", DefaultQubit2(), "adjoint", False],
-    ["autograd", DefaultQubit2(), "spsa", False],
-    ["autograd", DefaultQubit2(), "hadamard", False],
-    ["auto", DefaultQubit2(), "finite-diff", False],
-    ["auto", DefaultQubit2(), "parameter-shift", False],
-    ["auto", DefaultQubit2(), "backprop", True],
-    ["auto", DefaultQubit2(), "adjoint", True],
-    ["auto", DefaultQubit2(), "adjoint", False],
-    ["auto", DefaultQubit2(), "spsa", False],
-    ["auto", DefaultQubit2(), "hadamard", False],
+    ["autograd", DefaultQubit(), "finite-diff", False],
+    ["autograd", DefaultQubit(), "parameter-shift", False],
+    ["autograd", DefaultQubit(), "backprop", True],
+    ["autograd", DefaultQubit(), "adjoint", True],
+    ["autograd", DefaultQubit(), "adjoint", False],
+    ["autograd", DefaultQubit(), "spsa", False],
+    ["autograd", DefaultQubit(), "hadamard", False],
+    ["auto", DefaultQubit(), "finite-diff", False],
+    ["auto", DefaultQubit(), "parameter-shift", False],
+    ["auto", DefaultQubit(), "backprop", True],
+    ["auto", DefaultQubit(), "adjoint", True],
+    ["auto", DefaultQubit(), "adjoint", False],
+    ["auto", DefaultQubit(), "spsa", False],
+    ["auto", DefaultQubit(), "hadamard", False],
 ]
 
 pytestmark = pytest.mark.autograd
@@ -459,7 +459,7 @@ class TestShotsIntegration:
 
     def test_changing_shots(self):
         """Test that changing shots works on execution"""
-        dev = DefaultQubit2()
+        dev = DefaultQubit()
         a, b = np.array([0.543, -0.654], requires_grad=True)
 
         @qnode(dev, diff_method=qml.gradients.param_shift)
@@ -481,7 +481,7 @@ def circuit(a, b):
     def test_gradient_integration(self):
         """Test that temporarily setting the shots works
         for gradient computations"""
-        dev = DefaultQubit2()
+        dev = DefaultQubit()
         a, b = np.array([0.543, -0.654], requires_grad=True)
 
         @qnode(dev, diff_method=qml.gradients.param_shift)
@@ -508,7 +508,7 @@ def test_update_diff_method(self, mocker):
 
         spy = mocker.spy(qml, "execute")
 
-        @qnode(DefaultQubit2())
+        @qnode(DefaultQubit())
         def cost_fn(a, b):
             qml.RY(a, wires=0)
             qml.RX(b, wires=1)
@@ -1536,7 +1536,7 @@ class TestSample:
     @pytest.mark.xfail
     def test_backprop_error(self):
         """Test that sampling in backpropagation grad_on_execution raises an error"""
-        dev = DefaultQubit2()
+        dev = DefaultQubit()
 
         @qnode(dev, diff_method="backprop")
         def circuit():
@@ -1550,7 +1550,7 @@ def test_sample_dimension(self):
         """Test that the sample function outputs samples of the right size"""
         n_sample = 10
 
-        dev = DefaultQubit2()
+        dev = DefaultQubit()
 
         @qnode(dev, diff_method=None)
         def circuit():
@@ -1573,7 +1573,7 @@ def test_sample_combination(self):
 
         n_sample = 10
 
-        dev = DefaultQubit2()
+        dev = DefaultQubit()
 
         @qnode(dev, diff_method="parameter-shift")
         def circuit():
@@ -1595,7 +1595,7 @@ def test_single_wire_sample(self):
         """Test the return type and shape of sampling a single wire"""
         n_sample = 10
 
-        dev = DefaultQubit2()
+        dev = DefaultQubit()
 
         @qnode(dev, diff_method=None)
         def circuit():
@@ -1613,7 +1613,7 @@ def test_multi_wire_sample_regular_shape(self):
         where a rectangular array is expected"""
         n_sample = 10
 
-        dev = DefaultQubit2()
+        dev = DefaultQubit()
 
         @qnode(dev, diff_method=None)
         def circuit():
@@ -2123,7 +2123,7 @@ def test_no_ops():
     """Test that the return value of the QNode matches in the interface
     even if there are no ops"""
 
-    dev = DefaultQubit2()
+    dev = DefaultQubit()
 
     @qml.qnode(dev, interface="autograd")
     def circuit():
diff --git a/tests/interfaces/default_qubit_2_integration/test_execute_default_qubit_2.py b/tests/interfaces/default_qubit_2_integration/test_execute_default_qubit_2.py
index 68f1f5c8d53..d0cde46aaf2 100644
--- a/tests/interfaces/default_qubit_2_integration/test_execute_default_qubit_2.py
+++ b/tests/interfaces/default_qubit_2_integration/test_execute_default_qubit_2.py
@@ -18,7 +18,7 @@
 from pennylane import numpy as np
 from pennylane.interfaces.execution import _preprocess_expand_fn
 
-from pennylane.devices.experimental import DefaultQubit2
+from pennylane.devices import DefaultQubit
 
 
 class TestPreprocessExpandFn:
@@ -27,7 +27,7 @@ class TestPreprocessExpandFn:
     def test_provided_is_callable(self):
         """Test that if the expand_fn is not "device", it is simply returned."""
 
-        dev = DefaultQubit2()
+        dev = DefaultQubit()
 
         def f(tape):
             return tape
@@ -38,7 +38,7 @@ def f(tape):
     def test_new_device_blank_expand_fn(self):
         """Test that the expand_fn is blank if is new device."""
 
-        dev = DefaultQubit2()
+        dev = DefaultQubit()
 
         out = _preprocess_expand_fn("device", dev, 10)
 
@@ -59,7 +59,7 @@ class CustomOp(qml.operation.Operator):
             def decomposition(self):
                 return [qml.PauliX(self.wires[0])]
 
-        dev = DefaultQubit2()
+        dev = DefaultQubit()
 
         qs = qml.tape.QuantumScript([CustomOp(0)], [qml.expval(qml.PauliZ(0))])
 
@@ -92,8 +92,8 @@ def decomposition(self):
             qml.tape.QuantumScript(ops=ops, measurements=[measurements[1]]),
         ]
 
-        dev = DefaultQubit2()
-        config = qml.devices.experimental.ExecutionConfig(gradient_method="adjoint")
+        dev = DefaultQubit()
+        config = qml.devices.ExecutionConfig(gradient_method="adjoint")
 
         program, new_config = dev.preprocess(config)
         res_tapes, batch_fn = program(tapes)
@@ -129,7 +129,7 @@ class CustomOp(qml.operation.Operator):
         def decomposition(self):
             return [qml.PauliX(self.wires[0])]
 
-    dev = DefaultQubit2()
+    dev = DefaultQubit()
 
     qs = qml.tape.QuantumScript([CustomOp(0)], [qml.expval(qml.PauliZ(0))])
 
@@ -144,7 +144,7 @@ def decomposition(self):
 def test_caching(gradient_fn):
     """Test that cache execute returns the cached result if the same script is executed
     multiple times, both in multiple times in a batch and in separate batches."""
-    dev = DefaultQubit2()
+    dev = DefaultQubit()
 
     qs = qml.tape.QuantumScript([qml.PauliX(0)], [qml.expval(qml.PauliZ(0))])
 
diff --git a/tests/interfaces/default_qubit_2_integration/test_jax_default_qubit_2.py b/tests/interfaces/default_qubit_2_integration/test_jax_default_qubit_2.py
index 0e52e35e384..5d1ef9cf9a7 100644
--- a/tests/interfaces/default_qubit_2_integration/test_jax_default_qubit_2.py
+++ b/tests/interfaces/default_qubit_2_integration/test_jax_default_qubit_2.py
@@ -17,7 +17,7 @@
 
 
 import pennylane as qml
-from pennylane.devices.experimental import DefaultQubit2
+from pennylane.devices import DefaultQubit
 from pennylane.gradients import param_shift
 from pennylane.interfaces import execute
 
@@ -39,7 +39,7 @@ def test_caching_param_shift_hessian(self, num_params):
         """Test that, when using parameter-shift transform,
         caching reduces the number of evaluations to their optimum
         when computing Hessians."""
-        device = DefaultQubit2()
+        device = DefaultQubit()
         params = jnp.arange(1, num_params + 1) / 10
 
         N = len(params)
@@ -107,10 +107,10 @@ def cost(x, cache):
 # add tests for lightning 2 when possible
 # set rng for device when possible
 test_matrix = [
-    ({"gradient_fn": param_shift}, 100000, DefaultQubit2(seed=42)),
-    ({"gradient_fn": param_shift}, None, DefaultQubit2()),
-    ({"gradient_fn": "backprop"}, None, DefaultQubit2()),
-    ({"gradient_fn": "adjoint"}, None, DefaultQubit2()),
+    ({"gradient_fn": param_shift}, 100000, DefaultQubit(seed=42)),
+    ({"gradient_fn": param_shift}, None, DefaultQubit()),
+    ({"gradient_fn": "backprop"}, None, DefaultQubit()),
+    ({"gradient_fn": "adjoint"}, None, DefaultQubit()),
 ]
 
 
@@ -532,7 +532,7 @@ class TestHigherOrderDerivatives:
     def test_parameter_shift_hessian(self, params, tol):
         """Tests that the output of the parameter-shift transform
         can be differentiated using jax, yielding second derivatives."""
-        dev = DefaultQubit2()
+        dev = DefaultQubit()
 
         def cost_fn(x):
             ops1 = [qml.RX(x[0], 0), qml.RY(x[1], 1), qml.CNOT((0, 1))]
@@ -566,7 +566,7 @@ def cost_fn(x):
     def test_max_diff(self, tol):
         """Test that setting the max_diff parameter blocks higher-order
         derivatives"""
-        dev = DefaultQubit2()
+        dev = DefaultQubit()
         params = jnp.array([0.543, -0.654])
 
         def cost_fn(x):
diff --git a/tests/interfaces/default_qubit_2_integration/test_jax_jit_qnode_default_qubit_2.py b/tests/interfaces/default_qubit_2_integration/test_jax_jit_qnode_default_qubit_2.py
index e15e14bc19c..22b3bb88f37 100644
--- a/tests/interfaces/default_qubit_2_integration/test_jax_jit_qnode_default_qubit_2.py
+++ b/tests/interfaces/default_qubit_2_integration/test_jax_jit_qnode_default_qubit_2.py
@@ -19,16 +19,16 @@
 import pennylane as qml
 from pennylane import numpy as np
 from pennylane import qnode
-from pennylane.devices.experimental import DefaultQubit2
+from pennylane.devices import DefaultQubit
 
 qubit_device_and_diff_method = [
-    [DefaultQubit2(), "backprop", True],
-    [DefaultQubit2(), "finite-diff", False],
-    [DefaultQubit2(), "parameter-shift", False],
-    [DefaultQubit2(), "adjoint", True],
-    [DefaultQubit2(), "adjoint", False],
-    [DefaultQubit2(), "spsa", False],
-    [DefaultQubit2(), "hadamard", False],
+    [DefaultQubit(), "backprop", True],
+    [DefaultQubit(), "finite-diff", False],
+    [DefaultQubit(), "parameter-shift", False],
+    [DefaultQubit(), "adjoint", True],
+    [DefaultQubit(), "adjoint", False],
+    [DefaultQubit(), "spsa", False],
+    [DefaultQubit(), "hadamard", False],
 ]
 interface_and_qubit_device_and_diff_method = [
     ["auto"] + inner_list for inner_list in qubit_device_and_diff_method
@@ -739,7 +739,7 @@ class TestShotsIntegration:
 
     def test_diff_method_None(self, interface):
         """Test device works with diff_method=None."""
-        dev = DefaultQubit2()
+        dev = DefaultQubit()
 
         @jax.jit
         @qml.qnode(dev, diff_method=None, interface=interface)
@@ -754,7 +754,7 @@ def test_changing_shots(self, interface):
         """Test that changing shots works on execution"""
         a, b = jax.numpy.array([0.543, -0.654])
 
-        @qnode(DefaultQubit2(), diff_method=qml.gradients.param_shift, interface=interface)
+        @qnode(DefaultQubit(), diff_method=qml.gradients.param_shift, interface=interface)
         def circuit(a, b):
             qml.RY(a, wires=0)
             qml.RX(b, wires=1)
@@ -774,7 +774,7 @@ def test_gradient_integration(self, interface):
         for gradient computations"""
         a, b = jax.numpy.array([0.543, -0.654])
 
-        @qnode(DefaultQubit2(), diff_method=qml.gradients.param_shift, interface=interface)
+        @qnode(DefaultQubit(), diff_method=qml.gradients.param_shift, interface=interface)
         def cost_fn(a, b):
             qml.RY(a, wires=0)
             qml.RX(b, wires=1)
@@ -794,7 +794,7 @@ def test_update_diff_method(self, mocker, interface):
 
         spy = mocker.spy(qml, "execute")
 
-        @qnode(DefaultQubit2(), interface=interface)
+        @qnode(DefaultQubit(), interface=interface)
         def cost_fn(a, b):
             qml.RY(a, wires=0)
             qml.RX(b, wires=1)
@@ -2480,7 +2480,7 @@ def test_jax_device_hessian_shots(hessian, diff_method):
     """The hessian of multiple measurements with a multiple param array return a single array."""
 
     @partial(jax.jit, static_argnames="shots")
-    @qml.qnode(DefaultQubit2(), diff_method=diff_method, max_diff=2)
+    @qml.qnode(DefaultQubit(), diff_method=diff_method, max_diff=2)
     def circuit(x):
         qml.RY(x[0], wires=0)
         qml.RX(x[1], wires=0)
diff --git a/tests/interfaces/default_qubit_2_integration/test_jax_qnode_default_qubit_2.py b/tests/interfaces/default_qubit_2_integration/test_jax_qnode_default_qubit_2.py
index 89b96e4ebae..0a9500c6b9d 100644
--- a/tests/interfaces/default_qubit_2_integration/test_jax_qnode_default_qubit_2.py
+++ b/tests/interfaces/default_qubit_2_integration/test_jax_qnode_default_qubit_2.py
@@ -20,18 +20,18 @@
 
 import pennylane as qml
 from pennylane import qnode
-from pennylane.devices.experimental import DefaultQubit2
+from pennylane.devices import DefaultQubit
 
 device_seed = 42
 
 device_and_diff_method = [
-    [DefaultQubit2(seed=device_seed), "backprop", True],
-    [DefaultQubit2(seed=device_seed), "finite-diff", False],
-    [DefaultQubit2(seed=device_seed), "parameter-shift", False],
-    [DefaultQubit2(seed=device_seed), "adjoint", True],
-    [DefaultQubit2(seed=device_seed), "adjoint", False],
-    [DefaultQubit2(seed=device_seed), "spsa", False],
-    [DefaultQubit2(seed=device_seed), "hadamard", False],
+    [DefaultQubit(seed=device_seed), "backprop", True],
+    [DefaultQubit(seed=device_seed), "finite-diff", False],
+    [DefaultQubit(seed=device_seed), "parameter-shift", False],
+    [DefaultQubit(seed=device_seed), "adjoint", True],
+    [DefaultQubit(seed=device_seed), "adjoint", False],
+    [DefaultQubit(seed=device_seed), "spsa", False],
+    [DefaultQubit(seed=device_seed), "hadamard", False],
 ]
 
 interface_and_device_and_diff_method = [
@@ -685,7 +685,7 @@ class TestShotsIntegration:
     def test_diff_method_None(self, interface):
         """Test device works with diff_method=None."""
 
-        @qml.qnode(DefaultQubit2(), diff_method=None, interface=interface)
+        @qml.qnode(DefaultQubit(), diff_method=None, interface=interface)
         def circuit(x):
             qml.RX(x, wires=0)
             return qml.expval(qml.PauliZ(0))
@@ -696,7 +696,7 @@ def test_changing_shots(self, interface):
         """Test that changing shots works on execution"""
         a, b = jax.numpy.array([0.543, -0.654])
 
-        @qnode(DefaultQubit2(), diff_method=qml.gradients.param_shift, interface=interface)
+        @qnode(DefaultQubit(), diff_method=qml.gradients.param_shift, interface=interface)
         def circuit(a, b):
             qml.RY(a, wires=0)
             qml.RX(b, wires=1)
@@ -716,7 +716,7 @@ def test_gradient_integration(self, interface):
         for gradient computations"""
         a, b = jax.numpy.array([0.543, -0.654])
 
-        @qnode(DefaultQubit2(), diff_method=qml.gradients.param_shift, interface=interface)
+        @qnode(DefaultQubit(), diff_method=qml.gradients.param_shift, interface=interface)
         def cost_fn(a, b):
             qml.RY(a, wires=0)
             qml.RX(b, wires=1)
@@ -734,7 +734,7 @@ def test_update_diff_method(self, interface, mocker):
 
         spy = mocker.spy(qml, "execute")
 
-        @qnode(DefaultQubit2(), interface=interface)
+        @qnode(DefaultQubit(), interface=interface)
         def cost_fn(a, b):
             qml.RY(a, wires=0)
             qml.RX(b, wires=1)
@@ -2176,7 +2176,7 @@ def test_no_ops():
     """Test that the return value of the QNode matches in the interface
     even if there are no ops"""
 
-    @qml.qnode(DefaultQubit2(), interface="jax")
+    @qml.qnode(DefaultQubit(), interface="jax")
     def circuit():
         qml.Hadamard(wires=0)
         return qml.state()
diff --git a/tests/interfaces/default_qubit_2_integration/test_tensorflow_autograph_qnode_shot_vector_default_qubit_2.py b/tests/interfaces/default_qubit_2_integration/test_tensorflow_autograph_qnode_shot_vector_default_qubit_2.py
index ed6a7439d92..b24f223f248 100644
--- a/tests/interfaces/default_qubit_2_integration/test_tensorflow_autograph_qnode_shot_vector_default_qubit_2.py
+++ b/tests/interfaces/default_qubit_2_integration/test_tensorflow_autograph_qnode_shot_vector_default_qubit_2.py
@@ -18,7 +18,7 @@
 import pennylane as qml
 from pennylane import numpy as np
 from pennylane import qnode
-from pennylane.devices.experimental import DefaultQubit2
+from pennylane.devices import DefaultQubit
 
 pytestmark = pytest.mark.tf
 
@@ -28,10 +28,10 @@
 shots_and_num_copies_hess = [((10, (5, 1)), 2)]
 
 qubit_device_and_diff_method = [
-    [DefaultQubit2(), "finite-diff", {"h": 10e-2}],
-    [DefaultQubit2(), "parameter-shift", {}],
+    [DefaultQubit(), "finite-diff", {"h": 10e-2}],
+    [DefaultQubit(), "parameter-shift", {}],
     [
-        DefaultQubit2(),
+        DefaultQubit(),
         "spsa",
         {"h": 10e-2, "num_directions": 20, "sampler_rng": np.random.default_rng(42)},
     ],
diff --git a/tests/interfaces/default_qubit_2_integration/test_tensorflow_default_qubit_2.py b/tests/interfaces/default_qubit_2_integration/test_tensorflow_default_qubit_2.py
index 408e8fe359c..0312d5875ef 100644
--- a/tests/interfaces/default_qubit_2_integration/test_tensorflow_default_qubit_2.py
+++ b/tests/interfaces/default_qubit_2_integration/test_tensorflow_default_qubit_2.py
@@ -16,7 +16,7 @@
 import numpy as np
 
 import pennylane as qml
-from pennylane.devices.experimental import DefaultQubit2
+from pennylane.devices import DefaultQubit
 from pennylane.gradients import param_shift
 from pennylane.interfaces import execute
 
@@ -33,7 +33,7 @@ def test_caching_param_shift_hessian(self, num_params):
         """Test that, when using parameter-shift transform,
         caching reduces the number of evaluations to their optimum
         when computing Hessians."""
-        dev = DefaultQubit2()
+        dev = DefaultQubit()
         params = tf.Variable(tf.range(1, num_params + 1) / 10)
 
         N = num_params
@@ -109,14 +109,14 @@ def cost(x, cache):
 # add tests for lightning 2 when possible
 # set rng for device when possible
 test_matrix = [
-    ({"gradient_fn": param_shift, "interface": "tensorflow"}, 100000, DefaultQubit2(seed=42)),
-    ({"gradient_fn": param_shift, "interface": "tensorflow"}, None, DefaultQubit2()),
-    ({"gradient_fn": "backprop", "interface": "tensorflow"}, None, DefaultQubit2()),
-    ({"gradient_fn": "adjoint", "interface": "tensorflow"}, None, DefaultQubit2()),
-    ({"gradient_fn": param_shift, "interface": "tf-autograph"}, 100000, DefaultQubit2(seed=42)),
-    ({"gradient_fn": param_shift, "interface": "tf-autograph"}, None, DefaultQubit2()),
-    ({"gradient_fn": "backprop", "interface": "tf-autograph"}, None, DefaultQubit2()),
-    ({"gradient_fn": "adjoint", "interface": "tf-autograph"}, None, DefaultQubit2()),
+    ({"gradient_fn": param_shift, "interface": "tensorflow"}, 100000, DefaultQubit(seed=42)),
+    ({"gradient_fn": param_shift, "interface": "tensorflow"}, None, DefaultQubit()),
+    ({"gradient_fn": "backprop", "interface": "tensorflow"}, None, DefaultQubit()),
+    ({"gradient_fn": "adjoint", "interface": "tensorflow"}, None, DefaultQubit()),
+    ({"gradient_fn": param_shift, "interface": "tf-autograph"}, 100000, DefaultQubit(seed=42)),
+    ({"gradient_fn": param_shift, "interface": "tf-autograph"}, None, DefaultQubit()),
+    ({"gradient_fn": "backprop", "interface": "tf-autograph"}, None, DefaultQubit()),
+    ({"gradient_fn": "adjoint", "interface": "tf-autograph"}, None, DefaultQubit()),
 ]
 
 
@@ -591,7 +591,7 @@ class TestHigherOrderDerivatives:
     def test_parameter_shift_hessian(self, params, tol):
         """Tests that the output of the parameter-shift transform
         can be differentiated using tensorflow, yielding second derivatives."""
-        dev = DefaultQubit2()
+        dev = DefaultQubit()
 
         def cost_fn(x):
             ops1 = [qml.RX(x[0], 0), qml.RY(x[1], 1), qml.CNOT((0, 1))]
@@ -628,7 +628,7 @@ def cost_fn(x):
     def test_max_diff(self, tol):
         """Test that setting the max_diff parameter blocks higher-order
         derivatives"""
-        dev = DefaultQubit2()
+        dev = DefaultQubit()
         params = tf.Variable([0.543, -0.654])
 
         def cost_fn(x):
diff --git a/tests/interfaces/default_qubit_2_integration/test_tensorflow_qnode_default_qubit_2.py b/tests/interfaces/default_qubit_2_integration/test_tensorflow_qnode_default_qubit_2.py
index f7d722b3344..b894251fa5f 100644
--- a/tests/interfaces/default_qubit_2_integration/test_tensorflow_qnode_default_qubit_2.py
+++ b/tests/interfaces/default_qubit_2_integration/test_tensorflow_qnode_default_qubit_2.py
@@ -18,20 +18,20 @@
 
 import pennylane as qml
 from pennylane import qnode
-from pennylane.devices.experimental import DefaultQubit2
+from pennylane.devices import DefaultQubit
 
 pytestmark = pytest.mark.tf
 tf = pytest.importorskip("tensorflow")
 
 
 qubit_device_and_diff_method = [
-    [DefaultQubit2(), "finite-diff", False],
-    [DefaultQubit2(), "parameter-shift", False],
-    [DefaultQubit2(), "backprop", True],
-    [DefaultQubit2(), "adjoint", True],
-    [DefaultQubit2(), "adjoint", False],
-    [DefaultQubit2(), "spsa", False],
-    [DefaultQubit2(), "hadamard", False],
+    [DefaultQubit(), "finite-diff", False],
+    [DefaultQubit(), "parameter-shift", False],
+    [DefaultQubit(), "backprop", True],
+    [DefaultQubit(), "adjoint", True],
+    [DefaultQubit(), "adjoint", False],
+    [DefaultQubit(), "spsa", False],
+    [DefaultQubit(), "hadamard", False],
 ]
 
 TOL_FOR_SPSA = 1.0
@@ -435,7 +435,7 @@ class TestShotsIntegration:
 
     def test_changing_shots(self, interface):
         """Test that changing shots works on execution"""
-        dev = DefaultQubit2()
+        dev = DefaultQubit()
         a, b = [0.543, -0.654]
         weights = tf.Variable([a, b], dtype=tf.float64)
 
@@ -459,7 +459,7 @@ def test_gradient_integration(self, interface):
         """Test that temporarily setting the shots works
         for gradient computations"""
         # pylint: disable=unexpected-keyword-arg
-        dev = DefaultQubit2()
+        dev = DefaultQubit()
         a, b = [0.543, -0.654]
         weights = tf.Variable([a, b], dtype=tf.float64)
 
@@ -484,7 +484,7 @@ def test_multiple_gradient_integration(self, tol, interface):
         """Test that temporarily setting the shots works
         for gradient computations, even if the QNode has been re-evaluated
         with a different number of shots in the meantime."""
-        dev = DefaultQubit2()
+        dev = DefaultQubit()
         a, b = [0.543, -0.654]
         weights = tf.Variable([a, b], dtype=tf.float64)
 
@@ -509,7 +509,7 @@ def circuit(weights):
 
     def test_update_diff_method(self, mocker, interface):
         """Test that temporarily setting the shots updates the diff method"""
-        dev = DefaultQubit2()
+        dev = DefaultQubit()
         weights = tf.Variable([0.543, -0.654], dtype=tf.float64)
 
         spy = mocker.spy(qml, "execute")
@@ -1207,7 +1207,7 @@ class TestSample:
     def test_sample_dimension(self):
         """Test sampling works as expected"""
 
-        @qnode(DefaultQubit2(), diff_method="parameter-shift", interface="tf")
+        @qnode(DefaultQubit(), diff_method="parameter-shift", interface="tf")
         def circuit():
             qml.Hadamard(wires=[0])
             qml.CNOT(wires=[0, 1])
@@ -1226,7 +1226,7 @@ def circuit():
     def test_sampling_expval(self):
         """Test sampling works as expected if combined with expectation values"""
 
-        @qnode(DefaultQubit2(), diff_method="parameter-shift", interface="tf")
+        @qnode(DefaultQubit(), diff_method="parameter-shift", interface="tf")
         def circuit():
             qml.Hadamard(wires=[0])
             qml.CNOT(wires=[0, 1])
@@ -1244,7 +1244,7 @@ def circuit():
     def test_sample_combination(self):
         """Test the output of combining expval, var and sample"""
 
-        @qnode(DefaultQubit2(), diff_method="parameter-shift", interface="tf")
+        @qnode(DefaultQubit(), diff_method="parameter-shift", interface="tf")
         def circuit():
             qml.RX(0.54, wires=0)
 
@@ -1266,7 +1266,7 @@ def circuit():
     def test_single_wire_sample(self):
         """Test the return type and shape of sampling a single wire"""
 
-        @qnode(DefaultQubit2(), diff_method="parameter-shift", interface="tf")
+        @qnode(DefaultQubit(), diff_method="parameter-shift", interface="tf")
         def circuit():
             qml.RX(0.54, wires=0)
 
@@ -1281,7 +1281,7 @@ def test_multi_wire_sample_regular_shape(self):
         """Test the return type and shape of sampling multiple wires
         where a rectangular array is expected"""
 
-        @qnode(DefaultQubit2(), diff_method="parameter-shift", interface="tf")
+        @qnode(DefaultQubit(), diff_method="parameter-shift", interface="tf")
         def circuit():
             return qml.sample(qml.PauliZ(0)), qml.sample(qml.PauliZ(1)), qml.sample(qml.PauliZ(2))
 
@@ -1297,7 +1297,7 @@ def test_counts(self):
         """Test counts works as expected for TF"""
 
         # pylint:disable=unsubscriptable-object,no-member
-        @qnode(DefaultQubit2(), interface="tf")
+        @qnode(DefaultQubit(), interface="tf")
         def circuit():
             qml.Hadamard(wires=[0])
             qml.CNOT(wires=[0, 1])
@@ -1337,7 +1337,7 @@ def test_autograph_gradients(self, decorator, interface, tol):
         y = tf.Variable(-0.654, dtype=tf.float64)
 
         @decorator
-        @qnode(DefaultQubit2(), diff_method="parameter-shift", interface=interface)
+        @qnode(DefaultQubit(), diff_method="parameter-shift", interface=interface)
         def circuit(x, y):
             qml.RX(x, wires=[0])
             qml.RY(y, wires=[1])
@@ -1362,7 +1362,7 @@ def test_autograph_jacobian(self, decorator, interface, tol):
         y = tf.Variable(-0.654, dtype=tf.float64)
 
         @decorator
-        @qnode(DefaultQubit2(), diff_method="parameter-shift", max_diff=1, interface=interface)
+        @qnode(DefaultQubit(), diff_method="parameter-shift", max_diff=1, interface=interface)
         def circuit(x, y):
             qml.RX(x, wires=[0])
             qml.RY(y, wires=[1])
@@ -1403,7 +1403,7 @@ def test_autograph_adjoint_single_param(self, grad_on_execution, decorator, inte
 
         @decorator
         @qnode(
-            DefaultQubit2(),
+            DefaultQubit(),
             diff_method="adjoint",
             interface=interface,
             grad_on_execution=grad_on_execution,
@@ -1431,7 +1431,7 @@ def test_autograph_adjoint_multi_params(self, grad_on_execution, decorator, inte
 
         @decorator
         @qnode(
-            DefaultQubit2(),
+            DefaultQubit(),
             diff_method="adjoint",
             interface=interface,
             grad_on_execution=grad_on_execution,
@@ -1461,7 +1461,7 @@ def test_autograph_ragged_differentiation(self, decorator, interface, tol):
         y = tf.Variable(-0.654, dtype=tf.float64)
 
         @decorator
-        @qnode(DefaultQubit2(), diff_method="parameter-shift", max_diff=1, interface=interface)
+        @qnode(DefaultQubit(), diff_method="parameter-shift", max_diff=1, interface=interface)
         def circuit(x, y):
             qml.RX(x, wires=[0])
             qml.RY(y, wires=[1])
@@ -1497,7 +1497,7 @@ def test_autograph_hessian(self, decorator, interface, tol):
         b = tf.Variable(-0.654, dtype=tf.float64)
 
         @decorator
-        @qnode(DefaultQubit2(), diff_method="parameter-shift", max_diff=2, interface=interface)
+        @qnode(DefaultQubit(), diff_method="parameter-shift", max_diff=2, interface=interface)
         def circuit(x, y):
             qml.RY(x, wires=0)
             qml.RX(y, wires=0)
@@ -1531,7 +1531,7 @@ def test_autograph_state(self, decorator, interface, tol):
 
         # TODO: fix this for diff_method=None
         @decorator
-        @qnode(DefaultQubit2(), diff_method="parameter-shift", interface=interface)
+        @qnode(DefaultQubit(), diff_method="parameter-shift", interface=interface)
         def circuit(x, y):
             qml.RX(x, wires=[0])
             qml.RY(y, wires=[1])
@@ -1550,7 +1550,7 @@ def test_autograph_dimension(self, decorator, interface):
         """Test sampling works as expected"""
 
         @decorator
-        @qnode(DefaultQubit2(), diff_method="parameter-shift", interface=interface)
+        @qnode(DefaultQubit(), diff_method="parameter-shift", interface=interface)
         def circuit(**_):
             qml.Hadamard(wires=[0])
             qml.CNOT(wires=[0, 1])
@@ -2124,7 +2124,7 @@ def test_no_ops():
     """Test that the return value of the QNode matches in the interface
     even if there are no ops"""
 
-    @qml.qnode(DefaultQubit2(), interface="tf")
+    @qml.qnode(DefaultQubit(), interface="tf")
     def circuit():
         qml.Hadamard(wires=0)
         return qml.state()
diff --git a/tests/interfaces/default_qubit_2_integration/test_tensorflow_qnode_shot_vector_default_qubit_2.py b/tests/interfaces/default_qubit_2_integration/test_tensorflow_qnode_shot_vector_default_qubit_2.py
index 8c5cbe743fe..d4290493958 100644
--- a/tests/interfaces/default_qubit_2_integration/test_tensorflow_qnode_shot_vector_default_qubit_2.py
+++ b/tests/interfaces/default_qubit_2_integration/test_tensorflow_qnode_shot_vector_default_qubit_2.py
@@ -18,7 +18,7 @@
 import pennylane as qml
 from pennylane import numpy as np
 from pennylane import qnode
-from pennylane.devices.experimental import DefaultQubit2
+from pennylane.devices import DefaultQubit
 
 pytestmark = pytest.mark.tf
 
@@ -28,9 +28,9 @@
 shots_and_num_copies_hess = [((10, (5, 1)), 2)]
 
 qubit_device_and_diff_method = [
-    [DefaultQubit2(), "finite-diff", {"h": 10e-2}],
-    [DefaultQubit2(), "parameter-shift", {}],
-    [DefaultQubit2(), "spsa", {"h": 10e-2, "num_directions": 20}],
+    [DefaultQubit(), "finite-diff", {"h": 10e-2}],
+    [DefaultQubit(), "parameter-shift", {}],
+    [DefaultQubit(), "spsa", {"h": 10e-2, "num_directions": 20}],
 ]
 
 TOLS = {
diff --git a/tests/interfaces/default_qubit_2_integration/test_torch_default_qubit_2.py b/tests/interfaces/default_qubit_2_integration/test_torch_default_qubit_2.py
index ea169699fc8..34d9f82505c 100644
--- a/tests/interfaces/default_qubit_2_integration/test_torch_default_qubit_2.py
+++ b/tests/interfaces/default_qubit_2_integration/test_torch_default_qubit_2.py
@@ -16,7 +16,7 @@
 import pytest
 
 import pennylane as qml
-from pennylane.devices.experimental import DefaultQubit2
+from pennylane.devices import DefaultQubit
 from pennylane.gradients import param_shift
 from pennylane.interfaces import execute
 
@@ -42,7 +42,7 @@ def test_caching_param_shift_hessian(self, num_params):
         """Test that, when using parameter-shift transform,
         caching reduces the number of evaluations to their optimum
         when computing Hessians."""
-        dev = DefaultQubit2()
+        dev = DefaultQubit()
         params = torch.arange(1, num_params + 1, requires_grad=True, dtype=torch.float64)
 
         N = len(params)
@@ -128,11 +128,11 @@ def cost_cache(x):
 # add tests for lightning 2 when possible
 # set rng for device when possible
 test_matrix = [
-    ({"gradient_fn": param_shift}, 100000, DefaultQubit2(seed=42)),
-    ({"gradient_fn": param_shift}, None, DefaultQubit2()),
-    ({"gradient_fn": "backprop"}, None, DefaultQubit2()),
-    ({"gradient_fn": "adjoint", "grad_on_execution": True}, None, DefaultQubit2()),
-    ({"gradient_fn": "adjoint", "grad_on_execution": False}, None, DefaultQubit2()),
+    ({"gradient_fn": param_shift}, 100000, DefaultQubit(seed=42)),
+    ({"gradient_fn": param_shift}, None, DefaultQubit()),
+    ({"gradient_fn": "backprop"}, None, DefaultQubit()),
+    ({"gradient_fn": "adjoint", "grad_on_execution": True}, None, DefaultQubit()),
+    ({"gradient_fn": "adjoint", "grad_on_execution": False}, None, DefaultQubit()),
 ]
 
 
@@ -591,7 +591,7 @@ class TestHigherOrderDerivatives:
     def test_parameter_shift_hessian(self, params, tol):
         """Tests that the output of the parameter-shift transform
         can be differentiated using torch, yielding second derivatives."""
-        dev = DefaultQubit2()
+        dev = DefaultQubit()
 
         def cost_fn(x):
             ops1 = [qml.RX(x[0], 0), qml.RY(x[1], 1), qml.CNOT((0, 1))]
@@ -625,7 +625,7 @@ def cost_fn(x):
     def test_max_diff(self, tol):
         """Test that setting the max_diff parameter blocks higher-order
         derivatives"""
-        dev = DefaultQubit2()
+        dev = DefaultQubit()
         params = torch.tensor([0.543, -0.654], requires_grad=True)
 
         def cost_fn(x):
diff --git a/tests/interfaces/default_qubit_2_integration/test_torch_qnode_default_qubit_2.py b/tests/interfaces/default_qubit_2_integration/test_torch_qnode_default_qubit_2.py
index 5085a858276..a3f23a6a67c 100644
--- a/tests/interfaces/default_qubit_2_integration/test_torch_qnode_default_qubit_2.py
+++ b/tests/interfaces/default_qubit_2_integration/test_torch_qnode_default_qubit_2.py
@@ -18,7 +18,7 @@
 
 import pennylane as qml
 from pennylane import qnode
-from pennylane.devices.experimental import DefaultQubit2
+from pennylane.devices import DefaultQubit
 
 pytestmark = pytest.mark.torch
 
@@ -27,13 +27,13 @@
 hessian = torch.autograd.functional.hessian
 
 qubit_device_and_diff_method = [
-    [DefaultQubit2(), "finite-diff", False],
-    [DefaultQubit2(), "parameter-shift", False],
-    [DefaultQubit2(), "backprop", True],
-    [DefaultQubit2(), "adjoint", True],
-    [DefaultQubit2(), "adjoint", False],
-    [DefaultQubit2(), "spsa", False],
-    [DefaultQubit2(), "hadamard", False],
+    [DefaultQubit(), "finite-diff", False],
+    [DefaultQubit(), "parameter-shift", False],
+    [DefaultQubit(), "backprop", True],
+    [DefaultQubit(), "adjoint", True],
+    [DefaultQubit(), "adjoint", False],
+    [DefaultQubit(), "spsa", False],
+    [DefaultQubit(), "hadamard", False],
 ]
 
 interface_and_qubit_device_and_diff_method = [
@@ -497,7 +497,7 @@ class TestShotsIntegration:
 
     def test_changing_shots(self):
         """Test that changing shots works on execution"""
-        dev = DefaultQubit2()
+        dev = DefaultQubit()
         a, b = torch.tensor([0.543, -0.654], requires_grad=True, dtype=torch.float64)
 
         @qnode(dev, interface="torch", diff_method=qml.gradients.param_shift)
@@ -520,7 +520,7 @@ def circuit(a, b):
     def test_gradient_integration(self):
         """Test that temporarily setting the shots works
         for gradient computations"""
-        dev = DefaultQubit2()
+        dev = DefaultQubit()
         a, b = torch.tensor([0.543, -0.654], requires_grad=True)
 
         @qnode(dev, interface="torch", diff_method=qml.gradients.param_shift)
@@ -545,7 +545,7 @@ def test_multiple_gradient_integration(self, tol):
         weights = torch.tensor([0.543, -0.654], requires_grad=True)
         a, b = weights
 
-        @qnode(DefaultQubit2(), interface="torch", diff_method=qml.gradients.param_shift)
+        @qnode(DefaultQubit(), interface="torch", diff_method=qml.gradients.param_shift)
         def circuit(a, b):
             qml.RY(a, wires=0)
             qml.RX(b, wires=1)
@@ -569,7 +569,7 @@ def test_update_diff_method(self, mocker):
 
         spy = mocker.spy(qml, "execute")
 
-        @qnode(DefaultQubit2(), interface="torch")
+        @qnode(DefaultQubit(), interface="torch")
         def cost_fn(a, b):
             qml.RY(a, wires=0)
             qml.RX(b, wires=1)
@@ -1338,7 +1338,7 @@ class TestSample:
     def test_sample_dimension(self):
         """Test sampling works as expected"""
 
-        @qnode(DefaultQubit2(), diff_method="parameter-shift", interface="torch")
+        @qnode(DefaultQubit(), diff_method="parameter-shift", interface="torch")
         def circuit():
             qml.Hadamard(wires=[0])
             qml.CNOT(wires=[0, 1])
@@ -1358,7 +1358,7 @@ def circuit():
     def test_sampling_expval(self):
         """Test sampling works as expected if combined with expectation values"""
 
-        @qnode(DefaultQubit2(), diff_method="parameter-shift", interface="torch")
+        @qnode(DefaultQubit(), diff_method="parameter-shift", interface="torch")
         def circuit():
             qml.Hadamard(wires=[0])
             qml.CNOT(wires=[0, 1])
@@ -1377,7 +1377,7 @@ def circuit():
     def test_counts_expval(self):
         """Test counts works as expected if combined with expectation values"""
 
-        @qnode(DefaultQubit2(), diff_method="parameter-shift", interface="torch")
+        @qnode(DefaultQubit(), diff_method="parameter-shift", interface="torch")
         def circuit():
             qml.Hadamard(wires=[0])
             qml.CNOT(wires=[0, 1])
@@ -1395,7 +1395,7 @@ def circuit():
     def test_sample_combination(self):
         """Test the output of combining expval, var and sample"""
 
-        @qnode(DefaultQubit2(), diff_method="parameter-shift", interface="torch")
+        @qnode(DefaultQubit(), diff_method="parameter-shift", interface="torch")
         def circuit():
             qml.RX(0.54, wires=0)
 
@@ -1416,7 +1416,7 @@ def circuit():
     def test_single_wire_sample(self):
         """Test the return type and shape of sampling a single wire"""
 
-        @qnode(DefaultQubit2(), diff_method="parameter-shift", interface="torch")
+        @qnode(DefaultQubit(), diff_method="parameter-shift", interface="torch")
         def circuit():
             qml.RX(0.54, wires=0)
             return qml.sample(qml.PauliZ(0))
@@ -1430,7 +1430,7 @@ def test_multi_wire_sample_regular_shape(self):
         """Test the return type and shape of sampling multiple wires
         where a rectangular array is expected"""
 
-        @qnode(DefaultQubit2(), diff_method="parameter-shift", interface="torch")
+        @qnode(DefaultQubit(), diff_method="parameter-shift", interface="torch")
         def circuit():
             return qml.sample(qml.PauliZ(0)), qml.sample(qml.PauliZ(1)), qml.sample(qml.PauliZ(2))
 
@@ -1447,12 +1447,12 @@ def circuit():
 
 
 qubit_device_and_diff_method_and_grad_on_execution = [
-    [DefaultQubit2(), "backprop", True],
-    [DefaultQubit2(), "finite-diff", False],
-    [DefaultQubit2(), "parameter-shift", False],
-    [DefaultQubit2(), "adjoint", True],
-    [DefaultQubit2(), "adjoint", False],
-    [DefaultQubit2(), "hadamard", False],
+    [DefaultQubit(), "backprop", True],
+    [DefaultQubit(), "finite-diff", False],
+    [DefaultQubit(), "parameter-shift", False],
+    [DefaultQubit(), "adjoint", True],
+    [DefaultQubit(), "adjoint", False],
+    [DefaultQubit(), "hadamard", False],
 ]
 
 
@@ -2171,7 +2171,7 @@ def test_no_ops():
     """Test that the return value of the QNode matches in the interface
     even if there are no ops"""
 
-    @qml.qnode(DefaultQubit2(), interface="torch")
+    @qml.qnode(DefaultQubit(), interface="torch")
     def circuit():
         qml.Hadamard(wires=0)
         return qml.state()
diff --git a/tests/interfaces/test_autograd.py b/tests/interfaces/test_autograd.py
index 6732eb9316b..6659f351c8c 100644
--- a/tests/interfaces/test_autograd.py
+++ b/tests/interfaces/test_autograd.py
@@ -21,7 +21,7 @@
 from pennylane.operation import Observable, AnyWires
 
 import pennylane as qml
-from pennylane.devices import DefaultQubit
+from pennylane.devices import DefaultQubitLegacy
 from pennylane.gradients import finite_diff, param_shift
 
 pytestmark = pytest.mark.autograd
@@ -149,8 +149,8 @@ def cost(a):
     def test_no_gradients_on_execution(self, mocker):
         """Test that no grad on execution uses the `device.batch_execute` and `device.gradients` pathway"""
         dev = qml.device("default.qubit.legacy", wires=1)
-        spy_execute = mocker.spy(qml.devices.DefaultQubit, "batch_execute")
-        spy_gradients = mocker.spy(qml.devices.DefaultQubit, "gradients")
+        spy_execute = mocker.spy(qml.devices.DefaultQubitLegacy, "batch_execute")
+        spy_gradients = mocker.spy(qml.devices.DefaultQubitLegacy, "gradients")
 
         def cost(a):
             with qml.queuing.AnnotatedQueue() as q:
@@ -1318,7 +1318,7 @@ class TestCustomJacobian:
     def test_custom_jacobians(self):
         """Test custom Jacobian device methood"""
 
-        class CustomJacobianDevice(DefaultQubit):
+        class CustomJacobianDevice(DefaultQubitLegacy):
             @classmethod
             def capabilities(cls):
                 capabilities = super().capabilities()
@@ -1347,7 +1347,7 @@ def test_custom_jacobians_param_shift(self):
         """Test computing the gradient using the parameter-shift
         rule with a device that provides a jacobian"""
 
-        class MyQubit(DefaultQubit):
+        class MyQubit(DefaultQubitLegacy):
             @classmethod
             def capabilities(cls):
                 capabilities = super().capabilities().copy()
@@ -1428,10 +1428,10 @@ def diagonalizing_gates(self):
         return []
 
 
-class DeviceSupportingSpecialObservable(DefaultQubit):
+class DeviceSupportingSpecialObservable(DefaultQubitLegacy):
     name = "Device supporting SpecialObservable"
     short_name = "default.qubit.specialobservable"
-    observables = DefaultQubit.observables.union({"SpecialObservable"})
+    observables = DefaultQubitLegacy.observables.union({"SpecialObservable"})
 
     @staticmethod
     def _asarray(arr, dtype=None):
diff --git a/tests/interfaces/test_jax.py b/tests/interfaces/test_jax.py
index 9f519b733fe..8d774e6984e 100644
--- a/tests/interfaces/test_jax.py
+++ b/tests/interfaces/test_jax.py
@@ -151,8 +151,8 @@ def cost(params):
     def test_no_grad_on_execution(self, mocker):
         """Test that no grad on execution uses the `device.batch_execute` and `device.gradients` pathway"""
         dev = qml.device("default.qubit.legacy", wires=1)
-        spy_execute = mocker.spy(qml.devices.DefaultQubit, "batch_execute")
-        spy_gradients = mocker.spy(qml.devices.DefaultQubit, "gradients")
+        spy_execute = mocker.spy(qml.devices.DefaultQubitLegacy, "batch_execute")
+        spy_gradients = mocker.spy(qml.devices.DefaultQubitLegacy, "gradients")
 
         def cost(a):
             with qml.queuing.AnnotatedQueue() as q:
diff --git a/tests/interfaces/test_jax_jit.py b/tests/interfaces/test_jax_jit.py
index 5e375f8e912..2c48a1b97ff 100644
--- a/tests/interfaces/test_jax_jit.py
+++ b/tests/interfaces/test_jax_jit.py
@@ -142,8 +142,8 @@ def cost(a):
     def test_no_gradients_on_execution(self, mocker):
         """Test that no grad on execution uses the `device.batch_execute` and `device.gradients` pathway"""
         dev = qml.device("default.qubit.legacy", wires=1)
-        spy_execute = mocker.spy(qml.devices.DefaultQubit, "batch_execute")
-        spy_gradients = mocker.spy(qml.devices.DefaultQubit, "gradients")
+        spy_execute = mocker.spy(qml.devices.DefaultQubitLegacy, "batch_execute")
+        spy_gradients = mocker.spy(qml.devices.DefaultQubitLegacy, "gradients")
 
         def cost(a):
             with qml.queuing.AnnotatedQueue() as q:
diff --git a/tests/interfaces/test_set_shots.py b/tests/interfaces/test_set_shots.py
index f5233b4f565..1a2374e6575 100644
--- a/tests/interfaces/test_set_shots.py
+++ b/tests/interfaces/test_set_shots.py
@@ -26,7 +26,7 @@ def test_shots_new_device_interface():
     """Test that calling set_shots on a device implementing the new interface leaves it
     untouched.
     """
-    dev = qml.devices.experimental.DefaultQubit2()
+    dev = qml.devices.DefaultQubit()
     with pytest.raises(ValueError):
         with set_shots(dev, 10):
             pass
@@ -35,7 +35,7 @@ def test_shots_new_device_interface():
 def test_set_with_shots_class():
     """Test that shots can be set on the old device interface with a Shots class."""
 
-    dev = qml.devices.DefaultQubit(wires=1)
+    dev = qml.devices.DefaultQubitLegacy(wires=1)
     with set_shots(dev, Shots(10)):
         assert dev.shots == 10
 
@@ -53,6 +53,6 @@ def test_shots_not_altered_if_False():
     """Test a value of False can be passed to shots, indicating to not override
     shots on the device."""
 
-    dev = qml.devices.DefaultQubit(wires=1)
+    dev = qml.devices.DefaultQubitLegacy(wires=1)
     with set_shots(dev, False):
         assert dev.shots is None
diff --git a/tests/interfaces/test_tensorflow.py b/tests/interfaces/test_tensorflow.py
index d0ce671989c..45b7ac441f9 100644
--- a/tests/interfaces/test_tensorflow.py
+++ b/tests/interfaces/test_tensorflow.py
@@ -103,8 +103,8 @@ def test_grad_on_execution(self, mocker):
     def test_no_grad_execution(self, mocker):
         """Test that no grad on execution uses the `device.batch_execute` and `device.gradients` pathway"""
         dev = qml.device("default.qubit.legacy", wires=1)
-        spy_execute = mocker.spy(qml.devices.DefaultQubit, "batch_execute")
-        spy_gradients = mocker.spy(qml.devices.DefaultQubit, "gradients")
+        spy_execute = mocker.spy(qml.devices.DefaultQubitLegacy, "batch_execute")
+        spy_gradients = mocker.spy(qml.devices.DefaultQubitLegacy, "gradients")
         a = tf.Variable([0.1, 0.2])
 
         with tf.GradientTape() as t:
diff --git a/tests/interfaces/test_torch.py b/tests/interfaces/test_torch.py
index 5a2a0a93f84..2270182307d 100644
--- a/tests/interfaces/test_torch.py
+++ b/tests/interfaces/test_torch.py
@@ -135,8 +135,8 @@ def test_grad_on_execution(self, interface, mocker):
     def test_no_grad_on_execution(self, interface, mocker):
         """Test that no grad on execution uses the `device.batch_execute` and `device.gradients` pathway"""
         dev = qml.device("default.qubit.legacy", wires=1)
-        spy_execute = mocker.spy(qml.devices.DefaultQubit, "batch_execute")
-        spy_gradients = mocker.spy(qml.devices.DefaultQubit, "gradients")
+        spy_execute = mocker.spy(qml.devices.DefaultQubitLegacy, "batch_execute")
+        spy_gradients = mocker.spy(qml.devices.DefaultQubitLegacy, "gradients")
 
         a = torch.tensor([0.1, 0.2], requires_grad=True)
 
diff --git a/tests/interfaces/test_transform_program_integration.py b/tests/interfaces/test_transform_program_integration.py
index 3708c52552b..2ac697bcafe 100644
--- a/tests/interfaces/test_transform_program_integration.py
+++ b/tests/interfaces/test_transform_program_integration.py
@@ -26,7 +26,7 @@
 
 device_suite = (
     qml.device("default.qubit.legacy", wires=5),
-    qml.devices.experimental.DefaultQubit2(),
+    qml.devices.DefaultQubit(),
     qml.device("lightning.qubit", wires=5),
 )
 
@@ -39,7 +39,7 @@ class TestTransformProgram:
     def test_transform_program_none(self, interface):
         """Test that if no transform program is provided, null default behavior is used."""
 
-        dev = qml.devices.experimental.DefaultQubit2()
+        dev = qml.devices.DefaultQubit()
 
         tape0 = qml.tape.QuantumScript([qml.RX(1.0, 0)], [qml.expval(qml.PauliZ(0))])
         tape1 = qml.tape.QuantumScript([qml.RY(2.0, 0)], [qml.state()])
@@ -59,7 +59,7 @@ def test_transform_program_none(self, interface):
     def test_transform_program_modifies_circuit(self, interface):
         """Integration tests for a transform program that modifies the input tapes."""
 
-        dev = qml.devices.experimental.DefaultQubit2()
+        dev = qml.devices.DefaultQubit()
 
         def null_postprocessing(results):
             return results[0]
@@ -102,7 +102,7 @@ def test_shot_distributing_transform(self, interface):
 
         Note that this only works with the new device interface.
         """
-        dev = qml.devices.experimental.DefaultQubit2()
+        dev = qml.devices.DefaultQubit()
 
         def null_postprocessing(results):
             return results
diff --git a/tests/logging/test_logging_autograd.py b/tests/logging/test_logging_autograd.py
index 3ea8adfddc7..e278e674c1b 100644
--- a/tests/logging/test_logging_autograd.py
+++ b/tests/logging/test_logging_autograd.py
@@ -88,7 +88,7 @@ def circuit(params):
             (
                 "pennylane.interfaces.execution",
                 [
-                    "device=<DefaultQubit device (wires=2, shots=None)",
+                    "device=<DefaultQubitLegacy device (wires=2, shots=None)",
                     "gradient_fn=None, interface=None",
                 ],
             ),
@@ -129,7 +129,7 @@ def circuit(params):
                 "pennylane.qnode",
                 [
                     "Creating QNode(func=<function TestLogging.test_dq_qnode_execution_grad",
-                    "device=<DefaultQubit device (wires=2, shots=None)",
+                    "device=<DefaultQubitLegacy device (wires=2, shots=None)",
                     f"interface=auto, diff_method={diff_method[0]}, expansion_strategy=gradient, max_expansion=10, grad_on_execution=best,",
                 ],
             ),
diff --git a/tests/measurements/legacy/test_state_legacy.py b/tests/measurements/legacy/test_state_legacy.py
index a21cacfc9dd..b53ad74cd1b 100644
--- a/tests/measurements/legacy/test_state_legacy.py
+++ b/tests/measurements/legacy/test_state_legacy.py
@@ -17,7 +17,7 @@
 
 import pennylane as qml
 from pennylane import numpy as pnp
-from pennylane.devices import DefaultQubit
+from pennylane.devices import DefaultQubitLegacy
 from pennylane.measurements import (
     State,
     Shots,
@@ -192,7 +192,7 @@ def func():
             return state()
 
         with monkeypatch.context() as m:
-            m.setattr(DefaultQubit, "capabilities", lambda *args, **kwargs: capabilities)
+            m.setattr(DefaultQubitLegacy, "capabilities", lambda *args, **kwargs: capabilities)
             with pytest.raises(qml.QuantumFunctionError, match="The current device is not capable"):
                 func()
 
@@ -717,7 +717,7 @@ def func():
             return density_matrix(0)
 
         with monkeypatch.context() as m:
-            m.setattr(DefaultQubit, "capabilities", lambda *args, **kwargs: capabilities)
+            m.setattr(DefaultQubitLegacy, "capabilities", lambda *args, **kwargs: capabilities)
             with pytest.raises(
                 qml.QuantumFunctionError,
                 match="The current device is not capable" " of returning the state",
diff --git a/tests/measurements/test_state.py b/tests/measurements/test_state.py
index bf3b21b71e5..9aca8423e25 100644
--- a/tests/measurements/test_state.py
+++ b/tests/measurements/test_state.py
@@ -17,7 +17,7 @@
 
 import pennylane as qml
 from pennylane import numpy as pnp
-from pennylane.devices import DefaultQubit
+from pennylane.devices import DefaultQubitLegacy
 from pennylane.measurements import (
     State,
     StateMP,
@@ -399,7 +399,7 @@ def func():
             return state()
 
         with monkeypatch.context() as m:
-            m.setattr(DefaultQubit, "capabilities", lambda *args, **kwargs: capabilities)
+            m.setattr(DefaultQubitLegacy, "capabilities", lambda *args, **kwargs: capabilities)
             with pytest.raises(qml.QuantumFunctionError, match="The current device is not capable"):
                 func()
 
@@ -1050,7 +1050,7 @@ def func():
             return density_matrix(0)
 
         with monkeypatch.context() as m:
-            m.setattr(DefaultQubit, "capabilities", lambda *args, **kwargs: capabilities)
+            m.setattr(DefaultQubitLegacy, "capabilities", lambda *args, **kwargs: capabilities)
             with pytest.raises(
                 qml.QuantumFunctionError,
                 match="The current device is not capable" " of returning the state",
diff --git a/tests/test_debugging.py b/tests/test_debugging.py
index 42f59ca5f8e..38769399e42 100644
--- a/tests/test_debugging.py
+++ b/tests/test_debugging.py
@@ -57,7 +57,7 @@ def circuit():
     def test_default_qubit2(self):
         """Test that multiple snapshots are returned correctly on the new
         state-vector simulator."""
-        dev = qml.devices.experimental.DefaultQubit2()
+        dev = qml.devices.DefaultQubit()
 
         # TODO: add additional QNode test once the new device supports it
 
@@ -217,9 +217,7 @@ def circuit():
     def test_unsupported_device_new(self):
         """Test that an error is raised on unsupported devices."""
 
-        class DummyDevice(
-            qml.devices.experimental.Device
-        ):  # pylint: disable=too-few-public-methods
+        class DummyDevice(qml.devices.Device):  # pylint: disable=too-few-public-methods
             def execute(self, *args, **kwargs):
                 return args, kwargs
 
diff --git a/tests/test_qnode.py b/tests/test_qnode.py
index 65a7c95b784..d23aa1cbb99 100644
--- a/tests/test_qnode.py
+++ b/tests/test_qnode.py
@@ -26,7 +26,6 @@
 
 import pennylane as qml
 from pennylane import QNode
-from pennylane.devices import experimental
 from pennylane import numpy as pnp
 from pennylane import qnode
 from pennylane.tape import QuantumScript
@@ -242,7 +241,7 @@ def capabilities(cls):
             capabilities.update(model="None")
             return capabilities
 
-        monkeypatch.setattr(qml.devices.DefaultQubit, "capabilities", capabilities)
+        monkeypatch.setattr(qml.devices.DefaultQubitLegacy, "capabilities", capabilities)
         dev = qml.device("default.qubit", wires=1)
 
         with pytest.raises(
@@ -306,7 +305,7 @@ def capabilities(cls):
             return capabilities
 
         # finite differences is the fallback when we know nothing about the device
-        monkeypatch.setattr(qml.devices.DefaultQubit, "capabilities", capabilities)
+        monkeypatch.setattr(qml.devices.DefaultQubitLegacy, "capabilities", capabilities)
         res = QNode.get_best_method(dev, "another_interface")
         assert res == (qml.gradients.finite_diff, {}, dev)
 
@@ -365,7 +364,7 @@ def capabilities(cls):
             return capabilities
 
         # finite differences is the fallback when we know nothing about the device
-        monkeypatch.setattr(qml.devices.DefaultQubit, "capabilities", capabilities)
+        monkeypatch.setattr(qml.devices.DefaultQubitLegacy, "capabilities", capabilities)
         res = QNode.best_method_str(dev, "another_interface")
         assert res == "finite-diff"
 
@@ -822,7 +821,7 @@ def circuit1(x):
             qml.GlobalPhase(x, wires=0)
             return qml.expval(qml.PauliZ(0) @ qml.PauliZ(1))
 
-        dev = qml.devices.experimental.DefaultQubit2()  # TODO: add wires, change comment below
+        dev = qml.devices.DefaultQubit()  # TODO: add wires, change comment below
         qn = QNode(circuit1, dev)
 
         # fails when GlobalPhase is a strict subset of all tape wires
@@ -1710,7 +1709,7 @@ def test_scaling_shots_transform(self):
         """Test a transform that scales the number of shots used in an execution."""
 
         # note that this won't work with the old device interface :(
-        dev = qml.devices.experimental.DefaultQubit2()
+        dev = qml.devices.DefaultQubit()
 
         def num_of_shots_from_sample(results):
             return len(results[0])
@@ -1730,7 +1729,7 @@ def circuit():
 
 
 # pylint: disable=unused-argument
-class CustomDevice(experimental.Device):
+class CustomDevice(qml.devices.Device):
     """A null device that just returns 0."""
 
     def __repr__(self):
@@ -1741,7 +1740,7 @@ def execute(self, circuits, execution_config=None):
 
 
 class TestNewDeviceIntegration:
-    """Basic tests for integration of the experimental device interface and the QNode."""
+    """Basic tests for integration of the new device interface and the QNode."""
 
     dev = CustomDevice()
 
@@ -1772,7 +1771,7 @@ def test_get_gradient_fn_custom_device(self):
 
     def test_get_gradient_fn_default_qubit2(self):
         """Tests the get_gradient_fn is backprop for best for default qubit2."""
-        dev = experimental.DefaultQubit2()
+        dev = qml.devices.DefaultQubit()
         gradient_fn, kwargs, new_dev = QNode.get_gradient_fn(dev, "autograd", "best")
         assert gradient_fn == "backprop"
         assert not kwargs
@@ -1780,7 +1779,7 @@ def test_get_gradient_fn_default_qubit2(self):
 
     def test_get_gradient_fn_default_qubit2_adjoint(self):
         """Test that the get_gradient_fn and _validate_adjoint_methods work for default qubit 2."""
-        dev = experimental.DefaultQubit2()
+        dev = qml.devices.DefaultQubit()
         gradient_fn, kwargs, new_dev = QNode.get_gradient_fn(dev, "autograd", "adjoint")
         assert gradient_fn == "adjoint"
         assert len(kwargs) == 0
@@ -1815,7 +1814,7 @@ def test_custom_device_that_supports_backprop(self):
         """Test that a custom device and designate that it supports backprop derivatives."""
 
         # pylint: disable=unused-argument
-        class BackpropDevice(experimental.Device):
+        class BackpropDevice(qml.devices.Device):
             """A device that says it supports backpropagation."""
 
             def execute(self, circuits, execution_config=None):
@@ -1834,7 +1833,7 @@ def test_custom_device_with_device_derivative(self):
         """Test that a custom device can specify that it supports device derivatives."""
 
         # pylint: disable=unused-argument
-        class DerivativeDevice(experimental.Device):
+        class DerivativeDevice(qml.devices.Device):
             """A device that says it supports device derivatives."""
 
             def execute(self, circuits, execution_config=None):
@@ -1857,7 +1856,7 @@ def supports_derivatives(self, execution_config=None, circuit=None) -> bool:
     def test_shots_integration(self):
         """Test that shots provided at call time are passed through the workflow."""
 
-        dev = experimental.DefaultQubit2()
+        dev = qml.devices.DefaultQubit()
 
         @qml.qnode(dev, diff_method=None)
         def circuit():
diff --git a/tests/test_return_types_qnode.py b/tests/test_return_types_qnode.py
index 049833137f4..d4e2c11b86d 100644
--- a/tests/test_return_types_qnode.py
+++ b/tests/test_return_types_qnode.py
@@ -261,7 +261,7 @@ def test_sample(self, measurement, device, shots=100):
             if isinstance(measurement.obs, qml.PauliZ):
                 pytest.skip("DefaultQutrit doesn't support qubit observables.")
         elif isinstance(measurement.obs, qml.GellMann):
-            pytest.skip("DefaultQubit doesn't support qutrit observables.")
+            pytest.skip("DefaultQubitLegacy doesn't support qutrit observables.")
 
         dev = qml.device(device, wires=2, shots=shots)
         func = qutrit_ansatz if device == "default.qutrit" else qubit_ansatz
@@ -296,7 +296,7 @@ def test_counts(self, measurement, device, shots=100):
             if isinstance(measurement.obs, qml.PauliZ):
                 pytest.skip("DefaultQutrit doesn't support qubit observables.")
         elif isinstance(measurement.obs, qml.GellMann):
-            pytest.skip("DefaultQubit doesn't support qutrit observables.")
+            pytest.skip("DefaultQubitLegacy doesn't support qutrit observables.")
 
         dev = qml.device(device, wires=2, shots=shots)
         func = qutrit_ansatz if device == "default.qutrit" else qubit_ansatz
@@ -1234,7 +1234,7 @@ def test_expval_sample(self, measurement, device, shots=100):
             if isinstance(measurement.obs, qml.PauliZ):
                 pytest.skip("DefaultQutrit doesn't support qubit observables.")
         elif isinstance(measurement.obs, qml.GellMann):
-            pytest.skip("DefaultQubit doesn't support qutrit observables.")
+            pytest.skip("DefaultQubitLegacy doesn't support qutrit observables.")
 
         dev = qml.device(device, wires=2, shots=shots)
         func = qubit_ansatz if device != "default.qutrit" else qutrit_ansatz
@@ -1266,7 +1266,7 @@ def test_expval_counts(self, measurement, device, shots=100):
             if isinstance(measurement.obs, qml.PauliZ):
                 pytest.skip("DefaultQutrit doesn't support qubit observables.")
         elif isinstance(measurement.obs, qml.GellMann):
-            pytest.skip("DefaultQubit doesn't support qutrit observables.")
+            pytest.skip("DefaultQubitLegacy doesn't support qutrit observables.")
 
         dev = qml.device(device, wires=2, shots=shots)
         func = qubit_ansatz if device != "default.qutrit" else qutrit_ansatz
diff --git a/tests/transforms/test_adjoint_metric_tensor.py b/tests/transforms/test_adjoint_metric_tensor.py
index f4c37be342f..f79572cd074 100644
--- a/tests/transforms/test_adjoint_metric_tensor.py
+++ b/tests/transforms/test_adjoint_metric_tensor.py
@@ -261,7 +261,7 @@ def test_correct_output_tape_autograd(self, ansatz, params, interface):
         """Test that the output is correct when using Autograd and
         calling the adjoint metric tensor directly on a tape."""
         expected = autodiff_metric_tensor(ansatz, 3)(*params)
-        dev = qml.devices.experimental.DefaultQubit2()
+        dev = qml.devices.DefaultQubit()
 
         wires = ("a", "b", "c")
 
diff --git a/tests/transforms/test_optimization/test_cancel_inverses.py b/tests/transforms/test_optimization/test_cancel_inverses.py
index 14285fc7e83..10251c6ecf2 100644
--- a/tests/transforms/test_optimization/test_cancel_inverses.py
+++ b/tests/transforms/test_optimization/test_cancel_inverses.py
@@ -367,7 +367,7 @@ def qfunc_circuit(theta):
 
 
 ### QNode
-dev = qml.devices.experimental.DefaultQubit2()
+dev = qml.devices.DefaultQubit()
 
 
 @qml.qnode(device=dev)
diff --git a/tests/transforms/test_optimization/test_commute_controlled.py b/tests/transforms/test_optimization/test_commute_controlled.py
index 3327a27fcd3..4fd0c9e4b13 100644
--- a/tests/transforms/test_optimization/test_commute_controlled.py
+++ b/tests/transforms/test_optimization/test_commute_controlled.py
@@ -470,7 +470,7 @@ def qfunc_circuit():
 
 
 ### QNode
-dev = qml.devices.experimental.DefaultQubit2()
+dev = qml.devices.DefaultQubit()
 
 
 @qml.qnode(device=dev)
diff --git a/tests/transforms/test_optimization/test_merge_rotations.py b/tests/transforms/test_optimization/test_merge_rotations.py
index 18784f1bf0e..e941ab91444 100644
--- a/tests/transforms/test_optimization/test_merge_rotations.py
+++ b/tests/transforms/test_optimization/test_merge_rotations.py
@@ -469,7 +469,7 @@ def qfunc_circuit(theta):
 
 
 ### QNode
-dev = qml.devices.experimental.DefaultQubit2()
+dev = qml.devices.DefaultQubit()
 
 
 @qml.qnode(device=dev)
diff --git a/tests/transforms/test_qcut.py b/tests/transforms/test_qcut.py
index 092924907d3..e7b92977037 100644
--- a/tests/transforms/test_qcut.py
+++ b/tests/transforms/test_qcut.py
@@ -2457,9 +2457,7 @@ def func(x):
             )
 
 
-@pytest.mark.parametrize(
-    "dev_fn", [qml.devices.DefaultQubit, qml.devices.experimental.DefaultQubit2]
-)
+@pytest.mark.parametrize("dev_fn", [qml.devices.DefaultQubitLegacy, qml.devices.DefaultQubit])
 class TestCutCircuitMCTransform:
     """
     Tests that the `cut_circuit_mc` transform gives the correct results.

From c9c9fe31f71b59e1cb6654dbd126da554b194208 Mon Sep 17 00:00:00 2001
From: Matthew Silverman <matthews@xanadu.ai>
Date: Fri, 15 Sep 2023 11:02:32 -0400
Subject: [PATCH 058/131] get op batch size at runtime to match DQL (#4592)

**Context:**
Found this issue while checking if KerasLayer works with DQ2 yet
(spoiler alert - it does not).

We tried to be more elegant with DQ2 regarding batch sizes, and that was
good and all. The only problem is, `Operator._check_batching` has a
whole bunch of problems with TF-Autograph (just check [the inline
comments](https://github.com/PennyLaneAI/pennylane/blob/2d1f6a395fde6bf6a9bbe1aeba6945882760704c/pennylane/operation.py#L1073-L1090)).
The whole thing is just a mess, so after discussion with @rmoyard we
felt it would be best to mimic DQL's behaviour here.

tl;dr TF-Autograph has taught me that we _should_ check batch sizes at
run-time, not "compile" (aka preprocess)-time.

**Description of the Change:**
Check operator batch size at runtime during `apply_operation`.

**Benefits:**
It needs to be done to make `KerasLayer` work

**Possible Drawbacks:**
- This will slow things down fairly severely. Perhaps we can make a
smarter move in the future, but I just want something that works for
now. To be clear, this performance change puts us at parity with DQL
- I checked locally, and [one KerasLayer
test](https://github.com/PennyLaneAI/pennylane/blob/2d1f6a395fde6bf6a9bbe1aeba6945882760704c/tests/qnn/test_keras.py#L641)
is still failing with DQ2 whenever `output_dim` > 1 because for some
reason, when re-loading the QNN from file it assumes the abstract batch
size (`None`) is `1`, but it should be `None` (so Autograph can decide
on the value at runtime)

---------

Co-authored-by: Romain Moyard <rmoyard@gmail.com>
---
 pennylane/devices/qubit/apply_operation.py  | 16 +++++++++---
 pennylane/devices/qubit/measure.py          | 22 ++++++++++++++---
 pennylane/math/__init__.py                  |  2 +-
 pennylane/math/matrix_manipulation.py       | 27 +++++++++++++++++++++
 tests/devices/qubit/test_apply_operation.py | 11 +++++++++
 5 files changed, 69 insertions(+), 9 deletions(-)

diff --git a/pennylane/devices/qubit/apply_operation.py b/pennylane/devices/qubit/apply_operation.py
index 9083384ac81..8ec8035a321 100644
--- a/pennylane/devices/qubit/apply_operation.py
+++ b/pennylane/devices/qubit/apply_operation.py
@@ -87,9 +87,13 @@ def apply_operation_einsum(op: qml.operation.Operator, state, is_state_batched:
     )
 
     new_mat_shape = [2] * (num_indices * 2)
-    if op.batch_size is not None:
+    dim = 2**num_indices
+    batch_size = math.get_batch_size(mat, (dim, dim), dim**2)
+    if batch_size is not None:
         # Add broadcasting dimension to shape
-        new_mat_shape = [mat.shape[0]] + new_mat_shape
+        new_mat_shape = [batch_size] + new_mat_shape
+        if op.batch_size is None:
+            op._batch_size = batch_size  # pylint:disable=protected-access
     reshaped_mat = math.reshape(mat, new_mat_shape)
 
     return math.einsum(einsum_indices, reshaped_mat, state)
@@ -112,9 +116,13 @@ def apply_operation_tensordot(op: qml.operation.Operator, state, is_state_batche
     num_indices = len(op.wires)
 
     new_mat_shape = [2] * (num_indices * 2)
-    if is_mat_batched := op.batch_size is not None:
+    dim = 2**num_indices
+    batch_size = math.get_batch_size(mat, (dim, dim), dim**2)
+    if is_mat_batched := batch_size is not None:
         # Add broadcasting dimension to shape
-        new_mat_shape = [mat.shape[0]] + new_mat_shape
+        new_mat_shape = [batch_size] + new_mat_shape
+        if op.batch_size is None:
+            op._batch_size = batch_size  # pylint:disable=protected-access
     reshaped_mat = math.reshape(mat, new_mat_shape)
 
     mat_axes = list(range(-num_indices, 0))
diff --git a/pennylane/devices/qubit/measure.py b/pennylane/devices/qubit/measure.py
index 65aed196dd8..733f2b6f074 100644
--- a/pennylane/devices/qubit/measure.py
+++ b/pennylane/devices/qubit/measure.py
@@ -32,6 +32,23 @@
 from .apply_operation import apply_operation
 
 
+def flatten_state(state, num_wires):
+    """
+    Given a non-flat, potentially batched state, flatten it.
+
+    Args:
+        state (TensorLike): A state that needs flattening
+        num_wires (int): The number of wires the state represents
+
+    Returns:
+        A flat state, with an extra batch dimension if necessary
+    """
+    dim = 2**num_wires
+    batch_size = math.get_batch_size(state, (2,) * num_wires, dim)
+    shape = (batch_size, dim) if batch_size is not None else (dim,)
+    return math.reshape(state, shape)
+
+
 def state_diagonalizing_gates(
     measurementprocess: StateMeasurement, state: TensorLike, is_state_batched: bool = False
 ) -> TensorLike:
@@ -50,10 +67,7 @@ def state_diagonalizing_gates(
 
     total_indices = len(state.shape) - is_state_batched
     wires = Wires(range(total_indices))
-
-    flattened_state = (
-        math.reshape(state, (state.shape[0], -1)) if is_state_batched else math.flatten(state)
-    )
+    flattened_state = flatten_state(state, total_indices)
     return measurementprocess.process_state(flattened_state, wires)
 
 
diff --git a/pennylane/math/__init__.py b/pennylane/math/__init__.py
index 93b822f8289..9e8ec32837d 100644
--- a/pennylane/math/__init__.py
+++ b/pennylane/math/__init__.py
@@ -34,7 +34,7 @@
 import autoray as ar
 
 from .is_independent import is_independent
-from .matrix_manipulation import expand_matrix, reduce_matrices
+from .matrix_manipulation import expand_matrix, reduce_matrices, get_batch_size
 from .multi_dispatch import (
     add,
     array,
diff --git a/pennylane/math/matrix_manipulation.py b/pennylane/math/matrix_manipulation.py
index 4d3fb5d8b29..99a2ec811d5 100644
--- a/pennylane/math/matrix_manipulation.py
+++ b/pennylane/math/matrix_manipulation.py
@@ -306,3 +306,30 @@ def expand_and_reduce(op1_tuple: Tuple[np.ndarray, Wires], op2_tuple: Tuple[np.n
     reduced_mat, final_wires = reduce(expand_and_reduce, mats_and_wires_gen)
 
     return reduced_mat, final_wires
+
+
+def get_batch_size(tensor, expected_shape, expected_size):
+    """
+    Determine whether a tensor has an additional batch dimension for broadcasting,
+    compared to an expected_shape. Has support for abstract TF tensors.
+
+    Args:
+        tensor (TensorLike): A tensor to inspect for batching
+        expected_shape (Tuple[int]): The expected shape of the tensor if not batched
+        expected_size (int): The expected size of the tensor if not batched
+
+    Returns:
+        Optional[int]: The batch size of the tensor if there is one, otherwise None
+    """
+    try:
+        size = qml.math.size(tensor)
+        ndim = qml.math.ndim(tensor)
+        if ndim > len(expected_shape) or size > expected_size:
+            return size // expected_size
+
+    except Exception as err:  # pragma: no cover, pylint:disable=broad-except
+        # This except clause covers the usage of tf.function
+        if not qml.math.is_abstract(tensor):
+            raise err
+
+    return None
diff --git a/tests/devices/qubit/test_apply_operation.py b/tests/devices/qubit/test_apply_operation.py
index c88188c1688..1c332bb61c9 100644
--- a/tests/devices/qubit/test_apply_operation.py
+++ b/tests/devices/qubit/test_apply_operation.py
@@ -502,6 +502,17 @@ def test_broadcasted_op_broadcasted_state(self, op, method, ml_framework):
         assert qml.math.get_interface(res) == ml_framework
         assert qml.math.allclose(res, expected)
 
+    def test_batch_size_set_if_missing(self, method, ml_framework):
+        """Tests that the batch_size is set on an operator if it was missing before.
+        Mostly useful for TF-autograph since it may have batch size set to None."""
+        param = qml.math.asarray([0.1, 0.2, 0.3], like=ml_framework)
+        state = np.ones((2, 2)) / 2
+        op = qml.RX(param, 0)
+        op._batch_size = None  # pylint:disable=protected-access
+        state = method(op, state)
+        assert state.shape == (3, 2, 2)
+        assert op.batch_size == 3
+
 
 @pytest.mark.parametrize("method", methods)
 class TestLargerOperations:

From 115b7bd3b4cc98a963988f102538f5b28f79b1da Mon Sep 17 00:00:00 2001
From: Mudit Pandey <mudit.pandey@xanadu.ai>
Date: Fri, 15 Sep 2023 11:23:30 -0400
Subject: [PATCH 059/131] Removed unnecessary checks

---
 pennylane/_device.py       | 3 +--
 pennylane/_qubit_device.py | 2 +-
 2 files changed, 2 insertions(+), 3 deletions(-)

diff --git a/pennylane/_device.py b/pennylane/_device.py
index b3552f8f49a..bb46b77bf04 100644
--- a/pennylane/_device.py
+++ b/pennylane/_device.py
@@ -30,7 +30,6 @@
     CountsMP,
     Expectation,
     ExpectationMP,
-    MeasurementValue,
     MidMeasureMP,
     Probability,
     ProbabilityMP,
@@ -997,7 +996,7 @@ def check_validity(self, queue, observables):
                         raise DeviceError(
                             f"Observable {i.name} not supported on device {self.short_name}"
                         )
-            elif not isinstance(o, MeasurementValue):
+            else:
                 observable_name = o.name
 
                 if not self.supports_observable(observable_name):
diff --git a/pennylane/_qubit_device.py b/pennylane/_qubit_device.py
index cb6713cdac4..a9a1972b881 100644
--- a/pennylane/_qubit_device.py
+++ b/pennylane/_qubit_device.py
@@ -1480,7 +1480,7 @@ def sample(self, observable, shot_range=None, bin_size=None, counts=False):
 
         # translate to wire labels used by device
         device_wires = self.map_wires(observable.wires)
-        name = observable.name if not isinstance(observable, MeasurementValue) else None
+        name = observable.name
         # Select the samples from self._samples that correspond to ``shot_range`` if provided
         if shot_range is None:
             sub_samples = self._samples

From 6a31c3420ae4450c64749534d14cd1b57a60ec13 Mon Sep 17 00:00:00 2001
From: Mudit Pandey <mudit.pandey@xanadu.ai>
Date: Fri, 15 Sep 2023 17:03:48 -0400
Subject: [PATCH 060/131] Adding postselection support

---
 pennylane/devices/qubit/simulate.py        | 20 +++++++++++++++++++-
 pennylane/measurements/mid_measure.py      | 17 ++++++++++++++---
 pennylane/ops/qubit/observables.py         |  4 ++--
 pennylane/transforms/defer_measurements.py | 12 +++++++++++-
 4 files changed, 46 insertions(+), 7 deletions(-)

diff --git a/pennylane/devices/qubit/simulate.py b/pennylane/devices/qubit/simulate.py
index 55ba333a5df..3847d749555 100644
--- a/pennylane/devices/qubit/simulate.py
+++ b/pennylane/devices/qubit/simulate.py
@@ -13,7 +13,7 @@
 # limitations under the License.
 """Simulate a quantum script."""
 # pylint: disable=protected-access
-from numpy.random import default_rng
+from numpy.random import default_rng, binomial
 
 import pennylane as qml
 from pennylane.typing import Result
@@ -115,6 +115,24 @@ 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
 
+        if isinstance(op, qml.Projector):
+            # Handle postselection on mid-circuit measurements
+            if is_state_batched:
+                for i in qml.math.shape(state)[0]:
+                    state[i] = state[i] / qml.math.norm(state[i])
+            else:
+                norm = qml.math.norm(state)
+                if qml.math.isclose(norm, 0.0):
+                    raise ValueError("Requested postselection on state with 0 probability.")
+
+                state = state / norm
+
+                if circuit.shots:
+                    # defer_measurements will raise an error with batched shots or broadcasting
+                    # so we can assume that both the state and shots are unbatched.
+                    postselected_shots = binomial(circuit.shots.total_shots, norm)
+                    circuit._shots = qml.measurements.Shots(postselected_shots)
+
     if set(circuit.op_wires) < set(circuit.wires):
         state = expand_state_over_wires(
             state,
diff --git a/pennylane/measurements/mid_measure.py b/pennylane/measurements/mid_measure.py
index 6bbb90a14b8..691ba6f9b68 100644
--- a/pennylane/measurements/mid_measure.py
+++ b/pennylane/measurements/mid_measure.py
@@ -24,7 +24,7 @@
 from .measurements import MeasurementProcess, MidMeasure
 
 
-def measure(wires: Wires, reset: Optional[bool] = False):
+def measure(wires: Wires, reset: Optional[bool] = False, postselect: Optional[int] = None):
     r"""Perform a mid-circuit measurement in the computational basis on the
     supplied qubit.
 
@@ -86,6 +86,9 @@ def func():
         wires (Wires): The wire of the qubit the measurement process applies to.
         reset (Optional[bool]): Whether to reset the wire to the :math:`|0 \rangle`
             state after measurement.
+        postselect (Optional[int]): Which basis state to postselect after a mid-circuit
+            measurement. None by default. If postselection is requested, only the post-measurement
+            state that is used for postselection will be considered in the remaining circuit.
 
     Returns:
         MidMeasureMP: measurement process instance
@@ -102,7 +105,7 @@ def func():
 
     # Create a UUID and a map between MP and MV to support serialization
     measurement_id = str(uuid.uuid4())[:8]
-    mp = MidMeasureMP(wires=wire, reset=reset, id=measurement_id)
+    mp = MidMeasureMP(wires=wire, reset=reset, postselect=postselect, id=measurement_id)
     return MeasurementValue([mp], processing_fn=lambda v: v)
 
 
@@ -121,14 +124,22 @@ class MidMeasureMP(MeasurementProcess):
         wires (.Wires): The wires the measurement process applies to.
             This can only be specified if an observable was not provided.
         reset (bool): Whether to reset the wire after measurement.
+        postselect (Optional[int]): Which basis state to postselect after a mid-circuit
+            measurement. None by default. If postselection is requested, only the post-measurement
+            state that is used for postselection will be considered in the remaining circuit.
         id (str): Custom label given to a measurement instance.
     """
 
     def __init__(
-        self, wires: Optional[Wires] = None, reset: Optional[bool] = False, id: Optional[str] = None
+        self,
+        wires: Optional[Wires] = None,
+        reset: Optional[bool] = False,
+        postselect: Optional[int] = None,
+        id: Optional[str] = None,
     ):
         super().__init__(wires=Wires(wires), id=id)
         self.reset = reset
+        self.postselect = postselect
 
     @property
     def return_type(self):
diff --git a/pennylane/ops/qubit/observables.py b/pennylane/ops/qubit/observables.py
index 823e2250755..c8b92f86bf2 100644
--- a/pennylane/ops/qubit/observables.py
+++ b/pennylane/ops/qubit/observables.py
@@ -23,7 +23,7 @@
 from scipy.sparse import csr_matrix
 
 import pennylane as qml
-from pennylane.operation import AnyWires, Observable
+from pennylane.operation import AnyWires, Observable, Operation
 from pennylane.wires import Wires
 
 from .matrix_ops import QubitUnitary
@@ -321,7 +321,7 @@ def compute_sparse_matrix(H):  # pylint: disable=arguments-differ
         return H
 
 
-class Projector(Observable):
+class Projector(Operation, Observable):
     r"""Projector(state, wires, id=None)
     Observable corresponding to the state projector :math:`P=\ket{\phi}\bra{\phi}`.
 
diff --git a/pennylane/transforms/defer_measurements.py b/pennylane/transforms/defer_measurements.py
index 58034601188..07dc3172d1f 100644
--- a/pennylane/transforms/defer_measurements.py
+++ b/pennylane/transforms/defer_measurements.py
@@ -146,7 +146,13 @@ def func(x, y):
 
     for op in tape.operations:
         if isinstance(op, MidMeasureMP):
-            if op.reset is True:
+            if op.postselect is not None:
+                if tape.shots and (tape.shots.has_partitioned_shots or tape.batch_size):
+                    raise ValueError(
+                        "Cannot postselect on mid-circuit measurements when the circuit has finite "
+                        "shots if the circuit is batched or the shots are partitioned."
+                    )
+            if op.reset:
                 reused_measurement_wires.add(op.wires[0])
 
             if op.wires[0] in measured_wires:
@@ -169,6 +175,10 @@ def func(x, y):
         if isinstance(op, MidMeasureMP):
             _ = measured_wires.pop(0)
 
+            if op.postselect is not None:
+                with QueuingManager.stop_recording():
+                    new_operations.append(qml.Projector([op.postselect], wires=op.wires[0]))
+
             # Store measurement outcome in new wire if wire gets reused
             if op.wires[0] in reused_measurement_wires or op.wires[0] in measured_wires:
                 control_wires[op.id] = cur_wire

From a63c80293630ccd45279591224ba4692645aafbb Mon Sep 17 00:00:00 2001
From: Mudit Pandey <mudit.pandey@xanadu.ai>
Date: Fri, 15 Sep 2023 17:10:08 -0400
Subject: [PATCH 061/131] Fixed qnode

---
 tests/test_qnode.py                         |  2 +-
 tests/transforms/test_defer_measurements.py | 12 ++++++------
 2 files changed, 7 insertions(+), 7 deletions(-)

diff --git a/tests/test_qnode.py b/tests/test_qnode.py
index d23aa1cbb99..65254a331bf 100644
--- a/tests/test_qnode.py
+++ b/tests/test_qnode.py
@@ -1095,7 +1095,7 @@ def circuit():
         assert isinstance(circuit.tape.operations[1], qml.measurements.MidMeasureMP)
 
     @pytest.mark.parametrize(
-        "dev", [qml.device("default.qubit", wires=3), qml.devices.experimental.DefaultQubit2()]
+        "dev", [qml.device("default.qubit", wires=3), qml.device("default.qubit.legacy", wires=3)]
     )
     @pytest.mark.parametrize("first_par", np.linspace(0.15, np.pi - 0.3, 3))
     @pytest.mark.parametrize("sec_par", np.linspace(0.15, np.pi - 0.3, 3))
diff --git a/tests/transforms/test_defer_measurements.py b/tests/transforms/test_defer_measurements.py
index 25d100b2f6a..921be0abfcf 100644
--- a/tests/transforms/test_defer_measurements.py
+++ b/tests/transforms/test_defer_measurements.py
@@ -20,7 +20,7 @@
 
 import pennylane as qml
 import pennylane.numpy as np
-from pennylane.devices.experimental import DefaultQubit2
+from pennylane.devices import DefaultQubit
 
 
 class TestQNode:
@@ -157,7 +157,7 @@ def qnode2(phi, theta):
     def test_measurement_statistics_single_wire(self, shots):
         """Test that users can collect measurement statistics on
         a single mid-circuit measurement."""
-        dev = DefaultQubit2(seed=10)
+        dev = DefaultQubit(seed=10)
 
         @qml.qnode(dev)
         def circ1(x):
@@ -165,7 +165,7 @@ def circ1(x):
             m0 = qml.measure(0)
             return qml.probs(op=m0)
 
-        dev = DefaultQubit2(seed=10)
+        dev = DefaultQubit(seed=10)
 
         @qml.qnode(dev)
         def circ2(x):
@@ -179,8 +179,8 @@ def circ2(x):
     def test_terminal_measurements(self, shots):
         """Test that mid-circuit measurement statistics and terminal measurements
         can be made together."""
-        # Using DefaultQubit2 to allow non-commuting measurements
-        dev = DefaultQubit2(seed=10)
+        # Using DefaultQubit to allow non-commuting measurements
+        dev = DefaultQubit(seed=10)
 
         @qml.qnode(dev)
         def circ1(x, y):
@@ -189,7 +189,7 @@ def circ1(x, y):
             qml.RY(y, 1)
             return qml.expval(qml.PauliX(1)), qml.probs(op=m0)
 
-        dev = DefaultQubit2(seed=10)
+        dev = DefaultQubit(seed=10)
 
         @qml.qnode(dev)
         def circ2(x, y):

From 34cb02836819485587d7cb175670b2e175f558ee Mon Sep 17 00:00:00 2001
From: Mudit Pandey <mudit.pandey@xanadu.ai>
Date: Mon, 18 Sep 2023 10:09:42 -0400
Subject: [PATCH 062/131] Added MeasurementValue to DQL supported observables

---
 pennylane/devices/default_qubit_legacy.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/pennylane/devices/default_qubit_legacy.py b/pennylane/devices/default_qubit_legacy.py
index dc4983933c3..d0461886b1f 100644
--- a/pennylane/devices/default_qubit_legacy.py
+++ b/pennylane/devices/default_qubit_legacy.py
@@ -195,6 +195,7 @@ class DefaultQubitLegacy(QubitDevice):
         "Prod",
         "Exp",
         "Evolution",
+        "MeasurementValue",
     }
 
     def __init__(

From 33bfaa5e3664ca8a9e304aafee72f64d274cfee8 Mon Sep 17 00:00:00 2001
From: Mudit Pandey <mudit.pandey@xanadu.ai>
Date: Mon, 18 Sep 2023 10:15:41 -0400
Subject: [PATCH 063/131] update measure.py; add changelog for eigvals

---
 doc/releases/changelog-dev.md      |  4 ++++
 pennylane/devices/qubit/measure.py | 11 ++++-------
 2 files changed, 8 insertions(+), 7 deletions(-)

diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 5d5500d5de2..a3468276a3d 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -104,6 +104,10 @@
 
 <h3>Breaking changes 💔</h3>
 
+* `MeasurementProcess.eigvals()` now raises an `EigvalsUndefinedError` if the measurement observable
+  does not have eigenvalues.
+  [(#4544)](https://github.com/PennyLaneAI/pennylane/pull/4544)
+
 * The `__eq__` and `__hash__` methods of `Operator` and `MeasurementProcess` no longer rely on the
   object's address is memory. Using `==` with operators and measurement processes will now behave the
   same as `qml.equal`, and objects of the same type with the same data and hyperparameters will have
diff --git a/pennylane/devices/qubit/measure.py b/pennylane/devices/qubit/measure.py
index 733f2b6f074..3bc16d9b8d0 100644
--- a/pennylane/devices/qubit/measure.py
+++ b/pennylane/devices/qubit/measure.py
@@ -150,9 +150,10 @@ def get_measurement_function(
         Callable: function that returns the measurement result
     """
     if isinstance(measurementprocess, StateMeasurement):
+        if isinstance(measurementprocess.obs, MeasurementValue):
+            return state_diagonalizing_gates
+
         if isinstance(measurementprocess, ExpectationMP):
-            if isinstance(measurementprocess.obs, MeasurementValue):
-                return state_diagonalizing_gates
             if measurementprocess.obs.name == "SparseHamiltonian":
                 return csr_dot_products
 
@@ -174,11 +175,7 @@ def get_measurement_function(
 
                     return csr_dot_products
 
-        if (
-            measurementprocess.obs is None
-            or isinstance(measurementprocess.obs, MeasurementValue)
-            or measurementprocess.obs.has_diagonalizing_gates
-        ):
+        if measurementprocess.obs is None or measurementprocess.obs.has_diagonalizing_gates:
             return state_diagonalizing_gates
 
     raise NotImplementedError

From 0c3222300b451f690e1b12887286155282920586 Mon Sep 17 00:00:00 2001
From: Mudit Pandey <mudit.pandey@xanadu.ai>
Date: Fri, 15 Sep 2023 17:10:08 -0400
Subject: [PATCH 064/131] Fixed qnode

---
 tests/test_qnode.py                         |  2 +-
 tests/transforms/test_defer_measurements.py | 12 ++++++------
 2 files changed, 7 insertions(+), 7 deletions(-)

diff --git a/tests/test_qnode.py b/tests/test_qnode.py
index d23aa1cbb99..65254a331bf 100644
--- a/tests/test_qnode.py
+++ b/tests/test_qnode.py
@@ -1095,7 +1095,7 @@ def circuit():
         assert isinstance(circuit.tape.operations[1], qml.measurements.MidMeasureMP)
 
     @pytest.mark.parametrize(
-        "dev", [qml.device("default.qubit", wires=3), qml.devices.experimental.DefaultQubit2()]
+        "dev", [qml.device("default.qubit", wires=3), qml.device("default.qubit.legacy", wires=3)]
     )
     @pytest.mark.parametrize("first_par", np.linspace(0.15, np.pi - 0.3, 3))
     @pytest.mark.parametrize("sec_par", np.linspace(0.15, np.pi - 0.3, 3))
diff --git a/tests/transforms/test_defer_measurements.py b/tests/transforms/test_defer_measurements.py
index 25d100b2f6a..921be0abfcf 100644
--- a/tests/transforms/test_defer_measurements.py
+++ b/tests/transforms/test_defer_measurements.py
@@ -20,7 +20,7 @@
 
 import pennylane as qml
 import pennylane.numpy as np
-from pennylane.devices.experimental import DefaultQubit2
+from pennylane.devices import DefaultQubit
 
 
 class TestQNode:
@@ -157,7 +157,7 @@ def qnode2(phi, theta):
     def test_measurement_statistics_single_wire(self, shots):
         """Test that users can collect measurement statistics on
         a single mid-circuit measurement."""
-        dev = DefaultQubit2(seed=10)
+        dev = DefaultQubit(seed=10)
 
         @qml.qnode(dev)
         def circ1(x):
@@ -165,7 +165,7 @@ def circ1(x):
             m0 = qml.measure(0)
             return qml.probs(op=m0)
 
-        dev = DefaultQubit2(seed=10)
+        dev = DefaultQubit(seed=10)
 
         @qml.qnode(dev)
         def circ2(x):
@@ -179,8 +179,8 @@ def circ2(x):
     def test_terminal_measurements(self, shots):
         """Test that mid-circuit measurement statistics and terminal measurements
         can be made together."""
-        # Using DefaultQubit2 to allow non-commuting measurements
-        dev = DefaultQubit2(seed=10)
+        # Using DefaultQubit to allow non-commuting measurements
+        dev = DefaultQubit(seed=10)
 
         @qml.qnode(dev)
         def circ1(x, y):
@@ -189,7 +189,7 @@ def circ1(x, y):
             qml.RY(y, 1)
             return qml.expval(qml.PauliX(1)), qml.probs(op=m0)
 
-        dev = DefaultQubit2(seed=10)
+        dev = DefaultQubit(seed=10)
 
         @qml.qnode(dev)
         def circ2(x, y):

From e7b332bd2bb6095e435bc75ad660ecc54bb8f8b0 Mon Sep 17 00:00:00 2001
From: Mudit Pandey <mudit.pandey@xanadu.ai>
Date: Mon, 18 Sep 2023 10:09:42 -0400
Subject: [PATCH 065/131] Added MeasurementValue to DQL supported observables

---
 pennylane/devices/default_qubit_legacy.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/pennylane/devices/default_qubit_legacy.py b/pennylane/devices/default_qubit_legacy.py
index dc4983933c3..d0461886b1f 100644
--- a/pennylane/devices/default_qubit_legacy.py
+++ b/pennylane/devices/default_qubit_legacy.py
@@ -195,6 +195,7 @@ class DefaultQubitLegacy(QubitDevice):
         "Prod",
         "Exp",
         "Evolution",
+        "MeasurementValue",
     }
 
     def __init__(

From 67d9ede091ce1e4b15302d2188229ee87909e2cf Mon Sep 17 00:00:00 2001
From: Mudit Pandey <mudit.pandey@xanadu.ai>
Date: Mon, 18 Sep 2023 10:15:41 -0400
Subject: [PATCH 066/131] update measure.py; add changelog for eigvals

---
 doc/releases/changelog-dev.md      |  4 ++++
 pennylane/devices/qubit/measure.py | 11 ++++-------
 2 files changed, 8 insertions(+), 7 deletions(-)

diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 5d5500d5de2..a3468276a3d 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -104,6 +104,10 @@
 
 <h3>Breaking changes 💔</h3>
 
+* `MeasurementProcess.eigvals()` now raises an `EigvalsUndefinedError` if the measurement observable
+  does not have eigenvalues.
+  [(#4544)](https://github.com/PennyLaneAI/pennylane/pull/4544)
+
 * The `__eq__` and `__hash__` methods of `Operator` and `MeasurementProcess` no longer rely on the
   object's address is memory. Using `==` with operators and measurement processes will now behave the
   same as `qml.equal`, and objects of the same type with the same data and hyperparameters will have
diff --git a/pennylane/devices/qubit/measure.py b/pennylane/devices/qubit/measure.py
index 733f2b6f074..3bc16d9b8d0 100644
--- a/pennylane/devices/qubit/measure.py
+++ b/pennylane/devices/qubit/measure.py
@@ -150,9 +150,10 @@ def get_measurement_function(
         Callable: function that returns the measurement result
     """
     if isinstance(measurementprocess, StateMeasurement):
+        if isinstance(measurementprocess.obs, MeasurementValue):
+            return state_diagonalizing_gates
+
         if isinstance(measurementprocess, ExpectationMP):
-            if isinstance(measurementprocess.obs, MeasurementValue):
-                return state_diagonalizing_gates
             if measurementprocess.obs.name == "SparseHamiltonian":
                 return csr_dot_products
 
@@ -174,11 +175,7 @@ def get_measurement_function(
 
                     return csr_dot_products
 
-        if (
-            measurementprocess.obs is None
-            or isinstance(measurementprocess.obs, MeasurementValue)
-            or measurementprocess.obs.has_diagonalizing_gates
-        ):
+        if measurementprocess.obs is None or measurementprocess.obs.has_diagonalizing_gates:
             return state_diagonalizing_gates
 
     raise NotImplementedError

From 1e8c0de1ab9fc1ef61b78f779cc69da4b9b7c502 Mon Sep 17 00:00:00 2001
From: Mudit Pandey <mudit.pandey@xanadu.ai>
Date: Mon, 18 Sep 2023 11:39:50 -0400
Subject: [PATCH 067/131] Updated error message

---
 pennylane/devices/qubit/simulate.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pennylane/devices/qubit/simulate.py b/pennylane/devices/qubit/simulate.py
index 3847d749555..f55d85f392a 100644
--- a/pennylane/devices/qubit/simulate.py
+++ b/pennylane/devices/qubit/simulate.py
@@ -123,7 +123,7 @@ def get_final_state(circuit, debugger=None, interface=None):
             else:
                 norm = qml.math.norm(state)
                 if qml.math.isclose(norm, 0.0):
-                    raise ValueError("Requested postselection on state with 0 probability.")
+                    raise ValueError("Requested postselection on value with 0 probability.")
 
                 state = state / norm
 

From c0ef20bf360bfa83dae598593e73eeb1bb49f521 Mon Sep 17 00:00:00 2001
From: Mudit Pandey <mudit.pandey@xanadu.ai>
Date: Mon, 18 Sep 2023 18:15:06 -0400
Subject: [PATCH 068/131] Added tests

---
 pennylane/devices/qubit/simulate.py         |  14 ++-
 pennylane/transforms/defer_measurements.py  |   9 +-
 tests/devices/qubit/test_simulate.py        |  48 ++++++++
 tests/transforms/test_defer_measurements.py | 124 +++++++++++++++++++-
 4 files changed, 188 insertions(+), 7 deletions(-)

diff --git a/pennylane/devices/qubit/simulate.py b/pennylane/devices/qubit/simulate.py
index f55d85f392a..20fc9a0d5d3 100644
--- a/pennylane/devices/qubit/simulate.py
+++ b/pennylane/devices/qubit/simulate.py
@@ -118,18 +118,24 @@ def get_final_state(circuit, debugger=None, interface=None):
         if isinstance(op, qml.Projector):
             # Handle postselection on mid-circuit measurements
             if is_state_batched:
-                for i in qml.math.shape(state)[0]:
+                for i, _ in state:
+                    norm = qml.math.norm(state[i])
+                    if qml.math.isclose(norm, 0.0):
+                        raise ValueError("Requested postselection on value with zero probability.")
+
                     state[i] = state[i] / qml.math.norm(state[i])
             else:
                 norm = qml.math.norm(state)
                 if qml.math.isclose(norm, 0.0):
-                    raise ValueError("Requested postselection on value with 0 probability.")
+                    raise ValueError("Requested postselection on value with zero probability.")
 
                 state = state / norm
 
+                # defer_measurements will raise an error with batched shots or broadcasting so we can
+                # assume that both the state and shots are unbatched.
                 if circuit.shots:
-                    # defer_measurements will raise an error with batched shots or broadcasting
-                    # so we can assume that both the state and shots are unbatched.
+                    # Clip the number of shots using a binomial distribution using the probability of
+                    # measuring the postselected state.
                     postselected_shots = binomial(circuit.shots.total_shots, norm)
                     circuit._shots = qml.measurements.Shots(postselected_shots)
 
diff --git a/pennylane/transforms/defer_measurements.py b/pennylane/transforms/defer_measurements.py
index 07dc3172d1f..e0f4e56476c 100644
--- a/pennylane/transforms/defer_measurements.py
+++ b/pennylane/transforms/defer_measurements.py
@@ -23,7 +23,7 @@
 from pennylane.wires import Wires
 from pennylane.queuing import QueuingManager
 
-# pylint: disable=too-many-branches
+# pylint: disable=too-many-branches, too-many-statements
 
 
 @transform
@@ -188,7 +188,12 @@ def func(x, y):
 
                 if op.reset:
                     with QueuingManager.stop_recording():
-                        new_operations.append(qml.CNOT([cur_wire, op.wires[0]]))
+                        if op.postselect == 1:
+                            # We know that the measured wire will be in the |1> state if postselected
+                            # |1>. So we can just apply a PauliX instead of a CNOT to reset
+                            new_operations.append(qml.PauliX(op.wires[0]))
+                        else:
+                            new_operations.append(qml.CNOT([cur_wire, op.wires[0]]))
 
                 cur_wire += 1
             else:
diff --git a/tests/devices/qubit/test_simulate.py b/tests/devices/qubit/test_simulate.py
index eb9851b8e87..a9903e4e7ab 100644
--- a/tests/devices/qubit/test_simulate.py
+++ b/tests/devices/qubit/test_simulate.py
@@ -48,6 +48,39 @@ def compute_matrix():
     assert qml.math.allclose(result, -1.0)
 
 
+def test_projector_probability_error():
+    """Test that an error is raised if a Projector is applied in the operation queue
+    and the resulting state is null."""
+    tape = qml.tape.QuantumScript([qml.PauliX(0), qml.Projector([0], wires=0)], [qml.state()])
+    with pytest.raises(ValueError, match="Requested postselection on value with zero probability"):
+        _ = simulate(tape)
+
+    # Test with broadcasting
+    tape = qml.tape.QuantumScript(
+        [qml.RX([0.0, 0.0], wires=0), qml.Projector([1], 0)], [qml.state()]
+    )
+    with pytest.raises(ValueError, match="Requested postselection on value with zero probability"):
+        _ = simulate(tape)
+
+
+def test_projector_norm():
+    """Test that the norm of the state is maintained after applying a projector"""
+    tape = qml.tape.QuantumScript(
+        [qml.PauliX(0), qml.RX(0.123, 1), qml.Projector([0], wires=1)], [qml.state()]
+    )
+    res = simulate(tape)
+    assert np.isclose(np.linalg.norm(res), 1.0)
+
+    # Test with broadcasting
+    tape = qml.tape.QuantumScript(
+        [qml.PauliX(0), qml.RX([0.123, 0.456], 1), qml.Projector([0], wires=1)], [qml.state()]
+    )
+    res = simulate(tape)
+    assert len(res) == 2
+    for r in res:
+        assert np.isclose(np.linalg.norm(r), 1.0)
+
+
 # pylint: disable=too-few-public-methods
 class TestStatePrepBase:
     """Tests integration with various state prep methods."""
@@ -350,6 +383,21 @@ def test_broadcasting_with_extra_measurement_wires(self, mocker):
         assert np.allclose(res[2], -np.cos(x))
         assert spy.call_args_list[0].args == (qs, {2: 0, 1: 1, 0: 2})
 
+    def test_broadcasting_with_projector(self):
+        """Test that postselecting a broadcasted state works correctly"""
+        tape = qml.tape.QuantumScript(
+            [
+                qml.RX([0.1, 0.2], 0),
+                qml.Projector([0], wires=0),
+            ],
+            [qml.state()],
+        )
+
+        res = simulate(tape)
+        assert len(res) == 2
+        for r in res:
+            assert np.allclose(r, [1, 0])
+
 
 class TestDebugger:
     """Tests that the debugger works for a simple circuit"""
diff --git a/tests/transforms/test_defer_measurements.py b/tests/transforms/test_defer_measurements.py
index 921be0abfcf..ca72bda5e62 100644
--- a/tests/transforms/test_defer_measurements.py
+++ b/tests/transforms/test_defer_measurements.py
@@ -14,11 +14,12 @@
 """
 Tests for the transform implementing the deferred measurement principle.
 """
-# pylint: disable=too-few-public-methods
+# pylint: disable=too-few-public-methods, too-many-arguments
 import math
 import pytest
 
 import pennylane as qml
+from pennylane.measurements import MidMeasureMP, MeasurementValue
 import pennylane.numpy as np
 from pennylane.devices import DefaultQubit
 
@@ -153,6 +154,127 @@ def qnode2(phi, theta):
         assert len(deferred_tape2.wires) == 3
         assert len(deferred_tape2.operations) == 4
 
+    @pytest.mark.parametrize("shots", [None, 1000])
+    @pytest.mark.parametrize("phi", np.linspace(np.pi / 2, 7 * np.pi / 2, 6))
+    def test_postselection_qnode(self, phi, shots):
+        """Test that a Projector is queued when postselection is requested
+        on a mid-circuit measurement"""
+        dev = DefaultQubit()
+
+        @qml.qnode(dev)
+        def circ1(phi):
+            qml.RX(phi, wires=0)
+            # Postselecting on |1> on wire 0 means that the probability of measuring
+            # |1> on wire 0 is 1
+            m = qml.measure(0, postselect=1)
+            qml.cond(m, qml.PauliX)(wires=1)
+            # Probability of measuring |1> on wire 1 should be 1
+            return qml.probs(wires=1)
+
+        assert np.allclose(circ1(phi, shots=shots), [0, 1])
+
+        expected_circuit = [
+            qml.RX(phi, 0),
+            qml.Projector([1], wires=0),
+            qml.CNOT([0, 1]),
+            qml.probs(wires=1),
+        ]
+
+        for op, expected_op in zip(circ1.qtape, expected_circuit):
+            assert qml.equal(op, expected_op)
+
+    @pytest.mark.parametrize("shots", [None, 1000])
+    @pytest.mark.parametrize("phi", np.linspace(np.pi / 4, 4 * np.pi, 4))
+    @pytest.mark.parametrize("theta", np.linspace(np.pi / 3, 3 * np.pi, 4))
+    def test_multiple_postselection_qnode(self, phi, theta, shots, tol, tol_stochastic):
+        """Test that a qnode with mid-circuit measurements with postselection
+        is transformed correctly by defer_measurements"""
+        dev = DefaultQubit()
+
+        # Initializing mid circuit measurements here so that id can be controlled (affects
+        # wire ordering for qml.cond)
+        mp0 = MidMeasureMP(wires=0, postselect=0, id=0)
+        mv0 = MeasurementValue([mp0], lambda v: v)
+        mp1 = MidMeasureMP(wires=1, postselect=0, id=1)
+        mv1 = MeasurementValue([mp1], lambda v: v)
+        mp2 = MidMeasureMP(wires=2, reset=True, postselect=1, id=2)
+        mv2 = MeasurementValue([mp2], lambda v: v)
+
+        @qml.qnode(dev)
+        def circ1(phi, theta):
+            qml.RX(phi, 0)
+            qml.apply(mp0)
+            qml.CNOT([0, 1])
+            qml.apply(mp1)
+            qml.cond(~(mv0 & mv1), qml.RY)(theta, wires=2)
+            qml.apply(mp2)
+            qml.cond(mv2, qml.PauliX)(1)
+            return qml.probs(wires=[0, 1, 2])
+
+        @qml.qnode(dev)
+        def circ2():
+            # To add wire 0 to tape
+            qml.Identity(0)
+            qml.PauliX(1)
+            qml.Identity(2)
+            return qml.probs(wires=[0, 1, 2])
+
+        atol = tol if shots is None else tol_stochastic
+        assert np.allclose(circ1(phi, theta, shots=shots), circ2(), atol=atol, rtol=0)
+
+        expected_circuit = [
+            qml.RX(phi, wires=0),
+            qml.Projector([0], wires=0),
+            qml.CNOT([0, 3]),
+            qml.CNOT([0, 1]),
+            qml.Projector([0], wires=1),
+            qml.CNOT([1, 4]),
+            qml.ops.Controlled(
+                qml.RY(theta, wires=[2]), control_wires=[3, 4], control_values=[False, False]
+            ),
+            qml.ops.Controlled(
+                qml.RY(theta, wires=[2]), control_wires=[3, 4], control_values=[False, True]
+            ),
+            qml.ops.Controlled(
+                qml.RY(theta, wires=[2]), control_wires=[3, 4], control_values=[True, False]
+            ),
+            qml.Projector([1], wires=2),
+            qml.CNOT([2, 5]),
+            qml.PauliX(2),
+            qml.CNOT([5, 1]),
+            qml.probs(wires=[0, 1, 2]),
+        ]
+
+        for op, expected_op in zip(circ1.qtape, expected_circuit):
+            assert qml.equal(op, expected_op)
+
+    def test_postselection_error(self):
+        """Test that an error is raised if the circuit being transformed has batched dimensions
+        or partitioned shots."""
+        dev = DefaultQubit()
+
+        @qml.qnode(dev)
+        def circ1():
+            qml.RX([0.1, 0.2], wires=0)
+            _ = qml.measure(0, postselect=1)
+            return qml.probs()
+
+        # Test that no error is raised if broadcasting with analytic execution
+        _ = circ1()
+
+        with pytest.raises(ValueError, match="Cannot postselect on mid-circuit measurements"):
+            # Error with finite shots and broadcasting
+            _ = circ1(shots=10)
+
+        @qml.qnode(dev)
+        def circ2():
+            qml.measure(0, postselect=1)
+            return qml.probs()
+
+        with pytest.raises(ValueError, match="Cannot postselect on mid-circuit measurements"):
+            # Error with shot vector
+            _ = circ2(shots=[10, 10])
+
     @pytest.mark.parametrize("shots", [None, 1000, [1000, 1000]])
     def test_measurement_statistics_single_wire(self, shots):
         """Test that users can collect measurement statistics on

From b93563b511b8c2ca1075107feb863e7e71dcfc04 Mon Sep 17 00:00:00 2001
From: Mudit Pandey <mudit.pandey@xanadu.ai>
Date: Mon, 18 Sep 2023 18:24:35 -0400
Subject: [PATCH 069/131] Added changelog

---
 doc/releases/changelog-dev.md | 24 ++++++++++++++++++++++++
 1 file changed, 24 insertions(+)

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

From e8087a2d927b3133065b2a8e2a5f5a29359f3efb Mon Sep 17 00:00:00 2001
From: Mudit Pandey <mudit.pandey@xanadu.ai>
Date: Tue, 19 Sep 2023 09:44:39 -0400
Subject: [PATCH 070/131] Fixed error

---
 pennylane/devices/qubit/simulate.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pennylane/devices/qubit/simulate.py b/pennylane/devices/qubit/simulate.py
index 20fc9a0d5d3..fd38f0533bd 100644
--- a/pennylane/devices/qubit/simulate.py
+++ b/pennylane/devices/qubit/simulate.py
@@ -118,7 +118,7 @@ def get_final_state(circuit, debugger=None, interface=None):
         if isinstance(op, qml.Projector):
             # Handle postselection on mid-circuit measurements
             if is_state_batched:
-                for i, _ in state:
+                for i, _ in enumerate(state):
                     norm = qml.math.norm(state[i])
                     if qml.math.isclose(norm, 0.0):
                         raise ValueError("Requested postselection on value with zero probability.")

From 9117c8a0ac508c0699f6d7bdb2770a9f83616812 Mon Sep 17 00:00:00 2001
From: Mudit Pandey <mudit.pandey@xanadu.ai>
Date: Tue, 19 Sep 2023 14:05:58 -0400
Subject: [PATCH 071/131] Updated simulate; added docs

---
 doc/introduction/measurements.rst   | 30 +++++++++++++++++++++++++++++
 pennylane/devices/qubit/simulate.py |  9 +++++++--
 2 files changed, 37 insertions(+), 2 deletions(-)

diff --git a/doc/introduction/measurements.rst b/doc/introduction/measurements.rst
index befff75b137..bd5323d17b8 100644
--- a/doc/introduction/measurements.rst
+++ b/doc/introduction/measurements.rst
@@ -329,6 +329,36 @@ Executing this QNode:
 >>> func()
 tensor([0., 1.], requires_grad=True)
 
+PennyLane also supports postselecting on mid-circuit measurement outcomes such that only states matching the
+postselection are considered in the remaining execution by specifying the ``postselect`` keyword argument of
+:func:`~.pennylane.measure`:
+
+.. code-block:: python3
+
+    dev = qml.device("default.qubit", wires=2)
+
+    @qml.qnode(dev)
+    def func(x):
+        qml.RX(x, wires=0)
+        m0 = qml.measure(0, postselect=1)
+        qml.cond(m0, qml.PauliX)(wires=1)
+        return qml.sample(wires=1)
+
+By postselecting on ``1``, we only consider the ``1`` measurement outcome. So, the probability of measuring 1
+on wire 1 should be 100%. Executing this QNode with 10 shots:
+
+>>> func(np.pi / 2, shots=10)
+array([1, 1, 1, 1, 1, 1, 1])
+
+Note that only 7 samples are returned. This is because samples that do not meet the postselection criteria are
+thrown away.
+
+.. note::
+
+    Currently, postselection support is only available on ``"default.qubit"``. Requesting postselection on other
+    devices will not do anything.
+
+
 The deferred measurement principle provides a natural method to simulate the
 application of mid-circuit measurements and conditional operations in a
 differentiable and device-independent way. Performing true mid-circuit
diff --git a/pennylane/devices/qubit/simulate.py b/pennylane/devices/qubit/simulate.py
index fd38f0533bd..41e89e42ce9 100644
--- a/pennylane/devices/qubit/simulate.py
+++ b/pennylane/devices/qubit/simulate.py
@@ -121,13 +121,13 @@ def get_final_state(circuit, debugger=None, interface=None):
                 for i, _ in enumerate(state):
                     norm = qml.math.norm(state[i])
                     if qml.math.isclose(norm, 0.0):
-                        raise ValueError("Requested postselection on value with zero probability.")
+                        return qml.numpy.NaN, is_state_batched
 
                     state[i] = state[i] / qml.math.norm(state[i])
             else:
                 norm = qml.math.norm(state)
                 if qml.math.isclose(norm, 0.0):
-                    raise ValueError("Requested postselection on value with zero probability.")
+                    return qml.numpy.NaN, is_state_batched
 
                 state = state / norm
 
@@ -137,6 +137,8 @@ def get_final_state(circuit, debugger=None, interface=None):
                     # Clip the number of shots using a binomial distribution using the probability of
                     # measuring the postselected state.
                     postselected_shots = binomial(circuit.shots.total_shots, norm)
+                    if postselected_shots == 0:
+                        raise RuntimeError("None of the samples meet the postselection criteria")
                     circuit._shots = qml.measurements.Shots(postselected_shots)
 
     if set(circuit.op_wires) < set(circuit.wires):
@@ -167,6 +169,9 @@ def measure_final_state(circuit, state, is_state_batched, rng=None) -> Result:
     Returns:
         Tuple[TensorLike]: The measurement results
     """
+    if state is qml.numpy.NaN:
+        return state
+
     circuit = circuit.map_to_standard_wires()
 
     if not circuit.shots:

From 1d1230828a6a4d6c451a71689d046b26e4af3562 Mon Sep 17 00:00:00 2001
From: Mudit Pandey <mudit.pandey@xanadu.ai>
Date: Tue, 19 Sep 2023 14:23:32 -0400
Subject: [PATCH 072/131] Added docs to measurements.rst

---
 doc/introduction/measurements.rst | 26 ++++++++++++++++++++++++++
 1 file changed, 26 insertions(+)

diff --git a/doc/introduction/measurements.rst b/doc/introduction/measurements.rst
index befff75b137..e0b10c5f7c9 100644
--- a/doc/introduction/measurements.rst
+++ b/doc/introduction/measurements.rst
@@ -329,6 +329,32 @@ Executing this QNode:
 >>> func()
 tensor([0., 1.], requires_grad=True)
 
+Statistics can also be collected on mid-circuit measurements along with terminal measurement statistics.
+Currently, ``qml.probs``, ``qml.sample``, ``qml.expval``, ``qml.var``, and ``qml.counts`` are supported,
+and can be requested along with other measurements. The devices that currently support collecting such
+statistics are ``"default.qubit"``, ``"default.mixed"``, and ``"default.qubit.legacy"``.
+
+..code-block:: python3
+
+    dev = qml.devices.DefaultQubit()
+
+    @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.probs(wires=1), qml.probs(op=m0)
+
+Executing this QNode:
+
+>>> func(np.pi / 2, np.pi / 4)
+(tensor([0.9267767, 0.0732233], requires_grad=True),
+ tensor([0.5, 0.5], requires_grad=True))
+
+Currently, statistics can only be collected for single mid-circuit measurement values. Moreover, any
+measurement values manipulated using boolean or arithmetic operators cannot be used. These can lead to
+unexpected/incorrect behaviour.
+
 The deferred measurement principle provides a natural method to simulate the
 application of mid-circuit measurements and conditional operations in a
 differentiable and device-independent way. Performing true mid-circuit

From b3b490dec0f8be5425d9524491e593e8b2908ff6 Mon Sep 17 00:00:00 2001
From: Matthew Silverman <matthews@xanadu.ai>
Date: Mon, 18 Sep 2023 15:43:30 -0400
Subject: [PATCH 073/131] fix torch take with axis=-1 (#4605)

**Context:**
Basically, if you call `qml.math.take(x, indices=your_indices, axis=2,
like="torch")`, it returns `x[:, :, your_indices]` (with `axis`
placeholder colons, in this case 2). If `axis` is -1, we would return
`x[indices]` (because `[vals] * -1` gives an empty list). That's wrong.

**Description of the Change:**
`qml.math.take` for torch with `axis=-1` correctly returns `x[...,
indices]` instead of `x[indices]`.

**Benefits:**
`qml.math.take` works with `torch` and `axis=-1`.

**Related GitHub Issues:**
First reported in #4589
---
 doc/releases/changelog-dev.md     | 4 ++++
 pennylane/math/single_dispatch.py | 2 ++
 tests/math/test_functions.py      | 6 ++++++
 3 files changed, 12 insertions(+)

diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index e3a99bab095..c96e09f7362 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -264,6 +264,10 @@
   expectation.
   [(#4590)](https://github.com/PennyLaneAI/pennylane/pull/4590)
 
+* `qml.math.take` with torch now returns `tensor[..., indices]` when the user requests
+  the last axis (`axis=-1`). Without the fix, it would wrongly return `tensor[indices]`.
+  [(#4605)](https://github.com/PennyLaneAI/pennylane/pull/4605)
+
 <h3>Contributors ✍️</h3>
 
 This release contains contributions from (in alphabetical order):
diff --git a/pennylane/math/single_dispatch.py b/pennylane/math/single_dispatch.py
index 1ca2a96dfde..dd40be00e78 100644
--- a/pennylane/math/single_dispatch.py
+++ b/pennylane/math/single_dispatch.py
@@ -550,6 +550,8 @@ def _take_torch(tensor, indices, axis=None, **_):
 
         return torch.index_select(tensor, dim=axis, index=indices)
 
+    if axis == -1:
+        return tensor[..., indices]
     fancy_indices = [slice(None)] * axis + [indices]
     return tensor[fancy_indices]
 
diff --git a/tests/math/test_functions.py b/tests/math/test_functions.py
index 913ad85a407..5dceb524562 100644
--- a/tests/math/test_functions.py
+++ b/tests/math/test_functions.py
@@ -1596,6 +1596,12 @@ def cost_fn(t):
         expected = np.array([[[3, 3], [1, 1], [0, 0]], [[3, 3], [1, 1], [0, 0]]])
         assert fn.allclose(grad, expected)
 
+    @pytest.mark.torch
+    def test_last_axis_support_torch(self):
+        """Test that _torch_take correctly sets the last axis"""
+        x = fn.arange(8, like="torch").reshape((2, 4))
+        assert np.array_equal(fn.take(x, indices=3, axis=-1), [3, 7])
+
 
 where_data = [
     np.array([[[1, 2], [3, 4], [-1, 1]], [[5, 6], [0, -1], [2, 1]]]),

From 832106571bf9bf000c76c0cc3470ccb6e7e31461 Mon Sep 17 00:00:00 2001
From: Christina Lee <christina@xanadu.ai>
Date: Tue, 19 Sep 2023 10:21:57 -0400
Subject: [PATCH 074/131] [CI] Bump jax version to v0.4.16 (#4612)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

A new version of jax has been released. This PR attempts to bump the
version in the CI 🤞
---
 .github/workflows/interface-unit-tests.yml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.github/workflows/interface-unit-tests.yml b/.github/workflows/interface-unit-tests.yml
index 2b8d6a6c9d2..8ea39a1e5a5 100644
--- a/.github/workflows/interface-unit-tests.yml
+++ b/.github/workflows/interface-unit-tests.yml
@@ -10,7 +10,7 @@ on:
         description: The version of JAX to install for any job that requires JAX
         required: false
         type: string
-        default: 0.4.14
+        default: 0.4.16
       tensorflow_version:
         description: The version of TensorFlow to install for any job that requires TensorFlow
         required: false

From 8e421f532c547ab8f10c0e6547c4ce003f840071 Mon Sep 17 00:00:00 2001
From: Matthew Silverman <matthews@xanadu.ai>
Date: Tue, 19 Sep 2023 11:02:54 -0400
Subject: [PATCH 075/131] TmpPauliRot still decomposes if the theta of zero is
 trainable (#4585)

**Context:**
While porting DQ2, the integration tests failed. It took me waaaaay long
to find out why, but it's because DQL kept the `TmpPauliRot` ops around
despite their value originally being 0, while DQ2 did away with them,
and that broke differentiation. I figure this was supposed to be an
optimization, but it's breaking assumptions made by DQ2.

The break is because DQL reaches [jax's trainable-parameter evaluation
code](https://github.com/PennyLaneAI/pennylane/blob/cbc8d19e6a36f0d4d38b44eabd4a8b98db677327/pennylane/interfaces/jax.py#L121-L126)
_before_ filtering out the zeros, thereby enabling differentiation over
those zeros, while DQ2 reaches that code _after_ filtering them out (so
it has no trainable parameters and doesn't differentiate over anything).
The filtering logic assumed DQL's execution order, and the slight
re-arrangement of when we decompose unsupported ops broke this with DQ2.

**Description of the Change:**
`TmpPauliRot.compute_decomposition` does not become nothing if the theta
value is trainable.

**Benefits:**
`SpecialUnitary` differentiation works with DQ2

**Possible Drawbacks:**
I feel like this still might hurt performance, but this is a consequence
of `SpecialUnitary`-implementation-meets-execution-pipeline assumptions.
Fixing that requires a deeper change, and it should probably happen in
`SpecialUnitary`'s code.

---------

Co-authored-by: Romain Moyard <rmoyard@gmail.com>
---
 doc/releases/changelog-dev.md           |  4 ++
 pennylane/ops/qubit/special_unitary.py  |  2 +-
 tests/ops/qubit/test_special_unitary.py | 69 +++++++++++++++++++++----
 3 files changed, 64 insertions(+), 11 deletions(-)

diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index c96e09f7362..65f4fcbc367 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -208,6 +208,10 @@
   been removed. Please use ``QuantumScript.bind_new_parameters`` instead.
   [(#4548)](https://github.com/PennyLaneAI/pennylane/pull/4548)
 
+* The private `TmpPauliRot` operator used for `SpecialUnitary` no longer decomposes to nothing
+  when the theta value is trainable.
+  [(#4585)](https://github.com/PennyLaneAI/pennylane/pull/4585)
+
 <h3>Deprecations 👋</h3>
 
 * The ``prep`` keyword argument in ``QuantumScript`` is deprecated and will be removed from `QuantumScript`.
diff --git a/pennylane/ops/qubit/special_unitary.py b/pennylane/ops/qubit/special_unitary.py
index 501c2ff36e2..997e73551ca 100644
--- a/pennylane/ops/qubit/special_unitary.py
+++ b/pennylane/ops/qubit/special_unitary.py
@@ -718,7 +718,7 @@ def compute_decomposition(theta, wires, pauli_word):
             This operation is used in a differentiation pipeline of :class:`~.SpecialUnitary`
             and most likely should not be created manually by users.
         """
-        if qml.math.isclose(theta, theta * 0):
+        if qml.math.isclose(theta, theta * 0) and not qml.math.requires_grad(theta):
             return []
         return [PauliRot(theta, pauli_word, wires)]
 
diff --git a/tests/ops/qubit/test_special_unitary.py b/tests/ops/qubit/test_special_unitary.py
index ab2a05f7870..1b2fd4aa271 100644
--- a/tests/ops/qubit/test_special_unitary.py
+++ b/tests/ops/qubit/test_special_unitary.py
@@ -27,6 +27,7 @@
     _pauli_letters,
     _pauli_matrices,
 )
+from pennylane.transforms.convert_to_numpy_parameters import _convert_op_to_numpy_data
 from pennylane.wires import Wires
 
 
@@ -750,6 +751,7 @@ def comparison(x):
         assert np.allclose(jac, expected)
 
 
+@pytest.mark.parametrize("dev_fn", [qml.devices.DefaultQubitLegacy, qml.devices.DefaultQubit])
 class TestSpecialUnitaryIntegration:
     """Test that the operation SpecialUnitary is executable and
     differentiable in a QNode context, both with automatic differentiation
@@ -772,9 +774,9 @@ def paulirot_comp_circuit(x, word=None):
     state = qml.SpecialUnitary.compute_matrix(x, 2) @ np.eye(4)[0]
     exp = np.vdot(state, qml.matrix(qml.PauliZ(0) @ qml.PauliX(1)) @ state).real
 
-    def test_qnode_numpy(self):
+    def test_qnode_numpy(self, dev_fn):
         """Test that the QNode executes with Numpy."""
-        dev = qml.device("default.qubit", wires=2)
+        dev = dev_fn(wires=2)
         qnode = qml.QNode(self.circuit, dev, interface=None)
 
         res = qnode(self.x)
@@ -782,11 +784,11 @@ def test_qnode_numpy(self):
         assert qml.math.isclose(res, self.exp)
 
     @pytest.mark.autograd
-    def test_qnode_autograd(self):
+    def test_qnode_autograd(self, dev_fn):
         """Test that the QNode executes with Autograd.
         Neither hardware-ready nor autodiff gradients are available in Autograd."""
 
-        dev = qml.device("default.qubit", wires=2)
+        dev = dev_fn(wires=2)
         qnode = qml.QNode(self.circuit, dev, interface="autograd")
 
         x = qml.numpy.array(self.x, requires_grad=True)
@@ -797,7 +799,7 @@ def test_qnode_autograd(self):
     @pytest.mark.jax
     @pytest.mark.parametrize("use_jit", [False, True])
     @pytest.mark.parametrize("shots, atol", [(None, 1e-6), (10000, 1e-1)])
-    def test_qnode_jax(self, shots, atol, use_jit):
+    def test_qnode_jax(self, dev_fn, shots, atol, use_jit):
         """Test that the QNode executes and is differentiable with JAX. The shots
         argument controls whether autodiff or parameter-shift gradients are used."""
         if use_jit and shots is not None:
@@ -806,7 +808,7 @@ def test_qnode_jax(self, shots, atol, use_jit):
 
         jax.config.update("jax_enable_x64", True)
 
-        dev = qml.device("default.qubit", wires=2, shots=shots)
+        dev = dev_fn(wires=2, shots=shots)
         diff_method = "backprop" if shots is None else "parameter-shift"
         qnode = qml.QNode(self.circuit, dev, interface="jax", diff_method=diff_method)
         if use_jit:
@@ -838,12 +840,12 @@ def test_qnode_jax(self, shots, atol, use_jit):
 
     @pytest.mark.torch
     @pytest.mark.parametrize("shots, atol", [(None, 1e-6), (10000, 1e-1)])
-    def test_qnode_torch(self, shots, atol):
+    def test_qnode_torch(self, dev_fn, shots, atol):
         """Test that the QNode executes and is differentiable with Torch. The shots
         argument controls whether autodiff or parameter-shift gradients are used."""
         import torch
 
-        dev = qml.device("default.qubit", wires=2, shots=shots)
+        dev = dev_fn(wires=2, shots=shots)
         diff_method = "backprop" if shots is None else "parameter-shift"
         qnode = qml.QNode(self.circuit, dev, interface="torch", diff_method=diff_method)
 
@@ -870,12 +872,12 @@ def test_qnode_torch(self, shots, atol):
 
     @pytest.mark.tf
     @pytest.mark.parametrize("shots, atol", [(None, 1e-6), (10000, 1e-1)])
-    def test_qnode_tf(self, shots, atol):
+    def test_qnode_tf(self, dev_fn, shots, atol):
         """Test that the QNode executes and is differentiable with TensorFlow. The shots
         argument controls whether autodiff or parameter-shift gradients are used."""
         import tensorflow as tf
 
-        dev = qml.device("default.qubit", wires=2, shots=shots)
+        dev = dev_fn(wires=2, shots=shots)
         diff_method = "backprop" if shots is None else "parameter-shift"
         qnode = qml.QNode(self.circuit, dev, interface="tf", diff_method=diff_method)
 
@@ -912,6 +914,12 @@ def test_qnode_tf(self, shots, atol):
 class TestTmpPauliRot:
     """Tests for the helper Operation TmpPauliRot."""
 
+    @staticmethod
+    def get_decomposition(x):
+        """A nonsense function that can be trained, to convert x to a trainable value."""
+        decomp = TmpPauliRot.compute_decomposition(x, [0], "X")
+        return float(len(decomp))
+
     def test_has_matrix_false(self):
         """Test that TmpPauliRot reports to not have a matrix."""
         assert TmpPauliRot.has_matrix is False
@@ -936,6 +944,47 @@ def test_decomposition_at_zero(self, word):
         assert op.decomposition() == []
         assert TmpPauliRot.compute_decomposition(0.0, wires, word) == []
 
+    @pytest.mark.autograd
+    def test_decomposition_at_zero_autograd(self):
+        """Test that the decomposition is a PauliRot if the theta value is trainable."""
+        x = qml.numpy.array(0.0, requires_grad=True)
+        with qml.queuing.AnnotatedQueue() as q:
+            qml.grad(self.get_decomposition)(x)
+        assert q.queue[0] == qml.PauliRot(x, "X", [0])
+
+    @pytest.mark.jax
+    def test_decomposition_at_zero_jax(self):
+        """Test that the decomposition is a PauliRot if the theta value is trainable."""
+        import jax
+
+        x = jax.numpy.array(0.0)
+        with qml.queuing.AnnotatedQueue() as q:
+            jax.grad(self.get_decomposition)(x)
+        assert _convert_op_to_numpy_data(q.queue[0]) == qml.PauliRot(0.0, "X", [0])
+
+    @pytest.mark.tf
+    def test_decomposition_at_zero_tf(self):
+        """Test that the decomposition is a PauliRot if the theta value is trainable."""
+        import tensorflow as tf
+
+        x = tf.Variable(0.0)
+        with qml.queuing.AnnotatedQueue() as q:
+            with tf.GradientTape():
+                num_ops = self.get_decomposition(x)
+        assert num_ops == 1
+        assert q.queue[0] == qml.PauliRot(x, "X", [0])
+
+    @pytest.mark.torch
+    def test_decomposition_at_zero_torch(self):
+        """Test that the decomposition is a PauliRot if the theta value is trainable."""
+        import torch
+
+        x = torch.tensor(0.0, requires_grad=True)
+        with qml.queuing.AnnotatedQueue() as q:
+            num_ops = self.get_decomposition(x)
+        assert num_ops == 1
+        assert q.queue[0] == qml.PauliRot(x, "X", [0])
+
     @pytest.mark.parametrize("word", ["X", "IZ", "YYY"])
     @pytest.mark.parametrize("x", [1.2, 1e-4])
     def test_decomposition_nonzero(self, word, x):

From bf4dbb2c88ae1b5111368e35b7ae83e9f2435819 Mon Sep 17 00:00:00 2001
From: Matthew Silverman <matthews@xanadu.ai>
Date: Tue, 19 Sep 2023 11:42:00 -0400
Subject: [PATCH 076/131] Update shot_adaptive to not mutate device shots
 (#4599)

**Context:**
`shot_adaptive.py` would change device wires, run a QNode with the
device, then change the wires back. We can just run the QNode with
different shots without changing the device!

**Description of the Change:**
Instead of changing device shots, pass `shots=new_shots` as a kwarg to
any QNode execution.

I also removed a silly test for something hacky that users probably
should not even be able to do... it's not that hard to add
`@qml.qnode(qml.device("default.qubit", shots=1000))` to the top of any
function : )

**Benefits:**
- It will work on the new device API without any changes (I ran the
tests locally with DQ2, worked perfectly)
- No more unnecessary try-finally blocks
- Not mutating object properties is always cool and wise
---
 doc/releases/changelog-dev.md                 |  4 +
 pennylane/optimize/shot_adaptive.py           | 80 +++++++------------
 tests/optimize/test_optimize_shot_adaptive.py | 26 ------
 3 files changed, 32 insertions(+), 78 deletions(-)

diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 65f4fcbc367..9b26844ab3e 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -126,6 +126,10 @@
   when at the beginning of a circuit, thus behaving like `StatePrep`.
   [(#4583)](https://github.com/PennyLaneAI/pennylane/pull/4583)
 
+* `ShotAdaptiveOptimizer` has been updated to pass shots to QNode executions instead of overriding
+  device shots before execution. This makes it compatible with the new device API.
+  [(#4599)](https://github.com/PennyLaneAI/pennylane/pull/4599)
+
 <h3>Breaking changes 💔</h3>
 
 * `MeasurementProcess.eigvals()` now raises an `EigvalsUndefinedError` if the measurement observable
diff --git a/pennylane/optimize/shot_adaptive.py b/pennylane/optimize/shot_adaptive.py
index 7b629928ad7..7beb372f598 100644
--- a/pennylane/optimize/shot_adaptive.py
+++ b/pennylane/optimize/shot_adaptive.py
@@ -242,7 +242,7 @@ def weighted_random_sampling(qnodes, coeffs, shots, argnums, *args, **kwargs):
                 continue
 
             # set the QNode device shots
-            h.device.shots = 1 if s == 1 else [(1, int(s))]
+            new_shots = 1 if s == 1 else [(1, int(s))]
 
             jacs = []
             for i in argnums:
@@ -255,7 +255,7 @@ def cost(*args, **kwargs):
                 else:
                     cost = h
 
-                j = qml.jacobian(cost, argnum=i)(*args, **kwargs)
+                j = qml.jacobian(cost, argnum=i)(*args, **kwargs, shots=new_shots)
 
                 if s == 1:
                     j = np.expand_dims(j, 0)
@@ -277,7 +277,7 @@ def check_device(dev):
         Raises:
             ValueError: if the device is analytic
         """
-        if dev.analytic:
+        if not dev.shots:
             raise ValueError(
                 "The Rosalin optimizer can only be used with devices "
                 "that estimate expectation values with a finite number of shots."
@@ -303,58 +303,47 @@ def _single_shot_expval_gradients(self, expval_cost, args, kwargs):
         """Compute the single shot gradients of an ExpvalCost object"""
         qnodes = expval_cost.qnodes
         coeffs = expval_cost.hamiltonian.coeffs
-        device = qnodes[0].device
 
-        self.check_device(device)
-        original_shots = device.shots
+        self.check_device(qnodes[0].device)
 
         if self.lipschitz is None:
             self.check_learning_rate(coeffs)
 
-        try:
-            if self.term_sampling == "weighted_random_sampling":
-                grads = self.weighted_random_sampling(
-                    qnodes, coeffs, self.max_shots, self.trainable_args, *args, **kwargs
-                )
-            elif self.term_sampling is None:
-                device.shots = [(1, int(self.max_shots))]
-                # We iterate over each trainable argument, rather than using
-                # qml.jacobian(expval_cost), to take into account the edge case where
-                # different arguments have different shapes and cannot be stacked.
-                grads = [
-                    qml.jacobian(expval_cost, argnum=i)(*args, **kwargs)
-                    for i in self.trainable_args
-                ]
-            else:
-                raise ValueError(
-                    f"Unknown Hamiltonian term sampling method {self.term_sampling}. "
-                    "Only term_sampling='weighted_random_sampling' and "
-                    "term_sampling=None currently supported."
-                )
-        finally:
-            device.shots = original_shots
+        if self.term_sampling == "weighted_random_sampling":
+            grads = self.weighted_random_sampling(
+                qnodes, coeffs, self.max_shots, self.trainable_args, *args, **kwargs
+            )
+        elif self.term_sampling is None:
+            new_shots = [(1, int(self.max_shots))]
+            # We iterate over each trainable argument, rather than using
+            # qml.jacobian(expval_cost), to take into account the edge case where
+            # different arguments have different shapes and cannot be stacked.
+            grads = [
+                qml.jacobian(expval_cost, argnum=i)(*args, **kwargs, shots=new_shots)
+                for i in self.trainable_args
+            ]
+        else:
+            raise ValueError(
+                f"Unknown Hamiltonian term sampling method {self.term_sampling}. "
+                "Only term_sampling='weighted_random_sampling' and "
+                "term_sampling=None currently supported."
+            )
 
         return grads
 
     def _single_shot_qnode_gradients(self, qnode, args, kwargs):
         """Compute the single shot gradients of a QNode."""
-        device = qnode.device
-
         self.check_device(qnode.device)
-        original_shots = device.shots
 
         if self.lipschitz is None:
             self.check_learning_rate(1)
 
-        try:
-            device.shots = [(1, int(self.max_shots))]
+        new_shots = [(1, int(self.max_shots))]
 
-            def cost(*args, **kwargs):
-                return qml.math.stack(qnode(*args, **kwargs))
+        def cost(*args, **kwargs):
+            return qml.math.stack(qnode(*args, **kwargs, shots=new_shots))
 
-            grads = [qml.jacobian(cost, argnum=i)(*args, **kwargs) for i in self.trainable_args]
-        finally:
-            device.shots = original_shots
+        grads = [qml.jacobian(cost, argnum=i)(*args, **kwargs) for i in self.trainable_args]
 
         return grads
 
@@ -509,18 +498,5 @@ def step_and_cost(self, objective_fn, *args, **kwargs):
             If single arg is provided, list [array] is replaced by array.
         """
         new_args = self.step(objective_fn, *args, **kwargs)
-
-        if isinstance(objective_fn, qml.ExpvalCost):
-            device = objective_fn.qnodes[0].device
-        elif isinstance(objective_fn, qml.QNode) or hasattr(objective_fn, "device"):
-            device = objective_fn.device
-
-        original_shots = device.shots
-
-        try:
-            device.shots = int(self.max_shots)
-            forward = objective_fn(*args, **kwargs)
-        finally:
-            device.shots = original_shots
-
+        forward = objective_fn(*args, **kwargs, shots=int(self.max_shots))
         return new_args, forward
diff --git a/tests/optimize/test_optimize_shot_adaptive.py b/tests/optimize/test_optimize_shot_adaptive.py
index 125bef5de33..862fff19efc 100644
--- a/tests/optimize/test_optimize_shot_adaptive.py
+++ b/tests/optimize/test_optimize_shot_adaptive.py
@@ -82,32 +82,6 @@ def ansatz(x, **kwargs):
         with pytest.raises(ValueError, match=f"The learning rate must be less than {2 / 1}"):
             opt.step(expval_cost.qnodes[0], np.array(0.5, requires_grad=True))
 
-    def test_unknown_objective_function(self):
-        """Test that an exception is raised if an unknown objective function is passed"""
-        dev = qml.device("default.qubit", wires=1, shots=100)
-
-        @qml.qnode(dev)
-        def circuit(x):
-            qml.RX(x, wires=0)
-            return qml.expval(qml.PauliZ(0))
-
-        def cost(x):
-            return np.sin(circuit(x))
-
-        opt = qml.ShotAdaptiveOptimizer(min_shots=10)
-
-        # test expval cost
-        with pytest.raises(
-            ValueError, match="The objective function must either be encoded as a single QNode"
-        ):
-            opt.step(cost, np.array(0.5, requires_grad=True))
-
-        # defining the device attribute allows it to proceed
-        cost.device = circuit.device
-        new_x = opt.step(cost, np.array(0.5, requires_grad=True))
-
-        assert isinstance(new_x, np.tensor)
-
 
 def ansatz0(x, **kwargs):
     qml.RX(x, wires=0)

From d8d269fbe0bc3f214d9d7aab2e952c02b9f97e29 Mon Sep 17 00:00:00 2001
From: Christina Lee <christina@xanadu.ai>
Date: Tue, 19 Sep 2023 12:57:10 -0400
Subject: [PATCH 077/131] Tests always enable jax float64 (#4613)

This is just a quality-of-life improvement for developers. When running
the tests locally, I always get a bunch of failures not due to my
change, but due to using float32 with jax instead of float64.

This just makes it so that float64 mode is always enabled.
---
 tests/conftest.py                       | 6 ++++++
 tests/devices/test_default_qubit_jax.py | 1 +
 2 files changed, 7 insertions(+)

diff --git a/tests/conftest.py b/tests/conftest.py
index 26eb11463b5..8e0275c37ec 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -221,6 +221,12 @@ def tear_down_thermitian():
     jax_available = False
 
 
+# pylint: disable=unused-argument
+def pytest_generate_tests(metafunc):
+    if jax_available:
+        jax.config.update("jax_enable_x64", True)
+
+
 def pytest_collection_modifyitems(items, config):
     rootdir = pathlib.Path(config.rootdir)
     for item in items:
diff --git a/tests/devices/test_default_qubit_jax.py b/tests/devices/test_default_qubit_jax.py
index c7c3350c720..00f2b24023c 100644
--- a/tests/devices/test_default_qubit_jax.py
+++ b/tests/devices/test_default_qubit_jax.py
@@ -97,6 +97,7 @@ def test_float_precision(self, jax_enable_x64, c_dtype, r_dtype):
         dev = qml.device("default.qubit.jax", wires=2)
         assert dev.state.dtype == c_dtype
         assert dev.state.real.dtype == r_dtype
+        jax.config.update("jax_enable_x64", True)
 
     def test_qubit_circuit(self, tol):
         """Test that the device provides the correct

From de83a59a6ea3b11739eecab366e34432c5fdaae9 Mon Sep 17 00:00:00 2001
From: Matthew Silverman <matthews@xanadu.ai>
Date: Tue, 19 Sep 2023 13:28:55 -0400
Subject: [PATCH 078/131] `process_state` assumes flat, `state()` is always
 complex (#4602)

**Context:**
While porting to DQ2, we're discovering that some features are not at
parity with DQL. These are a few more. `ProbabilityMP.process_state()`
thought the input was not flat! What!!

**Description of the Change:**
- Change `ProbabilityMP.process_state` to assume the input is flat (I
just copy-pasted the contents of `marginal_prob` and used the new
`get_batch_size` helper, I didn't change `_count_samples`, Git just
thinks I did).
- Fix `StateMP.process_state` so it returns the correct shape when
batched
- Fix `StateMP.process_state` so it always returns complex values, even
if the inputs are not complex
- Update the docstring for `process_state` to state (haha) this
assumption

**Benefits:**
- Feature parity with DQL
- Better standardization, cleaner code, better code reuse
---
 doc/releases/changelog-dev.md                 |   9 ++
 pennylane/devices/qubit/sampling.py           |   4 +-
 pennylane/measurements/measurements.py        |   6 +-
 pennylane/measurements/probs.py               | 118 +++++-------------
 pennylane/measurements/state.py               |  10 +-
 tests/drawer/test_tape_mpl.py                 |   2 +-
 .../measurements/legacy/test_expval_legacy.py |   3 +-
 .../measurements/legacy/test_probs_legacy.py  |   3 +-
 tests/measurements/legacy/test_var_legacy.py  |   3 +-
 tests/measurements/test_state.py              |  39 +++++-
 10 files changed, 98 insertions(+), 99 deletions(-)

diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 9b26844ab3e..8963cb3be5a 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -106,6 +106,7 @@
   method will re-order the given state to go from the inputted wire-order to the process's wire-order.
   If the process's wire-order contains extra wires, it will assume those are in the zero-state.
   [(#4570)](https://github.com/PennyLaneAI/pennylane/pull/4570)
+  [(#4602)](https://github.com/PennyLaneAI/pennylane/pull/4602)
 
 * Improve builtin types support with `qml.pauli_decompose`.
   [(#4577)](https://github.com/PennyLaneAI/pennylane/pull/4577)
@@ -130,6 +131,10 @@
   device shots before execution. This makes it compatible with the new device API.
   [(#4599)](https://github.com/PennyLaneAI/pennylane/pull/4599)
 
+* `StateMeasurement.process_state` now assumes the input is flat. `ProbabilityMP.process_state` has
+  been updated to reflect this assumption and avoid redundant reshaping.
+  [(#4602)](https://github.com/PennyLaneAI/pennylane/pull/4602)
+
 <h3>Breaking changes 💔</h3>
 
 * `MeasurementProcess.eigvals()` now raises an `EigvalsUndefinedError` if the measurement observable
@@ -216,6 +221,10 @@
   when the theta value is trainable.
   [(#4585)](https://github.com/PennyLaneAI/pennylane/pull/4585)
 
+* `ProbabilityMP.marginal_prob` has been removed. Its contents have been moved into `process_state`,
+  which effectively just called `marginal_prob` with `np.abs(state) ** 2`.
+  [(#4602)](https://github.com/PennyLaneAI/pennylane/pull/4602)
+
 <h3>Deprecations 👋</h3>
 
 * The ``prep`` keyword argument in ``QuantumScript`` is deprecated and will be removed from `QuantumScript`.
diff --git a/pennylane/devices/qubit/sampling.py b/pennylane/devices/qubit/sampling.py
index 0a2acd9dad1..b1fc9b55060 100644
--- a/pennylane/devices/qubit/sampling.py
+++ b/pennylane/devices/qubit/sampling.py
@@ -27,6 +27,7 @@
 )
 from pennylane.typing import TensorLike
 from .apply_operation import apply_operation
+from .measure import flatten_state
 
 
 def _group_measurements(mps: List[Union[SampleMeasurement, ClassicalShadowMP, ShadowExpvalMP]]):
@@ -342,7 +343,8 @@ def sample_state(
     num_wires = len(wires_to_sample)
     basis_states = np.arange(2**num_wires)
 
-    probs = qml.probs(wires=wires_to_sample).process_state(state, state_wires)
+    flat_state = flatten_state(state, total_indices)
+    probs = qml.probs(wires=wires_to_sample).process_state(flat_state, state_wires)
 
     if is_state_batched:
         # rng.choice doesn't support broadcasting
diff --git a/pennylane/measurements/measurements.py b/pennylane/measurements/measurements.py
index 5096bee9739..e653705817b 100644
--- a/pennylane/measurements/measurements.py
+++ b/pennylane/measurements/measurements.py
@@ -519,7 +519,8 @@ class StateMeasurement(MeasurementProcess):
     Any class inheriting from ``StateMeasurement`` should define its own ``process_state`` method,
     which should have the following arguments:
 
-    * state (Sequence[complex]): quantum state
+    * state (Sequence[complex]): quantum state with a flat shape. It may also have an
+        optional batch dimension
     * wire_order (Wires): wires determining the subspace that ``state`` acts on; a matrix of
         dimension :math:`2^n` acts on a subspace of :math:`n` wires
 
@@ -551,7 +552,8 @@ def process_state(self, state: Sequence[complex], wire_order: Wires):
         """Process the given quantum state.
 
         Args:
-            state (Sequence[complex]): quantum state
+            state (Sequence[complex]): quantum state with a flat shape. It may also have an
+                optional batch dimension
             wire_order (Wires): wires determining the subspace that ``state`` acts on; a matrix of
                 dimension :math:`2^n` acts on a subspace of :math:`n` wires
         """
diff --git a/pennylane/measurements/probs.py b/pennylane/measurements/probs.py
index 17a6bd9c4f9..9b5cbf9a4a2 100644
--- a/pennylane/measurements/probs.py
+++ b/pennylane/measurements/probs.py
@@ -188,93 +188,7 @@ def process_samples(
         return qml.math.squeeze(prob) if bin_size is None else prob
 
     def process_state(self, state: Sequence[complex], wire_order: Wires):
-        num_wires = len(wire_order)
-        batch_size = self._compute_batch_size(state, num_wires)
-        shape = (2**num_wires,) if batch_size is None else (batch_size, 2**num_wires)
-        flat_state = qml.math.reshape(state, shape)
-        real_state = qml.math.real(flat_state)
-        imag_state = qml.math.imag(flat_state)
-        return self.marginal_prob(real_state**2 + imag_state**2, wire_order)
-
-    @staticmethod
-    def _compute_batch_size(state, num_wires):
-        """This logic is isolated because tf.function cannot compile without the try-except
-        clause below (it raises a tf.errors.OperatorNotAllowedInGraphError because it can't
-        handle a function that sometimes returns a tuple, and other times returns None)."""
-
-        expected_shape = [2] * num_wires
-        expected_size = 2**num_wires
-        size = qml.math.size(state)
-
-        try:
-            if qml.math.ndim(state) > len(expected_shape) or size > expected_size:
-                return size // expected_size
-        except Exception as err:  # pylint:disable=broad-except
-            if not qml.math.is_abstract(state):
-                raise err
-
-        return None
-
-    @staticmethod
-    def _count_samples(indices, batch_size, dim):
-        """Count the occurrences of sampled indices and convert them to relative
-        counts in order to estimate their occurrence probability."""
-        num_bins, bin_size = indices.shape[-2:]
-        if batch_size is None:
-            prob = qml.math.zeros((dim, num_bins), dtype="float64")
-            # count the basis state occurrences, and construct the probability vector for each bin
-            for b, idx in enumerate(indices):
-                basis_states, counts = qml.math.unique(idx, return_counts=True)
-                prob[basis_states, b] = counts / bin_size
-
-            return prob
-
-        prob = qml.math.zeros((batch_size, dim, num_bins), dtype="float64")
-        indices = indices.reshape((batch_size, num_bins, bin_size))
-
-        # count the basis state occurrences, and construct the probability vector
-        # for each bin and broadcasting index
-        for i, _indices in enumerate(indices):  # First iterate over broadcasting dimension
-            for b, idx in enumerate(_indices):  # Then iterate over bins dimension
-                basis_states, counts = qml.math.unique(idx, return_counts=True)
-                prob[i, basis_states, b] = counts / bin_size
-
-        return prob
-
-    def marginal_prob(self, prob, wire_order):
-        r"""Return the marginal probability of the computational basis
-        states by summing the probabilities on the non-specified wires.
-
-        If no wires are specified, then all the basis states representable by
-        the device are considered and no marginalization takes place.
-
-        .. note::
-
-            If the provided wires are not in the order as they appear on the device,
-            the returned marginal probabilities take this permutation into account.
-
-            For example, if the addressable wires on this device are ``Wires([0, 1, 2])`` and
-            this function gets passed ``wires=[2, 0]``, then the returned marginal
-            probability vector will take this 'reversal' of the two wires
-            into account:
-
-            .. math::
-
-                \mathbb{P}^{(2, 0)}
-                            = \left[
-                               |00\rangle, |10\rangle, |01\rangle, |11\rangle
-                              \right]
-
-        Args:
-            prob: The probabilities to return the marginal probabilities
-                for
-            wire_order (Iterable[Number, str], Number, str, Wires): wires to return
-                marginal probabilities for. Wires not provided
-                are traced out of the system.
-
-        Returns:
-            array[float]: array of the resulting marginal probabilities.
-        """
+        prob = qml.math.real(state) ** 2 + qml.math.imag(state) ** 2
         if self.wires == Wires([]):
             # no need to marginalize
             return prob
@@ -292,7 +206,9 @@ def marginal_prob(self, prob, wire_order):
         shape = [2] * num_device_wires
         desired_axes = np.argsort(np.argsort(mapped_wires))
         flat_shape = (-1,)
-        if (batch_size := self._compute_batch_size(prob, num_device_wires)) is not None:
+        expected_size = 2**num_device_wires
+        batch_size = qml.math.get_batch_size(prob, (expected_size,), expected_size)
+        if batch_size is not None:
             # prob now is reshaped to have self.num_wires+1 axes in the case of broadcasting
             shape.insert(0, batch_size)
             inactive_wires = [idx + 1 for idx in inactive_wires]
@@ -306,3 +222,29 @@ def marginal_prob(self, prob, wire_order):
         prob = qml.math.transpose(prob, desired_axes)
         # flatten and return probabilities
         return qml.math.reshape(prob, flat_shape)
+
+    @staticmethod
+    def _count_samples(indices, batch_size, dim):
+        """Count the occurrences of sampled indices and convert them to relative
+        counts in order to estimate their occurrence probability."""
+        num_bins, bin_size = indices.shape[-2:]
+        if batch_size is None:
+            prob = qml.math.zeros((dim, num_bins), dtype="float64")
+            # count the basis state occurrences, and construct the probability vector for each bin
+            for b, idx in enumerate(indices):
+                basis_states, counts = qml.math.unique(idx, return_counts=True)
+                prob[basis_states, b] = counts / bin_size
+
+            return prob
+
+        prob = qml.math.zeros((batch_size, dim, num_bins), dtype="float64")
+        indices = indices.reshape((batch_size, num_bins, bin_size))
+
+        # count the basis state occurrences, and construct the probability vector
+        # for each bin and broadcasting index
+        for i, _indices in enumerate(indices):  # First iterate over broadcasting dimension
+            for b, idx in enumerate(_indices):  # Then iterate over bins dimension
+                basis_states, counts = qml.math.unique(idx, return_counts=True)
+                prob[i, basis_states, b] = counts / bin_size
+
+        return prob
diff --git a/pennylane/measurements/state.py b/pennylane/measurements/state.py
index a1a85267fcf..952223ffb1a 100644
--- a/pennylane/measurements/state.py
+++ b/pennylane/measurements/state.py
@@ -158,7 +158,7 @@ def process_state(self, state: Sequence[complex], wire_order: Wires):
         # pylint:disable=redefined-outer-name
         wires = self.wires
         if not wires or wire_order == wires:
-            return state
+            return qml.math.cast(state, "complex128")
 
         if not wires.contains_wires(wire_order):
             raise WireError(
@@ -170,9 +170,12 @@ def process_state(self, state: Sequence[complex], wire_order: Wires):
         pad_width = 2 ** len(wires) - 2 ** len(wire_order)
         pad = (pad_width, 0) if qml.math.get_interface(state) == "torch" else (0, pad_width)
         shape = (2,) * len(wires)
+        flat_shape = (2 ** len(wires),)
         if is_state_batched:
+            batch_size = qml.math.shape(state)[0]
             pad = ((0, 0), pad)
-            shape = (qml.math.shape(state)[0],) + shape
+            shape = (batch_size,) + shape
+            flat_shape = (batch_size,) + flat_shape
         else:
             pad = (pad,)
 
@@ -185,7 +188,8 @@ def process_state(self, state: Sequence[complex], wire_order: Wires):
         if is_state_batched:
             desired_axes = [0] + [i + 1 for i in desired_axes]
         state = qml.math.transpose(state, desired_axes)
-        return qml.math.flatten(state)
+        state = qml.math.reshape(state, flat_shape)
+        return qml.math.cast(state, "complex128")
 
 
 class DensityMatrixMP(StateMP):
diff --git a/tests/drawer/test_tape_mpl.py b/tests/drawer/test_tape_mpl.py
index 3571b450ca3..a029aa48d79 100644
--- a/tests/drawer/test_tape_mpl.py
+++ b/tests/drawer/test_tape_mpl.py
@@ -521,7 +521,7 @@ def test_nested_control_values_bool(self):
 
     def check_tape_controlled_qubit_unitary(self, tape):
         """Checks the control symbols for a tape with some version of a controlled qubit unitary."""
-        _, ax = tape_mpl(tape, style=None)  # set style to None to use plt.rcParams values
+        _, ax = tape_mpl(tape, style="rcParams")  # use plt.rcParams values
         layer = 0
 
         # 5 wires -> 4 control, 1 target
diff --git a/tests/measurements/legacy/test_expval_legacy.py b/tests/measurements/legacy/test_expval_legacy.py
index 68fcecc8368..5ed64cea6d8 100644
--- a/tests/measurements/legacy/test_expval_legacy.py
+++ b/tests/measurements/legacy/test_expval_legacy.py
@@ -17,6 +17,7 @@
 
 import pennylane as qml
 from pennylane.measurements import Expectation, Shots
+from pennylane.devices.qubit.measure import flatten_state
 
 
 # TODO: Remove this when new CustomMP are the default
@@ -24,7 +25,7 @@ def custom_measurement_process(device, spy):
     assert len(spy.call_args_list) > 0  # make sure method is mocked properly
 
     samples = device._samples  # pylint: disable=protected-access
-    state = device._state  # pylint: disable=protected-access
+    state = flatten_state(device._state, device.num_wires)  # pylint: disable=protected-access
     call_args_list = list(spy.call_args_list)
     for call_args in call_args_list:
         obs = call_args.args[1]
diff --git a/tests/measurements/legacy/test_probs_legacy.py b/tests/measurements/legacy/test_probs_legacy.py
index ad08007fcc2..b8eaef7f669 100644
--- a/tests/measurements/legacy/test_probs_legacy.py
+++ b/tests/measurements/legacy/test_probs_legacy.py
@@ -18,6 +18,7 @@
 import pennylane as qml
 from pennylane import numpy as pnp
 from pennylane.measurements import ProbabilityMP, Shots
+from pennylane.devices.qubit.measure import flatten_state
 
 
 # TODO: Remove this when new CustomMP are the default
@@ -25,7 +26,7 @@ def custom_measurement_process(device, spy):
     assert len(spy.call_args_list) > 0  # make sure method is mocked properly
 
     samples = device._samples  # pylint: disable=protected-access
-    state = device._state  # pylint: disable=protected-access
+    state = flatten_state(device._state, device.num_wires)  # pylint: disable=protected-access
     call_args_list = list(spy.call_args_list)
     for call_args in call_args_list:
         wires, shot_range, bin_size = (
diff --git a/tests/measurements/legacy/test_var_legacy.py b/tests/measurements/legacy/test_var_legacy.py
index 5d5ed0ac523..07690eed6d4 100644
--- a/tests/measurements/legacy/test_var_legacy.py
+++ b/tests/measurements/legacy/test_var_legacy.py
@@ -17,6 +17,7 @@
 
 import pennylane as qml
 from pennylane.measurements import Variance, Shots
+from pennylane.devices.qubit.measure import flatten_state
 
 
 # TODO: Remove this when new CustomMP are the default
@@ -25,7 +26,7 @@ def custom_measurement_process(device, spy):
 
     # pylint: disable=protected-access
     samples = device._samples
-    state = device._state
+    state = flatten_state(device._state, device.num_wires)
     call_args_list = list(spy.call_args_list)
     for call_args in call_args_list:
         obs = call_args.args[1]
diff --git a/tests/measurements/test_state.py b/tests/measurements/test_state.py
index 9aca8423e25..a24c648052e 100644
--- a/tests/measurements/test_state.py
+++ b/tests/measurements/test_state.py
@@ -58,7 +58,9 @@ def test_process_state_vector(self, vec):
     def test_state_returns_itself_if_wires_match(self, interface):
         """Test that when wire_order matches the StateMP, the state is returned."""
         ket = qml.math.array([0.48j, 0.48, -0.64j, 0.36], like=interface)
-        assert StateMP(wires=[1, 0]).process_state(ket, wire_order=Wires([1, 0])) is ket
+        assert np.array_equal(
+            StateMP(wires=[1, 0]).process_state(ket, wire_order=Wires([1, 0])), ket
+        )
 
     @pytest.mark.all_interfaces
     @pytest.mark.parametrize("interface", ["numpy", "autograd", "jax", "torch", "tensorflow"])
@@ -91,6 +93,7 @@ def test_expand_state_over_wires(self, mp_wires, expected_state, custom_wire_lab
         mp = StateMP(wires=mp_wires)
         ket = np.arange(1, 5)
         result = mp.process_state(ket, wire_order=Wires(wire_order))
+        assert qml.math.get_dtype_name(result) == "complex128"
         assert np.array_equal(result, expected_state)
 
     @pytest.mark.all_interfaces
@@ -121,6 +124,7 @@ def test_expand_state_batched_all_interfaces(self, interface):
             like=interface,
         )
         result = mp.process_state(ket, wire_order=Wires([1, 2]))
+        assert qml.math.shape(result) == (3, 16)
         reshaped = qml.math.reshape(result, (3, 2, 2, 2, 2))
         assert qml.math.all(reshaped[:, 1, :, 1, :] == 0)
         assert qml.math.allclose(
@@ -157,6 +161,26 @@ def get_state(ket):
         assert qml.math.allclose(result, expected)
         assert isinstance(result, jax.Array)
 
+    @pytest.mark.tf
+    @pytest.mark.parametrize(
+        "wires,expected",
+        [
+            ([1, 0], np.array([0.48j, -0.64j, 0.48, 0.36])),
+            ([2, 1, 0], np.array([0.48j, -0.64j, 0.48, 0.36, 0.0, 0.0, 0.0, 0.0])),
+        ],
+    )
+    def test_state_tf_function(self, wires, expected):
+        """Test that re-ordering and expanding works with tf.function."""
+        import tensorflow as tf
+
+        @tf.function
+        def get_state(ket):
+            return StateMP(wires=wires).process_state(ket, wire_order=Wires([0, 1]))
+
+        result = get_state(tf.Variable([0.48j, 0.48, -0.64j, 0.36]))
+        assert qml.math.allclose(result, expected)
+        assert isinstance(result, tf.Tensor)
+
     def test_wire_ordering_error(self):
         """Test that a wire order error is raised when unknown wires are given."""
         with pytest.raises(WireError, match=r"Unexpected wires \{2\} found in wire order"):
@@ -534,6 +558,19 @@ def func():
 
         assert np.allclose(state_expected, state_val)
 
+    def test_return_type_is_complex(self):
+        """Test that state always returns a complex value."""
+        dev = qml.devices.DefaultQubit()
+
+        @qml.qnode(dev)
+        def func():
+            qml.StatePrep([1, 0, 0, 0], wires=[0, 1])
+            return state()
+
+        state_val = func()
+        assert state_val.dtype == np.complex128
+        assert np.array_equal(state_val, [1, 0, 0, 0])
+
     @pytest.mark.parametrize("shots", [None, 1, 10])
     def test_shape(self, shots):
         """Test that the shape is correct for qml.state."""

From 6eb3d57ad7d18baff034d50d7a2d090b02c1e025 Mon Sep 17 00:00:00 2001
From: lillian542 <38584660+lillian542@users.noreply.github.com>
Date: Tue, 19 Sep 2023 14:05:52 -0400
Subject: [PATCH 079/131] Register ParameterizedEvolution (#4598)

**Context:**
On the old device API, a `ParametrizedEvolution` uses the default
`apply_operation` behaviour if it is preferable to evolve the matrix and
apply (small matrix on a large state), but has special handling for if
the matrix gets bigger (specifically, if the wires for the matrix exceed
half the wires for the state, it evolves the state instead). Adding this
made things much faster and people were very happy.

**Description of the Change:**
Adds the same behaviour here, now based on the size of the state instead
of the number of device wires.

**Benefits:**
Presumably confers the speedups observed on DQL to DQ2.

**Possible Drawbacks:**
I tested whether the correct function was called based on the number of
calls to `math.einsum`, because `mocker.spy` can't access the
`devices.qubit.apply_operation` module due to ambiguity between the
module and the function of the same name. It's the best thing I could
come up with. It does not seem very reliable if anything changes in the
future.

---------

Co-authored-by: Matthew Silverman <matthews@xanadu.ai>
Co-authored-by: Mudit Pandey <mudit.pandey@xanadu.ai>
---
 doc/releases/changelog-dev.md               |   4 +
 pennylane/devices/qubit/apply_operation.py  |  85 +++++++
 pennylane/pulse/parametrized_evolution.py   |   5 +
 pennylane/pulse/parametrized_hamiltonian.py |  20 ++
 tests/devices/qubit/test_apply_operation.py | 258 ++++++++++++++++++++
 tests/pulse/test_parametrized_evolution.py  |  94 ++++++-
 6 files changed, 458 insertions(+), 8 deletions(-)

diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 8963cb3be5a..cccf9bfaf93 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -127,6 +127,10 @@
   when at the beginning of a circuit, thus behaving like `StatePrep`.
   [(#4583)](https://github.com/PennyLaneAI/pennylane/pull/4583)
 
+* DefaultQubit2 dispatches to a faster implementation for applying `ParameterizedEvolution` to a state
+  when it is more efficient to evolve the state than the operation matrix.
+  [(#4598)](https://github.com/PennyLaneAI/pennylane/pull/4598)
+
 * `ShotAdaptiveOptimizer` has been updated to pass shots to QNode executions instead of overriding
   device shots before execution. This makes it compatible with the new device API.
   [(#4599)](https://github.com/PennyLaneAI/pennylane/pull/4599)
diff --git a/pennylane/devices/qubit/apply_operation.py b/pennylane/devices/qubit/apply_operation.py
index 8ec8035a321..6800609e802 100644
--- a/pennylane/devices/qubit/apply_operation.py
+++ b/pennylane/devices/qubit/apply_operation.py
@@ -16,6 +16,7 @@
 
 from functools import singledispatch
 from string import ascii_letters as alphabet
+import numpy as np
 
 import pennylane as qml
 
@@ -194,6 +195,12 @@ def _(op: type_op, state):
         [1., 0.]], requires_grad=True)
 
     """
+    return _apply_operation_default(op, state, is_state_batched, debugger)
+
+
+def _apply_operation_default(op, state, is_state_batched, debugger):
+    """The default behaviour of apply_operation, accessed through the standard dispatch
+    of apply_operation, as well as conditionally in other dispatches."""
     if (
         len(op.wires) < EINSUM_OP_WIRECOUNT_PERF_THRESHOLD
         and math.ndim(state) < EINSUM_STATE_WIRECOUNT_PERF_THRESHOLD
@@ -266,3 +273,81 @@ def apply_snapshot(op: qml.Snapshot, state, is_state_batched: bool = False, debu
         else:
             debugger.snapshots[len(debugger.snapshots)] = flat_state
     return state
+
+
+# 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
+):
+    """Apply ParametrizedEvolution by evolving the state rather than the operator matrix
+    if we are operating on more than half of the subsystem"""
+    if is_state_batched:
+        raise RuntimeError(
+            "ParameterizedEvolution does not support batching, but received a batched state"
+        )
+
+    # shape(state) is static (not a tracer), we can use an if statement
+    num_wires = len(qml.math.shape(state))
+    state = qml.math.cast(state, complex)
+    if 2 * len(op.wires) > num_wires and not op.hyperparameters["complementary"]:
+        # the subsystem operated is more than half of the system based on the state vector --> evolve state
+        return _evolve_state_vector_under_parametrized_evolution(op, state, num_wires)
+    # otherwise --> evolve matrix
+    return _apply_operation_default(op, state, is_state_batched, debugger)
+
+
+def _evolve_state_vector_under_parametrized_evolution(
+    operation: qml.pulse.ParametrizedEvolution, state, num_wires
+):
+    """Uses an odeint solver to compute the evolution of the input ``state`` under the given
+    ``ParametrizedEvolution`` operation.
+
+    Args:
+        state (array[complex]): input state
+        operation (ParametrizedEvolution): operation to apply on the state
+
+    Raises:
+        ValueError: If the parameters and time windows of the ``ParametrizedEvolution`` are
+            not defined.
+
+    Returns:
+        TensorLike[complex]: output state
+    """
+
+    try:
+        import jax
+        from jax.experimental.ode import odeint
+
+        from pennylane.pulse.parametrized_hamiltonian_pytree import ParametrizedHamiltonianPytree
+
+    except ImportError as e:  # pragma: no cover
+        raise ImportError(
+            "Module jax is required for the ``ParametrizedEvolution`` class. "
+            "You can install jax via: pip install jax"
+        ) from e
+
+    if operation.data is None or operation.t is None:
+        raise ValueError(
+            "The parameters and the time window are required to execute a ParametrizedEvolution "
+            "You can update these values by calling the ParametrizedEvolution class: EV(params, t)."
+        )
+
+    state = state.flatten()
+
+    with jax.ensure_compile_time_eval():
+        H_jax = ParametrizedHamiltonianPytree.from_hamiltonian(  # pragma: no cover
+            operation.H,
+            dense=operation.dense,
+            wire_order=list(np.arange(num_wires)),
+        )
+
+    def fun(y, t):
+        """dy/dt = -i H(t) y"""
+        return (-1j * H_jax(operation.data, t=t)) @ y
+
+    result = odeint(fun, state, operation.t, **operation.odeint_kwargs)
+    out_shape = [2] * num_wires
+    if operation.hyperparameters["return_intermediate"]:
+        return qml.math.reshape(result, [-1] + out_shape)
+    return qml.math.reshape(result[-1], out_shape)
diff --git a/pennylane/pulse/parametrized_evolution.py b/pennylane/pulse/parametrized_evolution.py
index f1a4cea2177..8763a866136 100644
--- a/pennylane/pulse/parametrized_evolution.py
+++ b/pennylane/pulse/parametrized_evolution.py
@@ -452,6 +452,11 @@ def _check_time_batching(self):
         # subtract an additional 1 because the full time evolution is not being returned.
         self._batch_size = self.t.shape[0]
 
+    def map_wires(self, wire_map):
+        mapped_op = super().map_wires(wire_map)
+        mapped_op.H = self.H.map_wires(wire_map)
+        return mapped_op
+
     @property
     def hash(self):
         """int: Integer hash that uniquely represents the operator."""
diff --git a/pennylane/pulse/parametrized_hamiltonian.py b/pennylane/pulse/parametrized_hamiltonian.py
index 1326897ce2c..7ec4dcf99c5 100644
--- a/pennylane/pulse/parametrized_hamiltonian.py
+++ b/pennylane/pulse/parametrized_hamiltonian.py
@@ -15,6 +15,8 @@
 """
 This submodule contains the ParametrizedHamiltonian class
 """
+from copy import copy
+
 import pennylane as qml
 from pennylane.operation import Operator
 from pennylane.ops.qubit.hamiltonian import Hamiltonian
@@ -256,6 +258,24 @@ def __repr__(self):
 
         return "  " + "\n+ ".join(terms)
 
+    def map_wires(self, wire_map):
+        """Returns a copy of the current ParametrizedHamiltonian with its wires changed according
+        to the given wire map.
+
+        Args:
+            wire_map (dict): dictionary containing the old wires as keys and the new wires as values
+
+        Returns:
+            .ParametrizedHamiltonian: A new instance with mapped wires
+        """
+        new_ph = copy(self)
+        new_ph.ops_parametrized = [op.map_wires(wire_map) for op in self.ops_parametrized]
+        new_ph.ops_fixed = [op.map_wires(wire_map) for op in self.ops_fixed]
+        new_ph.wires = Wires.all_wires(
+            [op.wires for op in new_ph.ops_fixed] + [op.wires for op in new_ph.ops_parametrized]
+        )
+        return new_ph
+
     def H_fixed(self):
         """The fixed term(s) of the ``ParametrizedHamiltonian``. Returns a ``Sum`` operator of ``SProd``
         operators (or a single ``SProd`` operator in the event that there is only one term in ``H_fixed``).
diff --git a/tests/devices/qubit/test_apply_operation.py b/tests/devices/qubit/test_apply_operation.py
index 1c332bb61c9..9bc1755fc5d 100644
--- a/tests/devices/qubit/test_apply_operation.py
+++ b/tests/devices/qubit/test_apply_operation.py
@@ -14,6 +14,7 @@
 """
 Tests the apply_operation functions from devices/qubit
 """
+from functools import reduce
 import pytest
 
 import numpy as np
@@ -241,6 +242,263 @@ def test_globalphase(self, method, wire, ml_framework):
         assert qml.math.allclose(shift * initial_state, new_state_no_wire)
 
 
+# pylint:disable = unused-argument
+def time_independent_hamiltonian():
+    """Create a time-independent Hamiltonian on two qubits."""
+    ops = [qml.PauliX(0), qml.PauliZ(1), qml.PauliY(0), qml.PauliX(1)]
+
+    def f1(params, t):
+        return params  # constant
+
+    def f2(params, t):
+        return params  # constant
+
+    coeffs = [f1, f2, 4, 9]
+
+    return qml.pulse.ParametrizedHamiltonian(coeffs, ops)
+
+
+def time_dependent_hamiltonian():
+    """Create a time-dependent two-qubit Hamiltonian that takes two scalar parameters."""
+    import jax.numpy as jnp
+
+    ops = [qml.PauliX(0), qml.PauliZ(1), qml.PauliY(0), qml.PauliX(1)]
+
+    def f1(params, t):
+        return params * t
+
+    def f2(params, t):
+        return params * jnp.cos(t)
+
+    coeffs = [f1, f2, 4, 9]
+    return qml.pulse.ParametrizedHamiltonian(coeffs, ops)
+
+
+@pytest.mark.jax
+class TestApplyParameterizedEvolution:
+    @pytest.mark.parametrize("method", methods)
+    def test_parameterized_evolution_time_independent(self, method):
+        """Test that applying a ParameterizedEvolution gives the expected state
+        for a time-independent hamiltonian"""
+
+        import jax.numpy as jnp
+
+        initial_state = np.array(
+            [
+                [0.04624539 + 0.3895457j, 0.22399401 + 0.53870339j],
+                [-0.483054 + 0.2468498j, -0.02772249 - 0.45901669j],
+            ]
+        )
+
+        H = time_independent_hamiltonian()
+        params = jnp.array([1.0, 2.0])
+        t = 4
+
+        op = qml.pulse.ParametrizedEvolution(H=H, params=params, t=t)
+
+        true_mat = qml.math.expm(-1j * qml.matrix(H(params, t=t)) * t)
+        U = qml.QubitUnitary(U=true_mat, wires=[0, 1])
+
+        new_state = method(op, initial_state)
+        new_state_expected = apply_operation(U, initial_state)
+
+        assert np.allclose(new_state, new_state_expected, atol=0.002)
+
+    @pytest.mark.parametrize("method", methods)
+    def test_parameterized_evolution_time_dependent(self, method):
+        """Test that applying a ParameterizedEvolution gives the expected state
+        for a time dependent Hamiltonian"""
+
+        import jax
+        import jax.numpy as jnp
+
+        initial_state = np.array(
+            [
+                [0.04624539 + 0.3895457j, 0.22399401 + 0.53870339j],
+                [-0.483054 + 0.2468498j, -0.02772249 - 0.45901669j],
+            ]
+        )
+
+        H = time_dependent_hamiltonian()
+        params = jnp.array([1.0, 2.0])
+        t = 4
+
+        op = qml.pulse.ParametrizedEvolution(H=H, params=params, t=t)
+
+        def generator(params):
+            time_step = 1e-3
+            times = jnp.arange(0, t, step=time_step)
+            for ti in times:
+                yield jax.scipy.linalg.expm(-1j * time_step * qml.matrix(H(params, t=ti)))
+
+        true_mat = reduce(lambda x, y: y @ x, generator(params))
+        U = qml.QubitUnitary(U=true_mat, wires=[0, 1])
+
+        new_state = method(op, initial_state)
+        new_state_expected = apply_operation(U, initial_state)
+
+        assert np.allclose(new_state, new_state_expected, atol=0.002)
+
+    def test_large_state_small_matrix_evolves_matrix(self, mocker):
+        """Test that applying a ParameterizedEvolution operating on less
+        than half of the wires in the state uses the default function to evolve
+        the matrix"""
+
+        import jax.numpy as jnp
+
+        spy = mocker.spy(qml.math, "einsum")
+
+        initial_state = np.array(
+            [
+                [0.04624539 + 0.3895457j, 0.22399401 + 0.53870339j],
+                [-0.483054 + 0.2468498j, -0.02772249 - 0.45901669j],
+            ]
+        )
+
+        H = time_independent_hamiltonian()
+        params = jnp.array([1.0, 2.0])
+        t = 4
+
+        op = qml.pulse.ParametrizedEvolution(H=H, params=params, t=t)
+
+        true_mat = qml.math.expm(-1j * qml.matrix(H(params, t=t)) * t)
+        U = qml.QubitUnitary(U=true_mat, wires=[0, 1])
+
+        new_state = apply_operation(op, initial_state)
+        new_state_expected = apply_operation(U, initial_state)
+
+        assert np.allclose(new_state, new_state_expected, atol=0.002)
+
+        # seems like _evolve_state_vector_under_parametrized_evolution calls
+        # einsum twice, and the default apply_operation only once
+        assert spy.call_count == 1
+
+    def test_small_evolves_state(self, mocker):
+        """Test that applying a ParameterizedEvolution operating on less
+        than half of the wires in the state uses the default function to evolve
+        the matrix"""
+
+        import jax.numpy as jnp
+
+        spy = mocker.spy(qml.math, "einsum")
+
+        initial_state = np.array(
+            [
+                [
+                    [
+                        [
+                            [-0.02018048 + 0.0j, 0.0 + 0.05690523j],
+                            [0.0 + 0.01425524j, 0.04019714 + 0.0j],
+                        ],
+                        [
+                            [0.0 - 0.07174284j, -0.20230159 + 0.0j],
+                            [-0.05067824 + 0.0j, 0.0 + 0.14290331j],
+                        ],
+                    ],
+                    [
+                        [
+                            [0.0 + 0.05690523j, 0.16046226 + 0.0j],
+                            [0.04019714 + 0.0j, 0.0 - 0.11334853j],
+                        ],
+                        [
+                            [-0.20230159 + 0.0j, 0.0 + 0.57045322j],
+                            [0.0 + 0.14290331j, 0.402961 + 0.0j],
+                        ],
+                    ],
+                ],
+                [
+                    [
+                        [
+                            [0.0 + 0.01425524j, 0.04019714 + 0.0j],
+                            [0.01006972 + 0.0j, 0.0 - 0.02839476j],
+                        ],
+                        [
+                            [-0.05067824 + 0.0j, 0.0 + 0.14290331j],
+                            [0.0 + 0.03579848j, 0.10094511 + 0.0j],
+                        ],
+                    ],
+                    [
+                        [
+                            [0.04019714 + 0.0j, 0.0 - 0.11334853j],
+                            [0.0 - 0.02839476j, -0.08006798 + 0.0j],
+                        ],
+                        [
+                            [0.0 + 0.14290331j, 0.402961 + 0.0j],
+                            [0.10094511 + 0.0j, 0.0 - 0.2846466j],
+                        ],
+                    ],
+                ],
+            ]
+        )
+
+        H = time_independent_hamiltonian()
+        params = jnp.array([1.0, 2.0])
+        t = 4
+
+        op = qml.pulse.ParametrizedEvolution(H=H, params=params, t=t)
+
+        true_mat = qml.math.expm(-1j * qml.matrix(H(params, t=t)) * t)
+        U = qml.QubitUnitary(U=true_mat, wires=[0, 1])
+
+        new_state = apply_operation(op, initial_state)
+        new_state_expected = apply_operation(U, initial_state)
+
+        assert np.allclose(new_state, new_state_expected, atol=0.002)
+
+        # seems like _evolve_state_vector_under_parametrized_evolution calls
+        # einsum twice, and the default apply_operation only once
+        assert spy.call_count == 2
+
+    def test_parametrized_evolution_raises_error(self):
+        """Test applying a ParametrizedEvolution without params or t specified raises an error."""
+        import jax.numpy as jnp
+
+        state = jnp.array([[[1.0, 0.0], [0.0, 0.0]], [[0.0, 0.0], [0.0, 0.0]]], dtype=complex)
+        ev = qml.evolve(qml.pulse.ParametrizedHamiltonian([1], [qml.PauliX("a")]))
+        with pytest.raises(
+            ValueError,
+            match="The parameters and the time window are required to compute the matrix",
+        ):
+            apply_operation(ev, state)
+
+    def test_parametrized_evolution_state_vector_return_intermediate(self, mocker):
+        """Test that when executing a ParametrizedEvolution with ``num_wires >= device.num_wires/2``
+        and ``return_intermediate=True``, the ``_evolve_state_vector_under_parametrized_evolution``
+        method is used."""
+        import jax.numpy as jnp
+
+        H = qml.pulse.ParametrizedHamiltonian([1], [qml.PauliX(0)])
+        spy = mocker.spy(qml.math, "einsum")
+
+        phi = jnp.linspace(0.3, 0.7, 7)
+        phi_for_RX = phi - phi[0]
+        state = jnp.array([[[1.0, 0.0], [0.0, 0.0]], [[0.0, 0.0], [0.0, 0.0]]], dtype=complex)
+        ev = qml.evolve(H, return_intermediate=True)(params=[], t=phi / 2)
+        state_ev = apply_operation(ev, state)
+        state_rx = apply_operation(qml.RX(phi_for_RX, 0), state)
+
+        assert spy.call_count == 2
+        assert qml.math.allclose(state_ev, state_rx, atol=1e-6)
+
+    def test_batched_state_raises_an_error(self):
+        """Test that if is_state_batche=True, an error is raised"""
+        H = time_independent_hamiltonian()
+        params = np.array([1.0, 2.0])
+        t = 4
+
+        op = qml.pulse.ParametrizedEvolution(H=H, params=params, t=t)
+
+        initial_state = np.array(
+            [
+                [[0.81677345 + 0.0j, 0.0 + 0.0j], [0.0 - 0.57695852j, 0.0 + 0.0j]],
+                [[0.33894597 + 0.0j, 0.0 + 0.0j], [0.0 - 0.94080584j, 0.0 + 0.0j]],
+            ]
+        )
+
+        with pytest.raises(RuntimeError, match="does not support batching"):
+            _ = apply_operation(op, initial_state, is_state_batched=True)
+
+
 @pytest.mark.parametrize("ml_framework", ml_frameworks_list)
 class TestSnapshot:
     """Test that apply_operation works for Snapshot ops"""
diff --git a/tests/pulse/test_parametrized_evolution.py b/tests/pulse/test_parametrized_evolution.py
index 29d153a9189..710e76345b2 100644
--- a/tests/pulse/test_parametrized_evolution.py
+++ b/tests/pulse/test_parametrized_evolution.py
@@ -24,6 +24,7 @@
 from pennylane.ops import QubitUnitary
 from pennylane.pulse import ParametrizedEvolution, ParametrizedHamiltonian
 from pennylane.tape import QuantumTape
+from pennylane.devices import DefaultQubit, DefaultQubitLegacy
 
 
 class MyOp(qml.RX):  # pylint: disable=too-few-public-methods
@@ -476,10 +477,11 @@ def test_return_intermediate_and_complementary(self, comp, len_t):
 class TestIntegration:
     """Integration tests for the ParametrizedEvolution class."""
 
+    @pytest.mark.parametrize("device_class", [DefaultQubit, DefaultQubitLegacy])
     @pytest.mark.parametrize("time", [0.3, 1, [0, 2], [0.4, 2], (3, 3.1)])
     @pytest.mark.parametrize("time_interface", ["python", "numpy", "jax"])
     @pytest.mark.parametrize("use_jit", [False, True])
-    def test_time_input_formats(self, time, time_interface, use_jit):
+    def test_time_input_formats(self, device_class, time, time_interface, use_jit):
         import jax
         import jax.numpy as jnp
 
@@ -489,7 +491,7 @@ def test_time_input_formats(self, time, time_interface, use_jit):
             time = np.array(time)
         H = qml.pulse.ParametrizedHamiltonian([2], [qml.PauliX(0)])
 
-        dev = qml.device("default.qubit", wires=1)
+        dev = device_class(wires=1)
 
         @qml.qnode(dev, interface="jax")
         def circuit(t):
@@ -503,15 +505,16 @@ def circuit(t):
         duration = time if qml.math.ndim(time) == 0 else time[1] - time[0]
         assert qml.math.isclose(res, qml.math.cos(4 * duration))
 
+    @pytest.mark.parametrize("device_class", [DefaultQubit, DefaultQubitLegacy])
     # pylint: disable=unused-argument
-    def test_time_independent_hamiltonian(self):
+    def test_time_independent_hamiltonian(self, device_class):
         """Test the execution of a time independent hamiltonian."""
         import jax
         import jax.numpy as jnp
 
         H = time_independent_hamiltonian()
 
-        dev = qml.device("default.qubit", wires=2)
+        dev = device_class(wires=2)
 
         t = 4
 
@@ -543,8 +546,9 @@ def true_circuit(params):
             jax.grad(jitted_circuit)(params), jax.grad(true_circuit)(params), atol=1e-3
         )
 
+    @pytest.mark.parametrize("device_class", [DefaultQubit, DefaultQubitLegacy])
     @pytest.mark.slow
-    def test_time_dependent_hamiltonian(self):
+    def test_time_dependent_hamiltonian(self, device_class):
         """Test the execution of a time dependent hamiltonian. This test approximates the
         time-ordered exponential with a product of exponentials using small time steps.
         For more information, see https://en.wikipedia.org/wiki/Ordered_exponential."""
@@ -553,7 +557,7 @@ def test_time_dependent_hamiltonian(self):
 
         H = time_dependent_hamiltonian()
 
-        dev = qml.device("default.qubit", wires=2)
+        dev = device_class(wires=2)
         t = 2
 
         def generator(params):
@@ -590,7 +594,56 @@ def true_circuit(params):
             jax.grad(jitted_circuit)(params), jax.grad(true_circuit)(params), atol=5e-3
         )
 
-    def test_two_commuting_parametrized_hamiltonians(self):
+    @pytest.mark.slow
+    def test_map_wires_with_time_independent_hamiltonian(self):
+        """Test the wire mapping for custom wire labels works as expected with DefaultQubit"""
+        import jax
+        from jax import numpy as jnp
+
+        def f1(params, t):
+            return params  # constant
+
+        def f2(params, t):
+            return params  # constant
+
+        ops = [qml.PauliX("a"), qml.PauliZ("b"), qml.PauliY("a"), qml.PauliX("b")]
+        coeffs = [f1, f2, 4, 9]
+        H = ParametrizedHamiltonian(coeffs, ops)
+
+        dev = DefaultQubit()
+
+        t = 4
+
+        @qml.qnode(dev)
+        def circuit(params):
+            ParametrizedEvolution(H=H, params=params, t=t)
+            return qml.expval(qml.PauliX("a") @ qml.PauliX("b"))
+
+        @jax.jit
+        @qml.qnode(dev)
+        def jitted_circuit(params):
+            ParametrizedEvolution(H=H, params=params, t=t)
+            return qml.expval(qml.PauliX("a") @ qml.PauliX("b"))
+
+        @qml.qnode(dev)
+        def true_circuit(params):
+            true_mat = qml.math.expm(-1j * qml.matrix(H(params, t=t)) * t)
+            QubitUnitary(U=true_mat, wires=[0, 1])
+            return qml.expval(qml.PauliX(0) @ qml.PauliX(1))
+
+        params = jnp.array([1.0, 2.0])
+
+        assert qml.math.allclose(circuit(params), true_circuit(params), atol=1e-3)
+        assert qml.math.allclose(jitted_circuit(params), true_circuit(params), atol=1e-3)
+        assert qml.math.allclose(
+            jax.grad(circuit)(params), jax.grad(true_circuit)(params), atol=1e-3
+        )
+        assert qml.math.allclose(
+            jax.grad(jitted_circuit)(params), jax.grad(true_circuit)(params), atol=1e-3
+        )
+
+    @pytest.mark.parametrize("device_class", [DefaultQubit, DefaultQubitLegacy])
+    def test_two_commuting_parametrized_hamiltonians(self, device_class):
         """Test that the evolution of two parametrized hamiltonians that commute with each other
         is equal to evolve the two hamiltonians simultaneously. This test uses 8 wires for the device
         to test the case where 2 * n < N (the matrix is evolved instead of the state)."""
@@ -614,7 +667,7 @@ def f3(p, t):
         ops = [qml.PauliX(0), qml.PauliX(2)]
         H2 = qml.dot(coeffs, ops)
 
-        dev = qml.device("default.qubit", wires=8)
+        dev = device_class(wires=8)
 
         @jax.jit
         @qml.qnode(dev, interface="jax")
@@ -704,3 +757,28 @@ def U(params):
         jac_jit = jax.jacobian(jax.jit(U), holomorphic=True)(params)
 
         assert qml.math.allclose(jac, jac_jit)
+
+
+@pytest.mark.jax
+def test_map_wires():
+    """Test that map wires returns a new ParameterizedEvolution, with wires updated on
+    both the operator and the corresponding Hamiltonian"""
+
+    def f1(p, t):
+        return p * t
+
+    coeffs = [2, 4, f1]
+    ops = [qml.PauliX("a"), qml.PauliX("b"), qml.PauliX("c")]
+
+    H = qml.dot(coeffs, ops)
+
+    op = qml.evolve(H)([3], 2)
+
+    wire_map = {"a": 3, "b": 5, "c": 7}
+    new_op = op.map_wires(wire_map)
+
+    assert op.wires == qml.wires.Wires(["a", "b", "c"])
+    assert op.H.wires == qml.wires.Wires(["a", "b", "c"])
+
+    assert new_op.wires == qml.wires.Wires([3, 5, 7])
+    assert new_op.H.wires == qml.wires.Wires([3, 5, 7])

From a7c129409717a4afed805bf1aec973ee01d99074 Mon Sep 17 00:00:00 2001
From: Mudit Pandey <mudit.pandey@xanadu.ai>
Date: Tue, 19 Sep 2023 14:23:32 -0400
Subject: [PATCH 080/131] Added docs to measurements.rst

---
 doc/introduction/measurements.rst | 25 +++++++++++++++++++++++++
 1 file changed, 25 insertions(+)

diff --git a/doc/introduction/measurements.rst b/doc/introduction/measurements.rst
index bd5323d17b8..7ea6807c020 100644
--- a/doc/introduction/measurements.rst
+++ b/doc/introduction/measurements.rst
@@ -358,6 +358,31 @@ thrown away.
     Currently, postselection support is only available on ``"default.qubit"``. Requesting postselection on other
     devices will not do anything.
 
+Statistics can also be collected on mid-circuit measurements along with terminal measurement statistics.
+Currently, ``qml.probs``, ``qml.sample``, ``qml.expval``, ``qml.var``, and ``qml.counts`` are supported,
+and can be requested along with other measurements. The devices that currently support collecting such
+statistics are ``"default.qubit"``, ``"default.mixed"``, and ``"default.qubit.legacy"``.
+
+..code-block:: python3
+
+    dev = qml.devices.DefaultQubit()
+
+    @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.probs(wires=1), qml.probs(op=m0)
+
+Executing this QNode:
+
+>>> func(np.pi / 2, np.pi / 4)
+(tensor([0.9267767, 0.0732233], requires_grad=True),
+ tensor([0.5, 0.5], requires_grad=True))
+
+Currently, statistics can only be collected for single mid-circuit measurement values. Moreover, any
+measurement values manipulated using boolean or arithmetic operators cannot be used. These can lead to
+unexpected/incorrect behaviour.
 
 The deferred measurement principle provides a natural method to simulate the
 application of mid-circuit measurements and conditional operations in a

From 4bfcac5219b2af790885549f8f3c67af6258acc4 Mon Sep 17 00:00:00 2001
From: Mudit Pandey <mudit.pandey@xanadu.ai>
Date: Tue, 19 Sep 2023 14:31:05 -0400
Subject: [PATCH 081/131] Added docs for nan values

---
 doc/introduction/measurements.rst | 17 +++++++++++++++++
 1 file changed, 17 insertions(+)

diff --git a/doc/introduction/measurements.rst b/doc/introduction/measurements.rst
index 7ea6807c020..38d1195430a 100644
--- a/doc/introduction/measurements.rst
+++ b/doc/introduction/measurements.rst
@@ -353,6 +353,23 @@ array([1, 1, 1, 1, 1, 1, 1])
 Note that only 7 samples are returned. This is because samples that do not meet the postselection criteria are
 thrown away.
 
+If postselection is requested on a state with zero probability of being measured, ``NaN`` will be returned regardless
+of what and how many measurements are made:
+
+.. code-block:: python3
+
+    dev = qml.device("default.qubit", wires=2)
+
+    @qml.qnode(dev)
+    def func(x):
+        qml.RX(x, wires=0)
+        m0 = qml.measure(0, postselect=1)
+        qml.cond(m0, qml.PauliX)(wires=1)
+        return qml.probs(wires=1)
+
+>>> func(0.0)
+nan
+
 .. note::
 
     Currently, postselection support is only available on ``"default.qubit"``. Requesting postselection on other

From 25b9eff48ca1dbc21f0c19fcbe17d04801ba96f4 Mon Sep 17 00:00:00 2001
From: Mudit Pandey <mudit.pandey@xanadu.ai>
Date: Tue, 19 Sep 2023 15:55:37 -0400
Subject: [PATCH 082/131] Fixed doc error

---
 doc/introduction/measurements.rst | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/doc/introduction/measurements.rst b/doc/introduction/measurements.rst
index 38d1195430a..da14363f049 100644
--- a/doc/introduction/measurements.rst
+++ b/doc/introduction/measurements.rst
@@ -380,7 +380,7 @@ Currently, ``qml.probs``, ``qml.sample``, ``qml.expval``, ``qml.var``, and ``qml
 and can be requested along with other measurements. The devices that currently support collecting such
 statistics are ``"default.qubit"``, ``"default.mixed"``, and ``"default.qubit.legacy"``.
 
-..code-block:: python3
+.. code-block:: python3
 
     dev = qml.devices.DefaultQubit()
 

From cdf53d498b2172535e155df166c32a6b9d3566bb Mon Sep 17 00:00:00 2001
From: Mudit Pandey <mudit.pandey@xanadu.ai>
Date: Wed, 20 Sep 2023 10:26:51 -0400
Subject: [PATCH 083/131] Updated measurement process; updated tests

---
 pennylane/_qubit_device.py                    |  3 ++
 pennylane/devices/default_qubit_legacy.py     |  1 -
 pennylane/devices/qubit/measure.py            |  2 +-
 pennylane/measurements/counts.py              | 10 ++---
 pennylane/measurements/expval.py              |  3 +-
 pennylane/measurements/measurements.py        | 38 +++++++++++--------
 pennylane/measurements/sample.py              |  4 +-
 pennylane/measurements/var.py                 |  9 ++++-
 .../measurements/legacy/test_expval_legacy.py | 24 ++++++++++++
 .../measurements/legacy/test_probs_legacy.py  | 31 +++++++++++++++
 .../measurements/legacy/test_sample_legacy.py | 29 +++++++++++++-
 tests/measurements/legacy/test_var_legacy.py  | 25 ++++++++++++
 12 files changed, 150 insertions(+), 29 deletions(-)

diff --git a/pennylane/_qubit_device.py b/pennylane/_qubit_device.py
index a9a1972b881..67e68f4b0b9 100644
--- a/pennylane/_qubit_device.py
+++ b/pennylane/_qubit_device.py
@@ -616,6 +616,9 @@ def statistics(
             if m.obs is not None:
                 obs = m.obs
                 obs.return_type = m.return_type
+            elif m.mv is not None:
+                obs = m.mv
+                obs.return_type = m.return_type
             else:
                 obs = m
             # Check if there is an overriden version of the measurement process
diff --git a/pennylane/devices/default_qubit_legacy.py b/pennylane/devices/default_qubit_legacy.py
index d0461886b1f..dc4983933c3 100644
--- a/pennylane/devices/default_qubit_legacy.py
+++ b/pennylane/devices/default_qubit_legacy.py
@@ -195,7 +195,6 @@ class DefaultQubitLegacy(QubitDevice):
         "Prod",
         "Exp",
         "Evolution",
-        "MeasurementValue",
     }
 
     def __init__(
diff --git a/pennylane/devices/qubit/measure.py b/pennylane/devices/qubit/measure.py
index 3bc16d9b8d0..9b72def187d 100644
--- a/pennylane/devices/qubit/measure.py
+++ b/pennylane/devices/qubit/measure.py
@@ -150,7 +150,7 @@ def get_measurement_function(
         Callable: function that returns the measurement result
     """
     if isinstance(measurementprocess, StateMeasurement):
-        if isinstance(measurementprocess.obs, MeasurementValue):
+        if isinstance(measurementprocess.mv, MeasurementValue):
             return state_diagonalizing_gates
 
         if isinstance(measurementprocess, ExpectationMP):
diff --git a/pennylane/measurements/counts.py b/pennylane/measurements/counts.py
index 83e17a43b4b..9b5dc76ad21 100644
--- a/pennylane/measurements/counts.py
+++ b/pennylane/measurements/counts.py
@@ -214,12 +214,12 @@ def process_samples(
         shot_range: Tuple[int] = None,
         bin_size: int = None,
     ):
-        if not isinstance(self.obs, MeasurementValue):
-            samples = qml.sample(op=self.obs, wires=self._wires).process_samples(
+        if self.mv is not None:
+            samples = qml.sample(wires=self.mv.wires).process_samples(
                 samples, wire_order, shot_range, bin_size
             )
         else:
-            samples = qml.sample(wires=self.obs.wires).process_samples(
+            samples = qml.sample(op=self.obs, wires=self._wires).process_samples(
                 samples, wire_order, shot_range, bin_size
             )
 
@@ -229,7 +229,7 @@ def process_samples(
         num_wires = len(self.wires) if self.wires else len(wire_order)
         samples = (
             samples.reshape((num_wires, -1)).T.reshape(-1, bin_size, num_wires)
-            if self.obs is None or isinstance(self.obs, MeasurementValue)
+            if self.obs is None
             else samples.reshape((-1, bin_size))
         )
 
@@ -287,7 +287,7 @@ def circuit(x):
         batched_ndims = 2
         shape = qml.math.shape(samples)
 
-        if self.obs is None or isinstance(self.obs, MeasurementValue):
+        if self.obs is None:
             # convert samples and outcomes (if using) from arrays to str for dict keys
             samples = qml.math.cast_like(samples, qml.math.int8(0))
             samples = qml.math.apply_along_axis(_sample_to_str, -1, samples)
diff --git a/pennylane/measurements/expval.py b/pennylane/measurements/expval.py
index d9e60d3c232..f7e474f13f0 100644
--- a/pennylane/measurements/expval.py
+++ b/pennylane/measurements/expval.py
@@ -111,7 +111,8 @@ def process_samples(
             return probs[idx]
 
         # estimate the ev
-        samples = qml.sample(op=self.obs).process_samples(
+        op = self.mv if self.mv is not None else self.obs
+        samples = qml.sample(op=op).process_samples(
             samples=samples, wire_order=wire_order, shot_range=shot_range, bin_size=bin_size
         )
 
diff --git a/pennylane/measurements/measurements.py b/pennylane/measurements/measurements.py
index e653705817b..56e9227d193 100644
--- a/pennylane/measurements/measurements.py
+++ b/pennylane/measurements/measurements.py
@@ -133,7 +133,13 @@ def __init__(
         eigvals=None,
         id: Optional[str] = None,
     ):
-        self.obs = obs
+        if obs is not None and obs.name == "MeasurementValue":
+            self.mv = obs
+            self.obs = None
+        else:
+            self.obs = obs
+            self.mv = None
+
         self.id = id
 
         if wires is not None:
@@ -291,6 +297,9 @@ def wires(self):
 
         This is the union of all the Wires objects of the measurement.
         """
+        if self.mv is not None:
+            return self.mv.wires
+
         if self.obs is not None:
             return self.obs.wires
 
@@ -330,7 +339,7 @@ def eigvals(self):
         Returns:
             array: eigvals representation
         """
-        if self.obs.__class__.__name__ == "MeasurementValue":
+        if self.mv is not None:
             return qml.math.arange(0, 2 ** len(self.wires), 1)
 
         if self.obs is not None:
@@ -345,16 +354,12 @@ def has_decomposition(self):
         # If self.obs is not None, `expand` queues the diagonalizing gates of self.obs,
         # which we have to check to be defined. The subsequent creation of the new
         # `MeasurementProcess` within `expand` should never fail with the given parameters.
-        return (
-            self.obs.has_diagonalizing_gates
-            if self.obs is not None and self.obs.__class__.__name__ != "MeasurementValue"
-            else False
-        )
+        return self.obs.has_diagonalizing_gates if self.obs is not None else False
 
     @property
     def samples_computational_basis(self):
         r"""Bool: Whether or not the MeasurementProcess measures in the computational basis."""
-        return self.obs is None or self.obs.__class__.__name__ == "MeasurementValue"
+        return self.obs is None
 
     def expand(self):
         """Expand the measurement of an observable to a unitary
@@ -389,7 +394,7 @@ def expand(self):
         >>> print(tape.measurements[0].obs)
         None
         """
-        if self.obs is None or self.obs.__class__.__name__ == "MeasurementValue":
+        if self.obs is None:
             raise qml.operation.DecompositionUndefinedError
 
         with qml.queuing.AnnotatedQueue() as q:
@@ -400,7 +405,7 @@ def expand(self):
 
     def queue(self, context=qml.QueuingManager):
         """Append the measurement process to an annotated queue."""
-        if self.obs is not None and self.obs.__class__.__name__ != "MeasurementValue":
+        if self.obs is not None:
             context.remove(self.obs)
         context.append(self)
 
@@ -422,6 +427,7 @@ def hash(self):
         fingerprint = (
             self.__class__.__name__,
             getattr(self.obs, "hash", "None"),
+            getattr(self.mv, "hash", "None"),
             str(self._eigvals),  # eigvals() could be expensive to compute for large observables
             tuple(self.wires.tolist()),
         )
@@ -434,11 +440,7 @@ def simplify(self):
         Returns:
             .MeasurementProcess: A measurement process with a simplified observable.
         """
-        return (
-            self
-            if self.obs is None or self.obs.__class__.__name__ == "MeasurementValue"
-            else self.__class__(obs=self.obs.simplify())
-        )
+        return self if self.obs is None else self.__class__(obs=self.obs.simplify())
 
     # pylint: disable=protected-access
     def map_wires(self, wire_map: dict):
@@ -452,7 +454,11 @@ def map_wires(self, wire_map: dict):
             .MeasurementProcess: new measurement process
         """
         new_measurement = copy.copy(self)
-        if self.obs is not None and self.obs.__class__.__name__ != "MeasurementValue":
+        if self.mv is not None:
+            mv = copy.copy(self.mv)
+            mv.measurements = [m.map_wires(wire_map=wire_map) for m in mv.measurements]
+            new_measurement.mv = mv
+        if self.obs is not None:
             new_measurement.obs = self.obs.map_wires(wire_map=wire_map)
         elif self._wires is not None:
             new_measurement._wires = Wires([wire_map.get(wire, wire) for wire in self.wires])
diff --git a/pennylane/measurements/sample.py b/pennylane/measurements/sample.py
index dfaeb3ed805..529cc621ff8 100644
--- a/pennylane/measurements/sample.py
+++ b/pennylane/measurements/sample.py
@@ -149,7 +149,7 @@ def return_type(self):
     def numeric_type(self):
         # Note: we only assume an integer numeric type if the observable is a
         # built-in observable with integer eigenvalues or a tensor product thereof
-        if self.obs is None or isinstance(self.obs, MeasurementValue):
+        if self.obs is None:
             # Computational basis samples
             return int
         int_eigval_obs = {qml.PauliX, qml.PauliY, qml.PauliZ, qml.Hadamard, qml.Identity}
@@ -207,7 +207,7 @@ def process_samples(
 
         num_wires = samples.shape[-1]  # wires is the last dimension
 
-        if self.obs is None:
+        if self.obs is None and self.mv is None:
             # if no observable was provided then return the raw samples
             return samples if bin_size is None else samples.T.reshape(num_wires, bin_size, -1)
 
diff --git a/pennylane/measurements/var.py b/pennylane/measurements/var.py
index f3ab27cd4f4..cb714654b9c 100644
--- a/pennylane/measurements/var.py
+++ b/pennylane/measurements/var.py
@@ -113,9 +113,16 @@ def process_samples(
             return probs[idx] - probs[idx] ** 2
 
         # estimate the variance
-        samples = qml.sample(op=self.obs).process_samples(
+        op = self.mv if self.mv is not None else self.obs
+        samples = qml.sample(op=op).process_samples(
             samples=samples, wire_order=wire_order, shot_range=shot_range, bin_size=bin_size
         )
+        # if self.mv is not None:
+        #     num_wires = samples.shape[-1]
+        #     powers_of_two = 2 ** qml.math.arange(num_wires)[::-1]
+        #     indices = samples @ powers_of_two
+        #     indices = qml.math.array(indices)  # Add np.array here for Jax support.
+        #     samples = self.eigvals()[indices]
 
         # With broadcasting, we want to take the variance over axis 1, which is the -1st/-2nd with/
         # without bin_size. Without broadcasting, axis 0 is the -1st/-2nd with/without bin_size
diff --git a/tests/measurements/legacy/test_expval_legacy.py b/tests/measurements/legacy/test_expval_legacy.py
index 5ed64cea6d8..bde8fe72164 100644
--- a/tests/measurements/legacy/test_expval_legacy.py
+++ b/tests/measurements/legacy/test_expval_legacy.py
@@ -115,6 +115,30 @@ def circuit():
 
         custom_measurement_process(new_dev, spy)
 
+    @pytest.mark.parametrize("shots", [None, 10000, [10000, 10000]])
+    @pytest.mark.parametrize("phi", np.arange(0, 2 * np.pi, np.pi / 3))
+    def test_observable_is_measurement_value(
+        self, shots, phi, mocker, tol, tol_stochastic
+    ):  # pylint: disable=too-many-arguments
+        """Test that expectation values for mid-circuit measurement values
+        are correct for a single measurement value."""
+        dev = qml.device("default.qubit", wires=2, shots=shots)
+
+        @qml.qnode(dev)
+        def circuit(phi):
+            qml.RX(phi, 0)
+            m0 = qml.measure(0)
+            return qml.expval(m0)
+
+        new_dev = circuit.device
+        spy = mocker.spy(qml.QubitDevice, "expval")
+
+        res = circuit(phi)
+
+        atol = tol if shots is None else tol_stochastic
+        assert np.allclose(np.array(res), np.sin(phi / 2) ** 2, atol=atol, rtol=0)
+        custom_measurement_process(new_dev, spy)
+
     @pytest.mark.parametrize(
         "obs",
         [qml.PauliZ(0), qml.Hermitian(np.diag([1, 2]), 0), qml.Hermitian(np.diag([1.0, 2.0]), 0)],
diff --git a/tests/measurements/legacy/test_probs_legacy.py b/tests/measurements/legacy/test_probs_legacy.py
index b8eaef7f669..91bf12a1b40 100644
--- a/tests/measurements/legacy/test_probs_legacy.py
+++ b/tests/measurements/legacy/test_probs_legacy.py
@@ -173,6 +173,37 @@ def circuit():
 
         custom_measurement_process(dev, spy_probs)
 
+    @pytest.mark.parametrize("shots", [None, 10000, [10000, 10000]])
+    @pytest.mark.parametrize("phi", np.arange(0, 2 * np.pi, np.pi / 3))
+    def test_observable_is_measurement_value(
+        self, shots, phi, mocker, tol, tol_stochastic
+    ):  # pylint: disable=too-many-arguments
+        """Test that expectation values for mid-circuit measurement values
+        are correct for a single measurement value."""
+        dev = qml.device("default.qubit", wires=2, shots=shots)
+
+        @qml.qnode(dev)
+        def circuit(phi):
+            qml.RX(phi, 0)
+            m0 = qml.measure(0)
+            return qml.probs(op=m0)
+
+        new_dev = circuit.device
+        spy = mocker.spy(qml.QubitDevice, "probability")
+
+        res = circuit(phi)
+
+        atol = tol if shots is None else tol_stochastic
+        expected = np.array([np.cos(phi / 2) ** 2, np.sin(phi / 2) ** 2])
+
+        if not isinstance(shots, list):
+            assert np.allclose(np.array(res), expected, atol=atol, rtol=0)
+        else:
+            for r in res:  # pylint: disable=not-an-iterable
+                assert np.allclose(r, expected, atol=atol, rtol=0)
+
+        custom_measurement_process(new_dev, spy)
+
     def test_integration(self, tol, mocker):
         """Test the probability is correct for a known state preparation."""
         dev = qml.device("default.qubit.legacy", wires=2)
diff --git a/tests/measurements/legacy/test_sample_legacy.py b/tests/measurements/legacy/test_sample_legacy.py
index ee8bb3bca01..9630ad518bc 100644
--- a/tests/measurements/legacy/test_sample_legacy.py
+++ b/tests/measurements/legacy/test_sample_legacy.py
@@ -16,7 +16,7 @@
 import pytest
 
 import pennylane as qml
-from pennylane.measurements import MeasurementShapeError, Sample, Shots
+from pennylane.measurements import MeasurementShapeError, Sample, Shots, MeasurementValue
 from pennylane.operation import Operator
 
 # pylint: disable=protected-access, no-member
@@ -31,7 +31,7 @@ def custom_measurement_process(device, spy):
     for call_args in call_args_list:
         meas = call_args.args[1]
         shot_range, bin_size = (call_args.kwargs["shot_range"], call_args.kwargs["bin_size"])
-        if isinstance(meas, Operator):
+        if isinstance(meas, (Operator, MeasurementValue)):
             meas = qml.sample(op=meas)
         assert qml.math.allequal(
             device.sample(call_args.args[1], **call_args.kwargs),
@@ -95,6 +95,31 @@ def circuit():
 
         custom_measurement_process(dev, spy)
 
+    @pytest.mark.parametrize("shots", [5, [5, 5]])
+    @pytest.mark.parametrize("phi", np.arange(0, 2 * np.pi, np.pi / 2))
+    def test_observable_is_measurement_value(self, shots, phi, mocker):
+        """Test that expectation values for mid-circuit measurement values
+        are correct for a single measurement value."""
+        dev = qml.device("default.qubit", wires=2, shots=shots)
+
+        @qml.qnode(dev)
+        def circuit(phi):
+            qml.RX(phi, 0)
+            m0 = qml.measure(0)
+            return qml.sample(m0)
+
+        new_dev = circuit.device
+        spy = mocker.spy(qml.QubitDevice, "sample")
+
+        res = circuit(phi)
+
+        if isinstance(shots, list):
+            assert len(res) == len(shots)
+            assert all(r.shape == (s,) for r, s in zip(res, shots))
+        else:
+            assert res.shape == (shots,)
+        custom_measurement_process(new_dev, spy)
+
     def test_single_wire_sample(self, mocker):
         """Test the return type and shape of sampling a single wire"""
         n_sample = 10
diff --git a/tests/measurements/legacy/test_var_legacy.py b/tests/measurements/legacy/test_var_legacy.py
index 07690eed6d4..43eb58ff2d5 100644
--- a/tests/measurements/legacy/test_var_legacy.py
+++ b/tests/measurements/legacy/test_var_legacy.py
@@ -108,6 +108,31 @@ def circuit():
 
         custom_measurement_process(dev, spy)
 
+    @pytest.mark.parametrize("shots", [None, 10000, [10000, 10000]])
+    @pytest.mark.parametrize("phi", np.arange(0, 2 * np.pi, np.pi / 3))
+    def test_observable_is_measurement_value(
+        self, shots, phi, mocker, tol, tol_stochastic
+    ):  # pylint: disable=too-many-arguments
+        """Test that variances for mid-circuit measurement values
+        are correct for a single measurement value."""
+        dev = qml.device("default.qubit", wires=2, shots=shots)
+
+        @qml.qnode(dev)
+        def circuit(phi):
+            qml.RX(phi, 0)
+            m0 = qml.measure(0)
+            return qml.var(m0)
+
+        new_dev = circuit.device
+        spy = mocker.spy(qml.QubitDevice, "var")
+
+        res = circuit(phi)
+
+        atol = tol if shots is None else tol_stochastic
+        expected = np.sin(phi / 2) ** 2 - np.sin(phi / 2) ** 4
+        assert np.allclose(np.array(res), expected, atol=atol, rtol=0)
+        custom_measurement_process(new_dev, spy)
+
     @pytest.mark.parametrize(
         "obs",
         [qml.PauliZ(0), qml.Hermitian(np.diag([1, 2]), 0), qml.Hermitian(np.diag([1.0, 2.0]), 0)],

From 1b4a7ae0882b665a964a645c611c7d0dcf922f14 Mon Sep 17 00:00:00 2001
From: Mudit Pandey <mudit.pandey@xanadu.ai>
Date: Wed, 20 Sep 2023 13:38:39 -0400
Subject: [PATCH 084/131] Update tests

---
 pennylane/_qubit_device.py                    |  6 +-
 .../measurements/legacy/test_counts_legacy.py | 60 +++++++++++++++++++
 .../measurements/legacy/test_expval_legacy.py |  2 +-
 .../measurements/legacy/test_probs_legacy.py  |  4 +-
 .../measurements/legacy/test_sample_legacy.py |  4 +-
 tests/measurements/legacy/test_var_legacy.py  |  2 +-
 tests/measurements/test_probs.py              |  2 +-
 tests/measurements/test_sample.py             |  2 +-
 8 files changed, 73 insertions(+), 9 deletions(-)

diff --git a/pennylane/_qubit_device.py b/pennylane/_qubit_device.py
index 67e68f4b0b9..e1354655e3f 100644
--- a/pennylane/_qubit_device.py
+++ b/pennylane/_qubit_device.py
@@ -1442,7 +1442,11 @@ def circuit(x):
             if obs.all_outcomes:
                 outcomes = list(map(_sample_to_str, self.generate_basis_states(num_wires)))
         elif obs.return_type is AllCounts:
-            outcomes = qml.eigvals(obs)
+            outcomes = (
+                qml.eigvals(obs)
+                if not isinstance(obs, MeasurementValue)
+                else np.arange(0, 2 ** len(obs.wires), 1)
+            )
 
         batched = len(shape) == batched_ndims
         if not batched:
diff --git a/tests/measurements/legacy/test_counts_legacy.py b/tests/measurements/legacy/test_counts_legacy.py
index 58b0a95966d..57adfef7a74 100644
--- a/tests/measurements/legacy/test_counts_legacy.py
+++ b/tests/measurements/legacy/test_counts_legacy.py
@@ -171,6 +171,66 @@ def circuit():
         circuit()
         assert circuit._qfunc_output.return_type is Counts  # pylint: disable=protected-access
 
+    @pytest.mark.parametrize("shots", [1000, [1000, 1000]])
+    @pytest.mark.parametrize("phi", np.arange(np.pi / 4, 2 * np.pi, np.pi / 2))
+    def test_observable_is_measurement_value(self, shots, phi):
+        """Test that counts for mid-circuit measurement values
+        are correct for a single measurement value."""
+        dev = qml.device("default.qubit.legacy", wires=2, shots=shots)
+
+        @qml.qnode(dev)
+        def circuit(phi):
+            qml.RX(phi, 0)
+            m0 = qml.measure(0)
+            return qml.counts(m0)
+
+        res = circuit(phi)
+        if isinstance(shots, list):
+            assert isinstance(res, tuple)
+            assert len(res) == 2
+
+            for r in res:
+                assert isinstance(r, dict)
+                assert len(r) == 2
+                assert set(r.keys()) == {0, 1}
+
+        else:
+            assert isinstance(res, dict)
+            assert len(res) == 2
+            assert set(res.keys()) == {0, 1}
+
+    @pytest.mark.parametrize("shots", [5, [5, 5]])
+    def test_observable_is_measurement_value_all_outcomes(self, shots):
+        """Test that counts for mid-circuit measurement values
+        are correct for a single measurement value."""
+        dev = qml.device("default.qubit.legacy", wires=2, shots=shots)
+
+        @qml.qnode(dev)
+        def circuit():
+            qml.PauliX(0)
+            m0 = qml.measure(0)
+            return qml.counts(m0, all_outcomes=True)
+
+        res = circuit()
+        print(res)
+        if isinstance(shots, list):
+            assert isinstance(res, tuple)
+            assert len(res) == 2
+
+            for r in res:
+                assert isinstance(r, dict)
+                assert len(r) == 2
+                assert set(r.keys()) == {0, 1}
+                assert r[0] == 0
+                assert r[1] == 5
+
+        else:
+            assert isinstance(res, dict)
+            assert len(res) == 2
+            assert set(res.keys()) == {0, 1}
+            assert res[0] == 0
+            assert res[1] == 5
+
     def test_providing_no_observable_and_no_wires_counts(self, mocker):
         """Test that we can provide no observable and no wires to sample function"""
         dev = qml.device("default.qubit.legacy", wires=2, shots=1000)
diff --git a/tests/measurements/legacy/test_expval_legacy.py b/tests/measurements/legacy/test_expval_legacy.py
index bde8fe72164..08a5610d32b 100644
--- a/tests/measurements/legacy/test_expval_legacy.py
+++ b/tests/measurements/legacy/test_expval_legacy.py
@@ -122,7 +122,7 @@ def test_observable_is_measurement_value(
     ):  # pylint: disable=too-many-arguments
         """Test that expectation values for mid-circuit measurement values
         are correct for a single measurement value."""
-        dev = qml.device("default.qubit", wires=2, shots=shots)
+        dev = qml.device("default.qubit.legacy", wires=2, shots=shots)
 
         @qml.qnode(dev)
         def circuit(phi):
diff --git a/tests/measurements/legacy/test_probs_legacy.py b/tests/measurements/legacy/test_probs_legacy.py
index 91bf12a1b40..8742a737c6a 100644
--- a/tests/measurements/legacy/test_probs_legacy.py
+++ b/tests/measurements/legacy/test_probs_legacy.py
@@ -178,9 +178,9 @@ def circuit():
     def test_observable_is_measurement_value(
         self, shots, phi, mocker, tol, tol_stochastic
     ):  # pylint: disable=too-many-arguments
-        """Test that expectation values for mid-circuit measurement values
+        """Test that probs for mid-circuit measurement values
         are correct for a single measurement value."""
-        dev = qml.device("default.qubit", wires=2, shots=shots)
+        dev = qml.device("default.qubit.legacy", wires=2, shots=shots)
 
         @qml.qnode(dev)
         def circuit(phi):
diff --git a/tests/measurements/legacy/test_sample_legacy.py b/tests/measurements/legacy/test_sample_legacy.py
index 9630ad518bc..baea2216216 100644
--- a/tests/measurements/legacy/test_sample_legacy.py
+++ b/tests/measurements/legacy/test_sample_legacy.py
@@ -98,9 +98,9 @@ def circuit():
     @pytest.mark.parametrize("shots", [5, [5, 5]])
     @pytest.mark.parametrize("phi", np.arange(0, 2 * np.pi, np.pi / 2))
     def test_observable_is_measurement_value(self, shots, phi, mocker):
-        """Test that expectation values for mid-circuit measurement values
+        """Test that samples for mid-circuit measurement values
         are correct for a single measurement value."""
-        dev = qml.device("default.qubit", wires=2, shots=shots)
+        dev = qml.device("default.qubit.legacy", wires=2, shots=shots)
 
         @qml.qnode(dev)
         def circuit(phi):
diff --git a/tests/measurements/legacy/test_var_legacy.py b/tests/measurements/legacy/test_var_legacy.py
index 43eb58ff2d5..6245e45854e 100644
--- a/tests/measurements/legacy/test_var_legacy.py
+++ b/tests/measurements/legacy/test_var_legacy.py
@@ -115,7 +115,7 @@ def test_observable_is_measurement_value(
     ):  # pylint: disable=too-many-arguments
         """Test that variances for mid-circuit measurement values
         are correct for a single measurement value."""
-        dev = qml.device("default.qubit", wires=2, shots=shots)
+        dev = qml.device("default.qubit.legacy", wires=2, shots=shots)
 
         @qml.qnode(dev)
         def circuit(phi):
diff --git a/tests/measurements/test_probs.py b/tests/measurements/test_probs.py
index 9b89c186dae..c1a55296c6f 100644
--- a/tests/measurements/test_probs.py
+++ b/tests/measurements/test_probs.py
@@ -242,7 +242,7 @@ def circuit():
     def test_observable_is_measurement_value(
         self, shots, phi, tol, tol_stochastic
     ):  # pylint: disable=too-many-arguments
-        """Test that expectation values for mid-circuit measurement values
+        """Test that probs for mid-circuit measurement values
         are correct for a single measurement value."""
         dev = qml.device("default.qubit", wires=2, shots=shots)
 
diff --git a/tests/measurements/test_sample.py b/tests/measurements/test_sample.py
index becca8424a2..10a5fb070d9 100644
--- a/tests/measurements/test_sample.py
+++ b/tests/measurements/test_sample.py
@@ -156,7 +156,7 @@ def circuit():
     @pytest.mark.parametrize("shots", [5, [5, 5]])
     @pytest.mark.parametrize("phi", np.arange(0, 2 * np.pi, np.pi / 2))
     def test_observable_is_measurement_value(self, shots, phi):
-        """Test that expectation values for mid-circuit measurement values
+        """Test that samples for mid-circuit measurement values
         are correct for a single measurement value."""
         dev = qml.device("default.qubit", wires=2, shots=shots)
 

From d1ad001ec43a83148818e7c07ca49dc24dce8dfa Mon Sep 17 00:00:00 2001
From: Mudit Pandey <mudit.pandey@xanadu.ai>
Date: Wed, 20 Sep 2023 14:42:14 -0400
Subject: [PATCH 085/131] Update doc/introduction/measurements.rst

Co-authored-by: Matthew Silverman <matthews@xanadu.ai>
---
 doc/introduction/measurements.rst | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/doc/introduction/measurements.rst b/doc/introduction/measurements.rst
index da14363f049..9945273047b 100644
--- a/doc/introduction/measurements.rst
+++ b/doc/introduction/measurements.rst
@@ -382,7 +382,7 @@ statistics are ``"default.qubit"``, ``"default.mixed"``, and ``"default.qubit.le
 
 .. code-block:: python3
 
-    dev = qml.devices.DefaultQubit()
+    dev = qml.device("default.qubit", wires=2)
 
     @qml.qnode(dev)
     def func(x, y):

From f0baf8796d637dbd3ba63688117e94ea8a800e15 Mon Sep 17 00:00:00 2001
From: Mudit Pandey <mudit.pandey@xanadu.ai>
Date: Wed, 20 Sep 2023 14:42:30 -0400
Subject: [PATCH 086/131] Update doc/releases/changelog-dev.md

---
 doc/releases/changelog-dev.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index cccf9bfaf93..d9466aa4485 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -10,7 +10,7 @@
   [(#4544)](https://github.com/PennyLaneAI/pennylane/pull/4544)
 
   ```python
-  dev = qml.device("default.qubit", wires=3)
+  dev = qml.device("default.qubit", wires=2)
 
   @qml.qnode(dev)
   def circ(x, y):

From 227899b0f44cd413168ded9ade63427b6d468b8b Mon Sep 17 00:00:00 2001
From: Mudit Pandey <mudit.pandey@xanadu.ai>
Date: Wed, 20 Sep 2023 14:47:02 -0400
Subject: [PATCH 087/131] Update pennylane/tape/qscript.py

Co-authored-by: Matthew Silverman <matthews@xanadu.ai>
---
 pennylane/tape/qscript.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pennylane/tape/qscript.py b/pennylane/tape/qscript.py
index 23b3f20e2b4..49e30048536 100644
--- a/pennylane/tape/qscript.py
+++ b/pennylane/tape/qscript.py
@@ -310,7 +310,7 @@ def observables(self) -> List[Union[MeasurementProcess, Observable]]:
         obs = []
 
         for m in self.measurements:
-            if m.obs is not None and not isinstance(m.obs, MeasurementValue):
+            if m.obs is not None:
                 m.obs.return_type = m.return_type
                 obs.append(m.obs)
             else:

From 4728585b58babe99b4d15c0f6a18176688419a2f Mon Sep 17 00:00:00 2001
From: Mudit Pandey <mudit.pandey@xanadu.ai>
Date: Wed, 20 Sep 2023 14:57:50 -0400
Subject: [PATCH 088/131] Apply suggestions from code review

Co-authored-by: Matthew Silverman <matthews@xanadu.ai>
---
 pennylane/interfaces/jax.py                         |  2 +-
 pennylane/tape/qscript.py                           | 11 +++--------
 pennylane/transforms/convert_to_numpy_parameters.py |  2 +-
 tests/transforms/test_defer_measurements.py         |  4 +---
 4 files changed, 6 insertions(+), 13 deletions(-)

diff --git a/pennylane/interfaces/jax.py b/pennylane/interfaces/jax.py
index 8b7bd1cafc6..632be8a51a7 100644
--- a/pennylane/interfaces/jax.py
+++ b/pennylane/interfaces/jax.py
@@ -67,7 +67,7 @@ def get_jax_interface_name(tapes):
         for op in t:
             # Unwrap the observable from a MeasurementProcess
             op = op.obs if hasattr(op, "obs") else op
-            if op is not None and not isinstance(op, qml.measurements.MeasurementValue):
+            if op is not None:
                 # Some MeasurementProcess objects have op.obs=None
                 for param in op.data:
                     if qml.math.is_abstract(param):
diff --git a/pennylane/tape/qscript.py b/pennylane/tape/qscript.py
index 49e30048536..7d827a8f6c8 100644
--- a/pennylane/tape/qscript.py
+++ b/pennylane/tape/qscript.py
@@ -455,7 +455,7 @@ def _update_par_info(self):
 
         n_ops = len(self.operations)
         for idx, m in enumerate(self.measurements):
-            if m.obs is not None and not isinstance(m.obs, MeasurementValue):
+            if m.obs is not None:
                 self._par_info.extend(
                     {"op": m.obs, "op_idx": idx + n_ops, "p_idx": i}
                     for i, d in enumerate(m.obs.data)
@@ -481,12 +481,7 @@ def _update_observables(self):
             _obs_sharing_wires_id (list[int]): Indices of the measurements that contain
                 the observables in _obs_sharing_wires
         """
-        obs_wires = [
-            wire
-            for m in self.measurements
-            for wire in m.wires
-            if m.obs is not None and not isinstance(m.obs, MeasurementValue)
-        ]
+        obs_wires = [wire for m in self.measurements for wire in m.wires if m.obs is not None]
         self._obs_sharing_wires = []
         self._obs_sharing_wires_id = []
 
@@ -676,7 +671,7 @@ def get_parameters(
         if operations_only:
             return params
         for m in self.measurements:
-            if m.obs is not None and not isinstance(m.obs, MeasurementValue):
+            if m.obs is not None:
                 params.extend(m.obs.data)
         return params
 
diff --git a/pennylane/transforms/convert_to_numpy_parameters.py b/pennylane/transforms/convert_to_numpy_parameters.py
index a24554b847f..1c60e29bdb2 100644
--- a/pennylane/transforms/convert_to_numpy_parameters.py
+++ b/pennylane/transforms/convert_to_numpy_parameters.py
@@ -33,7 +33,7 @@ def _convert_op_to_numpy_data(op: qml.operation.Operator) -> qml.operation.Opera
 def _convert_measurement_to_numpy_data(
     m: qml.measurements.MeasurementProcess,
 ) -> qml.measurements.MeasurementProcess:
-    if m.obs is None or isinstance(m.obs, MeasurementValue):
+    if m.obs is None:
         if m.eigvals() is None or math.get_interface(m.eigvals()) == "numpy":
             return m
         return type(m)(wires=m.wires, eigvals=math.unwrap(m.eigvals()))
diff --git a/tests/transforms/test_defer_measurements.py b/tests/transforms/test_defer_measurements.py
index ca72bda5e62..2905ed10e20 100644
--- a/tests/transforms/test_defer_measurements.py
+++ b/tests/transforms/test_defer_measurements.py
@@ -1211,9 +1211,7 @@ def circ(x, y):
             qml.RZ(3.0, 1),
         ]
 
-        assert len(circ.qtape.operations) == 5
-        for op, exp in zip(circ.qtape.operations, expected):
-            assert qml.equal(op, exp)
+        assert circ.qtape.operations == expected
 
     def test_correct_cnot_for_reset(self):
         """Test that a CNOT is applied from the wire that stores the measurement

From b857dc0d96b8737d99b61304394d5157e4951a88 Mon Sep 17 00:00:00 2001
From: Mudit Pandey <mudit.pandey@xanadu.ai>
Date: Wed, 20 Sep 2023 14:58:21 -0400
Subject: [PATCH 089/131] Addressing PR review

---
 pennylane/measurements/mid_measure.py      | 3 +--
 pennylane/measurements/sample.py           | 4 +---
 pennylane/measurements/var.py              | 6 ------
 pennylane/transforms/defer_measurements.py | 7 ++++---
 4 files changed, 6 insertions(+), 14 deletions(-)

diff --git a/pennylane/measurements/mid_measure.py b/pennylane/measurements/mid_measure.py
index 691ba6f9b68..08dfeb3503e 100644
--- a/pennylane/measurements/mid_measure.py
+++ b/pennylane/measurements/mid_measure.py
@@ -180,7 +180,6 @@ class MeasurementValue(Generic[T]):
     def __init__(self, measurements, processing_fn):
         self.measurements = measurements
         self.processing_fn = processing_fn
-        self._wires = Wires([m.wires[0] for m in self.measurements])
 
     def _items(self):
         """A generator representing all the possible outcomes of the MeasurementValue."""
@@ -191,7 +190,7 @@ def _items(self):
     @property
     def wires(self):
         """Returns a list of wires corresponding to the mid-circuit measurements."""
-        return self._wires
+        return Wires([m.wires[0] for m in self.measurements])
 
     @property
     def branches(self):
diff --git a/pennylane/measurements/sample.py b/pennylane/measurements/sample.py
index 529cc621ff8..048ba97e279 100644
--- a/pennylane/measurements/sample.py
+++ b/pennylane/measurements/sample.py
@@ -26,9 +26,7 @@
 from .mid_measure import MeasurementValue
 
 
-def sample(
-    op: Optional[Union[Operator, Sequence[MeasurementValue]]] = None, wires=None
-) -> "SampleMP":
+def sample(op: Optional[Union[Operator, MeasurementValue]] = None, wires=None) -> "SampleMP":
     r"""Sample from the supplied observable, with the number of shots
     determined from the ``dev.shots`` attribute of the corresponding device,
     returning raw samples. If no observable is provided then basis state samples are returned
diff --git a/pennylane/measurements/var.py b/pennylane/measurements/var.py
index cb714654b9c..cafe07ec4ac 100644
--- a/pennylane/measurements/var.py
+++ b/pennylane/measurements/var.py
@@ -117,12 +117,6 @@ def process_samples(
         samples = qml.sample(op=op).process_samples(
             samples=samples, wire_order=wire_order, shot_range=shot_range, bin_size=bin_size
         )
-        # if self.mv is not None:
-        #     num_wires = samples.shape[-1]
-        #     powers_of_two = 2 ** qml.math.arange(num_wires)[::-1]
-        #     indices = samples @ powers_of_two
-        #     indices = qml.math.array(indices)  # Add np.array here for Jax support.
-        #     samples = self.eigvals()[indices]
 
         # With broadcasting, we want to take the variance over axis 1, which is the -1st/-2nd with/
         # without bin_size. Without broadcasting, axis 0 is the -1st/-2nd with/without bin_size
diff --git a/pennylane/transforms/defer_measurements.py b/pennylane/transforms/defer_measurements.py
index e0f4e56476c..27500db939d 100644
--- a/pennylane/transforms/defer_measurements.py
+++ b/pennylane/transforms/defer_measurements.py
@@ -14,7 +14,7 @@
 """Code for the tape transform implementing the deferred measurement principle."""
 from typing import Sequence, Callable
 import pennylane as qml
-from pennylane.measurements import MidMeasureMP, MeasurementValue
+from pennylane.measurements import MidMeasureMP
 from pennylane.ops.op_math import ctrl
 
 from pennylane.tape import QuantumTape
@@ -208,8 +208,9 @@ def func(x, y):
     new_measurements = []
 
     for mp in tape.measurements:
-        if isinstance(mp.obs, MeasurementValue):
-            mp.obs._wires = Wires([control_wires[mp.obs.measurements[0].id]])
+        if mp.mv is not None:
+            wire_map = {m.wires[0]: control_wires[m.id] for m in mp.mv.measurements}
+            mp = qml.map_wires(mp, wire_map=wire_map)
         new_measurements.append(mp)
 
     new_tape = QuantumTape(new_operations, new_measurements, shots=tape.shots)

From 275b355f9204cc50ae3aa6e3d7d20bbc8470e662 Mon Sep 17 00:00:00 2001
From: lillian542 <38584660+lillian542@users.noreply.github.com>
Date: Tue, 19 Sep 2023 16:18:59 -0400
Subject: [PATCH 090/131] Add prng_key kwarg to new device API (#4596)

**Context:**
The `default.qubit.jax` device can take a `prng_key`, and if it does, it
returns the same results from sampling any given state across repeated
measurements, rather than randomizing the sampling of the state.

**Description of the Change:**
Allow the `seed` on `DefaultQubit2` to be a `jax.random.PRNGKey`; if it
is, have the new device replicate the `DefaultQubitJax` sampling
behaviour.

**Benefits:**
You can get the same behaviour on the new device if you want a set PRNG
key when using `jax`.

**Possible Drawbacks:**
Previously, the jax device would use `jax.choice` to generate samples
regardless of whether or not a PRNG key was provided; the difference was
whether it used a new random PRNG key each time, or the same one over
and over.

For the new DefaultQubit, if you provide a PRNG key as the seed, the
behaviour will be the same as the `default.qubit.jax` device with a set
`prng_key`. If you leave the seed as None or an integer, even if the
interface is `jax`, it will use the `numpy.random.default_rng` to
generate the samples, and not `jax.choice`.

If you want a random PRNGKey and `jax.choice`, you can fix it by
reinitializing the device each time you run, with a randomly generated
PRNG key as the seed.

---------

Co-authored-by: Matthew Silverman <matthews@xanadu.ai>
---
 doc/releases/changelog-dev.md                 |   6 +
 pennylane/devices/default_qubit.py            |  89 ++++--
 pennylane/devices/qubit/sampling.py           | 104 ++++++-
 pennylane/devices/qubit/simulate.py           |  22 +-
 .../experimental/test_default_qubit_2.py      |  54 ++++
 tests/devices/qubit/test_sampling.py          | 260 ++++++++++++++++++
 6 files changed, 509 insertions(+), 26 deletions(-)

diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index d9466aa4485..a414ad74d50 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -127,6 +127,11 @@
   when at the beginning of a circuit, thus behaving like `StatePrep`.
   [(#4583)](https://github.com/PennyLaneAI/pennylane/pull/4583)
 
+* `DefaultQubit2` can now accept a `jax.random.PRNGKey` as a `seed`, to set the key for the JAX pseudo random 
+  number generator when using the JAX interface. This corresponds to the `prng_key` on 
+  `DefaultQubitJax` in the old API.
+  [(#4596)](https://github.com/PennyLaneAI/pennylane/pull/4596)
+
 * DefaultQubit2 dispatches to a faster implementation for applying `ParameterizedEvolution` to a state
   when it is more efficient to evolve the state than the operation matrix.
   [(#4598)](https://github.com/PennyLaneAI/pennylane/pull/4598)
@@ -139,6 +144,7 @@
   been updated to reflect this assumption and avoid redundant reshaping.
   [(#4602)](https://github.com/PennyLaneAI/pennylane/pull/4602)
 
+
 <h3>Breaking changes 💔</h3>
 
 * `MeasurementProcess.eigvals()` now raises an `EigvalsUndefinedError` if the measurement observable
diff --git a/pennylane/devices/default_qubit.py b/pennylane/devices/default_qubit.py
index e908996d879..6359a262269 100644
--- a/pennylane/devices/default_qubit.py
+++ b/pennylane/devices/default_qubit.py
@@ -21,6 +21,7 @@
 import concurrent.futures
 import numpy as np
 
+import pennylane as qml
 from pennylane.tape import QuantumTape, QuantumScript
 from pennylane.typing import Result, ResultBatch
 from pennylane.transforms import convert_to_numpy_parameters
@@ -50,11 +51,14 @@ class DefaultQubit(Device):
     Args:
         shots (int, Sequence[int], Sequence[Union[int, Sequence[int]]]): The default number of shots to use in executions involving
             this device.
-        seed (Union[str, None, int, array_like[int], SeedSequence, BitGenerator, Generator]): A
-            seed-like parameter matching that of ``seed`` for ``numpy.random.default_rng`` or
+        seed (Union[str, None, int, array_like[int], SeedSequence, BitGenerator, Generator, jax.random.PRNGKey]): A
+            seed-like parameter matching that of ``seed`` for ``numpy.random.default_rng``, or
             a request to seed from numpy's global random number generator.
             The default, ``seed="global"`` pulls a seed from NumPy's global generator. ``seed=None``
             will pull a seed from the OS entropy.
+            If a ``jax.random.PRNGKey`` is passed as the seed, a JAX-specific sampling function using
+            ``jax.random.choice`` and the ``PRNGKey`` will be used for sampling rather than
+            ``numpy.random.default_rng``.
         max_workers (int): A ``ProcessPoolExecutor`` executes tapes asynchronously
             using a pool of at most ``max_workers`` processes. If ``max_workers`` is ``None``,
             only the current process executes tapes. If you experience any
@@ -144,11 +148,23 @@ def name(self):
         """The name of the device."""
         return "default.qubit.2"
 
-    def __init__(self, wires=None, shots=None, seed="global", max_workers=None) -> None:
+    # pylint:disable = too-many-arguments
+    def __init__(
+        self,
+        wires=None,
+        shots=None,
+        seed="global",
+        max_workers=None,
+    ) -> None:
         super().__init__(wires=wires, shots=shots)
         self._max_workers = max_workers
         seed = np.random.randint(0, high=10000000) if seed == "global" else seed
-        self._rng = np.random.default_rng(seed)
+        if qml.math.get_interface(seed) == "jax":
+            self._prng_key = seed
+            self._rng = np.random.default_rng(None)
+        else:
+            self._prng_key = None
+            self._rng = np.random.default_rng(seed)
         self._debugger = None
 
     def supports_derivatives(
@@ -247,7 +263,13 @@ def execute(
         )
         if max_workers is None:
             results = tuple(
-                simulate(c, rng=self._rng, debugger=self._debugger, interface=interface)
+                simulate(
+                    c,
+                    rng=self._rng,
+                    prng_key=self._prng_key,
+                    debugger=self._debugger,
+                    interface=interface,
+                )
                 for c in circuits
             )
         else:
@@ -255,7 +277,12 @@ def execute(
             seeds = self._rng.integers(2**31 - 1, size=len(vanilla_circuits))
             _wrap_simulate = partial(simulate, debugger=None, interface=interface)
             with concurrent.futures.ProcessPoolExecutor(max_workers=max_workers) as executor:
-                exec_map = executor.map(_wrap_simulate, vanilla_circuits, seeds)
+                exec_map = executor.map(
+                    _wrap_simulate,
+                    vanilla_circuits,
+                    seeds,
+                    [self._prng_key] * len(vanilla_circuits),
+                )
                 results = tuple(exec_map)
 
             # reset _rng to mimic serial behavior
@@ -314,14 +341,24 @@ def execute_and_compute_derivatives(
         max_workers = execution_config.device_options.get("max_workers", self._max_workers)
         if max_workers is None:
             results = tuple(
-                _adjoint_jac_wrapper(c, rng=self._rng, debugger=self._debugger) for c in circuits
+                _adjoint_jac_wrapper(
+                    c, rng=self._rng, debugger=self._debugger, prng_key=self._prng_key
+                )
+                for c in circuits
             )
         else:
             vanilla_circuits = [convert_to_numpy_parameters(c) for c in circuits]
             seeds = self._rng.integers(2**31 - 1, size=len(vanilla_circuits))
 
             with concurrent.futures.ProcessPoolExecutor(max_workers=max_workers) as executor:
-                results = tuple(executor.map(_adjoint_jac_wrapper, vanilla_circuits, seeds))
+                results = tuple(
+                    executor.map(
+                        _adjoint_jac_wrapper,
+                        vanilla_circuits,
+                        seeds,
+                        [self._prng_key] * len(vanilla_circuits),
+                    )
+                )
 
             # reset _rng to mimic serial behavior
             self._rng = np.random.default_rng(self._rng.integers(2**31 - 1))
@@ -400,7 +437,9 @@ def execute_and_compute_jvp(
         max_workers = execution_config.device_options.get("max_workers", self._max_workers)
         if max_workers is None:
             results = tuple(
-                _adjoint_jvp_wrapper(c, t, rng=self._rng, debugger=self._debugger)
+                _adjoint_jvp_wrapper(
+                    c, t, rng=self._rng, debugger=self._debugger, prng_key=self._prng_key
+                )
                 for c, t in zip(circuits, tangents)
             )
         else:
@@ -409,7 +448,13 @@ def execute_and_compute_jvp(
 
             with concurrent.futures.ProcessPoolExecutor(max_workers=max_workers) as executor:
                 results = tuple(
-                    executor.map(_adjoint_jvp_wrapper, vanilla_circuits, tangents, seeds)
+                    executor.map(
+                        _adjoint_jvp_wrapper,
+                        vanilla_circuits,
+                        tangents,
+                        seeds,
+                        [self._prng_key] * len(vanilla_circuits),
+                    )
                 )
 
             # reset _rng to mimic serial behavior
@@ -489,7 +534,9 @@ def execute_and_compute_vjp(
         max_workers = execution_config.device_options.get("max_workers", self._max_workers)
         if max_workers is None:
             results = tuple(
-                _adjoint_vjp_wrapper(c, t, rng=self._rng, debugger=self._debugger)
+                _adjoint_vjp_wrapper(
+                    c, t, rng=self._rng, prng_key=self._prng_key, debugger=self._debugger
+                )
                 for c, t in zip(circuits, cotangents)
             )
         else:
@@ -498,7 +545,13 @@ def execute_and_compute_vjp(
 
             with concurrent.futures.ProcessPoolExecutor(max_workers=max_workers) as executor:
                 results = tuple(
-                    executor.map(_adjoint_vjp_wrapper, vanilla_circuits, cotangents, seeds)
+                    executor.map(
+                        _adjoint_vjp_wrapper,
+                        vanilla_circuits,
+                        cotangents,
+                        seeds,
+                        [self._prng_key] * len(vanilla_circuits),
+                    )
                 )
 
             # reset _rng to mimic serial behavior
@@ -508,22 +561,22 @@ def execute_and_compute_vjp(
         return (results[0], vjps[0]) if is_single_circuit else (results, vjps)
 
 
-def _adjoint_jac_wrapper(c, rng=None, debugger=None):
+def _adjoint_jac_wrapper(c, rng=None, prng_key=None, debugger=None):
     state, is_state_batched = get_final_state(c, debugger=debugger)
     jac = adjoint_jacobian(c, state=state)
-    res = measure_final_state(c, state, is_state_batched, rng=rng)
+    res = measure_final_state(c, state, is_state_batched, rng=rng, prng_key=prng_key)
     return res, jac
 
 
-def _adjoint_jvp_wrapper(c, t, rng=None, debugger=None):
+def _adjoint_jvp_wrapper(c, t, rng=None, prng_key=None, debugger=None):
     state, is_state_batched = get_final_state(c, debugger=debugger)
     jvp = adjoint_jvp(c, t, state=state)
-    res = measure_final_state(c, state, is_state_batched, rng=rng)
+    res = measure_final_state(c, state, is_state_batched, rng=rng, prng_key=prng_key)
     return res, jvp
 
 
-def _adjoint_vjp_wrapper(c, t, rng=None, debugger=None):
+def _adjoint_vjp_wrapper(c, t, rng=None, prng_key=None, debugger=None):
     state, is_state_batched = get_final_state(c, debugger=debugger)
     vjp = adjoint_vjp(c, t, state=state)
-    res = measure_final_state(c, state, is_state_batched, rng=rng)
+    res = measure_final_state(c, state, is_state_batched, rng=rng, prng_key=prng_key)
     return res, vjp
diff --git a/pennylane/devices/qubit/sampling.py b/pennylane/devices/qubit/sampling.py
index b1fc9b55060..0d418b95d51 100644
--- a/pennylane/devices/qubit/sampling.py
+++ b/pennylane/devices/qubit/sampling.py
@@ -108,12 +108,14 @@ def _apply_diagonalizing_gates(
     return state
 
 
+# pylint:disable = too-many-arguments
 def measure_with_samples(
     mps: List[Union[SampleMeasurement, ClassicalShadowMP, ShadowExpvalMP]],
     state: np.ndarray,
     shots: Shots,
     is_state_batched: bool = False,
     rng=None,
+    prng_key=None,
 ) -> List[TensorLike]:
     """
     Returns the samples of the measurement process performed on the given state.
@@ -129,6 +131,8 @@ def measure_with_samples(
         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. Only for simulation using JAX.
 
     Returns:
         List[TensorLike[Any]]: Sample measurement results
@@ -147,7 +151,11 @@ def measure_with_samples(
             # measure with the usual method (rotate into the measurement basis)
             measure_fn = _measure_with_samples_diagonalizing_gates
 
-        all_res.extend(measure_fn(group, state, shots, is_state_batched=is_state_batched, rng=rng))
+        all_res.extend(
+            measure_fn(
+                group, state, shots, is_state_batched=is_state_batched, rng=rng, prng_key=prng_key
+            )
+        )
 
     flat_indices = [_i for i in indices for _i in i]
 
@@ -169,6 +177,7 @@ def _measure_with_samples_diagonalizing_gates(
     shots: Shots,
     is_state_batched: bool = False,
     rng=None,
+    prng_key=None,
 ) -> TensorLike:
     """
     Returns the samples of the measurement process performed on the given state,
@@ -183,6 +192,8 @@ def _measure_with_samples_diagonalizing_gates(
         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. Only for simulation using JAX.
 
     Returns:
         TensorLike[Any]: Sample measurement results
@@ -212,7 +223,12 @@ def _process_single_shot(samples):
             # better to call sample_state just once with total_shots, then use
             # the shot_range keyword argument
             samples = sample_state(
-                state, shots=s, is_state_batched=is_state_batched, wires=wires, rng=rng
+                state,
+                shots=s,
+                is_state_batched=is_state_batched,
+                wires=wires,
+                rng=rng,
+                prng_key=prng_key,
             )
 
             processed_samples.append(_process_single_shot(samples))
@@ -225,6 +241,7 @@ def _process_single_shot(samples):
         is_state_batched=is_state_batched,
         wires=wires,
         rng=rng,
+        prng_key=prng_key,
     )
 
     return _process_single_shot(samples)
@@ -236,6 +253,7 @@ def _measure_classical_shadow(
     shots: Shots,
     is_state_batched: bool = False,
     rng=None,
+    prng_key=None,
 ):
     """
     Returns the result of a classical shadow measurement on the given state.
@@ -274,6 +292,7 @@ def _measure_hamiltonian_with_samples(
     shots: Shots,
     is_state_batched: bool = False,
     rng=None,
+    prng_key=None,
 ):
     # the list contains only one element based on how we group measurements
     mp = mp[0]
@@ -287,6 +306,7 @@ def _sum_for_single_shot(s):
             s,
             is_state_batched=is_state_batched,
             rng=rng,
+            prng_key=prng_key,
         )
         return sum(c * res for c, res in zip(mp.obs.terms()[0], results))
 
@@ -300,6 +320,7 @@ def _measure_sum_with_samples(
     shots: Shots,
     is_state_batched: bool = False,
     rng=None,
+    prng_key=None,
 ):
     # the list contains only one element based on how we group measurements
     mp = mp[0]
@@ -308,7 +329,12 @@ def _measure_sum_with_samples(
     # of the terms separately and sum
     def _sum_for_single_shot(s):
         results = measure_with_samples(
-            [ExpectationMP(t) for t in mp.obs], state, s, is_state_batched=is_state_batched, rng=rng
+            [ExpectationMP(t) for t in mp.obs],
+            state,
+            s,
+            is_state_batched=is_state_batched,
+            rng=rng,
+            prng_key=prng_key,
         )
         return sum(results)
 
@@ -317,7 +343,12 @@ def _sum_for_single_shot(s):
 
 
 def sample_state(
-    state, shots: int, is_state_batched: bool = False, wires=None, rng=None
+    state,
+    shots: int,
+    is_state_batched: bool = False,
+    wires=None,
+    rng=None,
+    prng_key=None,
 ) -> np.ndarray:
     """
     Returns a series of samples of a state.
@@ -330,10 +361,17 @@ def sample_state(
         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. Only for simulation using JAX.
 
     Returns:
         ndarray[int]: Sample values of the shape (shots, num_wires)
     """
+    if prng_key is not None:
+        return _sample_state_jax(
+            state, shots, prng_key, is_state_batched=is_state_batched, wires=wires
+        )
+
     rng = np.random.default_rng(rng)
 
     total_indices = len(state.shape) - is_state_batched
@@ -355,3 +393,61 @@ def sample_state(
     powers_of_two = 1 << np.arange(num_wires, dtype=np.int64)[::-1]
     states_sampled_base_ten = samples[..., None] & powers_of_two
     return (states_sampled_base_ten > 0).astype(np.int64)
+
+
+# pylint:disable = unused-argument
+def _sample_state_jax(
+    state,
+    shots: int,
+    prng_key,
+    is_state_batched: bool = False,
+    wires=None,
+) -> np.ndarray:
+    """
+    Returns a series of samples of a state for the JAX interface based on the PRNG.
+
+    Args:
+        state (array[complex]): A state vector to be sampled
+        shots (int): The number of samples to take
+        prng_key (jax.random.PRNGKey): A``jax.random.PRNGKey``. This is
+            the key to the JAX pseudo random number generator.
+        is_state_batched (bool): whether the state is batched or not
+        wires (Sequence[int]): The wires to sample
+
+    Returns:
+        ndarray[int]: Sample values of the shape (shots, num_wires)
+    """
+    # pylint: disable=import-outside-toplevel
+    import jax
+    import jax.numpy as jnp
+
+    key = prng_key
+
+    total_indices = len(state.shape) - is_state_batched
+    state_wires = qml.wires.Wires(range(total_indices))
+
+    wires_to_sample = wires or state_wires
+    num_wires = len(wires_to_sample)
+    basis_states = np.arange(2**num_wires)
+
+    flat_state = flatten_state(state, total_indices)
+    probs = qml.probs(wires=wires_to_sample).process_state(flat_state, state_wires)
+
+    if is_state_batched:
+        # Produce separate keys for each of the probabilities along the broadcasted axis
+        keys = []
+        for _ in state:
+            key, subkey = jax.random.split(key)
+            keys.append(subkey)
+        samples = jnp.array(
+            [
+                jax.random.choice(_key, basis_states, shape=(shots,), p=prob)
+                for _key, prob in zip(keys, probs)
+            ]
+        )
+    else:
+        samples = jax.random.choice(key, basis_states, shape=(shots,), p=probs)
+
+    powers_of_two = 1 << np.arange(num_wires, dtype=np.int64)[::-1]
+    states_sampled_base_ten = samples[..., None] & powers_of_two
+    return (states_sampled_base_ten > 0).astype(np.int64)
diff --git a/pennylane/devices/qubit/simulate.py b/pennylane/devices/qubit/simulate.py
index 41e89e42ce9..46a39cec74b 100644
--- a/pennylane/devices/qubit/simulate.py
+++ b/pennylane/devices/qubit/simulate.py
@@ -152,7 +152,7 @@ def get_final_state(circuit, debugger=None, interface=None):
     return state, is_state_batched
 
 
-def measure_final_state(circuit, state, is_state_batched, rng=None) -> Result:
+def measure_final_state(circuit, state, is_state_batched, rng=None, prng_key=None) -> Result:
     """
     Perform the measurements required by the circuit on the provided state.
 
@@ -165,6 +165,10 @@ def measure_final_state(circuit, state, is_state_batched, rng=None) -> Result:
         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. Only for simulation using JAX.
+            If None, the default ``sample_state`` function and a ``numpy.random.default_rng``
+            will be for sampling.
 
     Returns:
         Tuple[TensorLike]: The measurement results
@@ -188,7 +192,12 @@ def measure_final_state(circuit, state, is_state_batched, rng=None) -> Result:
 
     rng = default_rng(rng)
     results = measure_with_samples(
-        circuit.measurements, state, shots=circuit.shots, is_state_batched=is_state_batched, rng=rng
+        circuit.measurements,
+        state,
+        shots=circuit.shots,
+        is_state_batched=is_state_batched,
+        rng=rng,
+        prng_key=prng_key,
     )
 
     if len(circuit.measurements) == 1:
@@ -200,7 +209,9 @@ def measure_final_state(circuit, state, is_state_batched, rng=None) -> Result:
     return results
 
 
-def simulate(circuit: qml.tape.QuantumScript, rng=None, debugger=None, interface=None) -> Result:
+def simulate(
+    circuit: qml.tape.QuantumScript, rng=None, prng_key=None, debugger=None, interface=None
+) -> Result:
     """Simulate a single quantum script.
 
     This is an internal function that will be called by the successor to ``default.qubit``.
@@ -210,6 +221,9 @@ def simulate(circuit: qml.tape.QuantumScript, rng=None, debugger=None, interface
         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
 
@@ -227,4 +241,4 @@ def simulate(circuit: qml.tape.QuantumScript, rng=None, debugger=None, interface
 
     """
     state, is_state_batched = get_final_state(circuit, debugger=debugger, interface=interface)
-    return measure_final_state(circuit, state, is_state_batched, rng=rng)
+    return measure_final_state(circuit, state, is_state_batched, rng=rng, prng_key=prng_key)
diff --git a/tests/devices/experimental/test_default_qubit_2.py b/tests/devices/experimental/test_default_qubit_2.py
index 6f07f5c9786..849bc0006a2 100644
--- a/tests/devices/experimental/test_default_qubit_2.py
+++ b/tests/devices/experimental/test_default_qubit_2.py
@@ -1637,6 +1637,60 @@ def test_rng_as_seed(self):
         assert qml.math.allclose(first_num, second_num)
 
 
+@pytest.mark.jax
+class TestPRNGKeySeed:
+    """Test that the device behaves correctly when provided with a PRNG key and using the JAX interface"""
+
+    @pytest.mark.parametrize("max_workers", [None, 1, 2])
+    def test_same_prng_key(self, max_workers):
+        """Test that different devices given the same random jax.random.PRNGKey as a seed will produce
+        the same results for sample, even with different seeds"""
+        import jax
+
+        qs = qml.tape.QuantumScript([qml.Hadamard(0)], [qml.sample(wires=0)], shots=1000)
+        config = ExecutionConfig(interface="jax")
+
+        dev1 = DefaultQubit(max_workers=max_workers, seed=jax.random.PRNGKey(123))
+        result1 = dev1.execute(qs, config)
+
+        dev2 = DefaultQubit(max_workers=max_workers, seed=jax.random.PRNGKey(123))
+        result2 = dev2.execute(qs, config)
+
+        assert np.all(result1 == result2)
+
+    @pytest.mark.parametrize("max_workers", [None, 1, 2])
+    def test_different_prng_key(self, max_workers):
+        """Test that different devices given different jax.random.PRNGKey values will produce
+        different results"""
+        import jax
+
+        qs = qml.tape.QuantumScript([qml.Hadamard(0)], [qml.sample(wires=0)], shots=1000)
+        config = ExecutionConfig(interface="jax")
+
+        dev1 = DefaultQubit(max_workers=max_workers, seed=jax.random.PRNGKey(246))
+        result1 = dev1.execute(qs, config)
+
+        dev2 = DefaultQubit(max_workers=max_workers, seed=jax.random.PRNGKey(123))
+        result2 = dev2.execute(qs, config)
+
+        assert np.any(result1 != result2)
+
+    @pytest.mark.parametrize("max_workers", [None, 1, 2])
+    def test_different_executions_same_prng_key(self, max_workers):
+        """Test that the same device will produce the same results every execution if
+        the seed is a jax.random.PRNGKey"""
+        import jax
+
+        qs = qml.tape.QuantumScript([qml.Hadamard(0)], [qml.sample(wires=0)], shots=1000)
+        config = ExecutionConfig(interface="jax")
+
+        dev = DefaultQubit(max_workers=max_workers, seed=jax.random.PRNGKey(77))
+        result1 = dev.execute(qs, config)
+        result2 = dev.execute(qs, config)
+
+        assert np.all(result1 == result2)
+
+
 class TestHamiltonianSamples:
     """Test that the measure_with_samples function works as expected for
     Hamiltonian and Sum observables
diff --git a/tests/devices/qubit/test_sampling.py b/tests/devices/qubit/test_sampling.py
index 823bcf1ebae..8ce558d2b0f 100644
--- a/tests/devices/qubit/test_sampling.py
+++ b/tests/devices/qubit/test_sampling.py
@@ -14,12 +14,14 @@
 """Unit tests for sample_state in devices/qubit."""
 
 from random import shuffle
+
 import pytest
 
 import pennylane as qml
 from pennylane import numpy as np
 from pennylane.devices.qubit import simulate
 from pennylane.devices.qubit import sample_state, measure_with_samples
+from pennylane.devices.qubit.sampling import _sample_state_jax
 
 two_qubit_state = np.array([[0, 1j], [-1, 0]], dtype=np.complex128) / np.sqrt(2)
 APPROX_ATOL = 0.01
@@ -62,6 +64,48 @@ def test_sample_state_basic(self, interface):
         assert samples.dtype == np.int64
         assert all(qml.math.allequal(s, [0, 1]) or qml.math.allequal(s, [1, 0]) for s in samples)
 
+    @pytest.mark.jax
+    def test_prng_key_as_seed_uses_sample_state_jax(self, mocker):
+        """Tests that sample_state calls _sample_state_jax if the seed is a JAX PRNG key"""
+        import jax
+
+        spy = mocker.spy(qml.devices.qubit.sampling, "_sample_state_jax")
+        state = qml.math.array(two_qubit_state, like="jax")
+
+        # prng_key specified, should call _sample_state_jax
+        _ = sample_state(state, 10, prng_key=jax.random.PRNGKey(15))
+        # prng_key defaults to None, should NOT call _sample_state_jax
+        _ = sample_state(state, 10, rng=15)
+
+        spy.assert_called_once()
+
+    @pytest.mark.jax
+    def test_sample_state_jax(self):
+        """Tests that the returned samples are as expected when explicitly calling _sample_state_jax."""
+        import jax
+
+        state = qml.math.array(two_qubit_state, like="jax")
+
+        samples = _sample_state_jax(state, 10, prng_key=jax.random.PRNGKey(84))
+
+        assert samples.shape == (10, 2)
+        assert samples.dtype == np.int64
+        assert all(qml.math.allequal(s, [0, 1]) or qml.math.allequal(s, [1, 0]) for s in samples)
+
+    @pytest.mark.jax
+    def test_prng_key_determines_sample_state_jax_results(self):
+        """Test that setting the seed as a JAX PRNG key determines the results for _sample_state_jax"""
+        import jax
+
+        state = qml.math.array(two_qubit_state, like="jax")
+
+        samples = _sample_state_jax(state, shots=10, prng_key=jax.random.PRNGKey(12))
+        samples2 = _sample_state_jax(state, shots=10, prng_key=jax.random.PRNGKey(12))
+        samples3 = _sample_state_jax(state, shots=10, prng_key=jax.random.PRNGKey(13))
+
+        assert np.all(samples == samples2)
+        assert not np.allclose(samples, samples3)
+
     @pytest.mark.parametrize("wire_order", [[2], [2, 0], [0, 2, 1]])
     def test_marginal_sample_state(self, wire_order):
         """Tests that marginal states can be sampled as expected."""
@@ -451,6 +495,7 @@ class TestBroadcasting:
 
     def test_sample_measure(self):
         """Test that broadcasting works for qml.sample and single shots"""
+
         rng = np.random.default_rng(123)
         shots = qml.measurements.Shots(100)
 
@@ -489,6 +534,7 @@ def test_sample_measure(self):
     )
     def test_nonsample_measure(self, measurement, expected):
         """Test that broadcasting works for the other sample measurements and single shots"""
+
         rng = np.random.default_rng(123)
         shots = qml.measurements.Shots(10000)
 
@@ -514,6 +560,7 @@ def test_nonsample_measure(self, measurement, expected):
     )
     def test_sample_measure_shot_vector(self, shots):
         """Test that broadcasting works for qml.sample and shot vectors"""
+
         rng = np.random.default_rng(123)
         shots = qml.measurements.Shots(shots)
 
@@ -547,6 +594,7 @@ def test_sample_measure_shot_vector(self, shots):
             # third batch of samples can be any of |00>, |01>, |10>, or |11>
             assert np.all(np.logical_or(r[2] == 0, r[2] == 1))
 
+    # pylint:disable = too-many-arguments
     @pytest.mark.parametrize(
         "shots",
         [
@@ -570,6 +618,7 @@ def test_sample_measure_shot_vector(self, shots):
     )
     def test_nonsample_measure_shot_vector(self, shots, measurement, expected):
         """Test that broadcasting works for the other sample measurements and shot vectors"""
+
         rng = np.random.default_rng(123)
         shots = qml.measurements.Shots(shots)
 
@@ -594,6 +643,217 @@ def test_nonsample_measure_shot_vector(self, shots, measurement, expected):
             assert np.allclose(r, expected, atol=0.01)
 
 
+@pytest.mark.jax
+class TestBroadcastingPRNG:
+    """Test that measurements work and use _sample_state_jax when the state has a batch dim
+    and a PRNG key is provided"""
+
+    def test_sample_measure(self, mocker):
+        """Test that broadcasting works for qml.sample and single shots"""
+        import jax
+
+        spy = mocker.spy(qml.devices.qubit.sampling, "_sample_state_jax")
+
+        rng = np.random.default_rng(123)
+        shots = qml.measurements.Shots(100)
+
+        state = [
+            np.array([[0, 0], [0, 1]]),
+            np.array([[1, 0], [1, 0]]) / np.sqrt(2),
+            np.array([[1, 1], [1, 1]]) / 2,
+        ]
+        state = np.stack(state)
+
+        measurement = qml.sample(wires=[0, 1])
+        res = measure_with_samples(
+            [measurement],
+            state,
+            shots,
+            is_state_batched=True,
+            rng=rng,
+            prng_key=jax.random.PRNGKey(184),
+        )[0]
+
+        spy.assert_called()
+
+        assert res.shape == (3, shots.total_shots, 2)
+        assert res.dtype == np.int64
+
+        # convert to numpy array because prng_key -> JAX -> ArrayImpl -> angry vanilla numpy below
+        res = [np.array(r) for r in res]
+
+        # first batch of samples is always |11>
+        assert np.all(res[0] == 1)
+
+        # second batch of samples is either |00> or |10>
+        assert np.all(np.logical_or(res[1] == [0, 0], res[1] == [1, 0]))
+
+        # third batch of samples can be any of |00>, |01>, |10>, or |11>
+        assert np.all(np.logical_or(res[2] == 0, res[2] == 1))
+
+    @pytest.mark.parametrize(
+        "measurement, expected",
+        [
+            (
+                qml.probs(wires=[0, 1]),
+                np.array([[0, 0, 0, 1], [1 / 2, 0, 1 / 2, 0], [1 / 4, 1 / 4, 1 / 4, 1 / 4]]),
+            ),
+            (qml.expval(qml.PauliZ(1)), np.array([-1, 1, 0])),
+            (qml.var(qml.PauliZ(1)), np.array([0, 0, 1])),
+        ],
+    )
+    def test_nonsample_measure(self, mocker, measurement, expected):
+        """Test that broadcasting works for the other sample measurements and single shots"""
+        import jax
+
+        spy = mocker.spy(qml.devices.qubit.sampling, "_sample_state_jax")
+
+        rng = np.random.default_rng(123)
+        shots = qml.measurements.Shots(10000)
+
+        state = [
+            np.array([[0, 0], [0, 1]]),
+            np.array([[1, 0], [1, 0]]) / np.sqrt(2),
+            np.array([[1, 1], [1, 1]]) / 2,
+        ]
+        state = np.stack(state)
+
+        res = measure_with_samples(
+            [measurement],
+            state,
+            shots,
+            is_state_batched=True,
+            rng=rng,
+            prng_key=jax.random.PRNGKey(184),
+        )
+
+        spy.assert_called()
+        assert np.allclose(res, expected, atol=0.01)
+
+    @pytest.mark.parametrize(
+        "shots",
+        [
+            ((100, 2),),
+            (100, 100),
+            (100, 100),
+            (100, 100, 200),
+            (200, (100, 2)),
+        ],
+    )
+    def test_sample_measure_shot_vector(self, mocker, shots):
+        """Test that broadcasting works for qml.sample and shot vectors"""
+
+        import jax
+
+        spy = mocker.spy(qml.devices.qubit.sampling, "_sample_state_jax")
+
+        rng = np.random.default_rng(123)
+        shots = qml.measurements.Shots(shots)
+
+        state = [
+            np.array([[0, 0], [0, 1]]),
+            np.array([[1, 0], [1, 0]]) / np.sqrt(2),
+            np.array([[1, 1], [1, 1]]) / 2,
+        ]
+        state = np.stack(state)
+
+        measurement = qml.sample(wires=[0, 1])
+        res = measure_with_samples(
+            [measurement],
+            state,
+            shots,
+            is_state_batched=True,
+            rng=rng,
+            prng_key=jax.random.PRNGKey(184),
+        )
+
+        spy.assert_called()
+
+        assert isinstance(res, tuple)
+        assert len(res) == shots.num_copies
+
+        for s, r in zip(shots, res):
+            assert isinstance(r, tuple)
+            assert len(r) == 1
+            r = r[0]
+
+            assert r.shape == (3, s, 2)
+            assert r.dtype == np.int64
+
+            # convert to numpy array because prng_key -> JAX -> ArrayImpl -> angry vanilla numpy below
+            r = [np.array(i) for i in r]
+
+            # first batch of samples is always |11>
+            assert np.all(r[0] == 1)
+
+            # second batch of samples is either |00> or |10>
+            assert np.all(np.logical_or(r[1] == [0, 0], r[1] == [1, 0]))
+
+            # third batch of samples can be any of |00>, |01>, |10>, or |11>
+            assert np.all(np.logical_or(r[2] == 0, r[2] == 1))
+
+    # pylint:disable = too-many-arguments
+    @pytest.mark.parametrize(
+        "shots",
+        [
+            ((10000, 2),),
+            (10000, 10000),
+            (10000, 20000),
+            (10000, 10000, 20000),
+            (20000, (10000, 2)),
+        ],
+    )
+    @pytest.mark.parametrize(
+        "measurement, expected",
+        [
+            (
+                qml.probs(wires=[0, 1]),
+                np.array([[0, 0, 0, 1], [1 / 2, 0, 1 / 2, 0], [1 / 4, 1 / 4, 1 / 4, 1 / 4]]),
+            ),
+            (qml.expval(qml.PauliZ(1)), np.array([-1, 1, 0])),
+            (qml.var(qml.PauliZ(1)), np.array([0, 0, 1])),
+        ],
+    )
+    def test_nonsample_measure_shot_vector(self, mocker, shots, measurement, expected):
+        """Test that broadcasting works for the other sample measurements and shot vectors"""
+
+        import jax
+
+        spy = mocker.spy(qml.devices.qubit.sampling, "_sample_state_jax")
+
+        rng = np.random.default_rng(123)
+        shots = qml.measurements.Shots(shots)
+
+        state = [
+            np.array([[0, 0], [0, 1]]),
+            np.array([[1, 0], [1, 0]]) / np.sqrt(2),
+            np.array([[1, 1], [1, 1]]) / 2,
+        ]
+        state = np.stack(state)
+
+        res = measure_with_samples(
+            [measurement],
+            state,
+            shots,
+            is_state_batched=True,
+            rng=rng,
+            prng_key=jax.random.PRNGKey(184),
+        )
+
+        spy.assert_called()
+
+        assert isinstance(res, tuple)
+        assert len(res) == shots.num_copies
+
+        for r in res:
+            assert isinstance(r, tuple)
+            assert len(r) == 1
+            r = r[0]
+
+            assert r.shape == expected.shape
+            assert np.allclose(r, expected, atol=0.01)
+
+
 class TestHamiltonianSamples:
     """Test that the measure_with_samples function works as expected for
     Hamiltonian and Sum observables"""

From c36f08ca730d459965c99a43e070af40f02b439c Mon Sep 17 00:00:00 2001
From: Christina Lee <christina@xanadu.ai>
Date: Tue, 19 Sep 2023 17:23:50 -0400
Subject: [PATCH 091/131] Fix `qnn.TorchLayer` attributes (#4611)

Fixes bug #4584

`torch.nn.Module` has some special `__getattr__` logic that we are
ignoring if the qnode has been initialized. This bugfix makes it so the
`qnn.TorchLayer.__getattr__` uses `torch.nn.Module.__getattr__` if the
requested property is not on the QNode.

---------

Co-authored-by: Matthew Silverman <matthews@xanadu.ai>
---
 doc/releases/changelog-dev.md |  4 ++++
 pennylane/qnn/torch.py        | 19 ++++++++++---------
 tests/qnn/test_qnn_torch.py   |  3 +++
 3 files changed, 17 insertions(+), 9 deletions(-)

diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index a414ad74d50..248623bd5f7 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -291,6 +291,9 @@
   expectation.
   [(#4590)](https://github.com/PennyLaneAI/pennylane/pull/4590)
 
+* The `torch.nn.Module` properties are now accessible on a `pennylane.qnn.TorchLayer`.
+  [(#4611)](https://github.com/PennyLaneAI/pennylane/pull/4611)
+
 * `qml.math.take` with torch now returns `tensor[..., indices]` when the user requests
   the last axis (`axis=-1`). Without the fix, it would wrongly return `tensor[indices]`.
   [(#4605)](https://github.com/PennyLaneAI/pennylane/pull/4605)
@@ -302,6 +305,7 @@ This release contains contributions from (in alphabetical order):
 Utkarsh Azad,
 Diego Guala,
 Soran Jahangiri,
+Christina Lee,
 Lillian M. A. Frederiksen,
 Vincent Michaud-Rioux,
 Romain Moyard,
diff --git a/pennylane/qnn/torch.py b/pennylane/qnn/torch.py
index 069a5558d52..61b89462da3 100644
--- a/pennylane/qnn/torch.py
+++ b/pennylane/qnn/torch.py
@@ -13,6 +13,8 @@
 # limitations under the License.
 """This module contains the classes and functions for integrating QNodes with the Torch Module
 API."""
+
+import contextlib
 import functools
 import inspect
 import math
@@ -447,21 +449,20 @@ def construct(self, args, kwargs):
         self.qnode.construct((), kwargs)
 
     def __getattr__(self, item):
-        """If the given attribute does not exist in the class, look for it in the wrapped QNode."""
+        """If the qnode is initialized, first check to see if the attribute is on the qnode."""
         if self._initialized:
-            return getattr(self.qnode, item)
+            with contextlib.suppress(AttributeError):
+                return getattr(self.qnode, item)
 
-        try:
-            return self.__dict__[item]
-        except KeyError as exc:
-            raise AttributeError(item) from exc
+        return super().__getattr__(item)
 
     def __setattr__(self, item, val):
-        """If the given attribute does not exist in the class, try to set it in the wrapped QNode."""
-        if self._initialized:
+        """If the qnode is initialized and item is already a qnode property, update it on the qnode, else
+        just update the torch layer itself."""
+        if self._initialized and item in self.qnode.__dict__:
             setattr(self.qnode, item, val)
         else:
-            self.__dict__[item] = val
+            super().__setattr__(item, val)
 
     def _init_weights(
         self,
diff --git a/tests/qnn/test_qnn_torch.py b/tests/qnn/test_qnn_torch.py
index 46d6a4ec662..d1a4b385800 100644
--- a/tests/qnn/test_qnn_torch.py
+++ b/tests/qnn/test_qnn_torch.py
@@ -350,6 +350,9 @@ def test_qnode_weights_registered(self, get_circuit):  # pylint: disable=no-self
             )
             assert weight.requires_grad
 
+        assert layer.w1 is layer._parameters["w1"]  # pylint: disable=protected-access
+        assert layer.w2 is layer._parameters["w2"]  # pylint: disable=protected-access
+
     @pytest.mark.parametrize("n_qubits, output_dim", indices_up_to(2))
     def test_evaluate_qnode(self, get_circuit, n_qubits):  # pylint: disable=no-self-use
         """Test if the _evaluate_qnode() method works correctly, i.e., that it gives the same

From c782f781106594c384b65d22aeacc81cb5e38921 Mon Sep 17 00:00:00 2001
From: Christina Lee <christina@xanadu.ai>
Date: Wed, 20 Sep 2023 10:49:36 -0400
Subject: [PATCH 092/131] bump jax version for docs (#4618)

Hopefully CI passes now.
---
 doc/requirements.txt | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/doc/requirements.txt b/doc/requirements.txt
index 764d94467b4..fcc223767a5 100644
--- a/doc/requirements.txt
+++ b/doc/requirements.txt
@@ -3,8 +3,8 @@ pip
 appdirs
 autograd
 autoray
-jax==0.4.10
-jaxlib==0.4.10
+jax==0.4.16
+jaxlib==0.4.16
 mistune==0.8.4
 m2r2
 numpy

From 6c56a61c2fb0be7d7e3ee7fc1a476504558fb179 Mon Sep 17 00:00:00 2001
From: Romain Moyard <rmoyard@gmail.com>
Date: Wed, 20 Sep 2023 12:12:34 -0400
Subject: [PATCH 093/131] No qnode post processing on informative transforms
 (#4616)

**Description of the Change:**

Remove post processing applied by qnode when there are no parameters for
informative transforms.

---------

Co-authored-by: Rashid N H M <95639609+rashidnhm@users.noreply.github.com>
---
 pennylane/qnode.py          |  5 ++++-
 tests/transforms/test_zx.py | 15 +++++++++++++++
 2 files changed, 19 insertions(+), 1 deletion(-)

diff --git a/pennylane/qnode.py b/pennylane/qnode.py
index cccecc36485..6112915bd74 100644
--- a/pennylane/qnode.py
+++ b/pennylane/qnode.py
@@ -975,7 +975,10 @@ def __call__(self, *args, **kwargs) -> qml.typing.Result:
 
         # convert result to the interface in case the qfunc has no parameters
 
-        if len(self.tape.get_parameters(trainable_only=False)) == 0:
+        if (
+            len(self.tape.get_parameters(trainable_only=False)) == 0
+            and not self.transform_program.is_informative
+        ):
             res = _convert_to_interface(res, self.interface)
 
         if old_interface == "auto":
diff --git a/tests/transforms/test_zx.py b/tests/transforms/test_zx.py
index 34232d4267e..6b2d48f5254 100644
--- a/tests/transforms/test_zx.py
+++ b/tests/transforms/test_zx.py
@@ -664,3 +664,18 @@ def circuit(p):
         g = circuit(params)
 
         assert isinstance(g, pyzx.graph.graph_s.GraphS)
+
+    def test_qnode_decorator_no_params(self):
+        """Test the QNode decorator."""
+        dev = qml.device("default.qubit", wires=2)
+
+        @partial(qml.transforms.to_zx, expand_measurements=True)
+        @qml.qnode(device=dev)
+        def circuit():
+            qml.PauliZ(wires=0)
+            qml.PauliX(wires=1)
+            return qml.expval(qml.PauliZ(0) @ qml.PauliZ(1))
+
+        g = circuit()
+
+        assert isinstance(g, pyzx.graph.graph_s.GraphS)

From 569e1aa64f7eab817d28327c7f0760ce1772dc0d Mon Sep 17 00:00:00 2001
From: Mudit Pandey <mudit.pandey@xanadu.ai>
Date: Wed, 20 Sep 2023 15:08:22 -0400
Subject: [PATCH 094/131] Removed unused imports

---
 pennylane/tape/qscript.py                           | 1 -
 pennylane/transforms/convert_to_numpy_parameters.py | 1 -
 2 files changed, 2 deletions(-)

diff --git a/pennylane/tape/qscript.py b/pennylane/tape/qscript.py
index 7d827a8f6c8..442636069f9 100644
--- a/pennylane/tape/qscript.py
+++ b/pennylane/tape/qscript.py
@@ -28,7 +28,6 @@
     ClassicalShadowMP,
     CountsMP,
     MeasurementProcess,
-    MeasurementValue,
     ProbabilityMP,
     SampleMP,
     ShadowExpvalMP,
diff --git a/pennylane/transforms/convert_to_numpy_parameters.py b/pennylane/transforms/convert_to_numpy_parameters.py
index 1c60e29bdb2..92bebf4164f 100644
--- a/pennylane/transforms/convert_to_numpy_parameters.py
+++ b/pennylane/transforms/convert_to_numpy_parameters.py
@@ -17,7 +17,6 @@
 """
 import pennylane as qml
 from pennylane import math
-from pennylane.measurements import MeasurementValue
 from pennylane.tape import QuantumScript
 
 

From b5d9a5fefe7c37a2fb0d93167bbfe59c8919a2fa Mon Sep 17 00:00:00 2001
From: Matthew Silverman <matthews@xanadu.ai>
Date: Wed, 20 Sep 2023 15:19:16 -0400
Subject: [PATCH 095/131] `default.qubit` returns the new `DefaultQubit` device
 (#4436)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Note: kindly ignore the codecov failure, it seems totally unrelated to
this PR.

**Context:**
It's time for `qml.device("default.qubit")` to return the new device
(look, no wires!).

**Description of the Change:**
- 99% of this PR is changing `dev.batch_execute` to `dev.execute`
- I am sorry about all the pylint changes - it’s from before I silenced
them
- I couldn’t keep track but there were some little things that we
asserted didn’t work, but now they do! Like custom wires with some
measurement processes
- `test_state.py` used to try to test different dtypes, but the returned
state should now _always_ be complex, so that's what I test for.
- One stupid qcut torch test fails on CI only… idk why. it just says
`UNSUPPORTED DTYPE` (that's it because it's jit-compiled)
- I changed `math.size` for tensorflow because autoray [changed
it](https://github.com/jcmgray/autoray/commit/58f02d5e12abafeb7200310b6d846bc323b79633)
in 0.6.4 (I was still on 0.6.3 locally), but we actually want to use the
base `tf.size(x)`

Test formatting changes:
- Duplicated tests in `test_hadamard_gradient.py` for legacy devices
- Any test that used `default.qubit.some_interface` has been
parametrized to also test with `default.qubit` to make sure
interface-specific stuff also works with DQ2
- test_qnode.py uses legacy because it’s deeply steeped in legacy
- test_vqe.py uses the new device but the results are nonsense. I
commented the actual expected results out
- Pulse and qchem tests are being fixed in other PRs
- test_return_types.py uses legacy - test_return_types_dq2.py is a
copy-paste with DQ2
- The things in `test_tensor_measurements.py` just felt unnecessary and
too hard to move over

**Benefits:**
We use the new device!

**Possible Drawbacks:**
3 things that have changed afaik:
- Non-hermitian observables fail hard now, still needs investigation
- `expval(SparseHamiltonian)` used to fail nicely, now it raises a
`DiagGatesUndefined` error
- `qml.Snapshot()` instances are manually spotted and skipped while
doing `adjoint_jacobian` (when `diff_method="adjoint"`)

---------

Co-authored-by: Christina Lee <christina@xanadu.ai>
---
 .github/workflows/interface-unit-tests.yml    |    4 +-
 doc/releases/changelog-dev.md                 |   11 +-
 pennylane/devices/default_qubit.py            |    3 +-
 pennylane/devices/default_qubit_legacy.py     |    4 +-
 pennylane/devices/device_api.py               |    3 -
 pennylane/devices/qubit/adjoint_jacobian.py   |    2 +
 pennylane/math/single_dispatch.py             |    1 +
 pennylane/measurements/state.py               |    2 +-
 pennylane/optimize/spsa.py                    |   12 +-
 pennylane/qnode.py                            |    4 +-
 setup.py                                      |    2 +-
 .../experimental/test_default_qubit_2.py      |   59 +-
 tests/devices/test_default_qubit_autograd.py  |    4 +-
 tests/devices/test_default_qubit_legacy.py    |    6 +-
 .../test_default_qubit_legacy_broadcasting.py |    4 +-
 tests/devices/test_default_qubit_tf.py        |   14 +-
 tests/devices/test_default_qubit_torch.py     |   14 +-
 tests/docs/test_supported_confs.py            |  119 +-
 .../gradients/core/test_gradient_transform.py |    2 +-
 .../gradients/core/test_hadamard_gradient.py  |  130 +-
 .../core/test_hamiltonian_gradient.py         |    6 +-
 tests/gradients/core/test_jvp.py              |   52 +-
 .../core/test_pulse_generator_gradient.py     |   85 +-
 tests/gradients/core/test_pulse_gradient.py   |  120 +-
 tests/gradients/core/test_vjp.py              |   43 +-
 .../finite_diff/test_finite_difference.py     |   75 +-
 .../test_finite_difference_shot_vec.py        |  114 +-
 .../finite_diff/test_spsa_gradient.py         |   75 +-
 .../test_spsa_gradient_shot_vec.py            |   83 +-
 .../parameter_shift/test_parameter_shift.py   |  227 +--
 .../test_parameter_shift_shot_vec.py          |  167 ++-
 tests/logging/test_logging_autograd.py        |   34 +-
 tests/math/test_functions.py                  |   14 +-
 tests/measurements/test_classical_shadow.py   |    4 -
 tests/measurements/test_measurements.py       |    6 +-
 tests/measurements/test_mutual_info.py        |    2 -
 tests/measurements/test_probs.py              |    1 -
 tests/measurements/test_sample.py             |    1 -
 tests/measurements/test_state.py              |   42 +-
 tests/measurements/test_vn_entropy.py         |    2 -
 tests/ops/functions/test_map_wires.py         |   26 +-
 tests/ops/op_math/test_exp.py                 |   11 +-
 tests/ops/op_math/test_prod.py                |   28 +-
 tests/ops/op_math/test_sprod.py               |   12 +-
 tests/ops/qubit/test_hamiltonian.py           |   10 +-
 tests/ops/qubit/test_parametric_ops.py        |    4 +
 tests/ops/qubit/test_sparse.py                |   15 +-
 tests/pulse/test_hardware_hamiltonian.py      |    4 +-
 tests/qchem/test_dipole.py                    |    4 +-
 tests/qchem/test_hamiltonians.py              |    4 +-
 tests/qinfo/test_fisher.py                    |    2 +-
 tests/qnn/test_keras.py                       |    4 +-
 tests/qnn/test_qnn_torch.py                   |    6 +-
 tests/shadow/test_shadow_entropies.py         |    3 +-
 tests/shadow/test_shadow_transforms.py        |    9 +-
 tests/tape/test_qscript.py                    |   67 +-
 tests/tape/test_tape.py                       |   63 +-
 .../test_embeddings/test_amplitude.py         |   21 +-
 tests/templates/test_embeddings/test_angle.py |   11 +-
 tests/templates/test_embeddings/test_basis.py |   11 +-
 .../templates/test_embeddings/test_iqp_emb.py |   11 +-
 .../test_embeddings/test_qaoa_emb.py          |   11 +-
 .../test_layers/test_basic_entangler.py       |   11 +-
 .../templates/test_layers/test_gate_fabric.py |   21 +-
 .../test_particle_conserving_u1.py            |   17 +-
 .../test_particle_conserving_u2.py            |   17 +-
 tests/templates/test_layers/test_random.py    |   11 +-
 .../test_layers/test_simplified_twodesign.py  |   11 +-
 .../test_layers/test_strongly_entangling.py   |   11 +-
 .../test_arbitrary_state_prep.py              |   31 +-
 .../test_basis_state_prep.py                  |   27 +-
 .../test_mottonen_state_prep.py               |   45 +-
 .../test_all_singles_doubles.py               |   13 +-
 .../test_approx_time_evolution.py             |   13 +-
 .../test_arbitrary_unitary.py                 |   11 +-
 .../test_subroutines/test_basis_rotation.py   |   11 +-
 .../test_commuting_evolution.py               |   18 +-
 .../test_double_excitation.py                 |   11 +-
 .../test_subroutines/test_kupccgsd.py         |   15 +-
 .../test_subroutines/test_permute.py          |   21 +-
 tests/templates/test_subroutines/test_qft.py  |    5 +-
 tests/templates/test_subroutines/test_qpe.py  |   20 +-
 .../test_single_excitation.py                 |   11 +-
 .../templates/test_subroutines/test_uccsd.py  |   13 +-
 tests/test_debugging.py                       |   31 +-
 tests/test_qnode.py                           |  162 +--
 tests/test_return_types.py                    |   16 +-
 tests/test_return_types_dq2.py                | 1265 +++++++++++++++++
 tests/test_return_types_qnode.py              |  315 +++-
 tests/test_tensor_measurements.py             |   55 -
 tests/test_vqe.py                             |  127 +-
 .../transforms/test_adjoint_metric_tensor.py  |   27 +-
 tests/transforms/test_batch_input.py          |    6 +-
 tests/transforms/test_batch_params.py         |   20 +-
 tests/transforms/test_batch_transform.py      |    6 +-
 tests/transforms/test_hamiltonian_expand.py   |   44 +-
 tests/transforms/test_metric_tensor.py        |    6 +-
 tests/transforms/test_qcut.py                 |   11 +-
 tests/transforms/test_qmc_transform.py        |    2 +-
 tests/transforms/test_sign_expand.py          |    8 +-
 tests/transforms/test_specs.py                |   38 +-
 101 files changed, 3050 insertions(+), 1343 deletions(-)
 create mode 100644 tests/test_return_types_dq2.py

diff --git a/.github/workflows/interface-unit-tests.yml b/.github/workflows/interface-unit-tests.yml
index 8ea39a1e5a5..07bae8e236b 100644
--- a/.github/workflows/interface-unit-tests.yml
+++ b/.github/workflows/interface-unit-tests.yml
@@ -263,9 +263,9 @@ jobs:
     strategy:
       matrix:
         config:
-          - device: default.qubit
+          - device: default.qubit.legacy
             shots: None
-          - device: default.qubit
+          - device: default.qubit.legacy
             shots: 10000
           # - device: default.qubit.tf
           #   shots: None
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 248623bd5f7..053308c8193 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -64,9 +64,10 @@
 * Quantum information transforms are updated to the new transform program system.
   [(#4569)](https://github.com/PennyLaneAI/pennylane/pull/4569)
 
-* `qml.devices.DefaultQubit` now implements the new device API. The old version of `default.qubit`
-  is still accessible via `qml.devices.DefaultQubitLegacy`, or via short name `default.qubit.legacy`.
+* `default.qubit` now implements the new device API. The old version of the device is still
+  accessible by the short name `default.qubit.legacy`, or directly via `qml.devices.DefaultQubitLegacy`.
   [(#4594)](https://github.com/PennyLaneAI/pennylane/pull/4594)
+  [(#4436)](https://github.com/PennyLaneAI/pennylane/pull/4436)
 
 <h3>Improvements 🛠</h3>
 
@@ -235,6 +236,12 @@
   which effectively just called `marginal_prob` with `np.abs(state) ** 2`.
   [(#4602)](https://github.com/PennyLaneAI/pennylane/pull/4602)
 
+* `default.qubit` now implements the new device API. If you initialize a device
+  with `qml.device("default.qubit")`, all functions and properties that were tied to the old
+  device API will no longer be on the device. The legacy version can still be accessed with
+  `qml.device("default.qubit.legacy", wires=n_wires)`.
+  [(#4436)](https://github.com/PennyLaneAI/pennylane/pull/4436)
+
 <h3>Deprecations 👋</h3>
 
 * The ``prep`` keyword argument in ``QuantumScript`` is deprecated and will be removed from `QuantumScript`.
diff --git a/pennylane/devices/default_qubit.py b/pennylane/devices/default_qubit.py
index 6359a262269..ed029cc3653 100644
--- a/pennylane/devices/default_qubit.py
+++ b/pennylane/devices/default_qubit.py
@@ -146,7 +146,7 @@ def f(x):
     @property
     def name(self):
         """The name of the device."""
-        return "default.qubit.2"
+        return "default.qubit"
 
     # pylint:disable = too-many-arguments
     def __init__(
@@ -193,6 +193,7 @@ def supports_derivatives(
         if (
             execution_config.gradient_method == "backprop"
             and execution_config.device_options.get("max_workers", self._max_workers) is None
+            and execution_config.interface is not None
         ):
             return True
 
diff --git a/pennylane/devices/default_qubit_legacy.py b/pennylane/devices/default_qubit_legacy.py
index dc4983933c3..e7842935d43 100644
--- a/pennylane/devices/default_qubit_legacy.py
+++ b/pennylane/devices/default_qubit_legacy.py
@@ -104,8 +104,8 @@ class DefaultQubitLegacy(QubitDevice):
             returns analytical results.
     """
 
-    name = "Default qubit PennyLane plugin"
-    short_name = "default.qubit"
+    name = "Default qubit PennyLane plugin (Legacy)"
+    short_name = "default.qubit.legacy"
     pennylane_requires = __version__
     version = __version__
     author = "Xanadu Inc."
diff --git a/pennylane/devices/device_api.py b/pennylane/devices/device_api.py
index 1fb8dcb81b1..c8bfd53a825 100644
--- a/pennylane/devices/device_api.py
+++ b/pennylane/devices/device_api.py
@@ -41,9 +41,6 @@ class Device(abc.ABC):
     """A device driver that can control one or more backends. A backend can be either a physical
     Quantum Processing Unit or a virtual one such as a simulator.
 
-    Device drivers should be configured to run under :func:`~.enable_return`, the newer
-    return shape specification, as the old return shape specification is deprecated.
-
     Only the ``execute`` method must be defined to construct a device driver.
 
     .. details::
diff --git a/pennylane/devices/qubit/adjoint_jacobian.py b/pennylane/devices/qubit/adjoint_jacobian.py
index 1c6bc8ba9db..8be6a99909b 100644
--- a/pennylane/devices/qubit/adjoint_jacobian.py
+++ b/pennylane/devices/qubit/adjoint_jacobian.py
@@ -75,6 +75,8 @@ def adjoint_jacobian(tape: QuantumTape, state=None):
     param_number = len(tape.get_parameters(trainable_only=False, operations_only=True)) - 1
     trainable_param_number = len(tape.trainable_params) - 1
     for op in reversed(tape.operations[tape.num_preps :]):
+        if isinstance(op, qml.Snapshot):
+            continue
         adj_op = qml.adjoint(op)
         ket = apply_operation(adj_op, ket)
 
diff --git a/pennylane/math/single_dispatch.py b/pennylane/math/single_dispatch.py
index dd40be00e78..2bf96bef5be 100644
--- a/pennylane/math/single_dispatch.py
+++ b/pennylane/math/single_dispatch.py
@@ -454,6 +454,7 @@ def _cond_tf(pred, true_fn, false_fn, args):
     "vander",
     lambda *args, **kwargs: _i("tf").experimental.numpy.vander(*args, **kwargs),
 )
+ar.register_function("tensorflow", "size", lambda x: _i("tf").size(x))
 
 
 # -------------------------------- Torch --------------------------------- #
diff --git a/pennylane/measurements/state.py b/pennylane/measurements/state.py
index 952223ffb1a..2cc0c576b08 100644
--- a/pennylane/measurements/state.py
+++ b/pennylane/measurements/state.py
@@ -222,4 +222,4 @@ def process_state(self, state: Sequence[complex], wire_order: Wires):
         # pylint:disable=redefined-outer-name
         wire_map = dict(zip(wire_order, range(len(wire_order))))
         mapped_wires = [wire_map[w] for w in self.wires]
-        return qml.math.reduce_statevector(state, indices=mapped_wires, c_dtype=state.dtype)
+        return qml.math.reduce_statevector(state, indices=mapped_wires)
diff --git a/pennylane/optimize/spsa.py b/pennylane/optimize/spsa.py
index a8313735640..a3a1e5faae8 100644
--- a/pennylane/optimize/spsa.py
+++ b/pennylane/optimize/spsa.py
@@ -262,11 +262,13 @@ def compute_grad(self, objective_fn, args, kwargs):
         yminus = objective_fn(*thetaminus, **kwargs)
         try:
             # pylint: disable=protected-access
-            shots = (
-                Shots(objective_fn.device._raw_shot_sequence)
-                if objective_fn.device.shot_vector is not None
-                else Shots(None)
-            )
+            dev_shots = objective_fn.device.shots
+            if isinstance(dev_shots, Shots):
+                shots = dev_shots if dev_shots.has_partitioned_shots else Shots(None)
+            elif objective_fn.device.shot_vector is not None:
+                shots = Shots(objective_fn.device._raw_shot_sequence)  # pragma: no cover
+            else:
+                shots = Shots(None)
             if np.prod(objective_fn.func(*args).shape(objective_fn.device, shots)) > 1:
                 raise ValueError(
                     "The objective function must be a scalar function for the gradient "
diff --git a/pennylane/qnode.py b/pennylane/qnode.py
index 6112915bd74..8c0bddcd55d 100644
--- a/pennylane/qnode.py
+++ b/pennylane/qnode.py
@@ -705,7 +705,9 @@ def _validate_backprop_method(device, interface, shots=None):
             config = qml.devices.ExecutionConfig(gradient_method="backprop", interface=interface)
             if device.supports_derivatives(config):
                 return "backprop", {}, device
-            raise qml.QuantumFunctionError(f"Device {device.name} does not support backprop")
+            raise qml.QuantumFunctionError(
+                f"Device {device.name} does not support backprop with {config}"
+            )
 
         mapped_interface = INTERFACE_MAP.get(interface, interface)
 
diff --git a/setup.py b/setup.py
index 5bd049958f4..34021c1c337 100644
--- a/setup.py
+++ b/setup.py
@@ -47,7 +47,7 @@
         # TODO: rename entry point 'pennylane.plugins' to 'pennylane.devices'.
         # This requires a rename in the setup file of all devices, and is best done during another refactor
         "pennylane.plugins": [
-            "default.qubit = pennylane.devices:DefaultQubitLegacy",
+            "default.qubit = pennylane.devices:DefaultQubit",
             "default.qubit.legacy = pennylane.devices:DefaultQubitLegacy",
             "default.gaussian = pennylane.devices:DefaultGaussian",
             "default.qubit.tf = pennylane.devices.default_qubit_tf:DefaultQubitTF",
diff --git a/tests/devices/experimental/test_default_qubit_2.py b/tests/devices/experimental/test_default_qubit_2.py
index 849bc0006a2..c9233240776 100644
--- a/tests/devices/experimental/test_default_qubit_2.py
+++ b/tests/devices/experimental/test_default_qubit_2.py
@@ -27,7 +27,7 @@
 
 def test_name():
     """Tests the name of DefaultQubit."""
-    assert DefaultQubit().name == "default.qubit.2"
+    assert DefaultQubit().name == "default.qubit"
 
 
 def test_shots():
@@ -200,6 +200,56 @@ def test_tracking_resources(self):
         assert len(tracker.history["resources"]) == 1
         assert tracker.history["resources"][0] == expected_resources
 
+    def test_tracking_batched_execution(self):
+        """Test the number of times the device is executed over a QNode's
+        lifetime is tracked by the device's tracker."""
+
+        dev_1 = qml.device("default.qubit", wires=2)
+
+        def circuit_1(x, y):
+            qml.RX(x, wires=[0])
+            qml.RY(y, wires=[1])
+            qml.CNOT(wires=[0, 1])
+            return qml.expval(qml.PauliZ(0) @ qml.PauliX(1))
+
+        node_1 = qml.QNode(circuit_1, dev_1)
+        num_evals_1 = 10
+
+        with qml.Tracker(dev_1, persistent=True) as tracker1:
+            for _ in range(num_evals_1):
+                node_1(0.432, np.array([0.12, 0.5, 3.2]))
+        assert tracker1.totals["executions"] == num_evals_1
+
+        # test a second instance of a default qubit device
+        dev_2 = qml.device("default.qubit", wires=2)
+
+        def circuit_2(x):
+            qml.RX(x, wires=[0])
+            qml.CNOT(wires=[0, 1])
+            return qml.expval(qml.PauliZ(0) @ qml.PauliX(1))
+
+        node_2 = qml.QNode(circuit_2, dev_2)
+        num_evals_2 = 5
+
+        with qml.Tracker(dev_2) as tracker2:
+            for _ in range(num_evals_2):
+                node_2(np.array([0.432, 0.61, 8.2]))
+        assert tracker2.totals["executions"] == num_evals_2
+
+        # test a new circuit on an existing instance of a qubit device
+        def circuit_3(y):
+            qml.RY(y, wires=[1])
+            qml.CNOT(wires=[0, 1])
+            return qml.expval(qml.PauliZ(0) @ qml.PauliX(1))
+
+        node_3 = qml.QNode(circuit_3, dev_1)
+        num_evals_3 = 7
+
+        with tracker1:
+            for _ in range(num_evals_3):
+                node_3(np.array([0.12, 1.214]))
+        assert tracker1.totals["executions"] == num_evals_1 + num_evals_3
+
 
 # pylint: disable=too-few-public-methods
 class TestPreprocessing:
@@ -302,7 +352,7 @@ def test_supports_backprop(self):
         assert dev.supports_jvp() is True
         assert dev.supports_vjp() is True
 
-        config = ExecutionConfig(gradient_method="backprop")
+        config = ExecutionConfig(gradient_method="backprop", interface="auto")
         assert dev.supports_derivatives(config) is True
         assert dev.supports_jvp(config) is True
         assert dev.supports_vjp(config) is True
@@ -317,6 +367,11 @@ def test_supports_backprop(self):
         assert dev.supports_jvp(config) is False
         assert dev.supports_vjp(config) is False
 
+        config = ExecutionConfig(gradient_method="backprop", interface=None)
+        assert dev.supports_derivatives(config) is False
+        assert dev.supports_jvp(config) is False
+        assert dev.supports_vjp(config) is False
+
     def test_supports_adjoint(self):
         """Test that DefaultQubit says that it supports adjoint differentiation."""
         dev = DefaultQubit()
diff --git a/tests/devices/test_default_qubit_autograd.py b/tests/devices/test_default_qubit_autograd.py
index 78094c23191..150f95e6dad 100644
--- a/tests/devices/test_default_qubit_autograd.py
+++ b/tests/devices/test_default_qubit_autograd.py
@@ -487,7 +487,7 @@ def circuit(a, b):
 
         def cost(a, b):
             prob_wire_1 = circuit(a, b)
-            return prob_wire_1[1] - prob_wire_1[0]
+            return prob_wire_1[1] - prob_wire_1[0]  # pylint:disable=unsubscriptable-object
 
         res = cost(a, b)
         expected = -np.cos(a) * np.cos(b)
@@ -513,7 +513,7 @@ def circuit(a, b):
 
         def cost(a, b):
             prob_wire_1 = circuit(a, b)
-            return prob_wire_1[:, 1] - prob_wire_1[:, 0]
+            return prob_wire_1[:, 1] - prob_wire_1[:, 0]  # pylint:disable=unsubscriptable-object
 
         res = cost(a, b)
         expected = -np.cos(a) * np.cos(b)
diff --git a/tests/devices/test_default_qubit_legacy.py b/tests/devices/test_default_qubit_legacy.py
index ec98819a7d4..aaf2f8fe6ba 100644
--- a/tests/devices/test_default_qubit_legacy.py
+++ b/tests/devices/test_default_qubit_legacy.py
@@ -643,7 +643,7 @@ def test_apply_errors_qubit_state_vector(self, qubit_device_2_wires):
         with pytest.raises(
             DeviceError,
             match="Operation StatePrep cannot be used after other Operations have already been applied "
-            "on a default.qubit device.",
+            "on a default.qubit.legacy device.",
         ):
             qubit_device_2_wires.reset()
             qubit_device_2_wires.apply(
@@ -664,7 +664,7 @@ def test_apply_errors_basis_state(self, qubit_device_2_wires):
         with pytest.raises(
             DeviceError,
             match="Operation BasisState cannot be used after other Operations have already been applied "
-            "on a default.qubit device.",
+            "on a default.qubit.legacy device.",
         ):
             qubit_device_2_wires.reset()
             qubit_device_2_wires.apply(
@@ -2091,7 +2091,7 @@ def test_apply_parametrized_evolution_raises_error(self):
         param_ev = qml.evolve(ParametrizedHamiltonian([1], [qml.PauliX(0)]))
         with pytest.raises(
             NotImplementedError,
-            match="The device default.qubit cannot execute a ParametrizedEvolution operation",
+            match="The device default.qubit.legacy cannot execute a ParametrizedEvolution operation",
         ):
             self.dev._apply_parametrized_evolution(state=self.state, operation=param_ev)
 
diff --git a/tests/devices/test_default_qubit_legacy_broadcasting.py b/tests/devices/test_default_qubit_legacy_broadcasting.py
index 8040f505e4c..c737f194959 100644
--- a/tests/devices/test_default_qubit_legacy_broadcasting.py
+++ b/tests/devices/test_default_qubit_legacy_broadcasting.py
@@ -462,7 +462,7 @@ def test_apply_errors_qubit_state_vector_broadcasted(self, qubit_device_2_wires)
         with pytest.raises(
             DeviceError,
             match="Operation StatePrep cannot be used after other Operations have already been applied "
-            "on a default.qubit device.",
+            "on a default.qubit.legacy device.",
         ):
             qubit_device_2_wires.apply([qml.RZ(0.5, wires=[0]), vec])
 
@@ -491,7 +491,7 @@ def test_apply_errors_basis_state_broadcasted(self, qubit_device_2_wires):
         with pytest.raises(
             DeviceError,
             match="Operation BasisState cannot be used after other Operations have already been applied "
-            "on a default.qubit device.",
+            "on a default.qubit.legacy device.",
         ):
             qubit_device_2_wires.apply([vec])
 
diff --git a/tests/devices/test_default_qubit_tf.py b/tests/devices/test_default_qubit_tf.py
index f30d4d70f17..5fbfcb1820d 100644
--- a/tests/devices/test_default_qubit_tf.py
+++ b/tests/devices/test_default_qubit_tf.py
@@ -1873,7 +1873,7 @@ def circuit(a, b):
             # get the probability of wire 1
             prob_wire_1 = circuit(a, b)
             # compute Prob(|1>_1) - Prob(|0>_1)
-            res = prob_wire_1[1] - prob_wire_1[0]
+            res = prob_wire_1[1] - prob_wire_1[0]  # pylint:disable=unsubscriptable-object
 
         expected = -tf.cos(a) * tf.cos(b)
         assert np.allclose(res, expected, atol=tol, rtol=0)
@@ -1900,7 +1900,7 @@ def circuit(a, b):
             # get the probability of wire 1
             prob_wire_1 = circuit(a, b)
             # compute Prob(|1>_1) - Prob(|0>_1)
-            res = prob_wire_1[:, 1] - prob_wire_1[:, 0]
+            res = prob_wire_1[:, 1] - prob_wire_1[:, 0]  # pylint:disable=unsubscriptable-object
 
         expected = -tf.cos(a) * tf.cos(b)
         assert np.allclose(res, expected, atol=tol, rtol=0)
@@ -1938,6 +1938,7 @@ def circuit(a, b):
             [-0.5 * np.sin(a) * (np.cos(b) + 1), 0.5 * np.sin(b) * (1 - np.cos(a))]
         )
 
+        # pylint:disable=no-member
         assert np.allclose(res.numpy(), expected_cost, atol=tol, rtol=0)
 
         res = tape.gradient(res, [a_tf, b_tf])
@@ -1971,6 +1972,7 @@ def circuit(a, b):
             [-0.5 * np.sin(a) * (np.cos(b) + 1), 0.5 * np.sin(b) * (1 - np.cos(a))]
         )
 
+        # pylint:disable=no-member
         assert np.allclose(res.numpy(), expected_cost, atol=tol, rtol=0)
 
         jac = tape.jacobian(res, [a_tf, b_tf])
@@ -2112,8 +2114,8 @@ def circuit(a):
         res = circuit(a)
 
         assert isinstance(res, tf.Tensor)
-        assert res.shape == (shots,)
-        assert set(res.numpy()) == {-1, 1}
+        assert res.shape == (shots,)  # pylint:disable=comparison-with-callable
+        assert set(res.numpy()) == {-1, 1}  # pylint:disable=no-member
 
     def test_estimating_marginal_probability(self, tol):
         """Test that the probability of a subset of wires is accurately estimated."""
@@ -2190,8 +2192,8 @@ def circuit(a):
         res = circuit(a)
 
         assert isinstance(res, tf.Tensor)
-        assert res.shape == (3, shots)
-        assert set(res.numpy().flat) == {-1, 1}
+        assert res.shape == (3, shots)  # pylint:disable=comparison-with-callable
+        assert set(res.numpy().flat) == {-1, 1}  # pylint:disable=no-member
 
     @pytest.mark.parametrize("batch_size", [2, 3])
     def test_estimating_marginal_probability_broadcasted(self, batch_size, tol):
diff --git a/tests/devices/test_default_qubit_torch.py b/tests/devices/test_default_qubit_torch.py
index 9c2a4c53801..52f80820318 100644
--- a/tests/devices/test_default_qubit_torch.py
+++ b/tests/devices/test_default_qubit_torch.py
@@ -1808,7 +1808,7 @@ def circuit(p):
             return qml.expval(qml.PauliZ(0))
 
         res = circuit([x, y, z])
-        res.backward()
+        res.backward()  # pylint:disable=no-member
 
         expected = torch.cos(3 * x) * torch.cos(y) * torch.cos(z / 2) - torch.sin(
             3 * x
@@ -1895,7 +1895,7 @@ def circuit(x):
             return qml.expval(qml.PauliZ(0))
 
         res = circuit(p)
-        res.backward()
+        res.backward()  # pylint:disable=no-member
 
         expected = torch.cos(y) ** 2 - torch.sin(x) * torch.sin(y) ** 2
 
@@ -2033,7 +2033,7 @@ def circuit(a, b):
         # get the probability of wire 1
         prob_wire_1 = circuit(a, b)
         # compute Prob(|1>_1) - Prob(|0>_1)
-        res = prob_wire_1[1] - prob_wire_1[0]
+        res = prob_wire_1[1] - prob_wire_1[0]  # pylint:disable=unsubscriptable-object
         res.backward()
 
         expected = -torch.cos(a) * torch.cos(b)
@@ -2062,7 +2062,7 @@ def cost(a, b):
             # get the probability of wire 1
             prob_wire_1 = circuit(a, b)
             # compute Prob(|1>_1) - Prob(|0>_1)
-            res = prob_wire_1[:, 1] - prob_wire_1[:, 0]
+            res = prob_wire_1[:, 1] - prob_wire_1[:, 0]  # pylint:disable=unsubscriptable-object
             return res
 
         res = cost(a, b)
@@ -2087,7 +2087,7 @@ def circuit(a, b):
         b = torch.tensor(0.654, dtype=torch.float64, requires_grad=True, device=torch_device)
 
         res = circuit(a, b)
-        res.backward()
+        res.backward()  # pylint:disable=no-member
 
         # the analytic result of evaluating circuit(a, b)
         expected_cost = 0.5 * (torch.cos(a) * torch.cos(b) + torch.cos(a) - torch.cos(b) + 1)
@@ -2242,7 +2242,7 @@ def circuit(a):
         res = circuit(a)
 
         assert torch.is_tensor(res)
-        assert res.shape == (shots,)
+        assert res.shape == (shots,)  # pylint:disable=comparison-with-callable
         assert torch.allclose(
             torch.unique(res), torch.tensor([-1, 1], dtype=torch.int64, device=torch_device)
         )
@@ -2326,7 +2326,7 @@ def circuit(a):
         res = circuit(a)
 
         assert torch.is_tensor(res)
-        assert res.shape == (batch_size, shots)
+        assert res.shape == (batch_size, shots)  # pylint:disable=comparison-with-callable
         assert torch.allclose(
             torch.unique(res), torch.tensor([-1, 1], dtype=torch.int64, device=torch_device)
         )
diff --git a/tests/docs/test_supported_confs.py b/tests/docs/test_supported_confs.py
index b8d83a0d72f..84e38b40822 100644
--- a/tests/docs/test_supported_confs.py
+++ b/tests/docs/test_supported_confs.py
@@ -24,7 +24,6 @@
 A configuration is supported if gradients can be computed for the
 QNode without an exception being raised."""
 # pylint: disable=too-many-arguments
-import re
 import pytest
 
 import pennylane as qml
@@ -47,7 +46,6 @@
 jax = pytest.importorskip("jax")
 jnp = pytest.importorskip("jax.numpy")
 
-devices = ["default.qubit"]
 interfaces = [None, "autograd", "jax", "tf", "torch"]
 diff_interfaces = ["autograd", "jax", "tf", "torch"]
 shots_list = [None, 100]
@@ -277,13 +275,10 @@ def test_all_device(self, interface, return_type, shots, wire_specs):
     @pytest.mark.parametrize("wire_specs", wire_specs_list)
     def test_none_backprop(self, return_type, wire_specs):
         """Test interface=None and diff_method=backprop raises an error"""
-        msg = (
-            "Device default.qubit only supports diff_method='backprop' when "
-            "using the ['tf', 'torch', 'autograd', 'jax'] interfaces."
-        )
-        msg = re.escape(msg)
-
-        with pytest.raises(QuantumFunctionError, match=msg):
+        with pytest.raises(
+            QuantumFunctionError,
+            match=r"Device default\.qubit does not support backprop with .*gradient_method='backprop'.*interface=None",
+        ):
             get_qnode(None, "backprop", return_type, None, wire_specs)
 
     @pytest.mark.parametrize(
@@ -295,17 +290,33 @@ def test_none_backprop(self, return_type, wire_specs):
     def test_none_all(self, diff_method, return_type, shots, wire_specs):
         """Test interface=None and diff_method in [adjoint, parameter-shift,
         finite-diff, spsa, hadamard] has a working forward pass"""
-        warn_msg = (
-            "Requested adjoint differentiation to be computed with finite shots. "
-            "Adjoint differentiation always calculated exactly."
-        )
+        circuit = get_qnode(None, diff_method, return_type, shots, wire_specs)
+        x = get_variable(None, wire_specs)
+
+        msg = None
+        if not shots and return_type is Sample:
+            msg = "Analytic circuits must only contain StateMeasurements"
+        elif diff_method == "adjoint":
+            if shots:
+                msg = (
+                    "Circuits with finite shots must be executed with non-analytic gradient methods"
+                )
+            elif return_type not in ("Hermitian", "Projector", Expectation):
+                msg = "Adjoint differentiation method does not support measurement .*"
+        elif shots and return_type in (
+            VnEntropy,
+            MutualInfo,
+            "DensityMatrix",
+            "StateCost",
+            "StateVector",
+        ):
+            msg = "Circuits with finite shots must only contain"
 
-        if diff_method == "adjoint" and shots is not None:
-            # this warning is still raised in the forward pass
-            with pytest.warns(UserWarning, match=warn_msg):
-                get_qnode(None, diff_method, return_type, shots, wire_specs)
+        if msg is not None:
+            with pytest.raises(qml.DeviceError, match=msg):
+                circuit(x)
         else:
-            get_qnode(None, diff_method, return_type, shots, wire_specs)
+            circuit(x)
 
     @pytest.mark.parametrize("interface", diff_interfaces)
     @pytest.mark.parametrize(
@@ -348,18 +359,16 @@ def test_all_backprop_finite_shots(self, interface, return_type, wire_specs):
     def test_all_adjoint_nonexp(self, interface, return_type, shots, wire_specs):
         """Test diff_method=adjoint raises an error for non-expectation
         measurements for all interfaces"""
-        msg = "Adjoint differentiation method does not support measurement .*"
-
-        warn_msg = (
-            "Requested adjoint differentiation to be computed with finite shots. "
-            "Adjoint differentiation always calculated exactly."
+        msg = (
+            "Circuits with finite shots must be executed with non-analytic gradient methods"
+            if shots
+            else "Adjoint differentiation method does not support measurement .*"
         )
 
-        with pytest.raises(QuantumFunctionError, match=msg):
-            with pytest.warns(UserWarning, match=warn_msg):
-                circuit = get_qnode(interface, "adjoint", return_type, shots, wire_specs)
-                x = get_variable(interface, wire_specs)
-                compute_gradient(x, interface, circuit, return_type)
+        circuit = get_qnode(interface, "adjoint", return_type, shots, wire_specs)
+        x = get_variable(interface, wire_specs)
+        with pytest.raises(qml.DeviceError, match=msg):
+            compute_gradient(x, interface, circuit, return_type)
 
     @pytest.mark.parametrize("interface", diff_interfaces)
     @pytest.mark.parametrize("return_type", [Expectation, "Hermitian", "Projector"])
@@ -367,11 +376,6 @@ def test_all_adjoint_nonexp(self, interface, return_type, shots, wire_specs):
     @pytest.mark.parametrize("wire_specs", wire_specs_list)
     def test_all_adjoint_exp(self, interface, return_type, shots, wire_specs):
         """Test diff_method=adjoint works for expectation measurements for all interfaces"""
-        warn_msg = (
-            "Requested adjoint differentiation to be computed with finite shots. "
-            "Adjoint differentiation always calculated exactly."
-        )
-
         if shots is None:
             # test that everything runs
             # correctness is already tested in other test files
@@ -379,10 +383,12 @@ def test_all_adjoint_exp(self, interface, return_type, shots, wire_specs):
             x = get_variable(interface, wire_specs)
             compute_gradient(x, interface, circuit, return_type)
         else:
-            # test warning is raised when shots > 0
-            with pytest.warns(UserWarning, match=warn_msg):
-                circuit = get_qnode(interface, "adjoint", return_type, shots, wire_specs)
-                x = get_variable(interface, wire_specs)
+            circuit = get_qnode(interface, "adjoint", return_type, shots, wire_specs)
+            x = get_variable(interface, wire_specs)
+            with pytest.raises(
+                qml.DeviceError,
+                match="Circuits with finite shots must be executed with non-analytic gradient methods",
+            ):
                 compute_gradient(x, interface, circuit, return_type)
 
     @pytest.mark.parametrize("interface", diff_interfaces)
@@ -416,13 +422,16 @@ def test_all_paramshift_state(self, interface, return_type, shots, wire_specs):
         )
         complex = return_type == "StateVector"
 
-        with pytest.raises(ValueError, match=msg):
-            circuit = get_qnode(interface, "parameter-shift", return_type, shots, wire_specs)
-            x = get_variable(interface, wire_specs, complex=complex)
-            if shots is not None:
-                with pytest.warns(UserWarning, match="the returned result is analytic"):
-                    compute_gradient(x, interface, circuit, return_type, complex=complex)
-            else:
+        # with pytest.raises(ValueError, match=msg):
+        circuit = get_qnode(interface, "parameter-shift", return_type, shots, wire_specs)
+        x = get_variable(interface, wire_specs, complex=complex)
+        if shots is not None:
+            with pytest.raises(
+                qml.DeviceError, match="Circuits with finite shots must only contain"
+            ):
+                compute_gradient(x, interface, circuit, return_type, complex=complex)
+        else:
+            with pytest.raises(ValueError, match=msg):
                 compute_gradient(x, interface, circuit, return_type, complex=complex)
 
     @pytest.mark.parametrize("interface", diff_interfaces)
@@ -441,7 +450,9 @@ def test_all_finitediff_nonstate(self, interface, return_type, shots, wire_specs
         circuit = get_qnode(interface, diff_method, return_type, shots, wire_specs)
         x = get_variable(interface, wire_specs)
         if shots is not None and return_type in (VnEntropy, MutualInfo):
-            with pytest.warns(UserWarning, match="unaffected by sampling"):
+            with pytest.raises(
+                qml.DeviceError, match="Circuits with finite shots must only contain"
+            ):
                 compute_gradient(x, interface, circuit, return_type)
         else:
             compute_gradient(x, interface, circuit, return_type)
@@ -480,12 +491,8 @@ def test_all_finitediff_state(self, interface, return_type, shots, wire_specs, d
     def test_all_sample_none_shots(self, interface, diff_method, wire_specs):
         """Test sample measurement fails for all interfaces and diff_methods
         when shots=None"""
-        msg = (
-            "The number of shots has to be explicitly set on the device "
-            "when using sample-based measurements."
-        )
 
-        with pytest.raises(QuantumFunctionError, match=msg):
+        with pytest.raises(qml.DeviceError, match="Analytic circuits must only contain"):
             circuit = get_qnode(interface, diff_method, Sample, None, wire_specs)
             x = get_variable(interface, wire_specs)
             circuit(x)
@@ -540,11 +547,13 @@ def test_all_hadamard_nonstate_non_var(
         circuit = get_qnode(interface, diff_method, return_type, shots, wire_specs)
         x = get_variable(interface, wire_specs)
         if return_type in (VnEntropy, MutualInfo):
-            msg = (
-                "Computing the gradient of circuits that return the state with the "
-                "Hadamard test gradient transform is not supported."
-            )
-            with pytest.raises(ValueError, match=msg):
+            if shots:
+                err_cls = qml.DeviceError
+                msg = "Circuits with finite shots must only contain"
+            else:
+                err_cls = ValueError
+                msg = "Computing the gradient of circuits that return the state with the Hadamard test gradient transform is not supported"
+            with pytest.raises(err_cls, match=msg):
                 compute_gradient(x, interface, circuit, return_type)
         elif return_type == Variance:
             with pytest.raises(
diff --git a/tests/gradients/core/test_gradient_transform.py b/tests/gradients/core/test_gradient_transform.py
index 3aded27707f..f592a535129 100644
--- a/tests/gradients/core/test_gradient_transform.py
+++ b/tests/gradients/core/test_gradient_transform.py
@@ -616,7 +616,7 @@ def test_setting_shots(self):
         """Test that setting the number of shots works correctly for
         a gradient transform"""
 
-        dev = qml.device("default.qubit", wires=1, shots=1000)
+        dev = qml.device("default.qubit.legacy", wires=1, shots=1000)
 
         @qml.qnode(dev)
         def circuit(x):
diff --git a/tests/gradients/core/test_hadamard_gradient.py b/tests/gradients/core/test_hadamard_gradient.py
index 92bfbf67357..3b485b7aaab 100644
--- a/tests/gradients/core/test_hadamard_gradient.py
+++ b/tests/gradients/core/test_hadamard_gradient.py
@@ -25,7 +25,7 @@
 def grad_fn(tape, dev, fn=qml.gradients.hadamard_grad, **kwargs):
     """Utility function to automate execution and processing of gradient tapes"""
     tapes, fn = fn(tape, **kwargs)
-    return fn(dev.batch_execute(tapes)), tapes
+    return fn(dev.execute(tapes)), tapes
 
 
 def cost1(x):
@@ -548,6 +548,8 @@ def test_shots_attribute(self, shots):
 class TestHadamardGradEdgeCases:
     """Test the Hadamard gradient transform and edge cases such as non diff parameters, auxiliary wires, etc..."""
 
+    # pylint:disable=too-many-public-methods
+
     device_wires = [qml.wires.Wires([0, 1, "aux"])]
     device_wires_no_aux = [qml.wires.Wires([0, 1, 2])]
 
@@ -725,7 +727,7 @@ def test_no_trainable_params_qnode_autograd(self, mocker):
         """Test that the correct ouput and warning is generated in the absence of any trainable
         parameters"""
         dev = qml.device("default.qubit", wires=2)
-        spy = mocker.spy(dev, "expval")
+        spy = mocker.spy(qml.devices.qubit, "measure")
 
         @qml.qnode(dev, interface="autograd")
         def circuit(weights):
@@ -745,7 +747,7 @@ def test_no_trainable_params_qnode_torch(self, mocker):
         """Test that the correct ouput and warning is generated in the absence of any trainable
         parameters"""
         dev = qml.device("default.qubit", wires=2)
-        spy = mocker.spy(dev, "expval")
+        spy = mocker.spy(qml.devices.qubit, "measure")
 
         @qml.qnode(dev, interface="torch")
         def circuit(weights):
@@ -765,7 +767,7 @@ def test_no_trainable_params_qnode_tf(self, mocker):
         """Test that the correct ouput and warning is generated in the absence of any trainable
         parameters"""
         dev = qml.device("default.qubit", wires=2)
-        spy = mocker.spy(dev, "expval")
+        spy = mocker.spy(qml.devices.qubit, "measure")
 
         @qml.qnode(dev, interface="tf")
         def circuit(weights):
@@ -785,6 +787,86 @@ def test_no_trainable_params_qnode_jax(self, mocker):
         """Test that the correct ouput and warning is generated in the absence of any trainable
         parameters"""
         dev = qml.device("default.qubit", wires=2)
+        spy = mocker.spy(qml.devices.qubit, "measure")
+
+        @qml.qnode(dev, interface="jax")
+        def circuit(weights):
+            qml.RX(weights[0], wires=0)
+            qml.RY(weights[1], wires=0)
+            return qml.expval(qml.PauliZ(0) @ qml.PauliZ(1))
+
+        weights = [0.1, 0.2]
+        with pytest.warns(UserWarning, match="gradient of a QNode with no trainable parameters"):
+            res_hadamard = qml.gradients.hadamard_grad(circuit)(weights)
+
+        assert res_hadamard == ()
+        spy.assert_not_called()
+
+    @pytest.mark.autograd
+    def test_no_trainable_params_qnode_autograd_legacy(self, mocker):
+        """Test that the correct ouput and warning is generated in the absence of any trainable
+        parameters"""
+        dev = qml.device("default.qubit.autograd", wires=2)
+        spy = mocker.spy(dev, "expval")
+
+        @qml.qnode(dev, interface="autograd")
+        def circuit(weights):
+            qml.RX(weights[0], wires=0)
+            qml.RY(weights[1], wires=0)
+            return qml.expval(qml.PauliZ(0) @ qml.PauliZ(1))
+
+        weights = [0.1, 0.2]
+        with pytest.warns(UserWarning, match="gradient of a QNode with no trainable parameters"):
+            res_hadamard = qml.gradients.hadamard_grad(circuit)(weights)
+
+        assert res_hadamard == ()
+        spy.assert_not_called()
+
+    @pytest.mark.torch
+    def test_no_trainable_params_qnode_torch_legacy(self, mocker):
+        """Test that the correct ouput and warning is generated in the absence of any trainable
+        parameters"""
+        dev = qml.device("default.qubit.torch", wires=2)
+        spy = mocker.spy(dev, "expval")
+
+        @qml.qnode(dev, interface="torch")
+        def circuit(weights):
+            qml.RX(weights[0], wires=0)
+            qml.RY(weights[1], wires=0)
+            return qml.expval(qml.PauliZ(0) @ qml.PauliZ(1))
+
+        weights = [0.1, 0.2]
+        with pytest.warns(UserWarning, match="gradient of a QNode with no trainable parameters"):
+            res_hadamard = qml.gradients.hadamard_grad(circuit)(weights)
+
+        assert res_hadamard == ()
+        spy.assert_not_called()
+
+    @pytest.mark.tf
+    def test_no_trainable_params_qnode_tf_legacy(self, mocker):
+        """Test that the correct ouput and warning is generated in the absence of any trainable
+        parameters"""
+        dev = qml.device("default.qubit.tf", wires=2)
+        spy = mocker.spy(dev, "expval")
+
+        @qml.qnode(dev, interface="tf")
+        def circuit(weights):
+            qml.RX(weights[0], wires=0)
+            qml.RY(weights[1], wires=0)
+            return qml.expval(qml.PauliZ(0) @ qml.PauliZ(1))
+
+        weights = [0.1, 0.2]
+        with pytest.warns(UserWarning, match="gradient of a QNode with no trainable parameters"):
+            res_hadamard = qml.gradients.hadamard_grad(circuit)(weights)
+
+        assert res_hadamard == ()
+        spy.assert_not_called()
+
+    @pytest.mark.jax
+    def test_no_trainable_params_qnode_jax_legacy(self, mocker):
+        """Test that the correct ouput and warning is generated in the absence of any trainable
+        parameters"""
+        dev = qml.device("default.qubit.jax", wires=2)
         spy = mocker.spy(dev, "expval")
 
         @qml.qnode(dev, interface="jax")
@@ -971,10 +1053,12 @@ class TestHadamardTestGradDiff:
     """Test that the transform is differentiable"""
 
     @pytest.mark.autograd
-    def test_autograd(self):
+    @pytest.mark.parametrize("dev_name", ["default.qubit", "default.qubit.autograd"])
+    def test_autograd(self, dev_name):
         """Tests that the output of the hadamard gradient transform
         can be differentiated using autograd, yielding second derivatives."""
-        dev = qml.device("default.qubit.autograd", wires=3)
+        dev = qml.device(dev_name, wires=3)
+        execute_fn = dev.execute if dev_name == "default.qubit" else dev.batch_execute
         params = np.array([0.543, -0.654], requires_grad=True)
 
         def cost_fn_hadamard(x):
@@ -987,7 +1071,7 @@ def cost_fn_hadamard(x):
             tape = qml.tape.QuantumScript.from_queue(q)
             tape.trainable_params = {0, 1}
             tapes, fn = qml.gradients.hadamard_grad(tape)
-            jac = fn(dev.batch_execute(tapes))
+            jac = fn(execute_fn(tapes))
             return qml.math.stack(jac)
 
         def cost_fn_param_shift(x):
@@ -1000,7 +1084,7 @@ def cost_fn_param_shift(x):
             tape = qml.tape.QuantumScript.from_queue(q)
             tape.trainable_params = {0, 1}
             tapes, fn = qml.gradients.param_shift(tape)
-            jac = fn(dev.batch_execute(tapes))
+            jac = fn(execute_fn(tapes))
             return qml.math.stack(jac)
 
         res_hadamard = qml.jacobian(cost_fn_hadamard)(params)
@@ -1008,12 +1092,14 @@ def cost_fn_param_shift(x):
         assert np.allclose(res_hadamard, res_param_shift)
 
     @pytest.mark.tf
-    def test_tf(self):
+    @pytest.mark.parametrize("dev_name", ["default.qubit", "default.qubit.tf"])
+    def test_tf(self, dev_name):
         """Tests that the output of the hadamard gradient transform
         can be differentiated using TF, yielding second derivatives."""
         import tensorflow as tf
 
-        dev = qml.device("default.qubit.tf", wires=3)
+        dev = qml.device(dev_name, wires=3)
+        execute_fn = dev.execute if dev_name == "default.qubit" else dev.batch_execute
         params = tf.Variable([0.543, -0.654], dtype=tf.float64)
 
         with tf.GradientTape() as t_h:
@@ -1026,7 +1112,7 @@ def test_tf(self):
             tape = qml.tape.QuantumScript.from_queue(q)
             tape.trainable_params = {0, 1}
             tapes, fn = qml.gradients.hadamard_grad(tape)
-            jac_h = fn(dev.batch_execute(tapes))
+            jac_h = fn(execute_fn(tapes))
             jac_h = qml.math.stack(jac_h)
 
         with tf.GradientTape() as t_p:
@@ -1039,7 +1125,7 @@ def test_tf(self):
             tape = qml.tape.QuantumScript.from_queue(q)
             tape.trainable_params = {0, 1}
             tapes, fn = qml.gradients.param_shift(tape)
-            jac_p = fn(dev.batch_execute(tapes))
+            jac_p = fn(execute_fn(tapes))
             jac_p = qml.math.stack(jac_p)
 
         res_hadamard = t_h.jacobian(jac_h, params)
@@ -1048,12 +1134,14 @@ def test_tf(self):
         assert np.allclose(res_hadamard, res_param_shift)
 
     @pytest.mark.torch
-    def test_torch(self):
+    @pytest.mark.parametrize("dev_name", ["default.qubit", "default.qubit.torch"])
+    def test_torch(self, dev_name):
         """Tests that the output of the hadamard gradient transform
         can be differentiated using Torch, yielding second derivatives."""
         import torch
 
-        dev = qml.device("default.qubit.torch", wires=3)
+        dev = qml.device(dev_name, wires=3)
+        execute_fn = dev.execute if dev_name == "default.qubit" else dev.batch_execute
         params = torch.tensor([0.543, -0.654], dtype=torch.float64, requires_grad=True)
 
         def cost_h(x):
@@ -1066,7 +1154,7 @@ def cost_h(x):
             tape = qml.tape.QuantumScript.from_queue(q)
             tapes, fn = qml.gradients.hadamard_grad(tape)
 
-            jac = fn(dev.batch_execute(tapes))
+            jac = fn(execute_fn(tapes))
             return jac
 
         def cost_p(x):
@@ -1079,7 +1167,7 @@ def cost_p(x):
             tape = qml.tape.QuantumScript.from_queue(q)
             tapes, fn = qml.gradients.param_shift(tape)
 
-            jac = fn(dev.batch_execute(tapes))
+            jac = fn(execute_fn(tapes))
             return jac
 
         res_hadamard = torch.autograd.functional.jacobian(cost_h, params)
@@ -1089,7 +1177,8 @@ def cost_p(x):
         assert np.allclose(res_hadamard[1].detach(), res_param_shift[1].detach())
 
     @pytest.mark.jax
-    def test_jax(self):
+    @pytest.mark.parametrize("dev_name", ["default.qubit", "default.qubit.jax"])
+    def test_jax(self, dev_name):
         """Tests that the output of the hadamard gradient transform
         can be differentiated using JAX, yielding second derivatives."""
         import jax
@@ -1098,7 +1187,8 @@ def test_jax(self):
 
         config.update("jax_enable_x64", True)
 
-        dev = qml.device("default.qubit.jax", wires=3)
+        dev = qml.device(dev_name, wires=3)
+        execute_fn = dev.execute if dev_name == "default.qubit" else dev.batch_execute
         params = jnp.array([0.543, -0.654])
 
         def cost_h(x):
@@ -1112,7 +1202,7 @@ def cost_h(x):
             tape.trainable_params = {0, 1}
             tapes, fn = qml.gradients.hadamard_grad(tape)
 
-            jac = fn(dev.batch_execute(tapes))
+            jac = fn(execute_fn(tapes))
             return jac
 
         def cost_p(x):
@@ -1126,7 +1216,7 @@ def cost_p(x):
             tape.trainable_params = {0, 1}
             tapes, fn = qml.gradients.hadamard_grad(tape)
 
-            jac = fn(dev.batch_execute(tapes))
+            jac = fn(execute_fn(tapes))
             return jac
 
         res_hadamard = jax.jacobian(cost_h)(params)
diff --git a/tests/gradients/core/test_hamiltonian_gradient.py b/tests/gradients/core/test_hamiltonian_gradient.py
index 488b8e7dd8a..1bcb4bfc4fe 100644
--- a/tests/gradients/core/test_hamiltonian_gradient.py
+++ b/tests/gradients/core/test_hamiltonian_gradient.py
@@ -30,10 +30,10 @@ def test_behaviour():
     tape = qml.tape.QuantumScript.from_queue(q)
     tape.trainable_params = {2, 3}
     tapes, processing_fn = hamiltonian_grad(tape, idx=0)
-    res1 = processing_fn(dev.batch_execute(tapes))
+    res1 = processing_fn(dev.execute(tapes))
 
     tapes, processing_fn = hamiltonian_grad(tape, idx=1)
-    res2 = processing_fn(dev.batch_execute(tapes))
+    res2 = processing_fn(dev.execute(tapes))
 
     with qml.queuing.AnnotatedQueue() as q1:
         qml.RY(0.3, wires=0)
@@ -49,9 +49,7 @@ def test_behaviour():
         qml.expval(qml.PauliZ(1))
 
     tape2 = qml.tape.QuantumScript.from_queue(q2)
-    dev.reset()
     res_expected1 = qml.math.squeeze(dev.execute(tape1))
-    dev.reset()
     res_expected2 = qml.math.squeeze(dev.execute(tape2))
 
     assert res_expected1 == res1
diff --git a/tests/gradients/core/test_jvp.py b/tests/gradients/core/test_jvp.py
index beb2cc9bf8e..ac973027534 100644
--- a/tests/gradients/core/test_jvp.py
+++ b/tests/gradients/core/test_jvp.py
@@ -484,7 +484,7 @@ def test_single_expectation_value(self, tol, batch_dim):
         tapes, fn = qml.gradients.jvp(tape, tangent, param_shift)
         assert len(tapes) == 4
 
-        res = fn(dev.batch_execute(tapes))
+        res = fn(dev.execute(tapes))
         assert res.shape == () if batch_dim is None else (batch_dim,)
 
         exp = np.sum(np.array([-np.sin(y) * np.sin(x), np.cos(y) * np.cos(x)]), axis=0)
@@ -514,7 +514,7 @@ def test_multiple_expectation_values(self, tol, batch_dim):
         tapes, fn = qml.gradients.jvp(tape, tangent, param_shift)
         assert len(tapes) == 4
 
-        res = fn(dev.batch_execute(tapes))
+        res = fn(dev.execute(tapes))
         assert isinstance(res, tuple)
         assert len(res) == 2
         assert all(r.shape == () if batch_dim is None else (batch_dim,) for r in res)
@@ -546,7 +546,7 @@ def test_prob_expval_single_param(self, tol, batch_dim):
         tapes, fn = qml.gradients.jvp(tape, tangent, param_shift)
         assert len(tapes) == 2
 
-        res = fn(dev.batch_execute(tapes))
+        res = fn(dev.execute(tapes))
         assert isinstance(res, tuple)
         assert len(res) == 2
         assert res[0].shape == () if batch_dim is None else (batch_dim,)
@@ -583,7 +583,7 @@ def test_prob_expval_multi_param(self, tol, batch_dim):
         tapes, fn = qml.gradients.jvp(tape, tangent, param_shift)
         assert len(tapes) == 4
 
-        res = fn(dev.batch_execute(tapes))
+        res = fn(dev.execute(tapes))
         assert isinstance(res, tuple)
         assert len(res) == 2
 
@@ -655,10 +655,12 @@ class TestJVPGradients:
     # Include batch_dim!=None cases once #4462 is resolved
     @pytest.mark.autograd
     @pytest.mark.parametrize("batch_dim", [None])  # , 1, 3])
-    def test_autograd(self, tol, batch_dim):
+    @pytest.mark.parametrize("dev_name", ["default.qubit", "default.qubit.autograd"])
+    def test_autograd(self, tol, dev_name, batch_dim):
         """Tests that the output of the JVP transform
         can be differentiated using autograd."""
-        dev = qml.device("default.qubit.autograd", wires=2)
+        dev = qml.device(dev_name, wires=2)
+        execute_fn = dev.execute if dev_name == "default.qubit" else dev.batch_execute
         params = np.array([0.543, -0.654], requires_grad=True)
         if batch_dim is not None:
             params = np.outer(np.arange(1, 1 + batch_dim), params, requires_grad=True)
@@ -671,7 +673,7 @@ def cost_fn(params, tangent):
             tape = qml.tape.QuantumScript.from_queue(q)
             tape.trainable_params = {0, 1}
             tapes, fn = qml.gradients.jvp(tape, tangent, param_shift)
-            jvp = fn(dev.batch_execute(tapes))
+            jvp = fn(execute_fn(tapes))
             return jvp
 
         res = cost_fn(params, tangent)
@@ -685,12 +687,14 @@ def cost_fn(params, tangent):
     # Include batch_dim!=None cases once #4462 is resolved
     @pytest.mark.torch
     @pytest.mark.parametrize("batch_dim", [None])  # , 1, 3])
-    def test_torch(self, tol, batch_dim):
+    @pytest.mark.parametrize("dev_name", ["default.qubit", "default.qubit.torch"])
+    def test_torch(self, tol, dev_name, batch_dim):
         """Tests that the output of the JVP transform
         can be differentiated using Torch."""
         import torch
 
-        dev = qml.device("default.qubit.torch", wires=2)
+        dev = qml.device(dev_name, wires=2)
+        execute_fn = dev.execute if dev_name == "default.qubit" else dev.batch_execute
 
         params_np = np.array([0.543, -0.654], requires_grad=True)
         if batch_dim is not None:
@@ -706,7 +710,7 @@ def cost_fn(params, tangent):
             tape = qml.tape.QuantumScript.from_queue(q)
             tape.trainable_params = {0, 1}
             tapes, fn = qml.gradients.jvp(tape, tangent, param_shift)
-            jvp = fn(dev.batch_execute(tapes))
+            jvp = fn(execute_fn(tapes))
             return jvp
 
         res = cost_fn(params, tangent)
@@ -721,12 +725,14 @@ def cost_fn(params, tangent):
     @pytest.mark.tf
     @pytest.mark.slow
     @pytest.mark.parametrize("batch_dim", [None])  # , 1, 3])
-    def test_tf(self, tol, batch_dim):
+    @pytest.mark.parametrize("dev_name", ["default.qubit", "default.qubit.tf"])
+    def test_tf(self, tol, dev_name, batch_dim):
         """Tests that the output of the JVP transform
         can be differentiated using Tensorflow."""
         import tensorflow as tf
 
-        dev = qml.device("default.qubit.tf", wires=2)
+        dev = qml.device(dev_name, wires=2)
+        execute_fn = dev.execute if dev_name == "default.qubit" else dev.batch_execute
         params_np = np.array([0.543, -0.654], requires_grad=True)
         if batch_dim is not None:
             params_np = np.outer(np.arange(1, 1 + batch_dim), params_np, requires_grad=True)
@@ -741,7 +747,7 @@ def cost_fn(params, tangent):
             tape = qml.tape.QuantumScript.from_queue(q)
             tape.trainable_params = {0, 1}
             tapes, fn = qml.gradients.jvp(tape, tangent, param_shift)
-            jvp = fn(dev.batch_execute(tapes))
+            jvp = fn(execute_fn(tapes))
             return jvp
 
         with tf.GradientTape() as t:
@@ -757,13 +763,15 @@ def cost_fn(params, tangent):
     # Include batch_dim!=None cases once #4462 is resolved
     @pytest.mark.jax
     @pytest.mark.parametrize("batch_dim", [None])  # , 1, 3])
-    def test_jax(self, tol, batch_dim):
+    @pytest.mark.parametrize("dev_name", ["default.qubit", "default.qubit.jax"])
+    def test_jax(self, tol, dev_name, batch_dim):
         """Tests that the output of the JVP transform
         can be differentiated using JAX."""
         import jax
         from jax import numpy as jnp
 
-        dev = qml.device("default.qubit.jax", wires=2)
+        dev = qml.device(dev_name, wires=2)
+        execute_fn = dev.execute if dev_name == "default.qubit" else dev.batch_execute
         params_np = np.array([0.543, -0.654], requires_grad=True)
         if batch_dim is not None:
             params_np = np.outer(np.arange(1, 1 + batch_dim), params_np, requires_grad=True)
@@ -778,7 +786,7 @@ def cost_fn(params, tangent):
             tape = qml.tape.QuantumScript.from_queue(q)
             tape.trainable_params = {0, 1}
             tapes, fn = qml.gradients.jvp(tape, tangent, param_shift)
-            jvp = fn(dev.batch_execute(tapes))
+            jvp = fn(execute_fn(tapes))
             return jvp
 
         res = cost_fn(params, tangent)
@@ -821,7 +829,7 @@ def test_one_tape_no_trainable_parameters(self):
 
         # Even though there are 3 parameters, only two contribute
         # to the JVP, so only 2*2=4 quantum evals
-        res = fn(dev.batch_execute(v_tapes))
+        res = fn(dev.execute(v_tapes))
 
         assert res[0] is None
         assert res[1] is not None
@@ -877,7 +885,7 @@ def test_zero_tangent(self):
         tangents = [np.array([0.0]), np.array([1.0, 1.0])]
 
         v_tapes, fn = qml.gradients.batch_jvp(tapes, tangents, param_shift)
-        res = fn(dev.batch_execute(v_tapes))
+        res = fn(dev.execute(v_tapes))
 
         # Even though there are 3 parameters, only two contribute
         # to the JVP, so only 2*2=4 quantum evals
@@ -909,7 +917,7 @@ def test_reduction_append(self):
         tangents = [np.array([1.0]), np.array([1.0, 1.0])]
 
         v_tapes, fn = qml.gradients.batch_jvp(tapes, tangents, param_shift, reduction="append")
-        res = fn(dev.batch_execute(v_tapes))
+        res = fn(dev.execute(v_tapes))
 
         # Returned JVPs will be appended to a list, one JVP per tape
 
@@ -942,7 +950,7 @@ def test_reduction_extend(self):
         tangents = [np.array([1.0]), np.array([1.0])]
 
         v_tapes, fn = qml.gradients.batch_jvp(tapes, tangents, param_shift, reduction="extend")
-        res = fn(dev.batch_execute(v_tapes))
+        res = fn(dev.execute(v_tapes))
         assert len(res) == 4
 
     def test_reduction_extend_special(self):
@@ -977,7 +985,7 @@ def test_reduction_extend_special(self):
             if not isinstance(x, tuple) and x.shape == ()
             else jvps.extend(x),
         )
-        res = fn(dev.batch_execute(v_tapes))
+        res = fn(dev.execute(v_tapes))
 
         assert len(res) == 3
 
@@ -1008,6 +1016,6 @@ def test_reduction_callable(self):
         v_tapes, fn = qml.gradients.batch_jvp(
             tapes, tangents, param_shift, reduction=lambda jvps, x: jvps.append(x)
         )
-        res = fn(dev.batch_execute(v_tapes))
+        res = fn(dev.execute(v_tapes))
         # Returned JVPs will be appended to a list, one JVP per tape
         assert len(res) == 2
diff --git a/tests/gradients/core/test_pulse_generator_gradient.py b/tests/gradients/core/test_pulse_generator_gradient.py
index 1f6b5b255f1..3396fe02210 100644
--- a/tests/gradients/core/test_pulse_generator_gradient.py
+++ b/tests/gradients/core/test_pulse_generator_gradient.py
@@ -40,7 +40,7 @@
 def grad_fn(tape, dev, fn=pulse_generator, **kwargs):
     """Utility function to automate execution and processing of gradient tapes"""
     _tapes, fn = fn(tape, **kwargs)
-    return fn(dev.batch_execute(_tapes)), _tapes
+    return fn(dev.execute(_tapes)), _tapes
 
 
 def integral_of_polyval(params, t):
@@ -474,7 +474,7 @@ def test_output_properties(self, ops_and_meas, op_idx, ops):
         """Test that the input tape and inserted ops are taken into account correctly."""
         evolve_op = qml.evolve(qml.pulse.constant * Z("a"))([np.array(0.2)], 0.2)
         operations, measurements = ops_and_meas
-        operations = [evolve_op if op == "evolve_op" else op for op in operations]
+        operations = [evolve_op if isinstance(op, str) else op for op in operations]
         tape = qml.tape.QuantumScript(operations, measurements)
         new_tapes = _insert_op(tape, ops, op_idx)
         assert isinstance(new_tapes, list) and len(new_tapes) == len(ops)
@@ -870,7 +870,7 @@ def test_no_trainable_params_multiple_return_tape(self):
 
         with pytest.warns(UserWarning, match="gradient of a tape with no trainable parameters"):
             _tapes, fn = pulse_generator(tape)
-        res = fn(dev.batch_execute(_tapes))
+        res = fn(dev.execute(_tapes))
 
         assert _tapes == []
         assert isinstance(res, tuple)
@@ -965,19 +965,21 @@ def test_all_zero_diff_methods_multiple_returns_tape(self):
         assert np.allclose(res_pulse_gen[1][2], 0)
 
 
+# TODO: add default.qubit once it supports PRNG key
 @pytest.mark.jax
+@pytest.mark.parametrize("dev_name", ["default.qubit.jax"])
 class TestPulseGeneratorTape:
     """Test that differentiating tapes with ``pulse_generator`` works."""
 
     @pytest.mark.parametrize("shots, tol", [(None, 1e-7), (1000, 0.05), ([1000, 100], 0.05)])
-    def test_single_pulse_single_term(self, shots, tol):
+    def test_single_pulse_single_term(self, dev_name, shots, tol):
         """Test that a single pulse with a single Hamiltonian term is
         differentiated correctly."""
         import jax
         import jax.numpy as jnp
 
         prng_key = jax.random.PRNGKey(8251)
-        dev = qml.device("default.qubit.jax", wires=1, shots=shots, prng_key=prng_key)
+        dev = qml.device(dev_name, wires=1, shots=shots, prng_key=prng_key)
 
         H = jnp.polyval * X(0)
         x = jnp.array([0.4, 0.2, 0.1])
@@ -1005,15 +1007,15 @@ def test_single_pulse_single_term(self, shots, tol):
 
     @pytest.mark.slow
     @pytest.mark.parametrize("shots, tol", [(None, 1e-7), ([1000, 100], 0.05)])
-    def test_single_pulse_multi_term(self, shots, tol):
+    def test_single_pulse_multi_term(self, dev_name, shots, tol):
         """Test that a single pulse with multiple Hamiltonian terms is
         differentiated correctly."""
         import jax
         import jax.numpy as jnp
 
         prng_key = jax.random.PRNGKey(8251)
-        dev = qml.device("default.qubit.jax", wires=1, shots=None)
-        dev_shots = qml.device("default.qubit.jax", wires=1, shots=shots, prng_key=prng_key)
+        dev = qml.device(dev_name, wires=1, shots=None)
+        dev_shots = qml.device(dev_name, wires=1, shots=shots, prng_key=prng_key)
 
         H = 0.1 * Z(0) + jnp.polyval * X(0) + qml.pulse.constant * Y(0)
         x = jnp.array([0.4, 0.2, 0.1])
@@ -1045,13 +1047,13 @@ def circuit(par):
             assert all(qml.math.allclose(g, e, atol=tol) for g, e in zip(grad, exp_grad))
 
     @pytest.mark.parametrize("argnum", (0, [0], 1, [1]))
-    def test_single_pulse_multi_term_argnum(self, argnum):
+    def test_single_pulse_multi_term_argnum(self, dev_name, argnum):
         """Test that a single pulse with multiple Hamiltonian terms is
         differentiated correctly when setting ``argnum``."""
         import jax
         import jax.numpy as jnp
 
-        dev = qml.device("default.qubit.jax", wires=1)
+        dev = qml.device(dev_name, wires=1)
 
         H = jnp.polyval * X(0) + qml.pulse.constant * X(0)
         x = jnp.array([0.4, 0.2, 0.1])
@@ -1085,15 +1087,15 @@ def test_single_pulse_multi_term_argnum(self, argnum):
 
     @pytest.mark.slow
     @pytest.mark.parametrize("shots, tol", [(None, 1e-7), ([1000, 100], 0.05)])
-    def test_multi_pulse(self, shots, tol):
+    def test_multi_pulse(self, dev_name, shots, tol):
         """Test that a single pulse with multiple Hamiltonian terms is
         differentiated correctly."""
         import jax
         import jax.numpy as jnp
 
         prng_key = jax.random.PRNGKey(8251)
-        dev = qml.device("default.qubit.jax", wires=1, shots=None)
-        dev_shots = qml.device("default.qubit.jax", wires=1, shots=shots, prng_key=prng_key)
+        dev = qml.device(dev_name, wires=1, shots=None)
+        dev_shots = qml.device(dev_name, wires=1, shots=shots, prng_key=prng_key)
 
         H0 = 0.1 * Z(0) + jnp.polyval * X(0)
         H1 = 0.2 * Y(0) + qml.pulse.constant * Y(0) + jnp.polyval * Z(0)
@@ -1129,14 +1131,15 @@ def circuit(par):
 
 
 @pytest.mark.jax
+@pytest.mark.parametrize("dev_name", ["default.qubit", "default.qubit.jax"])
 class TestPulseGeneratorQNode:
     """Test that pulse_generator integrates correctly with QNodes."""
 
-    def test_raises_for_application_to_qnodes(self):
+    def test_raises_for_application_to_qnodes(self, dev_name):
         """Test that an error is raised when applying ``stoch_pulse_grad``
         to a QNode directly."""
 
-        dev = qml.device("default.qubit.jax", wires=1)
+        dev = qml.device(dev_name, wires=1)
         ham_single_q_const = qml.pulse.constant * qml.PauliY(0)
 
         @qml.qnode(dev, interface="jax")
@@ -1150,14 +1153,14 @@ def circuit(params):
 
     # TODO: include the following tests when #4225 is resolved.
     @pytest.mark.skip("Applying this gradient transform to QNodes directly is not supported.")
-    def test_qnode_expval_single_par(self):
+    def test_qnode_expval_single_par(self, dev_name):
         """Test that a simple qnode that returns an expectation value
         can be differentiated with pulse_generator."""
         import jax
         import jax.numpy as jnp
 
         jax.config.update("jax_enable_x64", True)
-        dev = qml.device("default.qubit.jax", wires=1)
+        dev = qml.device(dev_name, wires=1)
         T = 0.2
         ham_single_q_const = qml.pulse.constant * Y(0)
 
@@ -1176,14 +1179,14 @@ def circuit(params):
         assert tracker.totals["executions"] == 2  # two shifted tapes
 
     @pytest.mark.skip("Applying this gradient transform to QNodes directly is not supported.")
-    def test_qnode_expval_probs_single_par(self):
+    def test_qnode_expval_probs_single_par(self, dev_name):
         """Test that a simple qnode that returns an expectation value
         can be differentiated with pulse_generator."""
         import jax
         import jax.numpy as jnp
 
         jax.config.update("jax_enable_x64", True)
-        dev = qml.device("default.qubit.jax", wires=1)
+        dev = qml.device(dev_name, wires=1)
         T = 0.2
         ham_single_q_const = jnp.polyval * Y(0)
 
@@ -1208,14 +1211,14 @@ def circuit(params):
             assert qml.math.allclose(j, e)
 
     @pytest.mark.skip("Applying this gradient transform to QNodes directly is not supported.")
-    def test_qnode_probs_expval_multi_par(self):
+    def test_qnode_probs_expval_multi_par(self, dev_name):
         """Test that a simple qnode that returns probabilities
         can be differentiated with pulse_generator."""
         import jax
         import jax.numpy as jnp
 
         jax.config.update("jax_enable_x64", True)
-        dev = qml.device("default.qubit.jax", wires=1, shots=None)
+        dev = qml.device(dev_name, wires=1, shots=None)
         T = 0.2
         ham_single_q_const = jnp.polyval * Y(0) + qml.pulse.constant * Y(0)
 
@@ -1248,17 +1251,18 @@ def circuit(params, c):
 
 
 @pytest.mark.jax
+@pytest.mark.parametrize("dev_name", ["default.qubit", "default.qubit.jax"])
 class TestPulseGeneratorIntegration:
     """Test that pulse_generator integrates correctly with QNodes."""
 
-    def test_simple_qnode_expval(self):
+    def test_simple_qnode_expval(self, dev_name):
         """Test that a simple qnode that returns an expectation value
         can be differentiated with pulse_generator."""
         import jax
         import jax.numpy as jnp
 
         jax.config.update("jax_enable_x64", True)
-        dev = qml.device("default.qubit.jax", wires=1)
+        dev = qml.device(dev_name, wires=1)
         T = 0.2
         ham_single_q_const = qml.pulse.constant * Y(0)
 
@@ -1275,14 +1279,14 @@ def circuit(params):
         assert qml.math.allclose(grad, exp_grad)
         assert tracker.totals["executions"] == 1 + 2  # one forward pass, two shifted tapes
 
-    def test_simple_qnode_expval_two_evolves(self):
+    def test_simple_qnode_expval_two_evolves(self, dev_name):
         """Test that a simple qnode that returns an expectation value
         can be differentiated with pulse_generator."""
         import jax
         import jax.numpy as jnp
 
         jax.config.update("jax_enable_x64", True)
-        dev = qml.device("default.qubit.jax", wires=1)
+        dev = qml.device(dev_name, wires=1)
         T_x = 0.1
         T_y = 0.2
         ham_x = qml.pulse.constant * X(0)
@@ -1301,14 +1305,14 @@ def circuit(params):
         exp_grad = [[-2 * jnp.sin(2 * (p_x + p_y)) * T_x], [-2 * jnp.sin(2 * (p_x + p_y)) * T_y]]
         assert qml.math.allclose(grad, exp_grad)
 
-    def test_simple_qnode_probs(self):
+    def test_simple_qnode_probs(self, dev_name):
         """Test that a simple qnode that returns probabilities
         can be differentiated with pulse_generator."""
         import jax
         import jax.numpy as jnp
 
         jax.config.update("jax_enable_x64", True)
-        dev = qml.device("default.qubit.jax", wires=1)
+        dev = qml.device(dev_name, wires=1)
         T = 0.2
         ham_single_q_const = qml.pulse.constant * Y(0)
 
@@ -1323,14 +1327,14 @@ def circuit(params):
         exp_jac = jnp.array([-1, 1]) * jnp.sin(2 * p) * T
         assert qml.math.allclose(jac, exp_jac)
 
-    def test_simple_qnode_probs_expval(self):
+    def test_simple_qnode_probs_expval(self, dev_name):
         """Test that a simple qnode that returns probabilities
         can be differentiated with pulse_generator."""
         import jax
         import jax.numpy as jnp
 
         jax.config.update("jax_enable_x64", True)
-        dev = qml.device("default.qubit.jax", wires=1)
+        dev = qml.device(dev_name, wires=1)
         T = 0.2
         ham_single_q_const = jnp.polyval * Y(0)
 
@@ -1352,13 +1356,13 @@ def circuit(params):
 
     @pytest.mark.xfail
     @pytest.mark.parametrize("time_interface", ["python", "numpy", "jax"])
-    def test_simple_qnode_jit(self, time_interface):
+    def test_simple_qnode_jit(self, dev_name, time_interface):
         """Test that a simple qnode can be differentiated with pulse_generator."""
         import jax
         import jax.numpy as jnp
 
         jax.config.update("jax_enable_x64", True)
-        dev = qml.device("default.qubit.jax", wires=1)
+        dev = qml.device(dev_name, wires=1)
         T = {"python": 0.2, "numpy": np.array(0.2), "jax": jnp.array(0.2)}[time_interface]
         ham_single_q_const = qml.pulse.constant * Y(0)
 
@@ -1374,7 +1378,7 @@ def circuit(params, T=None):
         assert qml.math.isclose(jit_grad, exp_grad)
 
     @pytest.mark.slow
-    def test_advanced_qnode(self):
+    def test_advanced_qnode(self, dev_name):
         """Test that an advanced qnode can be differentiated with pulse_generator."""
         import jax
         import jax.numpy as jnp
@@ -1382,7 +1386,7 @@ def test_advanced_qnode(self):
         jax.config.update("jax_enable_x64", True)
 
         params = [jnp.array(0.21), jnp.array(-0.171), jnp.array([0.05, 0.03, -0.1])]
-        dev = qml.device("default.qubit.jax", wires=2)
+        dev = qml.device(dev_name, wires=2)
         ham = (
             qml.pulse.constant * X(0)
             + (lambda p, t: jnp.sin(p * t)) * Z(0)
@@ -1401,8 +1405,9 @@ def ansatz(params):
         )
         qnode_backprop = qml.QNode(ansatz, dev, interface="jax")
 
-        grad_pulse_grad = jax.grad(qnode_pulse_grad)(params)
-        assert dev.num_executions == 1 + 12  # one forward execution, dim(DLA)=6
+        with qml.Tracker(dev) as tracker:
+            grad_pulse_grad = jax.grad(qnode_pulse_grad)(params)
+        assert tracker.totals["executions"] == 1 + 12  # one forward execution, dim(DLA)=6
         grad_backprop = jax.grad(qnode_backprop)(params)
 
         assert all(
@@ -1410,14 +1415,14 @@ def ansatz(params):
         )
 
     @pytest.mark.parametrize("argnums", [[0, 1], 0, 1])
-    def test_simple_qnode_expval_multiple_params(self, argnums):
+    def test_simple_qnode_expval_multiple_params(self, dev_name, argnums):
         """Test that a simple qnode with two parameters
         can be differentiated with pulse_generator and `argnums` works as expected."""
         import jax
         import jax.numpy as jnp
 
         jax.config.update("jax_enable_x64", True)
-        dev = qml.device("default.qubit.jax", wires=1)
+        dev = qml.device(dev_name, wires=1)
         T = 0.2
         ham1 = qml.pulse.constant * Y(0)
         ham2 = qml.pulse.constant * Y(0)
@@ -1442,7 +1447,9 @@ def circuit(param1, param2):
             assert tracker.totals["executions"] == 1 + 2  # one forward pass, two shifted tapes
 
 
+# TODO: port ParametrizedEvolution to new default.qubit
 @pytest.mark.jax
+@pytest.mark.parametrize("dev_name", ["default.qubit.jax"])
 class TestPulseGeneratorDiff:
     """Test that pulse_generator is differentiable, i.e. that computing
     the derivative with pulse_generator is differentiable a second time,
@@ -1450,14 +1457,14 @@ class TestPulseGeneratorDiff:
 
     # pylint: disable=too-few-public-methods
     @pytest.mark.slow
-    def test_jax(self):
+    def test_jax(self, dev_name):
         """Test that pulse_generator is differentiable,
         allowing to compute the Hessian, with JAX.."""
         import jax
         import jax.numpy as jnp
 
         jax.config.update("jax_enable_x64", True)
-        dev = qml.device("default.qubit.jax", wires=1)
+        dev = qml.device(dev_name, wires=1)
         T = 0.5
         ham_single_q_const = qml.pulse.constant * Y(0)
 
diff --git a/tests/gradients/core/test_pulse_gradient.py b/tests/gradients/core/test_pulse_gradient.py
index ac6abd9d2cb..38ea9f3af65 100644
--- a/tests/gradients/core/test_pulse_gradient.py
+++ b/tests/gradients/core/test_pulse_gradient.py
@@ -823,6 +823,7 @@ def test_raises_for_invalid_reorder_fn(self, reorder_fn):
 
 
 @pytest.mark.jax
+@pytest.mark.parametrize("dev_name", ["default.qubit", "default.qubit.jax"])
 class TestStochPulseGrad:
     """Test working cases of stoch_pulse_grad."""
 
@@ -856,7 +857,7 @@ def sine(p, t):
             ),
         ),
     )
-    def test_all_zero_grads(self, ops, arg, exp_shapes):
+    def test_all_zero_grads(self, dev_name, ops, arg, exp_shapes):  # pylint:disable=unused-argument
         """Test that a zero gradient is returned when all trainable parameters are
         identified to have zero gradient in advance."""
         import jax
@@ -878,7 +879,7 @@ def test_all_zero_grads(self, ops, arg, exp_shapes):
                 assert qml.math.allclose(r, np.zeros(exp_shape))
         jax.clear_caches()
 
-    def test_some_zero_grads(self):
+    def test_some_zero_grads(self, dev_name):
         """Test that a zero gradient is returned for trainable parameters that are
         identified to have a zero gradient in advance."""
         import jax
@@ -894,7 +895,7 @@ def test_some_zero_grads(self):
         tapes, fn = stoch_pulse_grad(tape, num_split_times=3)
         assert len(tapes) == 2 * 3
 
-        dev = qml.device("default.qubit.jax", wires=2)
+        dev = qml.device(dev_name, wires=2)
         res = fn(qml.execute(tapes, dev, None))
         assert isinstance(res, tuple) and len(res) == 2
         assert qml.math.allclose(res[0][0], np.zeros(5))
@@ -903,7 +904,7 @@ def test_some_zero_grads(self):
 
     @pytest.mark.parametrize("num_split_times", [1, 3])
     @pytest.mark.parametrize("t", [2.0, 3, (0.5, 0.6), (0.1, 0.9, 1.2)])
-    def test_constant_ry(self, num_split_times, t):
+    def test_constant_ry(self, dev_name, num_split_times, t):
         """Test that the derivative of a pulse generated by a constant Hamiltonian,
         which is a Pauli word, is computed correctly."""
         import jax
@@ -916,7 +917,7 @@ def test_constant_ry(self, num_split_times, t):
         op = qml.evolve(ham_single_q_const)(params, t)
         tape = qml.tape.QuantumScript([op], [qml.expval(qml.PauliZ(0))])
 
-        dev = qml.device("default.qubit.jax", wires=1)
+        dev = qml.device(dev_name, wires=1)
         # Effective rotation parameter
         p = params[0] * (delta_t := (T[-1] - T[0]))
         r = qml.execute([tape], dev, None)
@@ -930,7 +931,7 @@ def test_constant_ry(self, num_split_times, t):
 
     @pytest.mark.parametrize("num_split_times", [1, 3])
     @pytest.mark.parametrize("t", [2.0, 3, (0.5, 0.6), (0.1, 0.9, 1.2)])
-    def test_constant_ry_rescaled(self, num_split_times, t):
+    def test_constant_ry_rescaled(self, dev_name, num_split_times, t):
         """Test that the derivative of a pulse generated by a constant Hamiltonian,
         which is a Pauli sentence, is computed correctly."""
         import jax
@@ -945,7 +946,7 @@ def test_constant_ry_rescaled(self, num_split_times, t):
         op = qml.evolve(ham_single_q_const)(params, t)
         tape = qml.tape.QuantumScript([op], [qml.expval(qml.PauliZ(0))])
 
-        dev = qml.device("default.qubit.jax", wires=1)
+        dev = qml.device(dev_name, wires=1)
         # Prefactor due to the generator being a Pauli sentence
         prefactor = np.sqrt(0.85)
         # Effective rotation parameter
@@ -960,7 +961,7 @@ def test_constant_ry_rescaled(self, num_split_times, t):
         jax.clear_caches()
 
     @pytest.mark.parametrize("t", [0.02, (0.5, 0.6)])
-    def test_sin_envelope_rz_expval(self, t):
+    def test_sin_envelope_rz_expval(self, dev_name, t):
         """Test that the derivative of a pulse with a sine wave envelope
         is computed correctly when returning an expectation value."""
         import jax
@@ -968,7 +969,7 @@ def test_sin_envelope_rz_expval(self, t):
 
         T = t if isinstance(t, tuple) else (0, t)
 
-        dev = qml.device("default.qubit.jax", wires=1)
+        dev = qml.device(dev_name, wires=1)
         params = [jnp.array([2.3, -0.245])]
 
         ham = self.sine * qml.PauliZ(0)
@@ -999,7 +1000,7 @@ def test_sin_envelope_rz_expval(self, t):
         jax.clear_caches()
 
     @pytest.mark.parametrize("t", [0.02, (0.5, 0.6)])
-    def test_sin_envelope_rx_probs(self, t):
+    def test_sin_envelope_rx_probs(self, dev_name, t):
         """Test that the derivative of a pulse with a sine wave envelope
         is computed correctly when returning probabilities."""
         import jax
@@ -1007,7 +1008,7 @@ def test_sin_envelope_rx_probs(self, t):
 
         T = t if isinstance(t, tuple) else (0, t)
 
-        dev = qml.device("default.qubit.jax", wires=1)
+        dev = qml.device(dev_name, wires=1)
         params = [jnp.array([2.3, -0.245])]
 
         ham = self.sine * qml.PauliX(0)
@@ -1040,7 +1041,7 @@ def test_sin_envelope_rx_probs(self, t):
         jax.clear_caches()
 
     @pytest.mark.parametrize("t", [0.02, (0.5, 0.6)])
-    def test_sin_envelope_rx_expval_probs(self, t):
+    def test_sin_envelope_rx_expval_probs(self, dev_name, t):
         """Test that the derivative of a pulse with a sine wave envelope
         is computed correctly when returning expectation."""
         import jax
@@ -1048,7 +1049,7 @@ def test_sin_envelope_rx_expval_probs(self, t):
 
         T = t if isinstance(t, tuple) else (0, t)
 
-        dev = qml.device("default.qubit.jax", wires=1)
+        dev = qml.device(dev_name, wires=1)
         params = [jnp.array([2.3, -0.245])]
 
         ham = self.sine * qml.PauliX(0)
@@ -1085,7 +1086,7 @@ def test_sin_envelope_rx_expval_probs(self, t):
         jax.clear_caches()
 
     @pytest.mark.parametrize("t", [0.02, (0.5, 0.6)])
-    def test_pwc_envelope_rx(self, t):
+    def test_pwc_envelope_rx(self, dev_name, t):
         """Test that the derivative of a pulse generated by a piecewise constant Hamiltonian
         is computed correctly."""
         import jax
@@ -1093,7 +1094,7 @@ def test_pwc_envelope_rx(self, t):
 
         T = t if isinstance(t, tuple) else (0, t)
 
-        dev = qml.device("default.qubit.jax", wires=1)
+        dev = qml.device(dev_name, wires=1)
         params = [jnp.array([0.24, 0.9, -0.1, 2.3, -0.245])]
         op = qml.evolve(qml.pulse.pwc(t) * qml.PauliZ(0))(params, t)
         tape = qml.tape.QuantumScript([qml.Hadamard(0), op], [qml.expval(qml.PauliX(0))])
@@ -1115,7 +1116,7 @@ def test_pwc_envelope_rx(self, t):
         jax.clear_caches()
 
     @pytest.mark.parametrize("t", [2.0, 3, (0.5, 0.6)])
-    def test_constant_commuting(self, t):
+    def test_constant_commuting(self, dev_name, t):
         """Test that the derivative of a pulse generated by two constant commuting Hamiltonians
         is computed correctly."""
         import jax
@@ -1129,7 +1130,7 @@ def test_constant_commuting(self, t):
         )
         tape = qml.tape.QuantumScript([op], [qml.expval(qml.PauliZ(0) @ qml.PauliZ(1))])
 
-        dev = qml.device("default.qubit.jax", wires=2)
+        dev = qml.device(dev_name, wires=2)
         r = qml.execute([tape], dev, None)
         # Effective rotation parameters
         p = [_p * (T[1] - T[0]) for _p in params]
@@ -1145,7 +1146,7 @@ def test_constant_commuting(self, t):
         assert qml.math.allclose(res, exp_grad)
         jax.clear_caches()
 
-    def test_advanced_pulse(self):
+    def test_advanced_pulse(self, dev_name):
         """Test the derivative of a more complex pulse."""
         import jax
         import jax.numpy as jnp
@@ -1159,7 +1160,7 @@ def test_advanced_pulse(self):
             * qml.dot([1.0, 0.4], [qml.PauliY(0) @ qml.PauliY(1), qml.PauliX(0) @ qml.PauliX(1)])
         )
         params = [jnp.array(1.51), jnp.array(-0.371), jnp.array([0.2, 0.2, -0.4])]
-        dev = qml.device("default.qubit.jax", wires=2)
+        dev = qml.device(dev_name, wires=2)
 
         @qml.qnode(dev, interface="jax")
         def qnode(params):
@@ -1183,7 +1184,7 @@ def qnode(params):
         assert all(qml.math.allclose(r, e, rtol=0.4) for r, e in zip(res, exp_grad))
         jax.clear_caches()
 
-    def test_randomness(self):
+    def test_randomness(self, dev_name):
         """Test that the derivative of a pulse is exactly the same when reusing a seed and
         that it differs when using a different seed."""
         import jax
@@ -1213,7 +1214,7 @@ def test_randomness(self):
                 else:
                     assert qml.equal(op_a_0, op_a_1) and qml.equal(op_a_0, op_b)
 
-        dev = qml.device("default.qubit.jax", wires=1)
+        dev = qml.device(dev_name, wires=1)
         res_a_0 = fn_a_0(qml.execute(tapes_a_0, dev, None))
         res_a_1 = fn_a_1(qml.execute(tapes_a_1, dev, None))
         res_b = fn_b(qml.execute(tapes_b, dev, None))
@@ -1222,7 +1223,7 @@ def test_randomness(self):
         assert not res_a_0 == res_b
         jax.clear_caches()
 
-    def test_two_pulses(self):
+    def test_two_pulses(self, dev_name):
         """Test that the derivatives of two pulses in a circuit are computed correctly."""
         import jax
         import jax.numpy as jnp
@@ -1233,7 +1234,7 @@ def test_two_pulses(self):
         ham_1 = qml.dot([0.3, jnp.polyval], [qml.PauliZ(0), qml.PauliY(0) @ qml.PauliY(1)])
         params_0 = [jnp.array(1.51), jnp.array(-0.371)]
         params_1 = [jnp.array([0.2, 0.2, -0.4])]
-        dev = qml.device("default.qubit.jax", wires=2)
+        dev = qml.device(dev_name, wires=2)
 
         @qml.qnode(dev, interface="jax")
         def qnode(params_0, params_1):
@@ -1262,13 +1263,13 @@ def qnode(params_0, params_1):
             (qml.Hamiltonian([0.25, 1.2], [qml.PauliX(0), qml.PauliX(0) @ qml.PauliZ(1)]), 8, 1.45),
         ],
     )
-    def test_with_jit(self, generator, exp_num_tapes, prefactor):
+    def test_with_jit(self, dev_name, generator, exp_num_tapes, prefactor):
         """Test that the stochastic parameter-shift rule works with JITting."""
         import jax
         import jax.numpy as jnp
 
         jax.config.update("jax_enable_x64", True)
-        dev = qml.device("default.qubit.jax", wires=len(generator.wires))
+        dev = qml.device(dev_name, wires=len(generator.wires))
         T = (0.2, 0.5)
         ham_single_q_const = qml.dot([qml.pulse.constant], [generator])
         meas = [qml.expval(qml.PauliZ(0))]
@@ -1280,7 +1281,7 @@ def fun(params):
             tape = qml.tape.QuantumScript([op], meas)
             tapes, fn = stoch_pulse_grad(tape)
             assert len(tapes) == exp_num_tapes
-            res = fn(qml.execute(tapes, dev, None))
+            res = fn(qml.execute(tapes, dev, "backprop"))
             return res
 
         params = [jnp.array(0.24)]
@@ -1293,7 +1294,7 @@ def fun(params):
         jax.clear_caches()
 
     @pytest.mark.parametrize("shots", [None, 100])
-    def test_shots_attribute(self, shots):
+    def test_shots_attribute(self, dev_name, shots):  # pylint:disable=unused-argument
         """Tests that the shots attribute is copied to the new tapes"""
         tape = qml.tape.QuantumTape([], [qml.expval(qml.PauliZ(0)), qml.probs([1, 2])], shots=shots)
         with pytest.warns(UserWarning, match="Attempted to compute the gradient of a tape with no"):
@@ -1303,13 +1304,14 @@ def test_shots_attribute(self, shots):
 
 
 @pytest.mark.jax
+@pytest.mark.parametrize("dev_name", ["default.qubit", "default.qubit.jax"])
 class TestStochPulseGradQNode:
     """Test that pulse_generator integrates correctly with QNodes."""
 
-    def test_raises_for_application_to_qnodes(self):
+    def test_raises_for_application_to_qnodes(self, dev_name):
         """Test that an error is raised when applying ``stoch_pulse_grad``
         to a QNode directly."""
-        dev = qml.device("default.qubit.jax", wires=1)
+        dev = qml.device(dev_name, wires=1)
         ham_single_q_const = qml.pulse.constant * qml.PauliY(0)
 
         @qml.qnode(dev, interface="jax")
@@ -1323,14 +1325,14 @@ def circuit(params):
 
     # TODO: include the following tests when #4225 is resolved.
     @pytest.mark.skip("Applying this gradient transform to QNodes directly is not supported.")
-    def test_qnode_expval_single_par(self):
+    def test_qnode_expval_single_par(self, dev_name):
         """Test that a simple qnode that returns an expectation value
         can be differentiated with pulse_generator."""
         import jax
         import jax.numpy as jnp
 
         jax.config.update("jax_enable_x64", True)
-        dev = qml.device("default.qubit.jax", wires=1)
+        dev = qml.device(dev_name, wires=1)
         T = 0.2
         ham_single_q_const = qml.pulse.constant * qml.PauliY(0)
 
@@ -1350,20 +1352,22 @@ def circuit(params):
         assert tracker.totals["executions"] == 4  # two shifted tapes, two splitting times
 
 
+# TODO: add default.qubit once it supports PRNG key
 @pytest.mark.jax
+@pytest.mark.parametrize("dev_name", ["default.qubit.jax"])
 class TestStochPulseGradIntegration:
     """Test that stoch_pulse_grad integrates correctly with QNodes and ML interfaces."""
 
     @pytest.mark.parametrize("shots, tol", [(None, 1e-4), (100, 0.1), ([100, 99], 0.1)])
     @pytest.mark.parametrize("num_split_times", [1, 2])
-    def test_simple_qnode_expval(self, num_split_times, shots, tol):
+    def test_simple_qnode_expval(self, dev_name, num_split_times, shots, tol):
         """Test that a simple qnode that returns an expectation value
         can be differentiated with stoch_pulse_grad."""
         import jax
         import jax.numpy as jnp
 
         jax.config.update("jax_enable_x64", True)
-        dev = qml.device("default.qubit.jax", wires=1, shots=shots, prng_key=jax.random.PRNGKey(74))
+        dev = qml.device(dev_name, wires=1, shots=shots, prng_key=jax.random.PRNGKey(74))
         T = 0.2
         ham_single_q_const = qml.pulse.constant * qml.PauliY(0)
 
@@ -1383,14 +1387,14 @@ def circuit(params):
 
     @pytest.mark.parametrize("shots, tol", [(None, 1e-4), (100, 0.1), ([100, 99], 0.1)])
     @pytest.mark.parametrize("num_split_times", [1, 2])
-    def test_simple_qnode_expval_two_evolves(self, num_split_times, shots, tol):
+    def test_simple_qnode_expval_two_evolves(self, dev_name, num_split_times, shots, tol):
         """Test that a simple qnode that returns an expectation value
         can be differentiated with stoch_pulse_grad."""
         import jax
         import jax.numpy as jnp
 
         jax.config.update("jax_enable_x64", True)
-        dev = qml.device("default.qubit.jax", wires=1, shots=shots, prng_key=jax.random.PRNGKey(74))
+        dev = qml.device(dev_name, wires=1, shots=shots, prng_key=jax.random.PRNGKey(74))
         T_x = 0.1
         T_y = 0.2
         ham_x = qml.pulse.constant * qml.PauliX(0)
@@ -1414,14 +1418,14 @@ def circuit(params):
 
     @pytest.mark.parametrize("shots, tol", [(None, 1e-4), (100, 0.1), ([100, 99], 0.1)])
     @pytest.mark.parametrize("num_split_times", [1, 2])
-    def test_simple_qnode_probs(self, num_split_times, shots, tol):
+    def test_simple_qnode_probs(self, dev_name, num_split_times, shots, tol):
         """Test that a simple qnode that returns an probabilities
         can be differentiated with stoch_pulse_grad."""
         import jax
         import jax.numpy as jnp
 
         jax.config.update("jax_enable_x64", True)
-        dev = qml.device("default.qubit.jax", wires=1, shots=shots, prng_key=jax.random.PRNGKey(74))
+        dev = qml.device(dev_name, wires=1, shots=shots, prng_key=jax.random.PRNGKey(74))
         T = 0.2
         ham_single_q_const = qml.pulse.constant * qml.PauliY(0)
 
@@ -1441,14 +1445,14 @@ def circuit(params):
 
     @pytest.mark.parametrize("shots, tol", [(None, 1e-4), (100, 0.1), ([100, 100], 0.1)])
     @pytest.mark.parametrize("num_split_times", [1, 2])
-    def test_simple_qnode_probs_expval(self, num_split_times, shots, tol):
+    def test_simple_qnode_probs_expval(self, dev_name, num_split_times, shots, tol):
         """Test that a simple qnode that returns an probabilities
         can be differentiated with stoch_pulse_grad."""
         import jax
         import jax.numpy as jnp
 
         jax.config.update("jax_enable_x64", True)
-        dev = qml.device("default.qubit.jax", wires=1, shots=shots, prng_key=jax.random.PRNGKey(74))
+        dev = qml.device(dev_name, wires=1, shots=shots, prng_key=jax.random.PRNGKey(74))
         T = 0.2
         ham_single_q_const = qml.pulse.constant * qml.PauliY(0)
 
@@ -1475,13 +1479,13 @@ def circuit(params):
     @pytest.mark.xfail
     @pytest.mark.parametrize("num_split_times", [1, 2])
     @pytest.mark.parametrize("time_interface", ["python", "numpy", "jax"])
-    def test_simple_qnode_jit(self, num_split_times, time_interface):
+    def test_simple_qnode_jit(self, dev_name, num_split_times, time_interface):
         """Test that a simple qnode can be differentiated with stoch_pulse_grad."""
         import jax
         import jax.numpy as jnp
 
         jax.config.update("jax_enable_x64", True)
-        dev = qml.device("default.qubit.jax", wires=1)
+        dev = qml.device(dev_name, wires=1)
         T = {"python": 0.2, "numpy": np.array(0.2), "jax": jnp.array(0.2)}[time_interface]
         ham_single_q_const = qml.pulse.constant * qml.PauliY(0)
 
@@ -1500,7 +1504,7 @@ def circuit(params, T=None):
         jax.clear_caches()
 
     @pytest.mark.slow
-    def test_advanced_qnode(self):
+    def test_advanced_qnode(self, dev_name):
         """Test that an advanced qnode can be differentiated with stoch_pulse_grad."""
         import jax
         import jax.numpy as jnp
@@ -1508,7 +1512,7 @@ def test_advanced_qnode(self):
         jax.config.update("jax_enable_x64", True)
 
         params = [jnp.array(0.21), jnp.array(-0.171), jnp.array([0.05, 0.03, -0.1])]
-        dev = qml.device("default.qubit.jax", wires=2)
+        dev = qml.device(dev_name, wires=2)
         ham = (
             qml.pulse.constant * qml.PauliX(0)
             + (lambda p, t: jnp.sin(p * t)) * qml.PauliZ(0)
@@ -1540,7 +1544,7 @@ def ansatz(params):
         )
         jax.clear_caches()
 
-    def test_multi_return_broadcasting_multi_shots_raises(self):
+    def test_multi_return_broadcasting_multi_shots_raises(self, dev_name):
         """Test that a simple qnode that returns an expectation value and probabilities
         can be differentiated with stoch_pulse_grad with use_broadcasting."""
         import jax
@@ -1548,7 +1552,7 @@ def test_multi_return_broadcasting_multi_shots_raises(self):
 
         jax.config.update("jax_enable_x64", True)
         shots = [100, 100]
-        dev = qml.device("default.qubit.jax", wires=1, shots=shots, prng_key=jax.random.PRNGKey(74))
+        dev = qml.device(dev_name, wires=1, shots=shots, prng_key=jax.random.PRNGKey(74))
         T = 0.2
         ham_single_q_const = qml.pulse.constant * qml.PauliY(0)
 
@@ -1571,14 +1575,14 @@ def circuit(params):
     # TODO: delete error test above and uncomment the following test case once #2690 is resolved.
     @pytest.mark.parametrize("shots, tol", [(None, 1e-4), (100, 0.1)])  # , ([100, 100], 0.1)])
     @pytest.mark.parametrize("num_split_times", [1, 2])
-    def test_qnode_probs_expval_broadcasting(self, num_split_times, shots, tol):
+    def test_qnode_probs_expval_broadcasting(self, dev_name, num_split_times, shots, tol):
         """Test that a simple qnode that returns an expectation value and probabilities
         can be differentiated with stoch_pulse_grad with use_broadcasting."""
         import jax
         import jax.numpy as jnp
 
         jax.config.update("jax_enable_x64", True)
-        dev = qml.device("default.qubit.jax", wires=1, shots=shots, prng_key=jax.random.PRNGKey(74))
+        dev = qml.device(dev_name, wires=1, shots=shots, prng_key=jax.random.PRNGKey(74))
         T = 0.2
         ham_single_q_const = qml.pulse.constant * qml.PauliY(0)
 
@@ -1607,13 +1611,13 @@ def circuit(params):
         jax.clear_caches()
 
     @pytest.mark.parametrize("num_split_times", [1, 2])
-    def test_broadcasting_coincides_with_nonbroadcasting(self, num_split_times):
+    def test_broadcasting_coincides_with_nonbroadcasting(self, dev_name, num_split_times):
         """Test that using broadcasting or not does not change the result."""
         import jax
         import jax.numpy as jnp
 
         jax.config.update("jax_enable_x64", True)
-        dev = qml.device("default.qubit.jax", wires=2)
+        dev = qml.device(dev_name, wires=2)
         T = 0.2
 
         def f(p, t):
@@ -1651,7 +1655,7 @@ def ansatz(params):
             assert qml.math.allclose(j0, j1)
         jax.clear_caches()
 
-    def test_with_drive_exact(self):
+    def test_with_drive_exact(self, dev_name):
         """Test that a HardwareHamiltonian only containing a drive is differentiated correctly
         for a constant amplitude and zero frequency and phase."""
         import jax
@@ -1660,7 +1664,7 @@ def test_with_drive_exact(self):
 
         H = qml.pulse.transmon_drive(qml.pulse.constant, 0.0, 0.0, wires=[0])
         atol = 1e-5
-        dev = qml.device("default.qubit.jax", wires=1)
+        dev = qml.device(dev_name, wires=1)
 
         def ansatz(params):
             qml.evolve(H, atol=atol)(params, t=timespan)
@@ -1676,7 +1680,7 @@ def ansatz(params):
         assert qml.math.allclose(res, exact, atol=6e-5)
         jax.clear_caches()
 
-    def test_with_drive_approx(self):
+    def test_with_drive_approx(self, dev_name):
         """Test that a HardwareHamiltonian only containing a drive is differentiated
         approximately correctly for a constant phase and zero frequency."""
         import jax
@@ -1685,7 +1689,7 @@ def test_with_drive_approx(self):
 
         H = qml.pulse.transmon_drive(1 / (2 * np.pi), qml.pulse.constant, 0.0, wires=[0])
         atol = 1e-5
-        dev = qml.device("default.qubit.jax", wires=1)
+        dev = qml.device(dev_name, wires=1)
 
         def ansatz(params):
             qml.evolve(H, atol=atol)(params, t=timespan)
@@ -1710,7 +1714,7 @@ def ansatz(params):
         jax.clear_caches()
 
     @pytest.mark.parametrize("num_params", [1, 2])
-    def test_with_two_drives(self, num_params):
+    def test_with_two_drives(self, dev_name, num_params):
         """Test that a HardwareHamiltonian only containing two drives
         is differentiated approximately correctly. The two cases
         of the parametrization test the cases where reordered parameters
@@ -1729,7 +1733,7 @@ def test_with_two_drives(self, num_params):
             amps[0], qml.pulse.constant, 0.0, wires=[0]
         ) + qml.pulse.rydberg_drive(amps[1], qml.pulse.constant, 0.0, wires=[1])
         atol = 1e-5
-        dev = qml.device("default.qubit.jax", wires=2)
+        dev = qml.device(dev_name, wires=2)
 
         def ansatz(params):
             qml.evolve(H, atol=atol)(params, t=timespan)
@@ -1753,19 +1757,21 @@ def ansatz(params):
         jax.clear_caches()
 
 
+# TODO: port ParametrizedEvolution to new default.qubit
 @pytest.mark.jax
+@pytest.mark.parametrize("dev_name", ["default.qubit.jax"])
 class TestStochPulseGradDiff:
     """Test that stoch_pulse_grad is differentiable."""
 
     # pylint: disable=too-few-public-methods
     @pytest.mark.slow
-    def test_jax(self):
+    def test_jax(self, dev_name):
         """Test that stoch_pulse_grad is differentiable with JAX."""
         import jax
         import jax.numpy as jnp
 
         jax.config.update("jax_enable_x64", True)
-        dev = qml.device("default.qubit.jax", wires=1)
+        dev = qml.device(dev_name, wires=1)
         T = 0.5
         ham_single_q_const = qml.pulse.constant * qml.PauliY(0)
 
diff --git a/tests/gradients/core/test_vjp.py b/tests/gradients/core/test_vjp.py
index 2e3824e99a3..0904d1a19c2 100644
--- a/tests/gradients/core/test_vjp.py
+++ b/tests/gradients/core/test_vjp.py
@@ -247,7 +247,7 @@ def test_single_expectation_value(self, tol):
         tapes, fn = qml.gradients.vjp(tape, dy, param_shift)
         assert len(tapes) == 4
 
-        res = fn(dev.batch_execute(tapes))
+        res = fn(dev.execute(tapes))
         assert res.shape == (2,)
 
         exp = np.array([-np.sin(y) * np.sin(x), np.cos(y) * np.cos(x)])
@@ -274,7 +274,7 @@ def test_multiple_expectation_values(self, tol):
         tapes, fn = qml.gradients.vjp(tape, dy, param_shift)
         assert len(tapes) == 4
 
-        res = fn(dev.batch_execute(tapes))
+        res = fn(dev.execute(tapes))
         assert res.shape == (2,)
 
         exp = np.array([-np.sin(x), 2 * np.cos(y)])
@@ -301,7 +301,7 @@ def test_prob_expectation_values(self, tol):
         tapes, fn = qml.gradients.vjp(tape, dy, param_shift)
         assert len(tapes) == 4
 
-        res = fn(dev.batch_execute(tapes))
+        res = fn(dev.execute(tapes))
         assert res.shape == (2,)
 
         exp = (
@@ -374,10 +374,12 @@ class TestVJPGradients:
     """Gradient tests for the vjp function"""
 
     @pytest.mark.autograd
-    def test_autograd(self, tol):
+    @pytest.mark.parametrize("dev_name", ["default.qubit", "default.qubit.autograd"])
+    def test_autograd(self, dev_name, tol):
         """Tests that the output of the VJP transform
         can be differentiated using autograd."""
-        dev = qml.device("default.qubit.autograd", wires=2)
+        dev = qml.device(dev_name, wires=2)
+        execute_fn = dev.execute if dev_name == "default.qubit" else dev.batch_execute
         params = np.array([0.543, -0.654], requires_grad=True)
 
         def cost_fn(x, dy):
@@ -387,7 +389,7 @@ def cost_fn(x, dy):
             tape = qml.tape.QuantumScript.from_queue(q)
             tape.trainable_params = {0, 1}
             tapes, fn = qml.gradients.vjp(tape, dy, param_shift)
-            vjp = fn(dev.batch_execute(tapes))
+            vjp = fn(execute_fn(tapes))
             return vjp
 
         dy = np.array([-1.0, 0.0, 0.0, 1.0], requires_grad=False)
@@ -398,12 +400,13 @@ def cost_fn(x, dy):
         assert np.allclose(res, qml.jacobian(expected)(params), atol=tol, rtol=0)
 
     @pytest.mark.torch
-    def test_torch(self, tol):
+    @pytest.mark.parametrize("dev_name", ["default.qubit", "default.qubit.torch"])
+    def test_torch(self, dev_name, tol):
         """Tests that the output of the VJP transform
         can be differentiated using Torch."""
         import torch
 
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device(dev_name, wires=2)
 
         params_np = np.array([0.543, -0.654], requires_grad=True)
         params = torch.tensor(params_np, requires_grad=True, dtype=torch.float64)
@@ -427,12 +430,14 @@ def test_torch(self, tol):
 
     @pytest.mark.tf
     @pytest.mark.slow
-    def test_tf(self, tol):
+    @pytest.mark.parametrize("dev_name", ["default.qubit", "default.qubit.tf"])
+    def test_tf(self, dev_name, tol):
         """Tests that the output of the VJP transform
         can be differentiated using TF."""
         import tensorflow as tf
 
-        dev = qml.device("default.qubit.tf", wires=2)
+        dev = qml.device(dev_name, wires=2)
+        execute_fn = dev.execute if dev_name == "default.qubit" else dev.batch_execute
 
         params_np = np.array([0.543, -0.654], requires_grad=True)
         params = tf.Variable(params_np, dtype=tf.float64)
@@ -445,7 +450,7 @@ def test_tf(self, tol):
             tape = qml.tape.QuantumScript.from_queue(q)
             tape.trainable_params = {0, 1}
             tapes, fn = qml.gradients.vjp(tape, dy, param_shift)
-            vjp = fn(dev.batch_execute(tapes))
+            vjp = fn(execute_fn(tapes))
 
         assert np.allclose(vjp, expected(params), atol=tol, rtol=0)
 
@@ -490,13 +495,15 @@ def test_tf(self, tol):
 
     @pytest.mark.jax
     @pytest.mark.slow
-    def test_jax(self, tol):
+    @pytest.mark.parametrize("dev_name", ["default.qubit", "default.qubit.jax"])
+    def test_jax(self, dev_name, tol):
         """Tests that the output of the VJP transform
         can be differentiated using JAX."""
         import jax
         from jax import numpy as jnp
 
-        dev = qml.device("default.qubit.jax", wires=2)
+        dev = qml.device(dev_name, wires=2)
+        execute_fn = dev.execute if dev_name == "default.qubit" else dev.batch_execute
         params_np = np.array([0.543, -0.654], requires_grad=True)
         params = jnp.array(params_np)
 
@@ -508,7 +515,7 @@ def cost_fn(x):
             dy = jax.numpy.array([-1.0, 0.0, 0.0, 1.0])
             tape.trainable_params = {0, 1}
             tapes, fn = qml.gradients.vjp(tape, dy, param_shift)
-            vjp = fn(dev.batch_execute(tapes))
+            vjp = fn(execute_fn(tapes))
             return vjp
 
         res = cost_fn(params)
@@ -550,7 +557,7 @@ def test_one_tape_no_trainable_parameters(self):
 
         # Even though there are 3 parameters, only two contribute
         # to the VJP, so only 2*2=4 quantum evals
-        res = fn(dev.batch_execute(v_tapes))
+        res = fn(dev.execute(v_tapes))
         assert res[0] is None
         assert res[1] is not None
 
@@ -604,7 +611,7 @@ def test_zero_dy(self):
         dys = [np.array(0.0), np.array(1.0)]
 
         v_tapes, fn = qml.gradients.batch_vjp(tapes, dys, param_shift)
-        res = fn(dev.batch_execute(v_tapes))
+        res = fn(dev.execute(v_tapes))
 
         # Even though there are 3 parameters, only two contribute
         # to the VJP, so only 2*2=4 quantum evals
@@ -635,7 +642,7 @@ def test_reduction_append(self):
         dys = [np.array(1.0), np.array(1.0)]
 
         v_tapes, fn = qml.gradients.batch_vjp(tapes, dys, param_shift, reduction="append")
-        res = fn(dev.batch_execute(v_tapes))
+        res = fn(dev.execute(v_tapes))
 
         # Returned VJPs will be appended to a list, one vjp per tape
         assert len(res) == 2
@@ -666,7 +673,7 @@ def test_reduction_extend(self):
         dys = [np.array(1.0), np.array(1.0)]
 
         v_tapes, fn = qml.gradients.batch_vjp(tapes, dys, param_shift, reduction="extend")
-        res = fn(dev.batch_execute(v_tapes))
+        res = fn(dev.execute(v_tapes))
 
         # Returned VJPs will be extended into a list. Each element of the returned
         # list will correspond to a single input parameter of the combined
diff --git a/tests/gradients/finite_diff/test_finite_difference.py b/tests/gradients/finite_diff/test_finite_difference.py
index 003a3cfe665..0e2a296d4c0 100644
--- a/tests/gradients/finite_diff/test_finite_difference.py
+++ b/tests/gradients/finite_diff/test_finite_difference.py
@@ -130,7 +130,7 @@ def test_non_differentiable_error(self):
         dev = qml.device("default.qubit", wires=2)
         tapes, fn = finite_diff(tape)
 
-        res = fn(dev.batch_execute(tapes))
+        res = fn(dev.execute(tapes))
         assert isinstance(res, tuple)
 
         assert isinstance(res[0], numpy.ndarray)
@@ -152,7 +152,7 @@ def test_independent_parameter(self, mocker):
         tape = qml.tape.QuantumScript.from_queue(q)
         dev = qml.device("default.qubit", wires=2)
         tapes, fn = finite_diff(tape)
-        res = fn(dev.batch_execute(tapes))
+        res = fn(dev.execute(tapes))
 
         assert isinstance(res, tuple)
         assert len(res) == 2
@@ -413,13 +413,14 @@ def test_independent_parameters(self):
 
         tape2 = qml.tape.QuantumScript.from_queue(q2)
         tapes, fn = finite_diff(tape1, approx_order=1)
-        j1 = fn(dev.batch_execute(tapes))
+        with qml.Tracker(dev) as tracker:
+            j1 = fn(dev.execute(tapes))
 
         # We should only be executing the device to differentiate 1 parameter (2 executions)
-        assert dev.num_executions == 2
+        assert tracker.latest["executions"] == 2
 
         tapes, fn = finite_diff(tape2, approx_order=1)
-        j2 = fn(dev.batch_execute(tapes))
+        j2 = fn(dev.execute(tapes))
 
         exp = -np.sin(1)
 
@@ -564,7 +565,7 @@ def test_ragged_output(self, approx_order, strategy, validate):
         tapes, fn = finite_diff(
             tape, approx_order=approx_order, strategy=strategy, validate_params=validate
         )
-        res = fn(dev.batch_execute(tapes))
+        res = fn(dev.execute(tapes))
 
         assert isinstance(res, tuple)
 
@@ -597,7 +598,7 @@ def test_single_expectation_value(self, approx_order, strategy, validate, tol):
         tapes, fn = finite_diff(
             tape, approx_order=approx_order, strategy=strategy, validate_params=validate
         )
-        res = fn(dev.batch_execute(tapes))
+        res = fn(dev.execute(tapes))
 
         assert isinstance(res, tuple)
         assert len(res) == 2
@@ -634,7 +635,7 @@ def test_single_expectation_value_with_argnum_all(self, approx_order, strategy,
             strategy=strategy,
             validate_params=validate,
         )
-        res = fn(dev.batch_execute(tapes))
+        res = fn(dev.execute(tapes))
 
         assert isinstance(res, tuple)
         assert len(res) == 2
@@ -671,7 +672,7 @@ def test_single_expectation_value_with_argnum_one(self, approx_order, strategy,
         tapes, fn = finite_diff(
             tape, argnum=1, approx_order=approx_order, strategy=strategy, validate_params=validate
         )
-        res = fn(dev.batch_execute(tapes))
+        res = fn(dev.execute(tapes))
 
         assert isinstance(res, tuple)
         assert len(res) == 2
@@ -712,7 +713,7 @@ def test_multiple_expectation_value_with_argnum_one(
         tapes, fn = finite_diff(
             tape, argnum=1, approx_order=approx_order, strategy=strategy, validate_params=validate
         )
-        res = fn(dev.batch_execute(tapes))
+        res = fn(dev.execute(tapes))
 
         assert isinstance(res, tuple)
         assert isinstance(res[0], tuple)
@@ -738,7 +739,7 @@ def test_multiple_expectation_values(self, approx_order, strategy, validate, tol
         tapes, fn = finite_diff(
             tape, approx_order=approx_order, strategy=strategy, validate_params=validate
         )
-        res = fn(dev.batch_execute(tapes))
+        res = fn(dev.execute(tapes))
 
         assert isinstance(res, tuple)
         assert len(res) == 2
@@ -773,7 +774,7 @@ def test_var_expectation_values(self, approx_order, strategy, validate, tol):
         tapes, fn = finite_diff(
             tape, approx_order=approx_order, strategy=strategy, validate_params=validate
         )
-        res = fn(dev.batch_execute(tapes))
+        res = fn(dev.execute(tapes))
 
         assert isinstance(res, tuple)
         assert len(res) == 2
@@ -808,7 +809,7 @@ def test_prob_expectation_values(self, approx_order, strategy, validate, tol):
         tapes, fn = finite_diff(
             tape, approx_order=approx_order, strategy=strategy, validate_params=validate
         )
-        res = fn(dev.batch_execute(tapes))
+        res = fn(dev.execute(tapes))
 
         assert isinstance(res, tuple)
         assert len(res) == 2
@@ -854,10 +855,12 @@ class TestFiniteDiffGradients:
     """Test that the transform is differentiable"""
 
     @pytest.mark.autograd
-    def test_autograd(self, approx_order, strategy, tol):
+    @pytest.mark.parametrize("dev_name", ["default.qubit", "default.qubit.autograd"])
+    def test_autograd(self, dev_name, approx_order, strategy, tol):
         """Tests that the output of the finite-difference transform
         can be differentiated using autograd, yielding second derivatives."""
-        dev = qml.device("default.qubit.autograd", wires=2)
+        dev = qml.device(dev_name, wires=2)
+        execute_fn = dev.execute if dev_name == "default.qubit" else dev.batch_execute
         params = np.array([0.543, -0.654], requires_grad=True)
 
         def cost_fn(x):
@@ -870,7 +873,7 @@ def cost_fn(x):
             tape = qml.tape.QuantumScript.from_queue(q)
             tape.trainable_params = {0, 1}
             tapes, fn = finite_diff(tape, n=1, approx_order=approx_order, strategy=strategy)
-            jac = np.array(fn(dev.batch_execute(tapes)))
+            jac = np.array(fn(execute_fn(tapes)))
             return jac
 
         res = qml.jacobian(cost_fn)(params)
@@ -885,10 +888,12 @@ def cost_fn(x):
         assert np.allclose(res, expected, atol=tol, rtol=0)
 
     @pytest.mark.autograd
-    def test_autograd_ragged(self, approx_order, strategy, tol):
+    @pytest.mark.parametrize("dev_name", ["default.qubit", "default.qubit.autograd"])
+    def test_autograd_ragged(self, dev_name, approx_order, strategy, tol):
         """Tests that the output of the finite-difference transform
         of a ragged tape can be differentiated using autograd, yielding second derivatives."""
-        dev = qml.device("default.qubit.autograd", wires=2)
+        dev = qml.device(dev_name, wires=2)
+        execute_fn = dev.execute if dev_name == "default.qubit" else dev.batch_execute
         params = np.array([0.543, -0.654], requires_grad=True)
 
         def cost_fn(x):
@@ -902,7 +907,7 @@ def cost_fn(x):
             tape = qml.tape.QuantumScript.from_queue(q)
             tape.trainable_params = {0, 1}
             tapes, fn = finite_diff(tape, n=1, approx_order=approx_order, strategy=strategy)
-            jac = fn(dev.batch_execute(tapes))
+            jac = fn(execute_fn(tapes))
             return jac[1][0]
 
         x, y = params
@@ -912,12 +917,14 @@ def cost_fn(x):
 
     @pytest.mark.tf
     @pytest.mark.slow
-    def test_tf(self, approx_order, strategy, tol):
+    @pytest.mark.parametrize("dev_name", ["default.qubit", "default.qubit.tf"])
+    def test_tf(self, dev_name, approx_order, strategy, tol):
         """Tests that the output of the finite-difference transform
         can be differentiated using TF, yielding second derivatives."""
         import tensorflow as tf
 
-        dev = qml.device("default.qubit.tf", wires=2)
+        dev = qml.device(dev_name, wires=2)
+        execute_fn = dev.execute if dev_name == "default.qubit" else dev.batch_execute
         params = tf.Variable([0.543, -0.654], dtype=tf.float64)
 
         with tf.GradientTape(persistent=True) as t:
@@ -930,7 +937,7 @@ def test_tf(self, approx_order, strategy, tol):
             tape = qml.tape.QuantumScript.from_queue(q)
             tape.trainable_params = {0, 1}
             tapes, fn = finite_diff(tape, n=1, approx_order=approx_order, strategy=strategy)
-            jac_0, jac_1 = fn(dev.batch_execute(tapes))
+            jac_0, jac_1 = fn(execute_fn(tapes))
 
         x, y = 1.0 * params
 
@@ -946,12 +953,14 @@ def test_tf(self, approx_order, strategy, tol):
         assert np.allclose([res_0, res_1], expected, atol=tol, rtol=0)
 
     @pytest.mark.tf
-    def test_tf_ragged(self, approx_order, strategy, tol):
+    @pytest.mark.parametrize("dev_name", ["default.qubit", "default.qubit.tf"])
+    def test_tf_ragged(self, dev_name, approx_order, strategy, tol):
         """Tests that the output of the finite-difference transform
         of a ragged tape can be differentiated using TF, yielding second derivatives."""
         import tensorflow as tf
 
-        dev = qml.device("default.qubit.tf", wires=2)
+        dev = qml.device(dev_name, wires=2)
+        execute_fn = dev.execute if dev_name == "default.qubit" else dev.batch_execute
         params = tf.Variable([0.543, -0.654], dtype=tf.float64)
 
         with tf.GradientTape(persistent=True) as t:
@@ -966,7 +975,7 @@ def test_tf_ragged(self, approx_order, strategy, tol):
             tape.trainable_params = {0, 1}
             tapes, fn = finite_diff(tape, n=1, approx_order=approx_order, strategy=strategy)
 
-            jac_01 = fn(dev.batch_execute(tapes))[1][0]
+            jac_01 = fn(execute_fn(tapes))[1][0]
 
         x, y = 1.0 * params
 
@@ -977,12 +986,14 @@ def test_tf_ragged(self, approx_order, strategy, tol):
         assert np.allclose(res_01[0], expected, atol=tol, rtol=0)
 
     @pytest.mark.torch
-    def test_torch(self, approx_order, strategy, tol):
+    @pytest.mark.parametrize("dev_name", ["default.qubit", "default.qubit.torch"])
+    def test_torch(self, dev_name, approx_order, strategy, tol):
         """Tests that the output of the finite-difference transform
         can be differentiated using Torch, yielding second derivatives."""
         import torch
 
-        dev = qml.device("default.qubit.torch", wires=2)
+        dev = qml.device(dev_name, wires=2)
+        execute_fn = dev.execute if dev_name == "default.qubit" else dev.batch_execute
         params = torch.tensor([0.543, -0.654], dtype=torch.float64, requires_grad=True)
 
         def cost_fn(params):
@@ -994,7 +1005,7 @@ def cost_fn(params):
 
             tape = qml.tape.QuantumScript.from_queue(q)
             tapes, fn = finite_diff(tape, n=1, approx_order=approx_order, strategy=strategy)
-            jac = fn(dev.batch_execute(tapes))
+            jac = fn(execute_fn(tapes))
             return jac
 
         hess = torch.autograd.functional.jacobian(cost_fn, params)
@@ -1012,7 +1023,8 @@ def cost_fn(params):
         assert np.allclose(hess[1].detach().numpy(), expected[1], atol=tol, rtol=0)
 
     @pytest.mark.jax
-    def test_jax(self, approx_order, strategy, tol):
+    @pytest.mark.parametrize("dev_name", ["default.qubit", "default.qubit.jax"])
+    def test_jax(self, dev_name, approx_order, strategy, tol):
         """Tests that the output of the finite-difference transform
         can be differentiated using JAX, yielding second derivatives."""
         import jax
@@ -1021,7 +1033,8 @@ def test_jax(self, approx_order, strategy, tol):
 
         config.update("jax_enable_x64", True)
 
-        dev = qml.device("default.qubit.jax", wires=2)
+        dev = qml.device(dev_name, wires=2)
+        execute_fn = dev.execute if dev_name == "default.qubit" else dev.batch_execute
         params = jnp.array([0.543, -0.654])
 
         def cost_fn(x):
@@ -1034,7 +1047,7 @@ def cost_fn(x):
             tape = qml.tape.QuantumScript.from_queue(q)
             tape.trainable_params = {0, 1}
             tapes, fn = finite_diff(tape, n=1, approx_order=approx_order, strategy=strategy)
-            jac = fn(dev.batch_execute(tapes))
+            jac = fn(execute_fn(tapes))
             return jac
 
         res = jax.jacobian(cost_fn)(params)
diff --git a/tests/gradients/finite_diff/test_finite_difference_shot_vec.py b/tests/gradients/finite_diff/test_finite_difference_shot_vec.py
index d335db4068f..b2179096206 100644
--- a/tests/gradients/finite_diff/test_finite_difference_shot_vec.py
+++ b/tests/gradients/finite_diff/test_finite_difference_shot_vec.py
@@ -60,7 +60,7 @@ def test_non_differentiable_error(self):
         dev = qml.device("default.qubit", wires=2, shots=default_shot_vector)
         tapes, fn = finite_diff(tape, h=h_val)
 
-        all_res = fn(dev.batch_execute(tapes))
+        all_res = fn(dev.execute(tapes))
 
         assert isinstance(all_res, tuple)
         assert len(all_res) == len(default_shot_vector)
@@ -87,7 +87,7 @@ def test_independent_parameter_skipped(self, mocker):
         tape = qml.tape.QuantumScript.from_queue(q, shots=default_shot_vector)
         dev = qml.device("default.qubit", wires=2, shots=default_shot_vector)
         tapes, fn = finite_diff(tape, h=h_val)
-        all_res = fn(dev.batch_execute(tapes))
+        all_res = fn(dev.execute(tapes))
 
         assert isinstance(all_res, tuple)
         assert len(all_res) == len(default_shot_vector)
@@ -270,7 +270,7 @@ def circuit(params):
         params = np.array([0.5, 0.5, 0.5], requires_grad=True)
 
         all_result = qml.gradients.finite_diff(circuit, h=h_val)(params)
-        assert len(all_result) == len(default_shot_vector)
+        assert len(all_result) == len(many_shots_shot_vector)
 
         for result in all_result:
             assert isinstance(result, tuple)
@@ -358,20 +358,21 @@ def test_independent_parameters(self):
 
         tape2 = qml.tape.QuantumScript.from_queue(q2, shots=many_shots_shot_vector)
         tapes, fn = finite_diff(tape1, approx_order=1, h=h_val)
-        j1 = fn(dev.batch_execute(tapes))
+        with qml.Tracker(dev) as tracker:
+            j1 = fn(dev.execute(tapes))
 
         # We should only be executing the device to differentiate 1 parameter (2 executions)
-        assert dev.num_executions == 2
+        assert tracker.latest["executions"] == 2
 
         tapes, fn = finite_diff(tape2, approx_order=1, h=h_val)
-        j2 = fn(dev.batch_execute(tapes))
+        j2 = fn(dev.execute(tapes))
 
         exp = -np.sin(1)
 
         assert isinstance(j1, tuple)
-        assert len(j1) == len(default_shot_vector)
+        assert len(j1) == len(many_shots_shot_vector)
         assert isinstance(j2, tuple)
-        assert len(j2) == len(default_shot_vector)
+        assert len(j2) == len(many_shots_shot_vector)
 
         for _j1, _j2 in zip(j1, j2):
             assert np.allclose(_j1, [exp, 0], atol=finite_diff_shot_vec_tol)
@@ -519,7 +520,7 @@ def test_ragged_output(self, approx_order, strategy, validate):
             strategy=strategy,
             validate_params=validate,
         )
-        all_res = fn(dev.batch_execute(tapes))
+        all_res = fn(dev.execute(tapes))
         assert isinstance(all_res, tuple)
         assert len(all_res) == len(many_shots_shot_vector)
 
@@ -551,7 +552,7 @@ def test_single_expectation_value(self, approx_order, strategy, validate):
             qml.CNOT(wires=[0, 1])
             qml.expval(qml.PauliZ(0) @ qml.PauliX(1))
 
-        tape = qml.tape.QuantumScript.from_queue(q, shots=default_shot_vector)
+        tape = qml.tape.QuantumScript.from_queue(q, shots=many_shots_shot_vector)
         tapes, fn = finite_diff(
             tape,
             approx_order=approx_order,
@@ -559,10 +560,10 @@ def test_single_expectation_value(self, approx_order, strategy, validate):
             validate_params=validate,
             h=h_val,
         )
-        all_res = fn(dev.batch_execute(tapes))
+        all_res = fn(dev.execute(tapes))
 
         assert isinstance(all_res, tuple)
-        assert len(all_res) == len(default_shot_vector)
+        assert len(all_res) == len(many_shots_shot_vector)
 
         expected = np.array([[-np.sin(y) * np.sin(x), np.cos(y) * np.cos(x)]])
         for res in all_res:
@@ -601,10 +602,10 @@ def test_single_expectation_value_with_argnum_all(self, approx_order, strategy,
             validate_params=validate,
             h=h_val,
         )
-        all_res = fn(dev.batch_execute(tapes))
+        all_res = fn(dev.execute(tapes))
 
         assert isinstance(all_res, tuple)
-        assert len(all_res) == len(default_shot_vector)
+        assert len(all_res) == len(many_shots_shot_vector)
 
         expected = np.array([[-np.sin(y) * np.sin(x), np.cos(y) * np.cos(x)]])
         for res in all_res:
@@ -647,10 +648,10 @@ def test_single_expectation_value_with_argnum_one(self, approx_order, strategy,
             validate_params=validate,
             h=h_val,
         )
-        all_res = fn(dev.batch_execute(tapes))
+        all_res = fn(dev.execute(tapes))
 
         assert isinstance(all_res, tuple)
-        assert len(all_res) == len(default_shot_vector)
+        assert len(all_res) == len(many_shots_shot_vector)
 
         expected = [0, np.cos(y) * np.cos(x)]
 
@@ -674,6 +675,9 @@ def test_probs_expval_with_argnum_one(self, approx_order, strategy, validate):
         This test relies on the fact that exactly one term of the estimated
         jacobian will match the expected analytical value.
         """
+        if approx_order == 4 and strategy != "center":
+            pytest.skip("The latest default.qubit is unreliable with an approx_order of 4")
+
         dev = qml.device("default.qubit", wires=2, shots=many_shots_shot_vector)
         x = 0.543
         y = -0.654
@@ -695,10 +699,10 @@ def test_probs_expval_with_argnum_one(self, approx_order, strategy, validate):
             validate_params=validate,
             h=h_val,
         )
-        all_res = fn(dev.batch_execute(tapes))
+        all_res = fn(dev.execute(tapes))
 
         assert isinstance(all_res, tuple)
-        assert len(all_res) == len(default_shot_vector)
+        assert len(all_res) == len(many_shots_shot_vector)
 
         cx, sx, cy, sy = np.cos(x / 2), np.sin(x / 2), np.cos(y / 2), np.sin(y / 2)
         # probability vector is [cx**2 * cy**2, cx**2 * sy**2, sx**2 * sy**2, sx**2 * cy**2]
@@ -734,10 +738,10 @@ def test_multiple_expectation_values(self, approx_order, strategy, validate):
             validate_params=validate,
             h=h_val,
         )
-        all_res = fn(dev.batch_execute(tapes))
+        all_res = fn(dev.execute(tapes))
 
         assert isinstance(all_res, tuple)
-        assert len(all_res) == len(default_shot_vector)
+        assert len(all_res) == len(many_shots_shot_vector)
 
         for res in all_res:
             assert isinstance(res, tuple)
@@ -777,10 +781,10 @@ def test_var_expectation_values(self, approx_order, strategy, validate):
             validate_params=validate,
             h=h_val,
         )
-        all_res = fn(dev.batch_execute(tapes))
+        all_res = fn(dev.execute(tapes))
 
         assert isinstance(all_res, tuple)
-        assert len(all_res) == len(default_shot_vector)
+        assert len(all_res) == len(many_shots_shot_vector)
 
         for res in all_res:
             assert isinstance(res, tuple)
@@ -822,10 +826,10 @@ def test_prob_expectation_values(self, approx_order, strategy, validate):
             validate_params=validate,
             h=h_val,
         )
-        all_res = fn(dev.batch_execute(tapes))
+        all_res = fn(dev.execute(tapes))
 
         assert isinstance(all_res, tuple)
-        assert len(all_res) == len(default_shot_vector)
+        assert len(all_res) == len(many_shots_shot_vector)
 
         for res in all_res:
             assert isinstance(res, tuple)
@@ -873,10 +877,12 @@ class TestFiniteDiffGradients:
     """Test that the transform is differentiable"""
 
     @pytest.mark.autograd
-    def test_autograd(self, approx_order, strategy):
+    @pytest.mark.parametrize("dev_name", ["default.qubit", "default.qubit.autograd"])
+    def test_autograd(self, dev_name, approx_order, strategy):
         """Tests that the output of the finite-difference transform
         can be differentiated using autograd, yielding second derivatives."""
-        dev = qml.device("default.qubit.autograd", wires=2, shots=many_shots_shot_vector)
+        dev = qml.device(dev_name, wires=2, shots=many_shots_shot_vector)
+        execute_fn = dev.execute if dev_name == "default.qubit" else dev.batch_execute
         params = np.array([0.543, -0.654], requires_grad=True)
 
         def cost_fn(x):
@@ -895,13 +901,13 @@ def cost_fn(x):
                 strategy=strategy,
                 h=h_val,
             )
-            jac = np.array(fn(dev.batch_execute(tapes)))
+            jac = np.array(fn(execute_fn(tapes)))
             return jac
 
         all_res = qml.jacobian(cost_fn)(params)
 
         assert isinstance(all_res, tuple)
-        assert len(all_res) == len(default_shot_vector)
+        assert len(all_res) == len(many_shots_shot_vector)
 
         for res in all_res:
             x, y = params
@@ -915,10 +921,12 @@ def cost_fn(x):
             assert np.allclose(res, expected, atol=finite_diff_shot_vec_tol, rtol=0)
 
     @pytest.mark.autograd
-    def test_autograd_ragged(self, approx_order, strategy):
+    @pytest.mark.parametrize("dev_name", ["default.qubit", "default.qubit.autograd"])
+    def test_autograd_ragged(self, dev_name, approx_order, strategy):
         """Tests that the output of the finite-difference transform
         of a ragged tape can be differentiated using autograd, yielding second derivatives."""
-        dev = qml.device("default.qubit.autograd", wires=2, shots=many_shots_shot_vector)
+        dev = qml.device(dev_name, wires=2, shots=many_shots_shot_vector)
+        execute_fn = dev.execute if dev_name == "default.qubit" else dev.batch_execute
         params = np.array([0.543, -0.654], requires_grad=True)
 
         def cost_fn(x):
@@ -938,14 +946,14 @@ def cost_fn(x):
                 strategy=strategy,
                 h=h_val,
             )
-            jac = fn(dev.batch_execute(tapes))
+            jac = fn(execute_fn(tapes))
             return jac[1][0]
 
         x, y = params
         all_res = qml.jacobian(cost_fn)(params)[0]
 
         assert isinstance(all_res, tuple)
-        assert len(all_res) == len(default_shot_vector)
+        assert len(all_res) == len(many_shots_shot_vector)
 
         for res in all_res:
             expected = np.array([-np.cos(x) * np.cos(y) / 2, np.sin(x) * np.sin(y) / 2])
@@ -953,12 +961,14 @@ def cost_fn(x):
 
     @pytest.mark.tf
     @pytest.mark.slow
-    def test_tf(self, approx_order, strategy):
+    @pytest.mark.parametrize("dev_name", ["default.qubit", "default.qubit.tf"])
+    def test_tf(self, dev_name, approx_order, strategy):
         """Tests that the output of the finite-difference transform
         can be differentiated using TF, yielding second derivatives."""
         import tensorflow as tf
 
-        dev = qml.device("default.qubit.tf", wires=2, shots=many_shots_shot_vector)
+        dev = qml.device(dev_name, wires=2, shots=many_shots_shot_vector)
+        execute_fn = dev.execute if dev_name == "default.qubit" else dev.batch_execute
         params = tf.Variable([0.543, -0.654], dtype=tf.float64)
 
         with tf.GradientTape(persistent=True) as t:
@@ -977,7 +987,7 @@ def test_tf(self, approx_order, strategy):
                 strategy=strategy,
                 h=h_val,
             )
-            jac_0, jac_1 = fn(dev.batch_execute(tapes))
+            jac_0, jac_1 = fn(execute_fn(tapes))
 
         x, y = 1.0 * params
 
@@ -993,12 +1003,14 @@ def test_tf(self, approx_order, strategy):
         assert np.allclose([res_0, res_1], expected, atol=finite_diff_shot_vec_tol, rtol=0)
 
     @pytest.mark.tf
-    def test_tf_ragged(self, approx_order, strategy):
+    @pytest.mark.parametrize("dev_name", ["default.qubit", "default.qubit.tf"])
+    def test_tf_ragged(self, dev_name, approx_order, strategy):
         """Tests that the output of the finite-difference transform
         of a ragged tape can be differentiated using TF, yielding second derivatives."""
         import tensorflow as tf
 
-        dev = qml.device("default.qubit.tf", wires=2, shots=many_shots_shot_vector)
+        dev = qml.device(dev_name, wires=2, shots=many_shots_shot_vector)
+        execute_fn = dev.execute if dev_name == "default.qubit" else dev.batch_execute
         params = tf.Variable([0.543, -0.654], dtype=tf.float64)
 
         with tf.GradientTape(persistent=True) as t:
@@ -1019,7 +1031,7 @@ def test_tf_ragged(self, approx_order, strategy):
                 h=h_val,
             )
 
-            jac_01 = fn(dev.batch_execute(tapes))[1][0]
+            jac_01 = fn(execute_fn(tapes))[1][0]
 
         x, y = 1.0 * params
 
@@ -1030,12 +1042,14 @@ def test_tf_ragged(self, approx_order, strategy):
         assert np.allclose(res_01[0], expected, atol=finite_diff_shot_vec_tol, rtol=0)
 
     @pytest.mark.torch
-    def test_torch(self, approx_order, strategy):
+    @pytest.mark.parametrize("dev_name", ["default.qubit", "default.qubit.torch"])
+    def test_torch(self, dev_name, approx_order, strategy):
         """Tests that the output of the finite-difference transform
         can be differentiated using Torch, yielding second derivatives."""
         import torch
 
-        dev = qml.device("default.qubit.torch", wires=2, shots=many_shots_shot_vector)
+        dev = qml.device(dev_name, wires=2, shots=many_shots_shot_vector)
+        execute_fn = dev.execute if dev_name == "default.qubit" else dev.batch_execute
         params = torch.tensor([0.543, -0.654], dtype=torch.float64, requires_grad=True)
 
         def cost_fn(params):
@@ -1053,7 +1067,7 @@ def cost_fn(params):
                 strategy=strategy,
                 h=h_val,
             )
-            jac = fn(dev.batch_execute(tapes))
+            jac = fn(execute_fn(tapes))
             return jac
 
         hess = torch.autograd.functional.jacobian(cost_fn, params)
@@ -1075,7 +1089,8 @@ def cost_fn(params):
         )
 
     @pytest.mark.jax
-    def test_jax(self, approx_order, strategy):
+    @pytest.mark.parametrize("dev_name", ["default.qubit", "default.qubit.jax"])
+    def test_jax(self, dev_name, approx_order, strategy):
         """Tests that the output of the finite-difference transform
         can be differentiated using JAX, yielding second derivatives."""
         import jax
@@ -1084,7 +1099,8 @@ def test_jax(self, approx_order, strategy):
 
         config.update("jax_enable_x64", True)
 
-        dev = qml.device("default.qubit.jax", wires=2, shots=many_shots_shot_vector)
+        dev = qml.device(dev_name, wires=2, shots=many_shots_shot_vector)
+        execute_fn = dev.execute if dev_name == "default.qubit" else dev.batch_execute
         params = jnp.array([0.543, -0.654])
 
         def cost_fn(x):
@@ -1103,13 +1119,13 @@ def cost_fn(x):
                 strategy=strategy,
                 h=h_val,
             )
-            jac = fn(dev.batch_execute(tapes))
+            jac = fn(execute_fn(tapes))
             return jac
 
         all_res = jax.jacobian(cost_fn)(params)
 
         assert isinstance(all_res, tuple)
-        assert len(all_res) == len(default_shot_vector)
+        assert len(all_res) == len(many_shots_shot_vector)
 
         x, y = params
         expected = np.array(
@@ -1176,7 +1192,7 @@ def test_1_1(self, shot_vec, meas, shape, op_wires):
         tape.trainable_params = {0}
 
         tapes, fn = qml.gradients.finite_diff(tape)
-        all_res = fn(dev.batch_execute(tapes))
+        all_res = fn(dev.execute(tapes))
 
         assert len(all_res) == grad_transform_shots.num_copies
         assert isinstance(all_res, tuple)
@@ -1210,7 +1226,7 @@ def test_1_N(self, shot_vec, op_wire):
         tape.trainable_params = {0}
 
         tapes, fn = qml.gradients.finite_diff(tape)
-        all_res = fn(dev.batch_execute(tapes))
+        all_res = fn(dev.execute(tapes))
 
         assert len(all_res) == grad_transform_shots.num_copies
         assert isinstance(all_res, tuple)
@@ -1244,7 +1260,7 @@ def test_N_1(self, shot_vec, meas, shape, op_wires):
         tape.trainable_params = {0, 1}
 
         tapes, fn = qml.gradients.finite_diff(tape)
-        all_res = fn(dev.batch_execute(tapes))
+        all_res = fn(dev.execute(tapes))
 
         assert len(all_res) == grad_transform_shots.num_copies
         assert isinstance(all_res, tuple)
@@ -1286,7 +1302,7 @@ def test_N_N(self, shot_vec, op_wires):
         tape.trainable_params = {0, 1, 2, 3, 4}
 
         tapes, fn = qml.gradients.finite_diff(tape)
-        all_res = fn(dev.batch_execute(tapes))
+        all_res = fn(dev.execute(tapes))
 
         assert len(all_res) == grad_transform_shots.num_copies
         assert isinstance(all_res, tuple)
diff --git a/tests/gradients/finite_diff/test_spsa_gradient.py b/tests/gradients/finite_diff/test_spsa_gradient.py
index 44a149041cb..5897aed4b88 100644
--- a/tests/gradients/finite_diff/test_spsa_gradient.py
+++ b/tests/gradients/finite_diff/test_spsa_gradient.py
@@ -201,7 +201,7 @@ def test_non_differentiable_error(self):
         dev = qml.device("default.qubit", wires=2)
         tapes, fn = spsa_grad(tape)
 
-        res = fn(dev.batch_execute(tapes))
+        res = fn(dev.execute(tapes))
         assert isinstance(res, tuple)
         assert len(res) == 2
 
@@ -225,7 +225,7 @@ def test_independent_parameter(self, num_directions, mocker):
         tape = qml.tape.QuantumScript.from_queue(q)
         dev = qml.device("default.qubit", wires=2)
         tapes, fn = spsa_grad(tape, num_directions=num_directions)
-        res = fn(dev.batch_execute(tapes))
+        res = fn(dev.execute(tapes))
 
         assert isinstance(res, tuple)
         assert len(res) == 2
@@ -498,13 +498,14 @@ def test_independent_parameters(self):
         tapes, fn = spsa_grad(
             tape1, approx_order=1, strategy="forward", num_directions=n1, sampler_rng=rng
         )
-        j1 = fn(dev.batch_execute(tapes))
+        with qml.Tracker(dev) as tracker:
+            j1 = fn(dev.execute(tapes))
 
-        assert len(tapes) == dev.num_executions == n1 + 1
+        assert len(tapes) == tracker.latest["executions"] == n1 + 1
 
         n2 = 11
         tapes, fn = spsa_grad(tape2, num_directions=n2, sampler_rng=rng)
-        j2 = fn(dev.batch_execute(tapes))
+        j2 = fn(dev.execute(tapes))
 
         assert len(tapes) == 2 * n2
 
@@ -657,7 +658,7 @@ def test_ragged_output(self, approx_order, strategy, validate):
             num_directions=11,
             validate_params=validate,
         )
-        res = fn(dev.batch_execute(tapes))
+        res = fn(dev.execute(tapes))
 
         assert isinstance(res, tuple)
 
@@ -696,7 +697,7 @@ def test_single_expectation_value(self, approx_order, strategy, validate, tol):
             sampler=coordinate_sampler,
             validate_params=validate,
         )
-        res = fn(dev.batch_execute(tapes))
+        res = fn(dev.execute(tapes))
 
         assert isinstance(res, tuple)
         assert len(res) == 2
@@ -740,7 +741,7 @@ def test_single_expectation_value_with_argnum_all(self, approx_order, strategy,
             sampler=coordinate_sampler,
             validate_params=validate,
         )
-        res = fn(dev.batch_execute(tapes))
+        res = fn(dev.execute(tapes))
 
         assert isinstance(res, tuple)
         assert len(res) == 2
@@ -789,7 +790,7 @@ def test_single_expectation_value_with_argnum_one(self, approx_order, strategy,
             sampler=coordinate_sampler,
             validate_params=validate,
         )
-        res = fn(dev.batch_execute(tapes))
+        res = fn(dev.execute(tapes))
 
         assert isinstance(res, tuple)
         assert len(res) == 2
@@ -834,7 +835,7 @@ def test_multiple_expectation_value_with_argnum_one(self, approx_order, strategy
             sampler=coordinate_sampler,
             validate_params=validate,
         )
-        res = fn(dev.batch_execute(tapes))
+        res = fn(dev.execute(tapes))
 
         assert isinstance(res, tuple)
         assert isinstance(res[0], tuple)
@@ -865,7 +866,7 @@ def test_multiple_expectation_values(self, approx_order, strategy, validate, tol
             sampler=coordinate_sampler,
             validate_params=validate,
         )
-        res = fn(dev.batch_execute(tapes))
+        res = fn(dev.execute(tapes))
 
         assert isinstance(res, tuple)
         assert len(res) == 2
@@ -910,7 +911,7 @@ def test_var_expectation_values(self, approx_order, strategy, validate, tol):
             sampler=coordinate_sampler,
             num_directions=2,
         )
-        res = fn(dev.batch_execute(tapes))
+        res = fn(dev.execute(tapes))
 
         assert isinstance(res, tuple)
         assert len(res) == 2
@@ -955,7 +956,7 @@ def test_prob_expectation_values(self, approx_order, strategy, validate, tol):
             sampler=coordinate_sampler,
             num_directions=2,
         )
-        res = fn(dev.batch_execute(tapes))
+        res = fn(dev.execute(tapes))
 
         assert isinstance(res, tuple)
         assert len(res) == 2
@@ -1007,10 +1008,12 @@ class TestSpsaGradientDifferentiation:
     """Test that the transform is differentiable"""
 
     @pytest.mark.autograd
-    def test_autograd(self, sampler, num_directions, atol):
+    @pytest.mark.parametrize("dev_name", ["default.qubit", "default.qubit.autograd"])
+    def test_autograd(self, dev_name, sampler, num_directions, atol):
         """Tests that the output of the SPSA gradient transform
         can be differentiated using autograd, yielding second derivatives."""
-        dev = qml.device("default.qubit.autograd", wires=2)
+        dev = qml.device(dev_name, wires=2)
+        execute_fn = dev.execute if dev_name == "default.qubit" else dev.batch_execute
         params = np.array([0.543, -0.654], requires_grad=True)
         rng = np.random.default_rng(42)
 
@@ -1026,7 +1029,7 @@ def cost_fn(x):
             tapes, fn = spsa_grad(
                 tape, n=1, num_directions=num_directions, sampler=sampler, sampler_rng=rng
             )
-            jac = np.array(fn(dev.batch_execute(tapes)))
+            jac = np.array(fn(execute_fn(tapes)))
             if sampler is coordinate_sampler:
                 jac *= 2
             return jac
@@ -1043,10 +1046,12 @@ def cost_fn(x):
         assert np.allclose(res, expected, atol=atol, rtol=0)
 
     @pytest.mark.autograd
-    def test_autograd_ragged(self, sampler, num_directions, atol):
+    @pytest.mark.parametrize("dev_name", ["default.qubit", "default.qubit.autograd"])
+    def test_autograd_ragged(self, dev_name, sampler, num_directions, atol):
         """Tests that the output of the SPSA gradient transform
         of a ragged tape can be differentiated using autograd, yielding second derivatives."""
-        dev = qml.device("default.qubit.autograd", wires=2)
+        dev = qml.device(dev_name, wires=2)
+        execute_fn = dev.execute if dev_name == "default.qubit" else dev.batch_execute
         params = np.array([0.543, -0.654], requires_grad=True)
         rng = np.random.default_rng(42)
 
@@ -1063,7 +1068,7 @@ def cost_fn(x):
             tapes, fn = spsa_grad(
                 tape, n=1, num_directions=num_directions, sampler=sampler, sampler_rng=rng
             )
-            jac = fn(dev.batch_execute(tapes))
+            jac = fn(execute_fn(tapes))
             if sampler is coordinate_sampler:
                 jac = tuple(tuple(2 * _j for _j in _jac) for _jac in jac)
             return jac[1][0]
@@ -1075,12 +1080,14 @@ def cost_fn(x):
 
     @pytest.mark.tf
     @pytest.mark.slow
-    def test_tf(self, sampler, num_directions, atol):
+    @pytest.mark.parametrize("dev_name", ["default.qubit", "default.qubit.tf"])
+    def test_tf(self, dev_name, sampler, num_directions, atol):
         """Tests that the output of the SPSA gradient transform
         can be differentiated using TF, yielding second derivatives."""
         import tensorflow as tf
 
-        dev = qml.device("default.qubit.tf", wires=2)
+        dev = qml.device(dev_name, wires=2)
+        execute_fn = dev.execute if dev_name == "default.qubit" else dev.batch_execute
         params = tf.Variable([0.543, -0.654], dtype=tf.float64)
         rng = np.random.default_rng(42)
 
@@ -1096,7 +1103,7 @@ def test_tf(self, sampler, num_directions, atol):
             tapes, fn = spsa_grad(
                 tape, n=1, num_directions=num_directions, sampler=sampler, sampler_rng=rng
             )
-            jac_0, jac_1 = fn(dev.batch_execute(tapes))
+            jac_0, jac_1 = fn(execute_fn(tapes))
             if sampler is coordinate_sampler:
                 jac_0 *= 2
                 jac_1 *= 2
@@ -1116,12 +1123,14 @@ def test_tf(self, sampler, num_directions, atol):
 
     @pytest.mark.tf
     @pytest.mark.slow
-    def test_tf_ragged(self, sampler, num_directions, atol):
+    @pytest.mark.parametrize("dev_name", ["default.qubit", "default.qubit.tf"])
+    def test_tf_ragged(self, dev_name, sampler, num_directions, atol):
         """Tests that the output of the SPSA gradient transform
         of a ragged tape can be differentiated using TF, yielding second derivatives."""
         import tensorflow as tf
 
-        dev = qml.device("default.qubit.tf", wires=2)
+        dev = qml.device(dev_name, wires=2)
+        execute_fn = dev.execute if dev_name == "default.qubit" else dev.batch_execute
         params = tf.Variable([0.543, -0.654], dtype=tf.float64)
         rng = np.random.default_rng(42)
 
@@ -1139,7 +1148,7 @@ def test_tf_ragged(self, sampler, num_directions, atol):
                 tape, n=1, num_directions=num_directions, sampler=sampler, sampler_rng=rng
             )
 
-            jac_01 = fn(dev.batch_execute(tapes))[1][0]
+            jac_01 = fn(execute_fn(tapes))[1][0]
             if sampler is coordinate_sampler:
                 jac_01 *= 2
 
@@ -1152,12 +1161,14 @@ def test_tf_ragged(self, sampler, num_directions, atol):
         assert np.allclose(res_01[0], expected, atol=atol, rtol=0)
 
     @pytest.mark.torch
-    def test_torch(self, sampler, num_directions, atol):
+    @pytest.mark.parametrize("dev_name", ["default.qubit", "default.qubit.torch"])
+    def test_torch(self, dev_name, sampler, num_directions, atol):
         """Tests that the output of the SPSA gradient transform
         can be differentiated using Torch, yielding second derivatives."""
         import torch
 
-        dev = qml.device("default.qubit.torch", wires=2)
+        dev = qml.device(dev_name, wires=2)
+        execute_fn = dev.execute if dev_name == "default.qubit" else dev.batch_execute
         params = torch.tensor([0.543, -0.654], dtype=torch.float64, requires_grad=True)
         rng = np.random.default_rng(42)
 
@@ -1172,7 +1183,7 @@ def cost_fn(params):
             tapes, fn = spsa_grad(
                 tape, n=1, num_directions=num_directions, sampler=sampler, sampler_rng=rng
             )
-            jac = fn(dev.batch_execute(tapes))
+            jac = fn(execute_fn(tapes))
             if sampler is coordinate_sampler:
                 jac = tuple(2 * _jac for _jac in jac)
             return jac
@@ -1192,7 +1203,8 @@ def cost_fn(params):
         assert np.allclose(hess[1].detach().numpy(), expected[1], atol=atol, rtol=0)
 
     @pytest.mark.jax
-    def test_jax(self, sampler, num_directions, atol):
+    @pytest.mark.parametrize("dev_name", ["default.qubit", "default.qubit.jax"])
+    def test_jax(self, dev_name, sampler, num_directions, atol):
         """Tests that the output of the SPSA gradient transform
         can be differentiated using JAX, yielding second derivatives."""
         import jax
@@ -1201,7 +1213,8 @@ def test_jax(self, sampler, num_directions, atol):
 
         config.update("jax_enable_x64", True)
 
-        dev = qml.device("default.qubit.jax", wires=2)
+        dev = qml.device(dev_name, wires=2)
+        execute_fn = dev.execute if dev_name == "default.qubit" else dev.batch_execute
         params = jnp.array([0.543, -0.654])
         rng = np.random.default_rng(42)
 
@@ -1217,7 +1230,7 @@ def cost_fn(x):
             tapes, fn = spsa_grad(
                 tape, n=1, num_directions=num_directions, sampler=sampler, sampler_rng=rng
             )
-            jac = fn(dev.batch_execute(tapes))
+            jac = fn(execute_fn(tapes))
             if sampler is coordinate_sampler:
                 jac = tuple(2 * _jac for _jac in jac)
             return jac
diff --git a/tests/gradients/finite_diff/test_spsa_gradient_shot_vec.py b/tests/gradients/finite_diff/test_spsa_gradient_shot_vec.py
index 3d8fa8017df..f04dcd7ca25 100644
--- a/tests/gradients/finite_diff/test_spsa_gradient_shot_vec.py
+++ b/tests/gradients/finite_diff/test_spsa_gradient_shot_vec.py
@@ -71,7 +71,7 @@ def test_non_differentiable_error(self):
         dev = qml.device("default.qubit", wires=2, shots=default_shot_vector)
         tapes, fn = spsa_grad(tape, h=h_val)
 
-        all_res = fn(dev.batch_execute(tapes))
+        all_res = fn(dev.execute(tapes))
 
         assert isinstance(all_res, tuple)
         assert len(all_res) == len(default_shot_vector)
@@ -99,7 +99,7 @@ def test_independent_parameter(self, num_directions, mocker):
         tape = qml.tape.QuantumScript.from_queue(q, shots=default_shot_vector)
         dev = qml.device("default.qubit", wires=2, shots=default_shot_vector)
         tapes, fn = spsa_grad(tape, h=h_val, num_directions=num_directions)
-        all_res = fn(dev.batch_execute(tapes))
+        all_res = fn(dev.execute(tapes))
 
         assert isinstance(all_res, tuple)
         assert len(all_res) == len(default_shot_vector)
@@ -409,9 +409,10 @@ def test_independent_parameters(self):
             h=h_val,
             sampler_rng=rng,
         )
-        j1 = fn(dev.batch_execute(tapes))
+        with qml.Tracker(dev) as tracker:
+            j1 = fn(dev.execute(tapes))
 
-        assert len(tapes) == dev.num_executions == n1 + 1
+        assert len(tapes) == tracker.latest["executions"] == n1 + 1
 
         n2 = 3
         tapes, fn = spsa_grad(
@@ -422,7 +423,7 @@ def test_independent_parameters(self):
             num_directions=n2,
             sampler_rng=rng,
         )
-        j2 = fn(dev.batch_execute(tapes))
+        j2 = fn(dev.execute(tapes))
 
         assert len(tapes) == n2 + 1
 
@@ -584,7 +585,7 @@ def test_ragged_output(self, approx_order, strategy, validate):
             num_directions=3,
             sampler_rng=rng,
         )
-        all_res = fn(dev.batch_execute(tapes))
+        all_res = fn(dev.execute(tapes))
         assert isinstance(all_res, tuple)
         assert len(all_res) == len(many_shots_shot_vector)
 
@@ -628,7 +629,7 @@ def test_single_expectation_value(self, approx_order, strategy, validate):
             sampler=coordinate_sampler,
             sampler_rng=rng,
         )
-        all_res = fn(dev.batch_execute(tapes))
+        all_res = fn(dev.execute(tapes))
 
         assert isinstance(all_res, tuple)
         assert len(all_res) == len(default_shot_vector)
@@ -677,7 +678,7 @@ def test_single_expectation_value_with_argnum_all(self, approx_order, strategy,
             h=h_val,
             sampler_rng=rng,
         )
-        all_res = fn(dev.batch_execute(tapes))
+        all_res = fn(dev.execute(tapes))
 
         assert isinstance(all_res, tuple)
         assert len(all_res) == len(default_shot_vector)
@@ -730,7 +731,7 @@ def test_single_expectation_value_with_argnum_one(self, approx_order, strategy,
             h=h_val,
             sampler_rng=rng,
         )
-        all_res = fn(dev.batch_execute(tapes))
+        all_res = fn(dev.execute(tapes))
 
         assert isinstance(all_res, tuple)
         assert len(all_res) == len(default_shot_vector)
@@ -784,7 +785,7 @@ def test_multiple_expectation_value_with_argnum_one(self, approx_order, strategy
             h=h_val,
             sampler_rng=rng,
         )
-        all_res = fn(dev.batch_execute(tapes))
+        all_res = fn(dev.execute(tapes))
 
         assert isinstance(all_res, tuple)
         assert len(all_res) == len(default_shot_vector)
@@ -824,7 +825,7 @@ def test_multiple_expectation_values(self, approx_order, strategy, validate):
             num_directions=20,
             sampler_rng=rng,
         )
-        all_res = fn(dev.batch_execute(tapes))
+        all_res = fn(dev.execute(tapes))
 
         assert isinstance(all_res, tuple)
         assert len(all_res) == len(default_shot_vector)
@@ -878,7 +879,7 @@ def test_var_expectation_values(self, approx_order, strategy, validate):
             sampler=coordinate_sampler,
             sampler_rng=rng,
         )
-        all_res = fn(dev.batch_execute(tapes))
+        all_res = fn(dev.execute(tapes))
 
         assert isinstance(all_res, tuple)
         assert len(all_res) == len(default_shot_vector)
@@ -933,7 +934,7 @@ def test_prob_expectation_values(self, approx_order, strategy, validate):
             h=h_val,
             sampler_rng=rng,
         )
-        all_res = fn(dev.batch_execute(tapes))
+        all_res = fn(dev.execute(tapes))
 
         assert isinstance(all_res, tuple)
         assert len(all_res) == len(default_shot_vector)
@@ -991,10 +992,12 @@ class TestSpsaGradientDifferentiation:
     """Test that the transform is differentiable"""
 
     @pytest.mark.autograd
-    def test_autograd(self, approx_order, strategy):
+    @pytest.mark.parametrize("dev_name", ["default.qubit", "default.qubit.autograd"])
+    def test_autograd(self, dev_name, approx_order, strategy):
         """Tests that the output of the SPSA gradient transform
         can be differentiated using autograd, yielding second derivatives."""
-        dev = qml.device("default.qubit.autograd", wires=2, shots=many_shots_shot_vector)
+        dev = qml.device(dev_name, wires=2, shots=many_shots_shot_vector)
+        execute_fn = dev.execute if dev_name == "default.qubit" else dev.batch_execute
         params = np.array([0.543, -0.654], requires_grad=True)
         rng = np.random.default_rng(42)
 
@@ -1015,7 +1018,7 @@ def cost_fn(x):
                 h=h_val,
                 sampler_rng=rng,
             )
-            jac = np.array(fn(dev.batch_execute(tapes)))
+            jac = np.array(fn(execute_fn(tapes)))
             return jac
 
         all_res = qml.jacobian(cost_fn)(params)
@@ -1035,10 +1038,12 @@ def cost_fn(x):
             assert np.allclose(res, expected, atol=spsa_shot_vec_tol, rtol=0)
 
     @pytest.mark.autograd
-    def test_autograd_ragged(self, approx_order, strategy):
+    @pytest.mark.parametrize("dev_name", ["default.qubit", "default.qubit.autograd"])
+    def test_autograd_ragged(self, dev_name, approx_order, strategy):
         """Tests that the output of the SPSA gradient transform
         of a ragged tape can be differentiated using autograd, yielding second derivatives."""
-        dev = qml.device("default.qubit.autograd", wires=2, shots=many_shots_shot_vector)
+        dev = qml.device(dev_name, wires=2, shots=many_shots_shot_vector)
+        execute_fn = dev.execute if dev_name == "default.qubit" else dev.batch_execute
         params = np.array([0.543, -0.654], requires_grad=True)
         rng = np.random.default_rng(42)
 
@@ -1060,7 +1065,7 @@ def cost_fn(x):
                 h=h_val,
                 sampler_rng=rng,
             )
-            jac = fn(dev.batch_execute(tapes))
+            jac = fn(execute_fn(tapes))
             return jac[1][0]
 
         x, y = params
@@ -1075,12 +1080,14 @@ def cost_fn(x):
 
     @pytest.mark.tf
     @pytest.mark.slow
-    def test_tf(self, approx_order, strategy):
+    @pytest.mark.parametrize("dev_name", ["default.qubit", "default.qubit.tf"])
+    def test_tf(self, dev_name, approx_order, strategy):
         """Tests that the output of the SPSA gradient transform
         can be differentiated using TF, yielding second derivatives."""
         import tensorflow as tf
 
-        dev = qml.device("default.qubit.tf", wires=2, shots=many_shots_shot_vector)
+        dev = qml.device(dev_name, wires=2, shots=many_shots_shot_vector)
+        execute_fn = dev.execute if dev_name == "default.qubit" else dev.batch_execute
         params = tf.Variable([0.543, -0.654], dtype=tf.float64)
         rng = np.random.default_rng(42)
 
@@ -1101,7 +1108,7 @@ def test_tf(self, approx_order, strategy):
                 h=h_val,
                 sampler_rng=rng,
             )
-            jac_0, jac_1 = fn(dev.batch_execute(tapes))
+            jac_0, jac_1 = fn(execute_fn(tapes))
 
         x, y = 1.0 * params
 
@@ -1117,12 +1124,14 @@ def test_tf(self, approx_order, strategy):
         assert np.allclose([res_0, res_1], expected, atol=spsa_shot_vec_tol, rtol=0)
 
     @pytest.mark.tf
-    def test_tf_ragged(self, approx_order, strategy):
+    @pytest.mark.parametrize("dev_name", ["default.qubit", "default.qubit.tf"])
+    def test_tf_ragged(self, dev_name, approx_order, strategy):
         """Tests that the output of the SPSA gradient transform
         of a ragged tape can be differentiated using TF, yielding second derivatives."""
         import tensorflow as tf
 
-        dev = qml.device("default.qubit.tf", wires=2, shots=many_shots_shot_vector)
+        dev = qml.device(dev_name, wires=2, shots=many_shots_shot_vector)
+        execute_fn = dev.execute if dev_name == "default.qubit" else dev.batch_execute
         params = tf.Variable([0.543, -0.654], dtype=tf.float64)
         rng = np.random.default_rng(42)
 
@@ -1145,7 +1154,7 @@ def test_tf_ragged(self, approx_order, strategy):
                 sampler_rng=rng,
             )
 
-            jac_01 = fn(dev.batch_execute(tapes))[1][0]
+            jac_01 = fn(execute_fn(tapes))[1][0]
 
         x, y = 1.0 * params
 
@@ -1156,12 +1165,14 @@ def test_tf_ragged(self, approx_order, strategy):
         assert np.allclose(res_01[0], expected, atol=spsa_shot_vec_tol, rtol=0)
 
     @pytest.mark.torch
-    def test_torch(self, approx_order, strategy):
+    @pytest.mark.parametrize("dev_name", ["default.qubit", "default.qubit.torch"])
+    def test_torch(self, dev_name, approx_order, strategy):
         """Tests that the output of the SPSA gradient transform
         can be differentiated using Torch, yielding second derivatives."""
         import torch
 
-        dev = qml.device("default.qubit.torch", wires=2, shots=many_shots_shot_vector)
+        dev = qml.device(dev_name, wires=2, shots=many_shots_shot_vector)
+        execute_fn = dev.execute if dev_name == "default.qubit" else dev.batch_execute
         params = torch.tensor([0.543, -0.654], dtype=torch.float64, requires_grad=True)
         rng = np.random.default_rng(42)
 
@@ -1181,7 +1192,7 @@ def cost_fn(params):
                 h=h_val,
                 sampler_rng=rng,
             )
-            jac = fn(dev.batch_execute(tapes))
+            jac = fn(execute_fn(tapes))
             return jac
 
         hess = torch.autograd.functional.jacobian(cost_fn, params)
@@ -1199,7 +1210,8 @@ def cost_fn(params):
         assert np.allclose(hess[1].detach().numpy(), expected[1], atol=spsa_shot_vec_tol, rtol=0)
 
     @pytest.mark.jax
-    def test_jax(self, approx_order, strategy):
+    @pytest.mark.parametrize("dev_name", ["default.qubit", "default.qubit.jax"])
+    def test_jax(self, dev_name, approx_order, strategy):
         """Tests that the output of the SPSA gradient transform
         can be differentiated using JAX, yielding second derivatives."""
         import jax
@@ -1208,7 +1220,8 @@ def test_jax(self, approx_order, strategy):
 
         config.update("jax_enable_x64", True)
 
-        dev = qml.device("default.qubit.jax", wires=2, shots=many_shots_shot_vector)
+        dev = qml.device(dev_name, wires=2, shots=many_shots_shot_vector)
+        execute_fn = dev.execute if dev_name == "default.qubit" else dev.batch_execute
         params = jnp.array([0.543, -0.654])
         rng = np.random.default_rng(42)
 
@@ -1229,7 +1242,7 @@ def cost_fn(x):
                 h=h_val,
                 sampler_rng=rng,
             )
-            jac = fn(dev.batch_execute(tapes))
+            jac = fn(execute_fn(tapes))
             return jac
 
         all_res = jax.jacobian(cost_fn)(params)
@@ -1302,7 +1315,7 @@ def test_1_1(self, shot_vec, meas, shape, op_wires):
         tape.trainable_params = {0}
 
         tapes, fn = spsa_grad(tape)
-        all_res = fn(dev.batch_execute(tapes))
+        all_res = fn(dev.execute(tapes))
 
         assert len(all_res) == grad_transform_shots.num_copies
         assert isinstance(all_res, tuple)
@@ -1336,7 +1349,7 @@ def test_1_N(self, shot_vec, op_wire):
         tape.trainable_params = {0}
 
         tapes, fn = spsa_grad(tape)
-        all_res = fn(dev.batch_execute(tapes))
+        all_res = fn(dev.execute(tapes))
 
         assert len(all_res) == grad_transform_shots.num_copies
         assert isinstance(all_res, tuple)
@@ -1370,7 +1383,7 @@ def test_N_1(self, shot_vec, meas, shape, op_wires):
         tape.trainable_params = {0, 1}
 
         tapes, fn = spsa_grad(tape)
-        all_res = fn(dev.batch_execute(tapes))
+        all_res = fn(dev.execute(tapes))
 
         assert len(all_res) == grad_transform_shots.num_copies
         assert isinstance(all_res, tuple)
@@ -1409,7 +1422,7 @@ def test_N_N(self, shot_vec, op_wires):
         tape.trainable_params = {0, 1, 2, 3, 4}
 
         tapes, fn = spsa_grad(tape)
-        all_res = fn(dev.batch_execute(tapes))
+        all_res = fn(dev.execute(tapes))
 
         assert len(all_res) == grad_transform_shots.num_copies
         assert isinstance(all_res, tuple)
diff --git a/tests/gradients/parameter_shift/test_parameter_shift.py b/tests/gradients/parameter_shift/test_parameter_shift.py
index a8e0aaa1858..d483c065837 100644
--- a/tests/gradients/parameter_shift/test_parameter_shift.py
+++ b/tests/gradients/parameter_shift/test_parameter_shift.py
@@ -315,7 +315,7 @@ def test_multi_measure_partitioned_shots_par_shapes(self, g, par_shapes):
 def grad_fn(tape, dev, fn=qml.gradients.param_shift, **kwargs):
     """Utility function to automate execution and processing of gradient tapes"""
     tapes, fn = fn(tape, **kwargs)
-    return fn(dev.batch_execute(tapes))
+    return fn(dev.execute(tapes))
 
 
 class TestParamShift:
@@ -370,7 +370,7 @@ def test_independent_parameter(self, mocker):
         assert len(tapes) == 2
         assert tapes[0].batch_size == tapes[1].batch_size == None
 
-        res = fn(dev.batch_execute(tapes))
+        res = fn(dev.execute(tapes))
         assert isinstance(res, tuple)
         assert len(res) == 2
         assert res[0].shape == ()
@@ -761,7 +761,7 @@ def test_f0_provided(self, y_wire):
         # one tape per parameter that impacts the expval
         assert len(tapes) == 2 if y_wire == 0 else 1
 
-        fn(dev.batch_execute(tapes))
+        fn(dev.execute(tapes))
 
     def test_op_with_custom_unshifted_term(self):
         """Test that an operation with a gradient recipe that depends on
@@ -798,7 +798,7 @@ class RX(qml.RX):
             assert tape.operations[0].data[0] == x[0] + expected[0]
             assert tape.operations[1].data[0] == x[1] + expected[1]
 
-        grad = fn(dev.batch_execute(tapes))
+        grad = fn(dev.execute(tapes))
         exp = np.stack([-np.sin(x[0] + x[1]), -np.sin(x[0] + x[1]) + 0.2 * np.cos(x[0] + x[1])])
         assert len(grad) == len(exp)
         for (
@@ -825,15 +825,16 @@ def test_independent_parameters_analytic(self):
 
         tape2 = qml.tape.QuantumScript.from_queue(q2)
         tapes, fn = qml.gradients.param_shift(tape1)
-        j1 = fn(dev.batch_execute(tapes))
+        with qml.Tracker(dev) as tracker:
+            j1 = fn(dev.execute(tapes))
 
         # We should only be executing the device twice: Two shifted evaluations to differentiate
         # one parameter overall, as the other parameter does not impact the returned measurement.
 
-        assert dev.num_executions == 2
+        assert tracker.latest["executions"] == 2
 
         tapes, fn = qml.gradients.param_shift(tape2)
-        j2 = fn(dev.batch_execute(tapes))
+        j2 = fn(dev.execute(tapes))
 
         exp = -np.sin(1)
 
@@ -863,7 +864,7 @@ def test_grad_recipe_parameter_dependent(self):
         assert qml.math.allclose(tapes[0].operations[0].data[0], 0)
         assert qml.math.allclose(tapes[1].operations[0].data[0], 2 * x)
 
-        grad = fn(dev.batch_execute(tapes))
+        grad = fn(dev.execute(tapes))
         assert np.allclose(grad, -np.sin(x))
 
     def test_error_no_diff_info(self):
@@ -946,7 +947,7 @@ def test_with_single_parameter_broadcasted(self, dim, pos):
         assert np.allclose([t.batch_size for t in tapes], dim)
 
         dev = qml.device("default.qubit", wires=2)
-        res = fn(dev.batch_execute(tapes))
+        res = fn(dev.execute(tapes))
         assert isinstance(res, tuple)
         assert len(res) == 2
 
@@ -974,7 +975,7 @@ def test_with_multiple_parameters_broadcasted(self, dim, argnum):
         assert np.allclose([t.batch_size for t in tapes], dim)
 
         dev = qml.device("default.qubit", wires=2)
-        res = fn(dev.batch_execute(tapes))
+        res = fn(dev.execute(tapes))
         assert isinstance(res, tuple)
         assert len(res) == 3
 
@@ -1001,7 +1002,7 @@ def test_independent_parameter(self, mocker):
         assert len(tapes) == 1
         assert tapes[0].batch_size == 2
 
-        res = fn(dev.batch_execute(tapes))
+        res = fn(dev.execute(tapes))
         assert len(res) == 2
         assert res[0].shape == ()
         assert res[1].shape == ()
@@ -1080,15 +1081,16 @@ def test_independent_parameters_analytic(self):
 
         tape2 = qml.tape.QuantumScript.from_queue(q2)
         tapes, fn = qml.gradients.param_shift(tape1, broadcast=True)
-        j1 = fn(dev.batch_execute(tapes))
+        with qml.Tracker(dev) as tracker:
+            j1 = fn(dev.execute(tapes))
 
         # We should only be executing the device to differentiate 1 parameter
         # (1 broadcasted execution)
 
-        assert dev.num_executions == 1
+        assert tracker.latest["executions"] == 1
 
         tapes, fn = qml.gradients.param_shift(tape2, broadcast=True)
-        j2 = fn(dev.batch_execute(tapes))
+        j2 = fn(dev.execute(tapes))
 
         exp = -np.sin(1)
 
@@ -1122,7 +1124,7 @@ def fail(*args, **kwargs):
         assert tapes[0].batch_size == 2
         assert qml.math.allclose(tapes[0].operations[0].data[0], [0, 2 * x])
 
-        grad = fn(dev.batch_execute(tapes))
+        grad = fn(dev.execute(tapes))
         assert np.allclose(grad, -np.sin(x))
 
 
@@ -1153,12 +1155,12 @@ def test_pauli_rotation_gradient(self, mocker, G, theta, shift, tol):
         tapes, fn = qml.gradients.param_shift(tape, shifts=[(shift,)])
         assert len(tapes) == 2
 
-        autograd_val = fn(dev.batch_execute(tapes))
+        autograd_val = fn(dev.execute(tapes))
 
         tape_fwd = tape.bind_new_parameters([theta + np.pi / 2], [1])
         tape_bwd = tape.bind_new_parameters([theta - np.pi / 2], [1])
 
-        manualgrad_val = np.subtract(*dev.batch_execute([tape_fwd, tape_bwd])) / 2
+        manualgrad_val = np.subtract(*dev.execute([tape_fwd, tape_bwd])) / 2
         assert np.allclose(autograd_val, manualgrad_val, atol=tol, rtol=0)
 
         assert isinstance(autograd_val, np.ndarray)
@@ -1167,7 +1169,7 @@ def test_pauli_rotation_gradient(self, mocker, G, theta, shift, tol):
         assert spy.call_args[1]["shifts"] == (shift,)
 
         tapes, fn = qml.gradients.finite_diff(tape)
-        numeric_val = fn(dev.batch_execute(tapes))
+        numeric_val = fn(dev.execute(tapes))
         assert np.allclose(autograd_val, numeric_val, atol=tol, rtol=0)
 
     @pytest.mark.parametrize("theta", np.linspace(-2 * np.pi, 2 * np.pi, 7))
@@ -1190,7 +1192,7 @@ def test_Rot_gradient(self, mocker, theta, shift, tol):
         num_params = len(tape.trainable_params)
         assert len(tapes) == 2 * num_params
 
-        autograd_val = fn(dev.batch_execute(tapes))
+        autograd_val = fn(dev.execute(tapes))
         assert isinstance(autograd_val, tuple)
         assert len(autograd_val) == num_params
 
@@ -1215,7 +1217,7 @@ def test_Rot_gradient(self, mocker, theta, shift, tol):
             assert spy.call_args[1]["shifts"] == (shift,)
 
         tapes, fn = qml.gradients.finite_diff(tape)
-        numeric_val = fn(dev.batch_execute(tapes))
+        numeric_val = fn(dev.execute(tapes))
         for a_val, n_val in zip(autograd_val, numeric_val):
             assert np.allclose(a_val, n_val, atol=tol, rtol=0)
 
@@ -1237,12 +1239,12 @@ def test_controlled_rotation_gradient(self, G, tol):
         assert np.allclose(res, -np.cos(b / 2), atol=tol, rtol=0)
 
         tapes, fn = qml.gradients.param_shift(tape)
-        grad = fn(dev.batch_execute(tapes))
+        grad = fn(dev.execute(tapes))
         expected = np.sin(b / 2) / 2
         assert np.allclose(grad, expected, atol=tol, rtol=0)
 
         tapes, fn = qml.gradients.finite_diff(tape)
-        numeric_val = fn(dev.batch_execute(tapes))
+        numeric_val = fn(dev.execute(tapes))
         assert np.allclose(grad, numeric_val, atol=tol, rtol=0)
 
     @pytest.mark.parametrize("theta", np.linspace(-2 * np.pi, np.pi, 7))
@@ -1267,7 +1269,7 @@ def test_CRot_gradient(self, theta, tol):
         tapes, fn = qml.gradients.param_shift(tape)
         assert len(tapes) == 4 * len(tape.trainable_params)
 
-        grad = fn(dev.batch_execute(tapes))
+        grad = fn(dev.execute(tapes))
         expected = np.array(
             [
                 0.5 * np.cos(b / 2) * np.sin(0.5 * (a + c)),
@@ -1281,7 +1283,7 @@ def test_CRot_gradient(self, theta, tol):
             assert np.allclose(g, expected[idx], atol=tol, rtol=0)
 
         tapes, fn = qml.gradients.finite_diff(tape)
-        numeric_val = fn(dev.batch_execute(tapes))
+        numeric_val = fn(dev.execute(tapes))
         for idx, g in enumerate(grad):
             assert np.allclose(g, numeric_val[idx], atol=tol, rtol=0)
 
@@ -1373,7 +1375,7 @@ def cost_fn(params):
             spy.assert_called()
             assert spy.call_args[1]["argnum"] == {1}
 
-            return fn(dev.batch_execute(tapes))
+            return fn(dev.execute(tapes))
 
         res = cost_fn(params)
 
@@ -1402,10 +1404,12 @@ def cost_fn(params):
         # assert np.allclose(jac[1, 1, 1], -2 * np.cos(2 * y), atol=tol, rtol=0)
 
     @pytest.mark.autograd
-    def test_fallback_single_meas(self, mocker):
+    @pytest.mark.parametrize("dev_name", ["default.qubit", "default.qubit.autograd"])
+    def test_fallback_single_meas(self, dev_name, mocker):
         """Test that fallback gradient functions are correctly used for a single measurement."""
         spy = mocker.spy(qml.gradients, "finite_diff")
-        dev = qml.device("default.qubit.autograd", wires=2)
+        dev = qml.device(dev_name, wires=2)
+        execute_fn = dev.execute if dev_name == "default.qubit" else dev.batch_execute
         x = 0.543
         y = -0.654
 
@@ -1425,7 +1429,7 @@ def cost_fn(params):
             spy.assert_called()
             assert spy.call_args[1]["argnum"] == {1}
 
-            return fn(dev.batch_execute(tapes))
+            return fn(execute_fn(tapes))
 
         res = cost_fn(params)
 
@@ -1442,10 +1446,12 @@ def cost_fn(params):
 
     @pytest.mark.autograd
     @pytest.mark.parametrize("RX, RY, argnum", [(RX_with_F, qml.RY, 0), (qml.RX, RY_with_F, 1)])
-    def test_fallback_probs(self, RX, RY, argnum, mocker):
+    @pytest.mark.parametrize("dev_name", ["default.qubit", "default.qubit.autograd"])
+    def test_fallback_probs(self, dev_name, RX, RY, argnum, mocker):
         """Test that fallback gradient functions are correctly used with probs"""
         spy = mocker.spy(qml.gradients, "finite_diff")
-        dev = qml.device("default.qubit.autograd", wires=2)
+        dev = qml.device(dev_name, wires=2)
+        execute_fn = dev.execute if dev_name == "default.qubit" else dev.batch_execute
         x = 0.543
         y = -0.654
 
@@ -1467,7 +1473,7 @@ def cost_fn(params):
             spy.assert_called()
             assert spy.call_args[1]["argnum"] == {argnum}
 
-            return fn(dev.batch_execute(tapes))
+            return fn(execute_fn(tapes))
 
         res = cost_fn(params)
 
@@ -1524,13 +1530,15 @@ def cost_fn(params):
         assert np.allclose(res[1][1], probs_expected[:, 1])
 
     @pytest.mark.autograd
-    def test_all_fallback(self, mocker, tol):
+    @pytest.mark.parametrize("dev_name", ["default.qubit", "default.qubit.autograd"])
+    def test_all_fallback(self, dev_name, mocker, tol):
         """Test that *only* the fallback logic is called if no parameters
         support the parameter-shift rule"""
         spy_fd = mocker.spy(qml.gradients, "finite_diff")
         spy_ps = mocker.spy(qml.gradients.parameter_shift, "expval_param_shift")
 
-        dev = qml.device("default.qubit.autograd", wires=2)
+        dev = qml.device(dev_name, wires=2)
+        execute_fn = dev.execute if dev_name == "default.qubit" else dev.batch_execute
         x = 0.543
         y = -0.654
 
@@ -1548,7 +1556,7 @@ def test_all_fallback(self, mocker, tol):
         spy_fd.assert_called()
         spy_ps.assert_not_called()
 
-        res = fn(dev.batch_execute(tapes))
+        res = fn(execute_fn(tapes))
 
         assert isinstance(res, tuple)
         assert res[0].shape == ()
@@ -1574,7 +1582,7 @@ def test_single_expectation_value(self, tol):
         tapes, fn = qml.gradients.param_shift(tape)
         assert len(tapes) == 4
 
-        res = fn(dev.batch_execute(tapes))
+        res = fn(dev.execute(tapes))
         assert len(res) == 2
         assert not isinstance(res[0], tuple)
         assert not isinstance(res[1], tuple)
@@ -1619,7 +1627,7 @@ def test_multiple_expectation_values(self, tol):
         tapes, fn = qml.gradients.param_shift(tape)
         assert len(tapes) == 4
 
-        res = fn(dev.batch_execute(tapes))
+        res = fn(dev.execute(tapes))
         assert len(res) == 2
         assert len(res[0]) == 2
         assert len(res[1]) == 2
@@ -1646,7 +1654,7 @@ def test_var_expectation_values(self, tol):
         tapes, fn = qml.gradients.param_shift(tape)
         assert len(tapes) == 5
 
-        res = fn(dev.batch_execute(tapes))
+        res = fn(dev.execute(tapes))
         assert len(res) == 2
         assert len(res[0]) == 2
         assert len(res[1]) == 2
@@ -1675,7 +1683,7 @@ def test_prob_expectation_values(self):
         tapes, fn = qml.gradients.param_shift(tape)
         assert len(tapes) == 4
 
-        res = fn(dev.batch_execute(tapes))
+        res = fn(dev.execute(tapes))
         assert len(res) == 2
 
         for r in res:
@@ -1732,13 +1740,13 @@ def test_involutory_variance_single_param(self, tol):
 
         # circuit jacobians
         tapes, fn = qml.gradients.param_shift(tape)
-        gradA = fn(dev.batch_execute(tapes))
+        gradA = fn(dev.execute(tapes))
         assert isinstance(gradA, np.ndarray)
         assert gradA.shape == ()
         assert len(tapes) == 1 + 2 * 1
 
         tapes, fn = qml.gradients.finite_diff(tape)
-        gradF = fn(dev.batch_execute(tapes))
+        gradF = fn(dev.execute(tapes))
         assert len(tapes) == 2
 
         expected = 2 * np.sin(a) * np.cos(a)
@@ -1765,7 +1773,7 @@ def test_involutory_variance_multi_param(self, tol):
 
         # circuit jacobians
         tapes, fn = qml.gradients.param_shift(tape)
-        gradA = fn(dev.batch_execute(tapes))
+        gradA = fn(dev.execute(tapes))
         assert isinstance(gradA, tuple)
 
         assert isinstance(gradA[0], np.ndarray)
@@ -1777,7 +1785,7 @@ def test_involutory_variance_multi_param(self, tol):
         assert len(tapes) == 1 + 2 * 2
 
         tapes, fn = qml.gradients.finite_diff(tape)
-        gradF = fn(dev.batch_execute(tapes))
+        gradF = fn(dev.execute(tapes))
         assert len(tapes) == 3
 
         expected = 2 * np.sin(a + b) * np.cos(a + b)
@@ -1806,13 +1814,13 @@ def test_non_involutory_variance_single_param(self, tol):
 
         # circuit jacobians
         tapes, fn = qml.gradients.param_shift(tape)
-        gradA = fn(dev.batch_execute(tapes))
+        gradA = fn(dev.execute(tapes))
         assert isinstance(gradA, np.ndarray)
         assert gradA.shape == ()
         assert len(tapes) == 1 + 4 * 1
 
         tapes, fn = qml.gradients.finite_diff(tape)
-        gradF = fn(dev.batch_execute(tapes))
+        gradF = fn(dev.execute(tapes))
         assert len(tapes) == 2
 
         expected = -35 * np.sin(2 * a) - 12 * np.cos(2 * a)
@@ -1840,7 +1848,7 @@ def test_non_involutory_variance_multi_param(self, tol):
 
         # circuit jacobians
         tapes, fn = qml.gradients.param_shift(tape)
-        gradA = fn(dev.batch_execute(tapes))
+        gradA = fn(dev.execute(tapes))
         assert isinstance(gradA, tuple)
 
         assert isinstance(gradA[0], np.ndarray)
@@ -1851,7 +1859,7 @@ def test_non_involutory_variance_multi_param(self, tol):
         assert len(tapes) == 1 + 4 * 2
 
         tapes, fn = qml.gradients.finite_diff(tape)
-        gradF = fn(dev.batch_execute(tapes))
+        gradF = fn(dev.execute(tapes))
         assert len(tapes) == 3
 
         expected = -35 * np.sin(2 * (a + b)) - 12 * np.cos(2 * (a + b))
@@ -1884,11 +1892,11 @@ def test_involutory_and_noninvolutory_variance_single_param(self, tol):
 
         # circuit jacobians
         tapes, fn = qml.gradients.param_shift(tape)
-        gradA = fn(dev.batch_execute(tapes))
+        gradA = fn(dev.execute(tapes))
         assert len(tapes) == 1 + 4
 
         tapes, fn = qml.gradients.finite_diff(tape)
-        gradF = fn(dev.batch_execute(tapes))
+        gradF = fn(dev.execute(tapes))
         assert len(tapes) == 1 + 1
 
         expected = [2 * np.sin(a) * np.cos(a), 0]
@@ -1937,7 +1945,7 @@ def test_var_and_probs_single_param(self, ind):
         # circuit jacobians
         tapes, fn = qml.gradients.param_shift(tape)
 
-        gradA = fn(dev.batch_execute(tapes))
+        gradA = fn(dev.execute(tapes))
 
         assert isinstance(gradA, tuple)
         assert len(gradA) == 3
@@ -1987,7 +1995,7 @@ def test_var_and_probs_multi_params(self):
 
         # circuit jacobians
         tapes, fn = qml.gradients.param_shift(tape)
-        gradA = fn(dev.batch_execute(tapes))
+        gradA = fn(dev.execute(tapes))
 
         assert isinstance(gradA, tuple)
         assert len(gradA) == 3
@@ -2117,10 +2125,10 @@ def test_expval_and_variance_single_param(self, tol):
 
         # # circuit jacobians
         tapes, fn = qml.gradients.param_shift(tape)
-        gradA = fn(dev.batch_execute(tapes))
+        gradA = fn(dev.execute(tapes))
 
         tapes, fn = qml.gradients.finite_diff(tape)
-        gradF = fn(dev.batch_execute(tapes))
+        gradF = fn(dev.execute(tapes))
 
         expected = np.array([2 * np.cos(a) * np.sin(a), -np.cos(b) * np.sin(a), 0])
         assert isinstance(gradA, tuple)
@@ -2164,10 +2172,10 @@ def test_expval_and_variance_multi_param(self, tol):
 
         # # circuit jacobians
         tapes, fn = qml.gradients.param_shift(tape)
-        gradA = fn(dev.batch_execute(tapes))
+        gradA = fn(dev.execute(tapes))
 
         tapes, fn = qml.gradients.finite_diff(tape)
-        gradF = fn(dev.batch_execute(tapes))
+        gradF = fn(dev.execute(tapes))
 
         expected = np.array(
             [
@@ -2238,10 +2246,10 @@ def test_projector_variance(self, state, tol):
 
         # # circuit jacobians
         tapes, fn = qml.gradients.param_shift(tape)
-        gradA = fn(dev.batch_execute(tapes))
+        gradA = fn(dev.execute(tapes))
 
         tapes, fn = qml.gradients.finite_diff(tape)
-        gradF = fn(dev.batch_execute(tapes))
+        gradF = fn(dev.execute(tapes))
 
         expected = np.array(
             [
@@ -2432,7 +2440,7 @@ def test_multi_measure_no_warning(self):
         tape = qml.tape.QuantumScript.from_queue(q)
         with warnings.catch_warnings(record=True) as record:
             tapes, fn = qml.gradients.param_shift(tape)
-            fn(dev.batch_execute(tapes))
+            fn(dev.execute(tapes))
 
         assert len(record) == 0
 
@@ -2463,18 +2471,18 @@ def test_pauli_rotation_gradient(self, mocker, G, theta, shift, tol):
         assert len(tapes) == 1
         assert tapes[0].batch_size == 2
 
-        autograd_val = fn(dev.batch_execute(tapes))
+        autograd_val = fn(dev.execute(tapes))
 
         tape_fwd = tape.bind_new_parameters([theta + np.pi / 2], [1])
         tape_bwd = tape.bind_new_parameters([theta - np.pi / 2], [1])
 
-        manualgrad_val = np.subtract(*dev.batch_execute([tape_fwd, tape_bwd])) / 2
+        manualgrad_val = np.subtract(*dev.execute([tape_fwd, tape_bwd])) / 2
         assert np.allclose(autograd_val, manualgrad_val, atol=tol, rtol=0)
 
         assert spy.call_args[1]["shifts"] == (shift,)
 
         tapes, fn = qml.gradients.finite_diff(tape)
-        numeric_val = fn(dev.batch_execute(tapes))
+        numeric_val = fn(dev.execute(tapes))
         assert np.allclose(autograd_val, numeric_val, atol=tol, rtol=0)
 
     @pytest.mark.parametrize("theta", np.linspace(-2 * np.pi, 2 * np.pi, 7))
@@ -2497,7 +2505,7 @@ def test_Rot_gradient(self, mocker, theta, shift, tol):
         assert len(tapes) == len(tape.trainable_params)
         assert [t.batch_size for t in tapes] == [2, 2, 2]
 
-        autograd_val = fn(dev.batch_execute(tapes))
+        autograd_val = fn(dev.execute(tapes))
         manualgrad_val = np.zeros_like(autograd_val)
 
         for idx in list(np.ndindex(*params.shape)):
@@ -2516,7 +2524,7 @@ def test_Rot_gradient(self, mocker, theta, shift, tol):
         assert spy.call_args[1]["shifts"] == (shift,)
 
         tapes, fn = qml.gradients.finite_diff(tape)
-        numeric_val = fn(dev.batch_execute(tapes))
+        numeric_val = fn(dev.execute(tapes))
 
         assert len(autograd_val) == len(numeric_val)
         for a, n in zip(autograd_val, numeric_val):
@@ -2540,12 +2548,12 @@ def test_controlled_rotation_gradient(self, G, tol):
         assert np.allclose(res, -np.cos(b / 2), atol=tol, rtol=0)
 
         tapes, fn = qml.gradients.param_shift(tape, broadcast=True)
-        grad = fn(dev.batch_execute(tapes))
+        grad = fn(dev.execute(tapes))
         expected = np.sin(b / 2) / 2
         assert np.allclose(grad, expected, atol=tol, rtol=0)
 
         tapes, fn = qml.gradients.finite_diff(tape)
-        numeric_val = fn(dev.batch_execute(tapes))
+        numeric_val = fn(dev.execute(tapes))
         assert np.allclose(grad, numeric_val, atol=tol, rtol=0)
 
     @pytest.mark.parametrize("theta", np.linspace(-2 * np.pi, np.pi, 7))
@@ -2571,7 +2579,7 @@ def test_CRot_gradient(self, theta, tol):
         assert len(tapes) == len(tape.trainable_params)
         assert [t.batch_size for t in tapes] == [4, 4, 4]
 
-        grad = fn(dev.batch_execute(tapes))
+        grad = fn(dev.execute(tapes))
         expected = np.array(
             [
                 0.5 * np.cos(b / 2) * np.sin(0.5 * (a + c)),
@@ -2584,7 +2592,7 @@ def test_CRot_gradient(self, theta, tol):
             assert np.allclose(g, e, atol=tol, rtol=0)
 
         tapes, fn = qml.gradients.finite_diff(tape)
-        numeric_val = fn(dev.batch_execute(tapes))
+        numeric_val = fn(dev.execute(tapes))
         assert np.allclose(grad, numeric_val, atol=tol, rtol=0)
 
     def test_gradients_agree_finite_differences(self, tol):
@@ -2650,10 +2658,12 @@ def test_variance_gradients_agree_finite_differences(self, tol):
                 assert np.allclose(g, grad_F2[idx1][idx2], atol=tol, rtol=0)
 
     @pytest.mark.autograd
-    def test_fallback(self, mocker):
+    @pytest.mark.parametrize("dev_name", ["default.qubit", "default.qubit.autograd"])
+    def test_fallback(self, dev_name, mocker):
         """Test that fallback gradient functions are correctly used"""
         spy = mocker.spy(qml.gradients, "finite_diff")
-        dev = qml.device("default.qubit.autograd", wires=2)
+        dev = qml.device(dev_name, wires=2)
+        execute_fn = dev.execute if dev_name == "default.qubit" else dev.batch_execute
         x = 0.543
         y = -0.654
 
@@ -2675,7 +2685,7 @@ def cost_fn(params):
             spy.assert_called()
             assert spy.call_args[1]["argnum"] == {1}
 
-            return fn(dev.batch_execute(tapes))
+            return fn(execute_fn(tapes))
 
         with pytest.raises(NotImplementedError, match="Broadcasting with multiple measurements"):
             cost_fn(params)
@@ -2691,13 +2701,15 @@ def cost_fn(params):
         # assert np.allclose(jac[1, 1, 1], -2 * np.cos(2 * y), atol=tol, rtol=0)
 
     @pytest.mark.autograd
-    def test_all_fallback(self, mocker, tol):
+    @pytest.mark.parametrize("dev_name", ["default.qubit", "default.qubit.autograd"])
+    def test_all_fallback(self, dev_name, mocker, tol):
         """Test that *only* the fallback logic is called if no parameters
         support the parameter-shift rule"""
         spy_fd = mocker.spy(qml.gradients, "finite_diff")
         spy_ps = mocker.spy(qml.gradients.parameter_shift, "expval_param_shift")
 
-        dev = qml.device("default.qubit.autograd", wires=2)
+        dev = qml.device(dev_name, wires=2)
+        execute_fn = dev.execute if dev_name == "default.qubit" else dev.batch_execute
         x = 0.543
         y = -0.654
 
@@ -2715,7 +2727,7 @@ def test_all_fallback(self, mocker, tol):
         spy_fd.assert_called()
         spy_ps.assert_not_called()
 
-        res = fn(dev.batch_execute(tapes))
+        res = fn(execute_fn(tapes))
         assert len(res) == 2
         assert res[0].shape == ()
         assert res[1].shape == ()
@@ -2741,7 +2753,7 @@ def test_single_expectation_value(self, tol):
         assert len(tapes) == 2
         assert tapes[0].batch_size == tapes[1].batch_size == 2
 
-        res = fn(dev.batch_execute(tapes))
+        res = fn(dev.execute(tapes))
         assert len(res) == 2
         assert res[0].shape == ()
         assert res[1].shape == ()
@@ -2773,7 +2785,7 @@ def test_multiple_expectation_values(self):
         # assert len(tapes) == 2
         # assert tapes[0].batch_size == tapes[1].batch_size == 2
 
-        # res = fn(dev.batch_execute(tapes))
+        # res = fn(dev.execute(tapes))
         # assert res.shape == (2, 2)
 
         # expected = np.array([[-np.sin(x), 0], [0, np.cos(y)]])
@@ -2802,7 +2814,7 @@ def test_var_expectation_values(self):
         # assert tapes[0].batch_size is None
         # assert tapes[1].batch_size == tapes[2].batch_size == 2
 
-        # res = fn(dev.batch_execute(tapes))
+        # res = fn(dev.execute(tapes))
         # assert res.shape == (2, 2)
 
         # expected = np.array([[-np.sin(x), 0], [0, -2 * np.cos(y) * np.sin(y)]])
@@ -2832,7 +2844,7 @@ def test_prob_expectation_values(self):
         # assert len(tapes) == 2
         # assert tapes[0].batch_size == tapes[1].batch_size == 2
 
-        # res = fn(dev.batch_execute(tapes))
+        # res = fn(dev.execute(tapes))
         # assert res.shape == (5, 2)
 
         # expected = (
@@ -2878,7 +2890,7 @@ def test_involutory_variance(self, tol):
 
         # circuit jacobians
         tapes, fn = qml.gradients.param_shift(tape, broadcast=True)
-        gradA = fn(dev.batch_execute(tapes))
+        gradA = fn(dev.execute(tapes))
         assert isinstance(gradA, np.ndarray)
         assert gradA.shape == ()
 
@@ -2887,7 +2899,7 @@ def test_involutory_variance(self, tol):
         assert tapes[1].batch_size == 2
 
         tapes, fn = qml.gradients.finite_diff(tape)
-        gradF = fn(dev.batch_execute(tapes))
+        gradF = fn(dev.execute(tapes))
         assert len(tapes) == 2
 
         expected = 2 * np.sin(a) * np.cos(a)
@@ -2914,7 +2926,7 @@ def test_non_involutory_variance(self, tol):
 
         # circuit jacobians
         tapes, fn = qml.gradients.param_shift(tape, broadcast=True)
-        gradA = fn(dev.batch_execute(tapes))
+        gradA = fn(dev.execute(tapes))
         assert isinstance(gradA, np.ndarray)
         assert gradA.shape == ()
 
@@ -2923,7 +2935,7 @@ def test_non_involutory_variance(self, tol):
         assert tapes[1].batch_size == tapes[2].batch_size == 2
 
         tapes, fn = qml.gradients.finite_diff(tape)
-        gradF = fn(dev.batch_execute(tapes))
+        gradF = fn(dev.execute(tapes))
         assert len(tapes) == 2
 
         expected = -35 * np.sin(2 * a) - 12 * np.cos(2 * a)
@@ -2955,11 +2967,11 @@ def test_involutory_and_noninvolutory_variance(self, tol):
             qml.gradients.param_shift(tape, broadcast=True)
         # TODO: Uncomment the following when #2693 is resolved.
         # tapes, fn = qml.gradients.param_shift(tape, broadcast=True)
-        # gradA = fn(dev.batch_execute(tapes))
+        # gradA = fn(dev.execute(tapes))
         # assert len(tapes) == 1 + 2 * 4
 
         # tapes, fn = qml.gradients.finite_diff(tape)
-        # gradF = fn(dev.batch_execute(tapes))
+        # gradF = fn(dev.execute(tapes))
         # assert len(tapes) == 1 + 2
 
         # expected = [2 * np.sin(a) * np.cos(a), -35 * np.sin(2 * a) - 12 * np.cos(2 * a)]
@@ -3001,10 +3013,10 @@ def test_expval_and_variance(self, tol):
             qml.gradients.param_shift(tape, broadcast=True)
         # TODO: Uncomment the following when #2693 is resolved.
         # tapes, fn = qml.gradients.param_shift(tape, broadcast=True)
-        # gradA = fn(dev.batch_execute(tapes))
+        # gradA = fn(dev.execute(tapes))
 
         # tapes, fn = qml.gradients.finite_diff(tape)
-        # gradF = fn(dev.batch_execute(tapes))
+        # gradF = fn(dev.execute(tapes))
         # ca, sa, cb, sb = np.cos(a), np.sin(a), np.cos(b), np.sin(b)
         # c2c, s2c = np.cos(2 * c), np.sin(2 * c)
         # expected = np.array(
@@ -3038,10 +3050,10 @@ def test_projector_variance(self, state, tol):
 
         # circuit jacobians
         tapes, fn = qml.gradients.param_shift(tape, broadcast=True)
-        gradA = fn(dev.batch_execute(tapes))
+        gradA = fn(dev.execute(tapes))
 
         tapes, fn = qml.gradients.finite_diff(tape)
-        gradF = fn(dev.batch_execute(tapes))
+        gradF = fn(dev.execute(tapes))
 
         expected = np.array(
             [
@@ -3113,10 +3125,12 @@ class TestParamShiftGradients:
     @pytest.mark.autograd
     # TODO: support Hessian with the new return types
     @pytest.mark.skip
-    def test_autograd(self, tol, broadcast, expected):
+    @pytest.mark.parametrize("dev_name", ["default.qubit", "default.qubit.autograd"])
+    def test_autograd(self, dev_name, tol, broadcast, expected):
         """Tests that the output of the parameter-shift transform
         can be differentiated using autograd, yielding second derivatives."""
-        dev = qml.device("default.qubit.autograd", wires=2)
+        dev = qml.device(dev_name, wires=2)
+        execute_fn = dev.execute if dev_name == "default.qubit" else dev.batch_execute
         params = np.array([0.543, -0.654], requires_grad=True)
         exp_num_tapes, exp_batch_sizes = expected
 
@@ -3132,7 +3146,7 @@ def cost_fn(x):
             tapes, fn = qml.gradients.param_shift(tape, broadcast=broadcast)
             assert len(tapes) == exp_num_tapes
             assert [t.batch_size for t in tapes] == exp_batch_sizes
-            jac = fn(dev.batch_execute(tapes))
+            jac = fn(execute_fn(tapes))
             return jac
 
         res = qml.jacobian(cost_fn)(params)
@@ -3194,7 +3208,7 @@ def test_no_trainable_coeffs(self, mocker, tol, broadcast):
         x, y = weights
         tape.trainable_params = {0, 1}
 
-        res = dev.batch_execute([tape])
+        res = dev.execute([tape])
         expected = -c * np.sin(x) * np.sin(y) + np.cos(x) * (a + b * np.sin(y))
         assert np.allclose(res, expected, atol=tol, rtol=0)
 
@@ -3204,7 +3218,7 @@ def test_no_trainable_coeffs(self, mocker, tol, broadcast):
         assert [t.batch_size for t in tapes] == ([2, 2] if broadcast else [None] * 4)
         spy.assert_not_called()
 
-        res = fn(dev.batch_execute(tapes))
+        res = fn(dev.execute(tapes))
         assert isinstance(res, tuple)
 
         assert len(res) == 2
@@ -3240,7 +3254,7 @@ def test_trainable_coeffs(self, mocker, tol, broadcast):
         x, y = weights
         tape.trainable_params = {0, 1, 2, 4}
 
-        res = dev.batch_execute([tape])
+        res = dev.execute([tape])
         expected = -c * np.sin(x) * np.sin(y) + np.cos(x) * (a + b * np.sin(y))
         assert np.allclose(res, expected, atol=tol, rtol=0)
 
@@ -3251,7 +3265,7 @@ def test_trainable_coeffs(self, mocker, tol, broadcast):
         assert [t.batch_size for t in tapes] == ([2, 2, None, None] if broadcast else [None] * 6)
         spy.assert_called()
 
-        res = fn(dev.batch_execute(tapes))
+        res = fn(dev.execute(tapes))
         assert isinstance(res, tuple)
         assert len(res) == 4
         assert res[0].shape == ()
@@ -3298,7 +3312,7 @@ def test_multiple_hamiltonians(self, mocker, tol, broadcast):
         tape = qml.tape.QuantumScript.from_queue(q)
         tape.trainable_params = {0, 1, 2, 4, 5}
 
-        res = dev.batch_execute([tape])
+        res = dev.execute([tape])
         expected = [-c * np.sin(x) * np.sin(y) + np.cos(x) * (a + b * np.sin(y)), d * np.cos(x)]
         assert np.allclose(res, expected, atol=tol, rtol=0)
 
@@ -3313,7 +3327,7 @@ def test_multiple_hamiltonians(self, mocker, tol, broadcast):
         assert len(tapes) == 2 * 2 + 3
         spy.assert_called()
 
-        res = fn(dev.batch_execute(tapes))
+        res = fn(dev.execute(tapes))
         assert isinstance(res, tuple)
         assert len(res) == 2
         assert len(res[0]) == 5
@@ -3351,7 +3365,8 @@ def cost_fn(weights, coeffs1, coeffs2, dev=None, broadcast=False):
         tape = qml.tape.QuantumScript.from_queue(q)
         tape.trainable_params = {0, 1, 2, 3, 4, 5}
         tapes, fn = qml.gradients.param_shift(tape, broadcast=broadcast)
-        jac = fn(dev.batch_execute(tapes))
+        execute_fn = dev.batch_execute if isinstance(dev, qml.Device) else dev.execute
+        jac = fn(execute_fn(tapes))
         return jac
 
     @staticmethod
@@ -3373,13 +3388,14 @@ def cost_fn_expected(weights, coeffs1, coeffs2):
         ]
 
     @pytest.mark.autograd
-    def test_autograd(self, tol, broadcast):
+    @pytest.mark.parametrize("dev_name", ["default.qubit", "default.qubit.autograd"])
+    def test_autograd(self, dev_name, tol, broadcast):
         """Test gradient of multiple trainable Hamiltonian coefficients
         using autograd"""
         coeffs1 = np.array([0.1, 0.2, 0.3], requires_grad=True)
         coeffs2 = np.array([0.7], requires_grad=True)
         weights = np.array([0.4, 0.5], requires_grad=True)
-        dev = qml.device("default.qubit.autograd", wires=2)
+        dev = qml.device(dev_name, wires=2)
 
         if broadcast:
             with pytest.raises(
@@ -3399,7 +3415,8 @@ def test_autograd(self, tol, broadcast):
         # assert np.allclose(res[2][:, -1], np.zeros([2, 1, 1]), atol=tol, rtol=0)
 
     @pytest.mark.tf
-    def test_tf(self, tol, broadcast):
+    @pytest.mark.parametrize("dev_name", ["default.qubit", "default.qubit.tf"])
+    def test_tf(self, dev_name, tol, broadcast):
         """Test gradient of multiple trainable Hamiltonian coefficients
         using tf"""
         import tensorflow as tf
@@ -3408,7 +3425,7 @@ def test_tf(self, tol, broadcast):
         coeffs2 = tf.Variable([0.7], dtype=tf.float64)
         weights = tf.Variable([0.4, 0.5], dtype=tf.float64)
 
-        dev = qml.device("default.qubit.tf", wires=2)
+        dev = qml.device(dev_name, wires=2)
 
         if broadcast:
             with pytest.raises(
@@ -3435,7 +3452,8 @@ def test_tf(self, tol, broadcast):
     # TODO: Torch support for param-shift
     @pytest.mark.torch
     @pytest.mark.xfail
-    def test_torch(self, tol, broadcast):
+    @pytest.mark.parametrize("dev_name", ["default.qubit", "default.qubit.torch"])
+    def test_torch(self, dev_name, tol, broadcast):
         """Test gradient of multiple trainable Hamiltonian coefficients
         using torch"""
         import torch
@@ -3444,7 +3462,7 @@ def test_torch(self, tol, broadcast):
         coeffs2 = torch.tensor([0.7], dtype=torch.float64, requires_grad=True)
         weights = torch.tensor([0.4, 0.5], dtype=torch.float64, requires_grad=True)
 
-        dev = qml.device("default.qubit.torch", wires=2)
+        dev = qml.device(dev_name, wires=2)
 
         if broadcast:
             with pytest.raises(
@@ -3466,7 +3484,8 @@ def test_torch(self, tol, broadcast):
         assert np.allclose(hess[2][:, -1], np.zeros([2, 1, 1]), atol=tol, rtol=0)
 
     @pytest.mark.jax
-    def test_jax(self, tol, broadcast):
+    @pytest.mark.parametrize("dev_name", ["default.qubit", "default.qubit.jax"])
+    def test_jax(self, dev_name, tol, broadcast):
         """Test gradient of multiple trainable Hamiltonian coefficients
         using JAX"""
         import jax
@@ -3476,7 +3495,7 @@ def test_jax(self, tol, broadcast):
         coeffs1 = jnp.array([0.1, 0.2, 0.3])
         coeffs2 = jnp.array([0.7])
         weights = jnp.array([0.4, 0.5])
-        dev = qml.device("default.qubit.jax", wires=2)
+        dev = qml.device(dev_name, wires=2)
 
         if broadcast:
             with pytest.raises(
diff --git a/tests/gradients/parameter_shift/test_parameter_shift_shot_vec.py b/tests/gradients/parameter_shift/test_parameter_shift_shot_vec.py
index 16499641081..563423a0416 100644
--- a/tests/gradients/parameter_shift/test_parameter_shift_shot_vec.py
+++ b/tests/gradients/parameter_shift/test_parameter_shift_shot_vec.py
@@ -40,7 +40,7 @@
 def grad_fn(tape, dev, fn=qml.gradients.param_shift, **kwargs):
     """Utility function to automate execution and processing of gradient tapes"""
     tapes, fn = fn(tape, **kwargs)
-    return fn(dev.batch_execute(tapes))
+    return fn(dev.execute(tapes))
 
 
 # pylint: disable=too-few-public-methods
@@ -90,7 +90,7 @@ def test_independent_parameter(self, mocker):
         assert len(tapes) == 2
         assert tapes[0].batch_size == tapes[1].batch_size == None
 
-        res = fn(dev.batch_execute(tapes))
+        res = fn(dev.execute(tapes))
         for r in res:
             assert isinstance(r, tuple)
             assert len(r) == 2
@@ -293,7 +293,7 @@ def test_all_zero_diff_methods_multiple_returns_tape(self):
         g_tapes, post_processing = qml.gradients.param_shift(tape)
         assert g_tapes == []
 
-        all_result = post_processing(dev.batch_execute(g_tapes))
+        all_result = post_processing(dev.execute(g_tapes))
 
         assert isinstance(all_result, tuple)
 
@@ -414,7 +414,7 @@ def test_f0_provided(self, y_wire):
         # one tape per parameter that impacts the expval
         assert len(tapes) == 2 if y_wire == 0 else 1
 
-        fn(dev.batch_execute(tapes))
+        fn(dev.execute(tapes))
 
     def test_op_with_custom_unshifted_term(self):
         """Test that an operation with a gradient recipe that depends on
@@ -451,7 +451,7 @@ class RX(qml.RX):
             assert tape.operations[0].data[0] == x[0] + expected[0]
             assert tape.operations[1].data[0] == x[1] + expected[1]
 
-        grad = fn(dev.batch_execute(tapes))
+        grad = fn(dev.execute(tapes))
         exp = np.stack([-np.sin(x[0] + x[1]), -np.sin(x[0] + x[1]) + 0.2 * np.cos(x[0] + x[1])])
         assert isinstance(grad, tuple)
         assert len(grad) == len(default_shot_vector)
@@ -484,15 +484,16 @@ def test_independent_parameters_analytic(self):
 
         tape2 = qml.tape.QuantumScript.from_queue(q2, shots=shot_vec)
         tapes, fn = qml.gradients.param_shift(tape1)
-        j1 = fn(dev.batch_execute(tapes))
+        with qml.Tracker(dev) as tracker:
+            j1 = fn(dev.execute(tapes))
 
         # We should only be executing the device twice: Two shifted evaluations to differentiate
         # one parameter overall, as the other parameter does not impact the returned measurement.
 
-        assert dev.num_executions == 2
+        assert tracker.latest["executions"] == 2
 
         tapes, fn = qml.gradients.param_shift(tape2)
-        j2 = fn(dev.batch_execute(tapes))
+        j2 = fn(dev.execute(tapes))
 
         exp = -np.sin(1)
 
@@ -532,7 +533,7 @@ def test_grad_recipe_parameter_dependent(self):
         assert qml.math.allclose(tapes[0].operations[0].data[0], 0)
         assert qml.math.allclose(tapes[1].operations[0].data[0], 2 * x)
 
-        grad = fn(dev.batch_execute(tapes))
+        grad = fn(dev.execute(tapes))
         assert np.allclose(grad, -np.sin(x), atol=shot_vec_tol)
 
     def test_error_no_diff_info(self):
@@ -606,14 +607,14 @@ def test_pauli_rotation_gradient(self, mocker, G, theta, shift):
         tapes, fn = qml.gradients.param_shift(tape, shifts=[(shift,)])
         assert len(tapes) == 2
 
-        autograd_val = fn(dev.batch_execute(tapes))
+        autograd_val = fn(dev.execute(tapes))
 
         tape_fwd = tape.bind_new_parameters([theta + np.pi / 2], [1])
         tape_bwd = tape.bind_new_parameters([theta - np.pi / 2], [1])
 
-        shot_vec_manual_res = dev.batch_execute([tape_fwd, tape_bwd])
+        shot_vec_manual_res = dev.execute([tape_fwd, tape_bwd])
 
-        # Parameter axis is the first - reorder the results from batch_execute
+        # Parameter axis is the first - reorder the results from execute
         shot_vec_len = len(many_shots_shot_vector)
         shot_vec_manual_res = [
             tuple(comp[l] for comp in shot_vec_manual_res) for l in range(shot_vec_len)
@@ -628,7 +629,7 @@ def test_pauli_rotation_gradient(self, mocker, G, theta, shift):
         assert spy.call_args[1]["shifts"] == (shift,)
 
         tapes, fn = qml.gradients.finite_diff(tape, h=h_val)
-        numeric_val = fn(dev.batch_execute(tapes))
+        numeric_val = fn(dev.execute(tapes))
         for a_val, n_val in zip(autograd_val, numeric_val):
             assert np.allclose(a_val, n_val, atol=finite_diff_tol, rtol=0)
 
@@ -654,7 +655,7 @@ def test_Rot_gradient(self, mocker, theta, shift):
         num_params = len(tape.trainable_params)
         assert len(tapes) == 2 * num_params
 
-        autograd_val = fn(dev.batch_execute(tapes))
+        autograd_val = fn(dev.execute(tapes))
         assert isinstance(autograd_val, tuple)
         assert len(autograd_val) == len(shot_vec)
 
@@ -687,7 +688,7 @@ def test_Rot_gradient(self, mocker, theta, shift):
             assert spy.call_args[1]["shifts"] == (shift,)
 
         tapes, fn = qml.gradients.finite_diff(tape, h=h_val)
-        numeric_val = fn(dev.batch_execute(tapes))
+        numeric_val = fn(dev.execute(tapes))
         for a_val, n_val in zip(autograd_val, numeric_val):
             assert np.allclose(a_val, n_val, atol=finite_diff_tol, rtol=0)
 
@@ -710,14 +711,14 @@ def test_controlled_rotation_gradient(self, G):
         assert np.allclose(res, -np.cos(b / 2), atol=shot_vec_tol, rtol=0)
 
         tapes, fn = qml.gradients.param_shift(tape)
-        grad = fn(dev.batch_execute(tapes))
+        grad = fn(dev.execute(tapes))
         expected = np.sin(b / 2) / 2
         assert isinstance(grad, tuple)
         assert len(grad) == len(many_shots_shot_vector)
         assert np.allclose(grad, expected, atol=shot_vec_tol, rtol=0)
 
         tapes, fn = qml.gradients.finite_diff(tape, h=h_val)
-        numeric_val = fn(dev.batch_execute(tapes))
+        numeric_val = fn(dev.execute(tapes))
         for a_val, n_val in zip(grad, numeric_val):
             assert np.allclose(a_val, n_val, atol=finite_diff_tol, rtol=0)
 
@@ -744,7 +745,7 @@ def test_CRot_gradient(self, theta):
         tapes, fn = qml.gradients.param_shift(tape)
         assert len(tapes) == 4 * len(tape.trainable_params)
 
-        grad = fn(dev.batch_execute(tapes))
+        grad = fn(dev.execute(tapes))
         expected = np.array(
             [
                 0.5 * np.cos(b / 2) * np.sin(0.5 * (a + c)),
@@ -762,7 +763,7 @@ def test_CRot_gradient(self, theta):
                 assert np.allclose(g, expected[idx], atol=shot_vec_tol, rtol=0)
 
         tapes, fn = qml.gradients.finite_diff(tape, h=h_val)
-        numeric_val = fn(dev.batch_execute(tapes))
+        numeric_val = fn(dev.execute(tapes))
         for a_val, n_val in zip(grad, numeric_val):
             assert np.allclose(a_val, n_val, atol=finite_diff_tol, rtol=0)
 
@@ -842,10 +843,12 @@ def test_variance_gradients_agree_finite_differences(self):
                 assert np.allclose(g, grad_F2[idx1][idx2], atol=finite_diff_tol, rtol=0)
 
     @pytest.mark.autograd
-    def test_fallback(self, mocker):
+    @pytest.mark.parametrize("dev_name", ["default.qubit", "default.qubit.autograd"])
+    def test_fallback(self, dev_name, mocker):
         """Test that fallback gradient functions are correctly used"""
         spy = mocker.spy(qml.gradients, "finite_diff")
-        dev = qml.device("default.qubit.autograd", wires=3, shots=fallback_shot_vec)
+        dev = qml.device(dev_name, wires=3, shots=fallback_shot_vec)
+        execute_fn = dev.execute if dev_name == "default.qubit" else dev.batch_execute
         x = 0.543
         y = -0.654
 
@@ -869,7 +872,7 @@ def cost_fn(params):
             spy.assert_called()
             assert spy.call_args[1]["argnum"] == {1}
 
-            return fn(dev.batch_execute(tapes))
+            return fn(execute_fn(tapes))
 
         all_res = cost_fn(params)
 
@@ -898,11 +901,13 @@ def cost_fn(params):
             # assert np.allclose(jac[1, 1, 1], -2 * np.cos(2 * y), atol=shot_vec_tol, rtol=0)
 
     @pytest.mark.autograd
-    def test_fallback_single_meas(self, mocker):
+    @pytest.mark.parametrize("dev_name", ["default.qubit", "default.qubit.autograd"])
+    def test_fallback_single_meas(self, dev_name, mocker):
         """Test that fallback gradient functions are correctly used for a single measurement."""
         spy = mocker.spy(qml.gradients, "finite_diff")
         shot_vec = tuple([1000000] * 4)
-        dev = qml.device("default.qubit.autograd", wires=3, shots=shot_vec)
+        dev = qml.device(dev_name, wires=3, shots=shot_vec)
+        execute_fn = dev.execute if dev_name == "default.qubit" else dev.batch_execute
         x = 0.543
         y = -0.654
 
@@ -923,7 +928,7 @@ def cost_fn(params):
             spy.assert_called()
             assert spy.call_args[1]["argnum"] == {1}
 
-            return fn(dev.batch_execute(tapes))
+            return fn(execute_fn(tapes))
 
         all_res = cost_fn(params)
 
@@ -941,10 +946,14 @@ def cost_fn(params):
 
     @pytest.mark.autograd
     @pytest.mark.parametrize("RX, RY, argnum", [(RX_with_F, qml.RY, 0), (qml.RX, RY_with_F, 1)])
-    def test_fallback_probs(self, RX, RY, argnum, mocker):
+    @pytest.mark.parametrize("dev_name", ["default.qubit", "default.qubit.autograd"])
+    def test_fallback_probs(
+        self, dev_name, RX, RY, argnum, mocker
+    ):  # pylint:disable=too-many-arguments
         """Test that fallback gradient functions are correctly used with probs"""
         spy = mocker.spy(qml.gradients, "finite_diff")
-        dev = qml.device("default.qubit.autograd", wires=3, shots=fallback_shot_vec)
+        dev = qml.device(dev_name, wires=3, shots=fallback_shot_vec)
+        execute_fn = dev.execute if dev_name == "default.qubit" else dev.batch_execute
         x = 0.543
         y = -0.654
 
@@ -967,7 +976,7 @@ def cost_fn(params):
             spy.assert_called()
 
             assert spy.call_args[1]["argnum"] == {argnum}
-            return fn(dev.batch_execute(tapes))
+            return fn(execute_fn(tapes))
 
         all_res = cost_fn(params)
         assert isinstance(all_res, tuple)
@@ -1028,13 +1037,15 @@ def cost_fn(params):
             assert np.allclose(res[1][1], probs_expected[:, 1], atol=finite_diff_tol)
 
     @pytest.mark.autograd
-    def test_all_fallback(self, mocker):
+    @pytest.mark.parametrize("dev_name", ["default.qubit", "default.qubit.autograd"])
+    def test_all_fallback(self, dev_name, mocker):
         """Test that *only* the fallback logic is called if no parameters
         support the parameter-shift rule"""
         spy_fd = mocker.spy(qml.gradients, "finite_diff")
         spy_ps = mocker.spy(qml.gradients.parameter_shift, "expval_param_shift")
 
-        dev = qml.device("default.qubit.autograd", wires=3, shots=fallback_shot_vec)
+        dev = qml.device(dev_name, wires=3, shots=fallback_shot_vec)
+        execute_fn = dev.execute if dev_name == "default.qubit" else dev.batch_execute
         x = 0.543
         y = -0.654
 
@@ -1053,7 +1064,7 @@ def test_all_fallback(self, mocker):
         spy_fd.assert_called()
         spy_ps.assert_not_called()
 
-        all_res = fn(dev.batch_execute(tapes))
+        all_res = fn(execute_fn(tapes))
         assert len(all_res) == len(fallback_shot_vec)
         assert isinstance(all_res, tuple)
 
@@ -1085,7 +1096,7 @@ def test_single_expectation_value(self):
         tapes, fn = qml.gradients.param_shift(tape)
         assert len(tapes) == 4
 
-        all_res = fn(dev.batch_execute(tapes))
+        all_res = fn(dev.execute(tapes))
 
         assert len(all_res) == len(many_shots_shot_vector)
         assert isinstance(all_res, tuple)
@@ -1118,7 +1129,7 @@ def test_multiple_expectation_values(self):
         tapes, fn = qml.gradients.param_shift(tape)
         assert len(tapes) == 4
 
-        all_res = fn(dev.batch_execute(tapes))
+        all_res = fn(dev.execute(tapes))
         assert len(all_res) == len(many_shots_shot_vector)
         assert isinstance(all_res, tuple)
 
@@ -1151,7 +1162,7 @@ def test_var_expectation_values(self):
         tapes, fn = qml.gradients.param_shift(tape)
         assert len(tapes) == 5
 
-        all_res = fn(dev.batch_execute(tapes))
+        all_res = fn(dev.execute(tapes))
         assert len(all_res) == len(many_shots_shot_vector)
         assert isinstance(all_res, tuple)
 
@@ -1185,7 +1196,7 @@ def test_prob_expectation_values(self):
         tapes, fn = qml.gradients.param_shift(tape)
         assert len(tapes) == 4
 
-        res = fn(dev.batch_execute(tapes))
+        res = fn(dev.execute(tapes))
         assert isinstance(res, tuple)
         assert len(res) == len(many_shots_shot_vector)
 
@@ -1272,7 +1283,7 @@ def test_involutory_variance_single_param(self):
 
         # circuit jacobians
         tapes, fn = qml.gradients.param_shift(tape)
-        gradA = fn(dev.batch_execute(tapes))
+        gradA = fn(dev.execute(tapes))
         for _gA in gradA:
             assert isinstance(_gA, np.ndarray)
             assert _gA.shape == ()
@@ -1280,7 +1291,7 @@ def test_involutory_variance_single_param(self):
         assert len(tapes) == 1 + 2 * 1
 
         tapes, fn = qml.gradients.finite_diff(tape, h=h_val)
-        all_gradF = fn(dev.batch_execute(tapes))
+        all_gradF = fn(dev.execute(tapes))
         assert len(tapes) == 2
 
         expected = 2 * np.sin(a) * np.cos(a)
@@ -1312,7 +1323,7 @@ def test_involutory_variance_multi_param(self):
 
         # circuit jacobians
         tapes, fn = qml.gradients.param_shift(tape)
-        all_res = fn(dev.batch_execute(tapes))
+        all_res = fn(dev.execute(tapes))
         assert len(all_res) == len(many_shots_shot_vector)
         assert isinstance(all_res, tuple)
 
@@ -1326,7 +1337,7 @@ def test_involutory_variance_multi_param(self):
             assert len(tapes) == 1 + 2 * 2
 
         tapes, fn = qml.gradients.finite_diff(tape, h=h_val)
-        all_Fres = fn(dev.batch_execute(tapes))
+        all_Fres = fn(dev.execute(tapes))
         for gradF, gradA in zip(all_Fres, all_res):
             assert len(tapes) == 3
 
@@ -1359,11 +1370,11 @@ def test_non_involutory_variance_single_param(self):
 
         # circuit jacobians
         tapes, fn = qml.gradients.param_shift(tape)
-        gradA = fn(dev.batch_execute(tapes))
+        gradA = fn(dev.execute(tapes))
         assert len(tapes) == 1 + 4 * 1
 
         tapes, fn = qml.gradients.finite_diff(tape, h=h_val)
-        all_gradF = fn(dev.batch_execute(tapes))
+        all_gradF = fn(dev.execute(tapes))
         assert len(tapes) == 2
 
         expected = -35 * np.sin(2 * a) - 12 * np.cos(2 * a)
@@ -1401,7 +1412,7 @@ def test_non_involutory_variance_multi_param(self):
 
         # circuit jacobians
         tapes, fn = qml.gradients.param_shift(tape)
-        all_res = fn(dev.batch_execute(tapes))
+        all_res = fn(dev.execute(tapes))
         assert len(all_res) == len(many_shots_shot_vector)
         assert isinstance(all_res, tuple)
 
@@ -1419,7 +1430,7 @@ def test_non_involutory_variance_multi_param(self):
             assert gradA[1] == pytest.approx(expected, abs=herm_shot_vec_tol)
 
         tapes, fn = qml.gradients.finite_diff(tape, h=h_val)
-        all_gradF = fn(dev.batch_execute(tapes))
+        all_gradF = fn(dev.execute(tapes))
         assert len(all_gradF) == len(many_shots_shot_vector)
         assert isinstance(all_gradF, tuple)
         for gradF in all_gradF:
@@ -1456,11 +1467,11 @@ def test_involutory_and_noninvolutory_variance_single_param(self):
 
         # circuit jacobians
         tapes, fn = qml.gradients.param_shift(tape)
-        gradA = fn(dev.batch_execute(tapes))
+        gradA = fn(dev.execute(tapes))
         assert len(tapes) == 1 + 4
 
         tapes, fn = qml.gradients.finite_diff(tape, h=h_val)
-        gradF = fn(dev.batch_execute(tapes))
+        gradF = fn(dev.execute(tapes))
         assert len(tapes) == 1 + 1
 
         expected = [2 * np.sin(a) * np.cos(a), 0]
@@ -1507,7 +1518,7 @@ def test_involutory_and_noninvolutory_variance_multi_param(self):
 
         # circuit jacobians
         tapes, fn = qml.gradients.param_shift(tape)
-        gradA = fn(dev.batch_execute(tapes))
+        gradA = fn(dev.execute(tapes))
 
         assert isinstance(gradA, tuple)
         assert len(gradA) == len(many_shots_shot_vector)
@@ -1524,7 +1535,7 @@ def test_involutory_and_noninvolutory_variance_multi_param(self):
         assert len(tapes) == 1 + 2 * 4
 
         tapes, fn = qml.gradients.finite_diff(tape, h=h_val)
-        gradF = fn(dev.batch_execute(tapes))
+        gradF = fn(dev.execute(tapes))
         assert len(tapes) == 1 + 2
 
         expected = [2 * np.sin(a) * np.cos(a), 0, 0, -35 * np.sin(2 * a) - 12 * np.cos(2 * a)]
@@ -1590,7 +1601,7 @@ def test_var_and_probs_single_param(self, ind):
         # circuit jacobians
         tapes, fn = qml.gradients.param_shift(tape)
 
-        all_res = fn(dev.batch_execute(tapes))
+        all_res = fn(dev.execute(tapes))
         assert len(all_res) == len(many_shots_shot_vector)
         assert isinstance(all_res, tuple)
 
@@ -1647,7 +1658,7 @@ def test_var_and_probs_multi_params(self):
 
         # circuit jacobians
         tapes, fn = qml.gradients.param_shift(tape)
-        all_res = fn(dev.batch_execute(tapes))
+        all_res = fn(dev.execute(tapes))
         assert len(all_res) == len(many_shots_shot_vector)
         assert isinstance(all_res, tuple)
 
@@ -1750,7 +1761,7 @@ def test_expval_and_variance_single_param(self):
 
         # # circuit jacobians
         tapes, fn = qml.gradients.param_shift(tape)
-        all_res = fn(dev.batch_execute(tapes))
+        all_res = fn(dev.execute(tapes))
 
         assert len(all_res) == len(many_shots_shot_vector)
         assert isinstance(all_res, tuple)
@@ -1764,7 +1775,7 @@ def test_expval_and_variance_single_param(self):
                 assert np.allclose(a_comp, e_comp, atol=shot_vec_tol, rtol=0)
 
         tapes, fn = qml.gradients.finite_diff(tape, h=h_val)
-        all_gradF = fn(dev.batch_execute(tapes))
+        all_gradF = fn(dev.execute(tapes))
         assert isinstance(all_gradF, tuple)
 
         for gradF in all_gradF:
@@ -1806,7 +1817,7 @@ def test_expval_and_variance_multi_param(self):
 
         # # circuit jacobians
         tapes, fn = qml.gradients.param_shift(tape)
-        all_res = fn(dev.batch_execute(tapes))
+        all_res = fn(dev.execute(tapes))
 
         assert len(all_res) == len(many_shots_shot_vector)
         assert isinstance(all_res, tuple)
@@ -1831,7 +1842,7 @@ def test_expval_and_variance_multi_param(self):
                     assert np.allclose(a_comp, e_comp, atol=shot_vec_tol, rtol=0)
 
         tapes, fn = qml.gradients.finite_diff(tape, h=h_val)
-        all_gradF = fn(dev.batch_execute(tapes))
+        all_gradF = fn(dev.execute(tapes))
         for gradF in all_gradF:
             assert gradF == pytest.approx(expected, abs=finite_diff_tol)
 
@@ -1861,7 +1872,7 @@ def test_projector_variance(self, state):
 
         # # circuit jacobians
         tapes, fn = qml.gradients.param_shift(tape)
-        all_res = fn(dev.batch_execute(tapes))
+        all_res = fn(dev.execute(tapes))
 
         assert len(all_res) == len(many_shots_shot_vector)
         assert isinstance(all_res, tuple)
@@ -1876,7 +1887,7 @@ def test_projector_variance(self, state):
             assert np.allclose(gradA, expected, atol=shot_vec_tol, rtol=0)
 
         tapes, fn = qml.gradients.finite_diff(tape, h=h_val)
-        all_gradF = fn(dev.batch_execute(tapes))
+        all_gradF = fn(dev.execute(tapes))
         for gradF in all_gradF:
             assert gradF == pytest.approx(expected, abs=finite_diff_tol)
 
@@ -2074,7 +2085,7 @@ def test_multi_measure_no_warning(self):
         tape = qml.tape.QuantumScript.from_queue(q, shots=shot_vec)
         with warnings.catch_warnings(record=True) as record:
             tapes, fn = qml.gradients.param_shift(tape)
-            fn(dev.batch_execute(tapes))
+            fn(dev.execute(tapes))
 
         assert len(record) == 0
 
@@ -2132,7 +2143,7 @@ def test_no_trainable_coeffs(self, mocker, broadcast, tol):
         x, y = weights
         tape.trainable_params = {0, 1}
 
-        res = dev.batch_execute([tape])
+        res = dev.execute([tape])
         expected = -c * np.sin(x) * np.sin(y) + np.cos(x) * (a + b * np.sin(y))
         assert np.allclose(res, expected, atol=shot_vec_tol, rtol=0)
 
@@ -2142,7 +2153,7 @@ def test_no_trainable_coeffs(self, mocker, broadcast, tol):
         assert [t.batch_size for t in tapes] == ([2, 2] if broadcast else [None] * 4)
         spy.assert_not_called()
 
-        all_res = fn(dev.batch_execute(tapes))
+        all_res = fn(dev.execute(tapes))
         assert isinstance(all_res, tuple)
         assert len(all_res) == len(many_shots_shot_vector)
 
@@ -2183,7 +2194,7 @@ def test_trainable_coeffs(self, mocker, broadcast, tol):
         x, y = weights
         tape.trainable_params = {0, 1, 2, 4}
 
-        res = dev.batch_execute([tape])
+        res = dev.execute([tape])
         expected = -c * np.sin(x) * np.sin(y) + np.cos(x) * (a + b * np.sin(y))
         assert np.allclose(res, expected, atol=tol, rtol=0)
 
@@ -2194,7 +2205,7 @@ def test_trainable_coeffs(self, mocker, broadcast, tol):
         assert [t.batch_size for t in tapes] == ([2, 2, None, None] if broadcast else [None] * 6)
         spy.assert_called()
 
-        res = fn(dev.batch_execute(tapes))
+        res = fn(dev.execute(tapes))
         assert isinstance(res, tuple)
         assert len(res) == 4
         assert res[0].shape == ()
@@ -2243,7 +2254,7 @@ def test_multiple_hamiltonians(self, mocker, broadcast, tol):
         tape = qml.tape.QuantumScript.from_queue(q, shots=shot_vec)
         tape.trainable_params = {0, 1, 2, 4, 5}
 
-        res = dev.batch_execute([tape])
+        res = dev.execute([tape])
         expected = [-c * np.sin(x) * np.sin(y) + np.cos(x) * (a + b * np.sin(y)), d * np.cos(x)]
         assert np.allclose(res, expected, atol=tol, rtol=0)
 
@@ -2258,7 +2269,7 @@ def test_multiple_hamiltonians(self, mocker, broadcast, tol):
         assert len(tapes) == 2 * 2 + 3
         spy.assert_called()
 
-        res = fn(dev.batch_execute(tapes))
+        res = fn(dev.execute(tapes))
         assert isinstance(res, tuple)
         assert len(res) == 2
         assert len(res[0]) == 5
@@ -2296,7 +2307,8 @@ def cost_fn(weights, coeffs1, coeffs2, dev=None, broadcast=False):
         tape = qml.tape.QuantumScript.from_queue(q, shots=dev.shots)
         tape.trainable_params = {0, 1, 2, 3, 4, 5}
         tapes, fn = qml.gradients.param_shift(tape, broadcast=broadcast)
-        return fn(dev.batch_execute(tapes))
+        execute_fn = dev.batch_execute if isinstance(dev, qml.Device) else dev.execute
+        return fn(execute_fn(tapes))
 
     @staticmethod
     def cost_fn_expected(weights, coeffs1, coeffs2):
@@ -2318,14 +2330,15 @@ def cost_fn_expected(weights, coeffs1, coeffs2):
 
     @pytest.mark.xfail(reason="TODO")
     @pytest.mark.autograd
-    def test_autograd(self, broadcast, tol):
+    @pytest.mark.parametrize("dev_name", ["default.qubit", "default.qubit.autograd"])
+    def test_autograd(self, dev_name, broadcast, tol):
         """Test gradient of multiple trainable Hamiltonian coefficients
         using autograd"""
         coeffs1 = np.array([0.1, 0.2, 0.3], requires_grad=True)
         coeffs2 = np.array([0.7], requires_grad=True)
         weights = np.array([0.4, 0.5], requires_grad=True)
         shot_vec = many_shots_shot_vector
-        dev = qml.device("default.qubit.autograd", wires=2, shots=shot_vec)
+        dev = qml.device(dev_name, wires=2, shots=shot_vec)
 
         if broadcast:
             with pytest.raises(
@@ -2346,7 +2359,8 @@ def test_autograd(self, broadcast, tol):
 
     @pytest.mark.xfail(reason="TODO")
     @pytest.mark.tf
-    def test_tf(self, broadcast, tol):
+    @pytest.mark.parametrize("dev_name", ["default.qubit", "default.qubit.tf"])
+    def test_tf(self, dev_name, broadcast, tol):
         """Test gradient of multiple trainable Hamiltonian coefficients
         using tf"""
         import tensorflow as tf
@@ -2356,7 +2370,7 @@ def test_tf(self, broadcast, tol):
         weights = tf.Variable([0.4, 0.5], dtype=tf.float64)
 
         shot_vec = many_shots_shot_vector
-        dev = qml.device("default.qubit.tf", wires=2, shots=shot_vec)
+        dev = qml.device(dev_name, wires=2, shots=shot_vec)
 
         if broadcast:
             with pytest.raises(
@@ -2383,7 +2397,8 @@ def test_tf(self, broadcast, tol):
     # TODO: Torch support for param-shift
     @pytest.mark.torch
     @pytest.mark.xfail
-    def test_torch(self, broadcast, tol):
+    @pytest.mark.parametrize("dev_name", ["default.qubit", "default.qubit.torch"])
+    def test_torch(self, dev_name, broadcast, tol):
         """Test gradient of multiple trainable Hamiltonian coefficients
         using torch"""
         import torch
@@ -2392,7 +2407,7 @@ def test_torch(self, broadcast, tol):
         coeffs2 = torch.tensor([0.7], dtype=torch.float64, requires_grad=True)
         weights = torch.tensor([0.4, 0.5], dtype=torch.float64, requires_grad=True)
 
-        dev = qml.device("default.qubit.torch", wires=2)
+        dev = qml.device(dev_name, wires=2)
 
         if broadcast:
             with pytest.raises(
@@ -2413,9 +2428,9 @@ def test_torch(self, broadcast, tol):
         assert np.allclose(hess[1][:, 2:5], np.zeros([2, 3, 3]), atol=tol, rtol=0)
         assert np.allclose(hess[2][:, -1], np.zeros([2, 1, 1]), atol=tol, rtol=0)
 
-    @pytest.mark.xfail(reason="TODO")
     @pytest.mark.jax
-    def test_jax(self, broadcast, tol):
+    @pytest.mark.parametrize("dev_name", ["default.qubit", "default.qubit.jax"])
+    def test_jax(self, dev_name, broadcast, tol):
         """Test gradient of multiple trainable Hamiltonian coefficients
         using JAX"""
         import jax
@@ -2425,7 +2440,7 @@ def test_jax(self, broadcast, tol):
         coeffs1 = jnp.array([0.1, 0.2, 0.3])
         coeffs2 = jnp.array([0.7])
         weights = jnp.array([0.4, 0.5])
-        dev = qml.device("default.qubit.jax", wires=2)
+        dev = qml.device(dev_name, wires=2)
 
         if broadcast:
             with pytest.raises(
@@ -2516,7 +2531,7 @@ def test_1_1(self, shot_vec, meas, shape, op_wires):
         tape.trainable_params = {0}
 
         tapes, fn = qml.gradients.param_shift(tape)
-        all_res = fn(dev.batch_execute(tapes))
+        all_res = fn(dev.execute(tapes))
 
         assert len(all_res) == grad_transform_shots.num_copies
         assert isinstance(all_res, tuple)
@@ -2550,7 +2565,7 @@ def test_1_N(self, shot_vec, op_wire):
         tape.trainable_params = {0}
 
         tapes, fn = qml.gradients.param_shift(tape)
-        all_res = fn(dev.batch_execute(tapes))
+        all_res = fn(dev.execute(tapes))
 
         assert len(all_res) == grad_transform_shots.num_copies
         assert isinstance(all_res, tuple)
@@ -2584,7 +2599,7 @@ def test_N_1(self, shot_vec, meas, shape, op_wires):
         tape.trainable_params = {0, 1}
 
         tapes, fn = qml.gradients.param_shift(tape)
-        all_res = fn(dev.batch_execute(tapes))
+        all_res = fn(dev.execute(tapes))
 
         assert len(all_res) == grad_transform_shots.num_copies
         assert isinstance(all_res, tuple)
@@ -2624,7 +2639,7 @@ def test_N_N(self, shot_vec, op_wires):
         tape.trainable_params = {0, 1, 2, 3, 4}
 
         tapes, fn = qml.gradients.param_shift(tape)
-        all_res = fn(dev.batch_execute(tapes))
+        all_res = fn(dev.execute(tapes))
 
         assert len(all_res) == grad_transform_shots.num_copies
         assert isinstance(all_res, tuple)
diff --git a/tests/logging/test_logging_autograd.py b/tests/logging/test_logging_autograd.py
index e278e674c1b..6dcafffb439 100644
--- a/tests/logging/test_logging_autograd.py
+++ b/tests/logging/test_logging_autograd.py
@@ -21,7 +21,7 @@
 
 
 _grad_log_map = {
-    "adjoint": "gradient_fn=device, interface=autograd, grad_on_execution=best, gradient_kwargs={'use_device_state': True, 'method': 'adjoint_jacobian'}",
+    "adjoint": "gradient_fn=adjoint, interface=autograd, grad_on_execution=best, gradient_kwargs={}",
     "backprop": "gradient_fn=backprop, interface=autograd, grad_on_execution=best, gradient_kwargs={}",
     "parameter-shift": "gradient_fn=<pennylane.gradients.gradient_transform.gradient_transform object",
 }
@@ -78,7 +78,7 @@ def circuit(params):
 
             circuit(params)
 
-        assert len(caplog.records) == 5
+        assert len(caplog.records) == 3
 
         log_records_expected = [
             (
@@ -88,21 +88,20 @@ def circuit(params):
             (
                 "pennylane.interfaces.execution",
                 [
-                    "device=<DefaultQubitLegacy device (wires=2, shots=None)",
+                    "device=<default.qubit device (wires=2)",
                     "gradient_fn=None, interface=None",
                 ],
             ),
         ]
 
-        for idx, r in enumerate(caplog.records[0:2]):
-            assert log_records_expected[idx][0] in r.name
-            for msg in log_records_expected[idx][1]:
-                assert msg in r.getMessage()
+        for expected, actual in zip(log_records_expected, caplog.records[:2]):
+            assert expected[0] in actual.name
+            assert all(msg in actual.getMessage() for msg in expected[1])
 
     @pytest.mark.parametrize(
-        "diff_method", [("parameter-shift", 12), ("backprop", 5), ("adjoint", 9)]
+        "diff_method,num_records", [("parameter-shift", 7), ("backprop", 3), ("adjoint", 7)]
     )
-    def test_dq_qnode_execution_grad(self, caplog, diff_method):
+    def test_dq_qnode_execution_grad(self, caplog, diff_method, num_records):
         "Test logging of QNode with parameterised gradients"
 
         enable_and_configure_logging()
@@ -115,34 +114,33 @@ def test_dq_qnode_execution_grad(self, caplog, diff_method):
             dev = qml.device("default.qubit", wires=2)
             params = qml.numpy.array(0.1234)
 
-            @qml.qnode(dev, diff_method=diff_method[0])
+            @qml.qnode(dev, diff_method=diff_method)
             def circuit(params):
                 qml.RX(params, wires=0)
                 return qml.expval(qml.PauliZ(0))
 
             qml.grad(circuit)(params)
 
-        assert len(caplog.records) == diff_method[1]
+        assert len(caplog.records) == num_records
 
         log_records_expected = [
             (
                 "pennylane.qnode",
                 [
                     "Creating QNode(func=<function TestLogging.test_dq_qnode_execution_grad",
-                    "device=<DefaultQubitLegacy device (wires=2, shots=None)",
-                    f"interface=auto, diff_method={diff_method[0]}, expansion_strategy=gradient, max_expansion=10, grad_on_execution=best,",
+                    "device=<default.qubit device (wires=2)",
+                    f"interface=auto, diff_method={diff_method}, expansion_strategy=gradient, max_expansion=10, grad_on_execution=best,",
                 ],
             ),
             (
                 "pennylane.interfaces.execution",
                 [
                     "Entry with args=(tapes=(<QuantumScript: wires=[0], params=1>,)",
-                    _grad_log_map[diff_method[0]],
+                    _grad_log_map[diff_method],
                 ],
             ),
         ]
 
-        for idx, r in enumerate(caplog.records[0:2]):
-            assert log_records_expected[idx][0] in r.name
-            for msg in log_records_expected[idx][1]:
-                assert msg in r.getMessage()
+        for expected, actual in zip(log_records_expected, caplog.records[:2]):
+            assert expected[0] in actual.name
+            assert all(msg in actual.getMessage() for msg in expected[1])
diff --git a/tests/math/test_functions.py b/tests/math/test_functions.py
index 5dceb524562..ddf8bffcc28 100644
--- a/tests/math/test_functions.py
+++ b/tests/math/test_functions.py
@@ -2569,11 +2569,17 @@ class TestSize:
         ([[0], [1], [2], [3], [4], [5]], 6),
     ]
 
-    @pytest.mark.torch
+    @pytest.mark.parametrize(
+        "interface",
+        [
+            pytest.param("torch", marks=pytest.mark.torch),
+            pytest.param("tensorflow", marks=pytest.mark.tf),
+        ],
+    )
     @pytest.mark.parametrize(("array", "size"), array_and_size)
-    def test_size_torch(self, array, size):
-        """Test size function with the torch interface."""
-        r = fn.size(torch.tensor(array))
+    def test_size_torch_and_tf(self, array, size, interface):
+        """Test size function with the torch and tf interfaces."""
+        r = fn.size(fn.asarray(array, like=interface))
         assert r == size
 
 
diff --git a/tests/measurements/test_classical_shadow.py b/tests/measurements/test_classical_shadow.py
index d67967c2ff9..47fb06f2b02 100644
--- a/tests/measurements/test_classical_shadow.py
+++ b/tests/measurements/test_classical_shadow.py
@@ -340,7 +340,6 @@ def test_return_distribution(self, wires, interface, circuit_fn, basis_recipe):
         ratios2 = np.unique(bits2, return_counts=True)[1] / bits2.shape[0]
         assert np.allclose(ratios2, 1 / 2, atol=1e-1)
 
-    @pytest.mark.xfail(reason="until DQ2 port")
     @pytest.mark.parametrize("seed", seed_recipes_list)
     def test_shots_none_error(self, wires, seed):
         """Test that an error is raised when a device with shots=None is used
@@ -351,7 +350,6 @@ def test_shots_none_error(self, wires, seed):
         with pytest.raises(qml.DeviceError, match=msg):
             circuit()
 
-    @pytest.mark.xfail(reason="until DQ2 port")
     @pytest.mark.parametrize("shots", shots_list)
     def test_multi_measurement_error(self, wires, shots):
         """Test that an error is raised when classical shadows is returned
@@ -458,7 +456,6 @@ def test_measurement_process_copy(self):
         assert copied_res.k == res.k
         assert copied_res.seed == res.seed
 
-    @pytest.mark.xfail(reason="until DQ2 port")
     def test_shots_none_error(self):
         """Test that an error is raised when a device with shots=None is used
         to obtain classical shadows"""
@@ -469,7 +466,6 @@ def test_shots_none_error(self):
         with pytest.raises(qml.DeviceError, match=msg):
             _ = circuit(H, k=10)
 
-    @pytest.mark.xfail(reason="until DQ2 port")
     def test_multi_measurement_allowed(self):
         """Test that no error is raised when classical shadows is returned
         with other measurement processes"""
diff --git a/tests/measurements/test_measurements.py b/tests/measurements/test_measurements.py
index 04ec964338e..84cccd5514e 100644
--- a/tests/measurements/test_measurements.py
+++ b/tests/measurements/test_measurements.py
@@ -502,7 +502,6 @@ def circuit():
 
         assert qml.math.allequal(circuit(), [1000, 0])
 
-    @pytest.mark.xfail(reason="until DQ2 port")
     def test_sample_measurement_without_shots(self):
         """Test that executing a sampled measurement with ``shots=None`` raises an error."""
 
@@ -547,7 +546,6 @@ def circuit():
 
         assert circuit() == 1
 
-    @pytest.mark.xfail(reason="until DQ2 port")
     def test_state_measurement_with_shots(self):
         """Test that executing a state measurement with shots raises an error."""
 
@@ -572,13 +570,13 @@ def circuit():
 class TestMeasurementTransform:
     """Tests for the MeasurementTransform class."""
 
-    @pytest.mark.xfail(reason="until DQ2 port")
     def test_custom_measurement(self):
         """Test the execution of a custom measurement."""
 
         class CountTapesMP(MeasurementTransform, SampleMeasurement):
             def process(self, tape, device):
-                tapes, _, _ = device.preprocess(tape)
+                program, _ = device.preprocess()
+                tapes, _ = program([tape])
                 return len(tapes)
 
             def process_samples(self, samples, wire_order, shot_range=None, bin_size=None):
diff --git a/tests/measurements/test_mutual_info.py b/tests/measurements/test_mutual_info.py
index e266d1e9df4..0424b97fd8a 100644
--- a/tests/measurements/test_mutual_info.py
+++ b/tests/measurements/test_mutual_info.py
@@ -109,7 +109,6 @@ def circuit():
         res = circuit()
         assert np.allclose(res, expected, atol=1e-6)
 
-    @pytest.mark.xfail(reason="until DQ2 port")
     @pytest.mark.parametrize("shots", [1000, [1, 10, 10, 1000]])
     def test_finite_shots_error(self, shots):
         """Test an error is raised when using shot vectors with mutual_info."""
@@ -491,7 +490,6 @@ def circuit(params):
         with pytest.raises(qml.QuantumFunctionError, match=msg):
             circuit(params)
 
-    @pytest.mark.xfail(reason="until DQ2 port")
     @pytest.mark.all_interfaces
     @pytest.mark.parametrize("interface", ["autograd", "jax", "tensorflow", "torch"])
     @pytest.mark.parametrize("params", [np.array([0.0, 0.0]), np.array([0.3, 0.4])])
diff --git a/tests/measurements/test_probs.py b/tests/measurements/test_probs.py
index c1a55296c6f..3cf2c177cc6 100644
--- a/tests/measurements/test_probs.py
+++ b/tests/measurements/test_probs.py
@@ -644,7 +644,6 @@ def test_estimate_probability_with_binsize_with_broadcasting(self, wires, expect
 
         assert np.allclose(res, expected)
 
-    @pytest.mark.xfail(reason="until DQ2 port")
     def test_non_commuting_probs_does_not_raises_error(self):
         """Tests that non-commuting probs with expval does not raise an error."""
         dev = qml.device("default.qubit", wires=5)
diff --git a/tests/measurements/test_sample.py b/tests/measurements/test_sample.py
index 10a5fb070d9..0aabe2675d3 100644
--- a/tests/measurements/test_sample.py
+++ b/tests/measurements/test_sample.py
@@ -46,7 +46,6 @@ def circuit():
             (n_sample,) if not n_sample == 1 else ()
         )
 
-    @pytest.mark.xfail(reason="until DQ2 port")
     def test_sample_combination(self):
         """Test the output of combining expval, var and sample"""
         n_sample = 10
diff --git a/tests/measurements/test_state.py b/tests/measurements/test_state.py
index a24c648052e..ab306b5b01e 100644
--- a/tests/measurements/test_state.py
+++ b/tests/measurements/test_state.py
@@ -250,22 +250,19 @@ def test_process_state_matrix_from_matrix(self, mat, wires):
 class TestState:
     """Tests for the state function"""
 
-    @pytest.mark.xfail(reason="until DQ2 port")
     @pytest.mark.parametrize("wires", range(2, 5))
-    @pytest.mark.parametrize("op,dtype", [(qml.PauliX, np.float64), (qml.PauliY, np.complex128)])
-    def test_state_shape_and_dtype(self, op, dtype, wires):
+    def test_state_shape_and_dtype(self, wires):
         """Test that the state is of correct size and dtype for a trivial circuit"""
 
         dev = qml.device("default.qubit", wires=wires)
 
         @qml.qnode(dev)
         def func():
-            op(0)
             return state()
 
         state_val = func()
         assert state_val.shape == (2**wires,)
-        assert state_val.dtype == dtype
+        assert state_val.dtype == np.complex128
 
     def test_return_type_is_state(self):
         """Test that the return type of the observable is State"""
@@ -300,7 +297,6 @@ def func():
         assert np.allclose(state_val[0], 1 / np.sqrt(2))
         assert np.allclose(state_val[-1], 1 / np.sqrt(2))
 
-    @pytest.mark.xfail(reason="until DQ2 port")
     def test_return_with_other_types_works(self):
         """Test that no exception is raised when a state is returned along with another return
         type"""
@@ -317,7 +313,6 @@ def func():
         assert np.allclose(res[0], np.array([1, 0, 1, 0]) / np.sqrt(2))
         assert np.isclose(res[1], 1)
 
-    @pytest.mark.xfail(reason="until DQ2 port")
     @pytest.mark.parametrize("wires", range(2, 5))
     def test_state_equal_to_expected_state(self, wires):
         """Test that the returned state is equal to the expected state for a template circuit"""
@@ -334,14 +329,14 @@ def func():
             return state()
 
         state_val = func()
-        scripts, _, _ = dev.preprocess(func.tape)
+        program, _ = dev.preprocess()
+        scripts, _ = program([func.tape])
         assert len(scripts) == 1
         expected_state, _ = qml.devices.qubit.get_final_state(scripts[0])
         assert np.allclose(state_val, expected_state.flatten())
 
     @pytest.mark.tf
-    @pytest.mark.parametrize("op", [qml.PauliX, qml.PauliY])
-    def test_interface_tf(self, op):
+    def test_interface_tf(self):
         """Test that the state correctly outputs in the tensorflow interface"""
         import tensorflow as tf
 
@@ -349,8 +344,6 @@ def test_interface_tf(self, op):
 
         @qml.qnode(dev, interface="tf")
         def func():
-            op(0)
-            op(0)
             for i in range(4):
                 qml.Hadamard(i)
             return state()
@@ -359,14 +352,12 @@ def func():
         state_val = func()
 
         assert isinstance(state_val, tf.Tensor)
-        assert state_val.dtype == tf.complex128 if op is qml.PauliY else tf.float64
+        assert state_val.dtype == tf.complex128
         assert np.allclose(state_expected, state_val.numpy())
         assert state_val.shape == (16,)
 
-    @pytest.mark.xfail(reason="until DQ2 port")
     @pytest.mark.torch
-    @pytest.mark.parametrize("op", [qml.PauliX, qml.PauliY])
-    def test_interface_torch(self, op):
+    def test_interface_torch(self):
         """Test that the state correctly outputs in the torch interface"""
         import torch
 
@@ -374,13 +365,11 @@ def test_interface_torch(self, op):
 
         @qml.qnode(dev, interface="torch")
         def func():
-            op(0)
-            op(0)
             for i in range(4):
                 qml.Hadamard(i)
             return state()
 
-        dtype = torch.complex128 if op is qml.PauliY else torch.float64
+        dtype = torch.complex128
         state_expected = 0.25 * torch.ones(16, dtype=dtype)
         state_val = func()
 
@@ -596,11 +585,9 @@ class TestDensityMatrix:
 
     # pylint: disable=too-many-public-methods
 
-    @pytest.mark.xfail(reason="until DQ2 port")
     @pytest.mark.parametrize("wires", range(2, 5))
     @pytest.mark.parametrize("dev_name", ["default.qubit", "default.mixed"])
-    @pytest.mark.parametrize("op,dtype", [(qml.PauliX, np.float64), (qml.PauliY, np.complex128)])
-    def test_density_matrix_shape_and_dtype(self, dev_name, op, dtype, wires):
+    def test_density_matrix_shape_and_dtype(self, dev_name, wires):
         """Test that the density matrix is of correct size and dtype for a
         trivial circuit"""
 
@@ -608,13 +595,12 @@ def test_density_matrix_shape_and_dtype(self, dev_name, op, dtype, wires):
 
         @qml.qnode(dev)
         def circuit():
-            op(0)
             return density_matrix([0])
 
         state_val = circuit()
 
         assert state_val.shape == (2, 2)
-        assert state_val.dtype == dtype if dev_name == "default.qubit" else np.complex128
+        assert state_val.dtype == np.complex128
 
     @pytest.mark.parametrize("dev_name", ["default.qubit", "default.mixed"])
     def test_return_type_is_state(self, dev_name):
@@ -681,7 +667,6 @@ def func():
 
         assert np.allclose(expected, density_mat)
 
-    @pytest.mark.xfail(reason="until DQ2 port")
     @pytest.mark.tf
     @pytest.mark.parametrize("diff_method", [None, "backprop"])
     def test_correct_density_matrix_tf_default_qubit(self, diff_method):
@@ -720,7 +705,6 @@ def func():
 
         assert np.allclose(expected, density_first)
 
-    @pytest.mark.xfail(reason="until DQ2 port")
     def test_correct_density_matrix_product_state_first_default_qubit(self):
         """Test that the correct density matrix is returned when
         tracing out a product state"""
@@ -758,7 +742,6 @@ def func():
         expected = np.array([[0.5 + 0.0j, 0.5 + 0.0j], [0.5 + 0.0j, 0.5 + 0.0j]])
         assert np.allclose(expected, density_second)
 
-    @pytest.mark.xfail(reason="until DQ2 port")
     def test_correct_density_matrix_product_state_second_default_qubit(self):
         """Test that the correct density matrix is returned when
         tracing out a product state"""
@@ -800,7 +783,6 @@ def func():
 
         assert np.allclose(expected, density_both)
 
-    @pytest.mark.xfail(reason="until DQ2 port")
     @pytest.mark.parametrize("return_wire_order", ([0, 1], [1, 0]))
     def test_correct_density_matrix_product_state_both_default_qubit(self, return_wire_order):
         """Test that the correct density matrix is returned
@@ -851,7 +833,6 @@ def func():
         )
         assert np.allclose(expected, density_full)
 
-    @pytest.mark.xfail(reason="until DQ2 port")
     def test_correct_density_matrix_three_wires_first_two_default_qubit(self):
         """Test that the correct density matrix is returned for an example with three wires,
         and tracing out the third wire."""
@@ -938,7 +919,6 @@ def func():
         expected = np.outer(exp_statevector.conj(), exp_statevector)
         assert np.allclose(expected, density_full)
 
-    @pytest.mark.xfail(reason="until DQ2 port")
     @pytest.mark.parametrize(
         "return_wire_order", ([0], [1], [2], [0, 1], [1, 0], [0, 2], [2, 0], [1, 2, 0], [2, 1, 0])
     )
@@ -1016,7 +996,6 @@ def func():
 
         assert np.allclose(expected, density)
 
-    @pytest.mark.xfail(reason="until DQ2 port")
     def test_correct_density_matrix_all_wires_default_qubit(self):
         """Test that the correct density matrix is returned when all wires are given"""
 
@@ -1044,7 +1023,6 @@ def func():
             qml.density_matrix(wires=[0, 1]).process_state(state=dev_state, wire_order=dev.wires),
         )
 
-    @pytest.mark.xfail(reason="until DQ2 port")
     def test_return_with_other_types_works(self):
         """Test that no exception is raised when a state is returned along with another return
         type"""
diff --git a/tests/measurements/test_vn_entropy.py b/tests/measurements/test_vn_entropy.py
index 432447227fe..b8dc7830ccf 100644
--- a/tests/measurements/test_vn_entropy.py
+++ b/tests/measurements/test_vn_entropy.py
@@ -136,7 +136,6 @@ class TestIntegration:
 
     diff_methods = ["backprop", "finite-diff"]
 
-    @pytest.mark.xfail(reason="until DQ2 port")
     @pytest.mark.parametrize("shots", [1000, [1, 10, 10, 1000]])
     def test_finite_shots_error(self, shots):
         """Test an error is raised when using shot vectors with vn_entropy."""
@@ -392,7 +391,6 @@ def circuit_entropy(x):
 
         assert qml.math.allclose(grad_entropy, grad_expected_entropy, rtol=1e-04, atol=1e-05)
 
-    @pytest.mark.xfail(reason="until DQ2 port")
     def test_qnode_entropy_custom_wires(self):
         """Test that entropy can be returned with custom wires."""
 
diff --git a/tests/ops/functions/test_map_wires.py b/tests/ops/functions/test_map_wires.py
index 90115b50b7e..785364faf2e 100644
--- a/tests/ops/functions/test_map_wires.py
+++ b/tests/ops/functions/test_map_wires.py
@@ -54,7 +54,7 @@ def test_map_wires_with_operator(self):
         op = build_op()
 
         m_op = qml.map_wires(op, wire_map=wire_map)
-        assert isinstance(m_op, qml.ops.Prod)
+        assert isinstance(m_op, qml.ops.Prod)  # pylint:disable=no-member
         assert m_op.data == mapped_op.data
         assert m_op.wires == mapped_op.wires
         assert m_op.arithmetic_depth == mapped_op.arithmetic_depth
@@ -112,13 +112,13 @@ def test_map_wires_tape(self, shots):
         assert len(s_tape) == 1
         assert s_tape.trainable_params == [0, 2]
         s_op = s_tape[0]
-        assert isinstance(s_op, qml.ops.Prod)
+        assert isinstance(s_op, qml.ops.Prod)  # pylint:disable=no-member
         assert s_op.data == mapped_op.data
         assert s_op.wires == mapped_op.wires
         assert s_op.arithmetic_depth == mapped_op.arithmetic_depth
         assert tape.shots == s_tape.shots
 
-    @pytest.mark.parametrize("shots", [None, 100])
+    @pytest.mark.parametrize("shots", [None, 5000])
     def test_execute_mapped_tape(self, shots):
         """Test the execution of a mapped tape."""
         dev = qml.device("default.qubit", wires=5)
@@ -134,13 +134,10 @@ def test_execute_mapped_tape(self, shots):
         assert m_tape._qfunc_output is tape._qfunc_output  # pylint: disable=protected-access
         m_op = m_tape.operations[0]
         m_obs = m_tape.observables[0]
-        assert isinstance(m_op, qml.ops.Prod)
-        assert m_op.data == mapped_op.data
-        assert m_op.wires == mapped_op.wires
-        assert m_op.arithmetic_depth == mapped_op.arithmetic_depth
+        assert qml.equal(m_op, mapped_op)
         assert tape.shots == m_tape.shots
         assert m_obs.wires == Wires(wire_map[1])
-        assert qml.math.allclose(dev.execute(tape), dev.execute(m_tape))
+        assert qml.math.allclose(dev.execute(tape), dev.execute(m_tape), atol=0.05)
 
 
 class TestMapWiresQNodes:
@@ -153,23 +150,18 @@ def test_map_wires_qnode(self):
         @qml.qnode(dev)
         def qnode():
             build_op()
-            return qml.expval(op=build_op())
+            return qml.expval(qml.prod(qml.PauliX(0), qml.PauliY(1), qml.PauliZ(2)))
 
         # TODO: Use qml.equal when supported
+        mapped_obs = qml.prod(qml.PauliX(4), qml.PauliY(3), qml.PauliZ(2))
 
         m_qnode = qml.map_wires(qnode, wire_map=wire_map)
         assert m_qnode() == qnode()
         assert len(m_qnode.tape) == 2
         m_op = m_qnode.tape.operations[0]
         m_obs = m_qnode.tape.observables[0]
-        assert isinstance(m_op, qml.ops.Prod)
-        assert m_op.data == mapped_op.data
-        assert m_op.wires == mapped_op.wires
-        assert m_op.arithmetic_depth == mapped_op.arithmetic_depth
-        assert isinstance(m_obs, qml.ops.Prod)
-        assert m_obs.data == mapped_op.data
-        assert m_obs.wires == mapped_op.wires
-        assert m_obs.arithmetic_depth == mapped_op.arithmetic_depth
+        assert qml.equal(m_op, mapped_op)
+        assert qml.equal(m_obs, mapped_obs)
 
 
 class TestMapWiresCallables:
diff --git a/tests/ops/op_math/test_exp.py b/tests/ops/op_math/test_exp.py
index 2e2848b8987..9e24dcebfca 100644
--- a/tests/ops/op_math/test_exp.py
+++ b/tests/ops/op_math/test_exp.py
@@ -647,7 +647,7 @@ def test_repr_tensor(self):
     def test_repr_deep_operator(self):
         """Test the __repr__ method when the base is any operator with arithmetic depth > 0."""
         base = qml.S(0) @ qml.PauliX(0)
-        op = qml.ops.Exp(base, 3)
+        op = qml.ops.Exp(base, 3)  # pylint:disable=no-member
 
         assert repr(op) == "Exp(3 S(wires=[0]) @ PauliX(wires=[0]))"
 
@@ -797,7 +797,7 @@ def circuit(phi):
         res = circuit(phi)
         assert qml.math.allclose(res, torch.cos(phi))
 
-        res.backward()
+        res.backward()  # pylint:disable=no-member
         assert qml.math.allclose(phi.grad, -torch.sin(phi))
 
     @pytest.mark.autograd
@@ -872,7 +872,7 @@ def circuit(x):
         expected = 0.5 * (torch.exp(x) + torch.exp(-x))
         assert qml.math.allclose(res, expected)
 
-        res.backward()
+        res.backward()  # pylint:disable=no-member
         expected_grad = 0.5 * (torch.exp(x) - torch.exp(-x))
         assert qml.math.allclose(x.grad, expected_grad)
 
@@ -901,9 +901,10 @@ def circuit(x):
     @pytest.mark.tf
     def test_tf_measurement(self):
         """Test Exp in a measurement with gradient and tensorflow."""
+        # pylint:disable=invalid-unary-operand-type
         import tensorflow as tf
 
-        x = tf.Variable(2.0)
+        x = tf.Variable(2.0, dtype=tf.float64)
 
         @qml.qnode(qml.device("default.qubit", wires=1))
         def circuit(x):
@@ -1011,7 +1012,7 @@ def test_parameter_frequency_with_parameters_in_base_operator(self):
         op2 = Evolution(base_op, 1)
 
         with pytest.raises(ParameterFrequenciesUndefinedError):
-            op1.parameter_frequencies()
+            _ = op1.parameter_frequencies
 
         assert op2.parameter_frequencies == [(4.0,)]
 
diff --git a/tests/ops/op_math/test_prod.py b/tests/ops/op_math/test_prod.py
index 6db76b23dab..07a5bd4644a 100644
--- a/tests/ops/op_math/test_prod.py
+++ b/tests/ops/op_math/test_prod.py
@@ -758,6 +758,7 @@ def test_is_hermitian(self, ops_lst, hermitian_status):
     @pytest.mark.tf
     def test_is_hermitian_tf(self):
         """Test that is_hermitian works when a tf type scalar is provided."""
+        # pylint:disable=invalid-unary-operand-type
         import tensorflow as tf
 
         theta = tf.Variable(1.23)
@@ -965,7 +966,7 @@ def test_simplify_method_product_of_sums(self):
             qml.RX(1, 0) @ qml.RX(1, 1),
         )
         simplified_op = prod_op.simplify()
-        assert isinstance(simplified_op, qml.ops.Sum)
+        assert isinstance(simplified_op, qml.ops.Sum)  # pylint:disable=no-member
         for s1, s2 in zip(final_op.operands, simplified_op.operands):
             assert s1.name == s2.name
             assert s1.wires == s2.wires
@@ -1016,7 +1017,7 @@ def test_simplify_method_with_nested_ops(self):
             qml.s_prod(5, qml.PauliX(1)),
         )
         simplified_op = prod_op.simplify()
-        assert isinstance(simplified_op, qml.ops.Sum)
+        assert isinstance(simplified_op, qml.ops.Sum)  # pylint:disable=no-member
         for s1, s2 in zip(final_op.operands, simplified_op.operands):
             assert s1.name == s2.name
             assert s1.wires.toset() == s2.wires.toset()
@@ -1050,7 +1051,7 @@ def test_simplify_method_with_pauli_words(self):
 
         # TODO: Use qml.equal when supported for nested operators
 
-        assert isinstance(simplified_op, qml.ops.Sum)
+        assert isinstance(simplified_op, qml.ops.Sum)  # pylint:disable=no-member
         for s1, s2 in zip(final_op.operands, simplified_op.operands):
             assert repr(s1) == repr(s2)
             assert s1.name == s2.name
@@ -1100,7 +1101,7 @@ def test_grouping_with_product_of_sum(self):
         final_op = qml.sum(qml.s_prod(1j, qml.PauliX(0)), qml.s_prod(-1, qml.PauliZ(0)))
         simplified_op = prod_op.simplify()
 
-        assert isinstance(simplified_op, qml.ops.Sum)
+        assert isinstance(simplified_op, qml.ops.Sum)  # pylint:disable=no-member
         for s1, s2 in zip(final_op.operands, simplified_op.operands):
             assert s1.name == s2.name
             assert s1.wires == s2.wires
@@ -1116,7 +1117,7 @@ def test_grouping_with_product_of_sums(self):
             qml.S(wires=[1]),
         )
         simplified_op = prod_op.simplify()
-        assert isinstance(simplified_op, qml.ops.Sum)
+        assert isinstance(simplified_op, qml.ops.Sum)  # pylint:disable=no-member
         for s1, s2 in zip(final_op.operands, simplified_op.operands):
             assert s1.name == s2.name
             assert s1.wires == s2.wires
@@ -1128,7 +1129,7 @@ def test_grouping_with_barriers(self):
         prod_op = qml.prod(qml.S(0), qml.Barrier(0), qml.S(0)).simplify()
         simplified_op = prod_op.simplify()
         assert isinstance(simplified_op, Prod)
-        for s1, s2 in zip(prod_op.operands, simplified_op.operands):
+        for s1, s2 in zip(prod_op.operands, simplified_op.operands):  # pylint:disable=no-member
             assert s1.name == s2.name
             assert s1.wires == s2.wires
             assert s1.data == s2.data
@@ -1321,7 +1322,7 @@ def circuit(weights):
         true_grad = -qnp.sqrt(2) * qnp.cos(weights[0] / 2) * qnp.sin(weights[0] / 2)
         assert qnp.allclose(grad, true_grad)
 
-    def test_non_hermitian_op_in_measurement_process(self):
+    def test_non_hermitian_obs_not_supported(self):
         """Test that non-hermitian ops in a measurement process will raise a warning."""
         wires = [0, 1]
         dev = qml.device("default.qubit", wires=wires)
@@ -1332,7 +1333,7 @@ def my_circ():
             qml.PauliX(0)
             return qml.expval(prod_op)
 
-        with pytest.warns(UserWarning, match="Prod might not be hermitian."):
+        with pytest.raises(NotImplementedError):
             my_circ()
 
     def test_operation_integration(self):
@@ -1387,11 +1388,14 @@ def test_params_can_be_considered_trainable(self):
         dev = qml.device("default.qubit", wires=2)
 
         @qml.qnode(dev)
-        def circuit():
-            return qml.expval(qml.prod(qml.RX(1.1, 0), qml.RY(qnp.array(2.2), 1)))
+        def circuit(x, U):
+            qml.RX(x, 0)
+            return qml.expval(qml.prod(qml.Hermitian(U, 0), qml.PauliX(1)))
+
+        x = qnp.array(0.1, requires_grad=False)
+        U = qnp.array([[1.0, 0.0], [0.0, -1.0]], requires_grad=True)
 
-        with pytest.warns(UserWarning):
-            circuit()
+        circuit(x, U)
         assert circuit.tape.trainable_params == [1]
 
 
diff --git a/tests/ops/op_math/test_sprod.py b/tests/ops/op_math/test_sprod.py
index bc25513af3e..014cde358da 100644
--- a/tests/ops/op_math/test_sprod.py
+++ b/tests/ops/op_math/test_sprod.py
@@ -790,7 +790,7 @@ def test_simplify_method(self):
         sprod_op = SProd(
             2, SProd(2, qml.RZ(1.32, wires=0)) + qml.Identity(wires=0) + qml.RX(1.9, wires=1)
         )
-        final_op = qml.ops.Sum(
+        final_op = qml.ops.Sum(  # pylint:disable=no-member
             SProd(4, qml.RZ(1.32, wires=0)),
             SProd(2, qml.Identity(wires=0)),
             SProd(2, qml.RX(1.9, wires=1)),
@@ -799,7 +799,7 @@ def test_simplify_method(self):
 
         # TODO: Use qml.equal when supported for nested operators
 
-        assert isinstance(simplified_op, qml.ops.Sum)
+        assert isinstance(simplified_op, qml.ops.Sum)  # pylint:disable=no-member
         for s1, s2 in zip(final_op.operands, simplified_op.operands):
             assert isinstance(s2, SProd)
             assert s1.name == s2.name
@@ -837,7 +837,7 @@ def test_simplify_with_sum_operator(self):
         final_op = s_prod(0 - 6j, qml.PauliX(0))
         simplified_op = sprod_op.simplify()
 
-        assert isinstance(simplified_op, qml.ops.SProd)
+        assert isinstance(simplified_op, qml.ops.SProd)  # pylint:disable=no-member
         assert simplified_op.name == final_op.name
         assert repr(simplified_op) == repr(final_op)
         assert simplified_op.wires == final_op.wires
@@ -1045,7 +1045,7 @@ def circuit(weights):
         true_grad = 100 * -qnp.sqrt(2) * qnp.cos(weights[0] / 2) * qnp.sin(weights[0] / 2)
         assert qnp.allclose(grad, true_grad)
 
-    def test_non_hermitian_op_in_measurement_process(self):
+    def test_non_hermitian_obs_not_supported(self):
         """Test that non-hermitian ops in a measurement process will raise a warning."""
         wires = [0, 1]
         dev = qml.device("default.qubit", wires=wires)
@@ -1056,7 +1056,7 @@ def my_circ():
             qml.PauliX(0)
             return qml.expval(sprod_op)
 
-        with pytest.warns(UserWarning, match="SProd might not be hermitian."):
+        with pytest.raises(NotImplementedError):
             my_circ()
 
     @pytest.mark.torch
@@ -1106,7 +1106,7 @@ def test_tensorflow_qnode(self, diff_method):
         def circuit(s):
             return qml.expval(qml.s_prod(s, qml.PauliZ(0)))
 
-        res = circuit(tf.Variable(2))
+        res = circuit(tf.Variable(2, dtype=tf.float64))
 
         assert qml.math.allclose(res, 2)
 
diff --git a/tests/ops/qubit/test_hamiltonian.py b/tests/ops/qubit/test_hamiltonian.py
index fe1961f71a6..3a9ede92715 100644
--- a/tests/ops/qubit/test_hamiltonian.py
+++ b/tests/ops/qubit/test_hamiltonian.py
@@ -1704,7 +1704,7 @@ def combine(coeffs, param):
     def test_nontrainable_coeffs_paramshift(self):
         """Test the parameter-shift method if the coefficients are explicitly set non-trainable
         by not passing them to the qnode."""
-        coeffs = pnp.array([-0.05, 0.17], requires_grad=False)
+        coeffs = np.array([-0.05, 0.17])
         param = pnp.array(1.7, requires_grad=True)
 
         # differentiating a circuit with measurement expval(H)
@@ -1898,7 +1898,7 @@ def circuit(coeffs, param):
             )
 
         res = circuit(coeffs, param)
-        res.backward()
+        res.backward()  # pylint:disable=no-member
         grad = (coeffs.grad, param.grad)
 
         # differentiating a cost that combines circuits with
@@ -1940,7 +1940,7 @@ def circuit(coeffs, param):
             )
 
         res = circuit(coeffs, param)
-        res.backward()
+        res.backward()  # pylint:disable=no-member
 
         # differentiating a cost that combines circuits with
         # measurements expval(Pauli)
@@ -2069,7 +2069,7 @@ def circuit(coeffs, param):
 
         grad_fn = qml.grad(circuit)
         with pytest.raises(
-            qml.QuantumFunctionError,
-            match="Adjoint differentiation method does not support Hamiltonian observables",
+            qml.DeviceError,
+            match="Adjoint differentiation method does not support observable Hamiltonian",
         ):
             grad_fn(coeffs, param)
diff --git a/tests/ops/qubit/test_parametric_ops.py b/tests/ops/qubit/test_parametric_ops.py
index ba71bd3208e..7e874328911 100644
--- a/tests/ops/qubit/test_parametric_ops.py
+++ b/tests/ops/qubit/test_parametric_ops.py
@@ -2554,6 +2554,7 @@ def test_globalphase_autograd_grad(self, tol, dev_name, diff_method):
 
         @qml.qnode(dev, diff_method=diff_method)
         def circuit(x):
+            qml.Identity(0)
             qml.Hadamard(1)
             qml.ctrl(qml.GlobalPhase(x), 1)
             qml.Hadamard(1)
@@ -2899,6 +2900,7 @@ def test_globalphase_tf_grad(self, tol, dev_name, diff_method):
 
         @qml.qnode(dev, diff_method=diff_method)
         def circuit(x):
+            qml.Identity(0)
             qml.Hadamard(1)
             qml.ctrl(qml.GlobalPhase(x), 1)
             qml.Hadamard(1)
@@ -3055,6 +3057,7 @@ def test_globalphase_jax_grad(self, tol, dev_name, diff_method):
 
         @qml.qnode(dev, diff_method=diff_method)
         def circuit(x):
+            qml.Identity(0)
             qml.Hadamard(1)
             qml.ctrl(qml.GlobalPhase(x), 1)
             qml.Hadamard(1)
@@ -3078,6 +3081,7 @@ def test_globalphase_torch_grad(self, tol, dev_name, diff_method):
 
         @qml.qnode(dev, diff_method=diff_method)
         def circuit(x):
+            qml.Identity(0)
             qml.Hadamard(1)
             qml.ctrl(qml.GlobalPhase(x), 1)
             qml.Hadamard(1)
diff --git a/tests/ops/qubit/test_sparse.py b/tests/ops/qubit/test_sparse.py
index 30039971588..883c96cec32 100644
--- a/tests/ops/qubit/test_sparse.py
+++ b/tests/ops/qubit/test_sparse.py
@@ -193,8 +193,10 @@ def test_sparse_hamiltonian_expval(self, qubits, operations, hamiltonian, expect
         hamiltonian = csr_matrix(hamiltonian)
 
         dev = qml.device("default.qubit", wires=qubits, shots=None)
-        dev.apply(operations)
-        expval = dev.expval(qml.SparseHamiltonian(hamiltonian, range(qubits)))[0]
+        qs = qml.tape.QuantumScript(
+            operations, [qml.expval((qml.SparseHamiltonian(hamiltonian, range(qubits))))]
+        )
+        expval = dev.execute(qs)
 
         assert np.allclose(expval, expected_output, atol=tol, rtol=0)
 
@@ -203,7 +205,10 @@ def test_sparse_expval_error(self):
         shots is requested."""
         hamiltonian = csr_matrix(np.array([[1.0, 0.0], [0.0, 1.0]]))
 
-        dev = qml.device("default.qubit", wires=1, shots=1)
+        dev = qml.device("default.qubit", wires=1)
+        qs = qml.tape.QuantumScript(
+            measurements=[qml.expval(qml.SparseHamiltonian(hamiltonian, [0]))], shots=1
+        )
 
-        with pytest.raises(AssertionError, match="SparseHamiltonian must be used with shots=None"):
-            dev.expval(qml.SparseHamiltonian(hamiltonian, [0]))
+        with pytest.raises(qml.operation.DiagGatesUndefinedError):
+            dev.execute(qs)
diff --git a/tests/pulse/test_hardware_hamiltonian.py b/tests/pulse/test_hardware_hamiltonian.py
index 566c3c22e13..09933ba2ccf 100644
--- a/tests/pulse/test_hardware_hamiltonian.py
+++ b/tests/pulse/test_hardware_hamiltonian.py
@@ -662,7 +662,7 @@ def qnode_jit(params):
         res_jit = qnode_jit(params)
 
         assert isinstance(res, jax.Array)
-        assert res == res_jit
+        assert qml.math.allclose(res, res_jit)
 
     @pytest.mark.jax
     def test_jitted_qnode_multidrive(self):
@@ -747,4 +747,4 @@ def qnode_jit(params):
         res_jit = qnode_jit(params)
 
         assert isinstance(res, jax.Array)
-        assert res == res_jit
+        assert qml.math.allclose(res, res_jit)
diff --git a/tests/qchem/test_dipole.py b/tests/qchem/test_dipole.py
index e19db8fe3f3..23554cf3518 100644
--- a/tests/qchem/test_dipole.py
+++ b/tests/qchem/test_dipole.py
@@ -285,7 +285,9 @@ def test_gradient_expvalD():
 
     mol = qchem.Molecule(symbols, geometry, charge=1, alpha=alpha)
     args = [mol.alpha]
-    dev = qml.device("default.qubit", wires=6)
+    # TODO: `d_qubit[0]` has coeff dtype complex, but is actually a real-valued Hamiltonian
+    # default.qubit.legacy casts Hamiltonian expectations to real, but default.qubit does not
+    dev = qml.device("default.qubit.legacy", wires=6)
 
     def dipole(mol):
         @qml.qnode(dev)
diff --git a/tests/qchem/test_hamiltonians.py b/tests/qchem/test_hamiltonians.py
index c6c9f3c091a..78caa5bf68a 100644
--- a/tests/qchem/test_hamiltonians.py
+++ b/tests/qchem/test_hamiltonians.py
@@ -284,7 +284,7 @@ def test_gradient_expvalH():
 
     mol = qchem.Molecule(symbols, geometry, alpha=alpha)
     args = [alpha]
-    dev = qml.device("default.qubit", wires=4)
+    dev = qml.device("default.qubit.legacy", wires=4)
 
     def energy(mol):
         @qml.qnode(dev)
@@ -336,7 +336,7 @@ def test_gradient_expvalH(self):
 
         mol = qchem.Molecule(symbols, geometry, alpha=alpha)
         args = [jax.numpy.array(alpha)]
-        dev = qml.device("default.qubit", wires=4)
+        dev = qml.device("default.qubit.legacy", wires=4)
 
         def energy(mol):
             @qml.qnode(dev, interface="jax")
diff --git a/tests/qinfo/test_fisher.py b/tests/qinfo/test_fisher.py
index 26a5456631d..7ccae1e9e29 100644
--- a/tests/qinfo/test_fisher.py
+++ b/tests/qinfo/test_fisher.py
@@ -157,7 +157,7 @@ def qfunc(params):
             qml.RX(params[0], wires=0)
             qml.RX(params[1], wires=0)
             qml.CNOT(wires=(0, 1))
-            return qml.state()
+            return qml.probs()
 
         params = pnp.random.random(2)
 
diff --git a/tests/qnn/test_keras.py b/tests/qnn/test_keras.py
index f8261f3af8f..34566606cf9 100644
--- a/tests/qnn/test_keras.py
+++ b/tests/qnn/test_keras.py
@@ -915,7 +915,7 @@ def circuit(inputs, w1, w2):
 
     assert info["num_observables"] == 2
     assert info["num_diagonalizing_gates"] == 0
-    assert info["num_device_wires"] == 3
+    assert info["num_device_wires"] == 2
     assert info["num_trainable_params"] == 2
     assert info["interface"] == "tf"
-    assert info["device_name"] == "default.qubit.tf"
+    assert info["device_name"] == "default.qubit"
diff --git a/tests/qnn/test_qnn_torch.py b/tests/qnn/test_qnn_torch.py
index d1a4b385800..ba2701063a9 100644
--- a/tests/qnn/test_qnn_torch.py
+++ b/tests/qnn/test_qnn_torch.py
@@ -666,8 +666,6 @@ def compute_matrix(
             z = params[0]
             return np.diag([z, z])
 
-    device.operations.add("DummyOp")
-
     @qml.qnode(device=device, interface="torch", diff_method="parameter-shift")
     def circ(inputs, w0):  # pylint: disable=unused-argument
         DummyOp(inputs[0], wires=0)
@@ -849,7 +847,7 @@ def circuit(inputs, w1, w2):
 
     assert info["num_observables"] == 2
     assert info["num_diagonalizing_gates"] == 0
-    assert info["num_device_wires"] == 3
+    assert info["num_device_wires"] == 2
     assert info["num_trainable_params"] == 0
     assert info["interface"] == "torch"
-    assert info["device_name"] == "default.qubit.torch"
+    assert info["device_name"] == "default.qubit"
diff --git a/tests/shadow/test_shadow_entropies.py b/tests/shadow/test_shadow_entropies.py
index 5987723c418..1de2af20454 100644
--- a/tests/shadow/test_shadow_entropies.py
+++ b/tests/shadow/test_shadow_entropies.py
@@ -77,9 +77,10 @@ def test_non_constant_distribution(
         """Test entropies match roughly with exact solution for a non-constant distribution using other PennyLane functionalities"""
         n_wires = 4
         # exact solution
+        dev_exact = qml.device("default.qubit", wires=range(n_wires), shots=None)
         dev = qml.device("default.qubit", wires=range(n_wires), shots=100000)
 
-        @qml.qnode(dev)
+        @qml.qnode(dev_exact)
         def qnode_exact(x):
             for i in range(n_wires):
                 qml.RY(x[i], wires=i)
diff --git a/tests/shadow/test_shadow_transforms.py b/tests/shadow/test_shadow_transforms.py
index c924a2dc65d..ef91d5956cf 100644
--- a/tests/shadow/test_shadow_transforms.py
+++ b/tests/shadow/test_shadow_transforms.py
@@ -197,9 +197,8 @@ def circuit_shadow():
             qml.CNOT(wires=[0, 1])
             return qml.classical_shadow(wires=[0, 1]), qml.expval(qml.PauliZ(0))
 
-        msg = "Classical shadows cannot be returned in combination with other return types"
-        with pytest.raises(qml.QuantumFunctionError, match=msg):
-            circuit_shadow()
+        res = circuit_shadow()
+        assert isinstance(res, tuple) and len(res) == 2
 
         @qml.qnode(dev)
         def circuit_expval():
@@ -207,8 +206,8 @@ def circuit_expval():
             qml.CNOT(wires=[0, 1])
             return qml.shadow_expval(qml.PauliZ(0)), qml.expval(qml.PauliZ(0))
 
-        with pytest.raises(qml.QuantumFunctionError, match=msg):
-            circuit_expval()
+        res = circuit_expval()
+        assert isinstance(res, tuple) and len(res) == 2
 
 
 @pytest.mark.all_interfaces
diff --git a/tests/tape/test_qscript.py b/tests/tape/test_qscript.py
index b8618d5afc5..9e5dbca64bc 100644
--- a/tests/tape/test_qscript.py
+++ b/tests/tape/test_qscript.py
@@ -925,12 +925,6 @@ def test_diagonalizing_gates_not_queued(self):
     ),
 ]
 
-warnings_matches = {
-    State: "Requested state or density matrix with finite shots",
-    VnEntropy: "Requested Von Neumann entropy with finite shots",
-    MutualInfo: "Requested mutual information with finite shots",
-}
-
 
 class TestMeasurementProcess:
     """Tests for the shape and numeric type of a measurement process"""
@@ -1027,7 +1021,7 @@ def test_output_shapes_single(self, measurement, expected_shape, shots):
         b = np.array(0.2)
 
         ops = [qml.RY(a, 0), qml.RX(b, 0)]
-        qs = QuantumScript(ops, [measurement])
+        qs = QuantumScript(ops, [measurement], shots=shots)
 
         shot_dim = len(shots) if isinstance(shots, tuple) else shots
         if expected_shape is None:
@@ -1048,10 +1042,8 @@ def test_output_shapes_single_qnode_check(self, measurement, expected_shape, sho
         shape of a QNode for a single measurement."""
         if shots is None and measurement.return_type is qml.measurements.Sample:
             pytest.skip("Sample doesn't support analytic computations.")
-        if isinstance(shots, tuple) and isinstance(
-            measurement, (qml.measurements.MutualInfoMP, qml.measurements.VnEntropyMP)
-        ):
-            pytest.skip("Shot vectors and entropies not supported.")
+        if shots and isinstance(measurement, qml.measurements.StateMeasurement):
+            pytest.skip("State measurements with finite shots not supported.")
 
         dev = qml.device("default.qubit", wires=3, shots=shots)
 
@@ -1059,15 +1051,10 @@ def test_output_shapes_single_qnode_check(self, measurement, expected_shape, sho
         b = np.array(0.2)
 
         ops = [qml.RY(a, 0), qml.RX(b, 0)]
-        qs = QuantumScript(ops, [measurement])
+        qs = QuantumScript(ops, [measurement], shots=shots)
 
-        w_match = warnings_matches.get(measurement.return_type, None)
-        if shots is not None and w_match is not None:
-            with pytest.warns(UserWarning, match=w_match):
-                res = qml.execute([qs], dev, gradient_fn=None)[0]
-        else:
-            # TODO: test gradient_fn is not None when the interface `execute` functions are implemented
-            res = qml.execute([qs], dev, gradient_fn=None)[0]
+        # TODO: test gradient_fn is not None when the interface `execute` functions are implemented
+        res = qml.execute([qs], dev, gradient_fn=None)[0]
 
         if isinstance(shots, tuple):
             res_shape = tuple(r.shape for r in res)
@@ -1129,7 +1116,7 @@ def test_multi_measure(self, measurements, expected, shots):
         expectation value, variance and probability measurements."""
         dev = qml.device("default.qubit", wires=3, shots=shots)
 
-        qs = QuantumScript(measurements=measurements)
+        qs = QuantumScript(measurements=measurements, shots=shots)
 
         if measurements[0].return_type is qml.measurements.Sample:
             expected[1] = shots
@@ -1163,7 +1150,7 @@ def test_multi_measure_shot_vector(self, measurements, expected):
         a = np.array(0.1)
         b = np.array(0.2)
         ops = [qml.RY(a, 0), qml.RX(b, 0)]
-        qs = QuantumScript(ops, measurements)
+        qs = QuantumScript(ops, measurements, shots=shots)
 
         # Update expected as we're using a shotvector
         expected = tuple(expected for _ in shots)
@@ -1188,7 +1175,9 @@ def test_multi_measure_sample(self, shots):
 
         num_samples = 3
         ops = [qml.RY(a, 0), qml.RX(b, 0)]
-        qs = QuantumScript(ops, [qml.sample(qml.PauliZ(i)) for i in range(num_samples)])
+        qs = QuantumScript(
+            ops, [qml.sample(qml.PauliZ(i)) for i in range(num_samples)], shots=shots
+        )
 
         expected = tuple(() if shots == 1 else (shots,) for _ in range(num_samples))
 
@@ -1226,7 +1215,7 @@ def test_broadcasting_single(self, measurement, expected_shape, shots):
             qml.RX(b, wires=0)
             qml.apply(measurement)
 
-        tape = qml.tape.QuantumScript.from_queue(q)
+        tape = qml.tape.QuantumScript.from_queue(q, shots=shots)
         expected_shape = qml.execute([tape], dev, gradient_fn=None)[0].shape
 
         assert tape.shape(dev) == expected_shape
@@ -1254,7 +1243,7 @@ def test_broadcasting_multi(self, measurement, expected, shots):
             for _ in range(2):
                 qml.apply(measurement)
 
-        tape = qml.tape.QuantumScript.from_queue(q)
+        tape = qml.tape.QuantumScript.from_queue(q, shots=shots)
         expected = qml.execute([tape], dev, gradient_fn=None)[0]
         actual = tape.shape(dev)
 
@@ -1273,7 +1262,9 @@ def test_multi_measure_sample_obs_shot_vector(self):
 
         num_samples = 3
         ops = [qml.RY(a, 0), qml.RX(b, 0)]
-        qs = QuantumScript(ops, [qml.sample(qml.PauliZ(i)) for i in range(num_samples)])
+        qs = QuantumScript(
+            ops, [qml.sample(qml.PauliZ(i)) for i in range(num_samples)], shots=shots
+        )
 
         expected = tuple(tuple(() if s == 1 else (s,) for _ in range(num_samples)) for s in shots)
 
@@ -1294,7 +1285,7 @@ def test_multi_measure_sample_wires_shot_vector(self):
 
         num_samples = 3
         ops = [qml.RY(0.3, 0), qml.RX(0.2, 0)]
-        qs = QuantumScript(ops, [qml.sample()] * num_samples)
+        qs = QuantumScript(ops, [qml.sample()] * num_samples, shots=shots)
 
         expected = tuple(
             tuple((3,) if s == 1 else (s, 3) for _ in range(num_samples)) for s in shots
@@ -1318,7 +1309,7 @@ def test_raises_broadcasting_shot_vector(self):
             qml.RX(np.array([0.3, 0.4]), wires=0)
             qml.expval(qml.PauliZ(0))
 
-        tape = qml.tape.QuantumScript.from_queue(q)
+        tape = qml.tape.QuantumScript.from_queue(q, shots=(1, 2, 3))
         msg = "Parameter broadcasting when using a shot vector is not supported yet"
 
         with pytest.raises(NotImplementedError, match=msg):
@@ -1343,21 +1334,13 @@ def test_float_measures(self, ret, shots):
         """Test that most measurements output floating point values and that
         the tape output domain correctly identifies this."""
         dev = qml.device("default.qubit", wires=3, shots=shots)
-        if isinstance(shots, tuple) and isinstance(
-            ret, (qml.measurements.MutualInfoMP, qml.measurements.VnEntropyMP)
-        ):
-            pytest.skip("Shot vectors and entropies not supported.")
+        if shots and isinstance(ret, (qml.measurements.MutualInfoMP, qml.measurements.VnEntropyMP)):
+            pytest.skip("Shots and entropies not supported.")
 
         a, b = 0.3, 0.2
         ops = [qml.RY(a, 0), qml.RZ(b, 0)]
-        qs = QuantumScript(ops, [ret])
-
-        w_match = warnings_matches.get(ret.return_type, None)
-        if shots is not None and w_match is not None:
-            with pytest.warns(UserWarning, match=w_match):
-                result = qml.execute([qs], dev, gradient_fn=None)[0]
-        else:
-            result = qml.execute([qs], dev, gradient_fn=None)[0]
+        qs = QuantumScript(ops, [ret], shots=shots)
+        result = qml.execute([qs], dev, gradient_fn=None)[0]
 
         if not isinstance(result, tuple):
             result = (result,)
@@ -1390,7 +1373,7 @@ def test_sample_int_eigvals(self, ret):
         sampling measurement with a Hermitian observable with integer
         eigenvalues."""
         dev = qml.device("default.qubit", wires=3, shots=5)
-        qs = QuantumScript([qml.RY(0.4, 0)], [ret])
+        qs = QuantumScript([qml.RY(0.4, 0)], [ret], shots=5)
 
         result = qml.execute([qs], dev, gradient_fn=None)[0]
 
@@ -1413,7 +1396,7 @@ def test_sample_real_eigvals(self):
         )
         herm = np.outer(arr, arr)
 
-        qs = QuantumScript([qml.RY(0.4, 0)], [qml.sample(qml.Hermitian(herm, wires=0))])
+        qs = QuantumScript([qml.RY(0.4, 0)], [qml.sample(qml.Hermitian(herm, wires=0))], shots=5)
 
         result = qml.execute([qs], dev, gradient_fn=None)[0]
 
@@ -1439,7 +1422,7 @@ def test_sample_real_and_int_eigvals(self):
         a, b = 0, 3
         ops = [qml.RY(a, 0), qml.RX(b, 0)]
         m = [qml.sample(qml.Hermitian(herm, wires=0)), qml.sample(qml.PauliZ(1))]
-        qs = QuantumScript(ops, m)
+        qs = QuantumScript(ops, m, shots=5)
 
         result = qml.execute([qs], dev, gradient_fn=None)[0]
 
diff --git a/tests/tape/test_tape.py b/tests/tape/test_tape.py
index a8c33a9d73f..e509694f117 100644
--- a/tests/tape/test_tape.py
+++ b/tests/tape/test_tape.py
@@ -1234,31 +1234,26 @@ def test_expand_does_not_affect_original_tape(self):
         assert qml.equal(expanded.measurements[1], qml.expval(qml.PauliZ(0)))
         assert expanded.shots is tape.shots
 
-    def test_is_sampled_reserved_after_expansion(self, monkeypatch):
+    def test_is_sampled_reserved_after_expansion(self):
         """Test that the is_sampled property is correctly set when tape
         expansion happens."""
         dev = qml.device("default.qubit", wires=1, shots=10)
 
-        # Remove support for an op to enforce decomposition & tape expansion
-        mock_ops = copy.copy(dev.operations)
-        mock_ops.remove("T")
+        class UnsupportedT(qml.operation.Operation):
+            """A T gate that provides a decomposition, but no matrix."""
 
-        with monkeypatch.context() as m:
-            m.setattr(dev, "operations", mock_ops)
+            @staticmethod
+            def compute_decomposition(wires):  # pylint:disable=arguments-differ
+                return [qml.PhaseShift(np.pi / 4, wires=wires)]
 
-            def circuit():
-                qml.T(wires=0)
-                return sample(qml.PauliZ(0))
-
-            # Choosing parameter-shift not to swap the device under the hood
-            qnode = qml.QNode(circuit, dev, diff_method="parameter-shift")
-            qnode()
+        @qml.qnode(dev, diff_method="parameter-shift")
+        def circuit():
+            UnsupportedT(wires=0)
+            return sample(qml.PauliZ(0))
 
-            # Double-checking that the T gate is not supported
-            assert "T" not in qnode.device.operations
-            assert "T" not in qnode._original_device.operations
+        circuit()
 
-            assert qnode.qtape.is_sampled
+        assert circuit.qtape.is_sampled
 
 
 class TestExecution:
@@ -1390,11 +1385,12 @@ def test_prob_expectation_values(self, tol):
         assert isinstance(res, tuple)
         assert len(res) == 2
 
-        assert isinstance(res[0], np.ndarray)
+        assert isinstance(res[0], np.float64)
         assert np.allclose(res[0], np.cos(x), atol=tol, rtol=0)
 
         assert isinstance(res[1], np.ndarray)
-        assert np.allclose(res[1], np.abs(dev.state) ** 2, atol=tol, rtol=0)
+        final_state, _ = qml.devices.qubit.get_final_state(tape)
+        assert np.allclose(res[1], np.abs(final_state.flatten()) ** 2, atol=tol, rtol=0)
 
     def test_single_mode_sample(self):
         """Test that there is only one array of values returned
@@ -1409,6 +1405,7 @@ def test_single_mode_sample(self):
             qml.CNOT(wires=[0, 1])
             qml.sample(qml.PauliZ(0) @ qml.PauliX(1))
 
+        tape._shots = qml.measurements.Shots(10)
         res = dev.execute(tape)
         assert res.shape == (10,)
 
@@ -1426,6 +1423,7 @@ def test_multiple_samples(self):
             qml.sample(qml.PauliZ(0))
             qml.sample(qml.PauliZ(1))
 
+        tape._shots = qml.measurements.Shots(10)
         res = dev.execute(tape)
         assert isinstance(res, tuple)
         assert isinstance(res[0], np.ndarray)
@@ -1447,22 +1445,27 @@ def test_samples_expval(self):
             qml.sample(qml.PauliZ(0))
             qml.expval(qml.PauliZ(1))
 
+        tape._shots = qml.measurements.Shots(10)
         res = dev.execute(tape)
         assert isinstance(res, tuple)
         assert isinstance(res[0], np.ndarray)
         assert res[0].shape == (10,)
-        assert isinstance(res[1], np.ndarray)
+        assert isinstance(res[1], np.float64)
         assert res[1].shape == ()
 
     def test_decomposition(self, tol):
         """Test decomposition onto a device's supported gate set"""
         dev = qml.device("default.qubit", wires=1)
+        from pennylane.devices.qubit.preprocess import _accepted_operator
 
         with QuantumTape() as tape:
             qml.U3(0.1, 0.2, 0.3, wires=[0])
             qml.expval(qml.PauliZ(0))
 
-        tape = tape.expand(stop_at=lambda obj: obj.name in dev.operations)
+        def stop_fn(op):
+            return isinstance(op, qml.measurements.MeasurementProcess) or _accepted_operator(op)
+
+        tape = tape.expand(stop_at=stop_fn)
         res = dev.execute(tape)
         assert np.allclose(res, np.cos(0.1), atol=tol, rtol=0)
 
@@ -1856,7 +1859,7 @@ def test_output_shapes_single(self, measurement, expected_shape, shots):
             qml.RX(b, wires=0)
             qml.apply(measurement)
 
-        tape = qml.tape.QuantumScript.from_queue(q)
+        tape = qml.tape.QuantumScript.from_queue(q, shots=shots)
         shot_dim = shots if not isinstance(shots, tuple) else len(shots)
         if expected_shape is None:
             expected_shape = shot_dim if shot_dim == 1 else (shot_dim,)
@@ -1900,7 +1903,7 @@ def test_output_shapes_single_qnode_check(self, measurement, _, shots):
             qml.RX(b, wires=0)
             qml.apply(measurement)
 
-        tape = qml.tape.QuantumScript.from_queue(q)
+        tape = qml.tape.QuantumScript.from_queue(q, shots=shots)
         res = qml.execute([tape], dev, gradient_fn=qml.gradients.param_shift)[0]
 
         if isinstance(res, tuple):
@@ -1976,7 +1979,7 @@ def test_multi_measure(self, measurements, expected, shots):
             for m in measurements:
                 qml.apply(m)
 
-        tape = qml.tape.QuantumScript.from_queue(q)
+        tape = qml.tape.QuantumScript.from_queue(q, shots=shots)
         if measurements[0].return_type is qml.measurements.Sample:
             expected[1] = shots
             expected = tuple(expected)
@@ -2015,7 +2018,7 @@ def test_multi_measure_shot_vector(self, measurements, expected):
             for m in measurements:
                 qml.apply(m)
 
-        tape = qml.tape.QuantumScript.from_queue(q)
+        tape = qml.tape.QuantumScript.from_queue(q, shots=shots)
         # Modify expected to account for shot vector
         expected = tuple(expected for _ in shots)
         res = tape.shape(dev)
@@ -2038,7 +2041,7 @@ def test_multi_measure_sample(self, shots):
             for i in range(num_samples):
                 qml.sample(qml.PauliZ(i))
 
-        tape = qml.tape.QuantumScript.from_queue(q)
+        tape = qml.tape.QuantumScript.from_queue(q, shots=shots)
         if shots == 1:
             expected = tuple(() for _ in range(num_samples))
         else:
@@ -2064,7 +2067,7 @@ def test_multi_measure_sample_shot_vector(self):
             for i in range(num_samples):
                 qml.sample(qml.PauliZ(i))
 
-        tape = qml.tape.QuantumScript.from_queue(q)
+        tape = qml.tape.QuantumScript.from_queue(q, shots=shots)
         expected = []
         for s in shots:
             if s == 1:
@@ -2101,7 +2104,7 @@ def test_broadcasting_single(self, measurement, _, shots):
             qml.RX(b, wires=0)
             qml.apply(measurement)
 
-        tape = qml.tape.QuantumScript.from_queue(q)
+        tape = qml.tape.QuantumScript.from_queue(q, shots=shots)
         expected = qml.execute([tape], dev, gradient_fn=None)[0]
         assert tape.shape(dev) == expected.shape
 
@@ -2128,7 +2131,7 @@ def test_broadcasting_multi(self, measurement, expected, shots):
             for _ in range(2):
                 qml.apply(measurement)
 
-        tape = qml.tape.QuantumScript.from_queue(q)
+        tape = qml.tape.QuantumScript.from_queue(q, shots=shots)
         expected = qml.execute([tape], dev, gradient_fn=None)[0]
         expected = tuple(i.shape for i in expected)
         assert tape.shape(dev) == expected
@@ -2231,7 +2234,7 @@ def test_sample_real_and_int_eigvals(self):
         """Test that the tape can correctly determine the output domain for
         multiple sampling measurements with a Hermitian observable with real
         eigenvalues and another one with integer eigenvalues."""
-        dev = qml.device("default.qubit", wires=3, shots=5, r_dtype=np.float64)
+        dev = qml.device("default.qubit", wires=3, shots=5)
 
         arr = np.array(
             [
diff --git a/tests/templates/test_embeddings/test_amplitude.py b/tests/templates/test_embeddings/test_amplitude.py
index 8b8d79ba1ab..64bdd9a6783 100644
--- a/tests/templates/test_embeddings/test_amplitude.py
+++ b/tests/templates/test_embeddings/test_amplitude.py
@@ -79,10 +79,9 @@ def test_prepares_correct_state(self, inpt, normalize):
         @qml.qnode(dev)
         def circuit(x=None):
             qml.AmplitudeEmbedding(features=x, wires=range(n_qubits), normalize=normalize)
-            return [qml.expval(qml.PauliZ(i)) for i in range(n_qubits)]
+            return qml.state()
 
-        circuit(x=inpt)
-        state = circuit.device.state.ravel()
+        state = circuit(x=inpt).ravel()
         assert np.allclose(state, inpt)
 
     @pytest.mark.parametrize("normalize", (True, False))
@@ -112,10 +111,9 @@ def test_prepares_padded_state(self, inpt, pad):
         @qml.qnode(dev)
         def circuit(x=None):
             qml.AmplitudeEmbedding(features=x, wires=range(n_qubits), pad_with=pad, normalize=False)
-            return [qml.expval(qml.PauliZ(i)) for i in range(n_qubits)]
+            return qml.state()
 
-        circuit(x=inpt)
-        state = circuit.device.state.ravel()
+        state = circuit(x=inpt).ravel()
         # Make sure all padded values are the same constant
         # by checking how many different values there are
         assert len(set(state[len(inpt) :])) == 1
@@ -148,17 +146,18 @@ def test_custom_wire_labels(self, tol):
         @qml.qnode(dev)
         def circuit():
             qml.AmplitudeEmbedding(features, wires=range(3))
-            return qml.expval(qml.Identity(0))
+            return qml.expval(qml.Identity(0)), qml.state()
 
         @qml.qnode(dev2)
         def circuit2():
             qml.AmplitudeEmbedding(features, wires=["z", "a", "k"])
-            return qml.expval(qml.Identity("z"))
+            return qml.expval(qml.Identity("z")), qml.state()
 
-        circuit()
-        circuit2()
+        res1, state1 = circuit()
+        res2, state2 = circuit2()
 
-        assert np.allclose(dev.state, dev2.state, atol=tol, rtol=0)
+        assert np.allclose(res1, res2, atol=tol, rtol=0)
+        assert np.allclose(state1, state2, atol=tol, rtol=0)
 
 
 class TestInputs:
diff --git a/tests/templates/test_embeddings/test_angle.py b/tests/templates/test_embeddings/test_angle.py
index ef531695f3d..2cb3677269e 100644
--- a/tests/templates/test_embeddings/test_angle.py
+++ b/tests/templates/test_embeddings/test_angle.py
@@ -130,17 +130,18 @@ def test_custom_wire_labels(self, tol):
         @qml.qnode(dev)
         def circuit():
             qml.AngleEmbedding(features, wires=range(3))
-            return qml.expval(qml.Identity(0))
+            return qml.expval(qml.Identity(0)), qml.state()
 
         @qml.qnode(dev2)
         def circuit2():
             qml.AngleEmbedding(features, wires=["z", "a", "k"])
-            return qml.expval(qml.Identity("z"))
+            return qml.expval(qml.Identity("z")), qml.state()
 
-        circuit()
-        circuit2()
+        res1, state1 = circuit()
+        res2, state2 = circuit2()
 
-        assert np.allclose(dev.state, dev2.state, atol=tol, rtol=0)
+        assert np.allclose(res1, res2, atol=tol, rtol=0)
+        assert np.allclose(state1, state2, atol=tol, rtol=0)
 
 
 class TestInputs:
diff --git a/tests/templates/test_embeddings/test_basis.py b/tests/templates/test_embeddings/test_basis.py
index 974c2e2e335..158448338bc 100644
--- a/tests/templates/test_embeddings/test_basis.py
+++ b/tests/templates/test_embeddings/test_basis.py
@@ -78,17 +78,18 @@ def test_custom_wire_labels(self, tol):
         @qml.qnode(dev)
         def circuit():
             qml.BasisEmbedding(features, wires=range(3))
-            return qml.expval(qml.Identity(0))
+            return qml.expval(qml.Identity(0)), qml.state()
 
         @qml.qnode(dev2)
         def circuit2():
             qml.BasisEmbedding(features, wires=["z", "a", "k"])
-            return qml.expval(qml.Identity("z"))
+            return qml.expval(qml.Identity("z")), qml.state()
 
-        circuit()
-        circuit2()
+        res1, state1 = circuit()
+        res2, state2 = circuit2()
 
-        assert np.allclose(dev.state, dev2.state, atol=tol, rtol=0)
+        assert np.allclose(res1, res2, atol=tol, rtol=0)
+        assert np.allclose(state1, state2, atol=tol, rtol=0)
 
 
 class TestInputs:
diff --git a/tests/templates/test_embeddings/test_iqp_emb.py b/tests/templates/test_embeddings/test_iqp_emb.py
index 5eb25401b43..b0cc165f07e 100644
--- a/tests/templates/test_embeddings/test_iqp_emb.py
+++ b/tests/templates/test_embeddings/test_iqp_emb.py
@@ -119,17 +119,18 @@ def test_custom_wire_labels(self, tol):
         @qml.qnode(dev)
         def circuit():
             qml.IQPEmbedding(features, wires=range(3))
-            return qml.expval(qml.Identity(0))
+            return qml.expval(qml.Identity(0)), qml.state()
 
         @qml.qnode(dev2)
         def circuit2():
             qml.IQPEmbedding(features, wires=["z", "a", "k"])
-            return qml.expval(qml.Identity("z"))
+            return qml.expval(qml.Identity("z")), qml.state()
 
-        circuit()
-        circuit2()
+        res1, state1 = circuit()
+        res2, state2 = circuit2()
 
-        assert np.allclose(dev.state, dev2.state, atol=tol, rtol=0)
+        assert np.allclose(res1, res2, atol=tol, rtol=0)
+        assert np.allclose(state1, state2, atol=tol, rtol=0)
 
 
 class TestInputs:
diff --git a/tests/templates/test_embeddings/test_qaoa_emb.py b/tests/templates/test_embeddings/test_qaoa_emb.py
index de49c972d67..a33dbac9c01 100644
--- a/tests/templates/test_embeddings/test_qaoa_emb.py
+++ b/tests/templates/test_embeddings/test_qaoa_emb.py
@@ -215,17 +215,18 @@ def test_custom_wire_labels(self, tol):
         @qml.qnode(dev)
         def circuit():
             qml.QAOAEmbedding(features, weights, wires=range(3))
-            return qml.expval(qml.Identity(0))
+            return qml.expval(qml.Identity(0)), qml.state()
 
         @qml.qnode(dev2)
         def circuit2():
             qml.QAOAEmbedding(features, weights, wires=["z", "a", "k"])
-            return qml.expval(qml.Identity("z"))
+            return qml.expval(qml.Identity("z")), qml.state()
 
-        circuit()
-        circuit2()
+        res1, state1 = circuit()
+        res2, state2 = circuit2()
 
-        assert np.allclose(dev.state, dev2.state, atol=tol, rtol=0)
+        assert np.allclose(res1, res2, atol=tol, rtol=0)
+        assert np.allclose(state1, state2, atol=tol, rtol=0)
 
 
 class TestInputs:
diff --git a/tests/templates/test_layers/test_basic_entangler.py b/tests/templates/test_layers/test_basic_entangler.py
index 4c2ca6a83bf..1e21703333d 100644
--- a/tests/templates/test_layers/test_basic_entangler.py
+++ b/tests/templates/test_layers/test_basic_entangler.py
@@ -92,17 +92,18 @@ def test_custom_wire_labels(self, tol):
         @qml.qnode(dev)
         def circuit():
             qml.BasicEntanglerLayers(weights, wires=range(3))
-            return qml.expval(qml.Identity(0))
+            return qml.expval(qml.Identity(0)), qml.state()
 
         @qml.qnode(dev2)
         def circuit2():
             qml.BasicEntanglerLayers(weights, wires=["z", "a", "k"])
-            return qml.expval(qml.Identity("z"))
+            return qml.expval(qml.Identity("z")), qml.state()
 
-        circuit()
-        circuit2()
+        res1, state1 = circuit()
+        res2, state2 = circuit2()
 
-        assert np.allclose(dev.state, dev2.state, atol=tol, rtol=0)
+        assert np.allclose(res1, res2, atol=tol, rtol=0)
+        assert np.allclose(state1, state2, atol=tol, rtol=0)
 
 
 class TestInputs:
diff --git a/tests/templates/test_layers/test_gate_fabric.py b/tests/templates/test_layers/test_gate_fabric.py
index 260090f79d1..fedad87a0bc 100644
--- a/tests/templates/test_layers/test_gate_fabric.py
+++ b/tests/templates/test_layers/test_gate_fabric.py
@@ -484,11 +484,9 @@ def test_decomposition_q(self, init_state, exp_state, tol):
         @qml.qnode(dev)
         def circuit(weight):
             qml.GateFabric(weight, wires, init_state=init_state)
-            return qml.expval(qml.PauliZ(0))
-
-        circuit(weight)
+            return qml.state()
 
-        assert qml.math.allclose(circuit.device.state, exp_state, atol=tol)
+        assert qml.math.allclose(circuit(weight), exp_state, atol=tol)
 
     @pytest.mark.parametrize(
         ("num_qubits", "layers", "exp_state"),
@@ -608,9 +606,7 @@ def circuit(weight):
             qml.GateFabric(weight, wires, init_state=init_state, include_pi=True)
             return qml.state()
 
-        circuit(weight)
-
-        assert qml.math.allclose(circuit.device.state, exp_state, atol=tol)
+        assert qml.math.allclose(circuit(weight), exp_state, atol=tol)
 
     def test_custom_wire_labels(self, tol):
         """Test that template can deal with non-numeric, nonconsecutive wire labels."""
@@ -623,17 +619,18 @@ def test_custom_wire_labels(self, tol):
         @qml.qnode(dev)
         def circuit():
             qml.GateFabric(weights, wires=range(4), init_state=init_state)
-            return qml.expval(qml.Identity(0))
+            return qml.expval(qml.Identity(0)), qml.state()
 
         @qml.qnode(dev2)
         def circuit2():
             qml.GateFabric(weights, wires=["z", "a", "k", "r"], init_state=init_state)
-            return qml.expval(qml.Identity("z"))
+            return qml.expval(qml.Identity("z")), qml.state()
 
-        circuit()
-        circuit2()
+        res1, state1 = circuit()
+        res2, state2 = circuit2()
 
-        assert qml.math.allclose(dev.state, dev2.state, atol=tol, rtol=0)
+        assert np.allclose(res1, res2, atol=tol, rtol=0)
+        assert np.allclose(state1, state2, atol=tol, rtol=0)
 
 
 class TestInputs:
diff --git a/tests/templates/test_layers/test_particle_conserving_u1.py b/tests/templates/test_layers/test_particle_conserving_u1.py
index 15fd62056d3..9a5915ca886 100644
--- a/tests/templates/test_layers/test_particle_conserving_u1.py
+++ b/tests/templates/test_layers/test_particle_conserving_u1.py
@@ -134,17 +134,18 @@ def test_custom_wire_labels(self, tol):
         @qml.qnode(dev)
         def circuit():
             qml.ParticleConservingU1(weights, wires=range(3), init_state=init_state)
-            return qml.expval(qml.Identity(0))
+            return qml.expval(qml.Identity(0)), qml.state()
 
         @qml.qnode(dev2)
         def circuit2():
             qml.ParticleConservingU1(weights, wires=["z", "a", "k"], init_state=init_state)
-            return qml.expval(qml.Identity("z"))
+            return qml.expval(qml.Identity("z")), qml.state()
 
-        circuit()
-        circuit2()
+        res1, state1 = circuit()
+        res2, state2 = circuit2()
 
-        assert np.allclose(dev.state, dev2.state, atol=tol, rtol=0)
+        assert np.allclose(res1, res2, atol=tol, rtol=0)
+        assert np.allclose(state1, state2, atol=tol, rtol=0)
 
     @pytest.mark.parametrize(
         ("init_state", "exp_state"),
@@ -198,11 +199,9 @@ def test_decomposition_u1ex(self, init_state, exp_state, tol):
         @qml.qnode(dev)
         def circuit(weights):
             qml.ParticleConservingU1(weights, wires, init_state=init_state)
-            return qml.expval(qml.PauliZ(0))
-
-        circuit(weights)
+            return qml.state()
 
-        assert np.allclose(circuit.device.state, exp_state, atol=tol)
+        assert np.allclose(circuit(weights), exp_state, atol=tol)
 
 
 class TestInputs:
diff --git a/tests/templates/test_layers/test_particle_conserving_u2.py b/tests/templates/test_layers/test_particle_conserving_u2.py
index 163aa458484..55ffd0a27ba 100644
--- a/tests/templates/test_layers/test_particle_conserving_u2.py
+++ b/tests/templates/test_layers/test_particle_conserving_u2.py
@@ -111,11 +111,9 @@ def test_decomposition_u2ex(self, init_state, exp_state, tol):
         def circuit(weight):
             qml.BasisState(init_state, wires=wires)
             qml.particle_conserving_u2.u2_ex_gate(weight, wires)
-            return qml.expval(qml.PauliZ(0))
-
-        circuit(weight)
+            return qml.state()
 
-        assert np.allclose(circuit.device.state, exp_state, atol=tol)
+        assert np.allclose(circuit(weight), exp_state, atol=tol)
 
     def test_custom_wire_labels(self, tol):
         """Test that template can deal with non-numeric, nonconsecutive wire labels."""
@@ -128,17 +126,18 @@ def test_custom_wire_labels(self, tol):
         @qml.qnode(dev)
         def circuit():
             qml.ParticleConservingU2(weights, wires=range(3), init_state=init_state)
-            return qml.expval(qml.Identity(0))
+            return qml.expval(qml.Identity(0)), qml.state()
 
         @qml.qnode(dev2)
         def circuit2():
             qml.ParticleConservingU2(weights, wires=["z", "a", "k"], init_state=init_state)
-            return qml.expval(qml.Identity("z"))
+            return qml.expval(qml.Identity("z")), qml.state()
 
-        circuit()
-        circuit2()
+        res1, state1 = circuit()
+        res2, state2 = circuit2()
 
-        assert np.allclose(dev.state, dev2.state, atol=tol, rtol=0)
+        assert np.allclose(res1, res2, atol=tol, rtol=0)
+        assert np.allclose(state1, state2, atol=tol, rtol=0)
 
 
 class TestInputs:
diff --git a/tests/templates/test_layers/test_random.py b/tests/templates/test_layers/test_random.py
index c8e5a034f20..37014185720 100644
--- a/tests/templates/test_layers/test_random.py
+++ b/tests/templates/test_layers/test_random.py
@@ -136,17 +136,18 @@ def test_custom_wire_labels(self, tol):
         @qml.qnode(dev)
         def circuit():
             qml.RandomLayers(weights, wires=range(3))
-            return qml.expval(qml.Identity(0))
+            return qml.expval(qml.Identity(0)), qml.state()
 
         @qml.qnode(dev2)
         def circuit2():
             qml.RandomLayers(weights, wires=["z", "a", "k"])
-            return qml.expval(qml.Identity("z"))
+            return qml.expval(qml.Identity("z")), qml.state()
 
-        circuit()
-        circuit2()
+        res1, state1 = circuit()
+        res2, state2 = circuit2()
 
-        assert np.allclose(dev.state, dev2.state, atol=tol, rtol=0)
+        assert np.allclose(res1, res2, atol=tol, rtol=0)
+        assert np.allclose(state1, state2, atol=tol, rtol=0)
 
 
 class TestInputs:
diff --git a/tests/templates/test_layers/test_simplified_twodesign.py b/tests/templates/test_layers/test_simplified_twodesign.py
index 7d0c8edb7eb..860d951f880 100644
--- a/tests/templates/test_layers/test_simplified_twodesign.py
+++ b/tests/templates/test_layers/test_simplified_twodesign.py
@@ -121,17 +121,18 @@ def test_custom_wire_labels(self, tol):
         @qml.qnode(dev)
         def circuit():
             qml.SimplifiedTwoDesign(initial_layer, weights, wires=range(3))
-            return qml.expval(qml.Identity(0))
+            return qml.expval(qml.Identity(0)), qml.state()
 
         @qml.qnode(dev2)
         def circuit2():
             qml.SimplifiedTwoDesign(initial_layer, weights, wires=["z", "a", "k"])
-            return qml.expval(qml.Identity("z"))
+            return qml.expval(qml.Identity("z")), qml.state()
 
-        circuit()
-        circuit2()
+        res1, state1 = circuit()
+        res2, state2 = circuit2()
 
-        assert np.allclose(dev.state, dev2.state, atol=tol, rtol=0)
+        assert np.allclose(res1, res2, atol=tol, rtol=0)
+        assert np.allclose(state1, state2, atol=tol, rtol=0)
 
 
 class TestInputs:
diff --git a/tests/templates/test_layers/test_strongly_entangling.py b/tests/templates/test_layers/test_strongly_entangling.py
index 4569a724a35..e75a203def4 100644
--- a/tests/templates/test_layers/test_strongly_entangling.py
+++ b/tests/templates/test_layers/test_strongly_entangling.py
@@ -76,17 +76,18 @@ def test_custom_wire_labels(self, tol):
         @qml.qnode(dev)
         def circuit():
             qml.StronglyEntanglingLayers(weights, wires=range(3))
-            return qml.expval(qml.Identity(0))
+            return qml.expval(qml.Identity(0)), qml.state()
 
         @qml.qnode(dev2)
         def circuit2():
             qml.StronglyEntanglingLayers(weights, wires=["z", "a", "k"])
-            return qml.expval(qml.Identity("z"))
+            return qml.expval(qml.Identity("z")), qml.state()
 
-        circuit()
-        circuit2()
+        res1, state1 = circuit()
+        res2, state2 = circuit2()
 
-        assert np.allclose(dev.state, dev2.state, atol=tol, rtol=0)
+        assert np.allclose(res1, res2, atol=tol, rtol=0)
+        assert np.allclose(state1, state2, atol=tol, rtol=0)
 
     @pytest.mark.parametrize(
         "n_layers, n_wires, ranges", [(2, 2, [1, 1]), (1, 3, [2]), (4, 4, [2, 3, 1, 3])]
diff --git a/tests/templates/test_state_preparations/test_arbitrary_state_prep.py b/tests/templates/test_state_preparations/test_arbitrary_state_prep.py
index d96bf25f69a..ba3a88ef399 100644
--- a/tests/templates/test_state_preparations/test_arbitrary_state_prep.py
+++ b/tests/templates/test_state_preparations/test_arbitrary_state_prep.py
@@ -118,23 +118,23 @@ def test_correct_gates_two_wires(self):
         assert queue[5].hyperparameters["pauli_word"] == "XY"
         assert queue[5].wires.labels == (0, 1)
 
-    def test_GHZ_generation(self, qubit_device_3_wires, tol):
+    def test_GHZ_generation(self, tol):
         """Test that the template prepares a GHZ state."""
         GHZ_state = np.array([1 / np.sqrt(2), 0, 0, 0, 0, 0, 0, 1 / np.sqrt(2)])
 
         weights = np.zeros(14)
         weights[13] = np.pi / 2
 
-        @qml.qnode(qubit_device_3_wires)
+        @qml.qnode(qml.device("default.qubit"))
         def circuit(weights):
             qml.ArbitraryStatePreparation(weights, [0, 1, 2])
-            return qml.expval(qml.PauliZ(0))
+            return qml.expval(qml.PauliZ(0)), qml.state()
 
-        circuit(weights)
+        _, state = circuit(weights)
 
-        assert np.allclose(circuit.device.state, GHZ_state, atol=tol, rtol=0)
+        assert np.allclose(state, GHZ_state, atol=tol, rtol=0)
 
-    def test_even_superposition_generation(self, qubit_device_3_wires, tol):
+    def test_even_superposition_generation(self, tol):
         """Test that the template prepares an even superposition state."""
         even_superposition_state = np.ones(8) / np.sqrt(8)
 
@@ -143,15 +143,15 @@ def test_even_superposition_generation(self, qubit_device_3_wires, tol):
         weights[3] = np.pi / 2
         weights[5] = np.pi / 2
 
-        @qml.qnode(qubit_device_3_wires)
+        @qml.qnode(qml.device("default.qubit"))
         def circuit(weights):
             qml.ArbitraryStatePreparation(weights, [0, 1, 2])
 
-            return qml.expval(qml.PauliZ(0))
+            return qml.expval(qml.PauliZ(0)), qml.state()
 
-        circuit(weights)
+        _, state = circuit(weights)
 
-        assert np.allclose(circuit.device.state, even_superposition_state, atol=tol, rtol=0)
+        assert np.allclose(state, even_superposition_state, atol=tol, rtol=0)
 
     def test_custom_wire_labels(self, tol):
         """Test that template can deal with non-numeric, nonconsecutive wire labels."""
@@ -163,17 +163,18 @@ def test_custom_wire_labels(self, tol):
         @qml.qnode(dev)
         def circuit():
             qml.ArbitraryStatePreparation(weights, wires=range(3))
-            return qml.expval(qml.Identity(0))
+            return qml.expval(qml.Identity(0)), qml.state()
 
         @qml.qnode(dev2)
         def circuit2():
             qml.ArbitraryStatePreparation(weights, wires=["z", "a", "k"])
-            return qml.expval(qml.Identity("z"))
+            return qml.expval(qml.Identity("z")), qml.state()
 
-        circuit()
-        circuit2()
+        res1, state1 = circuit()
+        res2, state2 = circuit2()
 
-        assert np.allclose(dev.state, dev2.state, atol=tol, rtol=0)
+        assert np.allclose(res1, res2, atol=tol, rtol=0)
+        assert np.allclose(state1, state2, atol=tol, rtol=0)
 
 
 class TestInputs:
diff --git a/tests/templates/test_state_preparations/test_basis_state_prep.py b/tests/templates/test_state_preparations/test_basis_state_prep.py
index 5819f867e2c..b63b01a33f7 100644
--- a/tests/templates/test_state_preparations/test_basis_state_prep.py
+++ b/tests/templates/test_state_preparations/test_basis_state_prep.py
@@ -62,10 +62,10 @@ def test_correct_pl_gates(self, basis_state, wires, target_wires):
         ([1, 0, 1], [0, 1, 2], [1, 0, 1]),
     ])
     # fmt: on
-    def test_state_preparation(self, tol, qubit_device_3_wires, basis_state, wires, target_state):
+    def test_state_preparation(self, tol, basis_state, wires, target_state):
         """Tests that the template produces the correct expectation values."""
 
-        @qml.qnode(qubit_device_3_wires)
+        @qml.qnode(qml.device("default.qubit"))
         def circuit():
             qml.BasisStatePreparation(basis_state, wires)
 
@@ -86,13 +86,11 @@ def circuit():
             ([1, 0, 1], [2, 0, 1], [0, 1, 1]),
         ],
     )
-    def test_state_preparation_jax_jit(
-        self, tol, qubit_device_3_wires, basis_state, wires, target_state
-    ):
+    def test_state_preparation_jax_jit(self, tol, basis_state, wires, target_state):
         """Tests that the template produces the correct expectation values."""
         import jax
 
-        @qml.qnode(qubit_device_3_wires, interface="jax")
+        @qml.qnode(qml.device("default.qubit"), interface="jax")
         def circuit(state):
             qml.BasisStatePreparation(state, wires)
 
@@ -115,14 +113,12 @@ def circuit(state):
             ([1, 0, 1], [2, 0, 1], [0, 1, 1]),
         ],
     )
-    def test_state_preparation_tf_autograph(
-        self, tol, qubit_device_3_wires, basis_state, wires, target_state
-    ):
+    def test_state_preparation_tf_autograph(self, tol, basis_state, wires, target_state):
         """Tests that the template produces the correct expectation values."""
         import tensorflow as tf
 
         @tf.function
-        @qml.qnode(qubit_device_3_wires, interface="tf")
+        @qml.qnode(qml.device("default.qubit"), interface="tf")
         def circuit(state):
             qml.BasisStatePreparation(state, wires)
 
@@ -144,17 +140,18 @@ def test_custom_wire_labels(self, tol):
         @qml.qnode(dev)
         def circuit():
             qml.BasisStatePreparation(basis_state, wires=range(3))
-            return qml.expval(qml.Identity(0))
+            return qml.expval(qml.Identity(0)), qml.state()
 
         @qml.qnode(dev2)
         def circuit2():
             qml.BasisStatePreparation(basis_state, wires=["z", "a", "k"])
-            return qml.expval(qml.Identity("z"))
+            return qml.expval(qml.Identity("z")), qml.state()
 
-        circuit()
-        circuit2()
+        res1, state1 = circuit()
+        res2, state2 = circuit2()
 
-        assert np.allclose(dev.state, dev2.state, atol=tol, rtol=0)
+        assert np.allclose(res1, res2, atol=tol, rtol=0)
+        assert np.allclose(state1, state2, atol=tol, rtol=0)
 
 
 class TestInputs:
diff --git a/tests/templates/test_state_preparations/test_mottonen_state_prep.py b/tests/templates/test_state_preparations/test_mottonen_state_prep.py
index 6548e44464e..f7ac7bb3844 100644
--- a/tests/templates/test_state_preparations/test_mottonen_state_prep.py
+++ b/tests/templates/test_state_preparations/test_mottonen_state_prep.py
@@ -122,19 +122,22 @@ class TestDecomposition:
         ([1 / 2, 0, 1j / 2, 1j / np.sqrt(2)], [0, 1], [1 / 2, 0, 0, 0, 1j / 2, 0, 1j / np.sqrt(2), 0]),
     ])
     # fmt: on
-    def test_state_preparation_fidelity(
-        self, tol, qubit_device_3_wires, state_vector, wires, target_state
-    ):
+    def test_state_preparation_fidelity(self, tol, state_vector, wires, target_state):
         """Tests that the template produces correct states with high fidelity."""
 
-        @qml.qnode(qubit_device_3_wires)
+        @qml.qnode(qml.device("default.qubit", wires=3))
         def circuit():
             qml.MottonenStatePreparation(state_vector, wires)
-            return qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliZ(1)), qml.expval(qml.PauliZ(2))
+            return (
+                qml.expval(qml.PauliZ(0)),
+                qml.expval(qml.PauliZ(1)),
+                qml.expval(qml.PauliZ(2)),
+                qml.state(),
+            )
 
-        circuit()
+        results = circuit()
 
-        state = circuit.device.state.ravel()
+        state = results[-1].ravel()
         fidelity = abs(np.vdot(state, target_state)) ** 2
 
         # We test for fidelity here, because the vector themselves will hardly match
@@ -207,18 +210,23 @@ def circuit():
     ])
     # fmt: on
     def test_state_preparation_probability_distribution(
-        self, tol, qubit_device_3_wires, state_vector, wires, target_state
+        self, tol, state_vector, wires, target_state
     ):
         """Tests that the template produces states with correct probability distribution."""
 
-        @qml.qnode(qubit_device_3_wires)
+        @qml.qnode(qml.device("default.qubit", wires=3))
         def circuit():
             qml.MottonenStatePreparation(state_vector, wires)
-            return qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliZ(1)), qml.expval(qml.PauliZ(2))
+            return (
+                qml.expval(qml.PauliZ(0)),
+                qml.expval(qml.PauliZ(1)),
+                qml.expval(qml.PauliZ(2)),
+                qml.state(),
+            )
 
-        circuit()
+        results = circuit()
 
-        state = circuit.device.state.ravel()
+        state = results[-1].ravel()
 
         probabilities = np.abs(state) ** 2
         target_probabilities = np.abs(target_state) ** 2
@@ -254,7 +262,7 @@ def circuit(state_vector):
         # when the RZ cascade is skipped, CNOT gates should only be those required for RY cascade
         spy = mocker.spy(circuit.device, "execute")
         circuit(state_vector)
-        tape = spy.call_args[0][0]
+        tape = spy.call_args[0][0][0]
 
         assert tape.specs["resources"].gate_types["CNOT"] == n_CNOT
 
@@ -268,17 +276,18 @@ def test_custom_wire_labels(self, tol):
         @qml.qnode(dev)
         def circuit():
             qml.MottonenStatePreparation(state, wires=range(3))
-            return qml.expval(qml.Identity(0))
+            return qml.expval(qml.Identity(0)), qml.state()
 
         @qml.qnode(dev2)
         def circuit2():
             qml.MottonenStatePreparation(state, wires=["z", "a", "k"])
-            return qml.expval(qml.Identity("z"))
+            return qml.expval(qml.Identity("z")), qml.state()
 
-        circuit()
-        circuit2()
+        res1, state1 = circuit()
+        res2, state2 = circuit2()
 
-        assert np.allclose(dev.state, dev2.state, atol=tol, rtol=0)
+        assert np.allclose(res1, res2, atol=tol, rtol=0)
+        assert np.allclose(state1, state2, atol=tol, rtol=0)
 
 
 class TestInputs:
diff --git a/tests/templates/test_subroutines/test_all_singles_doubles.py b/tests/templates/test_subroutines/test_all_singles_doubles.py
index cb37e63cd7a..a0e5413f521 100644
--- a/tests/templates/test_subroutines/test_all_singles_doubles.py
+++ b/tests/templates/test_subroutines/test_all_singles_doubles.py
@@ -120,7 +120,7 @@ def circuit():
                 singles=[[0, 1]],
                 doubles=[[0, 1, 2, 3]],
             )
-            return qml.expval(qml.Identity(0))
+            return qml.expval(qml.Identity(0)), qml.state()
 
         @qml.qnode(dev2)
         def circuit2():
@@ -131,12 +131,13 @@ def circuit2():
                 singles=[["z", "a"]],
                 doubles=[["z", "a", "k", "e"]],
             )
-            return qml.expval(qml.Identity("z"))
+            return qml.expval(qml.Identity("z")), qml.state()
 
-        circuit()
-        circuit2()
+        res1, state1 = circuit()
+        res2, state2 = circuit2()
 
-        assert np.allclose(dev.state, dev2.state, atol=tol, rtol=0)
+        assert np.allclose(res1, res2, atol=tol, rtol=0)
+        assert np.allclose(state1, state2, atol=tol, rtol=0)
 
 
 class TestInputs:
@@ -191,7 +192,7 @@ class TestInputs:
                 [[0, 2]],
                 [[0, 1, 2, 3]],
                 np.array([1, 1, 0, 0, 0]),
-                "BasisState parameter and wires",
+                "Basis states must be of length 4",
             ),
             (
                 np.array([-2.8, 1.6]),
diff --git a/tests/templates/test_subroutines/test_approx_time_evolution.py b/tests/templates/test_subroutines/test_approx_time_evolution.py
index 18e2c585f8a..82203dbd926 100644
--- a/tests/templates/test_subroutines/test_approx_time_evolution.py
+++ b/tests/templates/test_subroutines/test_approx_time_evolution.py
@@ -166,17 +166,18 @@ def test_custom_wire_labels(self, tol):
         @qml.qnode(dev)
         def circuit():
             qml.ApproxTimeEvolution(hamiltonian, 0.5, 2)
-            return qml.expval(qml.Identity(0))
+            return qml.expval(qml.Identity(0)), qml.state()
 
         @qml.qnode(dev2)
         def circuit2():
             qml.ApproxTimeEvolution(hamiltonian2, 0.5, 2)
-            return qml.expval(qml.Identity("z"))
+            return qml.expval(qml.Identity("z")), qml.state()
 
-        circuit()
-        circuit2()
+        res1, state1 = circuit()
+        res2, state2 = circuit2()
 
-        assert np.allclose(dev.state, dev2.state, atol=tol, rtol=0)
+        assert np.allclose(res1, res2, atol=tol, rtol=0)
+        assert np.allclose(state1, state2, atol=tol, rtol=0)
 
 
 class TestInputs:
@@ -410,7 +411,7 @@ def create_tape(coeffs, t):
     def cost(coeffs, t):
         tape = create_tape(coeffs, t)
 
-        if diff_method is qml.gradients.param_shift:
+        if diff_method is qml.gradients.param_shift and dev_name != "default.qubit":
             tape = dev.expand_fn(tape)
 
         return qml.execute([tape], dev, diff_method)[0]
diff --git a/tests/templates/test_subroutines/test_arbitrary_unitary.py b/tests/templates/test_subroutines/test_arbitrary_unitary.py
index afb45087a42..7f6ac28cd6a 100644
--- a/tests/templates/test_subroutines/test_arbitrary_unitary.py
+++ b/tests/templates/test_subroutines/test_arbitrary_unitary.py
@@ -158,17 +158,18 @@ def test_custom_wire_labels(self, tol):
         @qml.qnode(dev)
         def circuit():
             qml.ArbitraryUnitary(weights, wires=range(3))
-            return qml.expval(qml.Identity(0))
+            return qml.expval(qml.Identity(0)), qml.state()
 
         @qml.qnode(dev2)
         def circuit2():
             qml.ArbitraryUnitary(weights, wires=["z", "a", "k"])
-            return qml.expval(qml.Identity("z"))
+            return qml.expval(qml.Identity("z")), qml.state()
 
-        circuit()
-        circuit2()
+        res1, state1 = circuit()
+        res2, state2 = circuit2()
 
-        assert np.allclose(dev.state, dev2.state, atol=tol, rtol=0)
+        assert np.allclose(res1, res2, atol=tol, rtol=0)
+        assert np.allclose(state1, state2, atol=tol, rtol=0)
 
 
 class TestInputs:
diff --git a/tests/templates/test_subroutines/test_basis_rotation.py b/tests/templates/test_subroutines/test_basis_rotation.py
index 7767dd763c6..f969e001ba7 100644
--- a/tests/templates/test_subroutines/test_basis_rotation.py
+++ b/tests/templates/test_subroutines/test_basis_rotation.py
@@ -101,7 +101,7 @@ def circuit():
                 wires=range(2),
                 unitary_matrix=weights,
             )
-            return qml.expval(qml.Identity(0))
+            return qml.expval(qml.Identity(0)), qml.state()
 
         @qml.qnode(dev2)
         def circuit2():
@@ -110,12 +110,13 @@ def circuit2():
                 wires=["z", "a"],
                 unitary_matrix=weights,
             )
-            return qml.expval(qml.Identity("z"))
+            return qml.expval(qml.Identity("z")), qml.state()
 
-        circuit()
-        circuit2()
+        res1, state1 = circuit()
+        res2, state2 = circuit2()
 
-        assert np.allclose(dev.state, dev2.state, atol=tol, rtol=0)
+        assert np.allclose(res1, res2, atol=tol, rtol=0)
+        assert np.allclose(state1, state2, atol=tol, rtol=0)
 
     @pytest.mark.parametrize(
         ("unitary_matrix", "eigen_values", "exp_state"),
diff --git a/tests/templates/test_subroutines/test_commuting_evolution.py b/tests/templates/test_subroutines/test_commuting_evolution.py
index 7cb8c22c02a..e36ad81c76a 100644
--- a/tests/templates/test_subroutines/test_commuting_evolution.py
+++ b/tests/templates/test_subroutines/test_commuting_evolution.py
@@ -47,32 +47,32 @@ def test_adjoint():
     """Tests the CommutingEvolution.adjoint method provides the correct adjoint operation."""
 
     n_wires = 2
-    dev1 = qml.device("default.qubit", wires=n_wires)
-    dev2 = qml.device("default.qubit", wires=n_wires)
+    dev = qml.device("default.qubit", wires=n_wires)
 
     obs = [qml.PauliX(0) @ qml.PauliY(1), qml.PauliY(0) @ qml.PauliX(1)]
     coeffs = [1, -1]
     hamiltonian = qml.Hamiltonian(coeffs, obs)
     frequencies = (2,)
 
-    @qml.qnode(dev1)
+    @qml.qnode(dev)
     def adjoint_evolution_circuit(time):
         for i in range(n_wires):
             qml.Hadamard(i)
         qml.adjoint(qml.CommutingEvolution)(hamiltonian, time, frequencies)
-        return qml.expval(qml.PauliZ(1))
+        return qml.expval(qml.PauliZ(1)), qml.state()
 
-    @qml.qnode(dev2)
+    @qml.qnode(dev)
     def evolution_circuit(time):
         for i in range(n_wires):
             qml.Hadamard(i)
         qml.CommutingEvolution(hamiltonian, time, frequencies)
-        return qml.expval(qml.PauliZ(1))
+        return qml.expval(qml.PauliZ(1)), qml.state()
 
-    evolution_circuit(0.13)
-    adjoint_evolution_circuit(-0.13)
+    res1, state1 = evolution_circuit(0.13)
+    res2, state2 = adjoint_evolution_circuit(-0.13)
 
-    assert all(np.isclose(dev1.state, dev2.state))
+    assert res1 == res2
+    assert all(np.isclose(state1, state2))
 
 
 def test_decomposition_expand():
diff --git a/tests/templates/test_subroutines/test_double_excitation.py b/tests/templates/test_subroutines/test_double_excitation.py
index 76768b9a1fe..39b7c317428 100644
--- a/tests/templates/test_subroutines/test_double_excitation.py
+++ b/tests/templates/test_subroutines/test_double_excitation.py
@@ -242,17 +242,18 @@ def test_custom_wire_labels(self, tol):
         @qml.qnode(dev)
         def circuit():
             qml.FermionicDoubleExcitation(0.4, wires1=[0, 2], wires2=[1, 4, 3])
-            return qml.expval(qml.Identity(0))
+            return qml.expval(qml.Identity(0)), qml.state()
 
         @qml.qnode(dev2)
         def circuit2():
             qml.FermionicDoubleExcitation(0.4, wires1=["z", "k"], wires2=["a", "s", "t"])
-            return qml.expval(qml.Identity("z"))
+            return qml.expval(qml.Identity("z")), qml.state()
 
-        circuit()
-        circuit2()
+        res1, state1 = circuit()
+        res2, state2 = circuit2()
 
-        assert np.allclose(dev.state, dev2.state, atol=tol, rtol=0)
+        assert np.allclose(res1, res2, atol=tol, rtol=0)
+        assert np.allclose(state1, state2, atol=tol, rtol=0)
 
 
 class TestInputs:
diff --git a/tests/templates/test_subroutines/test_kupccgsd.py b/tests/templates/test_subroutines/test_kupccgsd.py
index 0343cb4d34f..4a1801985ae 100644
--- a/tests/templates/test_subroutines/test_kupccgsd.py
+++ b/tests/templates/test_subroutines/test_kupccgsd.py
@@ -137,7 +137,7 @@ def circuit():
                 delta_sz=0,
                 init_state=np.array([0, 1, 0, 1]),
             )
-            return qml.expval(qml.Identity(0))
+            return qml.expval(qml.Identity(0)), qml.state()
 
         @qml.qnode(dev2)
         def circuit2():
@@ -148,12 +148,13 @@ def circuit2():
                 delta_sz=0,
                 init_state=np.array([0, 1, 0, 1]),
             )
-            return qml.expval(qml.Identity("z"))
+            return qml.expval(qml.Identity("z")), qml.state()
 
-        circuit()
-        circuit2()
+        res1, state1 = circuit()
+        res2, state2 = circuit2()
 
-        assert np.allclose(dev.state, dev2.state, atol=tol, rtol=0)
+        assert np.allclose(res1, res2, atol=tol, rtol=0)
+        assert np.allclose(state1, state2, atol=tol, rtol=0)
 
     @pytest.mark.parametrize(
         ("num_qubits", "k", "exp_state"),
@@ -256,9 +257,9 @@ def circuit(weight):
             qml.kUpCCGSD(weight, wires=wires, k=k, delta_sz=0, init_state=init_state)
             return qml.state()
 
-        circuit(weight)
+        res = circuit(weight)
 
-        assert qml.math.allclose(circuit.device.state, exp_state, atol=tol)
+        assert qml.math.allclose(res, exp_state, atol=tol)
 
     @pytest.mark.parametrize(
         ("wires", "delta_sz", "generalized_singles_wires", "generalized_pair_doubles_wires"),
diff --git a/tests/templates/test_subroutines/test_permute.py b/tests/templates/test_subroutines/test_permute.py
index e2d61e75308..08f240d82e3 100644
--- a/tests/templates/test_subroutines/test_permute.py
+++ b/tests/templates/test_subroutines/test_permute.py
@@ -37,7 +37,7 @@ def identity_permutation():
         identity_permutation()
 
         # expand the Permute operation
-        tape = spy.call_args[0][0]
+        tape = spy.call_args[0][0][0]
 
         assert len(tape.operations) == 0
 
@@ -76,7 +76,7 @@ def two_cycle():
         spy = mocker.spy(two_cycle.device, "execute")
         two_cycle()
 
-        tape = spy.call_args[0][0]
+        tape = spy.call_args[0][0][0]
 
         # Check that the Permute operation was expanded to SWAPs when the QNode
         # is evaluated, and that the wires are the same
@@ -130,7 +130,7 @@ def cycle():
         spy = mocker.spy(cycle.device, "execute")
         cycle()
 
-        tape = spy.call_args[0][0]
+        tape = spy.call_args[0][0][0]
 
         # Check that the Permute operation was expanded to SWAPs when the QNode
         # is evaluated, and that the wires are the same
@@ -180,7 +180,7 @@ def arbitrary_perm():
         spy = mocker.spy(arbitrary_perm.device, "execute")
         arbitrary_perm()
 
-        tape = spy.call_args[0][0]
+        tape = spy.call_args[0][0][0]
 
         # Check that the Permute operation was expanded to SWAPs when the QNode
         # is evaluated, and that the wires are the same
@@ -240,7 +240,7 @@ def subset_perm():
         spy = mocker.spy(subset_perm.device, "execute")
         subset_perm()
 
-        tape = spy.call_args[0][0]
+        tape = spy.call_args[0][0][0]
 
         # Check that the Permute operation was expanded to SWAPs when the QNode
         # is evaluated, and that the wires are the same
@@ -290,17 +290,18 @@ def test_custom_wire_labels(self, tol):
         @qml.qnode(dev)
         def circuit():
             qml.Permute(permutation, wires=range(4))
-            return qml.expval(qml.Identity(0))
+            return qml.expval(qml.Identity(0)), qml.state()
 
         @qml.qnode(dev2)
         def circuit2():
             qml.Permute(permutation2, wires=["z", "a", "k", "o"])
-            return qml.expval(qml.Identity("z"))
+            return qml.expval(qml.Identity("z")), qml.state()
 
-        circuit()
-        circuit2()
+        res1, state1 = circuit()
+        res2, state2 = circuit2()
 
-        assert np.allclose(dev.state, dev2.state, atol=tol, rtol=0)
+        assert np.allclose(res1, res2, atol=tol, rtol=0)
+        assert np.allclose(state1, state2, atol=tol, rtol=0)
 
 
 class TestInputs:
diff --git a/tests/templates/test_subroutines/test_qft.py b/tests/templates/test_subroutines/test_qft.py
index 08ddec4bb29..9e17377f215 100644
--- a/tests/templates/test_subroutines/test_qft.py
+++ b/tests/templates/test_subroutines/test_qft.py
@@ -42,10 +42,9 @@ def test_QFT_decomposition(self, n_qubits):
 
         out_states = []
         for state in np.eye(2**n_qubits):
-            dev.reset()
             ops = [qml.StatePrep(state, wires=range(n_qubits))] + decomp
-            dev.apply(ops)
-            out_states.append(dev.state)
+            qs = qml.tape.QuantumScript(ops, [qml.state()])
+            out_states.append(dev.execute(qs))
 
         reconstructed_unitary = np.array(out_states).T
         expected_unitary = qml.QFT(wires=range(n_qubits)).matrix()
diff --git a/tests/templates/test_subroutines/test_qpe.py b/tests/templates/test_subroutines/test_qpe.py
index e3bd15d8b09..8037ae9c7b3 100644
--- a/tests/templates/test_subroutines/test_qpe.py
+++ b/tests/templates/test_subroutines/test_qpe.py
@@ -99,9 +99,10 @@ def test_phase_estimated(self, phase):
                 qml.probs(estimation_wires)
 
             tape = qml.tape.QuantumScript.from_queue(q)
-            tape = tape.expand(depth=2, stop_at=lambda obj: obj.name in dev.operations)
+            tapes, _ = dev.preprocess()[0]([tape])
+            assert len(tapes) == 1
 
-            res = dev.execute(tape).flatten()
+            res = dev.execute(tapes)[0].flatten()
             initial_estimate = np.argmax(res) / 2 ** (wires - 1)
 
             # We need to rescale because RX is exp(- i theta X / 2) and we expect a unitary of the
@@ -150,8 +151,9 @@ def test_phase_estimated_two_qubit(self):
                 qml.probs(estimation_wires)
 
             tape = qml.tape.QuantumScript.from_queue(q)
-            tape = tape.expand(depth=2, stop_at=lambda obj: obj.name in dev.operations)
-            res = dev.execute(tape).flatten()
+            tapes, _ = dev.preprocess()[0]([tape])
+            assert len(tapes) == 1
+            res = dev.execute(tapes)[0].flatten()
 
             if phase < 0:
                 estimate = np.argmax(res) / 2 ** (wires - 2) - 1
@@ -195,8 +197,9 @@ def test_phase_estimated_single_ops(self, param):
                 prep=[qml.StatePrep(eig_vec, wires=target_wires)],
             )
 
-            tape = tape.expand(depth=2, stop_at=lambda obj: obj.name in dev.operations)
-            res = dev.execute(tape).flatten()
+            tapes, _ = dev.preprocess()[0]([tape])
+            res = dev.execute(tapes)[0].flatten()
+            assert len(tapes) == 1
 
             estimate = np.argmax(res) / 2 ** (wires - 2)
             estimates.append(estimate)
@@ -237,8 +240,9 @@ def test_phase_estimated_ops(self, param):
                 prep=[qml.StatePrep(eig_vec, wires=target_wires)],
             )
 
-            tape = tape.expand(depth=2, stop_at=lambda obj: obj.name in dev.operations)
-            res = dev.execute(tape).flatten()
+            tapes, _ = dev.preprocess()[0]([tape])
+            assert len(tapes) == 1
+            res = dev.execute(tapes)[0].flatten()
 
             estimate = np.argmax(res) / 2 ** (wires - 2)
             estimates.append(estimate)
diff --git a/tests/templates/test_subroutines/test_single_excitation.py b/tests/templates/test_subroutines/test_single_excitation.py
index 69bf98d5719..76eda25a2b8 100644
--- a/tests/templates/test_subroutines/test_single_excitation.py
+++ b/tests/templates/test_subroutines/test_single_excitation.py
@@ -118,17 +118,18 @@ def test_custom_wire_labels(self, tol):
         @qml.qnode(dev)
         def circuit():
             qml.FermionicSingleExcitation(0.4, wires=[1, 0, 2])
-            return qml.expval(qml.Identity(0))
+            return qml.expval(qml.Identity(0)), qml.state()
 
         @qml.qnode(dev2)
         def circuit2():
             qml.FermionicSingleExcitation(0.4, wires=["a", "z", "k"])
-            return qml.expval(qml.Identity("z"))
+            return qml.expval(qml.Identity("z")), qml.state()
 
-        circuit()
-        circuit2()
+        res1, state1 = circuit()
+        res2, state2 = circuit2()
 
-        assert np.allclose(dev.state, dev2.state, atol=tol, rtol=0)
+        assert np.allclose(res1, res2, atol=tol, rtol=0)
+        assert np.allclose(state1, state2, atol=tol, rtol=0)
 
 
 class TestInputs:
diff --git a/tests/templates/test_subroutines/test_uccsd.py b/tests/templates/test_subroutines/test_uccsd.py
index 5aab270d5d6..22ed28965b3 100644
--- a/tests/templates/test_subroutines/test_uccsd.py
+++ b/tests/templates/test_subroutines/test_uccsd.py
@@ -162,7 +162,7 @@ def circuit():
                 d_wires=[[[0, 1], [2, 3]]],
                 init_state=np.array([0, 1, 0, 1]),
             )
-            return qml.expval(qml.Identity(0))
+            return qml.expval(qml.Identity(0)), qml.state()
 
         @qml.qnode(dev2)
         def circuit2():
@@ -173,12 +173,13 @@ def circuit2():
                 d_wires=[[["z", "a"], ["k", "e"]]],
                 init_state=np.array([0, 1, 0, 1]),
             )
-            return qml.expval(qml.Identity("z"))
+            return qml.expval(qml.Identity("z")), qml.state()
 
-        circuit()
-        circuit2()
+        res1, state1 = circuit()
+        res2, state2 = circuit2()
 
-        assert np.allclose(dev.state, dev2.state, atol=tol, rtol=0)
+        assert np.allclose(res1, res2, atol=tol, rtol=0)
+        assert np.allclose(state1, state2, atol=tol, rtol=0)
 
 
 class TestInputs:
@@ -213,7 +214,7 @@ class TestInputs:
                 [[0, 2]],
                 [],
                 np.array([1, 1, 0, 0, 0]),
-                "BasisState parameter and wires",
+                "Basis states must be of length 4",
             ),
             (
                 np.array([-2.8, 1.6]),
diff --git a/tests/test_debugging.py b/tests/test_debugging.py
index 38769399e42..1f95d33e3ee 100644
--- a/tests/test_debugging.py
+++ b/tests/test_debugging.py
@@ -54,30 +54,29 @@ def circuit():
         assert all(np.allclose(v1, v2) for v1, v2 in zip(result.values(), expected.values()))
 
     # pylint: disable=protected-access
-    def test_default_qubit2(self):
+    @pytest.mark.parametrize("method", [None, "backprop", "parameter-shift", "adjoint"])
+    def test_default_qubit2(self, method):
         """Test that multiple snapshots are returned correctly on the new
         state-vector simulator."""
-        dev = qml.devices.DefaultQubit()
+        dev = qml.device("default.qubit")
 
         # TODO: add additional QNode test once the new device supports it
 
-        ops = [
-            qml.Snapshot(),
-            qml.Hadamard(wires=0),
-            qml.Snapshot("very_important_state"),
-            qml.CNOT(wires=[0, 1]),
-            qml.Snapshot(),
-        ]
-        qs = qml.tape.QuantumScript(ops, [qml.expval(qml.PauliX(0))])
+        @qml.qnode(dev, diff_method=method)
+        def circuit():
+            qml.Snapshot()
+            qml.Hadamard(wires=0)
+            qml.Snapshot("very_important_state")
+            qml.CNOT(wires=[0, 1])
+            qml.Snapshot()
+            return qml.expval(qml.PauliX(0))
 
-        dev.execute(qs)
+        circuit()
         assert dev._debugger is None
+        if method is not None:
+            assert circuit.interface == "auto"
 
-        with qml.debugging._Debugger(dev) as dbg:
-            dev.execute(qs)
-
-        result = dbg.snapshots
-
+        result = qml.snapshots(circuit)()
         expected = {
             0: np.array([1, 0, 0, 0]),
             "very_important_state": np.array([1 / np.sqrt(2), 0, 1 / np.sqrt(2), 0]),
diff --git a/tests/test_qnode.py b/tests/test_qnode.py
index 65254a331bf..1c7d61273f7 100644
--- a/tests/test_qnode.py
+++ b/tests/test_qnode.py
@@ -43,7 +43,7 @@ class TestValidation:
 
     def test_invalid_interface(self):
         """Test that an exception is raised for an invalid interface"""
-        dev = qml.device("default.qubit", wires=1)
+        dev = qml.device("default.qubit.legacy", wires=1)
         test_interface = "something"
         expected_error = rf"Unknown interface {test_interface}\. Interface must be one of"
 
@@ -53,7 +53,7 @@ def test_invalid_interface(self):
     def test_changing_invalid_interface(self):
         """Test that an exception is raised for an invalid interface
         on a pre-existing QNode"""
-        dev = qml.device("default.qubit", wires=1)
+        dev = qml.device("default.qubit.legacy", wires=1)
         test_interface = "something"
 
         @qnode(dev)
@@ -72,7 +72,7 @@ def test_valid_interface(self):
         """Test that changing to a valid interface works as expected, and the
         diff method is updated as required."""
 
-        dev = qml.device("default.qubit", wires=1)
+        dev = qml.device("default.qubit.legacy", wires=1)
 
         @qnode(dev, interface="autograd", diff_method="best")
         def circuit(x):
@@ -95,7 +95,7 @@ def test_invalid_device(self):
     def test_validate_device_method(self, monkeypatch):
         """Test that the method for validating the device diff method
         tape works as expected"""
-        dev = qml.device("default.qubit", wires=1)
+        dev = qml.device("default.qubit.legacy", wires=1)
 
         with pytest.raises(
             qml.QuantumFunctionError,
@@ -124,7 +124,7 @@ def test_validate_backprop_method_invalid_device(self):
     def test_validate_backprop_method_invalid_interface(self, monkeypatch):
         """Test that the method for validating the backprop diff method
         tape raises an exception if the wrong interface is provided"""
-        dev = qml.device("default.qubit", wires=1)
+        dev = qml.device("default.qubit.legacy", wires=1)
         test_interface = "something"
 
         monkeypatch.setitem(dev._capabilities, "passthru_interface", test_interface)
@@ -136,7 +136,7 @@ def test_validate_backprop_method_invalid_interface(self, monkeypatch):
     def test_validate_backprop_method(self, monkeypatch):
         """Test that the method for validating the backprop diff method
         tape works as expected"""
-        dev = qml.device("default.qubit", wires=1)
+        dev = qml.device("default.qubit.legacy", wires=1)
         test_interface = "something"
         monkeypatch.setitem(dev._capabilities, "passthru_interface", test_interface)
 
@@ -154,7 +154,7 @@ def test_validate_backprop_method_all_interface_names(self, accepted_name, offic
         if accepted_name in {None, "auto", "scipy"}:
             pytest.skip("None is not a backprop interface.")
 
-        dev = qml.device("default.qubit", wires=1)
+        dev = qml.device("default.qubit.legacy", wires=1)
 
         diff_method, _, new_dev = QNode._validate_backprop_method(dev, accepted_name)
 
@@ -165,7 +165,7 @@ def test_validate_backprop_method_all_interface_names(self, accepted_name, offic
     def test_validate_backprop_child_method(self, monkeypatch):
         """Test that the method for validating the backprop diff method
         tape works as expected if a child device supports backprop"""
-        dev = qml.device("default.qubit", wires=1)
+        dev = qml.device("default.qubit.legacy", wires=1)
         test_interface = "something"
 
         orig_capabilities = dev.capabilities().copy()
@@ -182,7 +182,7 @@ def test_validate_backprop_child_method(self, monkeypatch):
     def test_validate_backprop_child_method_wrong_interface(self, monkeypatch):
         """Test that the method for validating the backprop diff method
         tape raises an error if a child device supports backprop but using a different interface"""
-        dev = qml.device("default.qubit", wires=1)
+        dev = qml.device("default.qubit.legacy", wires=1)
         test_interface = "something"
 
         orig_capabilities = dev.capabilities().copy()
@@ -196,7 +196,7 @@ def test_validate_backprop_child_method_wrong_interface(self, monkeypatch):
 
     # pylint: disable=protected-access
     @pytest.mark.autograd
-    @pytest.mark.parametrize("device_string", ("default.qubit", "default.qubit.autograd"))
+    @pytest.mark.parametrize("device_string", ("default.qubit.legacy", "default.qubit.autograd"))
     def test_validate_backprop_finite_shots(self, device_string):
         """Test that a device with finite shots cannot be used with backpropagation."""
         dev = qml.device(device_string, wires=1, shots=100)
@@ -209,7 +209,7 @@ def test_validate_backprop_finite_shots(self, device_string):
     def test_parameter_shift_qubit_device(self):
         """Test that the _validate_parameter_shift method
         returns the correct gradient transform for qubit devices."""
-        dev = qml.device("default.qubit", wires=1)
+        dev = qml.device("default.qubit.legacy", wires=1)
         gradient_fn = QNode._validate_parameter_shift(dev)
         assert gradient_fn[0] is qml.gradients.param_shift
 
@@ -242,7 +242,7 @@ def capabilities(cls):
             return capabilities
 
         monkeypatch.setattr(qml.devices.DefaultQubitLegacy, "capabilities", capabilities)
-        dev = qml.device("default.qubit", wires=1)
+        dev = qml.device("default.qubit.legacy", wires=1)
 
         with pytest.raises(
             qml.QuantumFunctionError, match="does not support the parameter-shift rule"
@@ -254,7 +254,7 @@ def capabilities(cls):
     def test_best_method_is_device(self, monkeypatch):
         """Test that the method for determining the best diff method
         for a given device and interface returns the device"""
-        dev = qml.device("default.qubit", wires=1)
+        dev = qml.device("default.qubit.legacy", wires=1)
         monkeypatch.setitem(dev._capabilities, "passthru_interface", "some_interface")
         monkeypatch.setitem(dev._capabilities, "provides_jacobian", True)
 
@@ -270,7 +270,7 @@ def test_best_method_is_device(self, monkeypatch):
     def test_best_method_is_backprop(self, monkeypatch):
         """Test that the method for determining the best diff method
         for a given device and interface returns backpropagation"""
-        dev = qml.device("default.qubit", wires=1)
+        dev = qml.device("default.qubit.legacy", wires=1)
         monkeypatch.setitem(dev._capabilities, "passthru_interface", "some_interface")
         monkeypatch.setitem(dev._capabilities, "provides_jacobian", False)
 
@@ -282,7 +282,7 @@ def test_best_method_is_backprop(self, monkeypatch):
     def test_best_method_is_param_shift(self, monkeypatch):
         """Test that the method for determining the best diff method
         for a given device and interface returns the parameter shift rule"""
-        dev = qml.device("default.qubit", wires=1)
+        dev = qml.device("default.qubit.legacy", wires=1)
         monkeypatch.setitem(dev._capabilities, "passthru_interface", "some_interface")
         monkeypatch.setitem(dev._capabilities, "provides_jacobian", False)
 
@@ -295,7 +295,7 @@ def test_best_method_is_param_shift(self, monkeypatch):
     def test_best_method_is_finite_diff(self, monkeypatch):
         """Test that the method for determining the best diff method
         for a given device and interface returns finite differences"""
-        dev = qml.device("default.qubit", wires=1)
+        dev = qml.device("default.qubit.legacy", wires=1)
         monkeypatch.setitem(dev._capabilities, "passthru_interface", "some_interface")
         monkeypatch.setitem(dev._capabilities, "provides_jacobian", False)
 
@@ -313,7 +313,7 @@ def capabilities(cls):
     def test_best_method_str_is_device(self, monkeypatch):
         """Test that the method for determining the best diff method string
         for a given device and interface returns 'device'"""
-        dev = qml.device("default.qubit", wires=1)
+        dev = qml.device("default.qubit.legacy", wires=1)
         monkeypatch.setitem(dev._capabilities, "passthru_interface", "some_interface")
         monkeypatch.setitem(dev._capabilities, "provides_jacobian", True)
 
@@ -329,7 +329,7 @@ def test_best_method_str_is_device(self, monkeypatch):
     def test_best_method_str_is_backprop(self, monkeypatch):
         """Test that the method for determining the best diff method string
         for a given device and interface returns 'backprop'"""
-        dev = qml.device("default.qubit", wires=1)
+        dev = qml.device("default.qubit.legacy", wires=1)
         monkeypatch.setitem(dev._capabilities, "passthru_interface", "some_interface")
         monkeypatch.setitem(dev._capabilities, "provides_jacobian", False)
 
@@ -341,7 +341,7 @@ def test_best_method_str_is_backprop(self, monkeypatch):
     def test_best_method_str_is_param_shift(self, monkeypatch):
         """Test that the method for determining the best diff method string
         for a given device and interface returns 'parameter-shift'"""
-        dev = qml.device("default.qubit", wires=1)
+        dev = qml.device("default.qubit.legacy", wires=1)
         monkeypatch.setitem(dev._capabilities, "passthru_interface", "some_interface")
         monkeypatch.setitem(dev._capabilities, "provides_jacobian", False)
 
@@ -354,7 +354,7 @@ def test_best_method_str_is_param_shift(self, monkeypatch):
     def test_best_method_str_is_finite_diff(self, monkeypatch):
         """Test that the method for determining the best diff method string
         for a given device and interface returns 'finite-diff'"""
-        dev = qml.device("default.qubit", wires=1)
+        dev = qml.device("default.qubit.legacy", wires=1)
         monkeypatch.setitem(dev._capabilities, "passthru_interface", "some_interface")
         monkeypatch.setitem(dev._capabilities, "provides_jacobian", False)
 
@@ -372,7 +372,7 @@ def capabilities(cls):
     def test_diff_method(self, mocker):
         """Test that a user-supplied diff method correctly returns the right
         diff method."""
-        dev = qml.device("default.qubit", wires=1)
+        dev = qml.device("default.qubit.legacy", wires=1)
 
         mock_best = mocker.patch("pennylane.QNode.get_best_method")
         mock_best.return_value = ("best", {}, dev)
@@ -437,7 +437,7 @@ def test_diff_method(self, mocker):
     @pytest.mark.autograd
     def test_gradient_transform(self, mocker):
         """Test passing a gradient transform directly to a QNode"""
-        dev = qml.device("default.qubit", wires=1)
+        dev = qml.device("default.qubit.legacy", wires=1)
         spy = mocker.spy(qml.gradients.finite_difference, "finite_diff_coeffs")
 
         @qnode(dev, diff_method=qml.gradients.finite_diff)
@@ -451,7 +451,7 @@ def circuit(x):
 
     def test_unknown_diff_method_string(self):
         """Test that an exception is raised for an unknown differentiation method string"""
-        dev = qml.device("default.qubit", wires=1)
+        dev = qml.device("default.qubit.legacy", wires=1)
 
         with pytest.raises(
             qml.QuantumFunctionError, match="Differentiation method hello not recognized"
@@ -460,7 +460,7 @@ def test_unknown_diff_method_string(self):
 
     def test_unknown_diff_method_type(self):
         """Test that an exception is raised for an unknown differentiation method type"""
-        dev = qml.device("default.qubit", wires=1)
+        dev = qml.device("default.qubit.legacy", wires=1)
 
         with pytest.raises(
             qml.QuantumFunctionError,
@@ -480,7 +480,7 @@ def test_validate_adjoint_invalid_device(self):
     def test_validate_adjoint_finite_shots(self):
         """Test that a UserWarning is raised when device has finite shots"""
 
-        dev = qml.device("default.qubit", wires=1, shots=1)
+        dev = qml.device("default.qubit.legacy", wires=1, shots=1)
 
         with pytest.warns(
             UserWarning, match="Requested adjoint differentiation to be computed with finite shots."
@@ -492,7 +492,7 @@ def test_adjoint_finite_shots(self):
         on QNode construction when the device has finite shots
         """
 
-        dev = qml.device("default.qubit", wires=1, shots=1)
+        dev = qml.device("default.qubit.legacy", wires=1, shots=1)
 
         @qnode(dev, diff_method="adjoint")
         def circ():
@@ -507,7 +507,7 @@ def circ():
     def test_sparse_diffmethod_error(self):
         """Test that an error is raised when the observable is SparseHamiltonian and the
         differentiation method is not parameter-shift."""
-        dev = qml.device("default.qubit", wires=2, shots=None)
+        dev = qml.device("default.qubit.legacy", wires=2, shots=None)
 
         @qnode(dev, diff_method="backprop")
         def circuit(param):
@@ -523,7 +523,7 @@ def circuit(param):
 
     def test_qnode_print(self):
         """Test that printing a QNode object yields the right information."""
-        dev = qml.device("default.qubit", wires=1)
+        dev = qml.device("default.qubit.legacy", wires=1)
 
         def func(x):
             qml.RX(x, wires=0)
@@ -533,7 +533,7 @@ def func(x):
 
         assert (
             repr(qn)
-            == "<QNode: wires=1, device='default.qubit', interface='auto', diff_method='best'>"
+            == "<QNode: wires=1, device='default.qubit.legacy', interface='auto', diff_method='best'>"
         )
 
         qn = QNode(func, dev, interface="autograd")
@@ -547,7 +547,7 @@ def func(x):
     def test_diff_method_none(self, tol):
         """Test that diff_method=None creates a QNode with no interface, and no
         device swapping."""
-        dev = qml.device("default.qubit", wires=1)
+        dev = qml.device("default.qubit.legacy", wires=1)
 
         @qnode(dev, diff_method=None)
         def circuit(x):
@@ -569,7 +569,7 @@ def circuit(x):
     # pylint: disable=unused-variable
     def test_unrecognized_kwargs_raise_warning(self):
         """Test that passing gradient_kwargs not included in qml.gradients.SUPPORTED_GRADIENT_KWARGS raises warning"""
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
 
         with warnings.catch_warnings(record=True) as w:
 
@@ -585,7 +585,7 @@ def circuit(params):
     def test_incorrect_diff_method_kwargs_raise_warning(self):
         """Tests that using one of the incorrect kwargs previously used in some examples in PennyLane
         (grad_method, gradient_fn) to set the qnode diff_method raises a warning"""
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
 
         with warnings.catch_warnings(record=True) as w:
 
@@ -605,7 +605,7 @@ def circuit2(params):
 
     def test_auto_interface_tracker_device_switched(self):
         """Test that checks that the tracker is switched to the new device."""
-        dev = qml.device("default.qubit", wires=1)
+        dev = qml.device("default.qubit.legacy", wires=1)
 
         @qml.qnode(dev)
         def circuit(params):
@@ -627,7 +627,7 @@ def circuit(params):
 
     def test_autograd_interface_device_switched_no_warnings(self):
         """Test that checks that no warning is raised for device switch when you define an interface."""
-        dev = qml.device("default.qubit", wires=1)
+        dev = qml.device("default.qubit.legacy", wires=1)
 
         @qml.qnode(dev, interface="autograd")
         def circuit(params):
@@ -642,7 +642,7 @@ def circuit(params):
     def test_not_giving_mode_kwarg_does_not_raise_warning(self):
         """Test that not providing a value for mode does not raise a warning."""
         with warnings.catch_warnings(record=True) as record:
-            _ = qml.QNode(lambda f: f, qml.device("default.qubit", wires=1))
+            _ = qml.QNode(lambda f: f, qml.device("default.qubit.legacy", wires=1))
 
         assert len(record) == 0
 
@@ -652,7 +652,7 @@ class TestTapeConstruction:
 
     def test_basic_tape_construction(self, tol):
         """Test that a quantum tape is properly constructed"""
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
 
         def func(x, y):
             qml.RX(x, wires=0)
@@ -685,7 +685,7 @@ def func(x, y):
 
     def test_jacobian(self):
         """Test the jacobian computation"""
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
 
         def func(x, y):
             qml.RX(x, wires=0)
@@ -709,7 +709,7 @@ def func(x, y):
     def test_returning_non_measurements(self):
         """Test that an exception is raised if a non-measurement
         is returned from the QNode."""
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
 
         def func0(x, y):
             qml.RX(x, wires=0)
@@ -753,7 +753,7 @@ def func3(x, y):
     def test_inconsistent_measurement_order(self):
         """Test that an exception is raised if measurements are returned in an
         order different to how they were queued on the tape"""
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
 
         def func(x, y):
             qml.RX(x, wires=0)
@@ -773,7 +773,7 @@ def func(x, y):
     def test_consistent_measurement_order(self):
         """Test evaluation proceeds as expected if measurements are returned in the
         same order to how they were queued on the tape"""
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
 
         contents = []
 
@@ -804,13 +804,13 @@ def circuit(x):
             qml.RX(x, wires=0)
             return qml.expval(qml.PauliZ(0))
 
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
         qn = QNode(circuit, dev)
 
         with pytest.raises(qml.QuantumFunctionError, match="Operator RX must act on all wires"):
             qn(0.5)
 
-        dev = qml.device("default.qubit", wires=1)
+        dev = qml.device("default.qubit.legacy", wires=1)
         qn = QNode(circuit, dev)
         assert np.allclose(qn(0.5), np.cos(0.5), atol=tol, rtol=0)
 
@@ -842,7 +842,7 @@ def test_jit_counts_raises_error(self):
         jitting raises an error."""
         import jax
 
-        dev = qml.device("default.qubit", wires=2, shots=5)
+        dev = qml.device("default.qubit.legacy", wires=2, shots=5)
 
         def circuit1(param):
             qml.Hadamard(0)
@@ -876,7 +876,7 @@ def circuit2(param):
 
 def test_decorator(tol):
     """Test that the decorator correctly creates a QNode."""
-    dev = qml.device("default.qubit", wires=2)
+    dev = qml.device("default.qubit.legacy", wires=2)
 
     @qnode(dev)
     def func(x, y):
@@ -922,7 +922,7 @@ def func():
             qml.CNOT(wires=[0, 1])
             return qml.expval(qml.PauliZ(0))
 
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
         qn = QNode(func, dev, interface="autograd")
 
         for _ in range(2):
@@ -946,7 +946,7 @@ def func():
             qml.CNOT(wires=[0, 1])
             return qml.expval(qml.PauliZ(0))
 
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
         qn = QNode(func, dev, interface=interface)
         for _ in range(2):
             qn()
@@ -975,7 +975,7 @@ def func():
             qml.CNOT(wires=[0, 1])
             return qml.expval(qml.PauliZ(0))
 
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
         qn = QNode(func, dev, interface=interface)
         for _ in range(2):
             qn()
@@ -998,7 +998,7 @@ def test_num_exec_caching_device_swap(self):
         """Tests that if we swapped the original device (e.g., when
         diff_method='backprop') then the number of executions recorded is
         correct."""
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
 
         cache = {}
 
@@ -1019,7 +1019,7 @@ def test_num_exec_caching_device_swap_two_exec(self):
         """Tests that if we swapped the original device (e.g., when
         diff_method='backprop') then the number of executions recorded is
         correct even with multiple QNode evaluations."""
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
 
         cache = {}
 
@@ -1054,7 +1054,7 @@ def test_single_expectation_value_with_argnum_one(self, diff_method, tol):
         This test relies on the fact that exactly one term of the estimated
         jacobian will match the expected analytical value.
         """
-        dev = qml.device("default.qubit", wires=3)
+        dev = qml.device("default.qubit.legacy", wires=3)
 
         x = pnp.array(0.543, requires_grad=True)
         y = pnp.array(-0.654, requires_grad=True)
@@ -1078,7 +1078,7 @@ def circuit(x, y):
     def test_no_defer_measurements_if_supported(self, mocker):
         """Test that the defer_measurements transform is not used during
         QNode construction if the device supports mid-circuit measurements."""
-        dev = qml.device("default.qubit", wires=3)
+        dev = qml.device("default.qubit.legacy", wires=3)
         mocker.patch.object(qml.Device, "_capabilities", {"supports_mid_measure": True})
         spy = mocker.spy(qml, "defer_measurements")
 
@@ -1146,7 +1146,7 @@ def conditional_ry_qnode(x, y):
     def test_drawing_has_deferred_measurements(self):
         """Test that `qml.draw` with qnodes uses defer_measurements
         to draw circuits with mid-circuit measurements."""
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
 
         @qml.qnode(dev)
         def circuit(x):
@@ -1164,7 +1164,7 @@ def circuit(x):
     def test_sampling_with_mcm(self, basis_state, mocker):
         """Tests that a QNode with qml.sample and mid-circuit measurements
         returns the expected results."""
-        dev = qml.device("default.qubit", wires=3, shots=1000)
+        dev = qml.device("default.qubit.legacy", wires=3, shots=1000)
 
         first_par = np.pi
 
@@ -1196,7 +1196,7 @@ def test_conditional_ops_tensorflow(self, interface):
         """Test conditional operations with TensorFlow."""
         import tensorflow as tf
 
-        dev = qml.device("default.qubit", wires=3)
+        dev = qml.device("default.qubit.legacy", wires=3)
 
         @qml.qnode(dev, interface=interface, diff_method="parameter-shift")
         def cry_qnode(x):
@@ -1239,7 +1239,7 @@ def test_conditional_ops_torch(self, interface):
         """Test conditional operations with Torch."""
         import torch
 
-        dev = qml.device("default.qubit", wires=3)
+        dev = qml.device("default.qubit.legacy", wires=3)
 
         @qml.qnode(dev, interface=interface, diff_method="parameter-shift")
         def cry_qnode(x):
@@ -1278,7 +1278,7 @@ def test_conditional_ops_jax(self, jax_interface):
         import jax
 
         jnp = jax.numpy
-        dev = qml.device("default.qubit", wires=3)
+        dev = qml.device("default.qubit.legacy", wires=3)
 
         @qml.qnode(dev, interface=jax_interface, diff_method="parameter-shift")
         def cry_qnode(x):
@@ -1309,7 +1309,7 @@ def conditional_ry_qnode(x):
 
     def test_qnode_does_not_support_nested_queuing(self):
         """Test that operators in QNodes are not queued to surrounding contexts."""
-        dev = qml.device("default.qubit", wires=1)
+        dev = qml.device("default.qubit.legacy", wires=1)
 
         @qml.qnode(dev)
         def circuit():
@@ -1329,7 +1329,7 @@ class TestShots:
     # pylint: disable=unexpected-keyword-arg
     def test_specify_shots_per_call_sample(self):
         """Tests that shots can be set per call for a sample return type."""
-        dev = qml.device("default.qubit", wires=1, shots=10)
+        dev = qml.device("default.qubit.legacy", wires=1, shots=10)
 
         @qnode(dev)
         def circuit(a):
@@ -1345,7 +1345,7 @@ def circuit(a):
     def test_specify_shots_per_call_expval(self):
         """Tests that shots can be set per call for an expectation value.
         Note: this test has a vanishingly small probability to fail."""
-        dev = qml.device("default.qubit", wires=1, shots=None)
+        dev = qml.device("default.qubit.legacy", wires=1, shots=None)
 
         @qnode(dev)
         def circuit():
@@ -1371,7 +1371,7 @@ def test_no_shots_per_call_if_user_has_shots_qfunc_kwarg(self):
         """Tests that the per-call shots overwriting is suspended if user
         has a shots keyword argument, but a warning is raised."""
 
-        dev = qml.device("default.qubit", wires=2, shots=10)
+        dev = qml.device("default.qubit.legacy", wires=2, shots=10)
 
         def circuit(a, shots=0):
             qml.RX(a, wires=shots)
@@ -1395,7 +1395,7 @@ def circuit(a, shots=0):
     def test_no_shots_per_call_if_user_has_shots_qfunc_arg(self):
         """Tests that the per-call shots overwriting is suspended
         if user has a shots argument, but a warning is raised."""
-        dev = qml.device("default.qubit", wires=[0, 1], shots=10)
+        dev = qml.device("default.qubit.legacy", wires=[0, 1], shots=10)
 
         def ansatz0(a, shots):
             qml.RX(a, wires=shots)
@@ -1410,7 +1410,7 @@ def ansatz0(a, shots):
         assert len(circuit(0.8, 1)) == 10
         assert circuit.qtape.operations[0].wires.labels == (1,)
 
-        dev = qml.device("default.qubit", wires=2, shots=10)
+        dev = qml.device("default.qubit.legacy", wires=2, shots=10)
 
         with pytest.warns(
             UserWarning, match="The 'shots' argument name is reserved for overriding"
@@ -1428,7 +1428,7 @@ def ansatz1(a, shots):
     def test_shots_setting_does_not_mutate_device(self):
         """Tests that per-call shots setting does not change the number of shots in the device."""
 
-        dev = qml.device("default.qubit", wires=1, shots=3)
+        dev = qml.device("default.qubit.legacy", wires=1, shots=3)
 
         @qnode(dev)
         def circuit(a):
@@ -1442,7 +1442,7 @@ def circuit(a):
 
     def test_warning_finite_shots_dev(self):
         """Tests that a warning is raised when caching is used with finite shots."""
-        dev = qml.device("default.qubit", wires=1, shots=5)
+        dev = qml.device("default.qubit.legacy", wires=1, shots=5)
 
         @qml.qnode(dev, cache={})
         def circuit(x):
@@ -1457,7 +1457,7 @@ def circuit(x):
     # pylint: disable=unexpected-keyword-arg
     def test_warning_finite_shots_override(self):
         """Tests that a warning is raised when caching is used with finite shots."""
-        dev = qml.device("default.qubit", wires=1, shots=5)
+        dev = qml.device("default.qubit.legacy", wires=1, shots=5)
 
         @qml.qnode(dev, cache={})
         def circuit(x):
@@ -1471,7 +1471,7 @@ def circuit(x):
 
     def test_warning_finite_shots_tape(self):
         """Tests that a warning is raised when caching is used with finite shots."""
-        dev = qml.device("default.qubit", wires=1, shots=5)
+        dev = qml.device("default.qubit.legacy", wires=1, shots=5)
 
         with qml.queuing.AnnotatedQueue() as q:
             qml.RZ(0.3, wires=0)
@@ -1486,7 +1486,7 @@ def test_warning_finite_shots_tape(self):
 
     def test_no_warning_infinite_shots(self):
         """Tests that no warning is raised when caching is used with infinite shots."""
-        dev = qml.device("default.qubit", wires=1)
+        dev = qml.device("default.qubit.legacy", wires=1)
 
         @qml.qnode(dev, cache={})
         def circuit(x):
@@ -1501,7 +1501,7 @@ def circuit(x):
     @pytest.mark.autograd
     def test_no_warning_internal_cache_reuse(self):
         """Tests that no warning is raised when only the internal cache is reused."""
-        dev = qml.device("default.qubit", wires=1, shots=5)
+        dev = qml.device("default.qubit.legacy", wires=1, shots=5)
 
         @qml.qnode(dev, cache=True)
         def circuit(x):
@@ -1524,7 +1524,7 @@ def circuit(x):
     )
     def test_tape_shots_set_on_call(self, shots, total_shots, shot_vector):
         """test that shots are placed on the tape if they are specified during a call."""
-        dev = qml.device("default.qubit", wires=2, shots=5)
+        dev = qml.device("default.qubit.legacy", wires=2, shots=5)
 
         def func(x, y):
             qml.RX(x, wires=0)
@@ -1562,7 +1562,7 @@ def qn2(x, y):
 class TestTransformProgramIntegration:
     def test_transform_program_modifies_circuit(self):
         """Test qnode integration with a transform that turns the circuit into just a pauli x."""
-        dev = qml.device("default.qubit", wires=1)
+        dev = qml.device("default.qubit.legacy", wires=1)
 
         def null_postprocessing(results):
             return results[0]
@@ -1595,7 +1595,7 @@ def circuit(x):
     def tet_transform_program_modifies_results(self):
         """Test integration with a transform that modifies the result output."""
 
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
 
         @qml.transforms.core.transform
         def pin_result(
@@ -1620,7 +1620,7 @@ def circuit(x):
     def test_transform_order_circuit_processing(self):
         """Test that transforms are applied in the correct order in integration."""
 
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
 
         def null_postprocessing(results):
             return results[0]
@@ -1667,7 +1667,7 @@ def circuit2(x):
     def test_transform_order_postprocessing(self):
         """Test that transform postprocessing is called in the right order."""
 
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
 
         def scale_by_factor(results, factor):
             return results[0] * factor
@@ -1883,7 +1883,7 @@ class TestTapeExpansion:
     )
     def test_device_expansion(self, diff_method, mode, mocker):
         """Test expansion of an unsupported operation on the device"""
-        dev = qml.device("default.qubit", wires=1)
+        dev = qml.device("default.qubit.legacy", wires=1)
 
         # pylint: disable=too-few-public-methods
         class UnsupportedOp(qml.operation.Operation):
@@ -1916,7 +1916,7 @@ def circuit(x):
     def test_no_gradient_expansion(self, mocker):
         """Test that an unsupported operation with defined gradient recipe is
         not expanded"""
-        dev = qml.device("default.qubit", wires=1)
+        dev = qml.device("default.qubit.legacy", wires=1)
 
         # pylint: disable=too-few-public-methods
         class UnsupportedOp(qml.operation.Operation):
@@ -1961,7 +1961,7 @@ def circuit(x):
     def test_gradient_expansion(self, mocker):
         """Test that a *supported* operation with no gradient recipe is
         expanded when applying the gradient transform, but not for execution."""
-        dev = qml.device("default.qubit", wires=1)
+        dev = qml.device("default.qubit.legacy", wires=1)
 
         # pylint: disable=too-few-public-methods
         class PhaseShift(qml.PhaseShift):
@@ -2006,7 +2006,7 @@ def circuit(x):
 
     def test_hamiltonian_expansion_analytic(self):
         """Test result if there are non-commuting groups and the number of shots is None"""
-        dev = qml.device("default.qubit", wires=3, shots=None)
+        dev = qml.device("default.qubit.legacy", wires=3, shots=None)
 
         obs = [qml.PauliX(0), qml.PauliX(0) @ qml.PauliZ(1), qml.PauliZ(0) @ qml.PauliZ(1)]
         c = np.array([-0.6543, 0.24, 0.54])
@@ -2025,7 +2025,7 @@ def circuit():
     def test_hamiltonian_expansion_finite_shots(self, mocker):
         """Test that the Hamiltonian is expanded if there
         are non-commuting groups and the number of shots is finite"""
-        dev = qml.device("default.qubit", wires=3, shots=50000)
+        dev = qml.device("default.qubit.legacy", wires=3, shots=50000)
 
         obs = [qml.PauliX(0), qml.PauliX(0) @ qml.PauliZ(1), qml.PauliZ(0) @ qml.PauliZ(1)]
         c = np.array([-0.6543, 0.24, 0.54])
@@ -2050,7 +2050,7 @@ def circuit():
     def test_invalid_hamiltonian_expansion_finite_shots(self):
         """Test that an error is raised if multiple expectations are requested
         when using finite shots"""
-        dev = qml.device("default.qubit", wires=3, shots=50000)
+        dev = qml.device("default.qubit.legacy", wires=3, shots=50000)
 
         obs = [qml.PauliX(0), qml.PauliX(0) @ qml.PauliZ(1), qml.PauliZ(0) @ qml.PauliZ(1)]
         c = np.array([-0.6543, 0.24, 0.54])
@@ -2071,7 +2071,7 @@ def circuit():
     def test_device_expansion_strategy(self, mocker):
         """Test that the device expansion strategy performs the device
         decomposition at construction time, and not at execution time"""
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
         x = pnp.array(0.5, requires_grad=True)
 
         @qnode(dev, diff_method="parameter-shift", expansion_strategy="device")
@@ -2097,7 +2097,7 @@ def circuit(x):
     def test_expansion_multiple_qwc_observables(self, mocker):
         """Test that the QNode correctly expands tapes that return
         multiple measurements of commuting observables"""
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
         obs = [qml.PauliX(0), qml.PauliX(0) @ qml.PauliY(1)]
 
         @qml.qnode(dev)
diff --git a/tests/test_return_types.py b/tests/test_return_types.py
index 733ff8e6119..6fcbe5298fd 100644
--- a/tests/test_return_types.py
+++ b/tests/test_return_types.py
@@ -22,7 +22,7 @@
 
 test_wires = [2, 3, 4]
 
-devices = ["default.qubit", "default.mixed"]
+devices = ["default.qubit.legacy", "default.mixed"]
 
 
 @pytest.mark.parametrize("interface, shots", [["autograd", None], ["auto", 100]])
@@ -32,7 +32,7 @@ class TestSingleReturnExecute:
     @pytest.mark.parametrize("wires", test_wires)
     def test_state_default(self, wires, interface, shots):
         """Return state with default.qubit."""
-        dev = qml.device("default.qubit", wires=wires, shots=shots)
+        dev = qml.device("default.qubit.legacy", wires=wires, shots=shots)
 
         def circuit(x):
             qml.Hadamard(wires=[0])
@@ -223,7 +223,7 @@ def test_sample(self, measurement, interface, shots):
         if shots is None:
             pytest.skip("Sample requires finite shots.")
 
-        dev = qml.device("default.qubit", wires=2, shots=shots)
+        dev = qml.device("default.qubit.legacy", wires=2, shots=shots)
 
         def circuit(x):
             qml.Hadamard(wires=[0])
@@ -244,7 +244,7 @@ def test_counts(self, measurement, interface, shots):
         if shots is None:
             pytest.skip("Counts requires finite shots.")
 
-        dev = qml.device("default.qubit", wires=2, shots=shots)
+        dev = qml.device("default.qubit.legacy", wires=2, shots=shots)
 
         def circuit(x):
             qml.Hadamard(wires=[0])
@@ -1221,7 +1221,7 @@ def return_type(self):
             DummyMeasurement(obs=qml.PauliZ(0))
 
         tape = qml.tape.QuantumScript.from_queue(q)
-        dev = qml.device("default.qubit", wires=3)
+        dev = qml.device("default.qubit.legacy", wires=3)
         with pytest.raises(
             qml.QuantumFunctionError, match="Unsupported return type specified for observable"
         ):
@@ -1231,7 +1231,7 @@ def test_state_return_with_other_types(self):
         """Test that an exception is raised when a state is returned along with another return
         type"""
 
-        dev = qml.device("default.qubit", wires=2)
+        dev = qml.device("default.qubit.legacy", wires=2)
 
         with qml.queuing.AnnotatedQueue() as q:
             qml.PauliX(wires=0)
@@ -1248,7 +1248,7 @@ def test_state_return_with_other_types(self):
     def test_entropy_no_custom_wires(self):
         """Test that entropy cannot be returned with custom wires."""
 
-        dev = qml.device("default.qubit", wires=["a", 1])
+        dev = qml.device("default.qubit.legacy", wires=["a", 1])
 
         with qml.queuing.AnnotatedQueue() as q:
             qml.PauliX(wires="a")
@@ -1264,7 +1264,7 @@ def test_entropy_no_custom_wires(self):
     def test_custom_wire_labels_error(self):
         """Tests that an error is raised when mutual information is measured
         with custom wire labels"""
-        dev = qml.device("default.qubit", wires=["a", "b"])
+        dev = qml.device("default.qubit.legacy", wires=["a", "b"])
 
         with qml.queuing.AnnotatedQueue() as q:
             qml.PauliX(wires="a")
diff --git a/tests/test_return_types_dq2.py b/tests/test_return_types_dq2.py
new file mode 100644
index 00000000000..3acae93dc59
--- /dev/null
+++ b/tests/test_return_types_dq2.py
@@ -0,0 +1,1265 @@
+# Copyright 2022 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 new return types.
+"""
+import numpy as np
+import pytest
+
+import pennylane as qml
+from pennylane.measurements import MeasurementProcess
+
+test_wires = [2, 3, 4]
+
+devices = ["default.qubit"]
+
+
+@pytest.mark.parametrize("interface, shots", [["autograd", None], ["auto", 100]])
+class TestSingleReturnExecute:
+    """Test that single measurements return behavior does not change."""
+
+    @pytest.mark.parametrize("wires", test_wires)
+    def test_state_default(self, wires, interface, shots):
+        """Return state with default.qubit."""
+        dev = qml.device("default.qubit", wires=wires, shots=shots)
+
+        def circuit(x):
+            qml.Hadamard(wires=[0])
+            qml.CRX(x, wires=[0, 1])
+            return qml.state()
+
+        qnode = qml.QNode(circuit, dev)
+        qnode.construct([0.5], {})
+
+        if dev.shots:
+            pytest.skip("cannot return analytic measurements with finite shots.")
+        res = qml.execute(tapes=[qnode.tape], device=dev, gradient_fn=None, interface=interface)
+
+        assert res[0].shape == (2**wires,)
+        assert isinstance(res[0], (np.ndarray, np.float64))
+
+    @pytest.mark.parametrize("device", devices)
+    @pytest.mark.parametrize("d_wires", test_wires)
+    def test_density_matrix(self, d_wires, device, interface, shots):
+        """Return density matrix."""
+        dev = qml.device(device, wires=4, shots=shots)
+
+        def circuit(x):
+            qml.Hadamard(wires=[0])
+            qml.CRX(x, wires=[0, 1])
+            return qml.density_matrix(wires=range(0, d_wires))
+
+        qnode = qml.QNode(circuit, dev)
+        qnode.construct([0.5], {})
+
+        if dev.shots:
+            pytest.skip("cannot return analytic measurements with finite shots.")
+        res = qml.execute(tapes=[qnode.tape], device=dev, gradient_fn=None, interface=interface)
+
+        assert res[0].shape == (2**d_wires, 2**d_wires)
+        assert isinstance(res[0], (np.ndarray, np.float64))
+
+    @pytest.mark.parametrize("device", devices)
+    def test_expval(self, device, interface, shots):
+        """Return a single expval."""
+        dev = qml.device(device, wires=2, shots=shots)
+
+        def circuit(x):
+            qml.Hadamard(wires=[0])
+            qml.CRX(x, wires=[0, 1])
+            return qml.expval(qml.PauliZ(wires=1))
+
+        qnode = qml.QNode(circuit, dev)
+        qnode.construct([0.5], {})
+
+        res = qml.execute(tapes=[qnode.tape], device=dev, gradient_fn=None, interface=interface)
+
+        assert res[0].shape == ()
+        assert isinstance(res[0], (np.ndarray, np.float64))
+
+    @pytest.mark.parametrize("device", devices)
+    def test_var(self, device, interface, shots):
+        """Return a single var."""
+        dev = qml.device(device, wires=2, shots=shots)
+
+        def circuit(x):
+            qml.Hadamard(wires=[0])
+            qml.CRX(x, wires=[0, 1])
+            return qml.var(qml.PauliZ(wires=1))
+
+        qnode = qml.QNode(circuit, dev)
+        qnode.construct([0.5], {})
+
+        res = qml.execute(tapes=[qnode.tape], device=dev, gradient_fn=None, interface=interface)
+
+        assert res[0].shape == ()
+        assert isinstance(res[0], (np.ndarray, np.float64))
+
+    @pytest.mark.parametrize("device", devices)
+    def test_vn_entropy(self, device, interface, shots):
+        """Return a single vn entropy."""
+        dev = qml.device(device, wires=2, shots=shots)
+
+        def circuit(x):
+            qml.Hadamard(wires=[0])
+            qml.CRX(x, wires=[0, 1])
+            return qml.vn_entropy(wires=0)
+
+        qnode = qml.QNode(circuit, dev)
+        qnode.construct([0.5], {})
+
+        if dev.shots:
+            pytest.skip("cannot return analytic measurements with finite shots.")
+        res = qml.execute(tapes=[qnode.tape], device=dev, gradient_fn=None, interface=interface)
+
+        assert res[0].shape == ()
+        assert isinstance(res[0], (np.ndarray, np.float64))
+
+    @pytest.mark.parametrize("device", devices)
+    def test_mutual_info(self, device, interface, shots):
+        """Return a single mutual information."""
+        dev = qml.device(device, wires=2, shots=shots)
+
+        def circuit(x):
+            qml.Hadamard(wires=[0])
+            qml.CRX(x, wires=[0, 1])
+            return qml.mutual_info(wires0=[0], wires1=[1])
+
+        qnode = qml.QNode(circuit, dev)
+        qnode.construct([0.5], {})
+
+        if dev.shots:
+            pytest.skip("cannot return analytic measurements with finite shots.")
+        res = qml.execute(tapes=[qnode.tape], device=dev, gradient_fn=None, interface=interface)
+
+        assert res[0].shape == ()
+        assert isinstance(res[0], (np.ndarray, np.float64))
+
+    herm = np.diag([1, 2, 3, 4])
+    probs_data = [
+        (None, [0]),
+        (None, [0, 1]),
+        (qml.PauliZ(0), None),
+        (qml.Hermitian(herm, wires=[1, 0]), None),
+    ]
+
+    # pylint: disable=too-many-arguments
+    @pytest.mark.parametrize("device", devices)
+    @pytest.mark.parametrize("op,wires", probs_data)
+    def test_probs(self, op, wires, device, interface, shots):
+        """Return a single prob."""
+        dev = qml.device(device, wires=3, shots=shots)
+
+        def circuit(x):
+            qml.Hadamard(wires=[0])
+            qml.CRX(x, wires=[0, 1])
+            return qml.probs(op=op, wires=wires)
+
+        qnode = qml.QNode(circuit, dev)
+        qnode.construct([0.5], {})
+
+        res = qml.execute(tapes=[qnode.tape], device=dev, gradient_fn=None, interface=interface)
+
+        if wires is None:
+            wires = op.wires
+
+        assert res[0].shape == (2 ** len(wires),)
+        assert isinstance(res[0], (np.ndarray, np.float64))
+
+    @pytest.mark.parametrize("measurement", [qml.sample(qml.PauliZ(0)), qml.sample(wires=[0])])
+    def test_sample(self, measurement, interface, shots):
+        """Test the sample measurement."""
+        if shots is None:
+            pytest.skip("Sample requires finite shots.")
+
+        dev = qml.device("default.qubit", wires=2, shots=shots)
+
+        def circuit(x):
+            qml.Hadamard(wires=[0])
+            qml.CRX(x, wires=[0, 1])
+            return qml.apply(measurement)
+
+        qnode = qml.QNode(circuit, dev)
+        qnode.construct([0.5], {})
+
+        res = qml.execute(tapes=[qnode.tape], device=dev, gradient_fn=None, interface=interface)
+
+        assert isinstance(res[0], (np.ndarray, np.float64))
+        assert res[0].shape == (shots,)
+
+    @pytest.mark.parametrize("measurement", [qml.counts(qml.PauliZ(0)), qml.counts(wires=[0])])
+    def test_counts(self, measurement, interface, shots):
+        """Test the counts measurement."""
+        if shots is None:
+            pytest.skip("Counts requires finite shots.")
+
+        dev = qml.device("default.qubit", wires=2, shots=shots)
+
+        def circuit(x):
+            qml.Hadamard(wires=[0])
+            qml.CRX(x, wires=[0, 1])
+            return qml.apply(measurement)
+
+        qnode = qml.QNode(circuit, dev)
+        qnode.construct([0.5], {})
+
+        res = qml.execute(tapes=[qnode.tape], device=dev, gradient_fn=None, interface=interface)
+
+        assert isinstance(res[0], dict)
+        assert sum(res[0].values()) == shots
+
+
+multi_return_wires = [([0], [1]), ([1], [0]), ([0], [0]), ([1], [1])]
+
+
+@pytest.mark.parametrize("shots", [None, 100])
+class TestMultipleReturns:
+    """Test the new return types for multiple measurements, it should always return a tuple containing the single
+    measurements.
+    """
+
+    @pytest.mark.parametrize("device", devices)
+    def test_multiple_expval(self, device, shots):
+        """Return multiple expvals."""
+        dev = qml.device(device, wires=2, shots=shots)
+
+        def circuit(x):
+            qml.Hadamard(wires=[0])
+            qml.CRX(x, wires=[0, 1])
+            return qml.expval(qml.PauliZ(wires=0)), qml.expval(qml.PauliZ(wires=1))
+
+        qnode = qml.QNode(circuit, dev)
+        qnode.construct([0.5], {})
+
+        res = qml.execute(tapes=[qnode.tape], device=dev, gradient_fn=None)
+
+        assert isinstance(res[0], tuple)
+        assert len(res[0]) == 2
+
+        assert isinstance(res[0][0], (np.ndarray, np.float64))
+        assert res[0][0].shape == ()
+
+        assert isinstance(res[0][1], (np.ndarray, np.float64))
+        assert res[0][1].shape == ()
+
+    @pytest.mark.parametrize("device", devices)
+    def test_multiple_var(self, device, shots):
+        """Return multiple vars."""
+        dev = qml.device(device, wires=2, shots=shots)
+
+        def circuit(x):
+            qml.Hadamard(wires=[0])
+            qml.CRX(x, wires=[0, 1])
+            return qml.var(qml.PauliZ(wires=0)), qml.var(qml.PauliZ(wires=1))
+
+        qnode = qml.QNode(circuit, dev)
+        qnode.construct([0.5], {})
+
+        res = qml.execute(tapes=[qnode.tape], device=dev, gradient_fn=None)
+
+        assert isinstance(res[0], tuple)
+        assert len(res[0]) == 2
+
+        assert isinstance(res[0][0], (np.ndarray, np.float64))
+        assert res[0][0].shape == ()
+
+        assert isinstance(res[0][1], (np.ndarray, np.float64))
+        assert res[0][1].shape == ()
+
+    # op1, wires1, op2, wires2
+    multi_probs_data = [
+        (None, [0], None, [0]),
+        (None, [0], None, [0, 1]),
+        (None, [0, 1], None, [0]),
+        (None, [0, 1], None, [0, 1]),
+        (qml.PauliZ(0), None, qml.PauliZ(1), None),
+        (None, [0], qml.PauliZ(1), None),
+        (qml.PauliZ(0), None, None, [0]),
+        (qml.PauliZ(1), None, qml.PauliZ(0), None),
+    ]
+
+    # pylint: disable=too-many-arguments
+    @pytest.mark.parametrize("device", devices)
+    @pytest.mark.parametrize("op1,wires1,op2,wires2", multi_probs_data)
+    def test_multiple_prob(self, op1, op2, wires1, wires2, device, shots):
+        """Return multiple probs."""
+        dev = qml.device(device, wires=2, shots=shots)
+
+        def circuit(x):
+            qml.Hadamard(wires=[0])
+            qml.CRX(x, wires=[0, 1])
+            return qml.probs(op=op1, wires=wires1), qml.probs(op=op2, wires=wires2)
+
+        qnode = qml.QNode(circuit, dev)
+        qnode.construct([0.5], {})
+
+        res = qml.execute(tapes=[qnode.tape], device=dev, gradient_fn=None)
+
+        assert isinstance(res[0], tuple)
+        assert len(res[0]) == 2
+
+        if wires1 is None:
+            wires1 = op1.wires
+
+        if wires2 is None:
+            wires2 = op2.wires
+
+        assert isinstance(res[0][0], (np.ndarray, np.float64))
+        assert res[0][0].shape == (2 ** len(wires1),)
+
+        assert isinstance(res[0][1], (np.ndarray, np.float64))
+        assert res[0][1].shape == (2 ** len(wires2),)
+
+    # pylint: disable=too-many-arguments
+    @pytest.mark.parametrize("device", devices)
+    @pytest.mark.parametrize("op1,wires1,op2,wires2", multi_probs_data)
+    @pytest.mark.parametrize("wires3, wires4", multi_return_wires)
+    def test_mix_meas(self, op1, wires1, op2, wires2, wires3, wires4, device, shots):
+        """Return multiple different measurements."""
+
+        dev = qml.device(device, wires=2, shots=shots)
+
+        def circuit(x):
+            qml.Hadamard(wires=[0])
+            qml.CRX(x, wires=[0, 1])
+            return (
+                qml.probs(op=op1, wires=wires1),
+                qml.vn_entropy(wires=wires3),
+                qml.probs(op=op2, wires=wires2),
+                qml.expval(qml.PauliZ(wires=wires4)),
+            )
+
+        qnode = qml.QNode(circuit, dev)
+        qnode.construct([0.5], {})
+
+        if dev.shots:
+            pytest.skip("cannot return analytic measurements with finite shots.")
+        res = qml.execute(tapes=[qnode.tape], device=dev, gradient_fn=None)
+
+        if wires1 is None:
+            wires1 = op1.wires
+
+        if wires2 is None:
+            wires2 = op2.wires
+
+        assert isinstance(res[0], tuple)
+        assert len(res[0]) == 4
+
+        assert isinstance(res[0][0], (np.ndarray, np.float64))
+        assert res[0][0].shape == (2 ** len(wires1),)
+
+        assert isinstance(res[0][1], (np.ndarray, np.float64))
+        assert res[0][1].shape == ()
+
+        assert isinstance(res[0][2], (np.ndarray, np.float64))
+        assert res[0][2].shape == (2 ** len(wires2),)
+
+        assert isinstance(res[0][3], (np.ndarray, np.float64))
+        assert res[0][3].shape == ()
+
+    wires = [2, 3, 4, 5]
+
+    @pytest.mark.parametrize("device", devices)
+    @pytest.mark.parametrize("wires", wires)
+    def test_list_multiple_expval(self, wires, device, shots):
+        """Return a comprehension list of multiple expvals."""
+        dev = qml.device(device, wires=wires, shots=shots)
+
+        def circuit(x):
+            qml.Hadamard(wires=[0])
+            qml.CRX(x, wires=[0, 1])
+            return [qml.expval(qml.PauliZ(wires=i)) for i in range(0, wires)]
+
+        qnode = qml.QNode(circuit, dev)
+        qnode.construct([0.5], {})
+
+        res = qml.execute(tapes=[qnode.tape], device=dev, gradient_fn=None)
+
+        assert isinstance(res[0], tuple)
+        assert len(res[0]) == wires
+
+        for i in range(0, wires):
+            assert isinstance(res[0][i], (np.ndarray, np.float64))
+            assert res[0][i].shape == ()
+
+    @pytest.mark.parametrize("device", devices)
+    @pytest.mark.parametrize("measurement", [qml.sample(qml.PauliZ(0)), qml.sample(wires=[0])])
+    def test_expval_sample(self, measurement, shots, device):
+        """Test the expval and sample measurements together."""
+        if shots is None:
+            pytest.skip("Sample requires finite shots.")
+
+        dev = qml.device(device, wires=2, shots=shots)
+
+        def circuit(x):
+            qml.Hadamard(wires=[0])
+            qml.CRX(x, wires=[0, 1])
+            return qml.expval(qml.PauliX(1)), qml.apply(measurement)
+
+        qnode = qml.QNode(circuit, dev)
+        qnode.construct([0.5], {})
+
+        res = qml.execute(tapes=[qnode.tape], device=dev, gradient_fn=None)
+
+        # Expval
+        assert isinstance(res[0][0], (np.ndarray, np.float64))
+        assert res[0][0].shape == ()
+
+        # Sample
+        assert isinstance(res[0][1], (np.ndarray, np.float64))
+        assert res[0][1].shape == (shots,)
+
+    @pytest.mark.parametrize("device", devices)
+    @pytest.mark.parametrize("measurement", [qml.counts(qml.PauliZ(0)), qml.counts(wires=[0])])
+    def test_expval_counts(self, measurement, shots, device):
+        """Test the expval and counts measurements together."""
+        if shots is None:
+            pytest.skip("Counts requires finite shots.")
+
+        dev = qml.device(device, wires=2, shots=shots)
+
+        def circuit(x):
+            qml.Hadamard(wires=[0])
+            qml.CRX(x, wires=[0, 1])
+            return qml.expval(qml.PauliX(1)), qml.apply(measurement)
+
+        qnode = qml.QNode(circuit, dev)
+        qnode.construct([0.5], {})
+
+        res = qml.execute(tapes=[qnode.tape], device=dev, gradient_fn=None)
+
+        # Expval
+        assert isinstance(res[0][0], (np.ndarray, np.float64))
+        assert res[0][0].shape == ()
+
+        # Counts
+        assert isinstance(res[0][1], dict)
+        assert sum(res[0][1].values()) == shots
+
+
+pauliz = qml.PauliZ(wires=1)
+proj = qml.Projector([1], wires=1)
+hermitian = qml.Hermitian(np.diag([1, 2]), wires=0)
+
+# Note: mutual info and vn_entropy do not support some shot vectors
+# qml.mutual_info(wires0=[0], wires1=[1]), qml.vn_entropy(wires=[0])]
+single_scalar_output_measurements = [
+    qml.expval(pauliz),
+    qml.var(pauliz),
+    qml.expval(proj),
+    qml.var(proj),
+    qml.expval(hermitian),
+    qml.var(hermitian),
+]
+
+herm = np.diag([1, 2, 3, 4])
+probs_data = [
+    (None, [0]),
+    (None, [0, 1]),
+    (qml.PauliZ(0), None),
+    (qml.Hermitian(herm, wires=[1, 0]), None),
+]
+
+shot_vectors = [[10, 1000], [1, 10, 10, 1000], [1, (10, 2), 1000]]
+
+
+@pytest.mark.parametrize("shot_vector", shot_vectors)
+@pytest.mark.parametrize("device", devices)
+class TestShotVector:
+    """Test the support for executing tapes with single measurements using a
+    device with shot vectors."""
+
+    @pytest.mark.parametrize("measurement", single_scalar_output_measurements)
+    def test_scalar(self, shot_vector, measurement, device):
+        """Test a single scalar-valued measurement."""
+        dev = qml.device(device, wires=2, shots=shot_vector)
+
+        def circuit(x):
+            qml.Hadamard(wires=[0])
+            qml.CRX(x, wires=[0, 1])
+            return qml.apply(measurement)
+
+        qnode = qml.QNode(circuit, dev)
+        qnode.construct([0.5], {})
+
+        res = qml.execute(tapes=[qnode.tape], device=dev, gradient_fn=None)
+
+        all_shots = sum(shot_tuple.copies for shot_tuple in dev.shots.shot_vector)
+
+        assert isinstance(res[0], tuple)
+        assert len(res[0]) == all_shots
+        assert all(r.shape == () for r in res[0])
+
+    @pytest.mark.parametrize("op,wires", probs_data)
+    def test_probs(self, shot_vector, op, wires, device):
+        """Test a single probability measurement."""
+        dev = qml.device(device, wires=2, shots=shot_vector)
+
+        def circuit(x):
+            qml.Hadamard(wires=[0])
+            qml.CRX(x, wires=[0, 1])
+            return qml.probs(op=op, wires=wires)
+
+        qnode = qml.QNode(circuit, dev)
+        qnode.construct([0.5], {})
+
+        res = qml.execute(tapes=[qnode.tape], device=dev, gradient_fn=None)
+
+        all_shots = sum(shot_tuple.copies for shot_tuple in dev.shots.shot_vector)
+
+        assert isinstance(res[0], tuple)
+        assert len(res[0]) == all_shots
+        wires_to_use = wires if wires else op.wires
+        assert all(r.shape == (2 ** len(wires_to_use),) for r in res[0])
+
+    @pytest.mark.parametrize("wires", [[0], [2, 0], [1, 0], [2, 0, 1]])
+    def test_density_matrix(self, shot_vector, wires, device):
+        """Test a density matrix measurement."""
+        dev = qml.device(device, wires=3, shots=shot_vector)
+
+        def circuit(x):
+            qml.Hadamard(wires=[0])
+            qml.CRX(x, wires=[0, 1])
+            return qml.density_matrix(wires=wires)
+
+        qnode = qml.QNode(circuit, dev)
+        qnode.construct([0.5], {})
+
+        if dev.shots:
+            pytest.skip("cannot return analytic measurements with finite shots.")
+        res = qml.execute(tapes=[qnode.tape], device=dev, gradient_fn=None)
+
+        all_shots = sum(shot_tuple.copies for shot_tuple in dev.shots.shot_vector)
+
+        assert isinstance(res[0], tuple)
+        assert len(res[0]) == all_shots
+        dim = 2 ** len(wires)
+        assert all(r.shape == (dim, dim) for r in res[0])
+
+    @pytest.mark.parametrize("measurement", [qml.sample(qml.PauliZ(0)), qml.sample(wires=[0])])
+    def test_samples(self, shot_vector, measurement, device):
+        """Test the sample measurement."""
+        dev = qml.device(device, wires=2, shots=shot_vector)
+
+        def circuit(x):
+            qml.Hadamard(wires=[0])
+            qml.CRX(x, wires=[0, 1])
+            return qml.apply(measurement)
+
+        qnode = qml.QNode(circuit, dev)
+        qnode.construct([0.5], {})
+
+        res = qml.execute(tapes=[qnode.tape], device=dev, gradient_fn=None)
+
+        all_shot_copies = [
+            shot_tuple.shots
+            for shot_tuple in dev.shots.shot_vector
+            for _ in range(shot_tuple.copies)
+        ]
+
+        assert len(res[0]) == len(all_shot_copies)
+        for r, shots in zip(res[0], all_shot_copies):
+            if shots == 1:
+                # Scalar tensors
+                assert r.shape == ()
+            else:
+                assert r.shape == (shots,)
+
+    @pytest.mark.parametrize("measurement", [qml.counts(qml.PauliZ(0)), qml.counts(wires=[0])])
+    def test_counts(self, shot_vector, measurement, device):
+        """Test the counts measurement."""
+        dev = qml.device(device, wires=2, shots=shot_vector)
+
+        def circuit(x):
+            qml.Hadamard(wires=[0])
+            qml.CRX(x, wires=[0, 1])
+            return qml.apply(measurement)
+
+        qnode = qml.QNode(circuit, dev)
+        qnode.construct([0.5], {})
+
+        res = qml.execute(tapes=[qnode.tape], device=dev, gradient_fn=None)
+
+        all_shots = sum(shot_tuple.copies for shot_tuple in dev.shots.shot_vector)
+
+        assert isinstance(res[0], tuple)
+        assert len(res[0]) == all_shots
+        assert all(isinstance(r, dict) for r in res[0])
+
+
+@pytest.mark.parametrize("shot_vector", shot_vectors)
+@pytest.mark.parametrize("device", devices)
+class TestSameMeasurementShotVector:
+    """Test the support for executing tapes with the same type of measurement
+    multiple times using a device with shot vectors"""
+
+    def test_scalar(self, shot_vector, device):
+        """Test multiple scalar-valued measurements."""
+        dev = qml.device(device, wires=2, shots=shot_vector)
+
+        def circuit(x):
+            qml.Hadamard(wires=[0])
+            qml.CRX(x, wires=[0, 1])
+            return qml.expval(qml.PauliX(0)), qml.var(qml.PauliZ(1))
+
+        qnode = qml.QNode(circuit, dev)
+        qnode.construct([0.5], {})
+
+        res = qml.execute(tapes=[qnode.tape], device=dev, gradient_fn=None)
+
+        all_shots = sum(shot_tuple.copies for shot_tuple in dev.shots.shot_vector)
+
+        assert isinstance(res[0], tuple)
+        assert len(res[0]) == all_shots
+        for r in res[0]:
+            assert len(r) == 2
+            assert all(r.shape == () for r in r)
+
+    probs_data2 = [
+        (None, [2]),
+        (None, [2, 3]),
+        (qml.PauliZ(2), None),
+        (qml.Hermitian(herm, wires=[3, 2]), None),
+    ]
+
+    # pylint: disable=too-many-arguments
+    @pytest.mark.parametrize("op1,wires1", probs_data)
+    @pytest.mark.parametrize("op2,wires2", reversed(probs_data2))
+    def test_probs(self, shot_vector, op1, wires1, op2, wires2, device):
+        """Test multiple probability measurements."""
+        dev = qml.device(device, wires=4, shots=shot_vector)
+
+        def circuit(x):
+            qml.Hadamard(wires=[0])
+            qml.CRX(x, wires=[0, 1])
+            return qml.probs(op=op1, wires=wires1), qml.probs(op=op2, wires=wires2)
+
+        qnode = qml.QNode(circuit, dev)
+        qnode.construct([0.5], {})
+
+        res = qml.execute(tapes=[qnode.tape], device=dev, gradient_fn=None)
+
+        all_shots = sum(shot_tuple.copies for shot_tuple in dev.shots.shot_vector)
+
+        assert isinstance(res[0], tuple)
+        assert len(res[0]) == all_shots
+
+        wires1 = wires1 if wires1 else op1.wires
+        wires2 = wires2 if wires2 else op2.wires
+        for r in res[0]:
+            assert len(r) == 2
+            assert r[0].shape == (2 ** len(wires1),)
+            assert r[1].shape == (2 ** len(wires2),)
+
+    @pytest.mark.parametrize("measurement1", [qml.sample(qml.PauliZ(0)), qml.sample(wires=[0])])
+    @pytest.mark.parametrize("measurement2", [qml.sample(qml.PauliX(1)), qml.sample(wires=[1])])
+    def test_samples(self, shot_vector, measurement1, measurement2, device):
+        """Test multiple sample measurements."""
+        dev = qml.device(device, wires=2, shots=shot_vector)
+
+        def circuit(x):
+            qml.Hadamard(wires=[0])
+            qml.CRX(x, wires=[0, 1])
+            return qml.apply(measurement1), qml.apply(measurement2)
+
+        qnode = qml.QNode(circuit, dev)
+        qnode.construct([0.5], {})
+
+        res = qml.execute(tapes=[qnode.tape], device=dev, gradient_fn=None)
+
+        all_shot_copies = [
+            shot_tuple.shots
+            for shot_tuple in dev.shots.shot_vector
+            for _ in range(shot_tuple.copies)
+        ]
+
+        assert len(res[0]) == len(all_shot_copies)
+        for r, shots in zip(res[0], all_shot_copies):
+            shape = () if shots == 1 else (shots,)
+            assert all(res_item.shape == shape for res_item in r)
+
+    @pytest.mark.parametrize("measurement1", [qml.counts(qml.PauliZ(0)), qml.counts(wires=[0])])
+    @pytest.mark.parametrize("measurement2", [qml.counts(qml.PauliZ(0)), qml.counts(wires=[0])])
+    def test_counts(self, shot_vector, measurement1, measurement2, device):
+        """Test multiple counts measurements."""
+        dev = qml.device(device, wires=2, shots=shot_vector)
+
+        def circuit(x):
+            qml.Hadamard(wires=[0])
+            qml.CRX(x, wires=[0, 1])
+            return qml.apply(measurement1), qml.apply(measurement2)
+
+        qnode = qml.QNode(circuit, dev)
+        qnode.construct([0.5], {})
+
+        res = qml.execute(tapes=[qnode.tape], device=dev, gradient_fn=None)
+
+        all_shots = sum(shot_tuple.copies for shot_tuple in dev.shots.shot_vector)
+
+        assert isinstance(res[0], tuple)
+        assert len(res[0]) == all_shots
+        for r in res[0]:
+            assert isinstance(r, tuple)
+            assert all(isinstance(res_item, dict) for res_item in r)
+
+
+# -------------------------------------------------
+# Shot vector multi measurement tests - test data
+# -------------------------------------------------
+
+pauliz_w2 = qml.PauliZ(wires=2)
+proj_w2 = qml.Projector([1], wires=2)
+hermitian = qml.Hermitian(np.diag([1, 2]), wires=0)
+tensor_product = qml.PauliZ(wires=2) @ qml.PauliX(wires=1)
+
+# Expval/Var with Probs
+
+scalar_probs_multi = [
+    # Expval
+    (qml.expval(pauliz_w2), qml.probs(wires=[2, 0])),
+    (qml.expval(proj_w2), qml.probs(wires=[1, 0])),
+    (qml.expval(tensor_product), qml.probs(wires=[2, 0])),
+    # Var
+    (qml.var(qml.PauliZ(wires=1)), qml.probs(wires=[0, 1])),
+    (qml.var(proj_w2), qml.probs(wires=[1, 0])),
+    (qml.var(tensor_product), qml.probs(wires=[2, 0])),
+]
+
+# Expval/Var with Sample
+
+scalar_sample_multi = [
+    # Expval
+    (qml.expval(pauliz_w2), qml.sample(op=qml.PauliZ(1) @ qml.PauliZ(0))),
+    (qml.expval(proj_w2), qml.sample(op=qml.PauliZ(1) @ qml.PauliZ(0))),
+    (qml.expval(tensor_product), qml.sample(op=qml.PauliZ(0))),
+    # Var
+    (qml.var(proj_w2), qml.sample(op=qml.PauliZ(1) @ qml.PauliZ(0))),
+    (qml.var(pauliz_w2), qml.sample(op=qml.PauliZ(1) @ qml.PauliZ(0))),
+    (qml.var(tensor_product), qml.sample(op=qml.PauliZ(0))),
+]
+
+scalar_sample_no_obs_multi = [
+    (qml.expval(qml.PauliZ(wires=1)), qml.sample()),
+    (qml.expval(qml.PauliZ(wires=1)), qml.sample(wires=[0, 1])),
+    (qml.var(qml.PauliZ(wires=1)), qml.sample(wires=[0, 1])),
+]
+
+# Expval/Var with Counts
+
+scalar_counts_multi = [
+    # Expval
+    (qml.expval(pauliz_w2), qml.counts(op=qml.PauliZ(1) @ qml.PauliZ(0))),
+    (qml.expval(proj_w2), qml.counts(op=qml.PauliZ(1) @ qml.PauliZ(0))),
+    (qml.expval(tensor_product), qml.counts(op=qml.PauliZ(0))),
+    # Var
+    (qml.var(proj_w2), qml.counts(op=qml.PauliZ(1) @ qml.PauliZ(0))),
+    (qml.var(pauliz_w2), qml.counts(op=qml.PauliZ(1) @ qml.PauliZ(0))),
+    (qml.var(tensor_product), qml.counts(op=qml.PauliZ(0))),
+]
+
+scalar_counts_no_obs_multi = [
+    (qml.expval(qml.PauliZ(wires=1)), qml.counts()),
+    (qml.expval(qml.PauliZ(wires=1)), qml.counts(wires=[0, 1])),
+    (qml.var(qml.PauliZ(wires=1)), qml.counts(wires=[0, 1])),
+]
+
+
+@pytest.mark.parametrize("shot_vector", shot_vectors)
+@pytest.mark.parametrize("device", devices)
+class TestMixMeasurementsShotVector:
+    """Test the support for executing tapes with multiple different
+    measurements using a device with shot vectors"""
+
+    @pytest.mark.parametrize("meas1,meas2", scalar_probs_multi)
+    def test_scalar_probs(self, shot_vector, meas1, meas2, device):
+        """Test scalar-valued and probability measurements"""
+        dev = qml.device(device, wires=3, shots=shot_vector)
+
+        def circuit(x):
+            qml.Hadamard(wires=[0])
+            qml.CRX(x, wires=[0, 1])
+            return qml.apply(meas1), qml.apply(meas2)
+
+        qnode = qml.QNode(circuit, dev)
+        qnode.construct([0.5], {})
+
+        res = qml.execute(tapes=[qnode.tape], device=dev, gradient_fn=None)
+
+        all_shots = sum(shot_tuple.copies for shot_tuple in dev.shots.shot_vector)
+
+        assert isinstance(res[0], tuple)
+        assert len(res[0]) == all_shots
+        assert all(isinstance(r, tuple) for r in res[0])
+        assert all(
+            isinstance(m, (np.ndarray, np.float64))
+            for measurement_res in res[0]
+            for m in measurement_res
+        )
+        for meas_res in res[0]:
+            for i, r in enumerate(meas_res):
+                if i % 2 == 0:
+                    # Scalar-val meas
+                    assert r.shape == ()
+                else:
+                    assert r.shape == (2**2,)
+
+                    # Probs add up to 1
+                    assert np.allclose(sum(r), 1)
+
+    @pytest.mark.parametrize("meas1,meas2", scalar_sample_multi)
+    def test_scalar_sample_with_obs(self, shot_vector, meas1, meas2, device):
+        """Test scalar-valued and sample measurements where sample takes an
+        observable."""
+        dev = qml.device(device, wires=3, shots=shot_vector)
+        raw_shot_vector = [
+            shot_tuple.shots
+            for shot_tuple in dev.shots.shot_vector
+            for _ in range(shot_tuple.copies)
+        ]
+
+        def circuit(x):
+            qml.Hadamard(wires=[0])
+            qml.CRX(x, wires=[0, 1])
+            return qml.apply(meas1), qml.apply(meas2)
+
+        qnode = qml.QNode(circuit, dev)
+        qnode.construct([0.5], {})
+
+        res = qml.execute(tapes=[qnode.tape], device=dev, gradient_fn=None)
+
+        all_shots = sum(shot_tuple.copies for shot_tuple in dev.shots.shot_vector)
+
+        assert isinstance(res[0], tuple)
+        assert len(res[0]) == all_shots
+        assert all(isinstance(r, tuple) for r in res[0])
+        assert all(
+            isinstance(m, (np.ndarray, np.float64))
+            for measurement_res in res[0]
+            for m in measurement_res
+        )
+
+        for idx, shots in enumerate(raw_shot_vector):
+            for i, r in enumerate(res[0][idx]):
+                if i % 2 == 0 or shots == 1:
+                    assert meas2.obs is not None
+                    expected_shape = ()
+                    assert r.shape == expected_shape
+                else:
+                    assert r.shape == (shots,)
+
+    @pytest.mark.parametrize("meas1,meas2", scalar_sample_no_obs_multi)
+    @pytest.mark.xfail
+    def test_scalar_sample_no_obs(self, shot_vector, meas1, meas2, device):
+        """Test scalar-valued and computational basis sample measurements."""
+        dev = qml.device(device, wires=3, shots=shot_vector)
+
+        def circuit(x):
+            qml.Hadamard(wires=[0])
+            qml.CRX(x, wires=[0, 1])
+            return qml.apply(meas1), qml.apply(meas2)
+
+        qnode = qml.QNode(circuit, dev)
+        qnode.construct([0.5], {})
+
+        res = qml.execute(tapes=[qnode.tape], device=dev, gradient_fn=None)
+
+        all_shots = sum(shot_tuple.copies for shot_tuple in dev.shots.shot_vector)
+
+        assert isinstance(res[0], tuple)
+        assert len(res[0]) == all_shots
+        assert all(isinstance(r, tuple) for r in res[0])
+        assert all(
+            isinstance(m, (np.ndarray, np.float64))
+            for measurement_res in res[0]
+            for m in measurement_res
+        )
+
+        for shot_tuple in dev.shots.shot_vector:
+            for idx in range(shot_tuple.copies):
+                for i, r in enumerate(res[0][idx]):
+                    if i % 2 == 0 or shot_tuple.shots == 1:
+                        assert meas2.obs is not None
+                        expected_shape = ()
+                        assert r.shape == expected_shape
+                    else:
+                        assert r.shape == (shot_tuple.shots,)
+
+    @pytest.mark.parametrize("meas1,meas2", scalar_counts_multi)
+    def test_scalar_counts_with_obs(self, shot_vector, meas1, meas2, device):
+        """Test scalar-valued and counts measurements where counts takes an
+        observable."""
+        dev = qml.device(device, wires=3, shots=shot_vector)
+        raw_shot_vector = [
+            shot_tuple.shots
+            for shot_tuple in dev.shots.shot_vector
+            for _ in range(shot_tuple.copies)
+        ]
+
+        def circuit(x):
+            qml.Hadamard(wires=[0])
+            qml.CRX(x, wires=[0, 1])
+            return qml.apply(meas1), qml.apply(meas2)
+
+        qnode = qml.QNode(circuit, dev)
+        qnode.construct([0.5], {})
+
+        res = qml.execute(tapes=[qnode.tape], device=dev, gradient_fn=None)
+
+        all_shots = sum(shot_tuple.copies for shot_tuple in dev.shots.shot_vector)
+
+        assert isinstance(res[0], tuple)
+        assert len(res[0]) == all_shots
+        assert all(isinstance(r, tuple) for r in res[0])
+
+        for r in res[0]:
+            assert isinstance(r[0], (np.ndarray, np.float64))
+            assert isinstance(r[1], dict)
+
+        expected_outcomes = {-1, 1}
+
+        for idx, shots in enumerate(raw_shot_vector):
+            for i, r in enumerate(res[0][idx]):
+                if i % 2 == 0:
+                    assert meas2.obs is not None
+                    expected_shape = ()
+                    assert r.shape == expected_shape
+                else:
+                    # Samples are either -1 or 1
+                    assert set(r.keys()).issubset(expected_outcomes)
+                    assert sum(r.values()) == shots
+
+    @pytest.mark.parametrize("meas1,meas2", scalar_counts_no_obs_multi)
+    def test_scalar_counts_no_obs(self, shot_vector, meas1, meas2, device):
+        """Test scalar-valued and computational basis counts measurements."""
+        dev = qml.device(device, wires=3, shots=shot_vector)
+        raw_shot_vector = [
+            shot_tuple.shots
+            for shot_tuple in dev.shots.shot_vector
+            for _ in range(shot_tuple.copies)
+        ]
+
+        def circuit(x):
+            qml.Hadamard(wires=[0])
+            qml.CRX(x, wires=[0, 1])
+            return qml.apply(meas1), qml.apply(meas2)
+
+        qnode = qml.QNode(circuit, dev)
+        qnode.construct([0.5], {})
+
+        res = qml.execute(tapes=[qnode.tape], device=dev, gradient_fn=None)
+
+        all_shots = sum(shot_tuple.copies for shot_tuple in dev.shots.shot_vector)
+
+        assert isinstance(res[0], tuple)
+        assert len(res[0]) == all_shots
+        assert all(isinstance(r, tuple) for r in res[0])
+
+        for idx, _ in enumerate(raw_shot_vector):
+            for i, r in enumerate(res[0][idx]):
+                if i % 2 == 0:
+                    assert isinstance(r, (np.ndarray, np.float64))
+                    assert meas2.obs is None
+                    expected_shape = ()
+                    assert r.shape == expected_shape
+                else:
+                    assert isinstance(r, dict)
+
+    @pytest.mark.parametrize("sample_obs", [qml.PauliZ, None])
+    def test_probs_sample(self, shot_vector, sample_obs, device):
+        """Test probs and sample measurements."""
+        dev = qml.device(device, wires=3, shots=shot_vector)
+        raw_shot_vector = [
+            shot_tuple.shots
+            for shot_tuple in dev.shots.shot_vector
+            for _ in range(shot_tuple.copies)
+        ]
+
+        meas1_wires = [0, 1]
+        meas2_wires = [2]
+
+        @qml.qnode(device=dev)
+        def circuit(x):
+            qml.Hadamard(wires=[0])
+            qml.CRX(x, wires=[0, 1])
+            if sample_obs is not None:
+                # Observable provided to sample
+                return qml.probs(wires=meas1_wires), qml.sample(sample_obs(meas2_wires))
+
+            # Only wires provided to sample
+            return qml.probs(wires=meas1_wires), qml.sample(wires=meas2_wires)
+
+        qnode = qml.QNode(circuit, dev)
+        qnode.construct([0.5], {})
+        res = qml.execute(tapes=[qnode.tape], device=dev, gradient_fn=None)
+
+        all_shots = sum(shot_tuple.copies for shot_tuple in dev.shots.shot_vector)
+
+        assert isinstance(res[0], tuple)
+        assert len(res[0]) == all_shots
+        assert all(isinstance(r, tuple) for r in res[0])
+        assert all(
+            isinstance(m, (np.ndarray, np.float64))
+            for measurement_res in res[0]
+            for m in measurement_res
+        )
+
+        for idx, shots in enumerate(raw_shot_vector):
+            for i, r in enumerate(res[0][idx]):
+                if i % 2 == 0:
+                    expected_shape = (len(meas1_wires) ** 2,)
+                    assert r.shape == expected_shape
+
+                    # Probs add up to 1
+                    assert np.allclose(sum(r), 1)
+                else:
+                    if shots == 1:
+                        assert r.shape == ()
+                    else:
+                        expected = (shots,)
+                        assert r.shape == expected
+
+    @pytest.mark.parametrize("sample_obs", [qml.PauliZ, None])
+    def test_probs_counts(self, shot_vector, sample_obs, device):
+        """Test probs and counts measurements."""
+        dev = qml.device(device, wires=3, shots=shot_vector)
+        raw_shot_vector = [
+            shot_tuple.shots
+            for shot_tuple in dev.shots.shot_vector
+            for _ in range(shot_tuple.copies)
+        ]
+
+        meas1_wires = [0, 1]
+        meas2_wires = [2]
+
+        @qml.qnode(device=dev)
+        def circuit(x):
+            qml.Hadamard(wires=[0])
+            qml.CRX(x, wires=[0, 1])
+            if sample_obs is not None:
+                # Observable provided to sample
+                return qml.probs(wires=meas1_wires), qml.counts(sample_obs(meas2_wires))
+
+            # Only wires provided to sample
+            return qml.probs(wires=meas1_wires), qml.counts(wires=meas2_wires)
+
+        qnode = qml.QNode(circuit, dev)
+        qnode.construct([0.5], {})
+        res = qml.execute(tapes=[qnode.tape], device=dev, gradient_fn=None)
+
+        all_shots = sum(shot_tuple.copies for shot_tuple in dev.shots.shot_vector)
+
+        assert isinstance(res[0], tuple)
+        assert len(res[0]) == all_shots
+        assert all(isinstance(r, tuple) for r in res[0])
+        assert all(
+            isinstance(measurement_res[0], (np.ndarray, np.float64)) for measurement_res in res[0]
+        )
+        assert all(isinstance(measurement_res[1], dict) for measurement_res in res[0])
+
+        expected_outcomes = {-1, 1} if sample_obs is not None else {"0", "1"}
+        for idx, shots in enumerate(raw_shot_vector):
+            for i, r in enumerate(res[0][idx]):
+                if i % 2 == 0:
+                    expected_shape = (len(meas1_wires) ** 2,)
+                    assert r.shape == expected_shape
+
+                    # Probs add up to 1
+                    assert np.allclose(sum(r), 1)
+                else:
+                    # Samples are -1 or 1
+                    assert set(r.keys()).issubset(expected_outcomes)
+                    assert sum(r.values()) == shots
+
+    @pytest.mark.parametrize("sample_wires", [[1], [0, 2]])
+    @pytest.mark.parametrize("counts_wires", [[4], [3, 5]])
+    def test_sample_counts(self, shot_vector, sample_wires, counts_wires, device):
+        """Test sample and counts measurements, each measurement with custom
+        samples or computational basis state samples."""
+        dev = qml.device(device, wires=6, shots=shot_vector)
+        raw_shot_vector = [
+            shot_tuple.shots
+            for shot_tuple in dev.shots.shot_vector
+            for _ in range(shot_tuple.copies)
+        ]
+
+        @qml.qnode(device=dev)
+        def circuit(x):
+            qml.Hadamard(wires=[0])
+            qml.CRX(x, wires=[0, 1])
+
+            # 1. Sample obs and Counts obs
+            if len(sample_wires) == 1 and len(counts_wires) == 1:
+                return qml.sample(qml.PauliY(sample_wires)), qml.counts(qml.PauliX(counts_wires))
+
+            # 2. Sample no obs and Counts obs
+            if len(sample_wires) > 1 and len(counts_wires) == 1:
+                return qml.sample(wires=sample_wires), qml.counts(qml.PauliX(counts_wires))
+
+            # 3. Sample obs and Counts no obs
+            if len(sample_wires) == 1 and len(counts_wires) > 1:
+                return qml.sample(qml.PauliY(sample_wires)), qml.counts(wires=counts_wires)
+
+            # 4. Sample no obs and Counts no obs
+            return qml.sample(wires=sample_wires), qml.counts(wires=counts_wires)
+
+        qnode = qml.QNode(circuit, dev)
+        qnode.construct([0.5], {})
+        res = qml.execute(tapes=[qnode.tape], device=dev, gradient_fn=None)
+
+        all_shots = sum(shot_tuple.copies for shot_tuple in dev.shots.shot_vector)
+
+        assert isinstance(res[0], tuple)
+        assert len(res[0]) == all_shots
+        assert all(isinstance(r, tuple) for r in res[0])
+        assert all(
+            isinstance(measurement_res[0], (np.ndarray, np.float64)) for measurement_res in res[0]
+        )
+        assert all(isinstance(measurement_res[1], dict) for measurement_res in res[0])
+
+        for idx, shots in enumerate(raw_shot_vector):
+            for i, r in enumerate(res[0][idx]):
+                num_wires = len(sample_wires)
+                if shots == 1 and i % 2 == 0:
+                    expected_shape = () if num_wires == 1 else (num_wires,)
+                    assert r.shape == expected_shape
+                elif i % 2 == 0:
+                    expected_shape = (shots,) if num_wires == 1 else (shots, num_wires)
+                    assert r.shape == expected_shape
+                else:
+                    assert isinstance(r, dict)
+
+    @pytest.mark.parametrize("meas1,meas2", scalar_probs_multi)
+    def test_scalar_probs_sample_counts(self, shot_vector, meas1, meas2, device):
+        """Test scalar-valued, probability, sample and counts measurements all
+        in a single qfunc."""
+        dev = qml.device(device, wires=5, shots=shot_vector)
+        raw_shot_vector = [
+            shot_tuple.shots
+            for shot_tuple in dev.shots.shot_vector
+            for _ in range(shot_tuple.copies)
+        ]
+
+        def circuit(x):
+            qml.Hadamard(wires=[0])
+            qml.CRX(x, wires=[0, 1])
+            return (
+                qml.apply(meas1),
+                qml.apply(meas2),
+                qml.sample(qml.PauliX(4)),
+                qml.counts(qml.PauliX(3)),
+            )
+
+        qnode = qml.QNode(circuit, dev)
+        qnode.construct([0.5], {})
+
+        res = qml.execute(tapes=[qnode.tape], device=dev, gradient_fn=None)
+
+        all_shots = sum(shot_tuple.copies for shot_tuple in dev.shots.shot_vector)
+
+        assert isinstance(res[0], tuple)
+        assert len(res[0]) == all_shots
+        assert all(isinstance(r, tuple) for r in res[0])
+
+        for res_idx, meas_res in enumerate(res[0]):
+            for i, r in enumerate(meas_res):
+                num_meas = i % 4
+                expval_or_var = num_meas == 0
+                probs = num_meas == 1
+                sample = num_meas == 2
+
+                if expval_or_var:
+                    assert r.shape == ()
+                elif probs:
+                    assert r.shape == (2**2,)
+
+                    # Probs add up to 1
+                    assert np.allclose(sum(r), 1)
+                elif sample:
+                    shots = raw_shot_vector[res_idx]
+                    if shots == 1:
+                        assert r.shape == ()
+                    else:
+                        expected = (shots,)
+                        assert r.shape == expected
+                else:
+                    # Return is Counts
+                    assert isinstance(r, dict)
+
+
+class TestDeviceNewUnits:
+    """Further unit tests for some new methods of Device."""
+
+    def test_unsupported_observable_return_type_raise_error(self):
+        """Check that an error is raised if the return type of an observable is unsupported"""
+        # pylint: disable=too-few-public-methods
+
+        class UnsupportedReturnType:
+            value = "unsupported"
+
+        class DummyMeasurement(MeasurementProcess):
+            @property
+            def return_type(self):
+                return UnsupportedReturnType
+
+        with qml.queuing.AnnotatedQueue() as q:
+            qml.PauliX(wires=0)
+            DummyMeasurement(obs=qml.PauliZ(0))
+
+        tape = qml.tape.QuantumScript.from_queue(q)
+        dev = qml.device("default.qubit", wires=3)
+        with pytest.raises(qml.DeviceError, match="Analytic circuits must only contain"):
+            qml.execute(tapes=[tape], device=dev, gradient_fn=None)
+
+    def test_state_return_with_other_types(self):
+        """Test that an exception is raised when a state is returned along with another return
+        type"""
+
+        dev = qml.device("default.qubit", wires=2)
+
+        with qml.queuing.AnnotatedQueue() as q:
+            qml.PauliX(wires=0)
+            qml.state()
+            qml.expval(qml.PauliZ(1))
+
+        tape = qml.tape.QuantumScript.from_queue(q)
+        res = qml.execute(tapes=[tape], device=dev, gradient_fn=None)[0]
+        assert isinstance(res, tuple) and len(res) == 2
+        assert np.array_equal(res[0], [0.0, 0.0, 1.0, 0.0])
+        assert res[1] == 1.0
+
+    def test_entropy_no_custom_wires(self):
+        """Test that entropy cannot be returned with custom wires."""
+
+        dev = qml.device("default.qubit", wires=["a", 1])
+
+        with qml.queuing.AnnotatedQueue() as q:
+            qml.PauliX(wires="a")
+            qml.vn_entropy(wires=["a"])
+
+        tape = qml.tape.QuantumScript.from_queue(q)
+        res = qml.execute(tapes=[tape], device=dev, gradient_fn=None)
+        assert res == (0,)
+
+    def test_custom_wire_labels_error(self):
+        """Tests that an error is raised when mutual information is measured
+        with custom wire labels"""
+        dev = qml.device("default.qubit", wires=["a", "b"])
+
+        with qml.queuing.AnnotatedQueue() as q:
+            qml.PauliX(wires="a")
+            qml.mutual_info(wires0=["a"], wires1=["b"])
+
+        tape = qml.tape.QuantumScript.from_queue(q)
+        res = qml.execute(tapes=[tape], device=dev, gradient_fn=None)
+        assert res == (0,)
diff --git a/tests/test_return_types_qnode.py b/tests/test_return_types_qnode.py
index d4e2c11b86d..3a3cdb93997 100644
--- a/tests/test_return_types_qnode.py
+++ b/tests/test_return_types_qnode.py
@@ -50,7 +50,7 @@ def circuit(x):
         res = qnode(0.5)
 
         assert res.shape == (2**wires,)
-        assert isinstance(res, np.ndarray)
+        assert isinstance(res, (np.ndarray, np.float64))
 
     @pytest.mark.parametrize("wires", test_wires)
     def test_state_mixed(self, wires):
@@ -65,7 +65,7 @@ def circuit(x):
         res = qnode(0.5)
 
         assert res.shape == (2**wires, 2**wires)
-        assert isinstance(res, np.ndarray)
+        assert isinstance(res, (np.ndarray, np.float64))
 
     @pytest.mark.parametrize("device", devices)
     @pytest.mark.parametrize("d_wires", test_wires)
@@ -83,7 +83,7 @@ def circuit(x):
 
         dim = 3 if device == "default.qutrit" else 2
         assert res.shape == (dim**d_wires, dim**d_wires)
-        assert isinstance(res, np.ndarray)
+        assert isinstance(res, (np.ndarray, np.float64))
 
     @pytest.mark.parametrize("device", devices)
     def test_expval(self, device):
@@ -101,7 +101,7 @@ def circuit(x):
         res = qnode(0.5)
 
         assert res.shape == ()
-        assert isinstance(res, np.ndarray)
+        assert isinstance(res, (np.ndarray, np.float64))
 
     @pytest.mark.parametrize("device", devices)
     def test_var(self, device):
@@ -119,7 +119,7 @@ def circuit(x):
         res = qnode(0.5)
 
         assert res.shape == ()
-        assert isinstance(res, np.ndarray)
+        assert isinstance(res, (np.ndarray, np.float64))
 
     @pytest.mark.parametrize("device", devices)
     def test_vn_entropy(self, device):
@@ -137,7 +137,7 @@ def circuit(x):
         res = qnode(0.5)
 
         assert res.shape == ()
-        assert isinstance(res, np.ndarray)
+        assert isinstance(res, (np.ndarray, np.float64))
 
     @pytest.mark.xfail(reason="qml.execute shot vec support required with new return types")
     @pytest.mark.filterwarnings("ignore:Requested Von Neumann entropy with finite shots")
@@ -171,7 +171,7 @@ def circuit(x):
         res = qnode(0.5)
 
         assert res.shape == ()
-        assert isinstance(res, np.ndarray)
+        assert isinstance(res, (np.ndarray, np.float64))
 
     @pytest.mark.xfail(reason="qml.execute shot vec support required with new return types")
     @pytest.mark.filterwarnings("ignore:Requested mutual information with finite shots")
@@ -218,7 +218,7 @@ def circuit(x):
             wires = op.wires
 
         assert res.shape == (2 ** len(wires),)
-        assert isinstance(res, np.ndarray)
+        assert isinstance(res, (np.ndarray, np.float64))
 
     probs_data_qutrit = [
         (qml.GellMann(0, 3), None),
@@ -243,7 +243,7 @@ def circuit(x):
             wires = op.wires
 
         assert res.shape == (3 ** len(wires),)
-        assert isinstance(res, np.ndarray)
+        assert isinstance(res, (np.ndarray, np.float64))
 
     @pytest.mark.parametrize("device", devices)
     @pytest.mark.parametrize(
@@ -273,7 +273,7 @@ def circuit(x):
         qnode = qml.QNode(circuit, dev, diff_method=None)
         res = qnode(0.5)
 
-        assert isinstance(res, np.ndarray)
+        assert isinstance(res, (np.ndarray, np.float64))
 
         if measurement.wires.tolist() != [0, 1]:
             assert res.shape == (shots,)
@@ -1020,10 +1020,10 @@ def circuit(x):
         assert isinstance(res, tuple)
         assert len(res) == 2
 
-        assert isinstance(res[0], np.ndarray)
+        assert isinstance(res[0], (np.ndarray, np.float64))
         assert res[0].shape == ()
 
-        assert isinstance(res[1], np.ndarray)
+        assert isinstance(res[1], (np.ndarray, np.float64))
         assert res[1].shape == ()
 
     @pytest.mark.parametrize("device", devices)
@@ -1049,10 +1049,10 @@ def circuit(x):
         assert isinstance(res, tuple)
         assert len(res) == 2
 
-        assert isinstance(res[0], np.ndarray)
+        assert isinstance(res[0], (np.ndarray, np.float64))
         assert res[0].shape == ()
 
-        assert isinstance(res[1], np.ndarray)
+        assert isinstance(res[1], (np.ndarray, np.float64))
         assert res[1].shape == ()
 
     # op1, wires1, op2, wires2
@@ -1093,10 +1093,10 @@ def circuit(x):
         if wires2 is None:
             wires2 = op2.wires
 
-        assert isinstance(res[0], np.ndarray)
+        assert isinstance(res[0], (np.ndarray, np.float64))
         assert res[0].shape == (2 ** len(wires1),)
 
-        assert isinstance(res[1], np.ndarray)
+        assert isinstance(res[1], (np.ndarray, np.float64))
         assert res[1].shape == (2 ** len(wires2),)
 
     multi_probs_data_qutrit = [
@@ -1131,10 +1131,10 @@ def circuit(x):
         if wires2 is None:
             wires2 = op2.wires
 
-        assert isinstance(res[0], np.ndarray)
+        assert isinstance(res[0], (np.ndarray, np.float64))
         assert res[0].shape == (3 ** len(wires1),)
 
-        assert isinstance(res[1], np.ndarray)
+        assert isinstance(res[1], (np.ndarray, np.float64))
         assert res[1].shape == (3 ** len(wires2),)
 
     # pylint: disable=too-many-arguments
@@ -1169,16 +1169,16 @@ def circuit(x):
         assert isinstance(res, tuple)
         assert len(res) == 4
 
-        assert isinstance(res[0], np.ndarray)
+        assert isinstance(res[0], (np.ndarray, np.float64))
         assert res[0].shape == (2 ** len(wires1),)
 
-        assert isinstance(res[1], np.ndarray)
+        assert isinstance(res[1], (np.ndarray, np.float64))
         assert res[1].shape == ()
 
-        assert isinstance(res[2], np.ndarray)
+        assert isinstance(res[2], (np.ndarray, np.float64))
         assert res[2].shape == (2 ** len(wires2),)
 
-        assert isinstance(res[3], np.ndarray)
+        assert isinstance(res[3], (np.ndarray, np.float64))
         assert res[3].shape == ()
 
     # pylint: disable=too-many-arguments
@@ -1211,16 +1211,16 @@ def circuit(x):
         assert isinstance(res, tuple)
         assert len(res) == 4
 
-        assert isinstance(res[0], np.ndarray)
+        assert isinstance(res[0], (np.ndarray, np.float64))
         assert res[0].shape == (3 ** len(wires1),)
 
-        assert isinstance(res[1], np.ndarray)
+        assert isinstance(res[1], (np.ndarray, np.float64))
         assert res[1].shape == ()
 
-        assert isinstance(res[2], np.ndarray)
+        assert isinstance(res[2], (np.ndarray, np.float64))
         assert res[2].shape == (3 ** len(wires2),)
 
-        assert isinstance(res[3], np.ndarray)
+        assert isinstance(res[3], (np.ndarray, np.float64))
         assert res[3].shape == ()
 
     @pytest.mark.parametrize("device", devices)
@@ -1248,11 +1248,11 @@ def circuit(x):
         res = qnode(0.5)
 
         # Expval
-        assert isinstance(res[0], np.ndarray)
+        assert isinstance(res[0], (np.ndarray, np.float64))
         assert res[0].shape == ()
 
         # Sample
-        assert isinstance(res[1], np.ndarray)
+        assert isinstance(res[1], (np.ndarray, np.float64))
         assert res[1].shape == (shots,)
 
     @pytest.mark.parametrize("device", devices)
@@ -1280,7 +1280,7 @@ def circuit(x):
         res = qnode(0.5)
 
         # Expval
-        assert isinstance(res[0], np.ndarray)
+        assert isinstance(res[0], (np.ndarray, np.float64))
         assert res[0].shape == ()
 
         # Counts
@@ -1306,7 +1306,7 @@ def circuit(x):
 
         assert isinstance(res, list)
         assert len(res) == 1
-        assert isinstance(res[0], np.ndarray)
+        assert isinstance(res[0], (np.ndarray, np.float64))
         assert res[0].shape == ()
 
     shot_vectors = [None, [10, 1000], [1, 10, 10, 1000], [1, (10, 2), 1000]]
@@ -1322,6 +1322,7 @@ def test_list_multiple_expval(self, wires, device, shot_vector):
 
         def circuit(x):
             func(x)
+            # pylint:disable=unexpected-keyword-arg
             return [
                 qml.expval(obs(wires=i) if device != "default.qutrit" else obs(wires=i, index=3))
                 for i in range(0, wires)
@@ -1334,7 +1335,7 @@ def circuit(x):
             assert isinstance(res, list)
             assert len(res) == wires
             for r in res:
-                assert isinstance(r, np.ndarray)
+                assert isinstance(r, (np.ndarray, np.float64))
                 assert r.shape == ()
 
         else:
@@ -1343,7 +1344,7 @@ def circuit(x):
                 assert len(r) == wires
 
                 for t in r:
-                    assert isinstance(t, np.ndarray)
+                    assert isinstance(t, (np.ndarray, np.float64))
                     assert t.shape == ()
 
     @pytest.mark.parametrize("device", devices)
@@ -1373,9 +1374,9 @@ def circuit(x):
         else:
             assert isinstance(res[0], dict)
 
-        assert isinstance(res[1], qml.numpy.ndarray)
+        assert isinstance(res[1], (qml.numpy.ndarray, qml.numpy.float64))
         assert res[1].shape == ()
-        assert isinstance(res[2], qml.numpy.ndarray)
+        assert isinstance(res[2], (qml.numpy.ndarray, qml.numpy.float64))
         assert res[2].shape == (2,)
 
 
@@ -2236,7 +2237,14 @@ def circuit(x):
         qnode = qml.QNode(circuit, dev, diff_method=None)
         res = qnode(0.5)
 
-        all_shots = sum([shot_tuple.copies for shot_tuple in dev.shot_vector])
+        all_shots = sum(
+            [
+                shot_tuple.copies
+                for shot_tuple in (
+                    dev.shot_vector if isinstance(dev, qml.Device) else dev.shots.shot_vector
+                )
+            ]
+        )
 
         assert isinstance(res, tuple)
         assert len(res) == all_shots
@@ -2257,7 +2265,14 @@ def circuit(x):
         qnode = qml.QNode(circuit, dev, diff_method=None)
         res = qnode(0.5)
 
-        all_shots = sum([shot_tuple.copies for shot_tuple in dev.shot_vector])
+        all_shots = sum(
+            [
+                shot_tuple.copies
+                for shot_tuple in (
+                    dev.shot_vector if isinstance(dev, qml.Device) else dev.shots.shot_vector
+                )
+            ]
+        )
 
         assert isinstance(res, tuple)
         assert len(res) == all_shots
@@ -2281,7 +2296,14 @@ def circuit(x):
         qnode = qml.QNode(circuit, dev, diff_method=None)
         res = qnode(0.5)
 
-        all_shots = sum([shot_tuple.copies for shot_tuple in dev.shot_vector])
+        all_shots = sum(
+            [
+                shot_tuple.copies
+                for shot_tuple in (
+                    dev.shot_vector if isinstance(dev, qml.Device) else dev.shots.shot_vector
+                )
+            ]
+        )
 
         assert isinstance(res, tuple)
         assert len(res) == all_shots
@@ -2303,7 +2325,11 @@ def circuit(x):
         res = qnode(0.5)
 
         all_shot_copies = [
-            shot_tuple.shots for shot_tuple in dev.shot_vector for _ in range(shot_tuple.copies)
+            shot_tuple.shots
+            for shot_tuple in (
+                dev.shot_vector if isinstance(dev, qml.Device) else dev.shots.shot_vector
+            )
+            for _ in range(shot_tuple.copies)
         ]
 
         assert len(res) == len(all_shot_copies)
@@ -2328,7 +2354,14 @@ def circuit(x):
         qnode = qml.QNode(circuit, dev, diff_method=None)
         res = qnode(0.5)
 
-        all_shots = sum([shot_tuple.copies for shot_tuple in dev.shot_vector])
+        all_shots = sum(
+            [
+                shot_tuple.copies
+                for shot_tuple in (
+                    dev.shot_vector if isinstance(dev, qml.Device) else dev.shots.shot_vector
+                )
+            ]
+        )
 
         assert isinstance(res, tuple)
         assert len(res) == all_shots
@@ -2353,7 +2386,14 @@ def circuit(x):
         qnode = qml.QNode(circuit, dev, diff_method=None)
         res = qnode(0.5)
 
-        all_shots = sum([shot_tuple.copies for shot_tuple in dev.shot_vector])
+        all_shots = sum(
+            [
+                shot_tuple.copies
+                for shot_tuple in (
+                    dev.shot_vector if isinstance(dev, qml.Device) else dev.shots.shot_vector
+                )
+            ]
+        )
 
         assert isinstance(res, tuple)
         assert len(res) == all_shots
@@ -2383,7 +2423,14 @@ def circuit(x):
         qnode = qml.QNode(circuit, dev, diff_method=None)
         res = qnode(0.5)
 
-        all_shots = sum([shot_tuple.copies for shot_tuple in dev.shot_vector])
+        all_shots = sum(
+            [
+                shot_tuple.copies
+                for shot_tuple in (
+                    dev.shot_vector if isinstance(dev, qml.Device) else dev.shots.shot_vector
+                )
+            ]
+        )
 
         assert isinstance(res, tuple)
         assert len(res) == all_shots
@@ -2410,7 +2457,11 @@ def circuit(x):
         res = qnode(0.5)
 
         all_shot_copies = [
-            shot_tuple.shots for shot_tuple in dev.shot_vector for _ in range(shot_tuple.copies)
+            shot_tuple.shots
+            for shot_tuple in (
+                dev.shot_vector if isinstance(dev, qml.Device) else dev.shots.shot_vector
+            )
+            for _ in range(shot_tuple.copies)
         ]
 
         assert len(res) == len(all_shot_copies)
@@ -2432,7 +2483,14 @@ def circuit(x):
         qnode = qml.QNode(circuit, dev, diff_method=None)
         res = qnode(0.5)
 
-        all_shots = sum([shot_tuple.copies for shot_tuple in dev.shot_vector])
+        all_shots = sum(
+            [
+                shot_tuple.copies
+                for shot_tuple in (
+                    dev.shot_vector if isinstance(dev, qml.Device) else dev.shots.shot_vector
+                )
+            ]
+        )
 
         assert isinstance(res, tuple)
         assert len(res) == all_shots
@@ -2522,12 +2580,23 @@ def circuit(x):
         qnode = qml.QNode(circuit, dev, diff_method=None)
         res = qnode(0.5)
 
-        all_shots = sum([shot_tuple.copies for shot_tuple in dev.shot_vector])
+        all_shots = sum(
+            [
+                shot_tuple.copies
+                for shot_tuple in (
+                    dev.shot_vector if isinstance(dev, qml.Device) else dev.shots.shot_vector
+                )
+            ]
+        )
 
         assert isinstance(res, tuple)
         assert len(res) == all_shots
         assert all(isinstance(r, tuple) for r in res)
-        assert all(isinstance(m, np.ndarray) for measurement_res in res for m in measurement_res)
+        assert all(
+            isinstance(m, (np.ndarray, np.float64))
+            for measurement_res in res
+            for m in measurement_res
+        )
         for meas_res in res:
             for i, r in enumerate(meas_res):
                 if i % 2 == 0:
@@ -2545,7 +2614,11 @@ def test_scalar_sample_with_obs(self, shot_vector, meas1, meas2, device):
         observable."""
         dev = qml.device(device, wires=3, shots=shot_vector)
         raw_shot_vector = [
-            shot_tuple.shots for shot_tuple in dev.shot_vector for _ in range(shot_tuple.copies)
+            shot_tuple.shots
+            for shot_tuple in (
+                dev.shot_vector if isinstance(dev, qml.Device) else dev.shots.shot_vector
+            )
+            for _ in range(shot_tuple.copies)
         ]
 
         def circuit(x):
@@ -2556,12 +2629,23 @@ def circuit(x):
         qnode = qml.QNode(circuit, dev, diff_method=None)
         res = qnode(0.5)
 
-        all_shots = sum([shot_tuple.copies for shot_tuple in dev.shot_vector])
+        all_shots = sum(
+            [
+                shot_tuple.copies
+                for shot_tuple in (
+                    dev.shot_vector if isinstance(dev, qml.Device) else dev.shots.shot_vector
+                )
+            ]
+        )
 
         assert isinstance(res, tuple)
         assert len(res) == all_shots
         assert all(isinstance(r, tuple) for r in res)
-        assert all(isinstance(m, np.ndarray) for measurement_res in res for m in measurement_res)
+        assert all(
+            isinstance(m, (np.ndarray, np.float64))
+            for measurement_res in res
+            for m in measurement_res
+        )
 
         for idx, shots in enumerate(raw_shot_vector):
             for i, r in enumerate(res[idx]):
@@ -2586,14 +2670,25 @@ def circuit(x):
         qnode = qml.QNode(circuit, dev, diff_method=None)
         res = qnode(0.5)
 
-        all_shots = sum([shot_tuple.copies for shot_tuple in dev.shot_vector])
+        all_shots = sum(
+            [
+                shot_tuple.copies
+                for shot_tuple in (
+                    dev.shot_vector if isinstance(dev, qml.Device) else dev.shots.shot_vector
+                )
+            ]
+        )
 
         assert isinstance(res, tuple)
         assert len(res) == all_shots
         assert all(isinstance(r, tuple) for r in res)
-        assert all(isinstance(m, np.ndarray) for measurement_res in res for m in measurement_res)
+        assert all(
+            isinstance(m, (np.ndarray, np.float64))
+            for measurement_res in res
+            for m in measurement_res
+        )
 
-        for shot_tuple in dev.shot_vector:
+        for shot_tuple in dev.shot_vector if isinstance(dev, qml.Device) else dev.shots.shot_vector:
             for idx in range(shot_tuple.copies):
                 for i, r in enumerate(res[idx]):
                     if i % 2 == 0 or shot_tuple.shots == 1:
@@ -2609,7 +2704,11 @@ def test_scalar_counts_with_obs(self, shot_vector, meas1, meas2, device):
         observable."""
         dev = qml.device(device, wires=3, shots=shot_vector)
         raw_shot_vector = [
-            shot_tuple.shots for shot_tuple in dev.shot_vector for _ in range(shot_tuple.copies)
+            shot_tuple.shots
+            for shot_tuple in (
+                dev.shot_vector if isinstance(dev, qml.Device) else dev.shots.shot_vector
+            )
+            for _ in range(shot_tuple.copies)
         ]
 
         def circuit(x):
@@ -2620,14 +2719,21 @@ def circuit(x):
         qnode = qml.QNode(circuit, dev, diff_method=None)
         res = qnode(0.5)
 
-        all_shots = sum([shot_tuple.copies for shot_tuple in dev.shot_vector])
+        all_shots = sum(
+            [
+                shot_tuple.copies
+                for shot_tuple in (
+                    dev.shot_vector if isinstance(dev, qml.Device) else dev.shots.shot_vector
+                )
+            ]
+        )
 
         assert isinstance(res, tuple)
         assert len(res) == all_shots
         assert all(isinstance(r, tuple) for r in res)
 
         for r in res:
-            assert isinstance(r[0], np.ndarray)
+            assert isinstance(r[0], (np.ndarray, np.float64))
             assert isinstance(r[1], dict)
 
         expected_outcomes = {-1, 1}
@@ -2650,7 +2756,11 @@ def test_scalar_counts_no_obs(self, shot_vector, meas1, meas2, device):
         dev = qml.device(device, wires=3, shots=shot_vector)
 
         raw_shot_vector = [
-            shot_tuple.shots for shot_tuple in dev.shot_vector for _ in range(shot_tuple.copies)
+            shot_tuple.shots
+            for shot_tuple in (
+                dev.shot_vector if isinstance(dev, qml.Device) else dev.shots.shot_vector
+            )
+            for _ in range(shot_tuple.copies)
         ]
 
         def circuit(x):
@@ -2661,12 +2771,23 @@ def circuit(x):
         qnode = qml.QNode(circuit, dev, diff_method=None)
         res = qnode(0.5)
 
-        all_shots = sum([shot_tuple.copies for shot_tuple in dev.shot_vector])
+        all_shots = sum(
+            [
+                shot_tuple.copies
+                for shot_tuple in (
+                    dev.shot_vector if isinstance(dev, qml.Device) else dev.shots.shot_vector
+                )
+            ]
+        )
 
         assert isinstance(res, tuple)
         assert len(res) == all_shots
         assert all(isinstance(r, tuple) for r in res)
-        assert all(isinstance(m, np.ndarray) for measurement_res in res for m in measurement_res)
+        assert all(
+            isinstance(m, (np.ndarray, np.float64))
+            for measurement_res in res
+            for m in measurement_res
+        )
 
         for idx, shots in enumerate(raw_shot_vector):
             for i, r in enumerate(res[idx]):
@@ -2683,7 +2804,11 @@ def test_probs_sample(self, shot_vector, sample_obs, device):
         dev = qml.device(device, wires=3, shots=shot_vector)
 
         raw_shot_vector = [
-            shot_tuple.shots for shot_tuple in dev.shot_vector for _ in range(shot_tuple.copies)
+            shot_tuple.shots
+            for shot_tuple in (
+                dev.shot_vector if isinstance(dev, qml.Device) else dev.shots.shot_vector
+            )
+            for _ in range(shot_tuple.copies)
         ]
 
         meas1_wires = [0, 1]
@@ -2702,12 +2827,23 @@ def circuit(x):
         qnode = qml.QNode(circuit, dev, diff_method=None)
         res = qnode(0.5)
 
-        all_shots = sum([shot_tuple.copies for shot_tuple in dev.shot_vector])
+        all_shots = sum(
+            [
+                shot_tuple.copies
+                for shot_tuple in (
+                    dev.shot_vector if isinstance(dev, qml.Device) else dev.shots.shot_vector
+                )
+            ]
+        )
 
         assert isinstance(res, tuple)
         assert len(res) == all_shots
         assert all(isinstance(r, tuple) for r in res)
-        assert all(isinstance(m, np.ndarray) for measurement_res in res for m in measurement_res)
+        assert all(
+            isinstance(m, (np.ndarray, np.float64))
+            for measurement_res in res
+            for m in measurement_res
+        )
 
         for idx, shots in enumerate(raw_shot_vector):
             for i, r in enumerate(res[idx]):
@@ -2729,7 +2865,11 @@ def test_probs_counts(self, shot_vector, sample_obs, device):
         """Test probs and counts measurements."""
         dev = qml.device(device, wires=3, shots=shot_vector)
         raw_shot_vector = [
-            shot_tuple.shots for shot_tuple in dev.shot_vector for _ in range(shot_tuple.copies)
+            shot_tuple.shots
+            for shot_tuple in (
+                dev.shot_vector if isinstance(dev, qml.Device) else dev.shots.shot_vector
+            )
+            for _ in range(shot_tuple.copies)
         ]
 
         meas1_wires = [0, 1]
@@ -2748,12 +2888,21 @@ def circuit(x):
         qnode = qml.QNode(circuit, dev, diff_method=None)
         res = qnode(0.5)
 
-        all_shots = sum([shot_tuple.copies for shot_tuple in dev.shot_vector])
+        all_shots = sum(
+            [
+                shot_tuple.copies
+                for shot_tuple in (
+                    dev.shot_vector if isinstance(dev, qml.Device) else dev.shots.shot_vector
+                )
+            ]
+        )
 
         assert isinstance(res, tuple)
         assert len(res) == all_shots
         assert all(isinstance(r, tuple) for r in res)
-        assert all(isinstance(measurement_res[0], np.ndarray) for measurement_res in res)
+        assert all(
+            isinstance(measurement_res[0], (np.ndarray, np.float64)) for measurement_res in res
+        )
         assert all(isinstance(measurement_res[1], dict) for measurement_res in res)
 
         expected_outcomes = {-1, 1} if sample_obs is not None else {"0", "1"}
@@ -2777,7 +2926,11 @@ def test_sample_counts(self, shot_vector, sample_wires, counts_wires, device):
         samples or computational basis state samples."""
         dev = qml.device(device, wires=6, shots=shot_vector)
         raw_shot_vector = [
-            shot_tuple.shots for shot_tuple in dev.shot_vector for _ in range(shot_tuple.copies)
+            shot_tuple.shots
+            for shot_tuple in (
+                dev.shot_vector if isinstance(dev, qml.Device) else dev.shots.shot_vector
+            )
+            for _ in range(shot_tuple.copies)
         ]
 
         def circuit(x):
@@ -2802,12 +2955,21 @@ def circuit(x):
         qnode = qml.QNode(circuit, dev, diff_method=None)
         res = qnode(0.5)
 
-        all_shots = sum([shot_tuple.copies for shot_tuple in dev.shot_vector])
+        all_shots = sum(
+            [
+                shot_tuple.copies
+                for shot_tuple in (
+                    dev.shot_vector if isinstance(dev, qml.Device) else dev.shots.shot_vector
+                )
+            ]
+        )
 
         assert isinstance(res, tuple)
         assert len(res) == all_shots
         assert all(isinstance(r, tuple) for r in res)
-        assert all(isinstance(measurement_res[0], np.ndarray) for measurement_res in res)
+        assert all(
+            isinstance(measurement_res[0], (np.ndarray, np.float64)) for measurement_res in res
+        )
         assert all(isinstance(measurement_res[1], dict) for measurement_res in res)
 
         for idx, shots in enumerate(raw_shot_vector):
@@ -2828,7 +2990,11 @@ def test_scalar_probs_sample_counts(self, shot_vector, meas1, meas2, device):
         in a single qfunc."""
         dev = qml.device(device, wires=5, shots=shot_vector)
         raw_shot_vector = [
-            shot_tuple.shots for shot_tuple in dev.shot_vector for _ in range(shot_tuple.copies)
+            shot_tuple.shots
+            for shot_tuple in (
+                dev.shot_vector if isinstance(dev, qml.Device) else dev.shots.shot_vector
+            )
+            for _ in range(shot_tuple.copies)
         ]
 
         def circuit(x):
@@ -2845,7 +3011,14 @@ def circuit(x):
 
         res = qnode(0.5)
 
-        all_shots = sum([shot_tuple.copies for shot_tuple in dev.shot_vector])
+        all_shots = sum(
+            [
+                shot_tuple.copies
+                for shot_tuple in (
+                    dev.shot_vector if isinstance(dev, qml.Device) else dev.shots.shot_vector
+                )
+            ]
+        )
 
         assert isinstance(res, tuple)
         assert len(res) == all_shots
@@ -2901,7 +3074,7 @@ def cost(a):
 
         res = qml.jacobian(cost)(x)
 
-        assert isinstance(res, np.ndarray)
+        assert isinstance(res, (np.ndarray, np.float64))
         assert res.shape == (2, 3)
 
     @pytest.mark.torch
@@ -3068,7 +3241,7 @@ def cost(a):
 
         res = qml.jacobian(cost)(x)
 
-        assert isinstance(res, np.ndarray)
+        assert isinstance(res, (np.ndarray, np.float64))
         assert res.shape == (2, 2, 3)
 
     @pytest.mark.torch
@@ -3207,7 +3380,7 @@ def cost(a):
 
         res = qml.jacobian(cost)(x)
 
-        assert isinstance(res, np.ndarray)
+        assert isinstance(res, (np.ndarray, np.float64))
         assert res.shape == (6, 3)
 
     @pytest.mark.torch
diff --git a/tests/test_tensor_measurements.py b/tests/test_tensor_measurements.py
index 4f64a68b52d..bca43eba57b 100644
--- a/tests/test_tensor_measurements.py
+++ b/tests/test_tensor_measurements.py
@@ -18,7 +18,6 @@
 import pytest
 
 import numpy as np
-from gate_data import I, Z, S, Rotx, Roty, H, CNOT
 import pennylane as qml
 from pennylane import expval, var, sample
 
@@ -389,25 +388,6 @@ def circuit(a, b, c):
         # s1 should only contain 1 and -1
         assert np.allclose(s1**2, 1, atol=tol_stochastic, rtol=0)
 
-        zero_state = np.zeros(2**3)
-        zero_state[0] = 1
-        psi = zero_state
-        psi = tensor_product([Rotx(theta), I, I]) @ zero_state
-        psi = tensor_product([I, Rotx(phi), I]) @ psi
-        psi = tensor_product([I, I, Rotx(varphi)]) @ psi
-        psi = tensor_product([CNOT, I]) @ psi
-        psi = tensor_product([I, CNOT]) @ psi
-
-        # Diagonalize according to the observable
-        psi = tensor_product([H, I, I]) @ psi
-        psi = tensor_product([I, I, Z]) @ psi
-        psi = tensor_product([I, I, S]) @ psi
-        psi = tensor_product([I, I, H]) @ psi
-
-        expected_probabilities = np.abs(psi) ** 2
-
-        assert np.allclose(dev.probability(), expected_probabilities, atol=tol_stochastic, rtol=0)
-
     def test_pauliz_tensor_hadamard(self, theta, phi, varphi, tol_stochastic):
         """Test that a tensor product involving PauliZ and hadamard works correctly"""
         dev = qml.device("default.qubit", wires=3, shots=int(1e6))
@@ -419,25 +399,6 @@ def circuit(a, b, c):
 
         s1 = circuit(theta, phi, varphi)
 
-        zero_state = np.zeros(2**3)
-        zero_state[0] = 1
-        psi = zero_state
-        psi = tensor_product([Rotx(theta), I, I]) @ zero_state
-        psi = tensor_product([I, Rotx(phi), I]) @ psi
-        psi = tensor_product([I, I, Rotx(varphi)]) @ psi
-        psi = tensor_product([CNOT, I]) @ psi
-        psi = tensor_product([I, CNOT]) @ psi
-
-        # Diagonalize according to the observable
-        psi = tensor_product([I, Roty(-np.pi / 4), I]) @ psi
-        psi = tensor_product([I, I, Z]) @ psi
-        psi = tensor_product([I, I, S]) @ psi
-        psi = tensor_product([I, I, H]) @ psi
-
-        expected_probabilities = np.abs(psi) ** 2
-
-        assert np.allclose(dev.probability(), expected_probabilities, atol=tol_stochastic, rtol=0)
-
         # s1 should only contain 1 and -1
         assert np.allclose(s1**2, 1, atol=tol_stochastic, rtol=0)
 
@@ -465,19 +426,3 @@ def circuit(a, b, c):
         # the hermitian matrix tensor product Z
         eigvals = np.linalg.eigvalsh(np.kron(Z, A))
         assert set(np.round(s1, 8)).issubset(set(np.round(eigvals, 8)))
-
-        zero_state = np.zeros(2**3)
-        zero_state[0] = 1
-        psi = tensor_product([Rotx(theta), I, I]) @ zero_state
-        psi = tensor_product([I, Rotx(phi), I]) @ psi
-        psi = tensor_product([I, I, Rotx(varphi)]) @ psi
-        psi = tensor_product([CNOT, I]) @ psi
-        psi = tensor_product([I, CNOT]) @ psi
-
-        # Diagonalize according to the observable
-        eigvals, eigvecs = np.linalg.eigh(A)
-        psi = tensor_product([I, eigvecs.conj().T]) @ psi
-
-        expected_probabilities = np.abs(psi) ** 2
-
-        assert np.allclose(dev.probability(), expected_probabilities, atol=tol_stochastic, rtol=0)
diff --git a/tests/test_vqe.py b/tests/test_vqe.py
index 33c0910b447..fd881ab0df5 100644
--- a/tests/test_vqe.py
+++ b/tests/test_vqe.py
@@ -345,15 +345,19 @@ def test_optimize_torch(self, shots):
         shape = qml.templates.StronglyEntanglingLayers.shape(n_layers=2, n_wires=4)
         w = np.random.random(shape)
 
-        c1 = cost(w)
-        exec_opt = dev.num_executions
-        dev._num_executions = 0
+        with qml.Tracker(dev) as tracker:
+            c1 = cost(w)
 
-        c2 = cost2(w)
-        exec_no_opt = dev.num_executions
+        exec_opt = tracker.latest["executions"]
 
-        assert exec_opt == 5  # Number of groups in the Hamiltonian
-        assert exec_no_opt == 15
+        with tracker:
+            c2 = cost2(w)
+
+        exec_no_opt = tracker.latest["executions"]
+
+        # was 5, 15 on old device
+        assert exec_opt == 1  # Number of groups in the Hamiltonian
+        assert exec_no_opt == 1
 
         assert np.allclose(c1, c2, atol=1e-1)
 
@@ -391,15 +395,17 @@ def test_optimize_tf(self, shots):
         shape = qml.templates.StronglyEntanglingLayers.shape(n_layers=2, n_wires=4)
         w = np.random.random(shape)
 
-        c1 = cost(w)
-        exec_opt = dev.num_executions
-        dev._num_executions = 0
+        with qml.Tracker(dev) as tracker:
+            c1 = cost(w)
+        exec_opt = tracker.latest["executions"]
 
-        c2 = cost2(w)
-        exec_no_opt = dev.num_executions
+        with tracker:
+            c2 = cost2(w)
+        exec_no_opt = tracker.latest["executions"]
 
-        assert exec_opt == 5  # Number of groups in the Hamiltonian
-        assert exec_no_opt == 15
+        # was 5, 15 on old device
+        assert exec_opt == 1  # Number of groups in the Hamiltonian
+        assert exec_no_opt == 1
 
         assert np.allclose(c1, c2, atol=1e-1)
 
@@ -437,15 +443,17 @@ def test_optimize_autograd(self, shots):
         shape = qml.templates.StronglyEntanglingLayers.shape(n_layers=2, n_wires=4)
         w = np.random.random(shape)
 
-        c1 = cost(w)
-        exec_opt = dev.num_executions
-        dev._num_executions = 0
+        with qml.Tracker(dev) as tracker:
+            c1 = cost(w)
+        exec_opt = tracker.latest["executions"]
 
-        c2 = cost2(w)
-        exec_no_opt = dev.num_executions
+        with tracker:
+            c2 = cost2(w)
+        exec_no_opt = tracker.latest["executions"]
 
-        assert exec_opt == 5  # Number of groups in the Hamiltonian
-        assert exec_no_opt == 15
+        # was 5, 15 on old device
+        assert exec_opt == 1  # Number of groups in the Hamiltonian
+        assert exec_no_opt == 1
 
         assert np.allclose(c1, c2, atol=1e-1)
 
@@ -493,15 +501,17 @@ def test_optimize_multiple_terms_autograd(self):
         shape = qml.templates.StronglyEntanglingLayers.shape(n_layers=2, n_wires=5)
         w = np.random.random(shape)
 
-        c1 = cost(w)
-        exec_opt = dev.num_executions
-        dev._num_executions = 0
+        with qml.Tracker(dev) as tracker:
+            c1 = cost(w)
+        exec_opt = tracker.latest["executions"]
 
-        c2 = cost2(w)
-        exec_no_opt = dev.num_executions
+        with tracker:
+            c2 = cost2(w)
+        exec_no_opt = tracker.latest["executions"]
 
+        # was 1, 8 on old device
         assert exec_opt == 1  # Number of groups in the Hamiltonian
-        assert exec_no_opt == 8
+        assert exec_no_opt == 1
 
         assert np.allclose(c1, c2)
 
@@ -549,15 +559,17 @@ def test_optimize_multiple_terms_torch(self):
         shape = qml.templates.StronglyEntanglingLayers.shape(n_layers=2, n_wires=5)
         w = np.random.random(shape)
 
-        c1 = cost(w)
-        exec_opt = dev.num_executions
-        dev._num_executions = 0
+        with qml.Tracker(dev) as tracker:
+            c1 = cost(w)
+        exec_opt = tracker.latest["executions"]
 
-        c2 = cost2(w)
-        exec_no_opt = dev.num_executions
+        with tracker:
+            c2 = cost2(w)
+        exec_no_opt = tracker.latest["executions"]
 
+        # was 1, 8 on old device
         assert exec_opt == 1  # Number of groups in the Hamiltonian
-        assert exec_no_opt == 8
+        assert exec_no_opt == 1
 
         assert np.allclose(c1, c2)
 
@@ -605,19 +617,22 @@ def test_optimize_multiple_terms_tf(self):
         shape = qml.templates.StronglyEntanglingLayers.shape(n_layers=2, n_wires=5)
         w = np.random.random(shape)
 
-        c1 = cost(w)
-        exec_opt = dev.num_executions
-        dev._num_executions = 0
+        with qml.Tracker(dev) as tracker:
+            c1 = cost(w)
+        exec_opt = tracker.latest["executions"]
 
-        c2 = cost2(w)
-        exec_no_opt = dev.num_executions
+        with tracker:
+            c2 = cost2(w)
+        exec_no_opt = tracker.latest["executions"]
 
+        # was 1, 8 on old device
         assert exec_opt == 1  # Number of groups in the Hamiltonian
-        assert exec_no_opt == 8
+        assert exec_no_opt == 1
 
         assert np.allclose(c1, c2)
 
     # pylint: disable=protected-access
+    @pytest.mark.skip("non-opt is less executions than opt with new device")
     @pytest.mark.autograd
     def test_optimize_grad(self):
         """Test that the gradient of ExpvalCost is accessible and correct when using observable
@@ -646,12 +661,13 @@ def test_optimize_grad(self):
         shape = qml.templates.StronglyEntanglingLayers.shape(n_layers=2, n_wires=4)
         w = pnp.random.uniform(low=0, high=2 * np.pi, size=shape, requires_grad=True)
 
-        dc = qml.grad(cost)(w)
-        exec_opt = dev.num_executions
-        dev._num_executions = 0
+        with qml.Tracker(dev) as tracker:
+            dc = qml.grad(cost)(w)
+        exec_opt = tracker.latest["executions"]
 
-        dc2 = qml.grad(cost2)(w)
-        exec_no_opt = dev.num_executions
+        with tracker:
+            dc2 = qml.grad(cost2)(w)
+        exec_no_opt = tracker.latest["executions"]
 
         assert exec_no_opt > exec_opt
         assert np.allclose(dc, big_hamiltonian_grad)
@@ -750,7 +766,7 @@ def ansatz(params, **kwargs):
 
         h = qml.Hamiltonian([1, 1], [qml.PauliZ(0), qml.PauliZ(1)])
         qnodes = catch_warn_ExpvalCost(ansatz, h, dev)
-        mt = qml.metric_tensor(qnodes, approx=approx)(p)
+        mt = qml.metric_tensor(qnodes, approx=approx)(p)  # pylint:disable=not-callable
         assert mt.shape == (3, 3)
         assert isinstance(mt, pnp.ndarray)
 
@@ -843,10 +859,14 @@ def circuit2():
 
         assert np.allclose(res1, res2, atol=tol)
 
-    @pytest.mark.autograd
+    @pytest.mark.jax
     @pytest.mark.parametrize("shots, dim", [([(1000, 2)], 2), ([30, 30], 2), ([2, 3, 4], 3)])
     def test_shot_distribution(self, shots, dim):
         """Tests that distributed shots work with the new VQE design."""
+        import jax
+
+        jax.config.update("jax_enable_x64", True)
+
         dev = qml.device("default.qubit", wires=2, shots=shots)
 
         @qml.qnode(dev)
@@ -857,12 +877,13 @@ def circuit(weights, coeffs):
 
         obs = [qml.PauliZ(0), qml.PauliX(0) @ qml.PauliZ(1)]
         coeffs = np.array([0.1, 0.2])
-        weights = pnp.random.random([2, 2, 3], requires_grad=True)
+        key = jax.random.PRNGKey(42)
+        weights = jax.random.uniform(key, [2, 2, 3])
 
         res = circuit(weights, coeffs)
-        grad = qml.jacobian(circuit, argnum=1)(weights, coeffs)
+        grad = jax.jacobian(circuit, argnums=[1])(weights, coeffs)
         assert len(res) == dim
-        assert grad.shape == (dim, 2)
+        assert qml.math.shape(grad) == (dim, 1, 2)
 
     def test_circuit_drawer(self):
         """Test that the circuit drawer displays Hamiltonians well."""
@@ -943,9 +964,7 @@ def test_error_var_measurement(self):
         def circuit():
             return qml.var(H)
 
-        with pytest.raises(
-            qml.operation.EigvalsUndefinedError, match="Cannot compute analytic variance"
-        ):
+        with pytest.raises(NotImplementedError):
             circuit()
 
     def test_error_sample_measurement(self):
@@ -959,7 +978,7 @@ def test_error_sample_measurement(self):
         def circuit():
             return qml.sample(H)
 
-        with pytest.raises(ValueError, match="Can only return the expectation of a single"):
+        with pytest.raises(qml.operation.DiagGatesUndefinedError):
             circuit()
 
     @pytest.mark.autograd
@@ -1010,7 +1029,7 @@ def circuit(w):
         w = torch.tensor(PARAMS, requires_grad=True)
 
         res = circuit(w)
-        res.backward()
+        res.backward()  # pylint:disable=no-member
         dc = w.grad.detach().numpy()
 
         assert np.allclose(dc, big_hamiltonian_grad, atol=tol)
diff --git a/tests/transforms/test_adjoint_metric_tensor.py b/tests/transforms/test_adjoint_metric_tensor.py
index f79572cd074..04aa867ddec 100644
--- a/tests/transforms/test_adjoint_metric_tensor.py
+++ b/tests/transforms/test_adjoint_metric_tensor.py
@@ -261,7 +261,7 @@ def test_correct_output_tape_autograd(self, ansatz, params, interface):
         """Test that the output is correct when using Autograd and
         calling the adjoint metric tensor directly on a tape."""
         expected = autodiff_metric_tensor(ansatz, 3)(*params)
-        dev = qml.devices.DefaultQubit()
+        dev = qml.device("default.qubit")
 
         wires = ("a", "b", "c")
 
@@ -281,7 +281,8 @@ def circuit(*params):
 
     @pytest.mark.jax
     @pytest.mark.skip("JAX does not support forward pass executiong of the metric tensor.")
-    def test_correct_output_tape_jax(self, ansatz, params):
+    @pytest.mark.parametrize("dev_name", ["default.qubit", "default.qubit.jax"])
+    def test_correct_output_tape_jax(self, dev_name, ansatz, params):
         """Test that the output is correct when using JAX and
         calling the adjoint metric tensor directly on a tape."""
 
@@ -292,7 +293,7 @@ def test_correct_output_tape_jax(self, ansatz, params):
 
         expected = autodiff_metric_tensor(ansatz, self.num_wires)(*params)
         j_params = tuple(jax.numpy.array(p) for p in params)
-        dev = qml.device("default.qubit.jax", wires=self.num_wires)
+        dev = qml.device(dev_name, wires=self.num_wires)
 
         @qml.qnode(dev, interface="jax")
         def circuit(*params):
@@ -312,7 +313,8 @@ def circuit(*params):
 
     @pytest.mark.torch
     @pytest.mark.parametrize("interface", interfaces)
-    def test_correct_output_tape_torch(self, ansatz, params, interface):
+    @pytest.mark.parametrize("dev_name", ["default.qubit", "default.qubit.torch"])
+    def test_correct_output_tape_torch(self, ansatz, params, interface, dev_name):
         """Test that the output is correct when using Torch and
         calling the adjoint metric tensor directly on a tape."""
 
@@ -320,7 +322,7 @@ def test_correct_output_tape_torch(self, ansatz, params, interface):
 
         expected = autodiff_metric_tensor(ansatz, self.num_wires)(*params)
         t_params = tuple(torch.tensor(p, requires_grad=True) for p in params)
-        dev = qml.device("default.qubit.torch", wires=self.num_wires)
+        dev = qml.device(dev_name, wires=self.num_wires)
 
         @qml.qnode(dev, interface=interface)
         def circuit(*params):
@@ -340,7 +342,8 @@ def circuit(*params):
 
     @pytest.mark.tf
     @pytest.mark.parametrize("interface", interfaces)
-    def test_correct_output_tape_tf(self, ansatz, params, interface):
+    @pytest.mark.parametrize("dev_name", ["default.qubit", "default.qubit.tf"])
+    def test_correct_output_tape_tf(self, ansatz, params, interface, dev_name):
         """Test that the output is correct when using TensorFlow and
         calling the adjoint metric tensor directly on a tape."""
 
@@ -348,7 +351,7 @@ def test_correct_output_tape_tf(self, ansatz, params, interface):
 
         expected = autodiff_metric_tensor(ansatz, self.num_wires)(*params)
         t_params = tuple(tf.Variable(p) for p in params)
-        dev = qml.device("default.qubit.tf", wires=self.num_wires)
+        dev = qml.device(dev_name, wires=self.num_wires)
 
         @qml.qnode(dev, interface=interface)
         def circuit(*params):
@@ -432,7 +435,8 @@ def circuit(*params):
     @pytest.mark.torch
     @pytest.mark.parametrize("ansatz, params", list(zip(fubini_ansatze, fubini_params)))
     @pytest.mark.parametrize("interface", interfaces)
-    def test_correct_output_qnode_torch(self, ansatz, params, interface):
+    @pytest.mark.parametrize("dev_name", ["default.qubit", "default.qubit.torch"])
+    def test_correct_output_qnode_torch(self, ansatz, params, interface, dev_name):
         """Test that the output is correct when using Torch and
         calling the adjoint metric tensor on a QNode."""
 
@@ -440,7 +444,7 @@ def test_correct_output_qnode_torch(self, ansatz, params, interface):
 
         expected = autodiff_metric_tensor(ansatz, self.num_wires)(*params)
         t_params = tuple(torch.tensor(p, requires_grad=True, dtype=torch.float64) for p in params)
-        dev = qml.device("default.qubit.torch", wires=self.num_wires)
+        dev = qml.device(dev_name, wires=self.num_wires)
 
         @qml.qnode(dev, interface=interface)
         def circuit(*params):
@@ -485,7 +489,8 @@ def circuit(*params):
             assert qml.math.allclose(mt, expected)
 
     @pytest.mark.autograd
-    def test_autograd_with_other_device(self):
+    @pytest.mark.parametrize("dev_name", ["default.qubit", "default.qubit.autograd"])
+    def test_autograd_with_other_device(self, dev_name):
         """Test passing an extra device to the QNode wrapper."""
         ansatz = fubini_ansatz2
         params = fubini_params[2]
@@ -493,7 +498,7 @@ def test_autograd_with_other_device(self):
         exp_fn = autodiff_metric_tensor(ansatz, self.num_wires)
         expected = qml.jacobian(exp_fn)(*params)
         dev = qml.device("default.qubit", wires=self.num_wires)
-        dev2 = qml.device("default.qubit.autograd", wires=self.num_wires)
+        dev2 = qml.device(dev_name, wires=self.num_wires)
 
         @qml.qnode(dev)
         def circuit(*params):
diff --git a/tests/transforms/test_batch_input.py b/tests/transforms/test_batch_input.py
index 9f02442d2ee..d30a12cc06c 100644
--- a/tests/transforms/test_batch_input.py
+++ b/tests/transforms/test_batch_input.py
@@ -14,7 +14,7 @@
 """
 Unit tests for the ``batch_inputs`` transform.
 """
-# pylint: disable=too-few-public-methods
+# pylint: disable=too-few-public-methods,no-value-for-parameter,comparison-with-callable
 from functools import partial
 
 import pytest
@@ -183,7 +183,7 @@ def circuit(data, weights):
     # weights is not batched
     weights = np.random.random((10, 3, 3), requires_grad=True)
 
-    spy = mocker.spy(circuit.device, "batch_execute")
+    spy = mocker.spy(circuit.device, "execute")
     res = circuit(data, weights)
     assert res.shape == (batch_size, 2**3)
     assert len(spy.call_args[0][0]) == batch_size
@@ -222,7 +222,7 @@ def circuit(data, weights):
     # weights is not batched
     weights = np.random.random((10, 3, 3), requires_grad=True)
 
-    spy = mocker.spy(circuit.device, "batch_execute")
+    spy = mocker.spy(circuit.device, "execute")
     res = circuit(data, weights)
     assert res.shape == (batch_size, 2**3)
     assert len(spy.call_args[0][0]) == batch_size
diff --git a/tests/transforms/test_batch_params.py b/tests/transforms/test_batch_params.py
index f2333c59eb2..36accf7bb25 100644
--- a/tests/transforms/test_batch_params.py
+++ b/tests/transforms/test_batch_params.py
@@ -14,6 +14,7 @@
 """
 Unit tests for the batch params transform.
 """
+# pylint:disable=comparison-with-callable
 import functools
 import pytest
 
@@ -39,7 +40,7 @@ def circuit(data, x, weights):
     x = np.linspace(0.1, 0.5, batch_size, requires_grad=True)
     weights = np.ones((batch_size, 10, 3, 3), requires_grad=True)
 
-    spy = mocker.spy(circuit.device, "batch_execute")
+    spy = mocker.spy(circuit.device, "execute")
     res = circuit(data, x, weights)
     assert res.shape == (batch_size, 4)
     assert len(spy.call_args[0][0]) == batch_size
@@ -63,7 +64,7 @@ def circuit(data, x, weights):
     x = np.linspace(0.1, 0.5, batch_size, requires_grad=True)
     weights = np.ones((batch_size, 10, 3, 3), requires_grad=True)
 
-    spy = mocker.spy(circuit.device, "batch_execute")
+    spy = mocker.spy(circuit.device, "execute")
     res = circuit(data, x, weights)
     assert res.shape == (batch_size, 4)
     assert len(spy.call_args[0][0]) == batch_size
@@ -90,7 +91,7 @@ def circuit(data, x, weights):
     x = np.linspace(0.1, 0.5, batch_size, requires_grad=True)
     weights = np.ones((batch_size, 10, 3, 3), requires_grad=True)
 
-    spy = mocker.spy(circuit.device, "batch_execute")
+    spy = mocker.spy(circuit.device, "execute")
     res = circuit(data, x, weights)
     assert res.shape == (batch_size, 4)
     assert len(spy.call_args[0][0]) == batch_size
@@ -110,7 +111,7 @@ def circuit(weights):
     batch_size = 5
     weights = np.random.random((batch_size, 2, 2))
 
-    spy = mocker.spy(circuit.device, "batch_execute")
+    spy = mocker.spy(circuit.device, "execute")
     res = circuit(weights)
     assert res.shape == (batch_size, 4)
     assert len(spy.call_args[0][0]) == batch_size
@@ -130,7 +131,7 @@ def circuit(data):
     batch_size = 5
     data = np.random.random((batch_size, 3))
 
-    spy = mocker.spy(circuit.device, "batch_execute")
+    spy = mocker.spy(circuit.device, "execute")
     res = circuit(data)
     assert res.shape == (batch_size, 4)
     assert len(spy.call_args[0][0]) == batch_size
@@ -154,7 +155,7 @@ def circuit(data, weights):
     data /= np.linalg.norm(data, axis=1).reshape(-1, 1)  # normalize
     weights = np.random.random((batch_size, 10, 3, 3))
 
-    spy = mocker.spy(circuit.device, "batch_execute")
+    spy = mocker.spy(circuit.device, "execute")
     res = circuit(data, weights)
     assert res.shape == (batch_size, 2**3)
     assert len(spy.call_args[0][0]) == batch_size
@@ -189,7 +190,7 @@ def circuit(data, weights):
     data = np.random.randint(2, size=(batch_size, 4))
     weights = np.random.random((batch_size, 10, 4, 3))
 
-    spy = mocker.spy(circuit.device, "batch_execute")
+    spy = mocker.spy(circuit.device, "execute")
     res = circuit(data, weights)
     assert res.shape == (batch_size, 2**4)
     assert len(spy.call_args[0][0]) == batch_size
@@ -225,7 +226,7 @@ def circuit(data, weights):
     data /= np.linalg.norm(data, axis=1).reshape(-1, 1)  # normalize
     weights = np.random.random((batch_size, 10, 3, 3))
 
-    spy = mocker.spy(circuit.device, "batch_execute")
+    spy = mocker.spy(circuit.device, "execute")
     res = circuit(data, weights)
     assert res.shape == (batch_size, 2**3)
     assert len(spy.call_args[0][0]) == batch_size
@@ -272,6 +273,7 @@ def circuit(data, x, weights):
 
 def test_shot_vector():
     """Test that batching works for a simple circuit with a shot vector"""
+    # pylint:disable=not-an-iterable
     dev = qml.device("default.qubit", wires=3, shots=(100, (200, 3), 300))
 
     @qml.batch_params
@@ -801,7 +803,7 @@ def circuit(x, weights):
     x = np.linspace(0.1, 0.5, batch_size, requires_grad=True)
     weights = np.ones((batch_size, 10, 3, 3), requires_grad=False)
 
-    spy = mocker.spy(circuit.device, "batch_execute")
+    spy = mocker.spy(circuit.device, "execute")
     res = circuit(x, weights)
     assert res.shape == (batch_size, 4)
     assert len(spy.call_args[0][0]) == batch_size
diff --git a/tests/transforms/test_batch_transform.py b/tests/transforms/test_batch_transform.py
index 742e15bc78b..eb2d671da38 100644
--- a/tests/transforms/test_batch_transform.py
+++ b/tests/transforms/test_batch_transform.py
@@ -428,7 +428,7 @@ def circuit(x):
         assert tapes[1].operations[1].name == "RZ"
         assert tapes[1].operations[1].parameters == [b * np.sin(x)]
 
-        expected = fn(dev.batch_execute(tapes))
+        expected = fn(dev.execute(tapes))
         assert res == expected
         assert circuit.interface == "auto"
 
@@ -464,7 +464,7 @@ def circuit(x):
         assert tapes[1].operations[1].name == "RZ"
         assert tapes[1].operations[1].parameters == [b * np.sin(x)]
 
-        expected = fn(dev.batch_execute(tapes))
+        expected = fn(dev.execute(tapes))
         assert res == expected
 
     def test_custom_qnode_wrapper(self):
@@ -667,7 +667,7 @@ def circuit(weights):
             qml.CNOT(wires=[0, 1])
             return qml.expval(H)
 
-        spy = mocker.spy(dev, "batch_transform")
+        spy = mocker.spy(dev, "preprocess")
 
         res = circuit(weights)
         spy.assert_called()
diff --git a/tests/transforms/test_hamiltonian_expand.py b/tests/transforms/test_hamiltonian_expand.py
index b6d4581e157..886623740fd 100644
--- a/tests/transforms/test_hamiltonian_expand.py
+++ b/tests/transforms/test_hamiltonian_expand.py
@@ -100,14 +100,14 @@ def test_hamiltonians(self, tape, output):
         """Tests that the hamiltonian_expand transform returns the correct value"""
 
         tapes, fn = hamiltonian_expand(tape)
-        results = dev.batch_execute(tapes)
+        results = dev.execute(tapes)
         expval = fn(results)
 
         assert np.isclose(output, expval)
 
         qs = QuantumScript(tape.operations, tape.measurements)
         tapes, fn = hamiltonian_expand(qs)
-        results = dev.batch_execute(tapes)
+        results = dev.execute(tapes)
         expval = fn(results)
         assert np.isclose(output, expval)
 
@@ -117,14 +117,14 @@ def test_hamiltonians_no_grouping(self, tape, output):
         if we switch grouping off"""
 
         tapes, fn = hamiltonian_expand(tape, group=False)
-        results = dev.batch_execute(tapes)
+        results = dev.execute(tapes)
         expval = fn(results)
 
         assert np.isclose(output, expval)
 
         qs = QuantumScript(tape.operations, tape.measurements)
         tapes, fn = hamiltonian_expand(qs, group=False)
-        results = dev.batch_execute(tapes)
+        results = dev.execute(tapes)
         expval = fn(results)
 
         assert np.isclose(output, expval)
@@ -274,10 +274,10 @@ def test_hamiltonian_dif_tensorflow(self):
 
         with tf.GradientTape() as gtape:
             with AnnotatedQueue() as q:
-                for i in range(2):
-                    qml.RX(var[i, 0], wires=0)
-                    qml.RX(var[i, 1], wires=1)
-                    qml.RX(var[i, 2], wires=2)
+                for _i in range(2):
+                    qml.RX(var[_i, 0], wires=0)
+                    qml.RX(var[_i, 1], wires=1)
+                    qml.RX(var[_i, 2], wires=2)
                     qml.CNOT(wires=[0, 1])
                     qml.CNOT(wires=[1, 2])
                     qml.CNOT(wires=[2, 0])
@@ -405,8 +405,21 @@ def test_observables_on_same_wires(self):
     @pytest.mark.parametrize(("qscript", "output"), zip(SUM_QSCRIPTS, SUM_OUTPUTS))
     def test_sums(self, qscript, output):
         """Tests that the sum_expand transform returns the correct value"""
+        processed, _ = dev.preprocess()[0]([qscript])
+        assert len(processed) == 1
+        qscript = processed[0]
         tapes, fn = sum_expand(qscript)
-        results = dev.batch_execute(tapes)
+        results = dev.execute(tapes)
+        expval = fn(results)
+
+        assert all(qml.math.allclose(o, e) for o, e in zip(output, expval))
+
+    @pytest.mark.parametrize(("qscript", "output"), zip(SUM_QSCRIPTS, SUM_OUTPUTS))
+    def test_sums_legacy(self, qscript, output):
+        """Tests that the sum_expand transform returns the correct value"""
+        dev_old = qml.device("default.qubit.legacy", wires=4)
+        tapes, fn = sum_expand(qscript)
+        results = dev_old.batch_execute(tapes)
         expval = fn(results)
 
         assert all(qml.math.allclose(o, e) for o, e in zip(output, expval))
@@ -415,8 +428,11 @@ def test_sums(self, qscript, output):
     def test_sums_no_grouping(self, qscript, output):
         """Tests that the sum_expand transform returns the correct value
         if we switch grouping off"""
+        processed, _ = dev.preprocess()[0]([qscript])
+        assert len(processed) == 1
+        qscript = processed[0]
         tapes, fn = sum_expand(qscript, group=False)
-        results = dev.batch_execute(tapes)
+        results = dev.execute(tapes)
         expval = fn(results)
 
         assert all(qml.math.allclose(o, e) for o, e in zip(output, expval))
@@ -559,10 +575,10 @@ def test_sum_dif_tensorflow(self):
 
         with tf.GradientTape() as gtape:
             with AnnotatedQueue() as q:
-                for i in range(2):
-                    qml.RX(var[i, 0], wires=0)
-                    qml.RX(var[i, 1], wires=1)
-                    qml.RX(var[i, 2], wires=2)
+                for _i in range(2):
+                    qml.RX(var[_i, 0], wires=0)
+                    qml.RX(var[_i, 1], wires=1)
+                    qml.RX(var[_i, 2], wires=2)
                     qml.CNOT(wires=[0, 1])
                     qml.CNOT(wires=[1, 2])
                     qml.CNOT(wires=[2, 0])
diff --git a/tests/transforms/test_metric_tensor.py b/tests/transforms/test_metric_tensor.py
index c28fcb23638..1b4c6db5807 100644
--- a/tests/transforms/test_metric_tensor.py
+++ b/tests/transforms/test_metric_tensor.py
@@ -14,7 +14,7 @@
 """
 Unit tests for the metric tensor transform.
 """
-# pylint: disable=too-many-arguments,too-many-public-methods,too-few-public-methods
+# pylint: disable=too-many-arguments,too-many-public-methods,too-few-public-methods,not-callable
 import pytest
 from scipy.linalg import block_diag
 
@@ -93,6 +93,7 @@ def circuit(a):
 
         circuit = qml.QNode(circuit, dev, diff_method=diff_method)
         params = np.array([0.1], requires_grad=True)
+        # pylint:disable=unexpected-keyword-arg
         result = qml.metric_tensor(circuit, hybrid=False, approx="block-diag")(*params)
         assert result.shape == (2, 2)
 
@@ -1192,6 +1193,7 @@ def circuit(*params):
             return qml.expval(qml.PauliZ(0))
 
         if len(params) > 1:
+            # pylint:disable=unexpected-keyword-arg
             mt = qml.metric_tensor(circuit, argnums=range(0, len(params)), approx=None)(*params)
         else:
             mt = qml.metric_tensor(circuit, approx=None)(*params)
@@ -1760,7 +1762,7 @@ def circuit(x):
         return qml.expval(qml.PauliZ(0))
 
     x = np.array([1.3, 0.2])
-    with pytest.raises(qml.wires.WireError, match=r"Did not find some of the wires \(0, 1\)"):
+    with pytest.raises(qml.wires.WireError, match=r"contain wires not found on the device: \{1\}"):
         qml.transforms.metric_tensor(circuit)(x)
 
 
diff --git a/tests/transforms/test_qcut.py b/tests/transforms/test_qcut.py
index e7b92977037..d6c87fd4191 100644
--- a/tests/transforms/test_qcut.py
+++ b/tests/transforms/test_qcut.py
@@ -15,7 +15,8 @@
 Unit tests for the `pennylane.qcut` package.
 """
 # pylint: disable=protected-access,too-few-public-methods,too-many-arguments
-# pylint: disable=too-many-public-methods
+# pylint: disable=too-many-public-methods,comparison-with-callable
+# pylint: disable=no-value-for-parameter,no-member,not-callable
 import copy
 import itertools
 import string
@@ -3718,7 +3719,7 @@ def test_qcut_processing_fn_tf(self, use_opt_einsum):
         x = tf.Variable(0.9, dtype=tf.float64)
 
         def f(x):
-            x = tf.cast(x, dtype=tf.float64)
+            x = tf.cast(x, dtype=tf.float64)  # pylint:disable=unexpected-keyword-arg
             t1 = x * tf.range(4, dtype=tf.float64)
             t2 = x**2 * tf.range(16, dtype=tf.float64)
             t3 = tf.sin(x * np.pi / 2) * tf.range(4, dtype=tf.float64)
@@ -4020,7 +4021,9 @@ def test_simple_cut_circuit_torch_trace(self, mocker, use_opt_einsum):
 
         import torch
 
-        dev = qml.device("default.qubit", wires=2)
+        # TODO: this passes with default.qubit locally, but fails on CI
+        # possibly an architecture-specific issue
+        dev = qml.device("default.qubit.legacy", wires=2)
 
         @qml.qnode(dev, interface="torch")
         def circuit(x):
@@ -4635,7 +4638,7 @@ def test_init_raises(self, devices, imbalance_tolerance, num_fragments_probed):
         """Test if ill-initialized instances throw errors."""
 
         if (
-            isinstance(devices, qml.Device)
+            isinstance(devices, (qml.Device, qml.devices.Device))
             and imbalance_tolerance is None
             and num_fragments_probed is None
         ):
diff --git a/tests/transforms/test_qmc_transform.py b/tests/transforms/test_qmc_transform.py
index 9e2542986b5..61446d0c6aa 100644
--- a/tests/transforms/test_qmc_transform.py
+++ b/tests/transforms/test_qmc_transform.py
@@ -85,7 +85,7 @@ def unitary_z(basis_state):
         return qml.state()
 
     bitstrings = list(itertools.product([0, 1], repeat=n_wires))
-    u = [unitary_z(np.array(bitstring)).numpy() for bitstring in bitstrings]
+    u = [unitary_z(np.array(bitstring)) for bitstring in bitstrings]
     u = np.array(u).T
     return u
 
diff --git a/tests/transforms/test_sign_expand.py b/tests/transforms/test_sign_expand.py
index 2641d749f1b..eec92d9ee9d 100644
--- a/tests/transforms/test_sign_expand.py
+++ b/tests/transforms/test_sign_expand.py
@@ -98,7 +98,7 @@ def test_hamiltonians(self, tape, output):
         """Tests that the sign_expand transform returns the correct value"""
 
         tapes, fn = qml.transforms.sign_expand(tape)
-        results = dev.batch_execute(tapes)
+        results = dev.execute(tapes)
         expval = fn(results)
 
         assert np.isclose(output, expval)
@@ -125,7 +125,7 @@ def test_hamiltonians_circuit_impl(self, tape, output):
         """
 
         tapes, fn = qml.transforms.sign_expand(tape, circuit=True)
-        results = dev.batch_execute(tapes)
+        results = dev.execute(tapes)
         expval = fn(results)
 
         assert np.isclose(output, expval, 1e-2)
@@ -186,7 +186,7 @@ def test_hamiltonians_vars(self, tape, output):
         """Tests that the sign_expand transform returns the correct value"""
 
         tapes, fn = qml.transforms.sign_expand(tape)
-        results = dev.batch_execute(tapes)
+        results = dev.execute(tapes)
         expval = fn(results)
 
         assert np.isclose(output, expval)
@@ -198,7 +198,7 @@ def test_hamiltonians_vars_circuit_impl(self, tape, output):
         """
 
         tapes, fn = qml.transforms.sign_expand(tape, circuit=True)
-        results = dev.batch_execute(tapes)
+        results = dev.execute(tapes)
         expval = fn(results)
 
         assert np.isclose(output, expval, 1e-1)
diff --git a/tests/transforms/test_specs.py b/tests/transforms/test_specs.py
index bcc8540318a..10c499be588 100644
--- a/tests/transforms/test_specs.py
+++ b/tests/transforms/test_specs.py
@@ -49,16 +49,12 @@ def circ():
         assert info["num_device_wires"] == 1
         assert info["diff_method"] == diff_method
         assert info["num_trainable_params"] == 0
+        assert info["device_name"] == dev.name
 
         if diff_method == "parameter-shift":
             assert info["num_gradient_executions"] == 0
             assert info["gradient_fn"] == "pennylane.gradients.parameter_shift.param_shift"
 
-        if diff_method != "backprop":
-            assert info["device_name"] == "default.qubit"
-        else:
-            assert info["device_name"] == "default.qubit.autograd"
-
     @pytest.mark.parametrize(
         "diff_method, len_info", [("backprop", 11), ("parameter-shift", 12), ("adjoint", 11)]
     )
@@ -96,18 +92,14 @@ def circuit(x, y, add_RY=True):
 
         assert info["num_observables"] == 2
         assert info["num_diagonalizing_gates"] == 1
-        assert info["num_device_wires"] == 4
+        assert info["num_device_wires"] == 3
         assert info["diff_method"] == diff_method
         assert info["num_trainable_params"] == 4
+        assert info["device_name"] == dev.name
 
         if diff_method == "parameter-shift":
             assert info["num_gradient_executions"] == 6
 
-        if diff_method != "backprop":
-            assert info["device_name"] == "default.qubit"
-        else:
-            assert info["device_name"] == "default.qubit.autograd"
-
     @pytest.mark.parametrize(
         "diff_method, len_info", [("backprop", 11), ("parameter-shift", 12), ("adjoint", 11)]
     )
@@ -143,10 +135,11 @@ def circuit(params):
 
         params_shape = qml.BasicEntanglerLayers.shape(n_layers=n_layers, n_wires=n_wires)
         rng = np.random.default_rng(seed=10)
-        params = rng.standard_normal(params_shape)
+        params = rng.standard_normal(params_shape)  # pylint:disable=no-member
 
         return circuit, params
 
+    @pytest.mark.xfail(reason="DefaultQubit2 does not support custom expansion depths")
     def test_max_expansion(self):
         """Test that a user can calculation specifications for a different max
         expansion parameter."""
@@ -167,7 +160,7 @@ def test_max_expansion(self):
         assert info["resources"] == expected_resources
         assert info["num_observables"] == 1
         assert info["num_device_wires"] == 5
-        assert info["device_name"] == "default.qubit.autograd"
+        assert info["device_name"] == "default.qubit"
         assert info["diff_method"] == "best"
         assert info["gradient_fn"] == "backprop"
 
@@ -209,3 +202,22 @@ def circuit():
         info = qml.specs(circuit)()
         assert info["diff_method"] == "test_specs.my_transform"
         assert info["gradient_fn"] == "test_specs.my_transform"
+
+    @pytest.mark.parametrize(
+        "device,num_wires",
+        [
+            (qml.device("default.qubit"), 1),
+            (qml.device("default.qubit", wires=2), 1),
+            (qml.device("default.qubit.legacy", wires=2), 2),
+        ],
+    )
+    def test_num_wires_source_of_truth(self, device, num_wires):
+        """Tests that num_wires behaves differently on old and new devices."""
+
+        @qml.qnode(device)
+        def circuit():
+            qml.PauliX(0)
+            return qml.state()
+
+        info = qml.specs(circuit)()
+        assert info["num_device_wires"] == num_wires

From 0a1b8c93a7c8de6348be38f3941a9beaf69e98a4 Mon Sep 17 00:00:00 2001
From: Mudit Pandey <mudit.pandey@xanadu.ai>
Date: Thu, 21 Sep 2023 10:15:01 -0400
Subject: [PATCH 096/131] [skip ci] Added qml.math.is_nan

---
 pennylane/devices/qubit/simulate.py | 59 ++++++++++++++++-------------
 pennylane/math/__init__.py          |  1 +
 pennylane/math/multi_dispatch.py    | 19 ++++++++++
 3 files changed, 53 insertions(+), 26 deletions(-)

diff --git a/pennylane/devices/qubit/simulate.py b/pennylane/devices/qubit/simulate.py
index 46a39cec74b..61f83456389 100644
--- a/pennylane/devices/qubit/simulate.py
+++ b/pennylane/devices/qubit/simulate.py
@@ -116,30 +116,36 @@ def get_final_state(circuit, debugger=None, interface=None):
         is_state_batched = is_state_batched or op.batch_size is not None
 
         if isinstance(op, qml.Projector):
-            # Handle postselection on mid-circuit measurements
-            if is_state_batched:
-                for i, _ in enumerate(state):
-                    norm = qml.math.norm(state[i])
-                    if qml.math.isclose(norm, 0.0):
-                        return qml.numpy.NaN, is_state_batched
-
-                    state[i] = state[i] / qml.math.norm(state[i])
-            else:
-                norm = qml.math.norm(state)
-                if qml.math.isclose(norm, 0.0):
-                    return qml.numpy.NaN, is_state_batched
-
-                state = state / norm
-
-                # defer_measurements will raise an error with batched shots or broadcasting so we can
-                # assume that both the state and shots are unbatched.
-                if circuit.shots:
-                    # Clip the number of shots using a binomial distribution using the probability of
-                    # measuring the postselected state.
-                    postselected_shots = binomial(circuit.shots.total_shots, norm)
-                    if postselected_shots == 0:
-                        raise RuntimeError("None of the samples meet the postselection criteria")
-                    circuit._shots = qml.measurements.Shots(postselected_shots)
+            # # Handle postselection on mid-circuit measurements
+            # if is_state_batched:
+            #     for i, _ in enumerate(state):
+            #         norm = qml.math.norm(state[i])
+            #         if qml.math.isclose(norm, 0.0):
+            #             state[i] = qml.math.asarray(
+            #                 [qml.numpy.NaN] * qml.math.size(state[i]),
+            #                 like=qml.math.get_interface(state),
+            #             )
+            #         else:
+            #             state[i] = state[i] / qml.math.norm(state[i])
+            norm = qml.math.norm(state)
+            if qml.math.isclose(norm, 0.0):
+                state = qml.math.asarray(
+                    [qml.numpy.NaN] * qml.math.size(state),
+                    like=qml.math.get_interface(state),
+                )
+                return state, is_state_batched
+
+            state = state / norm
+
+            # defer_measurements will raise an error with batched shots or broadcasting so we can
+            # assume that both the state and shots are unbatched.
+            if circuit.shots:
+                # Clip the number of shots using a binomial distribution using the probability of
+                # measuring the postselected state.
+                postselected_shots = binomial(circuit.shots.total_shots, norm)
+                if postselected_shots == 0:
+                    raise RuntimeError("None of the samples meet the postselection criteria")
+                circuit._shots = qml.measurements.Shots(postselected_shots)
 
     if set(circuit.op_wires) < set(circuit.wires):
         state = expand_state_over_wires(
@@ -173,8 +179,9 @@ def measure_final_state(circuit, state, is_state_batched, rng=None, prng_key=Non
     Returns:
         Tuple[TensorLike]: The measurement results
     """
-    if state is qml.numpy.NaN:
-        return state
+    if any(qml.math.isnan(state)):
+        # do something
+        state = 1
 
     circuit = circuit.map_to_standard_wires()
 
diff --git a/pennylane/math/__init__.py b/pennylane/math/__init__.py
index 9e8ec32837d..e57713e8d84 100644
--- a/pennylane/math/__init__.py
+++ b/pennylane/math/__init__.py
@@ -50,6 +50,7 @@
     gammainc,
     get_trainable_indices,
     iscomplex,
+    isnan,
     jax_argnums_to_tape_trainable,
     kron,
     matmul,
diff --git a/pennylane/math/multi_dispatch.py b/pennylane/math/multi_dispatch.py
index b676292a1a7..667333e1aca 100644
--- a/pennylane/math/multi_dispatch.py
+++ b/pennylane/math/multi_dispatch.py
@@ -852,6 +852,25 @@ def norm(tensor, like=None, **kwargs):
     return norm(tensor, **kwargs)
 
 
+@multi_dispatch()
+def isnan(tensor, like=None, **kwargs):
+    """Check a tensor element-wise for ``NaN`` values and return the result
+    as a boolean tensor with the same interface."""
+    if like == "jax":
+        from jax.numpy import isnan
+
+    elif like == "tensorflow":
+        from tensorflow.math import is_nan as isnan
+
+    elif like == "torch":
+        from torch import isnan
+
+    else:
+        from numpy import isnan
+
+    return isnan(tensor, **kwargs)
+
+
 @multi_dispatch(argnum=[1])
 def gammainc(m, t, like=None):
     r"""Return the lower incomplete Gamma function.

From 9cf3588afe70cdd128edac8fa3ade5c7d33ee695 Mon Sep 17 00:00:00 2001
From: Mudit Pandey <mudit.pandey@xanadu.ai>
Date: Thu, 21 Sep 2023 10:20:28 -0400
Subject: [PATCH 097/131] Update pennylane/transforms/defer_measurements.py

Co-authored-by: Matthew Silverman <matthews@xanadu.ai>
---
 pennylane/transforms/defer_measurements.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pennylane/transforms/defer_measurements.py b/pennylane/transforms/defer_measurements.py
index 27500db939d..a089f372191 100644
--- a/pennylane/transforms/defer_measurements.py
+++ b/pennylane/transforms/defer_measurements.py
@@ -213,7 +213,7 @@ def func(x, y):
             mp = qml.map_wires(mp, wire_map=wire_map)
         new_measurements.append(mp)
 
-    new_tape = QuantumTape(new_operations, new_measurements, shots=tape.shots)
+    new_tape = type(tape)(new_operations, new_measurements, shots=tape.shots)
     new_tape._qfunc_output = tape._qfunc_output  # pylint: disable=protected-access
 
     def null_postprocessing(results):

From 57d77bd67dd8fa751f7542286c0f3913457b0ed5 Mon Sep 17 00:00:00 2001
From: Christina Lee <christina@xanadu.ai>
Date: Wed, 20 Sep 2023 16:25:25 -0400
Subject: [PATCH 098/131] Make Measurements Pytrees (#4607)

This PR registers all `MeasurementProcess` objects as jax pytrees.
[sc-40588]
---
 doc/releases/changelog-dev.md              |  3 ++
 pennylane/measurements/__init__.py         | 17 ++++++++
 pennylane/measurements/classical_shadow.py | 15 +++++++
 pennylane/measurements/counts.py           |  4 ++
 pennylane/measurements/measurements.py     | 21 +++++++++-
 pennylane/measurements/mid_measure.py      |  4 ++
 pennylane/measurements/mutual_info.py      |  6 ++-
 pennylane/measurements/vn_entropy.py       |  4 ++
 pennylane/ops/functions/equal.py           |  8 +++-
 tests/measurements/test_measurements.py    | 47 ++++++++++++++++++++++
 10 files changed, 125 insertions(+), 4 deletions(-)

diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 053308c8193..83ed2a59d8d 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -71,6 +71,9 @@
 
 <h3>Improvements 🛠</h3>
 
+* `MeasurementProcess` objects are now registered as jax pytrees.
+  [(#4607)](https://github.com/PennyLaneAI/pennylane/pull/4607)
+
 * Tensor-network template `qml.MPS` now supports changing `offset` between subsequent blocks for more flexibility.
  [(#4531)](https://github.com/PennyLaneAI/pennylane/pull/4531)
 
diff --git a/pennylane/measurements/__init__.py b/pennylane/measurements/__init__.py
index db26c173a17..3728a0bf9a0 100644
--- a/pennylane/measurements/__init__.py
+++ b/pennylane/measurements/__init__.py
@@ -204,6 +204,23 @@ def circuit(x):
     functions will return an instance of the corresponding measurement process
     (e.g. :class:`CountsMP`). This decision is just for design purposes.
 
+.. details::
+    :title: Serialization and Pytree format
+    :href: serialization
+
+    PennyLane measurements are automatically registered as `Pytrees <https://jax.readthedocs.io/en/latest/pytrees.html>`_ .
+
+    The :class:`~.MeasurementProcess` definitions are sufficient for all PL measurements.
+
+    >>> H = 2.0 * qml.PauliX(0)
+    >>> mp = qml.expval(H)
+    >>> mp._flatten()
+    ((<Hamiltonian: terms=1, wires=[0]>, None), ())
+    >>> type(mp)._unflatten(*mp._flatten())
+    expval(  (2) [X0])
+    >>> jax.tree_util.tree_leaves(mp)
+    [2]
+
 Adding your new measurement to PennyLane
 ----------------------------------------
 
diff --git a/pennylane/measurements/classical_shadow.py b/pennylane/measurements/classical_shadow.py
index d7a6e997db0..4e88c215acf 100644
--- a/pennylane/measurements/classical_shadow.py
+++ b/pennylane/measurements/classical_shadow.py
@@ -232,6 +232,10 @@ def __init__(
         self.seed = seed
         super().__init__(wires=wires, id=id)
 
+    def _flatten(self):
+        metadata = (("wires", self.wires), ("seed", self.seed))
+        return (None, None), metadata
+
     @property
     def hash(self):
         """int: returns an integer hash uniquely representing the measurement process"""
@@ -469,6 +473,17 @@ class ShadowExpvalMP(MeasurementTransform):
             where the instance has to be identified
     """
 
+    def _flatten(self):
+        metadata = (
+            ("seed", self.seed),
+            ("k", self.k),
+        )
+        return (self.H,), metadata
+
+    @classmethod
+    def _unflatten(cls, data, metadata):
+        return cls(data[0], **dict(metadata))
+
     def __init__(
         self,
         H: Union[Operator, Sequence],
diff --git a/pennylane/measurements/counts.py b/pennylane/measurements/counts.py
index 9b5dc76ad21..b54b9b894c9 100644
--- a/pennylane/measurements/counts.py
+++ b/pennylane/measurements/counts.py
@@ -182,6 +182,10 @@ def __init__(
         self.all_outcomes = all_outcomes
         super().__init__(obs, wires, eigvals, id)
 
+    def _flatten(self):
+        metadata = (("wires", self.raw_wires), ("all_outcomes", self.all_outcomes))
+        return (self.obs, self._eigvals), metadata
+
     def __repr__(self):
         if self.obs is None:
             if self._eigvals is None:
diff --git a/pennylane/measurements/measurements.py b/pennylane/measurements/measurements.py
index 56e9227d193..26550bd02b5 100644
--- a/pennylane/measurements/measurements.py
+++ b/pennylane/measurements/measurements.py
@@ -24,6 +24,8 @@
 
 import pennylane as qml
 from pennylane.operation import Operator
+from pennylane.pytrees import register_pytree
+from pennylane.typing import TensorLike
 from pennylane.wires import Wires
 
 from .shots import Shots
@@ -125,12 +127,27 @@ class MeasurementProcess(ABC):
             where the instance has to be identified
     """
 
+    def __init_subclass__(cls, **_):
+        register_pytree(cls, cls._flatten, cls._unflatten)
+
+    def _flatten(self):
+        metadata = (("wires", self.raw_wires),)
+        return (self.obs, self._eigvals), metadata
+
+    @classmethod
+    def _unflatten(cls, data, metadata):
+        if data[0] is not None:
+            return cls(obs=data[0], **dict(metadata))
+        if data[1] is not None:
+            return cls(eigvals=data[1], **dict(metadata))
+        return cls(**dict(metadata))
+
     # pylint: disable=too-many-arguments
     def __init__(
         self,
         obs: Optional[Operator] = None,
         wires: Optional[Wires] = None,
-        eigvals=None,
+        eigvals: Optional[TensorLike] = None,
         id: Optional[str] = None,
     ):
         if obs is not None and obs.name == "MeasurementValue":
@@ -305,7 +322,7 @@ def wires(self):
 
         return (
             Wires.all_wires(self._wires)
-            if isinstance(self._wires, list)
+            if isinstance(self._wires, (tuple, list))
             else self._wires or Wires([])
         )
 
diff --git a/pennylane/measurements/mid_measure.py b/pennylane/measurements/mid_measure.py
index 08dfeb3503e..7e63eb9e831 100644
--- a/pennylane/measurements/mid_measure.py
+++ b/pennylane/measurements/mid_measure.py
@@ -130,6 +130,10 @@ class MidMeasureMP(MeasurementProcess):
         id (str): Custom label given to a measurement instance.
     """
 
+    def _flatten(self):
+        metadata = (("wires", self.raw_wires), ("reset", self.reset), ("id", self.id))
+        return (None, None), metadata
+
     def __init__(
         self,
         wires: Optional[Wires] = None,
diff --git a/pennylane/measurements/mutual_info.py b/pennylane/measurements/mutual_info.py
index e50f1bb9586..c5035bce77e 100644
--- a/pennylane/measurements/mutual_info.py
+++ b/pennylane/measurements/mutual_info.py
@@ -83,7 +83,7 @@ def circuit_mutual(x):
         raise qml.QuantumFunctionError(
             "Subsystems for computing mutual information must not overlap."
         )
-    return MutualInfoMP(wires=[wires0, wires1], log_base=log_base)
+    return MutualInfoMP(wires=(wires0, wires1), log_base=log_base)
 
 
 class MutualInfoMP(StateMeasurement):
@@ -99,6 +99,10 @@ class MutualInfoMP(StateMeasurement):
 
     """
 
+    def _flatten(self):
+        metadata = (("wires", tuple(self.raw_wires)), ("log_base", self.log_base))
+        return (None, None), metadata
+
     # pylint: disable=too-many-arguments
     def __init__(
         self,
diff --git a/pennylane/measurements/vn_entropy.py b/pennylane/measurements/vn_entropy.py
index 6e852674c32..34da3a66236 100644
--- a/pennylane/measurements/vn_entropy.py
+++ b/pennylane/measurements/vn_entropy.py
@@ -83,6 +83,10 @@ class VnEntropyMP(StateMeasurement):
         log_base (float): Base for the logarithm.
     """
 
+    def _flatten(self):
+        metadata = (("wires", self.raw_wires), ("log_base", self.log_base))
+        return (None, None), metadata
+
     # pylint: disable=too-many-arguments, unused-argument
     def __init__(
         self,
diff --git a/pennylane/ops/functions/equal.py b/pennylane/ops/functions/equal.py
index 52f3dee34b7..b7af0f8f422 100644
--- a/pennylane/ops/functions/equal.py
+++ b/pennylane/ops/functions/equal.py
@@ -24,6 +24,7 @@
 from pennylane.measurements.mid_measure import MidMeasureMP
 from pennylane.measurements.mutual_info import MutualInfoMP
 from pennylane.measurements.vn_entropy import VnEntropyMP
+from pennylane.measurements.counts import CountsMP
 from pennylane.pulse.parametrized_evolution import ParametrizedEvolution
 from pennylane.operation import Observable, Operator, Tensor
 from pennylane.ops import Hamiltonian, Controlled, Pow, Adjoint, Exp, SProd, CompositeOp
@@ -389,7 +390,7 @@ def _(op1: MutualInfoMP, op2: MutualInfoMP, **kwargs):
 
 @_equal.register
 # pylint: disable=unused-argument
-def _equal_shadow_measurements(op1: ShadowExpvalMP, op2: ShadowExpvalMP, **kwargs):
+def _equal_shadow_measurements(op1: ShadowExpvalMP, op2: ShadowExpvalMP, **_):
     """Determine whether two ShadowExpvalMP objects are equal"""
 
     wires_match = op1.wires == op2.wires
@@ -404,3 +405,8 @@ def _equal_shadow_measurements(op1: ShadowExpvalMP, op2: ShadowExpvalMP, **kwarg
     k_match = op1.k == op2.k
 
     return wires_match and H_match and k_match
+
+
+@_equal.register
+def _equal_counts(op1: CountsMP, op2: CountsMP, **kwargs):
+    return _equal_measurements(op1, op2, **kwargs) and op1.all_outcomes == op2.all_outcomes
diff --git a/tests/measurements/test_measurements.py b/tests/measurements/test_measurements.py
index 84cccd5514e..80ca82eda98 100644
--- a/tests/measurements/test_measurements.py
+++ b/tests/measurements/test_measurements.py
@@ -25,7 +25,9 @@
     MeasurementProcess,
     MeasurementTransform,
     MidMeasure,
+    MidMeasureMP,
     MutualInfoMP,
+    PurityMP,
     Probability,
     ProbabilityMP,
     Sample,
@@ -45,6 +47,7 @@
 )
 from pennylane.operation import DecompositionUndefinedError
 from pennylane.queuing import AnnotatedQueue
+from pennylane.wires import Wires
 
 # pylint: disable=too-few-public-methods, unused-argument
 
@@ -153,6 +156,50 @@ class DummyMP(MeasurementProcess):
     assert hash(mp1) == hash(mp2)
 
 
+valid_meausurements = [
+    ClassicalShadowMP(wires=Wires(0), seed=42),
+    ShadowExpvalMP(qml.s_prod(3.0, qml.PauliX(0)), seed=97, k=2),
+    ShadowExpvalMP([qml.PauliZ(0), 4.0 * qml.PauliX(0)], seed=86, k=4),
+    CountsMP(obs=2.0 * qml.PauliX(0), all_outcomes=True),
+    CountsMP(eigvals=[0.5, 0.6], wires=Wires(0), all_outcomes=False),
+    ExpectationMP(obs=qml.s_prod(2.0, qml.PauliX(0))),
+    ExpectationMP(eigvals=[0.5, 0.6], wires=Wires("a")),
+    MidMeasureMP(wires=Wires("a"), reset=True, id="abcd"),
+    MutualInfoMP(wires=(Wires("a"), Wires("b")), log_base=3),
+    ProbabilityMP(wires=Wires("a"), eigvals=[0.5, 0.6]),
+    ProbabilityMP(obs=3.0 * qml.PauliX(0)),
+    PurityMP(wires=Wires("a")),
+    SampleMP(obs=3.0 * qml.PauliY(0)),
+    SampleMP(wires=Wires("a"), eigvals=[0.5, 0.6]),
+    StateMP(),
+    StateMP(wires=("a", "b")),
+    VarianceMP(obs=qml.s_prod(0.5, qml.PauliX(0))),
+    VarianceMP(eigvals=[0.6, 0.7], wires=Wires(0)),
+    VnEntropyMP(wires=Wires("a"), log_base=3),
+]
+
+
+# pylint: disable=protected-access
+@pytest.mark.parametrize("mp", valid_meausurements)
+def test_flatten_unflatten(mp):
+    """Test flatten and unflatten methods."""
+
+    data, metadata = mp._flatten()
+    assert hash(metadata)
+
+    new_mp = type(mp)._unflatten(data, metadata)
+    assert qml.equal(new_mp, mp)
+
+
+@pytest.mark.jax
+@pytest.mark.parametrize("mp", valid_meausurements)
+def test_jax_pytree_integration(mp):
+    """Test that measurement processes are jax pytrees."""
+    import jax
+
+    jax.tree_util.tree_flatten(mp)
+
+
 @pytest.mark.parametrize(
     "stat_func,return_type", [(expval, Expectation), (var, Variance), (sample, Sample)]
 )

From a4ae4c0cecf580cc371332366d12bca1f35f83c9 Mon Sep 17 00:00:00 2001
From: Stepan Fomichev <stepan.fomichev95@gmail.com>
Date: Wed, 20 Sep 2023 17:07:25 -0400
Subject: [PATCH 099/131] Converter for MPS DMRG wavefunctions (#4523)

---
 doc/releases/changelog-dev.md        |   5 ++
 pennylane/qchem/convert.py           | 115 +++++++++++++++++++++++++++
 tests/qchem/of_tests/test_convert.py |  37 +++++++++
 3 files changed, 157 insertions(+)

diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 83ed2a59d8d..f5b6de8fd4f 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -71,6 +71,10 @@
 
 <h3>Improvements 🛠</h3>
 
+* Extended `qml.qchem.import_state` to import wavefunctions from MPS DMRG classical calculations 
+  performed with the Block2 library. 
+  [#4523](https://github.com/PennyLaneAI/pennylane/pull/4523)
+
 * `MeasurementProcess` objects are now registered as jax pytrees.
   [(#4607)](https://github.com/PennyLaneAI/pennylane/pull/4607)
 
@@ -313,6 +317,7 @@
 This release contains contributions from (in alphabetical order):
 
 Utkarsh Azad,
+Stepan Fomichev,
 Diego Guala,
 Soran Jahangiri,
 Christina Lee,
diff --git a/pennylane/qchem/convert.py b/pennylane/qchem/convert.py
index 87f1ec3d971..faa1f935b8f 100644
--- a/pennylane/qchem/convert.py
+++ b/pennylane/qchem/convert.py
@@ -1007,3 +1007,118 @@ def _uccsd_state(ccsd_solver, tol=1e-15):
     dict_fcimatr = {key: value for key, value in dict_fcimatr.items() if abs(value) > tol}
 
     return dict_fcimatr
+
+
+def _dmrg_state(wavefunction, tol=1e-15):
+    r"""Construct a wavefunction from the DMRG wavefunction obtained from the Block2 library.
+
+    The generated wavefunction is a dictionary where the keys represent a configuration, which
+    corresponds to a Slater determinant, and the values are the CI coefficients of the Slater
+    determinant. Each dictionary key is a tuple of two integers. The binary representation of these
+    integers correspond to a specific configuration: the first number represents the
+    configuration of the alpha electrons and the second number represents the configuration of the
+    beta electrons. For instance, the Hartree-Fock state :math:`|1 1 0 0 \rangle` will be
+    represented by the flipped binary string ``0011`` which is split to ``01`` and ``01`` for
+    the alpha and beta electrons. The integer corresponding to ``01`` is ``1`` and the dictionary
+    representation of the Hartree-Fock state will be ``{(1, 1): 1.0}``. The dictionary
+    representation of a state with ``0.99`` contribution from the Hartree-Fock state and ``0.01``
+    contribution from the doubly-excited state, i.e., :math:`|0 0 1 1 \rangle`, will be
+    ``{(1, 1): 0.99, (2, 2): 0.01}``.
+
+    The determinants and coefficients should be supplied externally. They should be calculated by
+    using Block2 DMRGDriver's `get_csf_coefficients()` method.
+
+    Args:
+        wavefunction tuple(array[int], array[float]): determinants and coefficients in physicist notation
+        tol (float): the tolerance for discarding Slater determinants with small coefficients
+
+    Returns:
+        dict: dictionary of the form `{(int_a, int_b) :coeff}`, with integers `int_a, int_b`
+        having binary representation corresponding to the Fock occupation vector in alpha and beta
+        spin sectors, respectively, and coeff being the CI coefficients of those configurations
+
+    **Example**
+
+    >>> from pyscf import gto, scf, mcscf
+    >>> from pyblock2.driver.core import DMRGDriver, SymmetryTypes
+    >>> from pyblock2._pyscf.ao2mo import integrals as itg
+    >>> mol = gto.M(atom=[['H', (0, 0, 0)], ['H', (0,0,0.71)]], basis='sto3g')
+    >>> mf = scf.RHF(mol).run()
+    >>> mc = mcscf.CASCI(mf, 2, 2)
+    >>> ncas, n_elec, spin, ecore, h1e, g2e, orb_sym = itg.get_rhf_integrals(mf, mc.ncore, mc.ncas, g2e_symm=8)
+    >>> driver = DMRGDriver(scratch="./tmp", symm_type=SymmetryTypes.SU2)
+    >>> driver.initialize_system(n_sites=ncas, n_elec=n_elec, spin=spin, orb_sym=orb_sym)
+    >>> mpo = driver.get_qc_mpo(h1e=h1e, g2e=g2e, ecore=ecore)
+    >>> ket = driver.get_random_mps(tag="GS")
+    >>> wavefunction = driver.get_csf_coefficients(ket)
+    >>> wf_dmrg = _dmrg_state(wavefunction, tol=1e-1)
+    >>> print(wf_dmrg)
+    {(2, 2): 0.6995876453534665, (1, 2): 0.7021101425511932, (1, 1): 0.13273460059658818}
+    """
+    dets, coeffs = wavefunction
+
+    row, col, dat = [], [], []
+
+    for ii, det in enumerate(dets):
+        stra, strb = _sitevec_to_fock(det.tolist(), format="dmrg")
+        row.append(stra)
+        col.append(strb)
+
+        # compute and fix parity to stick to pyscf notation
+        lsta = np.array(list(map(int, bin(stra)[2:])))[::-1]
+        lstb = np.array(list(map(int, bin(strb)[2:])))[::-1]
+
+        # pad the relevant list
+        maxlen = max([len(lsta), len(lstb)])
+        lsta = np.pad(lsta, (0, maxlen - len(lsta)))
+        lstb = np.pad(lstb, (0, maxlen - len(lstb)))
+
+        which_occ = np.where(lsta == 1)[0]
+        parity = (-1) ** np.sum([np.sum(lstb[: int(ind)]) for ind in which_occ])
+        dat.append(parity * coeffs[ii])
+
+    ## create the FCI matrix as a dict
+    dict_fcimatr = dict(zip(list(zip(row, col)), dat))
+
+    # filter based on tolerance cutoff
+    dict_fcimatr = {key: value for key, value in dict_fcimatr.items() if abs(value) > tol}
+
+    return dict_fcimatr
+
+
+def _sitevec_to_fock(det, format):
+    r"""Covert a Slater determinant from site vector to occupation number vector representation.
+
+    Args:
+        det (list): determinant in site vector representation
+        format (str): the format of the determinant
+
+    Returns:
+        tuple: tuple of integers representing binaries that correspond to occupation vectors in
+            alpha and beta spin sectors
+
+    **Example**
+
+    >>> det = [1, 2, 1, 0, 0, 2]
+    >>> _sitevec_to_fock(det, format = 'dmrg')
+    (5, 34)
+
+    >>> det = ["a", "b", "a", "0", "0", "b"]
+    >>> _sitevec_to_fock(det, format = 'shci')
+    (5, 34)
+    """
+
+    if format == "dmrg":
+        format_map = {0: "00", 1: "10", 2: "01", 3: "11"}
+    elif format == "shci":
+        format_map = {"0": "00", "a": "10", "b": "01", "2": "11"}
+
+    strab = [format_map[key] for key in det]
+
+    stra = "".join(i[0] for i in strab)
+    strb = "".join(i[1] for i in strab)
+
+    inta = int(stra[::-1], 2)
+    intb = int(strb[::-1], 2)
+
+    return inta, intb
diff --git a/tests/qchem/of_tests/test_convert.py b/tests/qchem/of_tests/test_convert.py
index 302c87c40f1..ff857557c36 100644
--- a/tests/qchem/of_tests/test_convert.py
+++ b/tests/qchem/of_tests/test_convert.py
@@ -1076,3 +1076,40 @@ def test_rccsd_state(molecule, basis, symm, tol, wf_ref):
 
     assert wf_ccsd.keys() == wf_ref.keys()
     assert np.allclose(abs(np.array(list(wf_ccsd.values()))), abs(np.array(list(wf_ref.values()))))
+
+
+@pytest.mark.parametrize(
+    ("sitevec", "format", "state_ref"),
+    [([1, 2, 1, 0, 0, 2], "dmrg", (5, 34)), (["a", "b", "a", "0", "0", "b"], "shci", (5, 34))],
+)
+def test_sitevec_to_fock(sitevec, format, state_ref):
+    r"""Test that _sitevec_to_fock returns the correct state."""
+
+    state = qml.qchem.convert._sitevec_to_fock(sitevec, format)
+
+    assert state == state_ref
+
+
+@pytest.mark.parametrize(
+    ("wavefunction", "state_ref"),
+    [
+        (
+            (np.array([[0, 3], [3, 0]]), np.array([-0.10660077, 0.9943019])),
+            {(2, 2): np.array([-0.10660077]), (1, 1): np.array([0.9943019])},
+        ),
+        (
+            (np.array([[0, 3], [1, 2], [3, 0]]), np.array([0.69958765, 0.70211014, 0.1327346])),
+            {
+                (2, 2): np.array([0.69958765]),
+                (1, 2): np.array([0.70211014]),
+                (1, 1): np.array([0.1327346]),
+            },
+        ),
+    ],
+)
+def test_dmrg_state(wavefunction, state_ref):
+    r"""Test that _dmrg_state returns the correct state."""
+
+    state = qml.qchem.convert._dmrg_state(wavefunction)
+
+    assert state == state_ref

From d0240d850fc256a7e433da545e9930f80aeed4be Mon Sep 17 00:00:00 2001
From: Stepan Fomichev <stepan.fomichev95@gmail.com>
Date: Wed, 20 Sep 2023 17:58:25 -0400
Subject: [PATCH 100/131] Converter for SHCI wavefunction (#4524)

---
 doc/releases/changelog-dev.md        |  7 +--
 pennylane/qchem/convert.py           | 71 ++++++++++++++++++++++++++++
 tests/qchem/of_tests/test_convert.py | 28 +++++++++++
 3 files changed, 103 insertions(+), 3 deletions(-)

diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index f5b6de8fd4f..ff0efbd7b38 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -71,15 +71,16 @@
 
 <h3>Improvements 🛠</h3>
 
-* Extended `qml.qchem.import_state` to import wavefunctions from MPS DMRG classical calculations 
-  performed with the Block2 library. 
+* Extended ``qml.qchem.import_state`` to import wavefunctions from MPS DMRG and SHCI classical
+  calculations performed with the Block2 and Dice libraries.
   [#4523](https://github.com/PennyLaneAI/pennylane/pull/4523)
+  [#4524](https://github.com/PennyLaneAI/pennylane/pull/4524)
 
 * `MeasurementProcess` objects are now registered as jax pytrees.
   [(#4607)](https://github.com/PennyLaneAI/pennylane/pull/4607)
 
 * Tensor-network template `qml.MPS` now supports changing `offset` between subsequent blocks for more flexibility.
- [(#4531)](https://github.com/PennyLaneAI/pennylane/pull/4531)
+  [(#4531)](https://github.com/PennyLaneAI/pennylane/pull/4531)
 
 * The qchem ``fermionic_dipole`` and ``particle_number`` functions are updated to use a
   ``FermiSentence``. The deprecated features for using tuples to represent fermionic operations are
diff --git a/pennylane/qchem/convert.py b/pennylane/qchem/convert.py
index faa1f935b8f..265e4666390 100644
--- a/pennylane/qchem/convert.py
+++ b/pennylane/qchem/convert.py
@@ -1122,3 +1122,74 @@ def _sitevec_to_fock(det, format):
     intb = int(strb[::-1], 2)
 
     return inta, intb
+
+
+def _shci_state(wavefunction, tol=1e-15):
+    r"""Construct a wavefunction from the SHCI wavefunction obtained from the Dice library.
+
+    The generated wavefunction is a dictionary where the keys represent a configuration, which
+    corresponds to a Slater determinant, and the values are the CI coefficients of the Slater
+    determinant. Each dictionary key is a tuple of two integers. The binary representation of these
+    integers correspond to a specific configuration: the first number represents the
+    configuration of the alpha electrons and the second number represents the configuration of the
+    beta electrons. For instance, the Hartree-Fock state :math:`|1 1 0 0 \rangle` will be
+    represented by the flipped binary string ``0011`` which is split to ``01`` and ``01`` for
+    the alpha and beta electrons. The integer corresponding to ``01`` is ``1`` and the dictionary
+    representation of the Hartree-Fock state will be ``{(1, 1): 1.0}``. The dictionary
+    representation of a state with ``0.99`` contribution from the Hartree-Fock state and ``0.01``
+    contribution from the doubly-excited state, i.e., :math:`|0 0 1 1 \rangle`, will be
+    ``{(1, 1): 0.99, (2, 2): 0.01}``.
+
+    The determinants and coefficients should be supplied externally. They are typically stored under
+    SHCI.outputfile.
+
+    Args:
+        wavefunction tuple(array[int], array[str]): determinants and coefficients in physicist notation
+        tol (float): the tolerance for discarding Slater determinants with small coefficients
+    Returns:
+        dict: dictionary of the form `{(int_a, int_b) :coeff}`, with integers `int_a, int_b`
+        having binary representation corresponding to the Fock occupation vector in alpha and beta
+        spin sectors, respectively, and coeff being the CI coefficients of those configurations
+
+    **Example**
+
+    >>> from pyscf import gto, scf, mcscf
+    >>> from pyscf.shciscf import shci
+    >>> import numpy as np
+    >>> mol = gto.M(atom=[['Li', (0, 0, 0)], ['Li', (0,0,0.71)]], basis='sto6g', symmetry="d2h")
+    >>> myhf = scf.RHF(mol).run()
+    >>> ncas, nelecas_a, nelecas_b = mol.nao, mol.nelectron // 2, mol.nelectron // 2
+    >>> myshci = mcscf.CASCI(myhf, ncas, (nelecas_a, nelecas_b))
+    >>> initialStates = myshci.getinitialStateSHCI(myhf, (nelecas_a, nelecas_b))
+    >>> output_file = f"shci_output.out"
+    >>> myshci.fcisolver = shci.SHCI(myhf.mol)
+    >>> myshci.internal_rotation = False
+    >>> myshci.fcisolver.initialStates = initialStates
+    >>> myshci.fcisolver.outputFile = output_file
+    >>> e_shci = np.atleast_1d(myshci.kernel(verbose=5))
+    >>> wf_shci = _shci_state(myshci, tol=1e-1)
+    >>> print(wf_shci)
+    {(7, 7): 0.8874167069, (11, 11): -0.3075774156, (19, 19): -0.3075774156, (35, 35): -0.1450474361}
+    """
+
+    dets, coefs = wavefunction
+
+    xa = []
+    xb = []
+    dat = []
+
+    for coef, det in zip(coefs, dets):
+        if abs(coef) > tol:
+            bin_a, bin_b = _sitevec_to_fock(list(det), "shci")
+
+            xa.append(bin_a)
+            xb.append(bin_b)
+            dat.append(coef)
+
+    ## create the FCI matrix as a dict
+    dict_fcimatr = dict(zip(list(zip(xa, xb)), dat))
+
+    # filter based on tolerance cutoff
+    dict_fcimatr = {key: value for key, value in dict_fcimatr.items() if abs(value) > tol}
+
+    return dict_fcimatr
diff --git a/tests/qchem/of_tests/test_convert.py b/tests/qchem/of_tests/test_convert.py
index ff857557c36..ab3a9e32789 100644
--- a/tests/qchem/of_tests/test_convert.py
+++ b/tests/qchem/of_tests/test_convert.py
@@ -1113,3 +1113,31 @@ def test_dmrg_state(wavefunction, state_ref):
     state = qml.qchem.convert._dmrg_state(wavefunction)
 
     assert state == state_ref
+
+
+@pytest.mark.parametrize(
+    ("wavefunction", "state_ref"),
+    [
+        (
+            (
+                ["02", "20"],
+                np.array([-0.10660077, 0.9943019]),
+            ),
+            {(2, 2): np.array([-0.10660077]), (1, 1): np.array([0.9943019])},
+        ),
+        (
+            (["02", "ab", "20"], np.array([0.69958765, 0.70211014, 0.1327346])),
+            {
+                (2, 2): np.array([0.69958765]),
+                (1, 2): np.array([0.70211014]),
+                (1, 1): np.array([0.1327346]),
+            },
+        ),
+    ],
+)
+def test_shci_state(wavefunction, state_ref):
+    r"""Test that _shci_state returns the correct state."""
+
+    state = qml.qchem.convert._shci_state(wavefunction)
+
+    assert state == state_ref

From 534601bee6ac09ddf5b19ce1eef22d003fdc9aca Mon Sep 17 00:00:00 2001
From: Mudit Pandey <mudit.pandey@xanadu.ai>
Date: Thu, 21 Sep 2023 10:47:01 -0400
Subject: [PATCH 101/131] Updated measurements pytree stuff

---
 pennylane/measurements/counts.py            |  2 +-
 pennylane/measurements/measurements.py      |  2 +-
 tests/measurements/test_measurements.py     |  7 +++++
 tests/transforms/test_defer_measurements.py | 33 +++++++++++++++++++++
 4 files changed, 42 insertions(+), 2 deletions(-)

diff --git a/pennylane/measurements/counts.py b/pennylane/measurements/counts.py
index b54b9b894c9..e1a8e42c257 100644
--- a/pennylane/measurements/counts.py
+++ b/pennylane/measurements/counts.py
@@ -184,7 +184,7 @@ def __init__(
 
     def _flatten(self):
         metadata = (("wires", self.raw_wires), ("all_outcomes", self.all_outcomes))
-        return (self.obs, self._eigvals), metadata
+        return (self.obs or self.mv, self._eigvals), metadata
 
     def __repr__(self):
         if self.obs is None:
diff --git a/pennylane/measurements/measurements.py b/pennylane/measurements/measurements.py
index 26550bd02b5..ab1ed6062b1 100644
--- a/pennylane/measurements/measurements.py
+++ b/pennylane/measurements/measurements.py
@@ -132,7 +132,7 @@ def __init_subclass__(cls, **_):
 
     def _flatten(self):
         metadata = (("wires", self.raw_wires),)
-        return (self.obs, self._eigvals), metadata
+        return (self.obs or self.mv, self._eigvals), metadata
 
     @classmethod
     def _unflatten(cls, data, metadata):
diff --git a/tests/measurements/test_measurements.py b/tests/measurements/test_measurements.py
index 80ca82eda98..9540c78efaa 100644
--- a/tests/measurements/test_measurements.py
+++ b/tests/measurements/test_measurements.py
@@ -156,25 +156,32 @@ class DummyMP(MeasurementProcess):
     assert hash(mp1) == hash(mp2)
 
 
+mv = qml.measure(0)
+
 valid_meausurements = [
     ClassicalShadowMP(wires=Wires(0), seed=42),
     ShadowExpvalMP(qml.s_prod(3.0, qml.PauliX(0)), seed=97, k=2),
     ShadowExpvalMP([qml.PauliZ(0), 4.0 * qml.PauliX(0)], seed=86, k=4),
     CountsMP(obs=2.0 * qml.PauliX(0), all_outcomes=True),
     CountsMP(eigvals=[0.5, 0.6], wires=Wires(0), all_outcomes=False),
+    CountsMP(obs=mv, all_outcomes=True),
     ExpectationMP(obs=qml.s_prod(2.0, qml.PauliX(0))),
     ExpectationMP(eigvals=[0.5, 0.6], wires=Wires("a")),
+    ExpectationMP(obs=mv),
     MidMeasureMP(wires=Wires("a"), reset=True, id="abcd"),
     MutualInfoMP(wires=(Wires("a"), Wires("b")), log_base=3),
     ProbabilityMP(wires=Wires("a"), eigvals=[0.5, 0.6]),
     ProbabilityMP(obs=3.0 * qml.PauliX(0)),
+    ProbabilityMP(obs=mv),
     PurityMP(wires=Wires("a")),
     SampleMP(obs=3.0 * qml.PauliY(0)),
     SampleMP(wires=Wires("a"), eigvals=[0.5, 0.6]),
+    SampleMP(obs=mv),
     StateMP(),
     StateMP(wires=("a", "b")),
     VarianceMP(obs=qml.s_prod(0.5, qml.PauliX(0))),
     VarianceMP(eigvals=[0.6, 0.7], wires=Wires(0)),
+    VarianceMP(obs=mv),
     VnEntropyMP(wires=Wires("a"), log_base=3),
 ]
 
diff --git a/tests/transforms/test_defer_measurements.py b/tests/transforms/test_defer_measurements.py
index 2905ed10e20..3075bcb6603 100644
--- a/tests/transforms/test_defer_measurements.py
+++ b/tests/transforms/test_defer_measurements.py
@@ -297,6 +297,39 @@ def circ2(x):
         param = 1.5
         assert np.allclose(circ1(param, shots=shots), circ2(param, shots=shots))
 
+    @pytest.mark.parametrize("shots", [None, 1000, [1000, 1000]])
+    def test_measured_value_wires_mapped(self, shots, tol, tol_stochastic):
+        """Test that collecting statistics on a measurement value works correctly
+        when the measured wire is reused."""
+        dev = DefaultQubit()
+
+        @qml.qnode(dev)
+        def circ1(x):
+            qml.RX(x, 0)
+            m0 = qml.measure(0)
+            qml.PauliX(0)
+            return qml.probs(op=m0)
+
+        dev = DefaultQubit()
+
+        @qml.qnode(dev)
+        def circ2(x):
+            qml.RX(x, 0)
+            return qml.probs(wires=[0])
+
+        param = 1.5
+        atol = tol if shots is None else tol_stochastic
+        assert np.allclose(circ1(param, shots=shots), circ2(param, shots=shots), atol=atol, rtol=0)
+
+        expected_ops = [qml.RX(param, 0), qml.CNOT([0, 1]), qml.PauliX(0)]
+        assert circ1.qtape.operations == expected_ops
+
+        assert len(circ1.qtape.measurements) == 1
+        mp = circ1.qtape.measurements[0]
+        assert isinstance(mp, qml.measurements.ProbabilityMP)
+        assert mp.mv is not None
+        assert mp.mv.wires == qml.wires.Wires([1])
+
     @pytest.mark.parametrize("shots", [None, 1000, [1000, 1000]])
     def test_terminal_measurements(self, shots):
         """Test that mid-circuit measurement statistics and terminal measurements

From e85674939a80216995e3c340e0a76798a5bff39a Mon Sep 17 00:00:00 2001
From: Mudit Pandey <mudit.pandey@xanadu.ai>
Date: Thu, 21 Sep 2023 12:31:32 -0400
Subject: [PATCH 102/131] Added tape splitting to defer_measurements

---
 pennylane/qnode.py                         |  8 +++++---
 pennylane/transforms/defer_measurements.py | 24 ++++++++++++----------
 2 files changed, 18 insertions(+), 14 deletions(-)

diff --git a/pennylane/qnode.py b/pennylane/qnode.py
index 8c0bddcd55d..5b4ecdb1702 100644
--- a/pennylane/qnode.py
+++ b/pennylane/qnode.py
@@ -899,11 +899,13 @@ def construct(self, args, kwargs):  # pylint: disable=too-many-branches
 
         # Apply the deferred measurement principle if the device doesn't
         # support mid-circuit measurements natively
-        expand_mid_measure = any(isinstance(op, MidMeasureMP) for op in self.tape.operations) and (
-            isinstance(self.device, qml.devices.Device)
-            or not self.device.capabilities().get("supports_mid_measure", False)
+        expand_mid_measure = (
+            any(isinstance(op, MidMeasureMP) for op in self.tape.operations)
+            and not isinstance(self.device, qml.devices.Device)
+            and not self.device.capabilities().get("supports_mid_measure", False)
         )
         if expand_mid_measure:
+            # Assume that tapes are not split if old device is used since postselection is not supported
             tapes, _ = qml.defer_measurements(self._tape)
             self._tape = tapes[0]
 
diff --git a/pennylane/transforms/defer_measurements.py b/pennylane/transforms/defer_measurements.py
index a089f372191..8f6bd808034 100644
--- a/pennylane/transforms/defer_measurements.py
+++ b/pennylane/transforms/defer_measurements.py
@@ -143,15 +143,12 @@ def func(x, y):
     measured_wires = []
     reused_measurement_wires = set()
     repeated_measurement_wire = False
+    is_postselecting = False
 
     for op in tape.operations:
         if isinstance(op, MidMeasureMP):
             if op.postselect is not None:
-                if tape.shots and (tape.shots.has_partitioned_shots or tape.batch_size):
-                    raise ValueError(
-                        "Cannot postselect on mid-circuit measurements when the circuit has finite "
-                        "shots if the circuit is batched or the shots are partitioned."
-                    )
+                is_postselecting = True
             if op.reset:
                 reused_measurement_wires.add(op.wires[0])
 
@@ -188,12 +185,13 @@ def func(x, y):
 
                 if op.reset:
                     with QueuingManager.stop_recording():
-                        if op.postselect == 1:
+                        # No need to manually reset if postselecting on |0>
+                        if op.postselect is None:
+                            new_operations.append(qml.CNOT([cur_wire, op.wires[0]]))
+                        elif op.postselect == 1:
                             # We know that the measured wire will be in the |1> state if postselected
                             # |1>. So we can just apply a PauliX instead of a CNOT to reset
                             new_operations.append(qml.PauliX(op.wires[0]))
-                        else:
-                            new_operations.append(qml.CNOT([cur_wire, op.wires[0]]))
 
                 cur_wire += 1
             else:
@@ -209,6 +207,7 @@ def func(x, y):
 
     for mp in tape.measurements:
         if mp.mv is not None:
+            # Update measurement value wires
             wire_map = {m.wires[0]: control_wires[m.id] for m in mp.mv.measurements}
             mp = qml.map_wires(mp, wire_map=wire_map)
         new_measurements.append(mp)
@@ -216,10 +215,13 @@ def func(x, y):
     new_tape = type(tape)(new_operations, new_measurements, shots=tape.shots)
     new_tape._qfunc_output = tape._qfunc_output  # pylint: disable=protected-access
 
+    if is_postselecting and new_tape.batch_size is not None:
+        # Split tapes if broadcasting with postselection
+        return qml.transforms.broadcast_expand(new_tape)
+
     def null_postprocessing(results):
-        """A postprocesing function returned by a transform that only converts the batch of results
-        into a result for a single ``QuantumTape``.
-        """
+        """A postprocessing function returned by a transform that only converts the batch of results
+        into a result for a single ``QuantumTape``."""
         return results[0]
 
     return [new_tape], null_postprocessing

From 7ee28f7ecbc90ac01d85c414d06ddece480f904a Mon Sep 17 00:00:00 2001
From: Mudit Pandey <mudit.pandey@xanadu.ai>
Date: Thu, 21 Sep 2023 15:13:39 -0400
Subject: [PATCH 103/131] Updated measure.py

---
 pennylane/devices/qubit/measure.py  | 13 +++++++++++--
 pennylane/devices/qubit/simulate.py | 24 +++++++++---------------
 2 files changed, 20 insertions(+), 17 deletions(-)

diff --git a/pennylane/devices/qubit/measure.py b/pennylane/devices/qubit/measure.py
index 9b72def187d..fb047f85d1c 100644
--- a/pennylane/devices/qubit/measure.py
+++ b/pennylane/devices/qubit/measure.py
@@ -19,6 +19,7 @@
 from scipy.sparse import csr_matrix
 
 from pennylane import math
+from pennylane.numpy import NaN
 from pennylane.ops import Sum, Hamiltonian
 from pennylane.measurements import (
     StateMeasurement,
@@ -62,12 +63,14 @@ def state_diagonalizing_gates(
     Returns:
         TensorLike: the result of the measurement
     """
-    for op in measurementprocess.diagonalizing_gates():
-        state = apply_operation(op, state, is_state_batched=is_state_batched)
+    if not any(math.isnan(state)):
+        for op in measurementprocess.diagonalizing_gates():
+            state = apply_operation(op, state, is_state_batched=is_state_batched)
 
     total_indices = len(state.shape) - is_state_batched
     wires = Wires(range(total_indices))
     flattened_state = flatten_state(state, total_indices)
+
     return measurementprocess.process_state(flattened_state, wires)
 
 
@@ -88,6 +91,12 @@ def csr_dot_products(
     total_wires = len(state.shape) - is_state_batched
     Hmat = measurementprocess.obs.sparse_matrix(wire_order=list(range(total_wires)))
 
+    if any(math.isnan(state)):
+        if is_state_batched:
+            return math.asarray([NaN] * math.shape(state)[0], like=math.get_interface(state))
+
+        return math.asarry(NaN, like=math.get_interface(state))
+
     if is_state_batched:
         state = math.toarray(state).reshape(math.shape(state)[0], -1)
 
diff --git a/pennylane/devices/qubit/simulate.py b/pennylane/devices/qubit/simulate.py
index 61f83456389..60dfd45f1b1 100644
--- a/pennylane/devices/qubit/simulate.py
+++ b/pennylane/devices/qubit/simulate.py
@@ -116,17 +116,15 @@ def get_final_state(circuit, debugger=None, interface=None):
         is_state_batched = is_state_batched or op.batch_size is not None
 
         if isinstance(op, qml.Projector):
-            # # Handle postselection on mid-circuit measurements
-            # if is_state_batched:
-            #     for i, _ in enumerate(state):
-            #         norm = qml.math.norm(state[i])
-            #         if qml.math.isclose(norm, 0.0):
-            #             state[i] = qml.math.asarray(
-            #                 [qml.numpy.NaN] * qml.math.size(state[i]),
-            #                 like=qml.math.get_interface(state),
-            #             )
-            #         else:
-            #             state[i] = state[i] / qml.math.norm(state[i])
+            # Handle postselection on mid-circuit measurements
+            if is_state_batched:
+                raise ValueError(
+                    "Cannot postselect on circuits with broadcasting. Use the "
+                    "qml.transforms.broadcast_expand transform to split a broadcasted "
+                    "tape into multiple non-broadcasted tapes before executing if "
+                    "postselection is used."
+                )
+
             norm = qml.math.norm(state)
             if qml.math.isclose(norm, 0.0):
                 state = qml.math.asarray(
@@ -179,10 +177,6 @@ def measure_final_state(circuit, state, is_state_batched, rng=None, prng_key=Non
     Returns:
         Tuple[TensorLike]: The measurement results
     """
-    if any(qml.math.isnan(state)):
-        # do something
-        state = 1
-
     circuit = circuit.map_to_standard_wires()
 
     if not circuit.shots:

From 240fbe2d906dc199578ab3a5047d66dd5b8aaf12 Mon Sep 17 00:00:00 2001
From: Tom Bromley <49409390+trbromley@users.noreply.github.com>
Date: Thu, 21 Sep 2023 13:26:08 -0400
Subject: [PATCH 104/131] Update returns.rst (#4617)

Update to the returns page after the removal of the old system.

Co-authored-by: Christina Lee <christina@xanadu.ai>
---
 doc/introduction/returns.rst | 8 +++++---
 1 file changed, 5 insertions(+), 3 deletions(-)

diff --git a/doc/introduction/returns.rst b/doc/introduction/returns.rst
index 64814afad84..97d7dd4b27e 100644
--- a/doc/introduction/returns.rst
+++ b/doc/introduction/returns.rst
@@ -18,6 +18,11 @@ provides additional details about the update and can be used to troubleshoot iss
       `bug report <https://github.com/PennyLaneAI/pennylane/issues/new?labels=bug+%3Abug%3A&template=bug_report.yml&title=[BUG]>`_
       on the PennyLane GitHub page.
 
+.. note::
+
+    The old return system has been completely removed from PennyLane as of version 0.33.0, along with the
+    ``enable_return()`` and ``disable_return()`` toggle functions.
+
 Summary of the update
 ---------------------
 
@@ -61,9 +66,6 @@ and probabilities separately:
     (tensor(0.87758256, requires_grad=True),
      tensor([0.93879128, 0.06120872], requires_grad=True))
 
-For a detailed explanation of this change, along with code examples, check out the
-:func:`~.enable_return` function documentation.
-
 Motivation
 ----------
 

From 09baad03dc3709007e693d2e9c56e46dfb29fc2a Mon Sep 17 00:00:00 2001
From: Matthew Silverman <matthews@xanadu.ai>
Date: Thu, 21 Sep 2023 15:44:55 -0400
Subject: [PATCH 105/131] fix and enable pulse gradient with broadcasting on
 new device (#4620)

**Context:**
Enabling this was missed in the DQ2 switch-over. We originally said that
`ParametrizedEvolution` should fail to apply to broadcasted states, but
`pulse_generator` and `stoch_pulse_grad` have custom handling for
certain cases, so we ought not be so decisive in `apply_operation`. That
said, if the inputted state is batched but the operator is not, then
we're in trouble.

**Description of the Change:**
- Only raise the batched-state-ParametrizedEvolution error when the op
is not itself batched
- put shots on tapes created by `pulse_generator` (missed in gradient
uses tape shots upgrade)
- change `test_jax` at the bottom of each gradient test file to use
`"backprop"` instead of `None` for the `gradient_fn`. This worked as-is
on `DefaultQubitJax` because it has a [passthru_interface
set](https://github.com/PennyLaneAI/pennylane/blob/e909e56a197cbdea772449e972a12da4931cf2b8/pennylane/interfaces/execution.py#L581),
but really it's using backprop. Maybe that boolean needs more work, but
this seemed like a niche-enough test case, so I fixed the test instead
of the code.

**Benefits:**
pulse gradients work with the new device and `use_broadcasting=True`!
---
 doc/releases/changelog-dev.md                 |  4 ++-
 pennylane/devices/qubit/apply_operation.py    |  4 +--
 .../gradients/pulse_generator_gradient.py     |  3 +-
 tests/devices/qubit/test_apply_operation.py   |  2 +-
 .../core/test_pulse_generator_gradient.py     | 17 ++++-----
 tests/gradients/core/test_pulse_gradient.py   | 35 ++++++++++++-------
 6 files changed, 39 insertions(+), 26 deletions(-)

diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index ff0efbd7b38..d3d2b6873a8 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -68,6 +68,7 @@
   accessible by the short name `default.qubit.legacy`, or directly via `qml.devices.DefaultQubitLegacy`.
   [(#4594)](https://github.com/PennyLaneAI/pennylane/pull/4594)
   [(#4436)](https://github.com/PennyLaneAI/pennylane/pull/4436)
+  [(#4620)](https://github.com/PennyLaneAI/pennylane/pull/4620)
 
 <h3>Improvements 🛠</h3>
 
@@ -141,9 +142,10 @@
   `DefaultQubitJax` in the old API.
   [(#4596)](https://github.com/PennyLaneAI/pennylane/pull/4596)
 
-* DefaultQubit2 dispatches to a faster implementation for applying `ParameterizedEvolution` to a state
+* DefaultQubit2 dispatches to a faster implementation for applying `ParametrizedEvolution` to a state
   when it is more efficient to evolve the state than the operation matrix.
   [(#4598)](https://github.com/PennyLaneAI/pennylane/pull/4598)
+  [(#4620)](https://github.com/PennyLaneAI/pennylane/pull/4620)
 
 * `ShotAdaptiveOptimizer` has been updated to pass shots to QNode executions instead of overriding
   device shots before execution. This makes it compatible with the new device API.
diff --git a/pennylane/devices/qubit/apply_operation.py b/pennylane/devices/qubit/apply_operation.py
index 6800609e802..a0a4e0be7fb 100644
--- a/pennylane/devices/qubit/apply_operation.py
+++ b/pennylane/devices/qubit/apply_operation.py
@@ -282,9 +282,9 @@ def apply_parametrized_evolution(
 ):
     """Apply ParametrizedEvolution by evolving the state rather than the operator matrix
     if we are operating on more than half of the subsystem"""
-    if is_state_batched:
+    if is_state_batched and op.batch_size is None:
         raise RuntimeError(
-            "ParameterizedEvolution does not support batching, but received a batched state"
+            "ParameterizedEvolution does not support standard broadcasting, but received a batched state"
         )
 
     # shape(state) is static (not a tracer), we can use an if statement
diff --git a/pennylane/gradients/pulse_generator_gradient.py b/pennylane/gradients/pulse_generator_gradient.py
index d0691c18183..ca213b602c3 100644
--- a/pennylane/gradients/pulse_generator_gradient.py
+++ b/pennylane/gradients/pulse_generator_gradient.py
@@ -174,7 +174,8 @@ def _insert_op(tape, ops, op_idx):
     ops_pre = tape.operations[:op_idx]
     ops_post = tape.operations[op_idx:]
     return [
-        qml.tape.QuantumScript(ops_pre + [insert] + ops_post, tape.measurements) for insert in ops
+        qml.tape.QuantumScript(ops_pre + [insert] + ops_post, tape.measurements, shots=tape.shots)
+        for insert in ops
     ]
 
 
diff --git a/tests/devices/qubit/test_apply_operation.py b/tests/devices/qubit/test_apply_operation.py
index 9bc1755fc5d..c19e3fa90c8 100644
--- a/tests/devices/qubit/test_apply_operation.py
+++ b/tests/devices/qubit/test_apply_operation.py
@@ -495,7 +495,7 @@ def test_batched_state_raises_an_error(self):
             ]
         )
 
-        with pytest.raises(RuntimeError, match="does not support batching"):
+        with pytest.raises(RuntimeError, match="does not support standard broadcasting"):
             _ = apply_operation(op, initial_state, is_state_batched=True)
 
 
diff --git a/tests/gradients/core/test_pulse_generator_gradient.py b/tests/gradients/core/test_pulse_generator_gradient.py
index 3396fe02210..6b36356a9d0 100644
--- a/tests/gradients/core/test_pulse_generator_gradient.py
+++ b/tests/gradients/core/test_pulse_generator_gradient.py
@@ -965,9 +965,8 @@ def test_all_zero_diff_methods_multiple_returns_tape(self):
         assert np.allclose(res_pulse_gen[1][2], 0)
 
 
-# TODO: add default.qubit once it supports PRNG key
 @pytest.mark.jax
-@pytest.mark.parametrize("dev_name", ["default.qubit.jax"])
+@pytest.mark.parametrize("dev_name", ["default.qubit", "default.qubit.jax"])
 class TestPulseGeneratorTape:
     """Test that differentiating tapes with ``pulse_generator`` works."""
 
@@ -979,7 +978,8 @@ def test_single_pulse_single_term(self, dev_name, shots, tol):
         import jax.numpy as jnp
 
         prng_key = jax.random.PRNGKey(8251)
-        dev = qml.device(dev_name, wires=1, shots=shots, prng_key=prng_key)
+        key = "prng_key" if dev_name == "default.qubit.jax" else "seed"
+        dev = qml.device(dev_name, wires=1, shots=shots, **{key: prng_key})
 
         H = jnp.polyval * X(0)
         x = jnp.array([0.4, 0.2, 0.1])
@@ -1015,7 +1015,8 @@ def test_single_pulse_multi_term(self, dev_name, shots, tol):
 
         prng_key = jax.random.PRNGKey(8251)
         dev = qml.device(dev_name, wires=1, shots=None)
-        dev_shots = qml.device(dev_name, wires=1, shots=shots, prng_key=prng_key)
+        key = "prng_key" if dev_name == "default.qubit.jax" else "seed"
+        dev_shots = qml.device(dev_name, wires=1, shots=shots, **{key: prng_key})
 
         H = 0.1 * Z(0) + jnp.polyval * X(0) + qml.pulse.constant * Y(0)
         x = jnp.array([0.4, 0.2, 0.1])
@@ -1095,7 +1096,8 @@ def test_multi_pulse(self, dev_name, shots, tol):
 
         prng_key = jax.random.PRNGKey(8251)
         dev = qml.device(dev_name, wires=1, shots=None)
-        dev_shots = qml.device(dev_name, wires=1, shots=shots, prng_key=prng_key)
+        key = "prng_key" if dev_name == "default.qubit.jax" else "seed"
+        dev_shots = qml.device(dev_name, wires=1, shots=shots, **{key: prng_key})
 
         H0 = 0.1 * Z(0) + jnp.polyval * X(0)
         H1 = 0.2 * Y(0) + qml.pulse.constant * Y(0) + jnp.polyval * Z(0)
@@ -1447,9 +1449,8 @@ def circuit(param1, param2):
             assert tracker.totals["executions"] == 1 + 2  # one forward pass, two shifted tapes
 
 
-# TODO: port ParametrizedEvolution to new default.qubit
 @pytest.mark.jax
-@pytest.mark.parametrize("dev_name", ["default.qubit.jax"])
+@pytest.mark.parametrize("dev_name", ["default.qubit", "default.qubit.jax"])
 class TestPulseGeneratorDiff:
     """Test that pulse_generator is differentiable, i.e. that computing
     the derivative with pulse_generator is differentiable a second time,
@@ -1473,7 +1474,7 @@ def fun(params):
             tape = qml.tape.QuantumScript([op], [qml.expval(Z(0))])
             tape.trainable_params = [0]
             _tapes, fn = pulse_generator(tape)
-            return fn(qml.execute(_tapes, dev, None))
+            return fn(qml.execute(_tapes, dev, "backprop"))
 
         params = [jnp.array(0.4)]
         p = params[0] * T
diff --git a/tests/gradients/core/test_pulse_gradient.py b/tests/gradients/core/test_pulse_gradient.py
index 38ea9f3af65..d02716a8163 100644
--- a/tests/gradients/core/test_pulse_gradient.py
+++ b/tests/gradients/core/test_pulse_gradient.py
@@ -1352,9 +1352,8 @@ def circuit(params):
         assert tracker.totals["executions"] == 4  # two shifted tapes, two splitting times
 
 
-# TODO: add default.qubit once it supports PRNG key
 @pytest.mark.jax
-@pytest.mark.parametrize("dev_name", ["default.qubit.jax"])
+@pytest.mark.parametrize("dev_name", ["default.qubit", "default.qubit.jax"])
 class TestStochPulseGradIntegration:
     """Test that stoch_pulse_grad integrates correctly with QNodes and ML interfaces."""
 
@@ -1367,7 +1366,8 @@ def test_simple_qnode_expval(self, dev_name, num_split_times, shots, tol):
         import jax.numpy as jnp
 
         jax.config.update("jax_enable_x64", True)
-        dev = qml.device(dev_name, wires=1, shots=shots, prng_key=jax.random.PRNGKey(74))
+        key = "prng_key" if dev_name == "default.qubit.jax" else "seed"
+        dev = qml.device(dev_name, wires=1, shots=shots, **{key: jax.random.PRNGKey(74)})
         T = 0.2
         ham_single_q_const = qml.pulse.constant * qml.PauliY(0)
 
@@ -1394,7 +1394,8 @@ def test_simple_qnode_expval_two_evolves(self, dev_name, num_split_times, shots,
         import jax.numpy as jnp
 
         jax.config.update("jax_enable_x64", True)
-        dev = qml.device(dev_name, wires=1, shots=shots, prng_key=jax.random.PRNGKey(74))
+        key = "prng_key" if dev_name == "default.qubit.jax" else "seed"
+        dev = qml.device(dev_name, wires=1, shots=shots, **{key: jax.random.PRNGKey(74)})
         T_x = 0.1
         T_y = 0.2
         ham_x = qml.pulse.constant * qml.PauliX(0)
@@ -1425,7 +1426,8 @@ def test_simple_qnode_probs(self, dev_name, num_split_times, shots, tol):
         import jax.numpy as jnp
 
         jax.config.update("jax_enable_x64", True)
-        dev = qml.device(dev_name, wires=1, shots=shots, prng_key=jax.random.PRNGKey(74))
+        key = "prng_key" if dev_name == "default.qubit.jax" else "seed"
+        dev = qml.device(dev_name, wires=1, shots=shots, **{key: jax.random.PRNGKey(74)})
         T = 0.2
         ham_single_q_const = qml.pulse.constant * qml.PauliY(0)
 
@@ -1452,7 +1454,8 @@ def test_simple_qnode_probs_expval(self, dev_name, num_split_times, shots, tol):
         import jax.numpy as jnp
 
         jax.config.update("jax_enable_x64", True)
-        dev = qml.device(dev_name, wires=1, shots=shots, prng_key=jax.random.PRNGKey(74))
+        key = "prng_key" if dev_name == "default.qubit.jax" else "seed"
+        dev = qml.device(dev_name, wires=1, shots=shots, **{key: jax.random.PRNGKey(74)})
         T = 0.2
         ham_single_q_const = qml.pulse.constant * qml.PauliY(0)
 
@@ -1535,8 +1538,13 @@ def ansatz(params):
         )
         qnode_backprop = qml.QNode(ansatz, dev, interface="jax")
 
-        grad_pulse_grad = jax.grad(qnode_pulse_grad)(params)
-        assert dev.num_executions == 1 + 2 * 3 * num_split_times
+        with qml.Tracker(dev) as tracker:
+            grad_pulse_grad = jax.grad(qnode_pulse_grad)(params)
+        assert (
+            tracker.totals["executions"] == (1 + 2 * 3 * num_split_times)
+            if dev_name == "default.qubit.jax"
+            else 1
+        )
         grad_backprop = jax.grad(qnode_backprop)(params)
 
         assert all(
@@ -1552,7 +1560,8 @@ def test_multi_return_broadcasting_multi_shots_raises(self, dev_name):
 
         jax.config.update("jax_enable_x64", True)
         shots = [100, 100]
-        dev = qml.device(dev_name, wires=1, shots=shots, prng_key=jax.random.PRNGKey(74))
+        key = "prng_key" if dev_name == "default.qubit.jax" else "seed"
+        dev = qml.device(dev_name, wires=1, shots=shots, **{key: jax.random.PRNGKey(74)})
         T = 0.2
         ham_single_q_const = qml.pulse.constant * qml.PauliY(0)
 
@@ -1582,7 +1591,8 @@ def test_qnode_probs_expval_broadcasting(self, dev_name, num_split_times, shots,
         import jax.numpy as jnp
 
         jax.config.update("jax_enable_x64", True)
-        dev = qml.device(dev_name, wires=1, shots=shots, prng_key=jax.random.PRNGKey(74))
+        key = "prng_key" if dev_name == "default.qubit.jax" else "seed"
+        dev = qml.device(dev_name, wires=1, shots=shots, **{key: jax.random.PRNGKey(74)})
         T = 0.2
         ham_single_q_const = qml.pulse.constant * qml.PauliY(0)
 
@@ -1757,9 +1767,8 @@ def ansatz(params):
         jax.clear_caches()
 
 
-# TODO: port ParametrizedEvolution to new default.qubit
 @pytest.mark.jax
-@pytest.mark.parametrize("dev_name", ["default.qubit.jax"])
+@pytest.mark.parametrize("dev_name", ["default.qubit", "default.qubit.jax"])
 class TestStochPulseGradDiff:
     """Test that stoch_pulse_grad is differentiable."""
 
@@ -1780,7 +1789,7 @@ def fun(params):
             tape = qml.tape.QuantumScript([op], [qml.expval(qml.PauliZ(0))])
             tape.trainable_params = [0]
             tapes, fn = stoch_pulse_grad(tape)
-            return fn(qml.execute(tapes, dev, None))
+            return fn(qml.execute(tapes, dev, "backprop"))
 
         params = [jnp.array(0.4)]
         p = params[0] * T

From b0a2b163567288ed068202dcf03dae24dcd2e4f2 Mon Sep 17 00:00:00 2001
From: Christina Lee <christina@xanadu.ai>
Date: Thu, 21 Sep 2023 17:14:21 -0400
Subject: [PATCH 106/131] Register `QuantumScript` and `QuantumTape` as pytrees
 (#4608)

`QuantumScript` and `QuantumTape` now have `_flatten` and `_unflatten`
methods and are registered as jax pytrees.

[sc-40589]
---
 doc/releases/changelog-dev.md                 |  3 +-
 pennylane/measurements/__init__.py            |  4 +-
 pennylane/tape/qscript.py                     | 13 +++++
 pennylane/tape/tape.py                        |  4 ++
 .../test_jax_default_qubit_2.py               | 15 ++++++
 tests/tape/test_qscript.py                    | 49 +++++++++++++++++++
 6 files changed, 85 insertions(+), 3 deletions(-)

diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index d3d2b6873a8..64d1a19fac0 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -77,8 +77,9 @@
   [#4523](https://github.com/PennyLaneAI/pennylane/pull/4523)
   [#4524](https://github.com/PennyLaneAI/pennylane/pull/4524)
 
-* `MeasurementProcess` objects are now registered as jax pytrees.
+* `MeasurementProcess` and `QuantumScript` objects are now registered as jax pytrees.
   [(#4607)](https://github.com/PennyLaneAI/pennylane/pull/4607)
+  [(#4608)](https://github.com/PennyLaneAI/pennylane/pull/4608)
 
 * Tensor-network template `qml.MPS` now supports changing `offset` between subsequent blocks for more flexibility.
   [(#4531)](https://github.com/PennyLaneAI/pennylane/pull/4531)
diff --git a/pennylane/measurements/__init__.py b/pennylane/measurements/__init__.py
index 3728a0bf9a0..a24611f3609 100644
--- a/pennylane/measurements/__init__.py
+++ b/pennylane/measurements/__init__.py
@@ -209,8 +209,8 @@ def circuit(x):
     :href: serialization
 
     PennyLane measurements are automatically registered as `Pytrees <https://jax.readthedocs.io/en/latest/pytrees.html>`_ .
-
-    The :class:`~.MeasurementProcess` definitions are sufficient for all PL measurements.
+    ``MeasurementProcess._flatten`` and ``MeasurementProcess._unflatten`` need to be overwritten if the measurement has additional
+    metadata, such as ``seed`` or ``all_outcomes``.
 
     >>> H = 2.0 * qml.PauliX(0)
     >>> mp = qml.expval(H)
diff --git a/pennylane/tape/qscript.py b/pennylane/tape/qscript.py
index 442636069f9..6e0ba7fd0a8 100644
--- a/pennylane/tape/qscript.py
+++ b/pennylane/tape/qscript.py
@@ -36,6 +36,7 @@
 )
 from pennylane.typing import TensorLike
 from pennylane.operation import Observable, Operator, Operation
+from pennylane.pytrees import register_pytree
 from pennylane.queuing import AnnotatedQueue, process_queue
 from pennylane.wires import Wires
 
@@ -173,6 +174,15 @@ class QuantumScript:
 
     """
 
+    def _flatten(self):
+        return (self._ops, self.measurements), (self.shots, tuple(self.trainable_params))
+
+    @classmethod
+    def _unflatten(cls, data, metadata):
+        new_tape = cls(*data, shots=metadata[0])
+        new_tape.trainable_params = metadata[1]
+        return new_tape
+
     def __init__(
         self,
         ops=None,
@@ -1256,3 +1266,6 @@ def wrapper(*args, **kwargs):
         return qscript
 
     return wrapper
+
+
+register_pytree(QuantumScript, QuantumScript._flatten, QuantumScript._unflatten)
diff --git a/pennylane/tape/tape.py b/pennylane/tape/tape.py
index d8e9cd4fbd8..d44da1b40ac 100644
--- a/pennylane/tape/tape.py
+++ b/pennylane/tape/tape.py
@@ -22,6 +22,7 @@
 from pennylane.measurements import CountsMP, ProbabilityMP, SampleMP
 from pennylane.operation import DecompositionUndefinedError, Operator, StatePrepBase
 from pennylane.queuing import AnnotatedQueue, QueuingManager, process_queue
+from pennylane.pytrees import register_pytree
 
 from .qscript import QuantumScript
 
@@ -483,3 +484,6 @@ def __setitem__(self, key, val):
 
     def __hash__(self):
         return QuantumScript.__hash__(self)
+
+
+register_pytree(QuantumTape, QuantumTape._flatten, QuantumTape._unflatten)
diff --git a/tests/interfaces/default_qubit_2_integration/test_jax_default_qubit_2.py b/tests/interfaces/default_qubit_2_integration/test_jax_default_qubit_2.py
index 5d1ef9cf9a7..249794ee193 100644
--- a/tests/interfaces/default_qubit_2_integration/test_jax_default_qubit_2.py
+++ b/tests/interfaces/default_qubit_2_integration/test_jax_default_qubit_2.py
@@ -29,6 +29,21 @@
 pytestmark = pytest.mark.jax
 
 
+def test_jit_execution():
+    """Test that qml.execute can be directly jitted."""
+    dev = qml.device("default.qubit")
+
+    tape = qml.tape.QuantumScript(
+        [qml.RX(jax.numpy.array(0.1), 0)], [qml.expval(qml.s_prod(2.0, qml.PauliZ(0)))]
+    )
+
+    out = jax.jit(qml.execute, static_argnames=("device", "gradient_fn"))(
+        (tape,), device=dev, gradient_fn=qml.gradients.param_shift
+    )
+    expected = 2.0 * jax.numpy.cos(jax.numpy.array(0.1))
+    assert qml.math.allclose(out[0], expected)
+
+
 # pylint: disable=too-few-public-methods
 class TestCaching:
     """Tests for caching behaviour"""
diff --git a/tests/tape/test_qscript.py b/tests/tape/test_qscript.py
index 9e5dbca64bc..3426d18741e 100644
--- a/tests/tape/test_qscript.py
+++ b/tests/tape/test_qscript.py
@@ -1430,3 +1430,52 @@ def test_sample_real_and_int_eigvals(self):
         assert np.issubdtype(result[0].dtype, float)
         assert np.issubdtype(result[1].dtype, np.int64)
         assert qs.numeric_type == (float, int)
+
+
+@pytest.mark.parametrize("qscript_type", (QuantumScript, qml.tape.QuantumTape))
+def test_flatten_unflatten(qscript_type):
+    """Test the flatten and unflatten methods."""
+    ops = [qml.RX(0.1, wires=0), qml.U3(0.2, 0.3, 0.4, wires=0)]
+    mps = [qml.expval(qml.PauliZ(0)), qml.state()]
+
+    tape = qscript_type(ops, mps, shots=100)
+    tape.trainable_params = {0}
+
+    data, metadata = tape._flatten()
+    assert all(o1 is o2 for o1, o2 in zip(ops, data[0]))
+    assert all(o1 is o2 for o1, o2 in zip(mps, data[1]))
+    assert metadata[0] == qml.measurements.Shots(100)
+    assert metadata[1] == (0,)
+    assert hash(metadata)
+
+    new_tape = qscript_type._unflatten(data, metadata)
+    assert all(o1 is o2 for o1, o2 in zip(new_tape.operations, tape.operations))
+    assert all(o1 is o2 for o1, o2 in zip(new_tape.measurements, tape.measurements))
+    assert new_tape.shots == qml.measurements.Shots(100)
+    assert new_tape.trainable_params == [0]
+
+
+@pytest.mark.jax
+@pytest.mark.parametrize("qscript_type", (QuantumScript, qml.tape.QuantumTape))
+def test_jax_pytree_integration(qscript_type):
+    """Test that QuantumScripts are integrated with jax pytress."""
+
+    eye_mat = np.eye(4)
+    ops = [qml.adjoint(qml.RY(0.5, wires=0)), qml.Rot(1.2, 2.3, 3.4, wires=0)]
+    mps = [
+        qml.var(qml.s_prod(2.0, qml.PauliX(0))),
+        qml.expval(qml.Hermitian(eye_mat, wires=(0, 1))),
+    ]
+
+    tape = qscript_type(ops, mps, shots=100)
+    tape.trainable_params = [2]
+
+    import jax
+
+    data, _ = jax.tree_util.tree_flatten(tape)
+    assert data[0] == 0.5
+    assert data[1] == 1.2
+    assert data[2] == 2.3
+    assert data[3] == 3.4
+    assert data[4] == 2.0
+    assert qml.math.allclose(data[5], eye_mat)

From c6f48e757e776217bcf7a16d7e6ba257f57479ef Mon Sep 17 00:00:00 2001
From: Mudit Pandey <mudit.pandey@xanadu.ai>
Date: Mon, 2 Oct 2023 16:19:58 -0400
Subject: [PATCH 107/131] Added nan measurement

---
 pennylane/devices/qubit/measure.py  | 13 ++------
 pennylane/devices/qubit/simulate.py | 51 +++++++++++++++++++++++++++--
 pennylane/math/multi_dispatch.py    | 19 -----------
 pennylane/math/single_dispatch.py   | 10 ++++++
 4 files changed, 61 insertions(+), 32 deletions(-)

diff --git a/pennylane/devices/qubit/measure.py b/pennylane/devices/qubit/measure.py
index fb047f85d1c..9b72def187d 100644
--- a/pennylane/devices/qubit/measure.py
+++ b/pennylane/devices/qubit/measure.py
@@ -19,7 +19,6 @@
 from scipy.sparse import csr_matrix
 
 from pennylane import math
-from pennylane.numpy import NaN
 from pennylane.ops import Sum, Hamiltonian
 from pennylane.measurements import (
     StateMeasurement,
@@ -63,14 +62,12 @@ def state_diagonalizing_gates(
     Returns:
         TensorLike: the result of the measurement
     """
-    if not any(math.isnan(state)):
-        for op in measurementprocess.diagonalizing_gates():
-            state = apply_operation(op, state, is_state_batched=is_state_batched)
+    for op in measurementprocess.diagonalizing_gates():
+        state = apply_operation(op, state, is_state_batched=is_state_batched)
 
     total_indices = len(state.shape) - is_state_batched
     wires = Wires(range(total_indices))
     flattened_state = flatten_state(state, total_indices)
-
     return measurementprocess.process_state(flattened_state, wires)
 
 
@@ -91,12 +88,6 @@ def csr_dot_products(
     total_wires = len(state.shape) - is_state_batched
     Hmat = measurementprocess.obs.sparse_matrix(wire_order=list(range(total_wires)))
 
-    if any(math.isnan(state)):
-        if is_state_batched:
-            return math.asarray([NaN] * math.shape(state)[0], like=math.get_interface(state))
-
-        return math.asarry(NaN, like=math.get_interface(state))
-
     if is_state_batched:
         state = math.toarray(state).reshape(math.shape(state)[0], -1)
 
diff --git a/pennylane/devices/qubit/simulate.py b/pennylane/devices/qubit/simulate.py
index 60dfd45f1b1..30b717b82ba 100644
--- a/pennylane/devices/qubit/simulate.py
+++ b/pennylane/devices/qubit/simulate.py
@@ -127,11 +127,12 @@ def get_final_state(circuit, debugger=None, interface=None):
 
             norm = qml.math.norm(state)
             if qml.math.isclose(norm, 0.0):
-                state = qml.math.asarray(
+                new_state = qml.math.asarray(
                     [qml.numpy.NaN] * qml.math.size(state),
                     like=qml.math.get_interface(state),
                 )
-                return state, is_state_batched
+                new_state = qml.math.reshape(new_state, qml.math.shape(state))
+                return new_state, is_state_batched
 
             state = state / norm
 
@@ -179,6 +180,9 @@ def measure_final_state(circuit, state, is_state_batched, rng=None, prng_key=Non
     """
     circuit = circuit.map_to_standard_wires()
 
+    if any(nan_value is True for nan_value in qml.math.flatten(qml.math.isnan(state))):
+        return _measure_nan_state(circuit, state, is_state_batched)
+
     if not circuit.shots:
         # analytic case
 
@@ -210,6 +214,49 @@ def measure_final_state(circuit, state, is_state_batched, rng=None, prng_key=Non
     return results
 
 
+def _measure_nan_state(circuit, state, is_state_batched):
+    """Helper function for creating NaN results with the expected shape."""
+    batch_size = qml.math.shape(state)[0] if is_state_batched else None
+    interface = qml.math.get_interface(state)
+
+    if circuit.shots.has_partitioned_shots:
+        res = (
+            _get_single_nan_res(circuit.measurements, batch_size, interface, s)
+            for s in circuit.shots
+        )
+    else:
+        res = _get_single_nan_res(
+            circuit.measurements, batch_size, interface, circuit.shots.total_shots
+        )
+
+    return res
+
+
+def _get_single_nan_res(measurements, batch_size, interface, shots):
+    """Helper to get NaN results for one item in a shot vector."""
+
+    res = []
+
+    for m in measurements:
+        shape = m.shape(None, qml.measurements.Shots(shots))
+        if batch_size is not None:
+            shape = (batch_size,) + shape
+
+        if shape == ():
+            out = qml.math.asarray(qml.numpy.NaN, like=interface)
+        else:
+            out = qml.math.full(shape, qml.numpy.NaN, like=interface)
+
+        res.append(out)
+
+    res = tuple(res)
+
+    if len(res) == 1:
+        res = res[0]
+
+    return res
+
+
 def simulate(
     circuit: qml.tape.QuantumScript, rng=None, prng_key=None, debugger=None, interface=None
 ) -> Result:
diff --git a/pennylane/math/multi_dispatch.py b/pennylane/math/multi_dispatch.py
index 667333e1aca..b676292a1a7 100644
--- a/pennylane/math/multi_dispatch.py
+++ b/pennylane/math/multi_dispatch.py
@@ -852,25 +852,6 @@ def norm(tensor, like=None, **kwargs):
     return norm(tensor, **kwargs)
 
 
-@multi_dispatch()
-def isnan(tensor, like=None, **kwargs):
-    """Check a tensor element-wise for ``NaN`` values and return the result
-    as a boolean tensor with the same interface."""
-    if like == "jax":
-        from jax.numpy import isnan
-
-    elif like == "tensorflow":
-        from tensorflow.math import is_nan as isnan
-
-    elif like == "torch":
-        from torch import isnan
-
-    else:
-        from numpy import isnan
-
-    return isnan(tensor, **kwargs)
-
-
 @multi_dispatch(argnum=[1])
 def gammainc(m, t, like=None):
     r"""Return the lower incomplete Gamma function.
diff --git a/pennylane/math/single_dispatch.py b/pennylane/math/single_dispatch.py
index 2bf96bef5be..f36952c9686 100644
--- a/pennylane/math/single_dispatch.py
+++ b/pennylane/math/single_dispatch.py
@@ -258,6 +258,16 @@ def _take_autograd(tensor, indices, axis=None):
 
 ar.register_function("tensorflow", "flatten", lambda x: _i("tf").reshape(x, [-1]))
 ar.register_function("tensorflow", "shape", lambda x: tuple(x.shape))
+ar.register_function(
+    "tensorflow",
+    "full",
+    lambda shape, fill_value, **kwargs: _i("tf").fill(
+        shape if isinstance(shape, (tuple, list)) else (shape), fill_value, **kwargs
+    ),
+)
+ar.register_function(
+    "tensorflow", "isnan", lambda tensor, **kwargs: _i("tf").math.is_nan(tensor, **kwargs)
+)
 ar.register_function(
     "tensorflow",
     "sqrt",

From b5ade012d6790fee1563da6f31ff2fbfa0231335 Mon Sep 17 00:00:00 2001
From: Mudit Pandey <mudit.pandey@xanadu.ai>
Date: Tue, 3 Oct 2023 12:39:28 -0400
Subject: [PATCH 108/131] Fixed simulate; added tests

---
 pennylane/devices/qubit/simulate.py         |  30 ++--
 pennylane/math/__init__.py                  |   1 -
 pennylane/math/single_dispatch.py           |   9 +-
 pennylane/measurements/__init__.py          |   2 +-
 pennylane/measurements/shots.py             |  47 +++++
 pennylane/transforms/defer_measurements.py  |   2 +-
 tests/devices/qubit/test_simulate.py        | 186 ++++++++++++++++----
 tests/transforms/test_defer_measurements.py |  27 ---
 8 files changed, 224 insertions(+), 80 deletions(-)

diff --git a/pennylane/devices/qubit/simulate.py b/pennylane/devices/qubit/simulate.py
index 30b717b82ba..8cf3ff9e3cc 100644
--- a/pennylane/devices/qubit/simulate.py
+++ b/pennylane/devices/qubit/simulate.py
@@ -126,12 +126,13 @@ def get_final_state(circuit, debugger=None, interface=None):
                 )
 
             norm = qml.math.norm(state)
-            if qml.math.isclose(norm, 0.0):
-                new_state = qml.math.asarray(
-                    [qml.numpy.NaN] * qml.math.size(state),
-                    like=qml.math.get_interface(state),
+            if qml.math.isclose(float(norm), 0.0):
+                new_state = qml.math.cast_like(
+                    qml.math.full(
+                        qml.math.shape(state), qml.numpy.NaN, like=qml.math.get_interface(state)
+                    ),
+                    state,
                 )
-                new_state = qml.math.reshape(new_state, qml.math.shape(state))
                 return new_state, is_state_batched
 
             state = state / norm
@@ -141,10 +142,10 @@ def get_final_state(circuit, debugger=None, interface=None):
             if circuit.shots:
                 # Clip the number of shots using a binomial distribution using the probability of
                 # measuring the postselected state.
-                postselected_shots = binomial(circuit.shots.total_shots, norm)
-                if postselected_shots == 0:
-                    raise RuntimeError("None of the samples meet the postselection criteria")
-                circuit._shots = qml.measurements.Shots(postselected_shots)
+                postselected_shots = [binomial(s, float(norm)) for s in circuit.shots]
+                # _FlexShots is used here since the binomial distribution could result in zero
+                # valid samples
+                circuit._shots = qml.measurements._FlexShots(postselected_shots)
 
     if set(circuit.op_wires) < set(circuit.wires):
         state = expand_state_over_wires(
@@ -180,7 +181,7 @@ def measure_final_state(circuit, state, is_state_batched, rng=None, prng_key=Non
     """
     circuit = circuit.map_to_standard_wires()
 
-    if any(nan_value is True for nan_value in qml.math.flatten(qml.math.isnan(state))):
+    if any(nan_value for nan_value in qml.math.flatten(qml.math.isnan(state))):
         return _measure_nan_state(circuit, state, is_state_batched)
 
     if not circuit.shots:
@@ -220,7 +221,7 @@ def _measure_nan_state(circuit, state, is_state_batched):
     interface = qml.math.get_interface(state)
 
     if circuit.shots.has_partitioned_shots:
-        res = (
+        res = tuple(
             _get_single_nan_res(circuit.measurements, batch_size, interface, s)
             for s in circuit.shots
         )
@@ -238,6 +239,13 @@ def _get_single_nan_res(measurements, batch_size, interface, shots):
     res = []
 
     for m in measurements:
+        if isinstance(m, qml.measurements.SampleMP):
+            res.append(qml.math.asarray([], like=interface))
+            continue
+        if isinstance(m, qml.measurements.CountsMP):
+            res.append({})
+            continue
+
         shape = m.shape(None, qml.measurements.Shots(shots))
         if batch_size is not None:
             shape = (batch_size,) + shape
diff --git a/pennylane/math/__init__.py b/pennylane/math/__init__.py
index e57713e8d84..9e8ec32837d 100644
--- a/pennylane/math/__init__.py
+++ b/pennylane/math/__init__.py
@@ -50,7 +50,6 @@
     gammainc,
     get_trainable_indices,
     iscomplex,
-    isnan,
     jax_argnums_to_tape_trainable,
     kron,
     matmul,
diff --git a/pennylane/math/single_dispatch.py b/pennylane/math/single_dispatch.py
index f36952c9686..1653666bd7d 100644
--- a/pennylane/math/single_dispatch.py
+++ b/pennylane/math/single_dispatch.py
@@ -266,7 +266,14 @@ def _take_autograd(tensor, indices, axis=None):
     ),
 )
 ar.register_function(
-    "tensorflow", "isnan", lambda tensor, **kwargs: _i("tf").math.is_nan(tensor, **kwargs)
+    "tensorflow",
+    "isnan",
+    lambda tensor, **kwargs: _i("tf").math.is_nan(
+        _i("tf").cast(tensor, "float64")
+        if tensor.dtype.name in ("int64", "int32", "complex128")
+        else tensor,
+        **kwargs,
+    ),
 )
 ar.register_function(
     "tensorflow",
diff --git a/pennylane/measurements/__init__.py b/pennylane/measurements/__init__.py
index a24611f3609..23697dca533 100644
--- a/pennylane/measurements/__init__.py
+++ b/pennylane/measurements/__init__.py
@@ -282,7 +282,7 @@ def circuit(x):
 from .purity import PurityMP, purity
 from .probs import ProbabilityMP, probs
 from .sample import SampleMP, sample
-from .shots import Shots, ShotCopies
+from .shots import Shots, ShotCopies, _FlexShots
 from .state import StateMP, DensityMatrixMP, density_matrix, state
 from .var import VarianceMP, var
 from .vn_entropy import VnEntropyMP, vn_entropy
diff --git a/pennylane/measurements/shots.py b/pennylane/measurements/shots.py
index 2d4a4bafe26..bbe8f9e2189 100644
--- a/pennylane/measurements/shots.py
+++ b/pennylane/measurements/shots.py
@@ -239,3 +239,50 @@ def has_partitioned_shots(self):
     def num_copies(self):
         """The total number of copies of any shot quantity."""
         return sum(s.copies for s in self.shot_vector)
+
+
+def valid_flex_int(s):
+    """Returns True if s is a non-negative integer."""
+    return isinstance(s, int) and s >= 0
+
+
+def valid_flex_tuple(s):
+    """Returns True if s is a tuple of the form (shots, copies)."""
+    return isinstance(s, tuple) and len(s) == 2 and valid_flex_int(s[0]) and valid_int(s[1])
+
+
+class _FlexShots(Shots):
+    """Shots class that allows zero shots."""
+
+    total_shots: int = None
+    """The total number of shots to be executed."""
+
+    shot_vector: Tuple[ShotCopies] = None
+    """The tuple of :class:`~ShotCopies` to be executed. Each element is of the form ``(shots, copies)``."""
+
+    _SHOT_ERROR = ValueError(
+        "Shots must be a single positive integer, a tuple pair of the form (shots, copies), or a sequence of these types."
+    )
+
+    _frozen = False
+
+    # pylint: disable=super-init-not-called
+    def __init__(self, shots=None):
+        if shots is None:
+            self.total_shots = None
+            self.shot_vector = ()
+        elif isinstance(shots, int):
+            if shots < 0:
+                raise self._SHOT_ERROR
+            self.total_shots = shots
+            self.shot_vector = (ShotCopies(shots, 1),)
+        elif isinstance(shots, Sequence):
+            if not all(valid_flex_int(s) or valid_flex_tuple(s) for s in shots):
+                raise self._SHOT_ERROR
+            self.__all_tuple_init__([s if isinstance(s, tuple) else (s, 1) for s in shots])
+        elif isinstance(shots, self.__class__):
+            return  # self already _is_ shots as defined by __new__
+        else:
+            raise self._SHOT_ERROR
+
+        self._frozen = True
diff --git a/pennylane/transforms/defer_measurements.py b/pennylane/transforms/defer_measurements.py
index 8f6bd808034..ae5b9f8ce7e 100644
--- a/pennylane/transforms/defer_measurements.py
+++ b/pennylane/transforms/defer_measurements.py
@@ -174,7 +174,7 @@ def func(x, y):
 
             if op.postselect is not None:
                 with QueuingManager.stop_recording():
-                    new_operations.append(qml.Projector([op.postselect], wires=op.wires[0]))
+                    new_operations.append(qml.Projector([float(op.postselect)], wires=op.wires[0]))
 
             # Store measurement outcome in new wire if wire gets reused
             if op.wires[0] in reused_measurement_wires or op.wires[0] in measured_wires:
diff --git a/tests/devices/qubit/test_simulate.py b/tests/devices/qubit/test_simulate.py
index a9903e4e7ab..9655aa86176 100644
--- a/tests/devices/qubit/test_simulate.py
+++ b/tests/devices/qubit/test_simulate.py
@@ -48,39 +48,6 @@ def compute_matrix():
     assert qml.math.allclose(result, -1.0)
 
 
-def test_projector_probability_error():
-    """Test that an error is raised if a Projector is applied in the operation queue
-    and the resulting state is null."""
-    tape = qml.tape.QuantumScript([qml.PauliX(0), qml.Projector([0], wires=0)], [qml.state()])
-    with pytest.raises(ValueError, match="Requested postselection on value with zero probability"):
-        _ = simulate(tape)
-
-    # Test with broadcasting
-    tape = qml.tape.QuantumScript(
-        [qml.RX([0.0, 0.0], wires=0), qml.Projector([1], 0)], [qml.state()]
-    )
-    with pytest.raises(ValueError, match="Requested postselection on value with zero probability"):
-        _ = simulate(tape)
-
-
-def test_projector_norm():
-    """Test that the norm of the state is maintained after applying a projector"""
-    tape = qml.tape.QuantumScript(
-        [qml.PauliX(0), qml.RX(0.123, 1), qml.Projector([0], wires=1)], [qml.state()]
-    )
-    res = simulate(tape)
-    assert np.isclose(np.linalg.norm(res), 1.0)
-
-    # Test with broadcasting
-    tape = qml.tape.QuantumScript(
-        [qml.PauliX(0), qml.RX([0.123, 0.456], 1), qml.Projector([0], wires=1)], [qml.state()]
-    )
-    res = simulate(tape)
-    assert len(res) == 2
-    for r in res:
-        assert np.isclose(np.linalg.norm(r), 1.0)
-
-
 # pylint: disable=too-few-public-methods
 class TestStatePrepBase:
     """Tests integration with various state prep methods."""
@@ -383,20 +350,163 @@ def test_broadcasting_with_extra_measurement_wires(self, mocker):
         assert np.allclose(res[2], -np.cos(x))
         assert spy.call_args_list[0].args == (qs, {2: 0, 1: 1, 0: 2})
 
-    def test_broadcasting_with_projector(self):
-        """Test that postselecting a broadcasted state works correctly"""
+
+class TestPostselection:
+    """Tests for applying projectors as operations."""
+
+    def test_projector_norm(self):
+        """Test that the norm of the state is maintained after applying a projector"""
+        tape = qml.tape.QuantumScript(
+            [qml.PauliX(0), qml.RX(0.123, 1), qml.Projector([0], wires=1)], [qml.state()]
+        )
+        res = simulate(tape)
+        assert np.isclose(np.linalg.norm(res), 1.0)
+
+    @pytest.mark.parametrize("shots", [None, 10, [10, 10]])
+    def test_broadcasting_with_projector(self, shots):
+        """Test that postselecting a broadcasted state raises an error"""
         tape = qml.tape.QuantumScript(
             [
                 qml.RX([0.1, 0.2], 0),
                 qml.Projector([0], wires=0),
             ],
             [qml.state()],
+            shots=shots,
+        )
+
+        with pytest.raises(ValueError, match="Cannot postselect on circuits with broadcasting"):
+            _ = simulate(tape)
+
+    @pytest.mark.all_interfaces
+    @pytest.mark.parametrize("mp", [qml.expval, qml.var])
+    @pytest.mark.parametrize("interface", ["numpy", "autograd", "torch", "tensorflow", "jax"])
+    @pytest.mark.parametrize("shots", [None, 10, [10, 10]])
+    def test_nan_float_result(self, mp, interface, shots):
+        """Test that the result of circuits with 0 probability postselections is NaN with the
+        expected shape."""
+
+        tape = qml.tape.QuantumScript(
+            [
+                qml.RX(qml.math.asarray(np.pi, like=interface, dtype="float64"), 0),
+                qml.Projector([0], wires=0),
+            ],
+            [mp(qml.PauliZ(0))],
+            shots=shots,
         )
 
         res = simulate(tape)
-        assert len(res) == 2
-        for r in res:
-            assert np.allclose(r, [1, 0])
+
+        if not isinstance(shots, list):
+            assert qml.math.ndim(res) == 0
+            assert qml.math.isnan(res)
+            assert qml.math.get_interface(res) == interface
+
+        else:
+            assert isinstance(res, tuple)
+            assert len(res) == 2
+            for r in res:
+                assert qml.math.ndim(r) == 0
+                assert qml.math.isnan(r)
+                assert qml.math.get_interface(r) == interface
+
+    @pytest.mark.all_interfaces
+    @pytest.mark.parametrize(
+        "mp", [qml.probs(wires=0), qml.probs(op=qml.PauliZ(0)), qml.probs(wires=[0, 1])]
+    )
+    @pytest.mark.parametrize("interface", ["numpy", "autograd", "torch", "tensorflow", "jax"])
+    @pytest.mark.parametrize("shots", [None, 10, [10, 10]])
+    def test_nan_probs(self, mp, interface, shots):
+        """Test that the result of circuits with 0 probability postselections is NaN with the
+        expected shape."""
+
+        tape = qml.tape.QuantumScript(
+            [
+                qml.RX(qml.math.asarray(np.pi, like=interface, dtype="float64"), 0),
+                qml.CNOT([0, 1]),
+                qml.Projector([0], wires=0),
+            ],
+            [mp],
+            shots=shots,
+        )
+
+        res = simulate(tape)
+
+        if not isinstance(shots, list):
+            assert qml.math.shape(res) == (2 ** len(mp.wires),)
+            assert all(nan_value for nan_value in qml.math.isnan(res))
+            assert qml.math.get_interface(res) == interface
+
+        else:
+            assert isinstance(res, tuple)
+            assert len(res) == 2
+            for r in res:
+                assert qml.math.shape(r) == (2 ** len(mp.wires),)
+                assert all(nan_value for nan_value in qml.math.isnan(r))
+                assert qml.math.get_interface(r) == interface
+
+    @pytest.mark.all_interfaces
+    @pytest.mark.parametrize(
+        "mp", [qml.sample(wires=0), qml.sample(op=qml.PauliZ(0)), qml.sample(wires=[0, 1])]
+    )
+    @pytest.mark.parametrize("interface", ["numpy", "autograd", "torch", "tensorflow", "jax"])
+    @pytest.mark.parametrize("shots", [None, 10, [10, 10]])
+    def test_nan_samples(self, mp, interface, shots):
+        """Test that the result of circuits with 0 probability postselections is NaN with the
+        expected shape."""
+
+        tape = qml.tape.QuantumScript(
+            [
+                qml.RX(qml.math.asarray(np.pi, like=interface, dtype="float64"), 0),
+                qml.CNOT([0, 1]),
+                qml.Projector([0], wires=0),
+            ],
+            [mp],
+            shots=shots,
+        )
+
+        res = simulate(tape)
+
+        if not isinstance(shots, list):
+            assert qml.math.shape(res) == (0,)
+            assert qml.math.get_interface(res) == interface
+
+        else:
+            assert isinstance(res, tuple)
+            assert len(res) == 2
+            for r in res:
+                assert qml.math.shape(r) == (0,)
+                assert qml.math.get_interface(r) == interface
+
+    @pytest.mark.all_interfaces
+    @pytest.mark.parametrize(
+        "mp", [qml.counts(wires=0), qml.counts(op=qml.PauliZ(0)), qml.counts()]
+    )
+    @pytest.mark.parametrize("interface", ["numpy", "autograd", "torch", "tensorflow", "jax"])
+    @pytest.mark.parametrize("shots", [None, 10, [10, 10]])
+    def test_nan_counts(self, mp, interface, shots):
+        """Test that the result of circuits with 0 probability postselections is NaN with the
+        expected shape."""
+
+        tape = qml.tape.QuantumScript(
+            [
+                qml.RX(qml.math.asarray(np.pi, like=interface, dtype="float64"), 0),
+                qml.CNOT([0, 1]),
+                qml.Projector([0], wires=0),
+            ],
+            [mp],
+            shots=shots,
+        )
+
+        res = simulate(tape)
+
+        if not isinstance(shots, list):
+            assert res == {}
+
+        else:
+            assert isinstance(res, tuple)
+            assert len(res) == 2
+            for r in res:
+                assert r == {}
 
 
 class TestDebugger:
diff --git a/tests/transforms/test_defer_measurements.py b/tests/transforms/test_defer_measurements.py
index 3075bcb6603..8bc2de2bb83 100644
--- a/tests/transforms/test_defer_measurements.py
+++ b/tests/transforms/test_defer_measurements.py
@@ -248,33 +248,6 @@ def circ2():
         for op, expected_op in zip(circ1.qtape, expected_circuit):
             assert qml.equal(op, expected_op)
 
-    def test_postselection_error(self):
-        """Test that an error is raised if the circuit being transformed has batched dimensions
-        or partitioned shots."""
-        dev = DefaultQubit()
-
-        @qml.qnode(dev)
-        def circ1():
-            qml.RX([0.1, 0.2], wires=0)
-            _ = qml.measure(0, postselect=1)
-            return qml.probs()
-
-        # Test that no error is raised if broadcasting with analytic execution
-        _ = circ1()
-
-        with pytest.raises(ValueError, match="Cannot postselect on mid-circuit measurements"):
-            # Error with finite shots and broadcasting
-            _ = circ1(shots=10)
-
-        @qml.qnode(dev)
-        def circ2():
-            qml.measure(0, postselect=1)
-            return qml.probs()
-
-        with pytest.raises(ValueError, match="Cannot postselect on mid-circuit measurements"):
-            # Error with shot vector
-            _ = circ2(shots=[10, 10])
-
     @pytest.mark.parametrize("shots", [None, 1000, [1000, 1000]])
     def test_measurement_statistics_single_wire(self, shots):
         """Test that users can collect measurement statistics on

From 8579cceea2e841b57606c19f22bdc7307d1d50b2 Mon Sep 17 00:00:00 2001
From: Mudit Pandey <mudit.pandey@xanadu.ai>
Date: Tue, 3 Oct 2023 18:28:27 -0400
Subject: [PATCH 109/131] Fixed mp errors

---
 pennylane/devices/qubit/simulate.py  | 240 ++++++++++++++++++---------
 pennylane/math/single_dispatch.py    |  11 +-
 pennylane/measurements/__init__.py   |   2 +-
 pennylane/measurements/counts.py     |  13 ++
 pennylane/measurements/probs.py      |   3 +
 pennylane/measurements/shots.py      |  47 ------
 tests/devices/qubit/test_simulate.py | 131 ++++++++++-----
 7 files changed, 265 insertions(+), 182 deletions(-)

diff --git a/pennylane/devices/qubit/simulate.py b/pennylane/devices/qubit/simulate.py
index 8cf3ff9e3cc..605ee268d86 100644
--- a/pennylane/devices/qubit/simulate.py
+++ b/pennylane/devices/qubit/simulate.py
@@ -13,6 +13,7 @@
 # limitations under the License.
 """Simulate a quantum script."""
 # pylint: disable=protected-access
+from typing import Sequence
 from numpy.random import default_rng, binomial
 
 import pennylane as qml
@@ -46,6 +47,59 @@
 }
 
 
+def _valid_flex_int(s):
+    """Returns True if s is a non-negative integer."""
+    return isinstance(s, int) and s >= 0
+
+
+def _valid_flex_tuple(s):
+    """Returns True if s is a tuple of the form (shots, copies)."""
+    return (
+        isinstance(s, tuple)
+        and len(s) == 2
+        and _valid_flex_int(s[0])
+        and isinstance(s[1], int)
+        and s[1] > 0
+    )
+
+
+class _FlexShots(qml.measurements.Shots):
+    """Shots class that allows zero shots."""
+
+    total_shots = None
+    """The total number of shots to be executed."""
+
+    shot_vector = None
+    """The tuple of :class:`~ShotCopies` to be executed. Each element is of the form ``(shots, copies)``."""
+
+    _SHOT_ERROR = ValueError(
+        "Shots must be a single positive integer, a tuple pair of the form (shots, copies), or a sequence of these types."
+    )
+
+    _frozen = False
+
+    # pylint: disable=super-init-not-called
+    def __init__(self, shots=None):
+        if shots is None:
+            self.total_shots = None
+            self.shot_vector = ()
+        elif isinstance(shots, int):
+            if shots < 0:
+                raise self._SHOT_ERROR
+            self.total_shots = shots
+            self.shot_vector = (qml.measurements.ShotCopies(shots, 1),)
+        elif isinstance(shots, Sequence):
+            if not all(_valid_flex_int(s) or _valid_flex_tuple(s) for s in shots):
+                raise self._SHOT_ERROR
+            self.__all_tuple_init__([s if isinstance(s, tuple) else (s, 1) for s in shots])
+        elif isinstance(shots, self.__class__):
+            return  # self already _is_ shots as defined by __new__
+        else:
+            raise self._SHOT_ERROR
+
+        self._frozen = True
+
+
 def expand_state_over_wires(state, state_wires, all_wires, is_state_batched):
     """
     Expand and re-order a state given some initial and target wire orders, setting
@@ -83,6 +137,49 @@ def expand_state_over_wires(state, state_wires, all_wires, is_state_batched):
     return qml.math.transpose(state, desired_axes)
 
 
+def _postselection_postprocess(state, is_state_batched, shots):
+    """Update state after projector is applied."""
+    if is_state_batched:
+        raise ValueError(
+            "Cannot postselect on circuits with broadcasting. Use the "
+            "qml.transforms.broadcast_expand transform to split a broadcasted "
+            "tape into multiple non-broadcasted tapes before executing if "
+            "postselection is used."
+        )
+
+    if qml.math.is_abstract(state):
+        return state, shots, False
+
+    norm = qml.math.norm(state)
+    if not qml.math.is_abstract(norm) and qml.math.isclose(float(norm), 0.0):
+        new_state = qml.math.cast_like(
+            qml.math.full(qml.math.shape(state), qml.numpy.NaN, like=qml.math.get_interface(state)),
+            state,
+        )
+        is_nan = True
+
+    else:
+        new_state = state / norm
+        is_nan = False
+
+        # defer_measurements will raise an error with batched shots or broadcasting so we can
+        # assume that both the state and shots are unbatched.
+        if shots:
+            # Clip the number of shots using a binomial distribution using the probability of
+            # measuring the postselected state.
+            postselected_shots = (
+                [binomial(s, float(norm)) for s in shots]
+                if not qml.math.is_abstract(norm)
+                else shots
+            )
+
+            # _FlexShots is used here since the binomial distribution could result in zero
+            # valid samples
+            shots = _FlexShots(postselected_shots)
+
+    return new_state, shots, is_nan
+
+
 def get_final_state(circuit, debugger=None, interface=None):
     """
     Get the final state that results from executing the given quantum script.
@@ -112,41 +209,17 @@ def get_final_state(circuit, debugger=None, interface=None):
     for op in circuit.operations[bool(prep) :]:
         state = apply_operation(op, state, is_state_batched=is_state_batched, debugger=debugger)
 
+        # Handle postselection on mid-circuit measurements
+        if isinstance(op, qml.Projector):
+            state, circuit._shots, state_nan = _postselection_postprocess(
+                state, is_state_batched, circuit.shots
+            )
+            if state_nan:
+                return state, is_state_batched
+
         # 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
 
-        if isinstance(op, qml.Projector):
-            # Handle postselection on mid-circuit measurements
-            if is_state_batched:
-                raise ValueError(
-                    "Cannot postselect on circuits with broadcasting. Use the "
-                    "qml.transforms.broadcast_expand transform to split a broadcasted "
-                    "tape into multiple non-broadcasted tapes before executing if "
-                    "postselection is used."
-                )
-
-            norm = qml.math.norm(state)
-            if qml.math.isclose(float(norm), 0.0):
-                new_state = qml.math.cast_like(
-                    qml.math.full(
-                        qml.math.shape(state), qml.numpy.NaN, like=qml.math.get_interface(state)
-                    ),
-                    state,
-                )
-                return new_state, is_state_batched
-
-            state = state / norm
-
-            # defer_measurements will raise an error with batched shots or broadcasting so we can
-            # assume that both the state and shots are unbatched.
-            if circuit.shots:
-                # Clip the number of shots using a binomial distribution using the probability of
-                # measuring the postselected state.
-                postselected_shots = [binomial(s, float(norm)) for s in circuit.shots]
-                # _FlexShots is used here since the binomial distribution could result in zero
-                # valid samples
-                circuit._shots = qml.measurements._FlexShots(postselected_shots)
-
     if set(circuit.op_wires) < set(circuit.wires):
         state = expand_state_over_wires(
             state,
@@ -158,6 +231,56 @@ def get_final_state(circuit, debugger=None, interface=None):
     return state, is_state_batched
 
 
+def _get_single_nan_res(measurements, batch_size, interface, shots):
+    """Helper to get NaN results for one item in a shot vector."""
+
+    res = []
+
+    for m in measurements:
+        if isinstance(m, qml.measurements.SampleMP):
+            res.append(qml.math.asarray([], like=interface))
+            continue
+        if isinstance(m, qml.measurements.CountsMP):
+            res.append({})
+            continue
+
+        shape = m.shape(qml.device("default.qubit", wires=m.wires), qml.measurements.Shots(shots))
+        if batch_size is not None:
+            shape = (batch_size,) + shape
+
+        if shape == ():
+            out = qml.math.asarray(qml.numpy.NaN, like=interface)
+        else:
+            out = qml.math.full(shape, qml.numpy.NaN, like=interface)
+
+        res.append(out)
+
+    res = tuple(res)
+
+    if len(res) == 1:
+        res = res[0]
+
+    return res
+
+
+def _measure_nan_state(circuit, state, is_state_batched):
+    """Helper function for creating NaN results with the expected shape."""
+    batch_size = qml.math.shape(state)[0] if is_state_batched else None
+    interface = qml.math.get_interface(state)
+
+    if circuit.shots.has_partitioned_shots:
+        res = tuple(
+            _get_single_nan_res(circuit.measurements, batch_size, interface, s)
+            for s in circuit.shots
+        )
+    else:
+        res = _get_single_nan_res(
+            circuit.measurements, batch_size, interface, circuit.shots.total_shots
+        )
+
+    return res
+
+
 def measure_final_state(circuit, state, is_state_batched, rng=None, prng_key=None) -> Result:
     """
     Perform the measurements required by the circuit on the provided state.
@@ -179,9 +302,10 @@ def measure_final_state(circuit, state, is_state_batched, rng=None, prng_key=Non
     Returns:
         Tuple[TensorLike]: The measurement results
     """
+
     circuit = circuit.map_to_standard_wires()
 
-    if any(nan_value for nan_value in qml.math.flatten(qml.math.isnan(state))):
+    if qml.math.any(qml.math.isnan(state)):
         return _measure_nan_state(circuit, state, is_state_batched)
 
     if not circuit.shots:
@@ -215,56 +339,6 @@ def measure_final_state(circuit, state, is_state_batched, rng=None, prng_key=Non
     return results
 
 
-def _measure_nan_state(circuit, state, is_state_batched):
-    """Helper function for creating NaN results with the expected shape."""
-    batch_size = qml.math.shape(state)[0] if is_state_batched else None
-    interface = qml.math.get_interface(state)
-
-    if circuit.shots.has_partitioned_shots:
-        res = tuple(
-            _get_single_nan_res(circuit.measurements, batch_size, interface, s)
-            for s in circuit.shots
-        )
-    else:
-        res = _get_single_nan_res(
-            circuit.measurements, batch_size, interface, circuit.shots.total_shots
-        )
-
-    return res
-
-
-def _get_single_nan_res(measurements, batch_size, interface, shots):
-    """Helper to get NaN results for one item in a shot vector."""
-
-    res = []
-
-    for m in measurements:
-        if isinstance(m, qml.measurements.SampleMP):
-            res.append(qml.math.asarray([], like=interface))
-            continue
-        if isinstance(m, qml.measurements.CountsMP):
-            res.append({})
-            continue
-
-        shape = m.shape(None, qml.measurements.Shots(shots))
-        if batch_size is not None:
-            shape = (batch_size,) + shape
-
-        if shape == ():
-            out = qml.math.asarray(qml.numpy.NaN, like=interface)
-        else:
-            out = qml.math.full(shape, qml.numpy.NaN, like=interface)
-
-        res.append(out)
-
-    res = tuple(res)
-
-    if len(res) == 1:
-        res = res[0]
-
-    return res
-
-
 def simulate(
     circuit: qml.tape.QuantumScript, rng=None, prng_key=None, debugger=None, interface=None
 ) -> Result:
diff --git a/pennylane/math/single_dispatch.py b/pennylane/math/single_dispatch.py
index 1653666bd7d..c5cd49c6d72 100644
--- a/pennylane/math/single_dispatch.py
+++ b/pennylane/math/single_dispatch.py
@@ -268,12 +268,11 @@ def _take_autograd(tensor, indices, axis=None):
 ar.register_function(
     "tensorflow",
     "isnan",
-    lambda tensor, **kwargs: _i("tf").math.is_nan(
-        _i("tf").cast(tensor, "float64")
-        if tensor.dtype.name in ("int64", "int32", "complex128")
-        else tensor,
-        **kwargs,
-    ),
+    lambda tensor, **kwargs: _i("tf").math.is_nan(_i("tf").math.real(tensor), **kwargs)
+    | _i("tf").math.is_nan(_i("tf").math.imag(tensor), **kwargs),
+)
+ar.register_function(
+    "tensorflow", "any", lambda tensor, **kwargs: _i("tf").reduce_any(tensor, **kwargs)
 )
 ar.register_function(
     "tensorflow",
diff --git a/pennylane/measurements/__init__.py b/pennylane/measurements/__init__.py
index 23697dca533..a24611f3609 100644
--- a/pennylane/measurements/__init__.py
+++ b/pennylane/measurements/__init__.py
@@ -282,7 +282,7 @@ def circuit(x):
 from .purity import PurityMP, purity
 from .probs import ProbabilityMP, probs
 from .sample import SampleMP, sample
-from .shots import Shots, ShotCopies, _FlexShots
+from .shots import Shots, ShotCopies
 from .state import StateMP, DensityMatrixMP, density_matrix, state
 from .var import VarianceMP, var
 from .vn_entropy import VnEntropyMP, vn_entropy
diff --git a/pennylane/measurements/counts.py b/pennylane/measurements/counts.py
index e1a8e42c257..5369af96bb2 100644
--- a/pennylane/measurements/counts.py
+++ b/pennylane/measurements/counts.py
@@ -218,6 +218,19 @@ def process_samples(
         shot_range: Tuple[int] = None,
         bin_size: int = None,
     ):
+        if qml.math.any(qml.math.isnan(samples)) or len(samples) == 0:
+            if self.all_outcomes:
+                if (eigs := self.eigvals()) is not None:
+                    keys = [str(e) for e in eigs]
+                    vals = [0] * len(eigs)
+                else:
+                    bits = len(wire_order)
+                    keys = [format(i, f"{bits}b") for i in range(2**bits)]
+                    vals = [0] * 2**bits
+                return dict(zip(keys, vals))
+
+            return {}
+
         if self.mv is not None:
             samples = qml.sample(wires=self.mv.wires).process_samples(
                 samples, wire_order, shot_range, bin_size
diff --git a/pennylane/measurements/probs.py b/pennylane/measurements/probs.py
index 9b5cbf9a4a2..9ba607ae131 100644
--- a/pennylane/measurements/probs.py
+++ b/pennylane/measurements/probs.py
@@ -159,6 +159,9 @@ def process_samples(
         shot_range: Tuple[int] = None,
         bin_size: int = None,
     ):
+        if qml.math.any(qml.math.isnan(samples)) or len(samples) == 0:
+            return qml.math.asarray([0.0] * 2 ** len(wire_order), dtype="float64")
+
         wire_map = dict(zip(wire_order, range(len(wire_order))))
         mapped_wires = [wire_map[w] for w in self.wires]
         if shot_range is not None:
diff --git a/pennylane/measurements/shots.py b/pennylane/measurements/shots.py
index bbe8f9e2189..2d4a4bafe26 100644
--- a/pennylane/measurements/shots.py
+++ b/pennylane/measurements/shots.py
@@ -239,50 +239,3 @@ def has_partitioned_shots(self):
     def num_copies(self):
         """The total number of copies of any shot quantity."""
         return sum(s.copies for s in self.shot_vector)
-
-
-def valid_flex_int(s):
-    """Returns True if s is a non-negative integer."""
-    return isinstance(s, int) and s >= 0
-
-
-def valid_flex_tuple(s):
-    """Returns True if s is a tuple of the form (shots, copies)."""
-    return isinstance(s, tuple) and len(s) == 2 and valid_flex_int(s[0]) and valid_int(s[1])
-
-
-class _FlexShots(Shots):
-    """Shots class that allows zero shots."""
-
-    total_shots: int = None
-    """The total number of shots to be executed."""
-
-    shot_vector: Tuple[ShotCopies] = None
-    """The tuple of :class:`~ShotCopies` to be executed. Each element is of the form ``(shots, copies)``."""
-
-    _SHOT_ERROR = ValueError(
-        "Shots must be a single positive integer, a tuple pair of the form (shots, copies), or a sequence of these types."
-    )
-
-    _frozen = False
-
-    # pylint: disable=super-init-not-called
-    def __init__(self, shots=None):
-        if shots is None:
-            self.total_shots = None
-            self.shot_vector = ()
-        elif isinstance(shots, int):
-            if shots < 0:
-                raise self._SHOT_ERROR
-            self.total_shots = shots
-            self.shot_vector = (ShotCopies(shots, 1),)
-        elif isinstance(shots, Sequence):
-            if not all(valid_flex_int(s) or valid_flex_tuple(s) for s in shots):
-                raise self._SHOT_ERROR
-            self.__all_tuple_init__([s if isinstance(s, tuple) else (s, 1) for s in shots])
-        elif isinstance(shots, self.__class__):
-            return  # self already _is_ shots as defined by __new__
-        else:
-            raise self._SHOT_ERROR
-
-        self._frozen = True
diff --git a/tests/devices/qubit/test_simulate.py b/tests/devices/qubit/test_simulate.py
index 9655aa86176..94e52de0dc2 100644
--- a/tests/devices/qubit/test_simulate.py
+++ b/tests/devices/qubit/test_simulate.py
@@ -380,21 +380,32 @@ def test_broadcasting_with_projector(self, shots):
     @pytest.mark.all_interfaces
     @pytest.mark.parametrize("mp", [qml.expval, qml.var])
     @pytest.mark.parametrize("interface", ["numpy", "autograd", "torch", "tensorflow", "jax"])
+    @pytest.mark.parametrize("jax_jit", [True, False])
     @pytest.mark.parametrize("shots", [None, 10, [10, 10]])
-    def test_nan_float_result(self, mp, interface, shots):
+    def test_nan_float_result(self, mp, interface, shots, jax_jit):
         """Test that the result of circuits with 0 probability postselections is NaN with the
         expected shape."""
+        if interface != "jax" and jax_jit:
+            pytest.skip("Can't jit without jax")
 
-        tape = qml.tape.QuantumScript(
-            [
-                qml.RX(qml.math.asarray(np.pi, like=interface, dtype="float64"), 0),
-                qml.Projector([0], wires=0),
-            ],
-            [mp(qml.PauliZ(0))],
-            shots=shots,
-        )
+        def f():
+            tape = qml.tape.QuantumScript(
+                [
+                    qml.RX(qml.math.asarray(np.pi, like=interface, dtype="float64"), 0),
+                    qml.Projector([0], wires=0),
+                ],
+                [mp(qml.PauliZ(0))],
+                shots=shots,
+            )
 
-        res = simulate(tape)
+            return simulate(tape)
+
+        if jax_jit:
+            import jax
+
+            f = jax.jit(f)
+
+        res = f()
 
         if not isinstance(shots, list):
             assert qml.math.ndim(res) == 0
@@ -414,22 +425,32 @@ def test_nan_float_result(self, mp, interface, shots):
         "mp", [qml.probs(wires=0), qml.probs(op=qml.PauliZ(0)), qml.probs(wires=[0, 1])]
     )
     @pytest.mark.parametrize("interface", ["numpy", "autograd", "torch", "tensorflow", "jax"])
+    @pytest.mark.parametrize("jax_jit", [True, False])
     @pytest.mark.parametrize("shots", [None, 10, [10, 10]])
-    def test_nan_probs(self, mp, interface, shots):
+    def test_nan_probs(self, mp, interface, shots, jax_jit):
         """Test that the result of circuits with 0 probability postselections is NaN with the
         expected shape."""
+        if interface != "jax" and jax_jit:
+            pytest.skip("Can't jit without jax")
 
-        tape = qml.tape.QuantumScript(
-            [
-                qml.RX(qml.math.asarray(np.pi, like=interface, dtype="float64"), 0),
-                qml.CNOT([0, 1]),
-                qml.Projector([0], wires=0),
-            ],
-            [mp],
-            shots=shots,
-        )
+        def f():
+            tape = qml.tape.QuantumScript(
+                [
+                    qml.RX(qml.math.asarray(np.pi, like=interface, dtype="float64"), 0),
+                    qml.CNOT([0, 1]),
+                    qml.Projector([0], wires=0),
+                ],
+                [mp],
+                shots=shots,
+            )
+            return simulate(tape)
 
-        res = simulate(tape)
+        if jax_jit:
+            import jax
+
+            f = jax.jit(f)
+
+        res = f()
 
         if not isinstance(shots, list):
             assert qml.math.shape(res) == (2 ** len(mp.wires),)
@@ -449,22 +470,32 @@ def test_nan_probs(self, mp, interface, shots):
         "mp", [qml.sample(wires=0), qml.sample(op=qml.PauliZ(0)), qml.sample(wires=[0, 1])]
     )
     @pytest.mark.parametrize("interface", ["numpy", "autograd", "torch", "tensorflow", "jax"])
-    @pytest.mark.parametrize("shots", [None, 10, [10, 10]])
-    def test_nan_samples(self, mp, interface, shots):
+    @pytest.mark.parametrize("jax_jit", [True, False])
+    @pytest.mark.parametrize("shots", [10, [10, 10]])
+    def test_nan_samples(self, mp, interface, shots, jax_jit):
         """Test that the result of circuits with 0 probability postselections is NaN with the
         expected shape."""
+        if interface != "jax" and jax_jit:
+            pytest.skip("Can't jit without jax")
 
-        tape = qml.tape.QuantumScript(
-            [
-                qml.RX(qml.math.asarray(np.pi, like=interface, dtype="float64"), 0),
-                qml.CNOT([0, 1]),
-                qml.Projector([0], wires=0),
-            ],
-            [mp],
-            shots=shots,
-        )
+        def f():
+            tape = qml.tape.QuantumScript(
+                [
+                    qml.RX(qml.math.asarray(np.pi, like=interface, dtype="float64"), 0),
+                    qml.CNOT([0, 1]),
+                    qml.Projector([0], wires=0),
+                ],
+                [mp],
+                shots=shots,
+            )
+            return simulate(tape)
 
-        res = simulate(tape)
+        if jax_jit:
+            import jax
+
+            f = jax.jit(f)
+
+        res = f()
 
         if not isinstance(shots, list):
             assert qml.math.shape(res) == (0,)
@@ -482,22 +513,32 @@ def test_nan_samples(self, mp, interface, shots):
         "mp", [qml.counts(wires=0), qml.counts(op=qml.PauliZ(0)), qml.counts()]
     )
     @pytest.mark.parametrize("interface", ["numpy", "autograd", "torch", "tensorflow", "jax"])
-    @pytest.mark.parametrize("shots", [None, 10, [10, 10]])
-    def test_nan_counts(self, mp, interface, shots):
+    @pytest.mark.parametrize("jax_jit", [True, False])
+    @pytest.mark.parametrize("shots", [10, [10, 10]])
+    def test_nan_counts(self, mp, interface, shots, jax_jit):
         """Test that the result of circuits with 0 probability postselections is NaN with the
         expected shape."""
+        if interface != "jax" and jax_jit:
+            pytest.skip("Can't jit without jax")
 
-        tape = qml.tape.QuantumScript(
-            [
-                qml.RX(qml.math.asarray(np.pi, like=interface, dtype="float64"), 0),
-                qml.CNOT([0, 1]),
-                qml.Projector([0], wires=0),
-            ],
-            [mp],
-            shots=shots,
-        )
+        def f():
+            tape = qml.tape.QuantumScript(
+                [
+                    qml.RX(qml.math.asarray(np.pi, like=interface, dtype="float64"), 0),
+                    qml.CNOT([0, 1]),
+                    qml.Projector([0], wires=0),
+                ],
+                [mp],
+                shots=shots,
+            )
+            return simulate(tape)
 
-        res = simulate(tape)
+        if jax_jit:
+            import jax
+
+            f = jax.jit(f)
+
+        res = f()
 
         if not isinstance(shots, list):
             assert res == {}

From 1ebc58d9cdd17997fff6de1b7284d6e3172f446b Mon Sep 17 00:00:00 2001
From: Mudit Pandey <mudit.pandey@xanadu.ai>
Date: Wed, 4 Oct 2023 17:36:43 -0400
Subject: [PATCH 110/131] Updated postselection measurement

---
 pennylane/devices/qubit/measure.py   |  36 +++++++++
 pennylane/devices/qubit/sampling.py  | 106 ++++++++++++++++++++++++++-
 pennylane/devices/qubit/simulate.py  |  53 --------------
 tests/devices/qubit/test_simulate.py |  40 +---------
 4 files changed, 143 insertions(+), 92 deletions(-)

diff --git a/pennylane/devices/qubit/measure.py b/pennylane/devices/qubit/measure.py
index 9b72def187d..d1cddb21780 100644
--- a/pennylane/devices/qubit/measure.py
+++ b/pennylane/devices/qubit/measure.py
@@ -18,6 +18,7 @@
 
 from scipy.sparse import csr_matrix
 
+import pennylane as qml
 from pennylane import math
 from pennylane.ops import Sum, Hamiltonian
 from pennylane.measurements import (
@@ -25,6 +26,7 @@
     MeasurementProcess,
     MeasurementValue,
     ExpectationMP,
+    Shots,
 )
 from pennylane.typing import TensorLike
 from pennylane.wires import Wires
@@ -136,6 +138,36 @@ def sum_of_terms_method(
     )
 
 
+def state_measure_nan(
+    measurement: StateMeasurement, state: TensorLike, is_state_batched: bool = False
+) -> TensorLike:
+    """Return measurement outcomes when the state contains NaN values.
+
+    Args:
+        measurementprocess (ExpectationMP): measurement process to apply to the state
+        state (TensorLike): the state to measure
+        is_state_batched (bool): whether the state is batched or not
+
+    Returns:
+        TensorLike: the result of the measurement
+    """
+
+    interface = math.get_interface(state)
+    batch_size = math.shape(state)[0] if is_state_batched else None
+
+    shape = measurement.shape(qml.device("default.qubit", wires=measurement.wires), Shots(None))
+    if batch_size is not None:
+        shape = (batch_size,) + shape
+
+    nan_val = (
+        measurement.numeric_type(qml.numpy.NaN)
+        if measurement.numeric_type is not int
+        else qml.numpy.NaN
+    )
+    return math.full(shape, measurement.numeric_type(nan_val), like=interface)
+
+
+# pylint: disable=too-many-return-statements
 def get_measurement_function(
     measurementprocess: MeasurementProcess, state: TensorLike
 ) -> Callable[[MeasurementProcess, TensorLike], TensorLike]:
@@ -149,7 +181,11 @@ def get_measurement_function(
     Returns:
         Callable: function that returns the measurement result
     """
+
     if isinstance(measurementprocess, StateMeasurement):
+        if not math.is_abstract(state) and math.any(math.isnan(state)):
+            return state_measure_nan
+
         if isinstance(measurementprocess.mv, MeasurementValue):
             return state_diagonalizing_gates
 
diff --git a/pennylane/devices/qubit/sampling.py b/pennylane/devices/qubit/sampling.py
index 0d418b95d51..7a45403d7d2 100644
--- a/pennylane/devices/qubit/sampling.py
+++ b/pennylane/devices/qubit/sampling.py
@@ -23,6 +23,7 @@
     ExpectationMP,
     ClassicalShadowMP,
     ShadowExpvalMP,
+    SampleMP,
     CountsMP,
 )
 from pennylane.typing import TensorLike
@@ -137,6 +138,9 @@ def measure_with_samples(
     Returns:
         List[TensorLike[Any]]: Sample measurement results
     """
+    if not qml.math.is_abstract(state) and qml.math.any(qml.math.isnan(state)):
+        return _measure_with_samples_nan(mps, state, shots, is_state_batched)
+
     groups, indices = _group_measurements(mps)
 
     all_res = []
@@ -171,6 +175,97 @@ def measure_with_samples(
     return sorted_res
 
 
+def _measure_with_samples_nan(
+    mps: List[SampleMeasurement], state: np.ndarray, shots: Shots, is_state_batched: bool = False
+) -> TensorLike:
+    batch_size = qml.math.shape(state)[0] if is_state_batched else None
+    interface = qml.math.get_interface(state)
+
+    if shots.has_partitioned_shots:
+        # res = []
+        # for s in shots:
+        #     res.extend(_get_single_shot_nan_res(mps, batch_size, interface, s))
+        # res = tuple(res)
+        res = tuple(_get_single_shot_nan_res(mps, batch_size, interface, s) for s in shots)
+    else:
+        res = _get_single_shot_nan_res(mps, batch_size, interface, shots.total_shots)
+
+    return res
+
+
+def _get_single_shot_nan_res(mps, batch_size, interface, shots):
+    """Helper to get NaN results for one item in a shot vector."""
+
+    res = []
+
+    for m in mps:
+        if isinstance(m, SampleMP):
+            out = [] if batch_size is None else [[]] * batch_size
+            res.append(qml.math.asarray(out, like=interface))
+            continue
+
+        if isinstance(m, CountsMP):
+            if m.all_outcomes:
+                if (eigs := m.eigvals()) is not None:
+                    keys = [str(e) for e in eigs]
+                    vals = [0] * len(eigs)
+                else:
+                    bits = len(m.wires)
+                    keys = [format(i, f"{bits}b") for i in range(2**bits)]
+                    vals = [0] * 2**bits
+                    res.append(
+                        dict(zip(keys, vals))
+                        if batch_size is None
+                        else [dict(zip(keys, vals)) for _ in range(batch_size)]
+                    )
+            else:
+                res.append({} if batch_size is None else [{} for _ in range(batch_size)])
+
+            continue
+
+        shape = m.shape(qml.device("default.qubit", wires=m.wires), Shots(shots))
+        if batch_size is not None:
+            shape = (batch_size,) + shape
+
+        nan_val = m.numeric_type(qml.numpy.NaN) if m.numeric_type is not int else qml.numpy.NaN
+        res.append(qml.math.full(shape, nan_val, like=interface))
+
+    return res
+
+
+def _get_zero_shot_res(mp, batch_size, interface):
+    """Helper to get measurement outcome when there are zero shots."""
+    if isinstance(mp, SampleMP):
+        out = [] if batch_size is None else [[]] * batch_size
+        return qml.math.asarray(out, like=interface)
+
+    if isinstance(mp, CountsMP):
+        if mp.all_outcomes:
+            if (eigs := mp.eigvals()) is not None:
+                keys = [str(e) for e in eigs]
+                vals = [0] * len(eigs)
+            else:
+                bits = len(mp.wires)
+                keys = [format(i, f"{bits}b") for i in range(2**bits)]
+                vals = [0] * 2**bits
+            return (
+                dict(zip(keys, vals))
+                if batch_size is None
+                else [dict(zip(keys, vals)) for _ in range(batch_size)]
+            )
+
+        return {} if batch_size is None else [{} for _ in range(batch_size)]
+
+    # Shots can be set to any positive integer to get the correct shape
+    # from the measurement process
+    shape = mp.shape(qml.device("default.qubit", wires=mp.wires), Shots(1))
+    if batch_size is not None:
+        shape = (batch_size,) + shape
+
+    nan_val = mp.numeric_type(qml.numpy.NaN) if mp.numeric_type is not int else qml.numpy.NaN
+    return qml.math.full(shape, nan_val, like=interface)
+
+
 def _measure_with_samples_diagonalizing_gates(
     mps: List[SampleMeasurement],
     state: np.ndarray,
@@ -207,9 +302,14 @@ def _measure_with_samples_diagonalizing_gates(
     def _process_single_shot(samples):
         processed = []
         for mp in mps:
-            res = mp.process_samples(samples, wires)
-            if not isinstance(mp, CountsMP):
-                res = qml.math.squeeze(res)
+            if len(samples) == 0:
+                batch_size = qml.math.shape(state)[0] if is_state_batched else None
+                interface = qml.math.get_interface(state)
+                res = _get_zero_shot_res(mp, batch_size, interface)
+            else:
+                res = mp.process_samples(samples, wires)
+                if not isinstance(mp, CountsMP):
+                    res = qml.math.squeeze(res)
 
             processed.append(res)
 
diff --git a/pennylane/devices/qubit/simulate.py b/pennylane/devices/qubit/simulate.py
index 605ee268d86..502404f010e 100644
--- a/pennylane/devices/qubit/simulate.py
+++ b/pennylane/devices/qubit/simulate.py
@@ -231,56 +231,6 @@ def get_final_state(circuit, debugger=None, interface=None):
     return state, is_state_batched
 
 
-def _get_single_nan_res(measurements, batch_size, interface, shots):
-    """Helper to get NaN results for one item in a shot vector."""
-
-    res = []
-
-    for m in measurements:
-        if isinstance(m, qml.measurements.SampleMP):
-            res.append(qml.math.asarray([], like=interface))
-            continue
-        if isinstance(m, qml.measurements.CountsMP):
-            res.append({})
-            continue
-
-        shape = m.shape(qml.device("default.qubit", wires=m.wires), qml.measurements.Shots(shots))
-        if batch_size is not None:
-            shape = (batch_size,) + shape
-
-        if shape == ():
-            out = qml.math.asarray(qml.numpy.NaN, like=interface)
-        else:
-            out = qml.math.full(shape, qml.numpy.NaN, like=interface)
-
-        res.append(out)
-
-    res = tuple(res)
-
-    if len(res) == 1:
-        res = res[0]
-
-    return res
-
-
-def _measure_nan_state(circuit, state, is_state_batched):
-    """Helper function for creating NaN results with the expected shape."""
-    batch_size = qml.math.shape(state)[0] if is_state_batched else None
-    interface = qml.math.get_interface(state)
-
-    if circuit.shots.has_partitioned_shots:
-        res = tuple(
-            _get_single_nan_res(circuit.measurements, batch_size, interface, s)
-            for s in circuit.shots
-        )
-    else:
-        res = _get_single_nan_res(
-            circuit.measurements, batch_size, interface, circuit.shots.total_shots
-        )
-
-    return res
-
-
 def measure_final_state(circuit, state, is_state_batched, rng=None, prng_key=None) -> Result:
     """
     Perform the measurements required by the circuit on the provided state.
@@ -305,9 +255,6 @@ def measure_final_state(circuit, state, is_state_batched, rng=None, prng_key=Non
 
     circuit = circuit.map_to_standard_wires()
 
-    if qml.math.any(qml.math.isnan(state)):
-        return _measure_nan_state(circuit, state, is_state_batched)
-
     if not circuit.shots:
         # analytic case
 
diff --git a/tests/devices/qubit/test_simulate.py b/tests/devices/qubit/test_simulate.py
index 94e52de0dc2..c49bd34c0ce 100644
--- a/tests/devices/qubit/test_simulate.py
+++ b/tests/devices/qubit/test_simulate.py
@@ -380,13 +380,10 @@ def test_broadcasting_with_projector(self, shots):
     @pytest.mark.all_interfaces
     @pytest.mark.parametrize("mp", [qml.expval, qml.var])
     @pytest.mark.parametrize("interface", ["numpy", "autograd", "torch", "tensorflow", "jax"])
-    @pytest.mark.parametrize("jax_jit", [True, False])
     @pytest.mark.parametrize("shots", [None, 10, [10, 10]])
-    def test_nan_float_result(self, mp, interface, shots, jax_jit):
+    def test_nan_float_result(self, mp, interface, shots):
         """Test that the result of circuits with 0 probability postselections is NaN with the
         expected shape."""
-        if interface != "jax" and jax_jit:
-            pytest.skip("Can't jit without jax")
 
         def f():
             tape = qml.tape.QuantumScript(
@@ -400,11 +397,6 @@ def f():
 
             return simulate(tape)
 
-        if jax_jit:
-            import jax
-
-            f = jax.jit(f)
-
         res = f()
 
         if not isinstance(shots, list):
@@ -425,13 +417,10 @@ def f():
         "mp", [qml.probs(wires=0), qml.probs(op=qml.PauliZ(0)), qml.probs(wires=[0, 1])]
     )
     @pytest.mark.parametrize("interface", ["numpy", "autograd", "torch", "tensorflow", "jax"])
-    @pytest.mark.parametrize("jax_jit", [True, False])
     @pytest.mark.parametrize("shots", [None, 10, [10, 10]])
-    def test_nan_probs(self, mp, interface, shots, jax_jit):
+    def test_nan_probs(self, mp, interface, shots):
         """Test that the result of circuits with 0 probability postselections is NaN with the
         expected shape."""
-        if interface != "jax" and jax_jit:
-            pytest.skip("Can't jit without jax")
 
         def f():
             tape = qml.tape.QuantumScript(
@@ -445,11 +434,6 @@ def f():
             )
             return simulate(tape)
 
-        if jax_jit:
-            import jax
-
-            f = jax.jit(f)
-
         res = f()
 
         if not isinstance(shots, list):
@@ -470,13 +454,10 @@ def f():
         "mp", [qml.sample(wires=0), qml.sample(op=qml.PauliZ(0)), qml.sample(wires=[0, 1])]
     )
     @pytest.mark.parametrize("interface", ["numpy", "autograd", "torch", "tensorflow", "jax"])
-    @pytest.mark.parametrize("jax_jit", [True, False])
     @pytest.mark.parametrize("shots", [10, [10, 10]])
-    def test_nan_samples(self, mp, interface, shots, jax_jit):
+    def test_nan_samples(self, mp, interface, shots):
         """Test that the result of circuits with 0 probability postselections is NaN with the
         expected shape."""
-        if interface != "jax" and jax_jit:
-            pytest.skip("Can't jit without jax")
 
         def f():
             tape = qml.tape.QuantumScript(
@@ -490,11 +471,6 @@ def f():
             )
             return simulate(tape)
 
-        if jax_jit:
-            import jax
-
-            f = jax.jit(f)
-
         res = f()
 
         if not isinstance(shots, list):
@@ -513,13 +489,10 @@ def f():
         "mp", [qml.counts(wires=0), qml.counts(op=qml.PauliZ(0)), qml.counts()]
     )
     @pytest.mark.parametrize("interface", ["numpy", "autograd", "torch", "tensorflow", "jax"])
-    @pytest.mark.parametrize("jax_jit", [True, False])
     @pytest.mark.parametrize("shots", [10, [10, 10]])
-    def test_nan_counts(self, mp, interface, shots, jax_jit):
+    def test_nan_counts(self, mp, interface, shots):
         """Test that the result of circuits with 0 probability postselections is NaN with the
         expected shape."""
-        if interface != "jax" and jax_jit:
-            pytest.skip("Can't jit without jax")
 
         def f():
             tape = qml.tape.QuantumScript(
@@ -533,11 +506,6 @@ def f():
             )
             return simulate(tape)
 
-        if jax_jit:
-            import jax
-
-            f = jax.jit(f)
-
         res = f()
 
         if not isinstance(shots, list):

From e862d49a3623df8d2945c2a7a0d589c66c65b20e Mon Sep 17 00:00:00 2001
From: Mudit Pandey <mudit.pandey@xanadu.ai>
Date: Thu, 5 Oct 2023 10:44:08 -0400
Subject: [PATCH 111/131] Updated defer_measurements tests

---
 pennylane/transforms/defer_measurements.py  |  7 +++++-
 tests/devices/qubit/test_simulate.py        | 26 +++++++++++++++++++--
 tests/transforms/test_defer_measurements.py |  8 +++++--
 3 files changed, 36 insertions(+), 5 deletions(-)

diff --git a/pennylane/transforms/defer_measurements.py b/pennylane/transforms/defer_measurements.py
index ae5b9f8ce7e..c92c06ecba9 100644
--- a/pennylane/transforms/defer_measurements.py
+++ b/pennylane/transforms/defer_measurements.py
@@ -213,7 +213,12 @@ def func(x, y):
         new_measurements.append(mp)
 
     new_tape = type(tape)(new_operations, new_measurements, shots=tape.shots)
-    new_tape._qfunc_output = tape._qfunc_output  # pylint: disable=protected-access
+
+    # pylint: disable=protected-access
+    if not isinstance(tape._qfunc_output, Sequence):
+        new_tape._qfunc_output = new_measurements[0]
+    else:
+        new_tape._qfunc_output = type(tape._qfunc_output)(new_measurements)
 
     if is_postselecting and new_tape.batch_size is not None:
         # Split tapes if broadcasting with postselection
diff --git a/tests/devices/qubit/test_simulate.py b/tests/devices/qubit/test_simulate.py
index c49bd34c0ce..6586e6d76f0 100644
--- a/tests/devices/qubit/test_simulate.py
+++ b/tests/devices/qubit/test_simulate.py
@@ -378,20 +378,42 @@ def test_broadcasting_with_projector(self, shots):
             _ = simulate(tape)
 
     @pytest.mark.all_interfaces
-    @pytest.mark.parametrize("mp", [qml.expval, qml.var])
+    @pytest.mark.parametrize(
+        "mp",
+        [
+            qml.expval(qml.PauliZ(0)),
+            qml.var(qml.PauliZ(0)),
+            qml.purity(0),
+            qml.vn_entropy(0),
+            qml.mutual_info(0, 1),
+        ],
+    )
     @pytest.mark.parametrize("interface", ["numpy", "autograd", "torch", "tensorflow", "jax"])
     @pytest.mark.parametrize("shots", [None, 10, [10, 10]])
     def test_nan_float_result(self, mp, interface, shots):
         """Test that the result of circuits with 0 probability postselections is NaN with the
         expected shape."""
+        if (
+            isinstance(
+                mp,
+                (
+                    qml.measurements.VnEntropyMP,
+                    qml.measurements.MutualInfoMP,
+                    qml.measurements.PurityMP,
+                ),
+            )
+            and shots is not None
+        ):
+            pytest.skip("QInfo measurements should be done with analytic execution.")
 
         def f():
             tape = qml.tape.QuantumScript(
                 [
                     qml.RX(qml.math.asarray(np.pi, like=interface, dtype="float64"), 0),
+                    qml.CNOT([0, 1]),
                     qml.Projector([0], wires=0),
                 ],
-                [mp(qml.PauliZ(0))],
+                [mp],
                 shots=shots,
             )
 
diff --git a/tests/transforms/test_defer_measurements.py b/tests/transforms/test_defer_measurements.py
index 8bc2de2bb83..5d118f818e2 100644
--- a/tests/transforms/test_defer_measurements.py
+++ b/tests/transforms/test_defer_measurements.py
@@ -37,8 +37,8 @@ def test_only_mcm(self):
         def qnode1():
             return qml.expval(qml.PauliZ(0))
 
-        @qml.defer_measurements
         @qml.qnode(dev)
+        @qml.defer_measurements
         def qnode2():
             qml.measure(1)
             return qml.expval(qml.PauliZ(0))
@@ -162,6 +162,7 @@ def test_postselection_qnode(self, phi, shots):
         dev = DefaultQubit()
 
         @qml.qnode(dev)
+        @qml.defer_measurements
         def circ1(phi):
             qml.RX(phi, wires=0)
             # Postselecting on |1> on wire 0 means that the probability of measuring
@@ -201,6 +202,7 @@ def test_multiple_postselection_qnode(self, phi, theta, shots, tol, tol_stochast
         mv2 = MeasurementValue([mp2], lambda v: v)
 
         @qml.qnode(dev)
+        @qml.defer_measurements
         def circ1(phi, theta):
             qml.RX(phi, 0)
             qml.apply(mp0)
@@ -270,13 +272,14 @@ def circ2(x):
         param = 1.5
         assert np.allclose(circ1(param, shots=shots), circ2(param, shots=shots))
 
-    @pytest.mark.parametrize("shots", [None, 1000, [1000, 1000]])
+    @pytest.mark.parametrize("shots", [None, 2000, [2000, 2000]])
     def test_measured_value_wires_mapped(self, shots, tol, tol_stochastic):
         """Test that collecting statistics on a measurement value works correctly
         when the measured wire is reused."""
         dev = DefaultQubit()
 
         @qml.qnode(dev)
+        @qml.defer_measurements
         def circ1(x):
             qml.RX(x, 0)
             m0 = qml.measure(0)
@@ -1198,6 +1201,7 @@ def test_new_wire_for_multiple_measurements(self):
         dev = qml.device("default.qubit", wires=4)
 
         @qml.qnode(dev)
+        @qml.defer_measurements
         def circ(x, y):
             qml.RX(x, 0)
             qml.measure(0)

From c1c763ba6764ed18d4cd5f80362faf9817e71ea5 Mon Sep 17 00:00:00 2001
From: Mudit Pandey <mudit.pandey@xanadu.ai>
Date: Fri, 6 Oct 2023 15:57:21 -0400
Subject: [PATCH 112/131] Reverted validation for zero probability. Need to
 update tests

---
 pennylane/devices/qubit/measure.py          |  34 ----
 pennylane/devices/qubit/sampling.py         | 143 +++--------------
 pennylane/devices/qubit/simulate.py         |  55 ++-----
 pennylane/measurements/counts.py            |  13 --
 pennylane/measurements/probs.py             |   3 -
 pennylane/transforms/defer_measurements.py  |  20 ++-
 tests/devices/qubit/test_measure.py         |  76 +++++++++
 tests/devices/qubit/test_sampling.py        | 135 ++++++++++++++++
 tests/devices/qubit/test_simulate.py        | 168 +-------------------
 tests/transforms/test_defer_measurements.py |  41 +++++
 10 files changed, 317 insertions(+), 371 deletions(-)

diff --git a/pennylane/devices/qubit/measure.py b/pennylane/devices/qubit/measure.py
index d1cddb21780..5d04f713946 100644
--- a/pennylane/devices/qubit/measure.py
+++ b/pennylane/devices/qubit/measure.py
@@ -18,7 +18,6 @@
 
 from scipy.sparse import csr_matrix
 
-import pennylane as qml
 from pennylane import math
 from pennylane.ops import Sum, Hamiltonian
 from pennylane.measurements import (
@@ -26,7 +25,6 @@
     MeasurementProcess,
     MeasurementValue,
     ExpectationMP,
-    Shots,
 )
 from pennylane.typing import TensorLike
 from pennylane.wires import Wires
@@ -138,35 +136,6 @@ def sum_of_terms_method(
     )
 
 
-def state_measure_nan(
-    measurement: StateMeasurement, state: TensorLike, is_state_batched: bool = False
-) -> TensorLike:
-    """Return measurement outcomes when the state contains NaN values.
-
-    Args:
-        measurementprocess (ExpectationMP): measurement process to apply to the state
-        state (TensorLike): the state to measure
-        is_state_batched (bool): whether the state is batched or not
-
-    Returns:
-        TensorLike: the result of the measurement
-    """
-
-    interface = math.get_interface(state)
-    batch_size = math.shape(state)[0] if is_state_batched else None
-
-    shape = measurement.shape(qml.device("default.qubit", wires=measurement.wires), Shots(None))
-    if batch_size is not None:
-        shape = (batch_size,) + shape
-
-    nan_val = (
-        measurement.numeric_type(qml.numpy.NaN)
-        if measurement.numeric_type is not int
-        else qml.numpy.NaN
-    )
-    return math.full(shape, measurement.numeric_type(nan_val), like=interface)
-
-
 # pylint: disable=too-many-return-statements
 def get_measurement_function(
     measurementprocess: MeasurementProcess, state: TensorLike
@@ -183,9 +152,6 @@ def get_measurement_function(
     """
 
     if isinstance(measurementprocess, StateMeasurement):
-        if not math.is_abstract(state) and math.any(math.isnan(state)):
-            return state_measure_nan
-
         if isinstance(measurementprocess.mv, MeasurementValue):
             return state_diagonalizing_gates
 
diff --git a/pennylane/devices/qubit/sampling.py b/pennylane/devices/qubit/sampling.py
index fe884e2347d..cf009aa3dd5 100644
--- a/pennylane/devices/qubit/sampling.py
+++ b/pennylane/devices/qubit/sampling.py
@@ -23,7 +23,6 @@
     ExpectationMP,
     ClassicalShadowMP,
     ShadowExpvalMP,
-    SampleMP,
     CountsMP,
 )
 from pennylane.typing import TensorLike
@@ -180,8 +179,6 @@ def measure_with_samples(
     Returns:
         List[TensorLike[Any]]: Sample measurement results
     """
-    if not qml.math.is_abstract(state) and qml.math.any(qml.math.isnan(state)):
-        return _measure_with_samples_nan(mps, state, shots, is_state_batched)
 
     groups, indices = _group_measurements(mps)
 
@@ -217,97 +214,6 @@ def measure_with_samples(
     return sorted_res
 
 
-def _measure_with_samples_nan(
-    mps: List[SampleMeasurement], state: np.ndarray, shots: Shots, is_state_batched: bool = False
-) -> TensorLike:
-    batch_size = qml.math.shape(state)[0] if is_state_batched else None
-    interface = qml.math.get_interface(state)
-
-    if shots.has_partitioned_shots:
-        # res = []
-        # for s in shots:
-        #     res.extend(_get_single_shot_nan_res(mps, batch_size, interface, s))
-        # res = tuple(res)
-        res = tuple(_get_single_shot_nan_res(mps, batch_size, interface, s) for s in shots)
-    else:
-        res = _get_single_shot_nan_res(mps, batch_size, interface, shots.total_shots)
-
-    return res
-
-
-def _get_single_shot_nan_res(mps, batch_size, interface, shots):
-    """Helper to get NaN results for one item in a shot vector."""
-
-    res = []
-
-    for m in mps:
-        if isinstance(m, SampleMP):
-            out = [] if batch_size is None else [[]] * batch_size
-            res.append(qml.math.asarray(out, like=interface))
-            continue
-
-        if isinstance(m, CountsMP):
-            if m.all_outcomes:
-                if (eigs := m.eigvals()) is not None:
-                    keys = [str(e) for e in eigs]
-                    vals = [0] * len(eigs)
-                else:
-                    bits = len(m.wires)
-                    keys = [format(i, f"{bits}b") for i in range(2**bits)]
-                    vals = [0] * 2**bits
-                    res.append(
-                        dict(zip(keys, vals))
-                        if batch_size is None
-                        else [dict(zip(keys, vals)) for _ in range(batch_size)]
-                    )
-            else:
-                res.append({} if batch_size is None else [{} for _ in range(batch_size)])
-
-            continue
-
-        shape = m.shape(qml.device("default.qubit", wires=m.wires), Shots(shots))
-        if batch_size is not None:
-            shape = (batch_size,) + shape
-
-        nan_val = m.numeric_type(qml.numpy.NaN) if m.numeric_type is not int else qml.numpy.NaN
-        res.append(qml.math.full(shape, nan_val, like=interface))
-
-    return res
-
-
-def _get_zero_shot_res(mp, batch_size, interface):
-    """Helper to get measurement outcome when there are zero shots."""
-    if isinstance(mp, SampleMP):
-        out = [] if batch_size is None else [[]] * batch_size
-        return qml.math.asarray(out, like=interface)
-
-    if isinstance(mp, CountsMP):
-        if mp.all_outcomes:
-            if (eigs := mp.eigvals()) is not None:
-                keys = [str(e) for e in eigs]
-                vals = [0] * len(eigs)
-            else:
-                bits = len(mp.wires)
-                keys = [format(i, f"{bits}b") for i in range(2**bits)]
-                vals = [0] * 2**bits
-            return (
-                dict(zip(keys, vals))
-                if batch_size is None
-                else [dict(zip(keys, vals)) for _ in range(batch_size)]
-            )
-
-        return {} if batch_size is None else [{} for _ in range(batch_size)]
-
-    # Shots can be set to any positive integer to get the correct shape
-    # from the measurement process
-    shape = mp.shape(qml.device("default.qubit", wires=mp.wires), Shots(1))
-    if batch_size is not None:
-        shape = (batch_size,) + shape
-
-    nan_val = mp.numeric_type(qml.numpy.NaN) if mp.numeric_type is not int else qml.numpy.NaN
-    return qml.math.full(shape, nan_val, like=interface)
-
-
 def _measure_with_samples_diagonalizing_gates(
     mps: List[SampleMeasurement],
     state: np.ndarray,
@@ -344,14 +250,9 @@ def _measure_with_samples_diagonalizing_gates(
     def _process_single_shot(samples):
         processed = []
         for mp in mps:
-            if len(samples) == 0:
-                batch_size = qml.math.shape(state)[0] if is_state_batched else None
-                interface = qml.math.get_interface(state)
-                res = _get_zero_shot_res(mp, batch_size, interface)
-            else:
-                res = mp.process_samples(samples, wires)
-                if not isinstance(mp, CountsMP):
-                    res = qml.math.squeeze(res)
+            res = mp.process_samples(samples, wires)
+            if not isinstance(mp, CountsMP):
+                res = qml.math.squeeze(res)
 
             processed.append(res)
 
@@ -364,27 +265,33 @@ def _process_single_shot(samples):
             # currently we call sample_state for each shot entry, but it may be
             # better to call sample_state just once with total_shots, then use
             # the shot_range keyword argument
-            samples = sample_state(
-                state,
-                shots=s,
-                is_state_batched=is_state_batched,
-                wires=wires,
-                rng=rng,
-                prng_key=prng_key,
-            )
+            try:
+                samples = sample_state(
+                    state,
+                    shots=s,
+                    is_state_batched=is_state_batched,
+                    wires=wires,
+                    rng=rng,
+                    prng_key=prng_key,
+                )
+            except ValueError:
+                samples = qml.math.full((s, len(wires)), np.NaN)
 
             processed_samples.append(_process_single_shot(samples))
 
         return tuple(zip(*processed_samples))
 
-    samples = sample_state(
-        state,
-        shots=shots.total_shots,
-        is_state_batched=is_state_batched,
-        wires=wires,
-        rng=rng,
-        prng_key=prng_key,
-    )
+    try:
+        samples = sample_state(
+            state,
+            shots=shots.total_shots,
+            is_state_batched=is_state_batched,
+            wires=wires,
+            rng=rng,
+            prng_key=prng_key,
+        )
+    except ValueError:
+        samples = qml.math.full((shots.total_shots, len(wires)), np.NaN)
 
     return _process_single_shot(samples)
 
diff --git a/pennylane/devices/qubit/simulate.py b/pennylane/devices/qubit/simulate.py
index 502404f010e..2da17f99ae3 100644
--- a/pennylane/devices/qubit/simulate.py
+++ b/pennylane/devices/qubit/simulate.py
@@ -66,18 +66,6 @@ def _valid_flex_tuple(s):
 class _FlexShots(qml.measurements.Shots):
     """Shots class that allows zero shots."""
 
-    total_shots = None
-    """The total number of shots to be executed."""
-
-    shot_vector = None
-    """The tuple of :class:`~ShotCopies` to be executed. Each element is of the form ``(shots, copies)``."""
-
-    _SHOT_ERROR = ValueError(
-        "Shots must be a single positive integer, a tuple pair of the form (shots, copies), or a sequence of these types."
-    )
-
-    _frozen = False
-
     # pylint: disable=super-init-not-called
     def __init__(self, shots=None):
         if shots is None:
@@ -147,37 +135,22 @@ def _postselection_postprocess(state, is_state_batched, shots):
             "postselection is used."
         )
 
-    if qml.math.is_abstract(state):
-        return state, shots, False
+    # if qml.math.is_abstract(state):
+    norm = qml.math.floor(qml.math.real(qml.math.norm(state)) * 1e8) * 1e-8
 
-    norm = qml.math.norm(state)
-    if not qml.math.is_abstract(norm) and qml.math.isclose(float(norm), 0.0):
-        new_state = qml.math.cast_like(
-            qml.math.full(qml.math.shape(state), qml.numpy.NaN, like=qml.math.get_interface(state)),
-            state,
+    if shots:
+        # Clip the number of shots using a binomial distribution using the probability of
+        # measuring the postselected state.
+        postselected_shots = (
+            [binomial(s, float(norm)) for s in shots] if not qml.math.is_abstract(norm) else shots
         )
-        is_nan = True
-
-    else:
-        new_state = state / norm
-        is_nan = False
-
-        # defer_measurements will raise an error with batched shots or broadcasting so we can
-        # assume that both the state and shots are unbatched.
-        if shots:
-            # Clip the number of shots using a binomial distribution using the probability of
-            # measuring the postselected state.
-            postselected_shots = (
-                [binomial(s, float(norm)) for s in shots]
-                if not qml.math.is_abstract(norm)
-                else shots
-            )
 
-            # _FlexShots is used here since the binomial distribution could result in zero
-            # valid samples
-            shots = _FlexShots(postselected_shots)
+        # _FlexShots is used here since the binomial distribution could result in zero
+        # valid samples
+        shots = _FlexShots(postselected_shots)
 
-    return new_state, shots, is_nan
+    state = state / qml.math.cast_like(norm, state)
+    return state, shots
 
 
 def get_final_state(circuit, debugger=None, interface=None):
@@ -211,11 +184,9 @@ def get_final_state(circuit, debugger=None, interface=None):
 
         # Handle postselection on mid-circuit measurements
         if isinstance(op, qml.Projector):
-            state, circuit._shots, state_nan = _postselection_postprocess(
+            state, circuit._shots = _postselection_postprocess(
                 state, is_state_batched, circuit.shots
             )
-            if state_nan:
-                return state, is_state_batched
 
         # 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
diff --git a/pennylane/measurements/counts.py b/pennylane/measurements/counts.py
index 5369af96bb2..e1a8e42c257 100644
--- a/pennylane/measurements/counts.py
+++ b/pennylane/measurements/counts.py
@@ -218,19 +218,6 @@ def process_samples(
         shot_range: Tuple[int] = None,
         bin_size: int = None,
     ):
-        if qml.math.any(qml.math.isnan(samples)) or len(samples) == 0:
-            if self.all_outcomes:
-                if (eigs := self.eigvals()) is not None:
-                    keys = [str(e) for e in eigs]
-                    vals = [0] * len(eigs)
-                else:
-                    bits = len(wire_order)
-                    keys = [format(i, f"{bits}b") for i in range(2**bits)]
-                    vals = [0] * 2**bits
-                return dict(zip(keys, vals))
-
-            return {}
-
         if self.mv is not None:
             samples = qml.sample(wires=self.mv.wires).process_samples(
                 samples, wire_order, shot_range, bin_size
diff --git a/pennylane/measurements/probs.py b/pennylane/measurements/probs.py
index 9ba607ae131..9b5cbf9a4a2 100644
--- a/pennylane/measurements/probs.py
+++ b/pennylane/measurements/probs.py
@@ -159,9 +159,6 @@ def process_samples(
         shot_range: Tuple[int] = None,
         bin_size: int = None,
     ):
-        if qml.math.any(qml.math.isnan(samples)) or len(samples) == 0:
-            return qml.math.asarray([0.0] * 2 ** len(wire_order), dtype="float64")
-
         wire_map = dict(zip(wire_order, range(len(wire_order))))
         mapped_wires = [wire_map[w] for w in self.wires]
         if shot_range is not None:
diff --git a/pennylane/transforms/defer_measurements.py b/pennylane/transforms/defer_measurements.py
index c92c06ecba9..24bb4793b48 100644
--- a/pennylane/transforms/defer_measurements.py
+++ b/pennylane/transforms/defer_measurements.py
@@ -27,7 +27,7 @@
 
 
 @transform
-def defer_measurements(tape: QuantumTape) -> (Sequence[QuantumTape], Callable):
+def defer_measurements(tape: QuantumTape, **kwargs) -> (Sequence[QuantumTape], Callable):
     """Quantum function transform that substitutes operations conditioned on
     measurement outcomes to controlled operations.
 
@@ -137,6 +137,8 @@ def func(x, y):
     if ops_cv or obs_cv:
         raise ValueError("Continuous variable operations and observables are not supported.")
 
+    device = kwargs.get("device", None)
+
     new_operations = []
 
     # Find wires that are reused after measurement
@@ -161,6 +163,9 @@ def func(x, y):
                 set(measured_wires).intersection(op.wires.toset())
             )
 
+    if is_postselecting and device is not None and not isinstance(device, qml.devices.DefaultQubit):
+        raise ValueError(f"Postselection is not supported on the {device} device.")
+
     # Apply controlled operations to store measurement outcomes and replace
     # classically controlled operations
     control_wires = {}
@@ -232,6 +237,19 @@ def null_postprocessing(results):
     return [new_tape], null_postprocessing
 
 
+@defer_measurements.custom_qnode_transform
+def _defer_measurements_qnode(self, qnode, targs, tkwargs):
+    """Custom qnode transform for ``defer_measurements``."""
+    if tkwargs.get("device", None):
+        raise ValueError(
+            "Cannot provide a 'device' value directly to the defer_measurements decorator "
+            "when transforming a QNode."
+        )
+
+    tkwargs.setdefault("device", qnode.device)
+    return self.default_qnode_transform(qnode, targs, tkwargs)
+
+
 def _add_control_gate(op, control_wires):
     """Helper function to add control gates"""
     control = [control_wires[m.id] for m in op.meas_val.measurements]
diff --git a/tests/devices/qubit/test_measure.py b/tests/devices/qubit/test_measure.py
index 9794ab16469..ac4f18653be 100644
--- a/tests/devices/qubit/test_measure.py
+++ b/tests/devices/qubit/test_measure.py
@@ -26,6 +26,7 @@
     csr_dot_products,
     get_measurement_function,
     sum_of_terms_method,
+    state_measure_nan,
 )
 
 
@@ -91,6 +92,12 @@ def test_sum_sum_of_terms_when_backprop(self):
         state = qml.numpy.zeros(2)
         assert get_measurement_function(qml.expval(S), state) is sum_of_terms_method
 
+    def test_measure_nan(self):
+        """Check that the nan measurement function is used if the state vector
+        has nan values."""
+        state = np.array([[0.0, 0.0], [1 / np.sqrt(2), np.NaN]])
+        assert get_measurement_function(qml.expval(qml.PauliZ(0)), state) is state_measure_nan
+
 
 class TestMeasurements:
     @pytest.mark.parametrize(
@@ -235,6 +242,75 @@ def test_sparse_hamiltonian(self):
         expected = np.array([-2, 2, 0])
         assert np.allclose(res, expected)
 
+    @pytest.mark.parametrize(
+        "mp, shape",
+        [
+            (qml.expval(qml.PauliZ(0)), (3,)),
+            (qml.var(qml.PauliZ(0)), (3,)),
+            (qml.probs(wires=1), (3, 2)),
+        ],
+    )
+    def test_nan_measurements(self, mp, shape):
+        """Test that broadcasting with nan values in the state vector creates
+        outcomes with the correct shape."""
+        state = qml.math.full((3, 2, 2), np.NaN, like="numpy")
+        res = state_measure_nan(mp, state, is_state_batched=True)
+
+        assert res.shape == shape
+
+
+class TestNaNMeasurements:
+    """Tests for state vectors containing nan values."""
+
+    @pytest.mark.all_interfaces
+    @pytest.mark.parametrize(
+        "mp",
+        [
+            qml.expval(qml.PauliZ(0)),
+            qml.expval(
+                qml.Hamiltonian(
+                    [1.0, 2.0, 3.0, 4.0],
+                    [qml.PauliZ(0) @ qml.PauliX(1), qml.PauliX(1), qml.PauliZ(1), qml.PauliY(1)],
+                )
+            ),
+            qml.expval(
+                qml.dot(
+                    [1.0, 2.0, 3.0, 4.0],
+                    [qml.PauliZ(0) @ qml.PauliX(1), qml.PauliX(1), qml.PauliZ(1), qml.PauliY(1)],
+                )
+            ),
+            qml.var(qml.PauliZ(0)),
+            qml.purity(0),
+            qml.vn_entropy(0),
+            qml.mutual_info(0, 1),
+        ],
+    )
+    @pytest.mark.parametrize("interface", ["numpy", "autograd", "torch", "tensorflow", "jax"])
+    def test_nan_float_result(self, mp, interface):
+        """Test that the result of circuits with 0 probability postselections is NaN with the
+        expected shape."""
+        state = qml.math.full((2, 2), np.NaN, like=interface)
+        res = measure(mp, state, is_state_batched=False)
+
+        assert qml.math.ndim(res) == 0
+        assert qml.math.isnan(res)
+        assert qml.math.get_interface(res) == interface
+
+    @pytest.mark.all_interfaces
+    @pytest.mark.parametrize(
+        "mp", [qml.probs(wires=0), qml.probs(op=qml.PauliZ(0)), qml.probs(wires=[0, 1])]
+    )
+    @pytest.mark.parametrize("interface", ["numpy", "autograd", "torch", "tensorflow", "jax"])
+    def test_nan_probs(self, mp, interface):
+        """Test that the result of circuits with 0 probability postselections is NaN with the
+        expected shape."""
+        state = qml.math.full((2, 2), np.NaN, like=interface)
+        res = measure(mp, state, is_state_batched=False)
+
+        assert qml.math.shape(res) == (2 ** len(mp.wires),)
+        assert qml.math.all(qml.math.isnan(res))
+        assert qml.math.get_interface(res) == interface
+
 
 class TestSumOfTermsDifferentiability:
     @staticmethod
diff --git a/tests/devices/qubit/test_sampling.py b/tests/devices/qubit/test_sampling.py
index 8ce558d2b0f..a475d988f9e 100644
--- a/tests/devices/qubit/test_sampling.py
+++ b/tests/devices/qubit/test_sampling.py
@@ -490,6 +490,141 @@ def test_measure_with_samples_one_shot_one_wire(self):
         assert result == -1.0
 
 
+class TestInvalidStateSamples:
+    """Tests for state vectors containing nan values or shot vectors with zero shots."""
+
+    @pytest.mark.all_interfaces
+    @pytest.mark.parametrize(
+        "mp",
+        [
+            qml.expval(qml.PauliZ(0)),
+            qml.expval(
+                qml.Hamiltonian(
+                    [1.0, 2.0, 3.0, 4.0],
+                    [qml.PauliZ(0) @ qml.PauliX(1), qml.PauliX(1), qml.PauliZ(1), qml.PauliY(1)],
+                )
+            ),
+            qml.expval(
+                qml.dot(
+                    [1.0, 2.0, 3.0, 4.0],
+                    [qml.PauliZ(0) @ qml.PauliX(1), qml.PauliX(1), qml.PauliZ(1), qml.PauliY(1)],
+                )
+            ),
+            qml.var(qml.PauliZ(0)),
+        ],
+    )
+    @pytest.mark.parametrize("interface", ["numpy", "autograd", "torch", "tensorflow", "jax"])
+    @pytest.mark.parametrize("shots", [10, [10, 10]])
+    def test_nan_float_result(self, mp, interface, shots):
+        """Test that the result of circuits with 0 probability postselections is NaN with the
+        expected shape."""
+        state = qml.math.full((2, 2), np.NaN, like=interface)
+        res = measure_with_samples([mp], state, is_state_batched=False)
+
+        if not isinstance(shots, list):
+            assert isinstance(res, tuple)
+            res = res[0]
+            assert qml.math.ndim(res) == 0
+            assert qml.math.isnan(res)
+            assert qml.math.get_interface(res) == interface
+
+        else:
+            assert isinstance(res, tuple)
+            res = res[0]
+            assert isinstance(res, tuple)
+            assert len(res) == 2
+            for r in res:
+                assert qml.math.ndim(r) == 0
+                assert qml.math.isnan(r)
+                assert qml.math.get_interface(r) == interface
+
+    @pytest.mark.all_interfaces
+    @pytest.mark.parametrize(
+        "mp", [qml.probs(wires=0), qml.probs(op=qml.PauliZ(0)), qml.probs(wires=[0, 1])]
+    )
+    @pytest.mark.parametrize("interface", ["numpy", "autograd", "torch", "tensorflow", "jax"])
+    @pytest.mark.parametrize("shots", [10, [10, 10]])
+    def test_nan_probs(self, mp, interface, shots):
+        """Test that the result of circuits with 0 probability postselections is NaN with the
+        expected shape."""
+        state = qml.math.full((2, 2), np.NaN, like=interface)
+        res = measure_with_samples([mp], state, is_state_batched=False)
+
+        if not isinstance(shots, list):
+            assert isinstance(res, tuple)
+            res = res[0]
+            assert qml.math.shape(res) == (2 ** len(mp.wires),)
+            assert all(nan_value for nan_value in qml.math.isnan(res))
+            assert qml.math.get_interface(res) == interface
+
+        else:
+            assert isinstance(res, tuple)
+            res = res[0]
+            assert isinstance(res, tuple)
+            assert len(res) == 2
+            for r in res:
+                assert qml.math.shape(r) == (2 ** len(mp.wires),)
+                assert all(nan_value for nan_value in qml.math.isnan(r))
+                assert qml.math.get_interface(r) == interface
+
+    @pytest.mark.all_interfaces
+    @pytest.mark.parametrize(
+        "mp", [qml.sample(wires=0), qml.sample(op=qml.PauliZ(0)), qml.sample(wires=[0, 1])]
+    )
+    @pytest.mark.parametrize("interface", ["numpy", "autograd", "torch", "tensorflow", "jax"])
+    @pytest.mark.parametrize("shots", [10, [10, 10]])
+    def test_nan_samples(self, mp, interface, shots):
+        """Test that the result of circuits with 0 probability postselections is NaN with the
+        expected shape."""
+        state = qml.math.full((2, 2), np.NaN, like=interface)
+        res = measure_with_samples([mp], state, is_state_batched=False)
+
+        if not isinstance(shots, list):
+            assert isinstance(res, tuple)
+            res = res[0]
+            assert qml.math.shape(res) == (0,)
+            assert qml.math.get_interface(res) == interface
+
+        else:
+            assert isinstance(res, tuple)
+            res = res[0]
+            assert isinstance(res, tuple)
+            assert len(res) == 2
+            for r in res:
+                assert qml.math.shape(r) == (0,)
+                assert qml.math.get_interface(r) == interface
+
+    @pytest.mark.all_interfaces
+    @pytest.mark.parametrize(
+        "mp",
+        [
+            qml.counts(wires=0),
+            qml.counts(op=qml.PauliZ(0)),
+            qml.counts(wires=[0, 1], all_outcomes=True),
+        ],
+    )
+    @pytest.mark.parametrize("interface", ["numpy", "autograd", "torch", "tensorflow", "jax"])
+    @pytest.mark.parametrize("shots", [10, [10, 10]])
+    def test_nan_counts(self, mp, interface, shots):
+        """Test that the result of circuits with 0 probability postselections is NaN with the
+        expected shape."""
+        state = qml.math.full((2, 2), np.NaN, like=interface)
+        res = measure_with_samples([mp], state, is_state_batched=False)
+
+        if not isinstance(shots, list):
+            assert isinstance(res, tuple)
+            res = res[0]
+            assert res == {}
+
+        else:
+            assert isinstance(res, tuple)
+            res = res[0]
+            assert isinstance(res, tuple)
+            assert len(res) == 2
+            for r in res:
+                assert r == {}
+
+
 class TestBroadcasting:
     """Test that measurements work when the state has a batch dim"""
 
diff --git a/tests/devices/qubit/test_simulate.py b/tests/devices/qubit/test_simulate.py
index 6586e6d76f0..bcddbc212b9 100644
--- a/tests/devices/qubit/test_simulate.py
+++ b/tests/devices/qubit/test_simulate.py
@@ -378,166 +378,14 @@ def test_broadcasting_with_projector(self, shots):
             _ = simulate(tape)
 
     @pytest.mark.all_interfaces
-    @pytest.mark.parametrize(
-        "mp",
-        [
-            qml.expval(qml.PauliZ(0)),
-            qml.var(qml.PauliZ(0)),
-            qml.purity(0),
-            qml.vn_entropy(0),
-            qml.mutual_info(0, 1),
-        ],
-    )
-    @pytest.mark.parametrize("interface", ["numpy", "autograd", "torch", "tensorflow", "jax"])
-    @pytest.mark.parametrize("shots", [None, 10, [10, 10]])
-    def test_nan_float_result(self, mp, interface, shots):
-        """Test that the result of circuits with 0 probability postselections is NaN with the
-        expected shape."""
-        if (
-            isinstance(
-                mp,
-                (
-                    qml.measurements.VnEntropyMP,
-                    qml.measurements.MutualInfoMP,
-                    qml.measurements.PurityMP,
-                ),
-            )
-            and shots is not None
-        ):
-            pytest.skip("QInfo measurements should be done with analytic execution.")
-
-        def f():
-            tape = qml.tape.QuantumScript(
-                [
-                    qml.RX(qml.math.asarray(np.pi, like=interface, dtype="float64"), 0),
-                    qml.CNOT([0, 1]),
-                    qml.Projector([0], wires=0),
-                ],
-                [mp],
-                shots=shots,
-            )
-
-            return simulate(tape)
-
-        res = f()
-
-        if not isinstance(shots, list):
-            assert qml.math.ndim(res) == 0
-            assert qml.math.isnan(res)
-            assert qml.math.get_interface(res) == interface
-
-        else:
-            assert isinstance(res, tuple)
-            assert len(res) == 2
-            for r in res:
-                assert qml.math.ndim(r) == 0
-                assert qml.math.isnan(r)
-                assert qml.math.get_interface(r) == interface
-
-    @pytest.mark.all_interfaces
-    @pytest.mark.parametrize(
-        "mp", [qml.probs(wires=0), qml.probs(op=qml.PauliZ(0)), qml.probs(wires=[0, 1])]
-    )
-    @pytest.mark.parametrize("interface", ["numpy", "autograd", "torch", "tensorflow", "jax"])
-    @pytest.mark.parametrize("shots", [None, 10, [10, 10]])
-    def test_nan_probs(self, mp, interface, shots):
-        """Test that the result of circuits with 0 probability postselections is NaN with the
-        expected shape."""
-
-        def f():
-            tape = qml.tape.QuantumScript(
-                [
-                    qml.RX(qml.math.asarray(np.pi, like=interface, dtype="float64"), 0),
-                    qml.CNOT([0, 1]),
-                    qml.Projector([0], wires=0),
-                ],
-                [mp],
-                shots=shots,
-            )
-            return simulate(tape)
-
-        res = f()
-
-        if not isinstance(shots, list):
-            assert qml.math.shape(res) == (2 ** len(mp.wires),)
-            assert all(nan_value for nan_value in qml.math.isnan(res))
-            assert qml.math.get_interface(res) == interface
-
-        else:
-            assert isinstance(res, tuple)
-            assert len(res) == 2
-            for r in res:
-                assert qml.math.shape(r) == (2 ** len(mp.wires),)
-                assert all(nan_value for nan_value in qml.math.isnan(r))
-                assert qml.math.get_interface(r) == interface
-
-    @pytest.mark.all_interfaces
-    @pytest.mark.parametrize(
-        "mp", [qml.sample(wires=0), qml.sample(op=qml.PauliZ(0)), qml.sample(wires=[0, 1])]
-    )
-    @pytest.mark.parametrize("interface", ["numpy", "autograd", "torch", "tensorflow", "jax"])
-    @pytest.mark.parametrize("shots", [10, [10, 10]])
-    def test_nan_samples(self, mp, interface, shots):
-        """Test that the result of circuits with 0 probability postselections is NaN with the
-        expected shape."""
-
-        def f():
-            tape = qml.tape.QuantumScript(
-                [
-                    qml.RX(qml.math.asarray(np.pi, like=interface, dtype="float64"), 0),
-                    qml.CNOT([0, 1]),
-                    qml.Projector([0], wires=0),
-                ],
-                [mp],
-                shots=shots,
-            )
-            return simulate(tape)
-
-        res = f()
-
-        if not isinstance(shots, list):
-            assert qml.math.shape(res) == (0,)
-            assert qml.math.get_interface(res) == interface
-
-        else:
-            assert isinstance(res, tuple)
-            assert len(res) == 2
-            for r in res:
-                assert qml.math.shape(r) == (0,)
-                assert qml.math.get_interface(r) == interface
-
-    @pytest.mark.all_interfaces
-    @pytest.mark.parametrize(
-        "mp", [qml.counts(wires=0), qml.counts(op=qml.PauliZ(0)), qml.counts()]
-    )
-    @pytest.mark.parametrize("interface", ["numpy", "autograd", "torch", "tensorflow", "jax"])
-    @pytest.mark.parametrize("shots", [10, [10, 10]])
-    def test_nan_counts(self, mp, interface, shots):
-        """Test that the result of circuits with 0 probability postselections is NaN with the
-        expected shape."""
-
-        def f():
-            tape = qml.tape.QuantumScript(
-                [
-                    qml.RX(qml.math.asarray(np.pi, like=interface, dtype="float64"), 0),
-                    qml.CNOT([0, 1]),
-                    qml.Projector([0], wires=0),
-                ],
-                [mp],
-                shots=shots,
-            )
-            return simulate(tape)
-
-        res = f()
-
-        if not isinstance(shots, list):
-            assert res == {}
-
-        else:
-            assert isinstance(res, tuple)
-            assert len(res) == 2
-            for r in res:
-                assert r == {}
+    @pytest.mark.parametrize("interface", ["numpy", "torch", "jax", "tensorflow", "autograd"])
+    def test_nan_state(self, interface):
+        """Test that a state with nan values is returned if the probability of a postselection state
+        is 0."""
+        tape = qml.tape.QuantumScript([qml.PauliX(0), qml.Projector([0], 0)])
+
+        res, _ = get_final_state(tape, interface=interface)
+        assert qml.math.all(qml.math.isnan(res))
 
 
 class TestDebugger:
diff --git a/tests/transforms/test_defer_measurements.py b/tests/transforms/test_defer_measurements.py
index 5d118f818e2..26fb3e722d6 100644
--- a/tests/transforms/test_defer_measurements.py
+++ b/tests/transforms/test_defer_measurements.py
@@ -24,6 +24,47 @@
 from pennylane.devices import DefaultQubit
 
 
+def test_broadcasted_postselection(mocker):
+    """Test that broadcast_expand is used iff broadcasting with postselection."""
+    spy = mocker.spy(qml.transforms, "broadcast_expand")
+
+    # Broadcasting with postselection
+    tape1 = qml.tape.QuantumScript(
+        [qml.RX([0.1, 0.2], 0), MidMeasureMP(0, postselect=1), qml.CNOT([0, 1])],
+        [qml.probs(wires=[0])],
+    )
+    _, _ = qml.defer_measurements(tape1)
+
+    assert spy.call_count == 1
+
+    # Broadcasting without postselection
+    tape2 = qml.tape.QuantumScript(
+        [qml.RX([0.1, 0.2], 0), MidMeasureMP(0), qml.CNOT([0, 1])],
+        [qml.probs(wires=[0])],
+    )
+    _, _ = qml.defer_measurements(tape2)
+
+    assert spy.call_count == 1
+
+    # Postselection without broadcasting
+    tape3 = qml.tape.QuantumScript(
+        [qml.RX(0.1, 0), MidMeasureMP(0, postselect=1), qml.CNOT([0, 1])],
+        [qml.probs(wires=[0])],
+    )
+    _, _ = qml.defer_measurements(tape3)
+
+    assert spy.call_count == 1
+
+    # No postselection, no broadcasting
+    tape4 = qml.tape.QuantumScript(
+        [qml.RX(0.1, 0), MidMeasureMP(0), qml.CNOT([0, 1])],
+        [qml.probs(wires=[0])],
+    )
+    _, _ = qml.defer_measurements(tape4)
+
+    assert spy.call_count == 1
+
+
 class TestQNode:
     """Test that the transform integrates well with QNodes."""
 

From b6279f0f377e4b0d702a3121632ea98c4c802c5a Mon Sep 17 00:00:00 2001
From: Mudit Pandey <mudit.pandey@xanadu.ai>
Date: Mon, 16 Oct 2023 11:06:23 -0400
Subject: [PATCH 113/131] Lots of changes

---
 pennylane/_device.py                        |   3 +
 pennylane/devices/qubit/measure.py          |   2 -
 pennylane/devices/qubit/sampling.py         |   8 +-
 pennylane/devices/qubit/simulate.py         |   7 +-
 tests/devices/qubit/test_measure.py         |  99 ++++++---
 tests/devices/qubit/test_sampling.py        | 234 +++++++++++++++++---
 tests/transforms/test_defer_measurements.py |  15 ++
 7 files changed, 298 insertions(+), 70 deletions(-)

diff --git a/pennylane/_device.py b/pennylane/_device.py
index c24400c2ec2..93e7af2488f 100644
--- a/pennylane/_device.py
+++ b/pennylane/_device.py
@@ -971,6 +971,9 @@ def check_validity(self, queue, observables):
                     "simulate the application of mid-circuit measurements on this device."
                 )
 
+            if isinstance(o, qml.Projector):
+                raise ValueError(f"Postselection is not supported on the {self.name} device.")
+
             if not self.stopping_condition(o):
                 raise DeviceError(
                     f"Gate {operation_name} not supported on device {self.short_name}"
diff --git a/pennylane/devices/qubit/measure.py b/pennylane/devices/qubit/measure.py
index 5d04f713946..9b72def187d 100644
--- a/pennylane/devices/qubit/measure.py
+++ b/pennylane/devices/qubit/measure.py
@@ -136,7 +136,6 @@ def sum_of_terms_method(
     )
 
 
-# pylint: disable=too-many-return-statements
 def get_measurement_function(
     measurementprocess: MeasurementProcess, state: TensorLike
 ) -> Callable[[MeasurementProcess, TensorLike], TensorLike]:
@@ -150,7 +149,6 @@ def get_measurement_function(
     Returns:
         Callable: function that returns the measurement result
     """
-
     if isinstance(measurementprocess, StateMeasurement):
         if isinstance(measurementprocess.mv, MeasurementValue):
             return state_diagonalizing_gates
diff --git a/pennylane/devices/qubit/sampling.py b/pennylane/devices/qubit/sampling.py
index cf009aa3dd5..b029d1a5fcf 100644
--- a/pennylane/devices/qubit/sampling.py
+++ b/pennylane/devices/qubit/sampling.py
@@ -274,7 +274,9 @@ def _process_single_shot(samples):
                     rng=rng,
                     prng_key=prng_key,
                 )
-            except ValueError:
+            except ValueError as e:
+                if str(e) != "probabilities contain NaN":
+                    raise e
                 samples = qml.math.full((s, len(wires)), np.NaN)
 
             processed_samples.append(_process_single_shot(samples))
@@ -290,7 +292,9 @@ def _process_single_shot(samples):
             rng=rng,
             prng_key=prng_key,
         )
-    except ValueError:
+    except ValueError as e:
+        if str(e) != "probabilities contain NaN":
+            raise e
         samples = qml.math.full((shots.total_shots, len(wires)), np.NaN)
 
     return _process_single_shot(samples)
diff --git a/pennylane/devices/qubit/simulate.py b/pennylane/devices/qubit/simulate.py
index 2da17f99ae3..1435ab913ce 100644
--- a/pennylane/devices/qubit/simulate.py
+++ b/pennylane/devices/qubit/simulate.py
@@ -135,8 +135,11 @@ def _postselection_postprocess(state, is_state_batched, shots):
             "postselection is used."
         )
 
-    # if qml.math.is_abstract(state):
-    norm = qml.math.floor(qml.math.real(qml.math.norm(state)) * 1e8) * 1e-8
+    # The floor function is being used here so that a norm very close to zero becomes exactly
+    # equal to zero so that the state can become invalid. This way, execution can continue, and
+    # bad postselection gives results that are invalid rather than results that look valid but
+    # are incorrect.
+    norm = qml.math.floor(qml.math.real(qml.math.norm(state)) * 1e15) * 1e-15
 
     if shots:
         # Clip the number of shots using a binomial distribution using the probability of
diff --git a/tests/devices/qubit/test_measure.py b/tests/devices/qubit/test_measure.py
index ac4f18653be..3f9573d1831 100644
--- a/tests/devices/qubit/test_measure.py
+++ b/tests/devices/qubit/test_measure.py
@@ -26,7 +26,6 @@
     csr_dot_products,
     get_measurement_function,
     sum_of_terms_method,
-    state_measure_nan,
 )
 
 
@@ -92,12 +91,6 @@ def test_sum_sum_of_terms_when_backprop(self):
         state = qml.numpy.zeros(2)
         assert get_measurement_function(qml.expval(S), state) is sum_of_terms_method
 
-    def test_measure_nan(self):
-        """Check that the nan measurement function is used if the state vector
-        has nan values."""
-        state = np.array([[0.0, 0.0], [1 / np.sqrt(2), np.NaN]])
-        assert get_measurement_function(qml.expval(qml.PauliZ(0)), state) is state_measure_nan
-
 
 class TestMeasurements:
     @pytest.mark.parametrize(
@@ -242,22 +235,6 @@ def test_sparse_hamiltonian(self):
         expected = np.array([-2, 2, 0])
         assert np.allclose(res, expected)
 
-    @pytest.mark.parametrize(
-        "mp, shape",
-        [
-            (qml.expval(qml.PauliZ(0)), (3,)),
-            (qml.var(qml.PauliZ(0)), (3,)),
-            (qml.probs(wires=1), (3, 2)),
-        ],
-    )
-    def test_nan_measurements(self, mp, shape):
-        """Test that broadcasting with nan values in the state vector creates
-        outcomes with the correct shape."""
-        state = qml.math.full((3, 2, 2), np.NaN, like="numpy")
-        res = state_measure_nan(mp, state, is_state_batched=True)
-
-        assert res.shape == shape
-
 
 class TestNaNMeasurements:
     """Tests for state vectors containing nan values."""
@@ -280,12 +257,15 @@ class TestNaNMeasurements:
                 )
             ),
             qml.var(qml.PauliZ(0)),
-            qml.purity(0),
-            qml.vn_entropy(0),
-            qml.mutual_info(0, 1),
+            qml.var(
+                qml.dot(
+                    [1.0, 2.0, 3.0, 4.0],
+                    [qml.PauliZ(0) @ qml.PauliX(1), qml.PauliX(1), qml.PauliZ(1), qml.PauliY(1)],
+                )
+            ),
         ],
     )
-    @pytest.mark.parametrize("interface", ["numpy", "autograd", "torch", "tensorflow", "jax"])
+    @pytest.mark.parametrize("interface", ["numpy", "autograd", "torch", "tensorflow"])
     def test_nan_float_result(self, mp, interface):
         """Test that the result of circuits with 0 probability postselections is NaN with the
         expected shape."""
@@ -296,11 +276,54 @@ def test_nan_float_result(self, mp, interface):
         assert qml.math.isnan(res)
         assert qml.math.get_interface(res) == interface
 
+    @pytest.mark.jax
+    @pytest.mark.parametrize(
+        "mp",
+        [
+            qml.expval(qml.PauliZ(0)),
+            qml.expval(
+                qml.Hamiltonian(
+                    [1.0, 2.0, 3.0, 4.0],
+                    [qml.PauliZ(0) @ qml.PauliX(1), qml.PauliX(1), qml.PauliZ(1), qml.PauliY(1)],
+                )
+            ),
+            qml.expval(
+                qml.dot(
+                    [1.0, 2.0, 3.0, 4.0],
+                    [qml.PauliZ(0) @ qml.PauliX(1), qml.PauliX(1), qml.PauliZ(1), qml.PauliY(1)],
+                )
+            ),
+            qml.var(qml.PauliZ(0)),
+            qml.var(
+                qml.dot(
+                    [1.0, 2.0, 3.0, 4.0],
+                    [qml.PauliZ(0) @ qml.PauliX(1), qml.PauliX(1), qml.PauliZ(1), qml.PauliY(1)],
+                )
+            ),
+        ],
+    )
+    @pytest.mark.parametrize("use_jit", [True, False])
+    def test_nan_float_result_jax(self, mp, use_jit):
+        """Test that the result of circuits with 0 probability postselections is NaN with the
+        expected shape."""
+        state = qml.math.full((2, 2), np.NaN, like="jax")
+        if use_jit:
+            import jax
+
+            res = jax.jit(measure, static_argnums=[0, 2])(mp, state, is_state_batched=False)
+        else:
+            res = measure(mp, state, is_state_batched=False)
+
+        assert qml.math.ndim(res) == 0
+
+        assert qml.math.isnan(res)
+        assert qml.math.get_interface(res) == "jax"
+
     @pytest.mark.all_interfaces
     @pytest.mark.parametrize(
         "mp", [qml.probs(wires=0), qml.probs(op=qml.PauliZ(0)), qml.probs(wires=[0, 1])]
     )
-    @pytest.mark.parametrize("interface", ["numpy", "autograd", "torch", "tensorflow", "jax"])
+    @pytest.mark.parametrize("interface", ["numpy", "autograd", "torch", "tensorflow"])
     def test_nan_probs(self, mp, interface):
         """Test that the result of circuits with 0 probability postselections is NaN with the
         expected shape."""
@@ -311,6 +334,26 @@ def test_nan_probs(self, mp, interface):
         assert qml.math.all(qml.math.isnan(res))
         assert qml.math.get_interface(res) == interface
 
+    @pytest.mark.jax
+    @pytest.mark.parametrize(
+        "mp", [qml.probs(wires=0), qml.probs(op=qml.PauliZ(0)), qml.probs(wires=[0, 1])]
+    )
+    @pytest.mark.parametrize("use_jit", [True, False])
+    def test_nan_probs_jax(self, mp, use_jit):
+        """Test that the result of circuits with 0 probability postselections is NaN with the
+        expected shape."""
+        state = qml.math.full((2, 2), np.NaN, like="jax")
+        if use_jit:
+            import jax
+
+            res = jax.jit(measure, static_argnums=[0, 2])(mp, state, is_state_batched=False)
+        else:
+            res = measure(mp, state, is_state_batched=False)
+
+        assert qml.math.shape(res) == (2 ** len(mp.wires),)
+        assert qml.math.all(qml.math.isnan(res))
+        assert qml.math.get_interface(res) == "jax"
+
 
 class TestSumOfTermsDifferentiability:
     @staticmethod
diff --git a/tests/devices/qubit/test_sampling.py b/tests/devices/qubit/test_sampling.py
index a475d988f9e..268759aed45 100644
--- a/tests/devices/qubit/test_sampling.py
+++ b/tests/devices/qubit/test_sampling.py
@@ -22,6 +22,7 @@
 from pennylane.devices.qubit import simulate
 from pennylane.devices.qubit import sample_state, measure_with_samples
 from pennylane.devices.qubit.sampling import _sample_state_jax
+from pennylane.measurements import Shots
 
 two_qubit_state = np.array([[0, 1j], [-1, 0]], dtype=np.complex128) / np.sqrt(2)
 APPROX_ATOL = 0.01
@@ -513,20 +514,82 @@ class TestInvalidStateSamples:
             qml.var(qml.PauliZ(0)),
         ],
     )
-    @pytest.mark.parametrize("interface", ["numpy", "autograd", "torch", "tensorflow", "jax"])
-    @pytest.mark.parametrize("shots", [10, [10, 10]])
+    @pytest.mark.parametrize("interface", ["numpy", "autograd", "torch", "tensorflow"])
+    @pytest.mark.parametrize("shots", [0, [0, 0]])
     def test_nan_float_result(self, mp, interface, shots):
         """Test that the result of circuits with 0 probability postselections is NaN with the
         expected shape."""
         state = qml.math.full((2, 2), np.NaN, like=interface)
-        res = measure_with_samples([mp], state, is_state_batched=False)
+        res = measure_with_samples((mp,), state, Shots(shots), is_state_batched=False)
+
+        if not isinstance(shots, list):
+            assert qml.math.ndim(res) == 0
+            assert qml.math.isnan(res)
+            assert qml.math.get_interface(res) == interface
+
+        else:
+            assert isinstance(res, tuple)
+            assert len(res) == 2
+            for r in res:
+                assert isinstance(r, tuple)
+                r = r[0]
+                assert qml.math.ndim(r) == 0
+                assert qml.math.isnan(r)
+                assert qml.math.get_interface(r) == interface
+
+    @pytest.mark.jax
+    @pytest.mark.parametrize(
+        "mp",
+        [
+            qml.expval(qml.PauliZ(0)),
+            qml.expval(
+                qml.Hamiltonian(
+                    [1.0, 2.0, 3.0, 4.0],
+                    [qml.PauliZ(0) @ qml.PauliX(1), qml.PauliX(1), qml.PauliZ(1), qml.PauliY(1)],
+                )
+            ),
+            qml.expval(
+                qml.dot(
+                    [1.0, 2.0, 3.0, 4.0],
+                    [qml.PauliZ(0) @ qml.PauliX(1), qml.PauliX(1), qml.PauliZ(1), qml.PauliY(1)],
+                )
+            ),
+            qml.var(qml.PauliZ(0)),
+            qml.var(
+                qml.Hamiltonian(
+                    [1.0, 2.0, 3.0, 4.0],
+                    [qml.PauliZ(0) @ qml.PauliX(1), qml.PauliX(1), qml.PauliZ(1), qml.PauliY(1)],
+                )
+            ),
+            qml.var(
+                qml.dot(
+                    [1.0, 2.0, 3.0, 4.0],
+                    [qml.PauliZ(0) @ qml.PauliX(1), qml.PauliX(1), qml.PauliZ(1), qml.PauliY(1)],
+                )
+            ),
+        ],
+    )
+    @pytest.mark.parametrize("use_jit", [True, False])
+    @pytest.mark.parametrize("shots", [0, [0, 0]])
+    def test_nan_float_result_jax(self, mp, use_jit, shots):
+        """Test that the result of circuits with 0 probability postselections is NaN with the
+        expected shape."""
+        state = qml.math.full((2, 2), np.NaN, like="jax")
+        if use_jit:
+            import jax
+
+            res = jax.jit(measure_with_samples, static_argnums=[0, 2, 3])(
+                (mp,), state, Shots(shots), is_state_batched=False
+            )
+        else:
+            res = measure_with_samples((mp,), state, Shots(shots), is_state_batched=False)
 
         if not isinstance(shots, list):
             assert isinstance(res, tuple)
             res = res[0]
             assert qml.math.ndim(res) == 0
             assert qml.math.isnan(res)
-            assert qml.math.get_interface(res) == interface
+            assert qml.math.get_interface(res) == "jax"
 
         else:
             assert isinstance(res, tuple)
@@ -536,25 +599,24 @@ def test_nan_float_result(self, mp, interface, shots):
             for r in res:
                 assert qml.math.ndim(r) == 0
                 assert qml.math.isnan(r)
-                assert qml.math.get_interface(r) == interface
+                assert qml.math.get_interface(r) == "jax"
 
     @pytest.mark.all_interfaces
     @pytest.mark.parametrize(
-        "mp", [qml.probs(wires=0), qml.probs(op=qml.PauliZ(0)), qml.probs(wires=[0, 1])]
+        "mp", [qml.sample(wires=0), qml.sample(op=qml.PauliZ(0)), qml.sample(wires=[0, 1])]
     )
-    @pytest.mark.parametrize("interface", ["numpy", "autograd", "torch", "tensorflow", "jax"])
-    @pytest.mark.parametrize("shots", [10, [10, 10]])
-    def test_nan_probs(self, mp, interface, shots):
+    @pytest.mark.parametrize("interface", ["numpy", "autograd", "torch", "tensorflow"])
+    @pytest.mark.parametrize("shots", [0, [0, 0]])
+    def test_nan_samples(self, mp, interface, shots):
         """Test that the result of circuits with 0 probability postselections is NaN with the
         expected shape."""
         state = qml.math.full((2, 2), np.NaN, like=interface)
-        res = measure_with_samples([mp], state, is_state_batched=False)
+        res = measure_with_samples((mp,), state, Shots(shots), is_state_batched=False)
 
         if not isinstance(shots, list):
             assert isinstance(res, tuple)
             res = res[0]
-            assert qml.math.shape(res) == (2 ** len(mp.wires),)
-            assert all(nan_value for nan_value in qml.math.isnan(res))
+            assert qml.math.shape(res) == (shots, 2)
             assert qml.math.get_interface(res) == interface
 
         else:
@@ -562,22 +624,54 @@ def test_nan_probs(self, mp, interface, shots):
             res = res[0]
             assert isinstance(res, tuple)
             assert len(res) == 2
-            for r in res:
-                assert qml.math.shape(r) == (2 ** len(mp.wires),)
-                assert all(nan_value for nan_value in qml.math.isnan(r))
+            for i, r in enumerate(res):
+                assert qml.math.shape(r) == (shots[i], 2)
                 assert qml.math.get_interface(r) == interface
 
-    @pytest.mark.all_interfaces
+    @pytest.mark.jax
     @pytest.mark.parametrize(
         "mp", [qml.sample(wires=0), qml.sample(op=qml.PauliZ(0)), qml.sample(wires=[0, 1])]
     )
-    @pytest.mark.parametrize("interface", ["numpy", "autograd", "torch", "tensorflow", "jax"])
-    @pytest.mark.parametrize("shots", [10, [10, 10]])
-    def test_nan_samples(self, mp, interface, shots):
+    @pytest.mark.parametrize("use_jit", [True, False])
+    @pytest.mark.parametrize("shots", [0, [0, 0]])
+    def test_nan_samples_jax(self, mp, use_jit, shots):
         """Test that the result of circuits with 0 probability postselections is NaN with the
         expected shape."""
+        state = qml.math.full((2, 2), np.NaN, like="jax")
+        if use_jit:
+            import jax
+
+            res = jax.jit(measure_with_samples, static_argnums=[0, 2, 3])(
+                (mp,), state, Shots(shots), is_state_batched=False
+            )
+        else:
+            res = measure_with_samples((mp,), state, Shots(shots), is_state_batched=False)
+
+        if not isinstance(shots, list):
+            assert isinstance(res, tuple)
+            res = res[0]
+            assert qml.math.shape(res) == (shots, 2)
+            assert qml.math.get_interface(res) == "jax"
+
+        else:
+            assert isinstance(res, tuple)
+            res = res[0]
+            assert isinstance(res, tuple)
+            assert len(res) == 2
+            for i, r in enumerate(res):
+                assert qml.math.shape(r) == (shots[i], 2)
+                assert qml.math.get_interface(r) == "jax"
+
+    @pytest.mark.all_interfaces
+    @pytest.mark.parametrize("interface", ["numpy", "autograd", "torch", "tensorflow"])
+    @pytest.mark.parametrize("shots", [0, [0, 0]])
+    def test_nan_classical_shadows(self, interface, shots):
+        """Test that classical_shadows returns an empty array when the state has
+        NaN values"""
         state = qml.math.full((2, 2), np.NaN, like=interface)
-        res = measure_with_samples([mp], state, is_state_batched=False)
+        res = measure_with_samples(
+            (qml.classical_shadow([0]),), state, Shots(shots), is_state_batched=False
+        )
 
         if not isinstance(shots, list):
             assert isinstance(res, tuple)
@@ -594,27 +688,93 @@ def test_nan_samples(self, mp, interface, shots):
                 assert qml.math.shape(r) == (0,)
                 assert qml.math.get_interface(r) == interface
 
+    @pytest.mark.jax
+    @pytest.mark.parametrize("use_jit", [True, False])
+    @pytest.mark.parametrize("shots", [0, [0, 0]])
+    def test_nan_classical_shadows_jax(self, use_jit, shots):
+        """Test that classical_shadows returns an empty array when the state has
+        NaN values"""
+        state = qml.math.full((2, 2), np.NaN, like="jax")
+        if use_jit:
+            import jax
+
+            res = jax.jit(measure_with_samples, static_argnums=[0, 2, 3])(
+                (qml.classical_shadow([0]),), state, Shots(shots), is_state_batched=False
+            )
+        else:
+            res = measure_with_samples(
+                (qml.classical_shadow([0]),), state, Shots(shots), is_state_batched=False
+            )
+
+        if not isinstance(shots, list):
+            assert isinstance(res, tuple)
+            res = res[0]
+            assert qml.math.shape(res) == (0,)
+            assert qml.math.get_interface(res) == "jax"
+
+        else:
+            assert isinstance(res, tuple)
+            res = res[0]
+            assert isinstance(res, tuple)
+            assert len(res) == 2
+            for r in res:
+                assert qml.math.shape(r) == (0,)
+                assert qml.math.get_interface(r) == "jax"
+
     @pytest.mark.all_interfaces
-    @pytest.mark.parametrize(
-        "mp",
-        [
-            qml.counts(wires=0),
-            qml.counts(op=qml.PauliZ(0)),
-            qml.counts(wires=[0, 1], all_outcomes=True),
-        ],
-    )
-    @pytest.mark.parametrize("interface", ["numpy", "autograd", "torch", "tensorflow", "jax"])
-    @pytest.mark.parametrize("shots", [10, [10, 10]])
-    def test_nan_counts(self, mp, interface, shots):
-        """Test that the result of circuits with 0 probability postselections is NaN with the
-        expected shape."""
+    @pytest.mark.parametrize("H", [qml.PauliZ(0), [qml.PauliZ(0), qml.PauliX(1)]])
+    @pytest.mark.parametrize("interface", ["numpy", "autograd", "torch", "tensorflow"])
+    @pytest.mark.parametrize("shots", [0, [0, 0]])
+    def test_nan_shadow_expval(self, H, interface, shots):
+        """Test that shadow_expval returns an empty array when the state has
+        NaN values"""
         state = qml.math.full((2, 2), np.NaN, like=interface)
-        res = measure_with_samples([mp], state, is_state_batched=False)
+        res = measure_with_samples(
+            (qml.shadow_expval(H),), state, Shots(shots), is_state_batched=False
+        )
+
+        if not isinstance(shots, list):
+            assert isinstance(res, tuple)
+            res = res[0]
+            assert qml.math.shape(res) == qml.math.shape(H)
+            assert qml.math.all(qml.math.isnan(res))
+            assert qml.math.get_interface(res) == interface
+
+        else:
+            assert isinstance(res, tuple)
+            res = res[0]
+            assert isinstance(res, tuple)
+            assert len(res) == 2
+            for r in res:
+                assert qml.math.shape(r) == qml.math.shape(H)
+                assert qml.math.all(qml.math.isnan(r))
+                assert qml.math.get_interface(r) == interface
+
+    @pytest.mark.jax
+    @pytest.mark.parametrize("H", [qml.PauliZ(0), [qml.PauliZ(0), qml.PauliX(1)]])
+    @pytest.mark.parametrize("use_jit", [True, False])
+    @pytest.mark.parametrize("shots", [0, [0, 0]])
+    def test_nan_shadow_expval_jax(self, H, use_jit, shots):
+        """Test that shadow_expval returns an empty array when the state has
+        NaN values"""
+        state = qml.math.full((2, 2), np.NaN, like="jax")
+        if use_jit:
+            import jax
+
+            res = jax.jit(measure_with_samples, static_argnums=[0, 2, 3])(
+                (qml.shadow_expval(H),), state, Shots(shots), is_state_batched=False
+            )
+        else:
+            res = measure_with_samples(
+                (qml.shadow_expval(H),), state, Shots(shots), is_state_batched=False
+            )
 
         if not isinstance(shots, list):
             assert isinstance(res, tuple)
             res = res[0]
-            assert res == {}
+            assert qml.math.shape(res) == qml.math.shape(H)
+            assert qml.math.all(qml.math.isnan(res))
+            assert qml.math.get_interface(res) == "jax"
 
         else:
             assert isinstance(res, tuple)
@@ -622,7 +782,9 @@ def test_nan_counts(self, mp, interface, shots):
             assert isinstance(res, tuple)
             assert len(res) == 2
             for r in res:
-                assert r == {}
+                assert qml.math.shape(r) == qml.math.shape(H)
+                assert qml.math.all(qml.math.isnan(r))
+                assert qml.math.get_interface(r) == "jax"
 
 
 class TestBroadcasting:
diff --git a/tests/transforms/test_defer_measurements.py b/tests/transforms/test_defer_measurements.py
index 26fb3e722d6..c4b55eedd6d 100644
--- a/tests/transforms/test_defer_measurements.py
+++ b/tests/transforms/test_defer_measurements.py
@@ -65,6 +65,21 @@ def test_broadcasted_postselection(mocker):
     assert spy.call_count == 1
 
 
+def test_postselection_error_with_wrong_device():
+    """Test that an error is raised when postselection is used with a device
+    other than `default.qubit`."""
+    dev = qml.device("default.mixed", wires=2)
+
+    @qml.defer_measurements
+    @qml.qnode(dev)
+    def circ():
+        qml.measure(0, postselect=1)
+        return qml.probs(wires=[0])
+
+    with pytest.raises(ValueError, match="Postselection is not supported"):
+        _ = circ()
+
+
 class TestQNode:
     """Test that the transform integrates well with QNodes."""
 

From ddb00d0acca3c8545939526211b99ca221c008b7 Mon Sep 17 00:00:00 2001
From: Mudit Pandey <mudit.pandey@xanadu.ai>
Date: Tue, 17 Oct 2023 12:51:04 -0400
Subject: [PATCH 114/131] Updated docs; added sampling tests

---
 doc/introduction/measurements.rst          |  95 +++++---
 pennylane/devices/qubit/sampling.py        |   8 +-
 pennylane/qnode.py                         |   2 +-
 pennylane/transforms/defer_measurements.py |   7 -
 tests/devices/qubit/test_sampling.py       | 261 ++++++---------------
 5 files changed, 135 insertions(+), 238 deletions(-)

diff --git a/doc/introduction/measurements.rst b/doc/introduction/measurements.rst
index 9945273047b..e9fccdcb7b0 100644
--- a/doc/introduction/measurements.rst
+++ b/doc/introduction/measurements.rst
@@ -272,6 +272,15 @@ PennyLane implements the deferred measurement principle to transform
 conditional operations with the :func:`~.defer_measurements` quantum
 function transform.
 
+The deferred measurement principle provides a natural method to simulate the
+application of mid-circuit measurements and conditional operations in a
+differentiable and device-independent way. Performing true mid-circuit
+measurements and conditional operations is dependent on the
+quantum hardware and PennyLane device capabilities.
+
+For more examples on applying quantum functions conditionally, refer to the
+:func:`~.pennylane.cond` transform.
+
 .. code-block:: python
 
     transformed_qfunc = qml.transforms.defer_measurements(my_quantum_function)
@@ -329,9 +338,36 @@ Executing this QNode:
 >>> func()
 tensor([0., 1.], requires_grad=True)
 
-PennyLane also supports postselecting on mid-circuit measurement outcomes such that only states matching the
-postselection are considered in the remaining execution by specifying the ``postselect`` keyword argument of
-:func:`~.pennylane.measure`:
+Statistics can also be collected on mid-circuit measurements along with terminal measurement statistics.
+Currently, ``qml.probs``, ``qml.sample``, ``qml.expval``, ``qml.var``, and ``qml.counts`` are supported,
+and can be requested along with other measurements. The devices that currently support collecting such
+statistics are ``"default.qubit"``, ``"default.mixed"``, and ``"default.qubit.legacy"``.
+
+.. code-block:: python3
+
+    dev = qml.device("default.qubit", wires=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.probs(wires=1), qml.probs(op=m0)
+
+Executing this QNode:
+
+>>> func(np.pi / 2, np.pi / 4)
+(tensor([0.9267767, 0.0732233], requires_grad=True),
+ tensor([0.5, 0.5], requires_grad=True))
+
+Currently, statistics can only be collected for single mid-circuit measurement values. Moreover, any
+measurement values manipulated using boolean or arithmetic operators cannot be used. These can lead to
+unexpected/incorrect behaviour.
+
+PennyLane also supports postselecting on mid-circuit measurement outcomes by specifying the ``postselect``
+keyword argument of :func:`~.pennylane.measure`. Postselection discards outcomes that do not meet the
+criteria provided by the ``postselect`` argument. For example, specifying ``postselect=1`` on wire 0 would
+be equivalent to projecting the state vector onto the :math:`|1\rangle` state on wire 0:
 
 .. code-block:: python3
 
@@ -353,8 +389,8 @@ array([1, 1, 1, 1, 1, 1, 1])
 Note that only 7 samples are returned. This is because samples that do not meet the postselection criteria are
 thrown away.
 
-If postselection is requested on a state with zero probability of being measured, ``NaN`` will be returned regardless
-of what and how many measurements are made:
+If postselection is requested on a state with zero probability of being measured, the result may contain ``NaN``
+or ``Inf`` values:
 
 .. code-block:: python3
 
@@ -368,48 +404,41 @@ of what and how many measurements are made:
         return qml.probs(wires=1)
 
 >>> func(0.0)
-nan
-
-.. note::
+tensor([nan, nan], requires_grad=True)
 
-    Currently, postselection support is only available on ``"default.qubit"``. Requesting postselection on other
-    devices will not do anything.
-
-Statistics can also be collected on mid-circuit measurements along with terminal measurement statistics.
-Currently, ``qml.probs``, ``qml.sample``, ``qml.expval``, ``qml.var``, and ``qml.counts`` are supported,
-and can be requested along with other measurements. The devices that currently support collecting such
-statistics are ``"default.qubit"``, ``"default.mixed"``, and ``"default.qubit.legacy"``.
+In the case of ``qml.sample``, an empty array will be returned:
 
 .. code-block:: python3
 
     dev = qml.device("default.qubit", wires=2)
 
     @qml.qnode(dev)
-    def func(x, y):
+    def func(x):
         qml.RX(x, wires=0)
-        m0 = qml.measure(0)
-        qml.cond(m0, qml.RY)(y, wires=1)
-        return qml.probs(wires=1), qml.probs(op=m0)
+        m0 = qml.measure(0, postselect=1)
+        qml.cond(m0, qml.PauliX)(wires=1)
+        return qml.sample()
 
-Executing this QNode:
+>>> func(0.0, shots=[10, 10])
+(array([], dtype=float64), array([], dtype=float64))
 
->>> func(np.pi / 2, np.pi / 4)
-(tensor([0.9267767, 0.0732233], requires_grad=True),
- tensor([0.5, 0.5], requires_grad=True))
+.. note::
 
-Currently, statistics can only be collected for single mid-circuit measurement values. Moreover, any
-measurement values manipulated using boolean or arithmetic operators cannot be used. These can lead to
-unexpected/incorrect behaviour.
+    Currently, postselection support is only available on ``"default.qubit"``. Using postselection
+    on other devices will raise an error.
 
-The deferred measurement principle provides a natural method to simulate the
-application of mid-circuit measurements and conditional operations in a
-differentiable and device-independent way. Performing true mid-circuit
-measurements and conditional operations is dependent on the
-quantum hardware and PennyLane device capabilities.
+.. warning::
 
-For more examples on applying quantum functions conditionally, refer to the
-:func:`~.pennylane.cond` transform.
+    All measurements are supported when using postselection. However, postselection on a zero probability
+    state can cause some measurements to break.
+
+    With finite shots, one must be careful when measuring ``qml.probs`` or ``qml.counts``, as these
+    measurements will raise errors if there are no valid samples after postselection. This will occur
+    with postselection states that have zero or close to zero probability.
 
+    With analytic execution, ``qml.mutual_info`` will raise errors when using any interfaces except
+    ``jax``, and ``qml.vn_entropy`` will raise an error with the ``tensorflow`` interface when the
+    postselection state has zero probability.
 
 Changing the number of shots
 ----------------------------
diff --git a/pennylane/devices/qubit/sampling.py b/pennylane/devices/qubit/sampling.py
index 8c3f48a80df..856d0cf592b 100644
--- a/pennylane/devices/qubit/sampling.py
+++ b/pennylane/devices/qubit/sampling.py
@@ -277,7 +277,7 @@ def _process_single_shot(samples):
             except ValueError as e:
                 if str(e) != "probabilities contain NaN":
                     raise e
-                samples = qml.math.full((s, len(wires)), np.NaN)
+                samples = qml.math.full((s, len(wires)), 0)
 
             processed_samples.append(_process_single_shot(samples))
 
@@ -295,7 +295,7 @@ def _process_single_shot(samples):
     except ValueError as e:
         if str(e) != "probabilities contain NaN":
             raise e
-        samples = qml.math.full((shots.total_shots, len(wires)), np.NaN)
+        samples = qml.math.full((shots.total_shots, len(wires)), 0)
 
     return _process_single_shot(samples)
 
@@ -363,7 +363,7 @@ def _sum_for_single_shot(s):
         )
         return sum(c * res for c, res in zip(mp.obs.terms()[0], results))
 
-    unsqueezed_results = tuple(_sum_for_single_shot(Shots(s)) for s in shots)
+    unsqueezed_results = tuple(_sum_for_single_shot(type(shots)(s)) for s in shots)
     return [unsqueezed_results] if shots.has_partitioned_shots else [unsqueezed_results[0]]
 
 
@@ -391,7 +391,7 @@ def _sum_for_single_shot(s):
         )
         return sum(results)
 
-    unsqueezed_results = tuple(_sum_for_single_shot(Shots(s)) for s in shots)
+    unsqueezed_results = tuple(_sum_for_single_shot(type(shots)(s)) for s in shots)
     return [unsqueezed_results] if shots.has_partitioned_shots else [unsqueezed_results[0]]
 
 
diff --git a/pennylane/qnode.py b/pennylane/qnode.py
index de657e8ed95..3676aa72107 100644
--- a/pennylane/qnode.py
+++ b/pennylane/qnode.py
@@ -906,7 +906,7 @@ def construct(self, args, kwargs):  # pylint: disable=too-many-branches
         )
         if expand_mid_measure:
             # Assume that tapes are not split if old device is used since postselection is not supported
-            tapes, _ = qml.defer_measurements(self._tape)
+            tapes, _ = qml.defer_measurements(self._tape, device=self.device)
             self._tape = tapes[0]
 
         if self.expansion_strategy == "device":
diff --git a/pennylane/transforms/defer_measurements.py b/pennylane/transforms/defer_measurements.py
index ca46bd5cb51..7720030e31f 100644
--- a/pennylane/transforms/defer_measurements.py
+++ b/pennylane/transforms/defer_measurements.py
@@ -232,13 +232,6 @@ def func(x, y):
 
     new_tape = type(tape)(new_operations, new_measurements, shots=tape.shots)
 
-    # pylint: disable=protected-access
-    # We can remove this once _qfunc_output is removed from tape
-    if not isinstance(tape._qfunc_output, Sequence):
-        new_tape._qfunc_output = new_measurements[0]
-    else:
-        new_tape._qfunc_output = type(tape._qfunc_output)(new_measurements)
-
     if is_postselecting and new_tape.batch_size is not None:
         # Split tapes if broadcasting with postselection
         return qml.transforms.broadcast_expand(new_tape)
diff --git a/tests/devices/qubit/test_sampling.py b/tests/devices/qubit/test_sampling.py
index 268759aed45..d8af614d37e 100644
--- a/tests/devices/qubit/test_sampling.py
+++ b/tests/devices/qubit/test_sampling.py
@@ -14,6 +14,7 @@
 """Unit tests for sample_state in devices/qubit."""
 
 from random import shuffle
+from typing import Sequence
 
 import pytest
 
@@ -41,6 +42,47 @@ def _init_state(n):
     return _init_state
 
 
+def _valid_flex_int(s):
+    """Returns True if s is a non-negative integer."""
+    return isinstance(s, int) and s >= 0
+
+
+def _valid_flex_tuple(s):
+    """Returns True if s is a tuple of the form (shots, copies)."""
+    return (
+        isinstance(s, tuple)
+        and len(s) == 2
+        and _valid_flex_int(s[0])
+        and isinstance(s[1], int)
+        and s[1] > 0
+    )
+
+
+class _FlexShots(Shots):
+    """Shots class that allows zero shots."""
+
+    # pylint: disable=super-init-not-called
+    def __init__(self, shots=None):
+        if shots is None:
+            self.total_shots = None
+            self.shot_vector = ()
+        elif isinstance(shots, int):
+            if shots < 0:
+                raise self._SHOT_ERROR
+            self.total_shots = shots
+            self.shot_vector = (qml.measurements.ShotCopies(shots, 1),)
+        elif isinstance(shots, Sequence):
+            if not all(_valid_flex_int(s) or _valid_flex_tuple(s) for s in shots):
+                raise self._SHOT_ERROR
+            self.__all_tuple_init__([s if isinstance(s, tuple) else (s, 1) for s in shots])
+        elif isinstance(shots, self.__class__):
+            return  # self already _is_ shots as defined by __new__
+        else:
+            raise self._SHOT_ERROR
+
+        self._frozen = True
+
+
 def samples_to_probs(samples, num_wires):
     """Converts samples to probs"""
     samples_decimal = [np.ravel_multi_index(sample, [2] * num_wires) for sample in samples]
@@ -514,223 +556,94 @@ class TestInvalidStateSamples:
             qml.var(qml.PauliZ(0)),
         ],
     )
-    @pytest.mark.parametrize("interface", ["numpy", "autograd", "torch", "tensorflow"])
+    @pytest.mark.parametrize("interface", ["numpy", "autograd", "torch", "tensorflow", "jax"])
     @pytest.mark.parametrize("shots", [0, [0, 0]])
     def test_nan_float_result(self, mp, interface, shots):
         """Test that the result of circuits with 0 probability postselections is NaN with the
         expected shape."""
         state = qml.math.full((2, 2), np.NaN, like=interface)
-        res = measure_with_samples((mp,), state, Shots(shots), is_state_batched=False)
-
-        if not isinstance(shots, list):
-            assert qml.math.ndim(res) == 0
-            assert qml.math.isnan(res)
-            assert qml.math.get_interface(res) == interface
-
-        else:
-            assert isinstance(res, tuple)
-            assert len(res) == 2
-            for r in res:
-                assert isinstance(r, tuple)
-                r = r[0]
-                assert qml.math.ndim(r) == 0
-                assert qml.math.isnan(r)
-                assert qml.math.get_interface(r) == interface
-
-    @pytest.mark.jax
-    @pytest.mark.parametrize(
-        "mp",
-        [
-            qml.expval(qml.PauliZ(0)),
-            qml.expval(
-                qml.Hamiltonian(
-                    [1.0, 2.0, 3.0, 4.0],
-                    [qml.PauliZ(0) @ qml.PauliX(1), qml.PauliX(1), qml.PauliZ(1), qml.PauliY(1)],
-                )
-            ),
-            qml.expval(
-                qml.dot(
-                    [1.0, 2.0, 3.0, 4.0],
-                    [qml.PauliZ(0) @ qml.PauliX(1), qml.PauliX(1), qml.PauliZ(1), qml.PauliY(1)],
-                )
-            ),
-            qml.var(qml.PauliZ(0)),
-            qml.var(
-                qml.Hamiltonian(
-                    [1.0, 2.0, 3.0, 4.0],
-                    [qml.PauliZ(0) @ qml.PauliX(1), qml.PauliX(1), qml.PauliZ(1), qml.PauliY(1)],
-                )
-            ),
-            qml.var(
-                qml.dot(
-                    [1.0, 2.0, 3.0, 4.0],
-                    [qml.PauliZ(0) @ qml.PauliX(1), qml.PauliX(1), qml.PauliZ(1), qml.PauliY(1)],
-                )
-            ),
-        ],
-    )
-    @pytest.mark.parametrize("use_jit", [True, False])
-    @pytest.mark.parametrize("shots", [0, [0, 0]])
-    def test_nan_float_result_jax(self, mp, use_jit, shots):
-        """Test that the result of circuits with 0 probability postselections is NaN with the
-        expected shape."""
-        state = qml.math.full((2, 2), np.NaN, like="jax")
-        if use_jit:
-            import jax
-
-            res = jax.jit(measure_with_samples, static_argnums=[0, 2, 3])(
-                (mp,), state, Shots(shots), is_state_batched=False
-            )
-        else:
-            res = measure_with_samples((mp,), state, Shots(shots), is_state_batched=False)
+        res = measure_with_samples((mp,), state, _FlexShots(shots), is_state_batched=False)
 
         if not isinstance(shots, list):
             assert isinstance(res, tuple)
             res = res[0]
             assert qml.math.ndim(res) == 0
             assert qml.math.isnan(res)
-            assert qml.math.get_interface(res) == "jax"
 
         else:
-            assert isinstance(res, tuple)
-            res = res[0]
             assert isinstance(res, tuple)
             assert len(res) == 2
             for r in res:
+                assert isinstance(r, tuple)
+                r = r[0]
                 assert qml.math.ndim(r) == 0
                 assert qml.math.isnan(r)
-                assert qml.math.get_interface(r) == "jax"
 
     @pytest.mark.all_interfaces
     @pytest.mark.parametrize(
         "mp", [qml.sample(wires=0), qml.sample(op=qml.PauliZ(0)), qml.sample(wires=[0, 1])]
     )
-    @pytest.mark.parametrize("interface", ["numpy", "autograd", "torch", "tensorflow"])
+    @pytest.mark.parametrize("interface", ["numpy", "autograd", "torch", "tensorflow", "jax"])
     @pytest.mark.parametrize("shots", [0, [0, 0]])
     def test_nan_samples(self, mp, interface, shots):
         """Test that the result of circuits with 0 probability postselections is NaN with the
         expected shape."""
         state = qml.math.full((2, 2), np.NaN, like=interface)
-        res = measure_with_samples((mp,), state, Shots(shots), is_state_batched=False)
+        res = measure_with_samples((mp,), state, _FlexShots(shots), is_state_batched=False)
 
         if not isinstance(shots, list):
             assert isinstance(res, tuple)
             res = res[0]
-            assert qml.math.shape(res) == (shots, 2)
-            assert qml.math.get_interface(res) == interface
+            assert qml.math.shape(res) == (shots,) if len(mp.wires) == 1 else (shots, len(mp.wires))
 
         else:
-            assert isinstance(res, tuple)
-            res = res[0]
             assert isinstance(res, tuple)
             assert len(res) == 2
             for i, r in enumerate(res):
-                assert qml.math.shape(r) == (shots[i], 2)
-                assert qml.math.get_interface(r) == interface
-
-    @pytest.mark.jax
-    @pytest.mark.parametrize(
-        "mp", [qml.sample(wires=0), qml.sample(op=qml.PauliZ(0)), qml.sample(wires=[0, 1])]
-    )
-    @pytest.mark.parametrize("use_jit", [True, False])
-    @pytest.mark.parametrize("shots", [0, [0, 0]])
-    def test_nan_samples_jax(self, mp, use_jit, shots):
-        """Test that the result of circuits with 0 probability postselections is NaN with the
-        expected shape."""
-        state = qml.math.full((2, 2), np.NaN, like="jax")
-        if use_jit:
-            import jax
-
-            res = jax.jit(measure_with_samples, static_argnums=[0, 2, 3])(
-                (mp,), state, Shots(shots), is_state_batched=False
-            )
-        else:
-            res = measure_with_samples((mp,), state, Shots(shots), is_state_batched=False)
-
-        if not isinstance(shots, list):
-            assert isinstance(res, tuple)
-            res = res[0]
-            assert qml.math.shape(res) == (shots, 2)
-            assert qml.math.get_interface(res) == "jax"
-
-        else:
-            assert isinstance(res, tuple)
-            res = res[0]
-            assert isinstance(res, tuple)
-            assert len(res) == 2
-            for i, r in enumerate(res):
-                assert qml.math.shape(r) == (shots[i], 2)
-                assert qml.math.get_interface(r) == "jax"
+                assert isinstance(r, tuple)
+                r = r[0]
+                assert (
+                    qml.math.shape(r) == (shots[i],)
+                    if len(mp.wires) == 1
+                    else (shots[i], len(mp.wires))
+                )
 
     @pytest.mark.all_interfaces
-    @pytest.mark.parametrize("interface", ["numpy", "autograd", "torch", "tensorflow"])
+    @pytest.mark.parametrize("interface", ["numpy", "autograd", "torch", "tensorflow", "jax"])
     @pytest.mark.parametrize("shots", [0, [0, 0]])
     def test_nan_classical_shadows(self, interface, shots):
         """Test that classical_shadows returns an empty array when the state has
         NaN values"""
         state = qml.math.full((2, 2), np.NaN, like=interface)
         res = measure_with_samples(
-            (qml.classical_shadow([0]),), state, Shots(shots), is_state_batched=False
+            (qml.classical_shadow([0]),), state, _FlexShots(shots), is_state_batched=False
         )
 
         if not isinstance(shots, list):
             assert isinstance(res, tuple)
             res = res[0]
-            assert qml.math.shape(res) == (0,)
-            assert qml.math.get_interface(res) == interface
-
-        else:
-            assert isinstance(res, tuple)
-            res = res[0]
-            assert isinstance(res, tuple)
-            assert len(res) == 2
-            for r in res:
-                assert qml.math.shape(r) == (0,)
-                assert qml.math.get_interface(r) == interface
-
-    @pytest.mark.jax
-    @pytest.mark.parametrize("use_jit", [True, False])
-    @pytest.mark.parametrize("shots", [0, [0, 0]])
-    def test_nan_classical_shadows_jax(self, use_jit, shots):
-        """Test that classical_shadows returns an empty array when the state has
-        NaN values"""
-        state = qml.math.full((2, 2), np.NaN, like="jax")
-        if use_jit:
-            import jax
-
-            res = jax.jit(measure_with_samples, static_argnums=[0, 2, 3])(
-                (qml.classical_shadow([0]),), state, Shots(shots), is_state_batched=False
-            )
-        else:
-            res = measure_with_samples(
-                (qml.classical_shadow([0]),), state, Shots(shots), is_state_batched=False
-            )
-
-        if not isinstance(shots, list):
-            assert isinstance(res, tuple)
-            res = res[0]
-            assert qml.math.shape(res) == (0,)
-            assert qml.math.get_interface(res) == "jax"
+            assert qml.math.shape(res) == (2, 0, 1)
+            assert qml.math.size(res) == 0
 
         else:
-            assert isinstance(res, tuple)
-            res = res[0]
             assert isinstance(res, tuple)
             assert len(res) == 2
             for r in res:
-                assert qml.math.shape(r) == (0,)
-                assert qml.math.get_interface(r) == "jax"
+                assert isinstance(r, tuple)
+                r = r[0]
+                assert qml.math.shape(r) == (2, 0, 1)
+                assert qml.math.size(r) == 0
 
     @pytest.mark.all_interfaces
     @pytest.mark.parametrize("H", [qml.PauliZ(0), [qml.PauliZ(0), qml.PauliX(1)]])
-    @pytest.mark.parametrize("interface", ["numpy", "autograd", "torch", "tensorflow"])
+    @pytest.mark.parametrize("interface", ["numpy", "autograd", "torch", "tensorflow", "jax"])
     @pytest.mark.parametrize("shots", [0, [0, 0]])
     def test_nan_shadow_expval(self, H, interface, shots):
         """Test that shadow_expval returns an empty array when the state has
         NaN values"""
         state = qml.math.full((2, 2), np.NaN, like=interface)
         res = measure_with_samples(
-            (qml.shadow_expval(H),), state, Shots(shots), is_state_batched=False
+            (qml.shadow_expval(H),), state, _FlexShots(shots), is_state_batched=False
         )
 
         if not isinstance(shots, list):
@@ -738,53 +651,15 @@ def test_nan_shadow_expval(self, H, interface, shots):
             res = res[0]
             assert qml.math.shape(res) == qml.math.shape(H)
             assert qml.math.all(qml.math.isnan(res))
-            assert qml.math.get_interface(res) == interface
-
-        else:
-            assert isinstance(res, tuple)
-            res = res[0]
-            assert isinstance(res, tuple)
-            assert len(res) == 2
-            for r in res:
-                assert qml.math.shape(r) == qml.math.shape(H)
-                assert qml.math.all(qml.math.isnan(r))
-                assert qml.math.get_interface(r) == interface
-
-    @pytest.mark.jax
-    @pytest.mark.parametrize("H", [qml.PauliZ(0), [qml.PauliZ(0), qml.PauliX(1)]])
-    @pytest.mark.parametrize("use_jit", [True, False])
-    @pytest.mark.parametrize("shots", [0, [0, 0]])
-    def test_nan_shadow_expval_jax(self, H, use_jit, shots):
-        """Test that shadow_expval returns an empty array when the state has
-        NaN values"""
-        state = qml.math.full((2, 2), np.NaN, like="jax")
-        if use_jit:
-            import jax
-
-            res = jax.jit(measure_with_samples, static_argnums=[0, 2, 3])(
-                (qml.shadow_expval(H),), state, Shots(shots), is_state_batched=False
-            )
-        else:
-            res = measure_with_samples(
-                (qml.shadow_expval(H),), state, Shots(shots), is_state_batched=False
-            )
-
-        if not isinstance(shots, list):
-            assert isinstance(res, tuple)
-            res = res[0]
-            assert qml.math.shape(res) == qml.math.shape(H)
-            assert qml.math.all(qml.math.isnan(res))
-            assert qml.math.get_interface(res) == "jax"
 
         else:
-            assert isinstance(res, tuple)
-            res = res[0]
             assert isinstance(res, tuple)
             assert len(res) == 2
             for r in res:
+                assert isinstance(r, tuple)
+                r = r[0]
                 assert qml.math.shape(r) == qml.math.shape(H)
                 assert qml.math.all(qml.math.isnan(r))
-                assert qml.math.get_interface(r) == "jax"
 
 
 class TestBroadcasting:

From 03807ac63c6df26fa6c19e29d21ad0624016c2b4 Mon Sep 17 00:00:00 2001
From: Mudit Pandey <mudit.pandey@xanadu.ai>
Date: Tue, 17 Oct 2023 17:20:31 -0400
Subject: [PATCH 115/131] [skip ci] adding tests; defer_measurements bug

---
 pennylane/devices/qubit/simulate.py           |  29 +--
 pennylane/transforms/defer_measurements.py    |   5 -
 .../default_qubit/test_default_qubit.py       | 209 +++++++++++++++++-
 3 files changed, 210 insertions(+), 33 deletions(-)

diff --git a/pennylane/devices/qubit/simulate.py b/pennylane/devices/qubit/simulate.py
index 1435ab913ce..dee0cefd545 100644
--- a/pennylane/devices/qubit/simulate.py
+++ b/pennylane/devices/qubit/simulate.py
@@ -47,43 +47,18 @@
 }
 
 
-def _valid_flex_int(s):
-    """Returns True if s is a non-negative integer."""
-    return isinstance(s, int) and s >= 0
-
-
-def _valid_flex_tuple(s):
-    """Returns True if s is a tuple of the form (shots, copies)."""
-    return (
-        isinstance(s, tuple)
-        and len(s) == 2
-        and _valid_flex_int(s[0])
-        and isinstance(s[1], int)
-        and s[1] > 0
-    )
-
-
 class _FlexShots(qml.measurements.Shots):
     """Shots class that allows zero shots."""
 
     # pylint: disable=super-init-not-called
     def __init__(self, shots=None):
-        if shots is None:
-            self.total_shots = None
-            self.shot_vector = ()
-        elif isinstance(shots, int):
-            if shots < 0:
-                raise self._SHOT_ERROR
+        if isinstance(shots, int):
             self.total_shots = shots
             self.shot_vector = (qml.measurements.ShotCopies(shots, 1),)
         elif isinstance(shots, Sequence):
-            if not all(_valid_flex_int(s) or _valid_flex_tuple(s) for s in shots):
-                raise self._SHOT_ERROR
             self.__all_tuple_init__([s if isinstance(s, tuple) else (s, 1) for s in shots])
-        elif isinstance(shots, self.__class__):
-            return  # self already _is_ shots as defined by __new__
         else:
-            raise self._SHOT_ERROR
+            return  # self already _is_ shots as defined by __new__
 
         self._frozen = True
 
diff --git a/pennylane/transforms/defer_measurements.py b/pennylane/transforms/defer_measurements.py
index 7720030e31f..e24ba6f0be8 100644
--- a/pennylane/transforms/defer_measurements.py
+++ b/pennylane/transforms/defer_measurements.py
@@ -236,11 +236,6 @@ def func(x, y):
         # Split tapes if broadcasting with postselection
         return qml.transforms.broadcast_expand(new_tape)
 
-    def null_postprocessing(results):
-        """A postprocessing function returned by a transform that only converts the batch of results
-        into a result for a single ``QuantumTape``."""
-        return results[0]
-
     return [new_tape], null_postprocessing
 
 
diff --git a/tests/devices/default_qubit/test_default_qubit.py b/tests/devices/default_qubit/test_default_qubit.py
index 1867834bf6d..65ef88f0e9a 100644
--- a/tests/devices/default_qubit/test_default_qubit.py
+++ b/tests/devices/default_qubit/test_default_qubit.py
@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 """Tests for default qubit."""
-# pylint: disable=import-outside-toplevel, no-member
+# pylint: disable=import-outside-toplevel, no-member, too-many-arguments
 
 import pytest
 
@@ -1637,6 +1637,213 @@ def test_projector_dynamic_type(max_workers, n_wires):
         assert np.isclose(res, 1 / 2**n_wires)
 
 
+@pytest.mark.all_interfaces
+@pytest.mark.parametrize("interface", ["numpy", "autograd", "torch", "jax", "tensorflow"])
+@pytest.mark.parametrize("use_jit", [True, False])
+class TestPostselection:
+    """Various integration tests for postselection of mid-circuit measurements."""
+
+    @pytest.mark.parametrize(
+        "mp",
+        [
+            qml.expval(qml.PauliZ(0)),
+            qml.var(qml.PauliZ(0)),
+            qml.probs(wires=[0, 1]),
+            qml.state(),
+            qml.density_matrix(wires=0),
+            qml.purity(0),
+            qml.vn_entropy(0),
+            qml.mutual_info(0, 1),
+        ],
+    )
+    @pytest.mark.parametrize("param", np.linspace(np.pi / 4, 3 * np.pi / 4, 3))
+    def test_postselection_valid_analytic(self, param, mp, interface, use_jit):
+        """Test that the results of a circuit with postselection is expected
+        with analytic execution."""
+        if use_jit and interface != "jax":
+            pytest.skip("Cannot JIT in non-JAX interfaces.")
+
+        dev = qml.device("default.qubit")
+        param = qml.math.asarray(param, like=interface)
+
+        @qml.qnode(dev, interface=interface)
+        def circ_postselect(theta):
+            qml.RX(theta, 0)
+            qml.CNOT([0, 1])
+            qml.measure(0, postselect=1)
+            return qml.apply(mp)
+
+        @qml.qnode(dev, interface=interface)
+        def circ_expected():
+            qml.RX(np.pi, 0)
+            qml.CNOT([0, 1])
+            return qml.apply(mp)
+
+        if use_jit:
+            import jax
+
+            circ_postselect = jax.jit(circ_postselect)
+
+        res = circ_postselect(param)
+        expected = circ_expected()
+
+        assert qml.math.allclose(res, expected)
+        assert qml.math.get_interface(res) == qml.math.get_interface(expected)
+
+    @pytest.mark.parametrize(
+        "mp",
+        [
+            qml.expval(qml.PauliZ(0)),
+            qml.var(qml.PauliZ(0)),
+            qml.probs(wires=[0, 1]),
+            qml.shadow_expval(qml.Hamiltonian([1.0, -1.0], [qml.PauliZ(0), qml.PauliX(0)])),
+            # qml.sample, qml.classical_shadow, qml.counts are not included because their
+            # shape/values are dependent on the number of shots, which will be changed
+            # randomly per the binomial distribution and the probability of the postselected
+            # state
+        ],
+    )
+    @pytest.mark.parametrize("param", np.linspace(np.pi / 4, 3 * np.pi / 4, 3))
+    @pytest.mark.parametrize("shots", [20000, (20000, 20000)])
+    def test_postselection_valid_finite_shots(
+        self, param, mp, shots, interface, use_jit, tol_stochastic
+    ):
+        """Test that the results of a circuit with postselection is expected with
+        finite shots."""
+        if use_jit and interface != "jax":
+            pytest.skip("Cannot JIT in non-JAX interfaces.")
+
+        dev = qml.device("default.qubit", shots=shots)
+        param = qml.math.asarray(param, like=interface)
+
+        @qml.qnode(dev, interface=interface)
+        def circ_postselect(theta):
+            qml.RX(theta, 0)
+            qml.CNOT([0, 1])
+            qml.measure(0, postselect=1)
+            return qml.apply(mp)
+
+        @qml.qnode(dev, interface=interface)
+        def circ_expected():
+            qml.RX(np.pi, 0)
+            qml.CNOT([0, 1])
+            return qml.apply(mp)
+
+        if use_jit:
+            import jax
+
+            circ_postselect = jax.jit(circ_postselect)
+
+        res = circ_postselect(param)
+        expected = circ_expected()
+
+        if not isinstance(shots, tuple):
+            assert qml.math.allclose(res, expected, atol=tol_stochastic, rtol=0)
+            assert qml.math.get_interface(res) == qml.math.get_interface(expected)
+
+        else:
+            assert isinstance(res, tuple)
+            for r, e in zip(res, expected):
+                assert qml.math.allclose(r, e, atol=tol_stochastic, rtol=0)
+                assert qml.math.get_interface(r) == qml.math.get_interface(e)
+
+    @pytest.mark.parametrize(
+        "mp, autograd_interface",
+        [
+            (qml.expval(qml.PauliZ(0)), "autograd"),
+            (qml.var(qml.PauliZ(0)), "autograd"),
+            (qml.probs(wires=[0, 1]), "autograd"),
+            (qml.state(), "autograd"),
+            (qml.density_matrix(wires=0), "autograd"),
+            (qml.purity(0), "numpy"),
+            # qml.vn_entropy(0),
+            # qml.mutual_info(0, 1),
+        ],
+    )
+    def test_postselection_invalid_analytic(self, mp, autograd_interface, interface, use_jit):
+        """Test that the results of a qnode are nan values of the correct shape if the state
+        that we are postselecting has a zero probability of occurring."""
+
+        if use_jit and interface != "jax":
+            pytest.skip("Can't jit with non-jax interfaces.")
+
+        # Wires are specified so that the shape for measurements can be determined correctly
+        dev = qml.device("default.qubit", wires=2)
+
+        @qml.qnode(dev, interface=interface)
+        def circ():
+            qml.RX(np.pi, 0)
+            qml.CNOT([0, 1])
+            qml.measure(0, postselect=0)
+            return qml.apply(mp)
+
+        if use_jit:
+            import jax
+
+            circ = jax.jit(circ)
+
+        res = circ()
+        if interface == "autograd":
+            assert qml.math.get_interface(res) == autograd_interface
+        else:
+            assert qml.math.get_interface(res) == interface
+        assert qml.math.shape(res) == mp.shape(dev, qml.measurements.Shots(None))
+        assert qml.math.all(qml.math.isnan(res))
+
+    @pytest.mark.parametrize(
+        "mp, expected_shape",
+        [
+            (qml.expval(qml.PauliZ(0)), ()),
+            (qml.var(qml.PauliZ(0)), ()),
+            (qml.sample(qml.PauliZ(0)), (0,)),
+            (qml.classical_shadow(wires=0), (2, 0, 1)),
+            (qml.shadow_expval(qml.Hamiltonian([1, 1], [qml.PauliZ(0), qml.PauliX(0)])), ()),
+            # qml.probs and qml.counts are not tested because they fail in this case
+        ],
+    )
+    @pytest.mark.parametrize("shots", [10, (10, 10)])
+    def test_postselection_invalid_finite_shots(
+        self, mp, expected_shape, shots, interface, use_jit
+    ):
+        """Test that the results of a qnode are nan values of the correct shape if the state
+        that we are postselecting has a zero probability of occurring with finite shots."""
+
+        if use_jit and interface != "jax":
+            pytest.skip("Can't jit with non-jax interfaces.")
+
+        # Wires are specified so that the shape for measurements can be determined correctly
+        dev = qml.device("default.qubit", wires=2)
+
+        @qml.qnode(dev, interface=interface)
+        def circ():
+            qml.RX(np.pi, 0)
+            qml.CNOT([0, 1])
+            qml.measure(0, postselect=0)
+            return qml.apply(mp)
+
+        if use_jit:
+            import jax
+
+            circ = jax.jit(circ, static_argnames=["shots"])
+
+        res = circ(shots=shots)
+
+        if not isinstance(shots, tuple):
+            assert qml.math.shape(res) == expected_shape
+            assert qml.math.get_interface(res) == interface if interface != "autograd" else "numpy"
+            if not 0 in expected_shape:  # No nan values if array is empty
+                assert qml.math.all(qml.math.isnan(res))
+        else:
+            assert isinstance(res, tuple)
+            for r in res:
+                assert qml.math.shape(r) == expected_shape
+                assert (
+                    qml.math.get_interface(r) == interface if interface != "autograd" else "numpy"
+                )
+                if not 0 in expected_shape:  # No nan values if array is empty
+                    assert qml.math.all(qml.math.isnan(r))
+
+
 class TestIntegration:
     """Various integration tests"""
 

From 160a77152048e49b618f2fec42a57de5cc1dec41 Mon Sep 17 00:00:00 2001
From: Mudit Pandey <mudit.pandey@xanadu.ai>
Date: Tue, 17 Oct 2023 17:32:19 -0400
Subject: [PATCH 116/131] Removed testing with shot vectors

---
 tests/devices/default_qubit/test_default_qubit.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/tests/devices/default_qubit/test_default_qubit.py b/tests/devices/default_qubit/test_default_qubit.py
index 65ef88f0e9a..b85e8b76bf3 100644
--- a/tests/devices/default_qubit/test_default_qubit.py
+++ b/tests/devices/default_qubit/test_default_qubit.py
@@ -1704,7 +1704,7 @@ def circ_expected():
         ],
     )
     @pytest.mark.parametrize("param", np.linspace(np.pi / 4, 3 * np.pi / 4, 3))
-    @pytest.mark.parametrize("shots", [20000, (20000, 20000)])
+    @pytest.mark.parametrize("shots", [20000])
     def test_postselection_valid_finite_shots(
         self, param, mp, shots, interface, use_jit, tol_stochastic
     ):
@@ -1742,6 +1742,8 @@ def circ_expected():
             assert qml.math.get_interface(res) == qml.math.get_interface(expected)
 
         else:
+            # No testing with shot vectors currently, but keeping this here so that
+            # we can just use it once shot vectors are supported
             assert isinstance(res, tuple)
             for r, e in zip(res, expected):
                 assert qml.math.allclose(r, e, atol=tol_stochastic, rtol=0)

From d336ef099b276d52d84169389219d77fe49c1638 Mon Sep 17 00:00:00 2001
From: Mudit Pandey <mudit.pandey@xanadu.ai>
Date: Tue, 17 Oct 2023 18:10:32 -0400
Subject: [PATCH 117/131] Updated tests; fixed docs

---
 doc/introduction/measurements.rst             |  4 +-
 .../default_qubit/test_default_qubit.py       | 40 +++++++++++--------
 2 files changed, 26 insertions(+), 18 deletions(-)

diff --git a/doc/introduction/measurements.rst b/doc/introduction/measurements.rst
index e9fccdcb7b0..b15b2bd6ceb 100644
--- a/doc/introduction/measurements.rst
+++ b/doc/introduction/measurements.rst
@@ -380,8 +380,8 @@ be equivalent to projecting the state vector onto the :math:`|1\rangle` state on
         qml.cond(m0, qml.PauliX)(wires=1)
         return qml.sample(wires=1)
 
-By postselecting on ``1``, we only consider the ``1`` measurement outcome. So, the probability of measuring 1
-on wire 1 should be 100%. Executing this QNode with 10 shots:
+By postselecting on ``1``, we only consider the ``1`` measurement outcome on wire 0. So, the probability of
+measuring ``1`` on wire 1 after postselection should also be 1. Executing this QNode with 10 shots:
 
 >>> func(np.pi / 2, shots=10)
 array([1, 1, 1, 1, 1, 1, 1])
diff --git a/tests/devices/default_qubit/test_default_qubit.py b/tests/devices/default_qubit/test_default_qubit.py
index b85e8b76bf3..e52595e51ec 100644
--- a/tests/devices/default_qubit/test_default_qubit.py
+++ b/tests/devices/default_qubit/test_default_qubit.py
@@ -1704,14 +1704,14 @@ def circ_expected():
         ],
     )
     @pytest.mark.parametrize("param", np.linspace(np.pi / 4, 3 * np.pi / 4, 3))
-    @pytest.mark.parametrize("shots", [20000])
+    @pytest.mark.parametrize("shots", [50000, (50000, 50000)])
     def test_postselection_valid_finite_shots(
         self, param, mp, shots, interface, use_jit, tol_stochastic
     ):
         """Test that the results of a circuit with postselection is expected with
         finite shots."""
-        if use_jit and interface != "jax":
-            pytest.skip("Cannot JIT in non-JAX interfaces.")
+        if use_jit and (interface != "jax" or isinstance(shots, tuple)):
+            pytest.skip("Cannot JIT in non-JAX interfaces, or with shot vectors.")
 
         dev = qml.device("default.qubit", shots=shots)
         param = qml.math.asarray(param, like=interface)
@@ -1742,30 +1742,35 @@ def circ_expected():
             assert qml.math.get_interface(res) == qml.math.get_interface(expected)
 
         else:
-            # No testing with shot vectors currently, but keeping this here so that
-            # we can just use it once shot vectors are supported
             assert isinstance(res, tuple)
             for r, e in zip(res, expected):
                 assert qml.math.allclose(r, e, atol=tol_stochastic, rtol=0)
                 assert qml.math.get_interface(r) == qml.math.get_interface(e)
 
     @pytest.mark.parametrize(
-        "mp, autograd_interface",
+        "mp, autograd_interface, is_nan",
         [
-            (qml.expval(qml.PauliZ(0)), "autograd"),
-            (qml.var(qml.PauliZ(0)), "autograd"),
-            (qml.probs(wires=[0, 1]), "autograd"),
-            (qml.state(), "autograd"),
-            (qml.density_matrix(wires=0), "autograd"),
-            (qml.purity(0), "numpy"),
-            # qml.vn_entropy(0),
-            # qml.mutual_info(0, 1),
+            (qml.expval(qml.PauliZ(0)), "autograd", True),
+            (qml.var(qml.PauliZ(0)), "autograd", True),
+            (qml.probs(wires=[0, 1]), "autograd", True),
+            (qml.state(), "autograd", True),
+            (qml.density_matrix(wires=0), "autograd", True),
+            (qml.purity(0), "numpy", True),
+            (qml.vn_entropy(0), "numpy", False),
+            (qml.mutual_info(0, 1), "numpy", False),
         ],
     )
-    def test_postselection_invalid_analytic(self, mp, autograd_interface, interface, use_jit):
+    def test_postselection_invalid_analytic(
+        self, mp, autograd_interface, is_nan, interface, use_jit
+    ):
         """Test that the results of a qnode are nan values of the correct shape if the state
         that we are postselecting has a zero probability of occurring."""
 
+        if (isinstance(mp, qml.measurements.MutualInfoMP) and interface != "jax") or (
+            isinstance(mp, qml.measurements.VnEntropyMP) and interface == "tensorflow"
+        ):
+            pytest.skip("Unsupported measurements and interfaces.")
+
         if use_jit and interface != "jax":
             pytest.skip("Can't jit with non-jax interfaces.")
 
@@ -1790,7 +1795,10 @@ def circ():
         else:
             assert qml.math.get_interface(res) == interface
         assert qml.math.shape(res) == mp.shape(dev, qml.measurements.Shots(None))
-        assert qml.math.all(qml.math.isnan(res))
+        if is_nan:
+            assert qml.math.all(qml.math.isnan(res))
+        else:
+            assert qml.math.allclose(res, 0.0)
 
     @pytest.mark.parametrize(
         "mp, expected_shape",

From 54be0ca9d68254220bc05cfa7c1119437e07cb4d Mon Sep 17 00:00:00 2001
From: Mudit Pandey <mudit.pandey@xanadu.ai>
Date: Wed, 18 Oct 2023 15:10:45 -0400
Subject: [PATCH 118/131] Fixing/adding tests

---
 pennylane/devices/default_qubit.py          |  2 +-
 pennylane/drawer/draw.py                    | 10 +++++++++-
 pennylane/qnode.py                          |  6 ++++--
 tests/drawer/test_draw.py                   | 21 ++++++++++++++++++++-
 tests/transforms/test_defer_measurements.py |  6 +++---
 5 files changed, 37 insertions(+), 8 deletions(-)

diff --git a/pennylane/devices/default_qubit.py b/pennylane/devices/default_qubit.py
index 2453fe8a284..702666efdd1 100644
--- a/pennylane/devices/default_qubit.py
+++ b/pennylane/devices/default_qubit.py
@@ -382,7 +382,7 @@ def preprocess(
         config = self._setup_execution_config(execution_config)
         transform_program = TransformProgram()
 
-        transform_program.add_transform(qml.defer_measurements)
+        transform_program.add_transform(qml.defer_measurements, device=self)
         transform_program.add_transform(validate_device_wires, self.wires, name=self.name)
         transform_program.add_transform(
             decompose, stopping_condition=stopping_condition, name=self.name
diff --git a/pennylane/drawer/draw.py b/pennylane/drawer/draw.py
index ce3939e2839..47e45651b58 100644
--- a/pennylane/drawer/draw.py
+++ b/pennylane/drawer/draw.py
@@ -265,7 +265,15 @@ def wrapper(*args, **kwargs):
                 tapes = qnode.construct(args, kwargs)
                 if isinstance(qnode.device, qml.devices.Device):
                     program = qnode.transform_program
-                    tapes = program([qnode.tape])
+                    if any(
+                        isinstance(op, qml.measurements.MidMeasureMP)
+                        for op in qnode.tape.operations
+                    ):
+                        tapes, _ = qml.defer_measurements(qnode.tape, device=qnode.device)
+                    else:
+                        tapes = [qnode.tape]
+
+                    tapes = program(tapes)
 
             finally:
                 qnode.expansion_strategy = original_expansion_strategy
diff --git a/pennylane/qnode.py b/pennylane/qnode.py
index 3676aa72107..6890554e9fd 100644
--- a/pennylane/qnode.py
+++ b/pennylane/qnode.py
@@ -898,14 +898,16 @@ def construct(self, args, kwargs):  # pylint: disable=too-many-branches
                 )
 
         # Apply the deferred measurement principle if the device doesn't
-        # support mid-circuit measurements natively
+        # support mid-circuit measurements natively.
+        # Only apply transform with old device API as postselection with
+        # broadcasting will split tapes.
         expand_mid_measure = (
             any(isinstance(op, MidMeasureMP) for op in self.tape.operations)
             and not isinstance(self.device, qml.devices.Device)
             and not self.device.capabilities().get("supports_mid_measure", False)
         )
         if expand_mid_measure:
-            # Assume that tapes are not split if old device is used since postselection is not supported
+            # Assume that tapes are not split if old device is used since postselection is not supported.
             tapes, _ = qml.defer_measurements(self._tape, device=self.device)
             self._tape = tapes[0]
 
diff --git a/tests/drawer/test_draw.py b/tests/drawer/test_draw.py
index a574ceac7c1..05333550320 100644
--- a/tests/drawer/test_draw.py
+++ b/tests/drawer/test_draw.py
@@ -99,7 +99,7 @@ def test_decimals_higher_value(self):
     def test_decimals_multiparameters(self):
         """Test decimals also displays parameters when the operation has multiple parameters."""
 
-        @qml.qnode(qml.device("default.qubit", wires=(0)))
+        @qml.qnode(qml.device("default.qubit", wires=[0]))
         def circ(x):
             qml.Rot(*x, wires=0)
             return qml.expval(qml.PauliZ(0))
@@ -281,6 +281,25 @@ def circ():
         assert draw(circ)() == expected
 
 
+@pytest.mark.parametrize("device_name", ["default.qubit"])
+def test_mid_circuit_measurement_device_api(device_name, mocker):
+    """Test that a circuit containing mid-circuit measurements is transformed by the drawer
+    to use deferred measurements if the device uses the new device API."""
+    dev = qml.device(device_name)
+
+    @qml.qnode(dev)
+    def circ():
+        qml.PauliX(0)
+        qml.measure(0)
+        return qml.probs(wires=0)
+
+    draw_qnode = qml.draw(circ)
+    spy = mocker.spy(qml.defer_measurements, "_transform")
+
+    _ = draw_qnode()
+    spy.assert_called_once()
+
+
 @pytest.mark.parametrize(
     "transform",
     [
diff --git a/tests/transforms/test_defer_measurements.py b/tests/transforms/test_defer_measurements.py
index d478832a5e6..63f876e72f2 100644
--- a/tests/transforms/test_defer_measurements.py
+++ b/tests/transforms/test_defer_measurements.py
@@ -156,7 +156,7 @@ def qnode2(phi):
 
         # Outputs should match
         assert np.isclose(qnode1(np.pi / 4), qnode2(np.pi / 4))
-        assert spy.call_count == 3  # once per device preprocessing, one for qnode
+        assert spy.call_count == 2  # once per device preprocessing
 
         deferred_tapes, _ = qml.defer_measurements(qnode1.qtape)
         deferred_tape = deferred_tapes[0]
@@ -196,7 +196,7 @@ def qnode2(phi, theta):
 
         res2 = qnode2(np.pi / 4, 3 * np.pi / 4)
 
-        assert spy.call_count == 4
+        assert spy.call_count == 2
 
         deferred_tapes1, _ = qml.defer_measurements(qnode1.qtape)
         deferred_tape1 = deferred_tapes1[0]
@@ -1359,7 +1359,7 @@ def qnode(p, x, y):
 
         spy = mocker.spy(qml.defer_measurements, "_transform")
         _ = qnode(0.123, 0.456, 0.789)
-        assert spy.call_count == 2
+        assert spy.call_count == 1
 
         expected_circuit = [
             qml.Hadamard(0),

From dbd2a589b7b50d0eabf4f7deee5ad750f89942a0 Mon Sep 17 00:00:00 2001
From: Mudit Pandey <mudit.pandey@xanadu.ai>
Date: Wed, 18 Oct 2023 15:24:01 -0400
Subject: [PATCH 119/131] Fixed docs; linting

---
 doc/introduction/measurements.rst     | 48 +----------------
 pennylane/measurements/mid_measure.py | 78 +++++++++++++++++++++++++++
 tests/devices/qubit/test_sampling.py  |  2 +-
 3 files changed, 80 insertions(+), 48 deletions(-)

diff --git a/doc/introduction/measurements.rst b/doc/introduction/measurements.rst
index b15b2bd6ceb..03a7d7980e6 100644
--- a/doc/introduction/measurements.rst
+++ b/doc/introduction/measurements.rst
@@ -371,7 +371,7 @@ be equivalent to projecting the state vector onto the :math:`|1\rangle` state on
 
 .. code-block:: python3
 
-    dev = qml.device("default.qubit", wires=2)
+    dev = qml.device("default.qubit")
 
     @qml.qnode(dev)
     def func(x):
@@ -389,57 +389,11 @@ array([1, 1, 1, 1, 1, 1, 1])
 Note that only 7 samples are returned. This is because samples that do not meet the postselection criteria are
 thrown away.
 
-If postselection is requested on a state with zero probability of being measured, the result may contain ``NaN``
-or ``Inf`` values:
-
-.. code-block:: python3
-
-    dev = qml.device("default.qubit", wires=2)
-
-    @qml.qnode(dev)
-    def func(x):
-        qml.RX(x, wires=0)
-        m0 = qml.measure(0, postselect=1)
-        qml.cond(m0, qml.PauliX)(wires=1)
-        return qml.probs(wires=1)
-
->>> func(0.0)
-tensor([nan, nan], requires_grad=True)
-
-In the case of ``qml.sample``, an empty array will be returned:
-
-.. code-block:: python3
-
-    dev = qml.device("default.qubit", wires=2)
-
-    @qml.qnode(dev)
-    def func(x):
-        qml.RX(x, wires=0)
-        m0 = qml.measure(0, postselect=1)
-        qml.cond(m0, qml.PauliX)(wires=1)
-        return qml.sample()
-
->>> func(0.0, shots=[10, 10])
-(array([], dtype=float64), array([], dtype=float64))
-
 .. note::
 
     Currently, postselection support is only available on ``"default.qubit"``. Using postselection
     on other devices will raise an error.
 
-.. warning::
-
-    All measurements are supported when using postselection. However, postselection on a zero probability
-    state can cause some measurements to break.
-
-    With finite shots, one must be careful when measuring ``qml.probs`` or ``qml.counts``, as these
-    measurements will raise errors if there are no valid samples after postselection. This will occur
-    with postselection states that have zero or close to zero probability.
-
-    With analytic execution, ``qml.mutual_info`` will raise errors when using any interfaces except
-    ``jax``, and ``qml.vn_entropy`` will raise an error with the ``tensorflow`` interface when the
-    postselection state has zero probability.
-
 Changing the number of shots
 ----------------------------
 
diff --git a/pennylane/measurements/mid_measure.py b/pennylane/measurements/mid_measure.py
index 7e63eb9e831..b98d5f7ae6f 100644
--- a/pennylane/measurements/mid_measure.py
+++ b/pennylane/measurements/mid_measure.py
@@ -90,6 +90,84 @@ def func():
             measurement. None by default. If postselection is requested, only the post-measurement
             state that is used for postselection will be considered in the remaining circuit.
 
+    .. details::
+        :title: Postselection
+
+        Postselection discards outcomes that do not meet the criteria provided by the ``postselect``
+        argument. For example, specifying ``postselect=1`` on wire 0 would be equivalent to projecting
+        the state vector onto the :math:`|1\rangle` state on wire 0:
+
+        .. code-block:: python3
+
+            dev = qml.device("default.qubit")
+
+            @qml.qnode(dev)
+            def func(x):
+                qml.RX(x, wires=0)
+                m0 = qml.measure(0, postselect=1)
+                qml.cond(m0, qml.PauliX)(wires=1)
+                return qml.sample(wires=1)
+
+        By postselecting on ``1``, we only consider the ``1`` measurement outcome on wire 0. So, the probability of
+        measuring ``1`` on wire 1 after postselection should also be 1. Executing this QNode with 10 shots:
+
+        >>> func(np.pi / 2, shots=10)
+        array([1, 1, 1, 1, 1, 1, 1])
+
+        Note that only 7 samples are returned. This is because samples that do not meet the postselection criteria are
+        thrown away.
+
+        If postselection is requested on a state with zero probability of being measured, the result may contain ``NaN``
+        or ``Inf`` values:
+
+        .. code-block:: python3
+
+            dev = qml.device("default.qubit")
+
+            @qml.qnode(dev)
+            def func(x):
+                qml.RX(x, wires=0)
+                m0 = qml.measure(0, postselect=1)
+                qml.cond(m0, qml.PauliX)(wires=1)
+                return qml.probs(wires=1)
+
+        >>> func(0.0)
+        tensor([nan, nan], requires_grad=True)
+
+        In the case of ``qml.sample``, an empty array will be returned:
+
+        .. code-block:: python3
+
+            dev = qml.device("default.qubit")
+
+            @qml.qnode(dev)
+            def func(x):
+                qml.RX(x, wires=0)
+                m0 = qml.measure(0, postselect=1)
+                qml.cond(m0, qml.PauliX)(wires=1)
+                return qml.sample()
+
+        >>> func(0.0, shots=[10, 10])
+        (array([], dtype=float64), array([], dtype=float64))
+
+        .. note::
+
+            Currently, postselection support is only available on ``"default.qubit"``. Using postselection
+            on other devices will raise an error.
+
+        .. warning::
+
+            All measurements are supported when using postselection. However, postselection on a zero probability
+            state can cause some measurements to break.
+
+            With finite shots, one must be careful when measuring ``qml.probs`` or ``qml.counts``, as these
+            measurements will raise errors if there are no valid samples after postselection. This will occur
+            with postselection states that have zero or close to zero probability.
+
+            With analytic execution, ``qml.mutual_info`` will raise errors when using any interfaces except
+            ``jax``, and ``qml.vn_entropy`` will raise an error with the ``tensorflow`` interface when the
+            postselection state has zero probability.
+
     Returns:
         MidMeasureMP: measurement process instance
 
diff --git a/tests/devices/qubit/test_sampling.py b/tests/devices/qubit/test_sampling.py
index d8af614d37e..1b3d42410c3 100644
--- a/tests/devices/qubit/test_sampling.py
+++ b/tests/devices/qubit/test_sampling.py
@@ -61,7 +61,7 @@ def _valid_flex_tuple(s):
 class _FlexShots(Shots):
     """Shots class that allows zero shots."""
 
-    # pylint: disable=super-init-not-called
+    # pylint: disable=super-init-not-called, too-few-methods
     def __init__(self, shots=None):
         if shots is None:
             self.total_shots = None

From 675a5f7cccd74e75ebab66d64f395d1cb4c9a205 Mon Sep 17 00:00:00 2001
From: Mudit Pandey <mudit.pandey@xanadu.ai>
Date: Wed, 18 Oct 2023 15:40:52 -0400
Subject: [PATCH 120/131] Fixed qnode tests; linting

---
 tests/devices/qubit/test_sampling.py |  4 ++--
 tests/test_qnode.py                  | 12 ++++++++----
 2 files changed, 10 insertions(+), 6 deletions(-)

diff --git a/tests/devices/qubit/test_sampling.py b/tests/devices/qubit/test_sampling.py
index 1b3d42410c3..aa1e14a285d 100644
--- a/tests/devices/qubit/test_sampling.py
+++ b/tests/devices/qubit/test_sampling.py
@@ -58,10 +58,10 @@ def _valid_flex_tuple(s):
     )
 
 
-class _FlexShots(Shots):
+class _FlexShots(Shots):  # pylint: disable=too-few-methods
     """Shots class that allows zero shots."""
 
-    # pylint: disable=super-init-not-called, too-few-methods
+    # pylint: disable=super-init-not-called
     def __init__(self, shots=None):
         if shots is None:
             self.total_shots = None
diff --git a/tests/test_qnode.py b/tests/test_qnode.py
index a528d54f43f..fc4ca54bc8a 100644
--- a/tests/test_qnode.py
+++ b/tests/test_qnode.py
@@ -848,7 +848,11 @@ def circuit(x, y):
         assert np.allclose(res, expected, atol=tol, rtol=0)
 
     @pytest.mark.parametrize(
-        "dev", [qml.device("default.qubit", wires=3), qml.device("default.qubit", wires=3)]
+        "dev, call_count",
+        [
+            (qml.device("default.qubit", wires=3), 2),
+            (qml.device("default.qubit.legacy", wires=3), 1),
+        ],
     )
     @pytest.mark.parametrize("first_par", np.linspace(0.15, np.pi - 0.3, 3))
     @pytest.mark.parametrize("sec_par", np.linspace(0.15, np.pi - 0.3, 3))
@@ -864,7 +868,7 @@ def circuit(x, y):
         ],
     )
     def test_defer_meas_if_mcm_unsupported(
-        self, dev, first_par, sec_par, return_type, mv_return, mv_res, mocker
+        self, dev, call_count, first_par, sec_par, return_type, mv_return, mv_res, mocker
     ):  # pylint: disable=too-many-arguments
         """Tests that the transform using the deferred measurement principle is
         applied if the device doesn't support mid-circuit measurements
@@ -894,7 +898,7 @@ def conditional_ry_qnode(x, y):
 
         assert np.allclose(r1, r2[0])
         assert np.allclose(r2[1], mv_res(first_par))
-        assert spy.call_count == 3  # once for each preprocessing, once for conditional qnode
+        assert spy.call_count == call_count  # once for each preprocessing
 
     def test_drawing_has_deferred_measurements(self):
         """Test that `qml.draw` with qnodes uses defer_measurements
@@ -941,7 +945,7 @@ def conditional_ry_qnode(x):
         r1 = cry_qnode(first_par)
         r2 = conditional_ry_qnode(first_par)
         assert np.allclose(r1, r2)
-        assert spy.call_count == 3  # once per device preprocessing, once for conditional qnode
+        assert spy.call_count == 2  # once per device preprocessing
 
     @pytest.mark.tf
     @pytest.mark.parametrize("interface", ["tf", "auto"])

From 6045241c7d20279a4464f91fc8fa5dcaadbb0d41 Mon Sep 17 00:00:00 2001
From: Mudit Pandey <mudit.pandey@xanadu.ai>
Date: Wed, 18 Oct 2023 17:53:18 -0400
Subject: [PATCH 121/131] Added coverage

---
 pennylane/devices/qubit/simulate.py         | 10 +++++++---
 tests/devices/qubit/test_sampling.py        | 12 +++++++++++-
 tests/test_device.py                        | 10 ++++++++++
 tests/transforms/test_defer_measurements.py | 15 +++++++++++++++
 4 files changed, 43 insertions(+), 4 deletions(-)

diff --git a/pennylane/devices/qubit/simulate.py b/pennylane/devices/qubit/simulate.py
index dee0cefd545..d076cf658b8 100644
--- a/pennylane/devices/qubit/simulate.py
+++ b/pennylane/devices/qubit/simulate.py
@@ -57,8 +57,6 @@ def __init__(self, shots=None):
             self.shot_vector = (qml.measurements.ShotCopies(shots, 1),)
         elif isinstance(shots, Sequence):
             self.__all_tuple_init__([s if isinstance(s, tuple) else (s, 1) for s in shots])
-        else:
-            return  # self already _is_ shots as defined by __new__
 
         self._frozen = True
 
@@ -120,7 +118,13 @@ def _postselection_postprocess(state, is_state_batched, shots):
         # Clip the number of shots using a binomial distribution using the probability of
         # measuring the postselected state.
         postselected_shots = (
-            [binomial(s, float(norm)) for s in shots] if not qml.math.is_abstract(norm) else shots
+            (
+                binomial(shots.total_shots, float(norm))
+                if not shots.has_partitioned_shots
+                else [binomial(s, float(norm)) for s in shots]
+            )
+            if not qml.math.is_abstract(norm)
+            else shots
         )
 
         # _FlexShots is used here since the binomial distribution could result in zero
diff --git a/tests/devices/qubit/test_sampling.py b/tests/devices/qubit/test_sampling.py
index aa1e14a285d..ccf6bb2f4d2 100644
--- a/tests/devices/qubit/test_sampling.py
+++ b/tests/devices/qubit/test_sampling.py
@@ -58,7 +58,7 @@ def _valid_flex_tuple(s):
     )
 
 
-class _FlexShots(Shots):  # pylint: disable=too-few-methods
+class _FlexShots(Shots):  # pylint: disable=too-few-public-methods
     """Shots class that allows zero shots."""
 
     # pylint: disable=super-init-not-called
@@ -536,6 +536,16 @@ def test_measure_with_samples_one_shot_one_wire(self):
 class TestInvalidStateSamples:
     """Tests for state vectors containing nan values or shot vectors with zero shots."""
 
+    @pytest.mark.parametrize("shots", [10, [10, 10]])
+    def test_only_catch_nan_errors(self, shots):
+        """Test that errors are only caught if they are raised due to nan values in the state."""
+        state = np.zeros((2, 2)).astype(np.complex128)
+        mp = qml.expval(qml.PauliZ(0))
+        _shots = Shots(shots)
+
+        with pytest.raises(Exception):
+            _ = measure_with_samples([mp], state, _shots)
+
     @pytest.mark.all_interfaces
     @pytest.mark.parametrize(
         "mp",
diff --git a/tests/test_device.py b/tests/test_device.py
index c3c8b0c253d..5b281ba9cad 100644
--- a/tests/test_device.py
+++ b/tests/test_device.py
@@ -377,6 +377,16 @@ def test_check_validity_on_invalid_observable(self, mock_device_supporting_pauli
         with pytest.raises(DeviceError, match="Observable Hadamard not supported on device"):
             dev.check_validity(queue, observables)
 
+    def test_check_validity_on_projector_as_operation(self, mock_device_with_operations):
+        """Test that an error is raised if the operation queue contains qml.Projector"""
+        dev = mock_device_with_operations(wires=1)
+
+        queue = [qml.PauliX(0), qml.Projector([0], wires=0), qml.PauliZ(0)]
+        observables = []
+
+        with pytest.raises(ValueError, match="Postselection is not supported"):
+            dev.check_validity(queue, observables)
+
     def test_args(self, mock_device):
         """Test that the device requires correct arguments"""
         with pytest.raises(
diff --git a/tests/transforms/test_defer_measurements.py b/tests/transforms/test_defer_measurements.py
index 63f876e72f2..f779c5a94cf 100644
--- a/tests/transforms/test_defer_measurements.py
+++ b/tests/transforms/test_defer_measurements.py
@@ -83,6 +83,21 @@ def circ():
 class TestQNode:
     """Test that the transform integrates well with QNodes."""
 
+    def test_custom_qnode_transform_error(self):
+        """Test that an error is raised if a user tries to give a device argument to the
+        transform when transformingn a qnode."""
+
+        dev = qml.device("default.qubit")
+
+        @qml.qnode(dev)
+        def circ():
+            qml.PauliX(0)
+            qml.measure(0)
+            return qml.probs()
+
+        with pytest.raises(ValueError, match="Cannot provide a 'device'"):
+            _ = qml.defer_measurements(circ, device=dev)
+
     def test_only_mcm(self):
         """Test that a quantum function that only contains one mid-circuit
         measurement yields the correct results and is transformed correctly."""

From 9e149fec13a7919ab2a2939a8655bbf687e9e501 Mon Sep 17 00:00:00 2001
From: Mudit Pandey <mudit.pandey@xanadu.ai>
Date: Wed, 18 Oct 2023 18:40:01 -0400
Subject: [PATCH 122/131] More tests

---
 pennylane/devices/qubit/simulate.py           |  7 +--
 .../default_qubit/test_default_qubit.py       | 50 +++++++++++++++++++
 2 files changed, 54 insertions(+), 3 deletions(-)

diff --git a/pennylane/devices/qubit/simulate.py b/pennylane/devices/qubit/simulate.py
index d076cf658b8..08f511d1723 100644
--- a/pennylane/devices/qubit/simulate.py
+++ b/pennylane/devices/qubit/simulate.py
@@ -14,7 +14,8 @@
 """Simulate a quantum script."""
 # pylint: disable=protected-access
 from typing import Sequence
-from numpy.random import default_rng, binomial
+from numpy.random import default_rng
+import numpy as np
 
 import pennylane as qml
 from pennylane.typing import Result
@@ -119,9 +120,9 @@ def _postselection_postprocess(state, is_state_batched, shots):
         # measuring the postselected state.
         postselected_shots = (
             (
-                binomial(shots.total_shots, float(norm))
+                np.random.binomial(shots.total_shots, float(norm))
                 if not shots.has_partitioned_shots
-                else [binomial(s, float(norm)) for s in shots]
+                else [np.random.binomial(s, float(norm)) for s in shots]
             )
             if not qml.math.is_abstract(norm)
             else shots
diff --git a/tests/devices/default_qubit/test_default_qubit.py b/tests/devices/default_qubit/test_default_qubit.py
index e52595e51ec..cbcbffcf49d 100644
--- a/tests/devices/default_qubit/test_default_qubit.py
+++ b/tests/devices/default_qubit/test_default_qubit.py
@@ -14,6 +14,7 @@
 """Tests for default qubit."""
 # pylint: disable=import-outside-toplevel, no-member, too-many-arguments
 
+from unittest import mock
 import pytest
 
 import numpy as np
@@ -1747,6 +1748,55 @@ def circ_expected():
                 assert qml.math.allclose(r, e, atol=tol_stochastic, rtol=0)
                 assert qml.math.get_interface(r) == qml.math.get_interface(e)
 
+    @pytest.mark.parametrize(
+        "mp, expected_shape",
+        [(qml.sample(wires=[0]), (5,)), (qml.classical_shadow(wires=[0]), (2, 5, 1))],
+    )
+    @pytest.mark.parametrize("param", np.linspace(np.pi / 4, 3 * np.pi / 4, 3))
+    @pytest.mark.parametrize("shots", [10, (10, 10)])
+    def test_postselection_valid_finite_shots_varied_shape(
+        self, mp, param, expected_shape, shots, interface, use_jit
+    ):
+        """Test that qml.sample and qml.classical_shadow work correctly.
+        Separate test because their shape is non-deterministic."""
+
+        if use_jit:
+            pytest.skip("Cannot JIT while mocking function.")
+
+        # with monkeypatch.context() as m:
+        #     m.setattr(np.random, "binomial", lambda *args, **kwargs: 5)
+
+        dev = qml.device("default.qubit", shots=shots, seed=42)
+        param = qml.math.asarray(param, like=interface)
+
+        with mock.patch("numpy.random.binomial", lambda *args, **kwargs: 5):
+
+            @qml.qnode(dev, interface=interface)
+            def circ_postselect(theta):
+                qml.RX(theta, 0)
+                qml.CNOT([0, 1])
+                qml.measure(0, postselect=1)
+                return qml.apply(mp)
+
+            if use_jit:
+                import jax
+
+                circ_postselect = jax.jit(circ_postselect)
+
+            res = circ_postselect(param)
+
+        if not isinstance(shots, tuple):
+            assert qml.math.get_interface(res) == interface if interface != "autograd" else "numpy"
+            assert qml.math.shape(res) == expected_shape
+
+        else:
+            assert isinstance(res, tuple)
+            for r in res:
+                assert (
+                    qml.math.get_interface(r) == interface if interface != "autograd" else "numpy"
+                )
+                assert qml.math.shape(r) == expected_shape
+
     @pytest.mark.parametrize(
         "mp, autograd_interface, is_nan",
         [

From d73a56d29dcbbede7493de6183a4953bbc50f597 Mon Sep 17 00:00:00 2001
From: Mudit Pandey <mudit.pandey@xanadu.ai>
Date: Thu, 19 Oct 2023 16:47:03 -0400
Subject: [PATCH 123/131] [skip ci] Updated measurement docs

---
 doc/introduction/measurements.rst | 113 +++++++++++++++++-------------
 1 file changed, 65 insertions(+), 48 deletions(-)

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

From 4f49f2bb6c1bdc0423f2480e64ffe7dfe18900b5 Mon Sep 17 00:00:00 2001
From: Mudit Pandey <mudit.pandey@xanadu.ai>
Date: Thu, 19 Oct 2023 17:24:49 -0400
Subject: [PATCH 124/131] Addressed PR review; tweaked docs

---
 doc/introduction/measurements.rst             |  9 +++++----
 pennylane/devices/qubit/simulate.py           | 13 ++----------
 pennylane/measurements/mid_measure.py         | 12 +++++------
 pennylane/ops/qubit/observables.py            |  4 ++--
 pennylane/transforms/defer_measurements.py    |  2 +-
 .../default_qubit/test_default_qubit.py       | 20 ++++++++-----------
 6 files changed, 24 insertions(+), 36 deletions(-)

diff --git a/doc/introduction/measurements.rst b/doc/introduction/measurements.rst
index e530bf73997..dc52d0474b1 100644
--- a/doc/introduction/measurements.rst
+++ b/doc/introduction/measurements.rst
@@ -272,7 +272,7 @@ measurement on qubit 1 yielded ``1`` as an outcome, otherwise doing nothing
 for the ``0`` measurement outcome.
 
 PennyLane implements the deferred measurement principle to transform
-conditional operations with the :func:`~.pennylane.defer_measurements` quantum
+conditional operations with the :mod:`~.pennylane.defer_measurements` quantum
 function transform. The deferred measurement principle provides a natural method
 to simulate the application of mid-circuit measurements and conditional operations
 in a differentiable and device-independent way. Performing true mid-circuit
@@ -377,8 +377,8 @@ discarded.
 
 .. note::
 
-    Currently, postselection support is only available on ``"default.qubit"``. Using postselection
-    on other devices will raise an error.
+    Currently, postselection support is only available on :class:`~.pennylane.devices.DefaultQubit`. Using
+    postselection on other devices will raise an error.
 
 Mid-circuit measurement statistics
 **********************************
@@ -386,7 +386,8 @@ Mid-circuit measurement statistics
 Statistics can be collected on mid-circuit measurements along with terminal measurement statistics.
 Currently, ``qml.probs``, ``qml.sample``, ``qml.expval``, ``qml.var``, and ``qml.counts`` are supported,
 and can be requested along with other measurements. The devices that currently support collecting such
-statistics are ``"default.qubit"``, ``"default.mixed"``, and ``"default.qubit.legacy"``.
+statistics are :class:`~.pennylane.devices.DefaultQubit`, :class:`~.pennylane.devices.DefaultMixed`, and
+:class:`~.pennylane.devices.DefaultQubitLegacy`.
 
 .. code-block:: python3
 
diff --git a/pennylane/devices/qubit/simulate.py b/pennylane/devices/qubit/simulate.py
index 08f511d1723..96814f0e9b2 100644
--- a/pennylane/devices/qubit/simulate.py
+++ b/pennylane/devices/qubit/simulate.py
@@ -13,7 +13,6 @@
 # limitations under the License.
 """Simulate a quantum script."""
 # pylint: disable=protected-access
-from typing import Sequence
 from numpy.random import default_rng
 import numpy as np
 
@@ -53,11 +52,7 @@ class _FlexShots(qml.measurements.Shots):
 
     # pylint: disable=super-init-not-called
     def __init__(self, shots=None):
-        if isinstance(shots, int):
-            self.total_shots = shots
-            self.shot_vector = (qml.measurements.ShotCopies(shots, 1),)
-        elif isinstance(shots, Sequence):
-            self.__all_tuple_init__([s if isinstance(s, tuple) else (s, 1) for s in shots])
+        self.__all_tuple_init__([s if isinstance(s, tuple) else (s, 1) for s in shots])
 
         self._frozen = True
 
@@ -119,11 +114,7 @@ def _postselection_postprocess(state, is_state_batched, shots):
         # Clip the number of shots using a binomial distribution using the probability of
         # measuring the postselected state.
         postselected_shots = (
-            (
-                np.random.binomial(shots.total_shots, float(norm))
-                if not shots.has_partitioned_shots
-                else [np.random.binomial(s, float(norm)) for s in shots]
-            )
+            [np.random.binomial(s, float(norm)) for s in shots]
             if not qml.math.is_abstract(norm)
             else shots
         )
diff --git a/pennylane/measurements/mid_measure.py b/pennylane/measurements/mid_measure.py
index b98d5f7ae6f..b36bdb1c285 100644
--- a/pennylane/measurements/mid_measure.py
+++ b/pennylane/measurements/mid_measure.py
@@ -90,6 +90,12 @@ def func():
             measurement. None by default. If postselection is requested, only the post-measurement
             state that is used for postselection will be considered in the remaining circuit.
 
+    Returns:
+        MidMeasureMP: measurement process instance
+
+    Raises:
+        QuantumFunctionError: if multiple wires were specified
+
     .. details::
         :title: Postselection
 
@@ -167,12 +173,6 @@ def func(x):
             With analytic execution, ``qml.mutual_info`` will raise errors when using any interfaces except
             ``jax``, and ``qml.vn_entropy`` will raise an error with the ``tensorflow`` interface when the
             postselection state has zero probability.
-
-    Returns:
-        MidMeasureMP: measurement process instance
-
-    Raises:
-        QuantumFunctionError: if multiple wires were specified
     """
 
     wire = Wires(wires)
diff --git a/pennylane/ops/qubit/observables.py b/pennylane/ops/qubit/observables.py
index c8b92f86bf2..2dbc515d6e1 100644
--- a/pennylane/ops/qubit/observables.py
+++ b/pennylane/ops/qubit/observables.py
@@ -321,7 +321,7 @@ def compute_sparse_matrix(H):  # pylint: disable=arguments-differ
         return H
 
 
-class Projector(Operation, Observable):
+class Projector(Observable):
     r"""Projector(state, wires, id=None)
     Observable corresponding to the state projector :math:`P=\ket{\phi}\bra{\phi}`.
 
@@ -412,7 +412,7 @@ def pow(self, z):
         return [copy(self)] if (isinstance(z, int) and z > 0) else super().pow(z)
 
 
-class BasisStateProjector(Projector):
+class BasisStateProjector(Projector, Operation):
     r"""Observable corresponding to the state projector :math:`P=\ket{\phi}\bra{\phi}`, where
     :math:`\phi` denotes a basis state."""
 
diff --git a/pennylane/transforms/defer_measurements.py b/pennylane/transforms/defer_measurements.py
index e24ba6f0be8..1acf9ec29ec 100644
--- a/pennylane/transforms/defer_measurements.py
+++ b/pennylane/transforms/defer_measurements.py
@@ -192,7 +192,7 @@ def func(x, y):
 
             if op.postselect is not None:
                 with QueuingManager.stop_recording():
-                    new_operations.append(qml.Projector([float(op.postselect)], wires=op.wires[0]))
+                    new_operations.append(qml.Projector([op.postselect], wires=op.wires[0]))
 
             # Store measurement outcome in new wire if wire gets reused
             if op.wires[0] in reused_measurement_wires or op.wires[0] in measured_wires:
diff --git a/tests/devices/default_qubit/test_default_qubit.py b/tests/devices/default_qubit/test_default_qubit.py
index cbcbffcf49d..c56966d4942 100644
--- a/tests/devices/default_qubit/test_default_qubit.py
+++ b/tests/devices/default_qubit/test_default_qubit.py
@@ -1714,7 +1714,7 @@ def test_postselection_valid_finite_shots(
         if use_jit and (interface != "jax" or isinstance(shots, tuple)):
             pytest.skip("Cannot JIT in non-JAX interfaces, or with shot vectors.")
 
-        dev = qml.device("default.qubit", shots=shots)
+        dev = qml.device("default.qubit")
         param = qml.math.asarray(param, like=interface)
 
         @qml.qnode(dev, interface=interface)
@@ -1733,10 +1733,10 @@ def circ_expected():
         if use_jit:
             import jax
 
-            circ_postselect = jax.jit(circ_postselect)
+            circ_postselect = jax.jit(circ_postselect, static_argnames=["shots"])
 
-        res = circ_postselect(param)
-        expected = circ_expected()
+        res = circ_postselect(param, shots=shots)
+        expected = circ_expected(shots=shots)
 
         if not isinstance(shots, tuple):
             assert qml.math.allclose(res, expected, atol=tol_stochastic, rtol=0)
@@ -1763,10 +1763,7 @@ def test_postselection_valid_finite_shots_varied_shape(
         if use_jit:
             pytest.skip("Cannot JIT while mocking function.")
 
-        # with monkeypatch.context() as m:
-        #     m.setattr(np.random, "binomial", lambda *args, **kwargs: 5)
-
-        dev = qml.device("default.qubit", shots=shots, seed=42)
+        dev = qml.device("default.qubit", seed=42)
         param = qml.math.asarray(param, like=interface)
 
         with mock.patch("numpy.random.binomial", lambda *args, **kwargs: 5):
@@ -1781,9 +1778,9 @@ def circ_postselect(theta):
             if use_jit:
                 import jax
 
-                circ_postselect = jax.jit(circ_postselect)
+                circ_postselect = jax.jit(circ_postselect, static_argnames=["shots"])
 
-            res = circ_postselect(param)
+            res = circ_postselect(param, shots=shots)
 
         if not isinstance(shots, tuple):
             assert qml.math.get_interface(res) == interface if interface != "autograd" else "numpy"
@@ -1871,8 +1868,7 @@ def test_postselection_invalid_finite_shots(
         if use_jit and interface != "jax":
             pytest.skip("Can't jit with non-jax interfaces.")
 
-        # 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.qnode(dev, interface=interface)
         def circ():

From 4fe2e25e59bd4ce2f684a96f46ad1b34cd482932 Mon Sep 17 00:00:00 2001
From: Mudit Pandey <mudit.pandey@xanadu.ai>
Date: Thu, 19 Oct 2023 17:46:19 -0400
Subject: [PATCH 125/131] Updated tests

---
 pennylane/devices/qubit/simulate.py  |  6 +++++-
 tests/devices/qubit/test_sampling.py | 27 +--------------------------
 2 files changed, 6 insertions(+), 27 deletions(-)

diff --git a/pennylane/devices/qubit/simulate.py b/pennylane/devices/qubit/simulate.py
index 96814f0e9b2..2716712b863 100644
--- a/pennylane/devices/qubit/simulate.py
+++ b/pennylane/devices/qubit/simulate.py
@@ -52,7 +52,11 @@ class _FlexShots(qml.measurements.Shots):
 
     # pylint: disable=super-init-not-called
     def __init__(self, shots=None):
-        self.__all_tuple_init__([s if isinstance(s, tuple) else (s, 1) for s in shots])
+        if isinstance(shots, int):
+            self.total_shots = shots
+            self.shot_vector = (qml.measurements.ShotCopies(shots, 1),)
+        else:
+            self.__all_tuple_init__([s if isinstance(s, tuple) else (s, 1) for s in shots])
 
         self._frozen = True
 
diff --git a/tests/devices/qubit/test_sampling.py b/tests/devices/qubit/test_sampling.py
index ccf6bb2f4d2..27e578b05ca 100644
--- a/tests/devices/qubit/test_sampling.py
+++ b/tests/devices/qubit/test_sampling.py
@@ -14,13 +14,13 @@
 """Unit tests for sample_state in devices/qubit."""
 
 from random import shuffle
-from typing import Sequence
 
 import pytest
 
 import pennylane as qml
 from pennylane import numpy as np
 from pennylane.devices.qubit import simulate
+from pennylane.devices.qubit.simulate import _FlexShots
 from pennylane.devices.qubit import sample_state, measure_with_samples
 from pennylane.devices.qubit.sampling import _sample_state_jax
 from pennylane.measurements import Shots
@@ -58,31 +58,6 @@ def _valid_flex_tuple(s):
     )
 
 
-class _FlexShots(Shots):  # pylint: disable=too-few-public-methods
-    """Shots class that allows zero shots."""
-
-    # pylint: disable=super-init-not-called
-    def __init__(self, shots=None):
-        if shots is None:
-            self.total_shots = None
-            self.shot_vector = ()
-        elif isinstance(shots, int):
-            if shots < 0:
-                raise self._SHOT_ERROR
-            self.total_shots = shots
-            self.shot_vector = (qml.measurements.ShotCopies(shots, 1),)
-        elif isinstance(shots, Sequence):
-            if not all(_valid_flex_int(s) or _valid_flex_tuple(s) for s in shots):
-                raise self._SHOT_ERROR
-            self.__all_tuple_init__([s if isinstance(s, tuple) else (s, 1) for s in shots])
-        elif isinstance(shots, self.__class__):
-            return  # self already _is_ shots as defined by __new__
-        else:
-            raise self._SHOT_ERROR
-
-        self._frozen = True
-
-
 def samples_to_probs(samples, num_wires):
     """Converts samples to probs"""
     samples_decimal = [np.ravel_multi_index(sample, [2] * num_wires) for sample in samples]

From f90a889f49a32f7d01dc178a58f1ff818cdd3262 Mon Sep 17 00:00:00 2001
From: Mudit Pandey <mudit.pandey@xanadu.ai>
Date: Fri, 20 Oct 2023 11:02:14 -0400
Subject: [PATCH 126/131] [skip ci] Fixing docs

---
 doc/introduction/measurements.rst | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/doc/introduction/measurements.rst b/doc/introduction/measurements.rst
index dc52d0474b1..9353c8a0c7e 100644
--- a/doc/introduction/measurements.rst
+++ b/doc/introduction/measurements.rst
@@ -272,7 +272,7 @@ measurement on qubit 1 yielded ``1`` as an outcome, otherwise doing nothing
 for the ``0`` measurement outcome.
 
 PennyLane implements the deferred measurement principle to transform
-conditional operations with the :mod:`~.pennylane.defer_measurements` quantum
+conditional operations with the :func:`~.pennylane.defer_measurements` quantum
 function transform. The deferred measurement principle provides a natural method
 to simulate the application of mid-circuit measurements and conditional operations
 in a differentiable and device-independent way. Performing true mid-circuit

From d2877981edeed378427049d67b82354960408207 Mon Sep 17 00:00:00 2001
From: Mudit Pandey <mudit.pandey@xanadu.ai>
Date: Fri, 20 Oct 2023 13:38:42 -0400
Subject: [PATCH 127/131] [skip ci] Added errors

---
 doc/releases/changelog-dev.md               |   6 +-
 pennylane/transforms/defer_measurements.py  | 111 +++++++++++++-------
 tests/transforms/test_defer_measurements.py |  23 +++-
 3 files changed, 99 insertions(+), 41 deletions(-)

diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 7bba23f3c5c..8043266f2b6 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -296,10 +296,12 @@
   decomposition.
   [(#4675)](https://github.com/PennyLaneAI/pennylane/pull/4675)
 
-
-
 <h3>Breaking changes 💔</h3>
 
+* ``qml.defer_measurements`` now raises an error if a transformed circuit measures ``qml.probs``,
+  ``qml.sample``, or ``qml.counts`` without any wires or obsrvable, or if it measures ``qml.state``.
+  [(#)]()
+
 * The device test suite now converts device kwargs to integers or floats if they can be converted to integers or floats.
   [(#4640)](https://github.com/PennyLaneAI/pennylane/pull/4640)
 
diff --git a/pennylane/transforms/defer_measurements.py b/pennylane/transforms/defer_measurements.py
index 1acf9ec29ec..ba676c73338 100644
--- a/pennylane/transforms/defer_measurements.py
+++ b/pennylane/transforms/defer_measurements.py
@@ -14,7 +14,7 @@
 """Code for the tape transform implementing the deferred measurement principle."""
 from typing import Sequence, Callable
 import pennylane as qml
-from pennylane.measurements import MidMeasureMP
+from pennylane.measurements import MidMeasureMP, ProbabilityMP, SampleMP, CountsMP
 from pennylane.ops.op_math import ctrl
 
 from pennylane.tape import QuantumTape
@@ -23,7 +23,62 @@
 from pennylane.wires import Wires
 from pennylane.queuing import QueuingManager
 
-# pylint: disable=too-many-branches, too-many-statements
+# pylint: disable=too-many-branches, protected-access
+
+
+def _check_tape_validity(tape: QuantumTape):
+    """Helper function to check that the tape is valid."""
+    cv_types = (qml.operation.CVOperation, qml.operation.CVObservable)
+    ops_cv = any(isinstance(op, cv_types) and op.name != "Identity" for op in tape.operations)
+    obs_cv = any(
+        isinstance(getattr(op, "obs", None), cv_types)
+        and not isinstance(getattr(op, "obs", None), qml.Identity)
+        for op in tape.measurements
+    )
+    if ops_cv or obs_cv:
+        raise ValueError("Continuous variable operations and observables are not supported.")
+
+    for mp in tape.measurements:
+        if isinstance(mp, (CountsMP, ProbabilityMP, SampleMP)) and not (
+            mp.obs or mp._wires or mp.mv
+        ):
+            raise ValueError(
+                f"Cannot use {mp.__class__.__name__} as a measurement without specifying wires "
+                "when using qml.defer_measurements."
+            )
+
+        if mp.__class__.__name__ == "StateMP":
+            raise ValueError(
+                "Cannot use StateMP as a measurement when using qml.defer_measurements."
+            )
+
+
+def _collect_mid_measure_info(tape: QuantumTape):
+    """Helper function to collect information related to mid-circuit measurements in the tape."""
+
+    # Find wires that are reused after measurement
+    measured_wires = []
+    reused_measurement_wires = set()
+    any_repeated_measurements = False
+    is_postselecting = False
+
+    for op in tape.operations:
+        if isinstance(op, MidMeasureMP):
+            if op.postselect is not None:
+                is_postselecting = True
+            if op.reset:
+                reused_measurement_wires.add(op.wires[0])
+
+            if op.wires[0] in measured_wires:
+                any_repeated_measurements = True
+            measured_wires.append(op.wires[0])
+
+        else:
+            reused_measurement_wires = reused_measurement_wires.union(
+                set(measured_wires).intersection(op.wires.toset())
+            )
+
+    return measured_wires, reused_measurement_wires, any_repeated_measurements, is_postselecting
 
 
 def null_postprocessing(results):
@@ -68,12 +123,18 @@ def defer_measurements(tape: QuantumTape, **kwargs) -> (Sequence[QuantumTape], C
 
     .. note::
 
-        When applying the transform on a quantum function that returns
-        :func:`~pennylane.state` as the terminal measurement or contains the
+        When applying the transform on a quantum function that contains the
         :class:`~.Snapshot` instruction, state information corresponding to
         simulating the transformed circuit will be obtained. No
         post-measurement states are considered.
 
+    .. warning::
+
+        :func:`~.pennylane.state` is not supported with the ``defer_measurements`` transform.
+        Additionally, :func:`~.pennylane.probs`, :func:`~.pennylane.sample` and
+        :func:`~.pennylane.counts` can only be used with ``defer_measurements`` if wires
+        or an observable are explicitly specified.
+
     Args:
         tape (.QuantumTape): a quantum tape
 
@@ -136,45 +197,23 @@ def func(x, y):
     >>> func(*pars)
     tensor([0.76960924, 0.13204407, 0.08394415, 0.01440254], requires_grad=True)
     """
-    # pylint: disable=protected-access
+
     if not any(isinstance(o, MidMeasureMP) for o in tape.operations):
         return (tape,), null_postprocessing
 
-    cv_types = (qml.operation.CVOperation, qml.operation.CVObservable)
-    ops_cv = any(isinstance(op, cv_types) and op.name != "Identity" for op in tape.operations)
-    obs_cv = any(
-        isinstance(getattr(op, "obs", None), cv_types)
-        and not isinstance(getattr(op, "obs", None), qml.Identity)
-        for op in tape.measurements
-    )
-    if ops_cv or obs_cv:
-        raise ValueError("Continuous variable operations and observables are not supported.")
+    _check_tape_validity(tape)
 
     device = kwargs.get("device", None)
 
     new_operations = []
 
     # Find wires that are reused after measurement
-    measured_wires = []
-    reused_measurement_wires = set()
-    repeated_measurement_wire = False
-    is_postselecting = False
-
-    for op in tape.operations:
-        if isinstance(op, MidMeasureMP):
-            if op.postselect is not None:
-                is_postselecting = True
-            if op.reset:
-                reused_measurement_wires.add(op.wires[0])
-
-            if op.wires[0] in measured_wires:
-                repeated_measurement_wire = True
-            measured_wires.append(op.wires[0])
-
-        else:
-            reused_measurement_wires = reused_measurement_wires.union(
-                set(measured_wires).intersection(op.wires.toset())
-            )
+    (
+        measured_wires,
+        reused_measurement_wires,
+        any_repeated_measurements,
+        is_postselecting,
+    ) = _collect_mid_measure_info(tape)
 
     if is_postselecting and device is not None and not isinstance(device, qml.devices.DefaultQubit):
         raise ValueError(f"Postselection is not supported on the {device} device.")
@@ -183,7 +222,7 @@ def func(x, y):
     # classically controlled operations
     control_wires = {}
     cur_wire = (
-        max(tape.wires) + 1 if reused_measurement_wires or repeated_measurement_wire else None
+        max(tape.wires) + 1 if reused_measurement_wires or any_repeated_measurements else None
     )
 
     for op in tape.operations:
@@ -257,7 +296,7 @@ def _add_control_gate(op, control_wires):
     control = [control_wires[m.id] for m in op.meas_val.measurements]
     new_ops = []
 
-    for branch, value in op.meas_val._items():  # pylint: disable=protected-access
+    for branch, value in op.meas_val._items():
         if value:
             qscript = qml.tape.make_qscript(
                 ctrl(
diff --git a/tests/transforms/test_defer_measurements.py b/tests/transforms/test_defer_measurements.py
index f779c5a94cf..ec81062f610 100644
--- a/tests/transforms/test_defer_measurements.py
+++ b/tests/transforms/test_defer_measurements.py
@@ -80,6 +80,23 @@ def circ():
         _ = circ()
 
 
+@pytest.mark.parametrize(
+    "mp, err_msg",
+    [
+        (qml.state(), "Cannot use StateMP as a measurement when"),
+        (qml.probs(), "Cannot use ProbabilityMP as a measurement without"),
+        (qml.sample(), "Cannot use SampleMP as a measurement without"),
+        (qml.counts(), "Cannot use CountsMP as a measurement without"),
+    ],
+)
+def test_unsupported_measurements(mp, err_msg):
+    """Test that using unsupported measurements raises an error."""
+    tape = qml.tape.QuantumScript([MidMeasureMP(0)], [mp])
+
+    with pytest.raises(ValueError, match=err_msg):
+        _, _ = qml.defer_measurements(tape)
+
+
 class TestQNode:
     """Test that the transform integrates well with QNodes."""
 
@@ -1467,7 +1484,7 @@ def test_custom_wire_labels_allowed_without_reset():
         qml.Hadamard("a")
         ma = qml.measure("a", reset=False)
         qml.cond(ma, qml.PauliX)("b")
-        qml.state()
+        qml.probs(wires="a")
 
     tape = qml.tape.QuantumScript.from_queue(q)
     tapes, _ = qml.defer_measurements(tape)
@@ -1476,7 +1493,7 @@ def test_custom_wire_labels_allowed_without_reset():
     assert len(tape) == 3
     assert qml.equal(tape[0], qml.Hadamard("a"))
     assert qml.equal(tape[1], qml.CNOT(["a", "b"]))
-    assert qml.equal(tape[2], qml.state())
+    assert qml.equal(tape[2], qml.probs(wires="a"))
 
 
 def test_custom_wire_labels_fails_with_reset():
@@ -1485,7 +1502,7 @@ def test_custom_wire_labels_fails_with_reset():
         qml.Hadamard("a")
         ma = qml.measure("a", reset=True)
         qml.cond(ma, qml.PauliX)("b")
-        qml.state()
+        qml.probs(wires="a")
 
     tape = qml.tape.QuantumScript.from_queue(q)
     with pytest.raises(TypeError, match="can only concatenate str"):

From 5066c28a5cce2bba4a0b41b990fb45d4ae5e3af9 Mon Sep 17 00:00:00 2001
From: Mudit Pandey <mudit.pandey@xanadu.ai>
Date: Fri, 20 Oct 2023 14:49:40 -0400
Subject: [PATCH 128/131] [skip ci] Removed test case with `qml.state()`

---
 tests/devices/default_qubit/test_default_qubit.py | 2 --
 1 file changed, 2 deletions(-)

diff --git a/tests/devices/default_qubit/test_default_qubit.py b/tests/devices/default_qubit/test_default_qubit.py
index c56966d4942..257d5708984 100644
--- a/tests/devices/default_qubit/test_default_qubit.py
+++ b/tests/devices/default_qubit/test_default_qubit.py
@@ -1650,7 +1650,6 @@ class TestPostselection:
             qml.expval(qml.PauliZ(0)),
             qml.var(qml.PauliZ(0)),
             qml.probs(wires=[0, 1]),
-            qml.state(),
             qml.density_matrix(wires=0),
             qml.purity(0),
             qml.vn_entropy(0),
@@ -1800,7 +1799,6 @@ def circ_postselect(theta):
             (qml.expval(qml.PauliZ(0)), "autograd", True),
             (qml.var(qml.PauliZ(0)), "autograd", True),
             (qml.probs(wires=[0, 1]), "autograd", True),
-            (qml.state(), "autograd", True),
             (qml.density_matrix(wires=0), "autograd", True),
             (qml.purity(0), "numpy", True),
             (qml.vn_entropy(0), "numpy", False),

From 1c5724c04a7f47875a14e696c76e49108eccc005 Mon Sep 17 00:00:00 2001
From: Mudit Pandey <mudit.pandey@xanadu.ai>
Date: Fri, 20 Oct 2023 15:57:54 -0400
Subject: [PATCH 129/131] [skip ci] Added link to changelog

---
 doc/releases/changelog-dev.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 8043266f2b6..885a332fc6e 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -300,7 +300,7 @@
 
 * ``qml.defer_measurements`` now raises an error if a transformed circuit measures ``qml.probs``,
   ``qml.sample``, or ``qml.counts`` without any wires or obsrvable, or if it measures ``qml.state``.
-  [(#)]()
+  [(#4701)](https://github.com/PennyLaneAI/pennylane/pull/4701)
 
 * The device test suite now converts device kwargs to integers or floats if they can be converted to integers or floats.
   [(#4640)](https://github.com/PennyLaneAI/pennylane/pull/4640)

From 0c53d59307add2d00363570433133c804b8ea371 Mon Sep 17 00:00:00 2001
From: Mudit Pandey <mudit.pandey@xanadu.ai>
Date: Fri, 20 Oct 2023 17:00:53 -0400
Subject: [PATCH 130/131] [skip ci] Reordered preprocessing steps

---
 pennylane/devices/default_qubit.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pennylane/devices/default_qubit.py b/pennylane/devices/default_qubit.py
index 52bcd90470b..c1a21cbb4fe 100644
--- a/pennylane/devices/default_qubit.py
+++ b/pennylane/devices/default_qubit.py
@@ -387,8 +387,8 @@ def preprocess(
         config = self._setup_execution_config(execution_config)
         transform_program = TransformProgram()
 
-        transform_program.add_transform(qml.defer_measurements, device=self)
         transform_program.add_transform(validate_device_wires, self.wires, name=self.name)
+        transform_program.add_transform(qml.defer_measurements, device=self)
         transform_program.add_transform(
             decompose, stopping_condition=stopping_condition, name=self.name
         )

From 48f56bf1e5e92d1c8c9ababa5ef1f673b33dd3c6 Mon Sep 17 00:00:00 2001
From: Mudit Pandey <mudit.pandey@xanadu.ai>
Date: Mon, 23 Oct 2023 14:20:45 -0400
Subject: [PATCH 131/131] [skip ci] Apply suggestions from code review

Co-authored-by: Tom Bromley <49409390+trbromley@users.noreply.github.com>
---
 pennylane/transforms/defer_measurements.py | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/pennylane/transforms/defer_measurements.py b/pennylane/transforms/defer_measurements.py
index ba676c73338..180df6cf432 100644
--- a/pennylane/transforms/defer_measurements.py
+++ b/pennylane/transforms/defer_measurements.py
@@ -44,12 +44,16 @@ def _check_tape_validity(tape: QuantumTape):
         ):
             raise ValueError(
                 f"Cannot use {mp.__class__.__name__} as a measurement without specifying wires "
-                "when using qml.defer_measurements."
+                "when using qml.defer_measurements. Deferred measurements can occur "
+                "automatically when using mid-circuit measurements on a device that does not "
+                "support them."
             )
 
         if mp.__class__.__name__ == "StateMP":
             raise ValueError(
-                "Cannot use StateMP as a measurement when using qml.defer_measurements."
+                "Cannot use StateMP as a measurement when using qml.defer_measurements. "
+                "Deferred measurements can occur automatically when using mid-circuit "
+                "measurements on a device that does not support them."
             )