Repository: PennyLaneAI/pennylane
Pull request #4727 (2023-10-25 02:51:27)

Daily rc sync to master
Automatic sync from the release candidate to master during a feature freeze.

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4727?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Comparison is base [(`6ecea1f`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/6ecea1f27ec939af4dd9441963dbde7c81b974db?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) 99.64% compared to head [(`742c448`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4727?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) 99.64%.

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


```diff
@@            Coverage Diff             @@
##           master    #4727      +/-   ##
==========================================
- Coverage   99.64%   99.64%   -0.01%     
==========================================
  Files         380      380              
  Lines       34245    33994     -251     
==========================================
- Hits        34124    33872     -252     
- Misses        121      122       +1     
```


| [Files](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4727?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | Coverage Δ | |
|---|---|---|
| [pennylane/\_device.py](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4727?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL19kZXZpY2UucHk=) | `99.41% <ø> (-0.01%)` | :arrow_down: |
| [pennylane/\_grad.py](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4727?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL19ncmFkLnB5) | `100.00% <ø> (ø)` | |
| [pennylane/debugging.py](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4727?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2RlYnVnZ2luZy5weQ==) | `100.00% <ø> (ø)` | |
| [pennylane/devices/\_\_init\_\_.py](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4727?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2RldmljZXMvX19pbml0X18ucHk=) | `100.00% <ø> (ø)` | |
| [pennylane/devices/default\_qubit.py](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4727?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2RldmljZXMvZGVmYXVsdF9xdWJpdC5weQ==) | `100.00% <ø> (ø)` | |
| [pennylane/devices/default\_qubit\_legacy.py](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4727?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2RldmljZXMvZGVmYXVsdF9xdWJpdF9sZWdhY3kucHk=) | `99.71% <ø> (ø)` | |
| [pennylane/devices/device\_api.py](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4727?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2RldmljZXMvZGV2aWNlX2FwaS5weQ==) | `98.57% <ø> (-0.03%)` | :arrow_down: |
| [pennylane/fourier/reconstruct.py](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4727?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2ZvdXJpZXIvcmVjb25zdHJ1Y3QucHk=) | `100.00% <ø> (ø)` | |
| [pennylane/fourier/visualize.py](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4727?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2ZvdXJpZXIvdmlzdWFsaXplLnB5) | `100.00% <ø> (ø)` | |
| [pennylane/gradients/\_\_init\_\_.py](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4727?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2dyYWRpZW50cy9fX2luaXRfXy5weQ==) | `100.00% <ø> (ø)` | |
| ... and [50 more](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4727?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | |

... and [33 files with indirect coverage changes](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4727/indirect-changes?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI)


</details>

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


Patch:
diff --git a/doc/development/guide/installation.rst b/doc/development/guide/installation.rst
index 85842cd99..3d3ba45ce 100644
--- a/doc/development/guide/installation.rst
+++ b/doc/development/guide/installation.rst
@@ -18,7 +18,7 @@ be installed alongside PennyLane:
 * `toml <https://github.com/uiri/toml>`_
 * `appdirs <https://github.com/ActiveState/appdirs>`_
 * `semantic-version <https://github.com/rbarrois/python-semanticversion>`_ >= 2.7
-* `autoray <https://github.com/jcmgray/autoray>`__
+* `autoray <https://github.com/jcmgray/autoray>`__ >= 0.6.1
 
 The following Python packages are optional:
 
diff --git a/doc/releases/changelog-0.33.0.md b/doc/releases/changelog-0.33.0.md
index 771ba4c2a..553087421 100644
--- a/doc/releases/changelog-0.33.0.md
+++ b/doc/releases/changelog-0.33.0.md
@@ -537,6 +537,8 @@
 
 <h3>Bug fixes 🐛</h3>
 
+* Fixes `LocalHilbertSchmidt.compute_decomposition` so the template can be used in a qnode.
+
 * Providing `work_wires=None` to `qml.GroverOperator` no longer interprets `None` as a wire.
   [(#4668)](https://github.com/PennyLaneAI/pennylane/pull/4668)
 
diff --git a/pennylane/_device.py b/pennylane/_device.py
index 724ae14f8..ba088955e 100644
--- a/pennylane/_device.py
+++ b/pennylane/_device.py
@@ -293,7 +293,7 @@ class Device(abc.ABC):
 
         **Example**
 
-        >>> dev = qml.device("default.qubit", wires=2, shots=[3, 1, 2, 2, 2, 2, 6, 1, 1, 5, 12, 10, 10])
+        >>> dev = qml.device("default.qubit.legacy", wires=2, shots=[3, 1, 2, 2, 2, 2, 6, 1, 1, 5, 12, 10, 10])
         >>> dev.shots
         57
         >>> dev.shot_vector
@@ -618,7 +618,7 @@ class Device(abc.ABC):
 
         .. code-block:: python
 
-            dev = qml.device("default.qubit", wires=2)
+            dev = qml.device("default.qubit.legacy", wires=2)
 
             @dev.custom_expand
             def my_expansion_function(self, tape, max_expansion=10):
diff --git a/pennylane/_grad.py b/pennylane/_grad.py
index 6769548d6..205d1bc2c 100644
--- a/pennylane/_grad.py
+++ b/pennylane/_grad.py
@@ -189,18 +189,18 @@ def jacobian(func, argnum=None):
             qml.RX(weights[0, 0, 0], wires=0)
             qml.RY(weights[0, 0, 1], wires=1)
             qml.RZ(weights[1, 0, 2], wires=0)
-            return tuple(qml.expval(qml.PauliZ(w)) for w in dev.wires)
+            return qml.probs()
 
         weights = np.array(
             [[[0.2, 0.9, -1.4]], [[0.5, 0.2, 0.1]]], requires_grad=True
         )
 
     It has a single array-valued QNode argument with shape ``(2, 1, 3)`` and outputs
-    a tuple of two expectation values. Therefore, the Jacobian of this QNode
-    will be a single array with shape ``(2, 2, 1, 3)``:
+    the probability of each 2-wire basis state, of which there are ``2**num_wires`` = 4.
+    Therefore, the Jacobian of this QNode will be a single array with shape ``(2, 2, 1, 3)``:
 
     >>> qml.jacobian(circuit)(weights).shape
-    (2, 2, 1, 3)
+    (4, 2, 1, 3)
 
     On the other hand, consider the following QNode for the same circuit
     structure:
@@ -220,18 +220,18 @@ def jacobian(func, argnum=None):
         y = np.array(0.9, requires_grad=True)
         z = np.array(-1.4, requires_grad=True)
 
-    It has three scalar QNode arguments and outputs a tuple of two expectation
-    values. Consequently, its Jacobian will be a three-tuple of arrays with the
-    shape ``(2,)``:
+    It has three scalar QNode arguments and outputs the probability for each of
+    the 4 basis states. Consequently, its Jacobian will be a three-tuple of
+    arrays with the shape ``(4,)``:
 
     >>> jac = qml.jacobian(circuit)(x, y, z)
     >>> type(jac)
     tuple
     >>> for sub_jac in jac:
     ...     print(sub_jac.shape)
-    (2,)
-    (2,)
-    (2,)
+    (4,)
+    (4,)
+    (4,)
 
     For a more advanced setting of QNode arguments, consider the QNode
 
@@ -256,14 +256,14 @@ def jacobian(func, argnum=None):
     >>> print(type(jac), len(jac))
     <class 'tuple'> 2
     >>> qml.math.shape(jac[0])
-    (3, 2)
+    (8, 2)
     >>> qml.math.shape(jac[1])
-    (3, 2, 4)
+    (8, 2, 4)
 
     As we can see, there are two entries in the output, one Jacobian for each
-    QNode argument. The shape ``(3, 2)`` of the first Jacobian is the combination
-    of the QNode output shape (``(3,)``) and the shape of ``x`` (``(2,)``).
-    Similarly, the shape ``(2, 4)`` of ``y`` leads to a Jacobian shape ``(3, 2, 4)``.
+    QNode argument. The shape ``(8, 2)`` of the first Jacobian is the combination
+    of the QNode output shape (``(8,)``) and the shape of ``x`` (``(2,)``).
+    Similarly, the shape ``(2, 4)`` of ``y`` leads to a Jacobian shape ``(8, 2, 4)``.
 
     Instead we may choose the output to contain only one of the two
     entries by providing an iterable as ``argnum``:
@@ -272,7 +272,7 @@ def jacobian(func, argnum=None):
     >>> print(type(jac), len(jac))
     <class 'tuple'> 1
     >>> qml.math.shape(jac)
-    (1, 3, 2, 4)
+    (1, 8, 2, 4)
 
     Here we included the size of the tuple in the shape analysis, corresponding to the
     first dimension of size ``1``.
@@ -282,9 +282,9 @@ def jacobian(func, argnum=None):
 
     >>> jac = qml.jacobian(circuit, argnum=1)(x, y)
     >>> print(type(jac), len(jac))
-    <class 'numpy.ndarray'> 3
+    <class 'numpy.ndarray'> 8
     >>> qml.math.shape(jac)
-    (3, 2, 4)
+    (8, 2, 4)
 
     As expected, the tuple was unpacked and we directly received the Jacobian of the
     QNode with respect to ``y``.
diff --git a/pennylane/debugging.py b/pennylane/debugging.py
index eded61f7e..c8327f2c9 100644
--- a/pennylane/debugging.py
+++ b/pennylane/debugging.py
@@ -80,10 +80,10 @@ def snapshots(qnode):
             return qml.expval(qml.PauliX(0))
 
     >>> qml.snapshots(circuit)()
-    {0: array([1.+0.j, 0.+0.j, 0.+0.j, 0.+0.j]),
-    'very_important_state': array([0.70710678+0.j, 0.+0.j, 0.70710678+0.j, 0.+0.j]),
-    2: array([0.70710678+0.j, 0.+0.j, 0.+0.j, 0.70710678+0.j]),
-    'execution_results': array(0.)}
+    {0: array([1., 0., 0., 0.]),
+    'very_important_state': array([0.70710678, 0.        , 0.70710678, 0.        ]),
+    2: array([0.70710678, 0.        , 0.        , 0.70710678]),
+    'execution_results': 0.0}
     """
 
     def get_snapshots(*args, **kwargs):
diff --git a/pennylane/devices/__init__.py b/pennylane/devices/__init__.py
index 49f282119..1f34318ad 100644
--- a/pennylane/devices/__init__.py
+++ b/pennylane/devices/__init__.py
@@ -25,6 +25,7 @@ to verify and test quantum gradient computations.
     :toctree: api
 
 
+    default_qubit
     default_qubit_legacy
     default_qubit_jax
     default_qubit_torch
diff --git a/pennylane/devices/default_qubit.py b/pennylane/devices/default_qubit.py
index c1a21cbb4..8ad55608e 100644
--- a/pennylane/devices/default_qubit.py
+++ b/pennylane/devices/default_qubit.py
@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 """
-This module contains the next generation successor to default qubit
+The default.qubit device is PennyLane's standard qubit-based device.
 """
 
 from dataclasses import replace
diff --git a/pennylane/devices/default_qubit_legacy.py b/pennylane/devices/default_qubit_legacy.py
index e7842935d..3f2c5498c 100644
--- a/pennylane/devices/default_qubit_legacy.py
+++ b/pennylane/devices/default_qubit_legacy.py
@@ -12,7 +12,7 @@
 # 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.
+This module contains the legacy implementation of default.qubit.
 
 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
diff --git a/pennylane/devices/device_api.py b/pennylane/devices/device_api.py
index 9c07e6c4d..749f0a186 100644
--- a/pennylane/devices/device_api.py
+++ b/pennylane/devices/device_api.py
@@ -567,7 +567,7 @@ class Device(abc.ABC):
         Returns:
             Tuple, Tuple: A numeric result of execution and of computing the jacobian vector product
 
-        .. seealso:: :meth:`~.Device.execute` and :meth:`~.Device.compute_jvp`
+        .. seealso:: :meth:`~pennylane.devices.Device.execute` and :meth:`~.Device.compute_jvp`
         """
         return self.execute(circuits, execution_config), self.compute_jvp(
             circuits, tangents, execution_config
@@ -647,7 +647,7 @@ class Device(abc.ABC):
         Returns:
             Tuple, Tuple: the result of executing the scripts and the numeric result of computing the vector jacobian product
 
-        .. seealso:: :meth:`~.Device.execute` and :meth:`~.Device.compute_vjp`
+        .. seealso:: :meth:`~pennylane.devices.Device.execute` and :meth:`~.Device.compute_vjp`
         """
         return self.execute(circuits, execution_config), self.compute_vjp(
             circuits, cotangents, execution_config
diff --git a/pennylane/fourier/reconstruct.py b/pennylane/fourier/reconstruct.py
index b2c55fa65..6fd56576c 100644
--- a/pennylane/fourier/reconstruct.py
+++ b/pennylane/fourier/reconstruct.py
@@ -448,7 +448,7 @@ def reconstruct(qnode, ids=None, nums_frequency=None, spectra=None, shifts=None)
     reconstructions:
 
     >>> tracker.totals
-    {'executions': 15}
+    {'batches': 15, 'simulations': 15, 'executions': 15}
 
     The example above used that we already knew the frequency spectra of the
     QNode of interest. However, this is in general not the case and we may need
diff --git a/pennylane/fourier/visualize.py b/pennylane/fourier/visualize.py
index b63f07544..733f01e59 100644
--- a/pennylane/fourier/visualize.py
+++ b/pennylane/fourier/visualize.py
@@ -174,8 +174,9 @@ def violin(coeffs, n_inputs, ax, colour_dict=None, show_freqs=True):
     to the plotting function:
 
     >>> import matplotlib.pyplot as plt
+    >>> from pennylane.fourier.visualize import violin
     >>> fig, ax = plt.subplots(2, 1, sharey=True, figsize=(15, 4))
-    >>> violinplt(coeffs, n_inputs, ax, show_freqs=True)
+    >>> violin(coeffs, n_inputs, ax, show_freqs=True)
 
     .. image:: ../../_static/fourier_vis_violin.png
         :align: center
@@ -279,6 +280,7 @@ def box(coeffs, n_inputs, ax, colour_dict=None, show_freqs=True, show_fliers=Tru
     to the plotting function:
 
     >>> import matplotlib.pyplot as plt
+    >>> from pennylane.fourier.visualize import box
     >>> fig, ax = plt.subplots(2, 1, sharey=True, figsize=(15, 4))
     >>> box(coeffs, n_inputs, ax, show_freqs=True)
 
@@ -391,6 +393,7 @@ def bar(coeffs, n_inputs, ax, colour_dict=None, show_freqs=True):
     to the plotting function:
 
     >>> import matplotlib.pyplot as plt
+    >>> from pennylane.fourier.visualize import bar
     >>> fig, ax = plt.subplots(2, 1, sharey=True, figsize=(15, 4))
     >>> bar(coeffs, n_inputs, ax, colour_dict={"real" : "red", "imag" : "blue"})
 
@@ -498,6 +501,7 @@ def panel(coeffs, n_inputs, ax, colour=None):
     degree 2, we need a 5x5 grid.
 
     >>> import matplotlib.pyplot as plt
+    >>> from pennylane.fourier.visualize import panel
     >>> fig, ax = plt.subplots(5, 5, figsize=(12, 10), sharex=True, sharey=True)
     >>> panel(coeffs, n_inputs, ax)
 
@@ -630,6 +634,7 @@ def radial_box(coeffs, n_inputs, ax, show_freqs=True, colour_dict=None, show_fli
     .. code-block:: python
 
         import matplotlib.pyplot as plt
+        from pennylane.fourier.visualize import radial_box
 
         fig, ax = plt.subplots(
             1, 2, sharex=True, sharey=True,
diff --git a/pennylane/gradients/__init__.py b/pennylane/gradients/__init__.py
index af928852a..1bbd0dbf4 100644
--- a/pennylane/gradients/__init__.py
+++ b/pennylane/gradients/__init__.py
@@ -149,8 +149,9 @@ Alternatively, quantum gradient transforms can be applied manually to QNodes.
 >>> circuit(weights)
 tensor([0.9658079, 0.0341921], requires_grad=True)
 >>> qml.gradients.param_shift(circuit)(weights)
-tensor([[-0.04673668, -0.09442394, -0.14409127],
-        [ 0.04673668,  0.09442394,  0.14409127]], requires_grad=True)
+(tensor([-0.04673668,  0.04673668], requires_grad=True),
+ tensor([-0.09442394,  0.09442394], requires_grad=True),
+ tensor([-0.14409127,  0.14409127], requires_grad=True))
 
 Comparing this to autodifferentiation:
 
@@ -291,7 +292,9 @@ the gradient:
 
 >>> dev = qml.device("default.qubit", wires=2)
 >>> fn(qml.execute(gradient_tapes, dev, None))
-(array(-0.09347337), array(-0.18884787), array(-0.28818254))
+(tensor(-0.09347337, requires_grad=True),
+ tensor(-0.18884787, requires_grad=True),
+ tensor(-0.28818254, requires_grad=True))
 
 Note that the post-processing function ``fn`` returned by the
 gradient transform is applied to the flat list of results returned
diff --git a/pennylane/gradients/finite_difference.py b/pennylane/gradients/finite_difference.py
index d925c5785..643f21333 100644
--- a/pennylane/gradients/finite_difference.py
+++ b/pennylane/gradients/finite_difference.py
@@ -327,8 +327,12 @@ def finite_diff(
 
         >>> dev = qml.device("default.qubit", wires=2)
         >>> fn(qml.execute(gradient_tapes, dev, None))
-        ((array(-0.38751724), array(-0.18884792), array(-0.38355709)),
-         (array(0.69916868), array(0.34072432), array(0.69202366)))
+        ((tensor(-0.56464251, requires_grad=True),
+         tensor(-0.56464251, requires_grad=True),
+         tensor(-0.56464251, requires_grad=True)),
+        (tensor(0.93203912, requires_grad=True),
+         tensor(0.93203912, requires_grad=True),
+         tensor(0.93203912, requires_grad=True)))
 
         This gradient transform is compatible with devices that use shot vectors for execution.
 
diff --git a/pennylane/gradients/hadamard_gradient.py b/pennylane/gradients/hadamard_gradient.py
index edac693b9..05f9081e2 100644
--- a/pennylane/gradients/hadamard_gradient.py
+++ b/pennylane/gradients/hadamard_gradient.py
@@ -125,16 +125,16 @@ def hadamard_grad(
     ...     return qml.expval(qml.PauliZ(0))
     >>> params = np.array([0.1, 0.2, 0.3], requires_grad=True)
     >>> qml.gradients.hadamard_grad(circuit)(params)
-    (tensor([-0.3875172], requires_grad=True),
-     tensor([-0.18884787], requires_grad=True),
-     tensor([-0.38355704], requires_grad=True))
+    (tensor(-0.3875172, requires_grad=True),
+     tensor(-0.18884787, requires_grad=True),
+     tensor(-0.38355704, requires_grad=True))
 
     This quantum gradient transform can also be applied to low-level
     :class:`~.QuantumTape` objects. This will result in no implicit quantum
     device evaluation. Instead, the processed tapes, and post-processing
     function, which together define the gradient are directly returned:
 
-    >>> ops = [qml.RX(p, wires=0) for p in params]
+    >>> ops = [qml.RX(params[0], 0), qml.RY(params[1], 0), qml.RX(params[2], 0)]
     >>> measurements = [qml.expval(qml.PauliZ(0))]
     >>> tape = qml.tape.QuantumTape(ops, measurements)
     >>> gradient_tapes, fn = qml.gradients.hadamard_grad(tape)
@@ -151,7 +151,9 @@ def hadamard_grad(
 
     >>> dev = qml.device("default.qubit", wires=2)
     >>> fn(qml.execute(gradient_tapes, dev, None))
-    (array(-0.3875172), array(-0.18884787), array(-0.38355704))
+    (tensor(-0.3875172, requires_grad=True),
+     tensor(-0.18884787, requires_grad=True),
+     tensor(-0.38355704, requires_grad=True))
 
     This transform can be registered directly as the quantum gradient transform
     to use during autodifferentiation:
@@ -165,7 +167,7 @@ def hadamard_grad(
     ...     return qml.expval(qml.PauliZ(0))
     >>> params = jax.numpy.array([0.1, 0.2, 0.3])
     >>> jax.jacobian(circuit)(params)
-    [-0.3875172  -0.18884787 -0.38355704]
+    Array([-0.3875172 , -0.18884787, -0.38355704], dtype=float64)
 
     If you use custom wires on your device, you need to pass an auxiliary wire.
 
@@ -179,7 +181,7 @@ def hadamard_grad(
     ...    return qml.expval(qml.PauliZ("a"))
     >>> params = jax.numpy.array([0.1, 0.2, 0.3])
     >>> jax.jacobian(circuit)(params)
-    [-0.3875172  -0.18884787 -0.38355704]
+    Array([-0.3875172 , -0.18884787, -0.38355704], dtype=float64)
 
     .. note::
 
diff --git a/pennylane/gradients/jvp.py b/pennylane/gradients/jvp.py
index a6cd26579..3d7c9a4dc 100644
--- a/pennylane/gradients/jvp.py
+++ b/pennylane/gradients/jvp.py
@@ -287,7 +287,7 @@ def jvp(tape, tangent, gradient_fn, gradient_kwargs=None):
     Executing the JVP tapes, and applying the processing function:
 
     >>> dev = qml.device("default.qubit", wires=2)
-    >>> jvp = fn(dev.batch_execute(jvp_tapes))
+    >>> jvp = fn(dev.execute(jvp_tapes))
     >>> jvp
     (Array(-0.62073976, dtype=float32), Array([-0.3259707 ,  0.32597077], dtype=float32))
     """
@@ -397,9 +397,9 @@ def batch_jvp(tapes, tangents, gradient_fn, reduction="append", gradient_kwargs=
     >>> jvp_tapes, fn = qml.gradients.batch_jvp(tapes, tangents, qml.gradients.param_shift)
 
     >>> dev = qml.device("default.qubit", wires=2)
-    >>> jvps = fn(dev.batch_execute(jvp_tapes))
+    >>> jvps = fn(dev.execute(jvp_tapes))
     >>> jvps
-    [(Array(-0.62073976, dtype=float32), Array([-0.3259707 ,  0.32597077], dtype=float32)), Array(-0.6900841, dtype=float32)]
+    ((Array(-0.62073976, dtype=float32), Array([-0.3259707 ,  0.32597077], dtype=float32)), Array(-0.6900841, dtype=float32))
 
     We have two JVPs; one per tape. Each one corresponds to the shape of the output of their respective tape.
     """
diff --git a/pennylane/gradients/parameter_shift.py b/pennylane/gradients/parameter_shift.py
index d7877bb89..de540e4d6 100644
--- a/pennylane/gradients/parameter_shift.py
+++ b/pennylane/gradients/parameter_shift.py
@@ -876,7 +876,7 @@ def param_shift(
     ...     return qml.expval(qml.PauliZ(0))
     >>> params = np.array([0.1, 0.2, 0.3], requires_grad=True)
     >>> qml.jacobian(circuit)(params)
-    tensor([-0.38751725, -0.18884792, -0.38355708], requires_grad=True)
+    array([-0.3875172 , -0.18884787, -0.38355704])
 
     When differentiating QNodes with multiple measurements using Autograd or TensorFlow, the outputs of the QNode first
     need to be stacked. The reason is that those two frameworks only allow differentiating functions with array or
@@ -940,16 +940,14 @@ def param_shift(
         ((tensor(-0.38751724, requires_grad=True),
           tensor(-0.18884792, requires_grad=True),
           tensor(-0.38355709, requires_grad=True)),
-         (tensor(0.69916868, requires_grad=True),
-          tensor(0.34072432, requires_grad=True),
-          tensor(0.69202366, requires_grad=True)))
+         (array(0.69916862), array(0.34072424), array(0.69202359)))
 
         This quantum gradient transform can also be applied to low-level
         :class:`~.QuantumTape` objects. This will result in no implicit quantum
         device evaluation. Instead, the processed tapes, and post-processing
         function, which together define the gradient are directly returned:
 
-        >>> ops = [qml.RX(p, wires=0) for p in params]
+        >>> ops = [qml.RX(params[0], 0), qml.RY(params[1], 0), qml.RX(params[2], 0)]
         >>> measurements = [qml.expval(qml.PauliZ(0)), qml.var(qml.PauliZ(0))]
         >>> tape = qml.tape.QuantumTape(ops, measurements)
         >>> gradient_tapes, fn = qml.gradients.param_shift(tape)
@@ -969,7 +967,9 @@ def param_shift(
 
         >>> dev = qml.device("default.qubit", wires=2)
         >>> fn(qml.execute(gradient_tapes, dev, None))
-        ((array(-0.3875172), array(-0.18884787), array(-0.38355704)),
+        ((tensor(-0.3875172, requires_grad=True),
+          tensor(-0.18884787, requires_grad=True),
+          tensor(-0.38355704, requires_grad=True)),
          (array(0.69916862), array(0.34072424), array(0.69202359)))
 
         This gradient transform is compatible with devices that use shot vectors for execution.
@@ -999,7 +999,7 @@ def param_shift(
 
         >>> params = np.array([0.1, 0.2, 0.3], requires_grad=True)
         >>> ops = [qml.RX(p, wires=0) for p in params]
-        >>> measurements = [qml.expval(qml.PauliZ(0)), qml.var(qml.PauliZ(0))]
+        >>> measurements = [qml.expval(qml.PauliZ(0))]
         >>> tape = qml.tape.QuantumTape(ops, measurements)
         >>> gradient_tapes, fn = qml.gradients.param_shift(tape, broadcast=True)
         >>> len(gradient_tapes)
diff --git a/pennylane/gradients/spsa_gradient.py b/pennylane/gradients/spsa_gradient.py
index f49454fe4..6bddef769 100644
--- a/pennylane/gradients/spsa_gradient.py
+++ b/pennylane/gradients/spsa_gradient.py
@@ -241,7 +241,7 @@ def spsa_grad(
         device evaluation. Instead, the processed tapes, and post-processing
         function, which together define the gradient are directly returned:
 
-        >>> ops = [qml.RX(p, wires=0) for p in params]
+        >>> ops = [qml.RX(params[0], 0), qml.RY(params[1], 0), qml.RX(params[2], 0)]
         >>> measurements = [qml.expval(qml.PauliZ(0)), qml.var(qml.PauliZ(0))]
         >>> tape = qml.tape.QuantumTape(ops, measurements)
         >>> gradient_tapes, fn = qml.gradients.spsa_grad(tape)
diff --git a/pennylane/interfaces/set_shots.py b/pennylane/interfaces/set_shots.py
index 652a7af89..bb463233f 100644
--- a/pennylane/interfaces/set_shots.py
+++ b/pennylane/interfaces/set_shots.py
@@ -31,7 +31,7 @@ def set_shots(device, shots):
 
     As a standard context manager:
 
-    >>> dev = qml.device("default.qubit", wires=2, shots=None)
+    >>> dev = qml.device("default.qubit.legacy", wires=2, shots=None)
     >>> with set_shots(dev, shots=100):
     ...     print(dev.shots)
     100
diff --git a/pennylane/math/quantum.py b/pennylane/math/quantum.py
index a27f660ed..6098fa6df 100644
--- a/pennylane/math/quantum.py
+++ b/pennylane/math/quantum.py
@@ -79,14 +79,13 @@ def cov_matrix(prob, obs, wires=None, diag_approx=False):
     >>> weights = np.random.random(shape, requires_grad=True)
     >>> cov = qml.math.cov_matrix(circuit(weights), obs_list)
     >>> cov
-    array([[0.98707611, 0.03665537],
-         [0.03665537, 0.99998377]])
+    tensor([[0.9275379 , 0.05233832], [0.05233832, 0.99335545]], requires_grad=True)
 
     Autodifferentiation is fully supported using all interfaces.
     Here we use autograd:
 
     >>> cost_fn = lambda weights: qml.math.cov_matrix(circuit(weights), obs_list)[0, 1]
-    >>> qml.grad(cost_fn)(weights)[0]
+    >>> qml.grad(cost_fn)(weights)
     array([[[ 4.94240914e-17, -2.33786398e-01, -1.54193959e-01],
             [-3.05414996e-17,  8.40072236e-04,  5.57884080e-04],
             [ 3.01859411e-17,  8.60411436e-03,  6.15745204e-04]],
diff --git a/pennylane/measurements/__init__.py b/pennylane/measurements/__init__.py
index a24611f36..7e48c2f3a 100644
--- a/pennylane/measurements/__init__.py
+++ b/pennylane/measurements/__init__.py
@@ -118,7 +118,7 @@ ability to collect statistics on single measurement values.
 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]))
+(0.6, array([1, 1, 1, 0, 1]))
 
 Creating custom measurements
 ----------------------------
@@ -138,10 +138,10 @@ obtained of a given state:
             wires = list(range(len(state)))
             super().__init__(wires=wires)
 
-        def process_samples(self, samples, wire_order, shot_range, bin_size):
+        def process_samples(self, samples, wire_order, shot_range=None, bin_size=None):
             counts_mp = qml.counts(wires=self._wires)
             counts = counts_mp.process_samples(samples, wire_order, shot_range, bin_size)
-            return counts.get(self.state, 0)
+            return float(counts.get(self.state, 0))
 
         def __copy__(self):
             return CountState(state=self.state)
@@ -181,7 +181,7 @@ When :math:`\theta = 1.23`, the probability of obtaining the state
 we should obtain the excited state 3333 times approximately.
 
 >>> circuit(1.23)
-tensor(3303., requires_grad=True)
+array(3303.)
 
 Given that the measurement process returns a real scalar value, we can differentiate it
 using the analytic method.
diff --git a/pennylane/measurements/classical_shadow.py b/pennylane/measurements/classical_shadow.py
index 4e88c215a..f8470d35a 100644
--- a/pennylane/measurements/classical_shadow.py
+++ b/pennylane/measurements/classical_shadow.py
@@ -72,7 +72,7 @@ def shadow_expval(H, k=1, seed=None):
     We can compute the expectation value of H as well as its gradient in the usual way.
 
     >>> circuit(x, H)
-    tensor(1.827, requires_grad=True)
+    array(1.8774)
     >>> qml.grad(circuit)(x, H)
     -0.44999999999999984
 
@@ -81,9 +81,9 @@ def shadow_expval(H, k=1, seed=None):
 
     >>> Hs = [H, qml.PauliX(0), qml.PauliY(0), qml.PauliZ(0)]
     >>> circuit(x, Hs)
-    [ 1.88586e+00,  4.50000e-03,  1.32000e-03, -1.92000e-03]
+    array([ 1.881 , -0.0312, -0.0027, -0.0087])
     >>> qml.jacobian(circuit)(x, Hs)
-    [-0.48312, -0.00198, -0.00375,  0.00168]
+    array([-0.4518,  0.0174, -0.0216, -0.0063])
     """
     seed = seed or np.random.randint(2**30)
     return ShadowExpvalMP(H=H, seed=seed, k=k)
@@ -178,7 +178,7 @@ def classical_shadow(wires, seed=None):
 
             ops = [qml.Hadamard(wires=0), qml.CNOT(wires=(0,1))]
             measurements = [qml.classical_shadow(wires=(0,1))]
-            tape = qml.tape.QuantumTape(ops, measurements)
+            tape = qml.tape.QuantumTape(ops, measurements, shots=5)
 
         >>> bits1, recipes1 = qml.execute([tape], device=dev, gradient_fn=None)[0]
         >>> bits2, recipes2 = qml.execute([tape], device=dev, gradient_fn=None)[0]
@@ -195,10 +195,10 @@ def classical_shadow(wires, seed=None):
             dev = qml.device("default.qubit", wires=2, shots=5)
 
             measurements1 = [qml.classical_shadow(wires=(0,1), seed=10)]
-            tape1 = qml.tape.QuantumTape(ops, measurements1)
+            tape1 = qml.tape.QuantumTape(ops, measurements1, shots=5)
 
             measurements2 = [qml.classical_shadow(wires=(0,1), seed=15)]
-            tape2 = qml.tape.QuantumTape(ops, measurements2)
+            tape2 = qml.tape.QuantumTape(ops, measurements2, shots=5)
 
         >>> bits1, recipes1 = qml.execute([tape1], device=dev, gradient_fn=None)[0]
         >>> bits2, recipes2 = qml.execute([tape2], device=dev, gradient_fn=None)[0]
diff --git a/pennylane/measurements/measurements.py b/pennylane/measurements/measurements.py
index ab1ed6062..3b01a5217 100644
--- a/pennylane/measurements/measurements.py
+++ b/pennylane/measurements/measurements.py
@@ -512,7 +512,7 @@ class SampleMeasurement(MeasurementProcess):
     ...     qml.PauliX(0)
     ...     return MyMeasurement(wires=[0]), MyMeasurement(wires=[1])
     >>> circuit()
-    tensor([1000,    0], requires_grad=True)
+    (tensor(1000, requires_grad=True), tensor(0, requires_grad=True))
     """
 
     @abstractmethod
diff --git a/pennylane/measurements/mutual_info.py b/pennylane/measurements/mutual_info.py
index c5035bce7..8704e5f39 100644
--- a/pennylane/measurements/mutual_info.py
+++ b/pennylane/measurements/mutual_info.py
@@ -65,7 +65,7 @@ def mutual_info(wires0, wires1, log_base=None):
 
     >>> param = np.array(np.pi/4, requires_grad=True)
     >>> qml.grad(circuit_mutual)(param)
-    1.2464504802804612
+    tensor(1.24645048, requires_grad=True)
 
     .. note::
 
diff --git a/pennylane/measurements/state.py b/pennylane/measurements/state.py
index 2cc0c576b..78dda1cc4 100644
--- a/pennylane/measurements/state.py
+++ b/pennylane/measurements/state.py
@@ -67,16 +67,16 @@ def state() -> "StateMP":
         is then differentiated:
 
         >>> dev = qml.device('default.qubit', wires=2)
-        >>> qml.qnode(dev, diff_method="backprop")
+        >>> @qml.qnode(dev, diff_method="backprop")
         ... def test(x):
         ...     qml.RY(x, wires=[0])
         ...     return qml.state()
         >>> def cost(x):
         ...     return np.abs(test(x)[0])
         >>> cost(x)
-        tensor(0.98877108, requires_grad=True)
+        0.9987502603949663
         >>> qml.grad(cost)(x)
-        -0.07471906623679961
+        tensor(-0.02498958, requires_grad=True)
     """
     return StateMP()
 
diff --git a/pennylane/measurements/vn_entropy.py b/pennylane/measurements/vn_entropy.py
index 34da3a662..1a21af26b 100644
--- a/pennylane/measurements/vn_entropy.py
+++ b/pennylane/measurements/vn_entropy.py
@@ -56,7 +56,7 @@ def vn_entropy(wires, log_base=None) -> "VnEntropyMP":
 
     >>> param = np.array(np.pi/4, requires_grad=True)
     >>> qml.grad(circuit_entropy)(param)
-    0.6232252401402305
+    tensor(0.62322524, requires_grad=True)
 
     .. note::
 
diff --git a/pennylane/operation.py b/pennylane/operation.py
index 6fc285eaa..619fabbb6 100644
--- a/pennylane/operation.py
+++ b/pennylane/operation.py
@@ -517,7 +517,7 @@ class Operator(abc.ABC):
 
     >>> a = np.array(3.14)
     >>> circuit(a)
-    -0.9999987318946099
+    tensor(-0.99999873, requires_grad=True)
 
     .. details::
         :title: Serialization and Pytree format
diff --git a/pennylane/ops/op_math/controlled.py b/pennylane/ops/op_math/controlled.py
index 165d72466..de3212140 100644
--- a/pennylane/ops/op_math/controlled.py
+++ b/pennylane/ops/op_math/controlled.py
@@ -69,7 +69,7 @@ def ctrl(op, control, control_values=None, work_wires=None):
     >>> circuit(x)
     tensor(0.36235775, requires_grad=True)
     >>> qml.grad(circuit)(x)
-    -0.9320390859672264
+    tensor(-0.93203909, requires_grad=True)
 
     :func:`~.ctrl` works on both callables like ``qml.RX`` or a quantum function
     and individual :class:`~.operation.Operator`'s.
diff --git a/pennylane/ops/op_math/sum.py b/pennylane/ops/op_math/sum.py
index 3bbdf8536..b0c52f7ba 100644
--- a/pennylane/ops/op_math/sum.py
+++ b/pennylane/ops/op_math/sum.py
@@ -154,7 +154,7 @@ class Sum(CompositeOp):
 
         >>> weights = qnp.array([0.1, 0.2, 0.3], requires_grad=True)
         >>> qml.grad(circuit)(weights)
-        tensor([-0.09347337, -0.18884787, -0.28818254], requires_grad=True)
+        array([-0.09347337, -0.18884787, -0.28818254])
     """
 
     _op_symbol = "+"
diff --git a/pennylane/ops/qubit/arithmetic_ops.py b/pennylane/ops/qubit/arithmetic_ops.py
index aaee6cae9..8d03016cf 100644
--- a/pennylane/ops/qubit/arithmetic_ops.py
+++ b/pennylane/ops/qubit/arithmetic_ops.py
@@ -347,9 +347,9 @@ class IntegerComparator(Operation):
     ...     qml.IntegerComparator(value, geq=geq, wires=range(3))
     ...     return qml.state()
     >>> circuit([1, 0, 1], 1, True).reshape(2, 2, 2)[1, 0, 0]
-    (1+0j)
+    tensor(1.+0.j, requires_grad=True)
     >>> circuit([0, 1, 0], 3, False).reshape(2, 2, 2)[0, 1, 1]
-    (1+0j)
+    tensor(1.+0.j, requires_grad=True)
     """
     is_self_inverse = True
     num_wires = AnyWires
diff --git a/pennylane/ops/qubit/observables.py b/pennylane/ops/qubit/observables.py
index 2dbc515d6..7d7958240 100644
--- a/pennylane/ops/qubit/observables.py
+++ b/pennylane/ops/qubit/observables.py
@@ -574,14 +574,14 @@ class StateVectorProjector(Projector):
         **Example:**
 
         >>> state_vector = np.array([0, 1, 1, 0])/np.sqrt(2)
-        >>> StateVectorProjector(state_vector, wires=(0, 1)).label()
+        >>> qml.Projector(state_vector, wires=(0, 1)).label()
         'P'
-        >>> StateVectorProjector(state_vector, wires=(0, 1)).label(base_label="hi!")
+        >>> qml.Projector(state_vector, wires=(0, 1)).label(base_label="hi!")
         'hi!'
         >>> dev = qml.device("default.qubit", wires=1)
         >>> @qml.qnode(dev)
         >>> def circuit(state):
-        ...     return qml.expval(StateVectorProjector(state, [0]))
+        ...     return qml.expval(qml.Projector(state, [0]))
         >>> print(qml.draw(circuit)([1, 0]))
         0: ───┤  <|0⟩⟨0|>
         >>> print(qml.draw(circuit)(np.array([1, 1]) / np.sqrt(2)))
diff --git a/pennylane/ops/snapshot.py b/pennylane/ops/snapshot.py
index 10d050c6a..134eaa8f2 100644
--- a/pennylane/ops/snapshot.py
+++ b/pennylane/ops/snapshot.py
@@ -50,10 +50,10 @@ class Snapshot(Operation):
             return qml.expval(qml.PauliX(0))
 
     >>> qml.snapshots(circuit)()
-    {0: array([1.+0.j, 0.+0.j, 0.+0.j, 0.+0.j]),
-    'very_important_state': array([0.70710678+0.j, 0.+0.j, 0.70710678+0.j, 0.+0.j]),
-    2: array([0.70710678+0.j, 0.+0.j, 0.+0.j, 0.70710678+0.j]),
-    'execution_results': array(0.)}
+    {0: array([1., 0., 0., 0.]),
+     'very_important_state': array([0.70710678, 0.        , 0.70710678, 0.        ]),
+     2: array([0.70710678, 0.        , 0.        , 0.70710678]),
+     'execution_results': 0.0}
 
     .. seealso:: :func:`~.snapshots`
     """
diff --git a/pennylane/optimize/adaptive.py b/pennylane/optimize/adaptive.py
index cbc986d82..c4edced9a 100644
--- a/pennylane/optimize/adaptive.py
+++ b/pennylane/optimize/adaptive.py
@@ -125,32 +125,32 @@ class AdaptiveOptimizer:
 
     .. code-block :: pycon
 
-        Energy: -1.246549938420637
-        0: ─╭BasisState(M0)─╭G²(0.20)─┤ ╭<𝓗>
-        1: ─├BasisState(M0)─├G²(0.20)─┤ ├<𝓗>
-        2: ─├BasisState(M0)─│─────────┤ ├<𝓗>
-        3: ─├BasisState(M0)─│─────────┤ ├<𝓗>
-        4: ─├BasisState(M0)─├G²(0.20)─┤ ├<𝓗>
-        5: ─╰BasisState(M0)─╰G²(0.20)─┤ ╰<𝓗>
-        Largest Gradient: 0.14399872776755085
-
-        Energy: -1.2613740231529604
-        0: ─╭BasisState(M0)─╭G²(0.20)─╭G²(0.19)─┤ ╭<𝓗>
-        1: ─├BasisState(M0)─├G²(0.20)─├G²(0.19)─┤ ├<𝓗>
-        2: ─├BasisState(M0)─│─────────├G²(0.19)─┤ ├<𝓗>
-        3: ─├BasisState(M0)─│─────────╰G²(0.19)─┤ ├<𝓗>
-        4: ─├BasisState(M0)─├G²(0.20)───────────┤ ├<𝓗>
-        5: ─╰BasisState(M0)─╰G²(0.20)───────────┤ ╰<𝓗>
-        Largest Gradient: 0.1349349562423238
-
-        Energy: -1.2743971719780331
-        0: ─╭BasisState(M0)─╭G²(0.20)─╭G²(0.19)──────────┤ ╭<𝓗>
-        1: ─├BasisState(M0)─├G²(0.20)─├G²(0.19)─╭G(0.00)─┤ ├<𝓗>
-        2: ─├BasisState(M0)─│─────────├G²(0.19)─│────────┤ ├<𝓗>
-        3: ─├BasisState(M0)─│─────────╰G²(0.19)─╰G(0.00)─┤ ├<𝓗>
-        4: ─├BasisState(M0)─├G²(0.20)────────────────────┤ ├<𝓗>
-        5: ─╰BasisState(M0)─╰G²(0.20)────────────────────┤ ╰<𝓗>
-        Largest Gradient: 0.00040841755397108586
+        Energy: -1.2465499384199699
+        0: ─╭|Ψ⟩─╭G²(0.20)─┤ ╭<𝓗>
+        1: ─├|Ψ⟩─├G²(0.20)─┤ ├<𝓗>
+        2: ─├|Ψ⟩─│─────────┤ ├<𝓗>
+        3: ─├|Ψ⟩─│─────────┤ ├<𝓗>
+        4: ─├|Ψ⟩─├G²(0.20)─┤ ├<𝓗>
+        5: ─╰|Ψ⟩─╰G²(0.20)─┤ ╰<𝓗>
+        Largest Gradient: 0.1439987277673651
+
+        Energy: -1.2613740231522532
+        0: ─╭|Ψ⟩─╭G²(0.20)─╭G²(0.19)─┤ ╭<𝓗>
+        1: ─├|Ψ⟩─├G²(0.20)─├G²(0.19)─┤ ├<𝓗>
+        2: ─├|Ψ⟩─│─────────├G²(0.19)─┤ ├<𝓗>
+        3: ─├|Ψ⟩─│─────────╰G²(0.19)─┤ ├<𝓗>
+        4: ─├|Ψ⟩─├G²(0.20)───────────┤ ├<𝓗>
+        5: ─╰|Ψ⟩─╰G²(0.20)───────────┤ ╰<𝓗>
+        Largest Gradient: 0.13493495624211427
+
+        Energy: -1.2743971719772815
+        0: ─╭|Ψ⟩─╭G²(0.20)─╭G²(0.19)──────────┤ ╭<𝓗>
+        1: ─├|Ψ⟩─├G²(0.20)─├G²(0.19)─╭G(0.00)─┤ ├<𝓗>
+        2: ─├|Ψ⟩─│─────────├G²(0.19)─│────────┤ ├<𝓗>
+        3: ─├|Ψ⟩─│─────────╰G²(0.19)─╰G(0.00)─┤ ├<𝓗>
+        4: ─├|Ψ⟩─├G²(0.20)────────────────────┤ ├<𝓗>
+        5: ─╰|Ψ⟩─╰G²(0.20)────────────────────┤ ╰<𝓗>
+        Largest Gradient: 0.0004084175253678331
     """
 
     def __init__(self, param_steps=10, stepsize=0.5):
diff --git a/pennylane/pulse/transmon.py b/pennylane/pulse/transmon.py
index 537283b31..0db919c7b 100644
--- a/pennylane/pulse/transmon.py
+++ b/pennylane/pulse/transmon.py
@@ -79,10 +79,11 @@ def transmon_interaction(
         :func:`~.transmon_drive`
 
     Args:
-        qubit_freq (Union[float, list[float]]): List of dressed qubit frequencies. This should be in units
+        qubit_freq (Union[float, list[float], Callable]): List of dressed qubit frequencies. This should be in units
             of frequency (GHz), and will be converted to angular frequency :math:`\omega` internally where
             needed, i.e. multiplied by :math:`2 \pi`. When passing a single float all qubits are assumed to
-            have that same frequency.
+            have that same frequency. When passing a parametrized function, it must have two
+            arguments, the first one being the trainable parameters and the second one being time.
         connections (list[tuple(int)]): List of connections ``(i, j)`` between qubits i and j.
             When the wires in ``connections`` are not contained in ``wires``, a warning is raised.
         coupling (Union[float, list[float]]): List of coupling strengths. This should be in units
@@ -158,9 +159,9 @@ def transmon_interaction(
         anharmonicity = [0.0] * n_wires
 
     # TODO: make coefficients callable / trainable. Currently not supported
-    if qml.math.ndim(qubit_freq) == 0:
+    if callable(qubit_freq) or qml.math.ndim(qubit_freq) == 0:
         qubit_freq = [qubit_freq] * n_wires
-    if len(qubit_freq) != n_wires:
+    elif len(qubit_freq) != n_wires:
         raise ValueError(
             f"Number of qubit frequencies in {qubit_freq} does not match the provided wires = {wires}"
         )
diff --git a/pennylane/qaoa/layers.py b/pennylane/qaoa/layers.py
index e2625e617..0dba911e2 100644
--- a/pennylane/qaoa/layers.py
+++ b/pennylane/qaoa/layers.py
@@ -80,7 +80,7 @@ def cost_layer(gamma, hamiltonian):
                 for i in range(2):
                     qml.Hadamard(wires=i)
 
-                cost_layer(gamma, cost_h)
+                qaoa.cost_layer(gamma, cost_h)
 
                 return [qml.expval(qml.PauliZ(wires=i)) for i in range(2)]
 
@@ -90,8 +90,8 @@ def cost_layer(gamma, hamiltonian):
         0: ──H─╭ApproxTimeEvolution(1.00,1.00,0.50)─┤  <Z>
         1: ──H─╰ApproxTimeEvolution(1.00,1.00,0.50)─┤  <Z>
         >>> print(qml.draw(circuit, expansion_strategy="device")(0.5))
-        0: ──H──MultiRZ(1.00)─╭MultiRZ(1.00)─┤  <Z>
-        1: ──H────────────────╰MultiRZ(1.00)─┤  <Z>
+        0: ──H──RZ(1.00)─╭RZZ(1.00)─┤  <Z>
+        1: ──H───────────╰RZZ(1.00)─┤  <Z>
 
     """
     if not isinstance(hamiltonian, qml.Hamiltonian):
@@ -152,9 +152,8 @@ def mixer_layer(alpha, hamiltonian):
         0: ──H─╭ApproxTimeEvolution(1.00,1.00,0.50)─┤  <Z>
         1: ──H─╰ApproxTimeEvolution(1.00,1.00,0.50)─┤  <Z>
         >>> print(qml.draw(circuit, expansion_strategy="device")(0.5))
-        0: ──H──H──MultiRZ(1.00)──H──H─╭MultiRZ(1.00)──H─┤  <Z>
-        1: ──H──H──────────────────────╰MultiRZ(1.00)──H─┤  <Z>
-
+        0: ──H──RX(1.00)─╭RXX(1.00)─┤  <Z>
+        1: ──H───────────╰RXX(1.00)─┤  <Z>
 
     """
     if not isinstance(hamiltonian, qml.Hamiltonian):
diff --git a/pennylane/qchem/tapering.py b/pennylane/qchem/tapering.py
index 218629389..23bb2eafb 100644
--- a/pennylane/qchem/tapering.py
+++ b/pennylane/qchem/tapering.py
@@ -140,7 +140,7 @@ def symmetry_generators(h):
     **Example**
 
     >>> symbols = ["H", "H"]
-    >>> geometry = np.array([[0.0, 0.0, 0.0], [0.0, 0.0, 1.0]])
+    >>> coordinates = np.array([[0.0, 0.0, 0.0], [0.0, 0.0, 1.0]])
     >>> H, qubits = qml.qchem.molecular_hamiltonian(symbols, coordinates)
     >>> t = symmetry_generators(H)
     >>> t
@@ -468,7 +468,7 @@ def taper_hf(generators, paulixops, paulix_sector, num_electrons, num_wires):
     >>> symbols = ['He', 'H']
     >>> geometry = np.array([[0.0, 0.0, 0.0], [0.0, 0.0, 1.4588684632]])
     >>> mol = qml.qchem.Molecule(symbols, geometry, charge=1)
-    >>> H, n_qubits = qml.qchem.molecular_hamiltonian(symbols, geometry)
+    >>> H, n_qubits = qml.qchem.molecular_hamiltonian(symbols, geometry, charge=1)
     >>> n_elec = mol.n_electrons
     >>> generators = qml.qchem.symmetry_generators(H)
     >>> paulixops = qml.qchem.paulix_ops(generators, 4)
@@ -650,14 +650,14 @@ def taper_operation(
 
     >>> symbols, geometry = ['He', 'H'], np.array([[0.0, 0.0, 0.0], [0.0, 0.0, 1.4589]])
     >>> mol = qchem.Molecule(symbols, geometry, charge=1)
-    >>> H, n_qubits = qchem.molecular_hamiltonian(symbols, geometry)
+    >>> H, n_qubits = qchem.molecular_hamiltonian(symbols, geometry, charge=1)
     >>> generators = qchem.symmetry_generators(H)
     >>> paulixops = qchem.paulix_ops(generators, n_qubits)
     >>> paulix_sector = qchem.optimal_sector(H, generators, mol.n_electrons)
     >>> tap_op = qchem.taper_operation(qml.SingleExcitation, generators, paulixops,
     ...                                paulix_sector, wire_order=H.wires, op_wires=[0, 2])
     >>> tap_op(3.14159)
-    [Exp(1.570795j PauliY)]
+    [Exp(1.5707949999999993j PauliY), Exp(0j Identity)]
 
     The obtained tapered operation function can then be used within a :class:`~.pennylane.QNode`:
 
diff --git a/pennylane/qinfo/transforms.py b/pennylane/qinfo/transforms.py
index 3e6f83c0b..ebf7c762d 100644
--- a/pennylane/qinfo/transforms.py
+++ b/pennylane/qinfo/transforms.py
@@ -510,15 +510,15 @@ def classical_fisher(qnode, argnums=0):
 
     >>> params = pnp.random.random(2)
     >>> circ(params)
-    tensor([0.77708372, 0.        , 0.        , 0.22291628], requires_grad=True)
+    [0.61281668 0.         0.         0.38718332]
 
     We can obtain its ``(2, 2)`` classical fisher information matrix (CFIM) by simply calling the function returned
     by ``classical_fisher()``:
 
     >>> cfim_func = qml.qinfo.classical_fisher(circ)
     >>> cfim_func(params)
-    tensor([[1., 1.],
-        [1., 1.]], requires_grad=True)
+    [[1. 1.]
+     [1. 1.]]
 
     This function has the same signature as the :class:`.QNode`. Here is a small example with multiple arguments:
 
@@ -532,10 +532,9 @@ def classical_fisher(qnode, argnums=0):
 
     >>> x, y = pnp.array([0.5, 0.6], requires_grad=True)
     >>> circ(x, y)
-    (tensor([0.87380224, 0.        , 0.12619776, 0.        ], requires_grad=True)
+    [0.86215007 0.         0.13784993 0.        ]
     >>> qml.qinfo.classical_fisher(circ)(x, y)
-     [tensor([[0.15828019]], requires_grad=True),
-      tensor([[0.74825326]], requires_grad=True)])
+    [array([[0.32934729]]), array([[0.51650396]])]
 
     Note how in the case of multiple variables we get a list of matrices with sizes
     ``[(n_params0, n_params0), (n_params1, n_params1)]``, which in this case is simply two ``(1, 1)`` matrices.
@@ -591,8 +590,8 @@ def classical_fisher(qnode, argnums=0):
         params = pnp.random.random(2)
 
     >>> qml.qinfo.classical_fisher(circ)(params)
-    tensor([[0.28096197, 0.36228429],
-            [0.36228429, 0.46714473]], requires_grad=True)
+    [[4.18575068e-06 2.34443943e-03]
+     [2.34443943e-03 1.31312079e+00]]
     >>> qml.jacobian(qml.qinfo.classical_fisher(circ))(params)
     array([[[9.98030491e-01, 3.46944695e-18],
             [1.36541817e-01, 5.15248592e-01]],
@@ -706,12 +705,12 @@ def quantum_fisher(
 
     >>> grad = qml.grad(circ)(params)
     >>> grad
-    array([ 0.59422561, -0.02615095, -0.05146226])
+    [ 0.59422561 -0.02615095 -0.05146226]
     >>> qfim = qml.qinfo.quantum_fisher(circ)(params)
     >>> qfim
-    tensor([[1.        , 0.        , 0.        ],
-            [0.        , 1.        , 0.        ],
-            [0.        , 0.        , 0.77517241]], requires_grad=True)
+    [[1.         0.         0.        ]
+     [0.         1.         0.        ]
+     [0.         0.         0.77517241]]
     >>> qfim @ grad
     tensor([ 0.59422561, -0.02615095, -0.03989212], requires_grad=True)
 
@@ -839,7 +838,7 @@ def fidelity(qnode0, qnode1, wires0, wires1):
             qml.RY(x, wires=0)
             return qml.state()
 
-    >>> qml.qinfo.fidelity(circuit_rx, circuit_ry, wires0=[0], wires1=[0])((), (0.2))
+    >>> qml.qinfo.fidelity(circuit_rx, circuit_ry, wires0=[0], wires1=[0])(None, (0.2))
     0.9900332889206207
 
     On the other hand, if the second QNode is the one that does not depend on parameters then a single tuple can also be
@@ -866,7 +865,7 @@ def fidelity(qnode0, qnode1, wires0, wires1):
             return qml.state()
 
     >>> jax.grad(qml.qinfo.fidelity(circuit0, circuit1, wires0=[0], wires1=[0]))((jax.numpy.array(0.3)))
-    -0.14776011
+    Array(-0.14776011, dtype=float64, weak_type=True)
 
     There is also the possibility to pass a single dictionary at the end of the tuple for fixing args,
     you can follow this example:
@@ -997,7 +996,7 @@ def relative_entropy(qnode0, qnode1, wires0, wires1):
 
     >>> x, y = np.array(0.4), np.array(0.6)
     >>> relative_entropy_circuit((x,), (y,))
-    0.017750012490703237
+    tensor(0.01775001, requires_grad=True)
 
     This transform is fully differentiable:
 
@@ -1007,7 +1006,7 @@ def relative_entropy(qnode0, qnode1, wires0, wires1):
             return relative_entropy_circuit((x,), (y,))
 
     >>> wrapper(x, y)
-    0.017750012490703237
+    tensor(0.01775001, requires_grad=True)
     >>> qml.grad(wrapper)(x, y)
     (tensor(-0.16458856, requires_grad=True),
      tensor(0.16953273, requires_grad=True))
diff --git a/pennylane/qnn/cost.py b/pennylane/qnn/cost.py
index ba0b1839d..243385631 100644
--- a/pennylane/qnn/cost.py
+++ b/pennylane/qnn/cost.py
@@ -81,7 +81,7 @@ class SquaredErrorLoss:
     >>> loss = qml.qnn.cost.SquaredErrorLoss(ansatz, obs, dev, interface="torch")
     >>> phis = np.ones(num_qubits)
     >>> loss(phis, target=np.array([1.0, 0.5, 0.1]))
-    tensor([0.2113, 0.2500, 0.0368], dtype=torch.float64)
+    array([0.21132197, 0.25      , 0.03683581])
 
     The loss function can be minimized using any gradient descent-based
     :doc:`optimizer </introduction/interfaces>`.
diff --git a/pennylane/shadows/__init__.py b/pennylane/shadows/__init__.py
index 5b6da2798..e7ef9a2ee 100644
--- a/pennylane/shadows/__init__.py
+++ b/pennylane/shadows/__init__.py
@@ -89,7 +89,7 @@ The easiest way of computing expectation values with classical shadows in PennyL
 The big advantage of this way of computing expectation values is that it is differentiable.
 
 >>> qnode(x)
-tensor(1.827, requires_grad=True)
+array(1.9242)
 >>> qml.grad(qnode)(x)
 -0.44999999999999984
 
diff --git a/pennylane/shadows/classical_shadow.py b/pennylane/shadows/classical_shadow.py
index d8ea39281..62e9d370a 100644
--- a/pennylane/shadows/classical_shadow.py
+++ b/pennylane/shadows/classical_shadow.py
@@ -204,8 +204,8 @@ class ClassicalShadow:
 
             bell_state = np.array([[0.5, 0, 0, 0.5], [0, 0, 0, 0], [0, 0, 0, 0], [0.5, 0, 0, 0.5]])
 
-            >>> np.allclose(bell_state, shadow_state, atol=1e-1)
-            True
+        >>> np.allclose(bell_state, shadow_state, atol=1e-1)
+        True
 
         """
         local_snapshot = self.local_snapshots(wires, snapshots)
diff --git a/pennylane/shadows/transforms.py b/pennylane/shadows/transforms.py
index 92ea0d5b4..f488a4795 100644
--- a/pennylane/shadows/transforms.py
+++ b/pennylane/shadows/transforms.py
@@ -85,7 +85,7 @@ def shadow_expval(tape: QuantumTape, H, k=1) -> (Sequence[QuantumTape], Callable
 
     >>> x = np.array(1.2)
     >>> circuit(x)
-    tensor(0.3069, requires_grad=True)
+    [array(0.3528)]
     >>> qml.grad(circuit)(x)
     -0.9323999999999998
     """
@@ -200,15 +200,10 @@ def shadow_state(tape: QuantumTape, wires, diffable=False) -> (Sequence[QuantumT
 
     >>> x = np.array(1.2)
     >>> circuit(x)
-    tensor([[ 0.33714998+0.j        ,  0.007875  +0.228825j  ,
-             -0.010575  +0.22642499j,  0.33705002+0.01125j   ],
-            [ 0.007875  -0.228825j  ,  0.16104999+0.j        ,
-              0.17055   -0.0126j    ,  0.011025  -0.232575j  ],
-            [-0.010575  -0.22642499j,  0.17055   +0.0126j    ,
-              0.16704999+0.j        , -0.006075  -0.225225j  ],
-            [ 0.33705002-0.01125j   ,  0.011025  +0.232575j  ,
-             -0.006075  +0.225225j  ,  0.33475   +0.j        ]],
-           dtype=complex64, requires_grad=True)
+    array([[ 0.33835   +0.j     , -0.01215   +0.2241j , -0.00465   +0.237j  ,  0.35504997-0.01755j],
+       [-0.01215   -0.2241j ,  0.1528    +0.j     ,  0.16919999-0.0036j , -0.00285   -0.22065j],
+       [-0.00465   -0.237j  ,  0.16919999+0.0036j ,  0.17529999+0.j     ,  0.0099    -0.2358j ],
+       [ 0.35504997+0.01755j, -0.00285   +0.22065j,  0.0099    +0.2358j ,  0.33355   +0.j     ]], dtype=complex64)
     >>> qml.jacobian(lambda x: np.real(circuit(x)))(x)
     array([[-0.245025, -0.005325,  0.004275, -0.2358  ],
            [-0.005325,  0.235275,  0.2358  , -0.004275],
diff --git a/pennylane/templates/embeddings/amplitude.py b/pennylane/templates/embeddings/amplitude.py
index 96009db8d..67d50f212 100644
--- a/pennylane/templates/embeddings/amplitude.py
+++ b/pennylane/templates/embeddings/amplitude.py
@@ -67,14 +67,14 @@ class AmplitudeEmbedding(StatePrep):
             @qml.qnode(dev)
             def circuit(f=None):
                 qml.AmplitudeEmbedding(features=f, wires=range(2))
-                return qml.expval(qml.PauliZ(0))
+                return qml.expval(qml.PauliZ(0)), qml.state()
 
-            circuit(f=[1/2, 1/2, 1/2, 1/2])
+            res, state = circuit(f=[1/2, 1/2, 1/2, 1/2])
 
         The final state of the device is - up to a global phase - equivalent to the input passed to the circuit:
 
-        >>> dev.state
-        [0.5+0.j 0.5+0.j 0.5+0.j 0.5+0.j]
+        >>> state
+        tensor([0.5+0.j, 0.5+0.j, 0.5+0.j, 0.5+0.j], requires_grad=True)
 
         **Differentiating with respect to the features**
 
@@ -91,12 +91,12 @@ class AmplitudeEmbedding(StatePrep):
             @qml.qnode(dev)
             def circuit(f=None):
                 qml.AmplitudeEmbedding(features=f, wires=range(2), normalize=True)
-                return qml.expval(qml.PauliZ(0))
+                return qml.expval(qml.PauliZ(0)), qml.state()
 
-            circuit(f=[15, 15, 15, 15])
+            res, state = circuit(f=[15, 15, 15, 15])
 
-        >>> dev.state
-        [0.5 + 0.j, 0.5 + 0.j, 0.5 + 0.j, 0.5 + 0.j]
+        >>> state
+        tensor([0.5+0.j, 0.5+0.j, 0.5+0.j, 0.5+0.j], requires_grad=True)
 
         **Padding**
 
@@ -110,12 +110,12 @@ class AmplitudeEmbedding(StatePrep):
             @qml.qnode(dev)
             def circuit(f=None):
                 qml.AmplitudeEmbedding(features=f, wires=range(2), pad_with=0.)
-                return qml.expval(qml.PauliZ(0))
+                return qml.expval(qml.PauliZ(0)), qml.state()
 
-            circuit(f=[1/sqrt(2), 1/sqrt(2)])
+            res, state = circuit(f=[1/sqrt(2), 1/sqrt(2)])
 
-        >>> dev.state
-        [0.70710678 + 0.j, 0.70710678 + 0.j, 0.0 + 0.j, 0.0 + 0.j]
+        >>> state
+        tensor([0.70710678+0.j, 0.70710678+0.j, 0.        +0.j, 0.        +0.j], requires_grad=True)
 
     """
 
diff --git a/pennylane/templates/subroutines/hilbert_schmidt.py b/pennylane/templates/subroutines/hilbert_schmidt.py
index f05920b5e..54e4b456e 100644
--- a/pennylane/templates/subroutines/hilbert_schmidt.py
+++ b/pennylane/templates/subroutines/hilbert_schmidt.py
@@ -248,7 +248,7 @@ class LocalHilbertSchmidt(HilbertSchmidt):
 
         Now that the cost function has been defined it can be called for specific parameters:
 
-        >>> cost_lhst([3*np.pi/2, 3*np.pi/2, np.pi/2], v_function = v_function, v_wires = [1], u_tape = u_tape)
+        >>> cost_lhst([3*np.pi/2, 3*np.pi/2, np.pi/2], v_function = v_function, v_wires = [2,3], u_tape = u_tape)
         0.5
     """
 
@@ -268,9 +268,10 @@ class LocalHilbertSchmidt(HilbertSchmidt):
         )
 
         # Unitary U
-        for op_u in u_tape.operations:
-            qml.apply(op_u)
-            decomp_ops.append(op_u)
+        if qml.QueuingManager.recording():
+            decomp_ops.extend(qml.apply(op_u) for op_u in u_tape.operations)
+        else:
+            decomp_ops.extend(u_tape.operations)
 
         # Unitary V conjugate
         decomp_ops.extend(qml.adjoint(qml.apply, lazy=False)(op_v) for op_v in v_tape.operations)
diff --git a/pennylane/templates/subroutines/trotter.py b/pennylane/templates/subroutines/trotter.py
index f843994f5..2a4fa2bf1 100644
--- a/pennylane/templates/subroutines/trotter.py
+++ b/pennylane/templates/subroutines/trotter.py
@@ -121,7 +121,7 @@ class TrotterProduct(Operation):
             return qml.state()
 
     >>> my_circ()
-    [-0.13259524+0.59790098j  0.        +0.j         -0.13259524-0.77932754j  0.        +0.j        ]
+    array([-0.13259524+0.59790098j,  0.        +0.j        , -0.13259524-0.77932754j,  0.        +0.j        ])
 
     .. details::
         :title: Usage Details
diff --git a/pennylane/tracker.py b/pennylane/tracker.py
index a1e7d7521..6383dee6a 100644
--- a/pennylane/tracker.py
+++ b/pennylane/tracker.py
@@ -72,13 +72,21 @@ class Tracker:
     You can then access the tabulated information through ``totals``, ``history``, and ``latest``:
 
     >>> tracker.totals
-    {'executions': 3, 'shots': 300, 'batches': 2, 'batch_len': 3}
+    {'batches': 2, 'simulations': 3, 'executions': 3, 'shots': 300}
     >>> tracker.latest
-    {'batches': 1, 'batch_len': 2}
+    {'simulations': 1,
+     'executions': 1,
+     'results': array(-0.08),
+     'shots': 100,
+     'resources': Resources(num_wires=1, num_gates=1,
+                            gate_types=defaultdict(<class 'int'>, {'RX': 1}),
+                            gate_sizes=defaultdict(<class 'int'>, {1: 1}),
+                            depth=1,
+                            shots=Shots(total_shots=100, shot_vector=(ShotCopies(100 shots x 1),)))}
     >>> tracker.history.keys()
-    dict_keys(['executions', 'shots', 'results', 'resources', 'batches', 'batch_len'])
+    dict_keys(['batches', 'simulations', 'executions', 'results', 'shots', 'resources'])
     >>> tracker.history['results']
-    [array([1.]), array([-0.06]), array([0.18])]
+    [array(1.), array(0.02), array(-0.08)]
     >>> print(tracker.history['resources'][0])
     wires: 1
     gates: 1
diff --git a/pennylane/transforms/batch_transform.py b/pennylane/transforms/batch_transform.py
index 45be45f44..cb25919c0 100644
--- a/pennylane/transforms/batch_transform.py
+++ b/pennylane/transforms/batch_transform.py
@@ -120,7 +120,7 @@ class batch_transform:
     >>> dev = qml.device("default.qubit", wires=1)
     >>> res = qml.execute(tapes, dev, interface="autograd", gradient_fn=qml.gradients.param_shift)
     >>> print(res)
-    [tensor([0.94765073], requires_grad=True), tensor([0.31532236], requires_grad=True)]
+    [0.9476507264148154, 0.31532236239526856]
 
     Applying the processing function, we retrieve the end result of the transform:
 
@@ -482,7 +482,7 @@ def map_batch_transform(
     >>> tapes, fn = map_batch_transform(qml.transforms.hamiltonian_expand, [tape1, tape2])
     >>> dev = qml.device("default.qubit", wires=2)
     >>> fn(qml.execute(tapes, dev, qml.gradients.param_shift))
-    [0.9950041652780257, 0.8150893013179248]
+    [array(0.99500417), array(0.8150893)]
     """
     execution_tapes = []
     batch_fns = []
diff --git a/pennylane/transforms/broadcast_expand.py b/pennylane/transforms/broadcast_expand.py
index c44254c7a..2309bb705 100644
--- a/pennylane/transforms/broadcast_expand.py
+++ b/pennylane/transforms/broadcast_expand.py
@@ -116,7 +116,7 @@ def broadcast_expand(tape: qml.tape.QuantumTape) -> (Sequence[qml.tape.QuantumTa
     >>> tapes
     [<QuantumTape: wires=[0], params=1>, <QuantumTape: wires=[0], params=1>, <QuantumTape: wires=[0], params=1>]
     >>> fn(qml.execute(tapes, qml.device("default.qubit", wires=1), None))
-    array([0.98006658, 0.82533561, 0.54030231])
+    tensor([0.98006658, 0.82533561, 0.54030231], requires_grad=True)
     """
     # pylint: disable=protected-access
     if tape.batch_size is None:
diff --git a/pennylane/transforms/core/transform.py b/pennylane/transforms/core/transform.py
index 0feeb1ca9..c656ca97a 100644
--- a/pennylane/transforms/core/transform.py
+++ b/pennylane/transforms/core/transform.py
@@ -59,7 +59,7 @@ def transform(
 
     .. code-block:: python
 
-        def my_quantum_transform(tape: qml.tape.QuantumTape) -> (Sequence[qml.tape.QuantumTape], callable):
+        def my_quantum_transform(tape: qml.tape.QuantumTape) -> (Sequence[qml.tape.QuantumTape], Callable):
             tape1 = tape
             tape2 = tape.copy()
 
@@ -98,8 +98,8 @@ def transform(
     For QNodes, the dispatched transform populates the ``TransformProgram`` of your QNode. The transform and its
     processing function are applied in the execution.
 
-    >>> transformed_qnode = dispatched_transform(qfunc_circuit)
-    <QNode: wires=2, device='default.qubit', interface='auto', diff_method='best'>
+    >>> transformed_qnode = dispatched_transform(qnode_circuit); transformed_qnode
+    <QNode: wires=2, device='<default.qubit device (wires=2) at 0x2a09261c0>', interface='auto', diff_method='best'>
 
     One subtlety here, this transform would not work for a qfunc because our transform return more than one case. If
     it was not the case you would be able to dispatch on quantum functions.
diff --git a/pennylane/transforms/core/transform_dispatcher.py b/pennylane/transforms/core/transform_dispatcher.py
index 93ae8ed9a..0c1004b1f 100644
--- a/pennylane/transforms/core/transform_dispatcher.py
+++ b/pennylane/transforms/core/transform_dispatcher.py
@@ -68,13 +68,22 @@ class TransformDispatcher:
 
         if isinstance(obj, qml.tape.QuantumScript):
             if self._expand_transform:
-                transformed_tapes, _ = self._expand_transform(obj, *targs, **tkwargs)
-                transformed_tapes, transform_processing_fn = self._transform(
-                    transformed_tapes[0], *targs, **tkwargs
-                )
+                expanded_tapes, expand_processing = self._expand_transform(obj, *targs, **tkwargs)
+                transformed_tapes = []
+                processing_and_sclices = []
+                start = 0
+                for tape in expanded_tapes:
+                    intermediate_tapes, post_processing_fn = self._transform(
+                        tape, *targs, **tkwargs
+                    )
+                    transformed_tapes.extend(intermediate_tapes)
+                    end = start + len(intermediate_tapes)
+                    processing_and_sclices.append(tuple([post_processing_fn, slice(start, end)]))
+                    start = end
 
                 def processing_fn(results):
-                    return transform_processing_fn(results)
+                    processed_results = [fn(results[slice]) for fn, slice in processing_and_sclices]
+                    return expand_processing(processed_results)
 
             else:
                 transformed_tapes, processing_fn = self._transform(obj, *targs, **tkwargs)
diff --git a/pennylane/transforms/hamiltonian_expand.py b/pennylane/transforms/hamiltonian_expand.py
index 8c3dd49e1..139238c19 100644
--- a/pennylane/transforms/hamiltonian_expand.py
+++ b/pennylane/transforms/hamiltonian_expand.py
@@ -64,7 +64,7 @@ def hamiltonian_expand(tape: QuantumTape, group: bool = True) -> (Sequence[Quant
     We can evaluate these tapes on a device:
 
     >>> dev = qml.device("default.qubit", wires=3)
-    >>> res = dev.batch_execute(tapes)
+    >>> res = dev.execute(tapes)
 
     Applying the processing function results in the expectation value of the Hamiltonian:
 
@@ -264,7 +264,7 @@ def sum_expand(tape: QuantumTape, group: bool = True) -> (Sequence[QuantumTape],
     We can evaluate these tapes on a device:
 
     >>> dev = qml.device("default.qubit", wires=3)
-    >>> res = dev.batch_execute(tapes)
+    >>> res = dev.execute(tapes)
 
     Applying the processing function results in the expectation value of the Hamiltonian:
 
diff --git a/pennylane/transforms/metric_tensor.py b/pennylane/transforms/metric_tensor.py
index 479e21d7e..a65cae2af 100644
--- a/pennylane/transforms/metric_tensor.py
+++ b/pennylane/transforms/metric_tensor.py
@@ -290,9 +290,9 @@ def metric_tensor(  # pylint:disable=too-many-arguments
 
         >>> dev = qml.device("default.qubit", wires=3)
         >>> fn(qml.execute(tapes, dev, None))
-        array([[ 0.25      ,  0.        ,  0.42073549],
-               [ 0.        ,  0.00415023, -0.26517488],
-               [ 0.42073549, -0.26517488,  0.24878844]])
+        tensor([[ 0.25      ,  0.        ,  0.42073549],
+                [ 0.        ,  0.00415023, -0.26517488],
+                [ 0.42073549, -0.26517488,  0.24878844]], requires_grad=True)
 
         The first term of the off block-diagonal entries of the full metric tensor are
         computed with Hadamard tests. This first term reads
@@ -340,9 +340,9 @@ def metric_tensor(  # pylint:disable=too-many-arguments
             >>> mt = qml.metric_tensor(circuit, argnum=(0, 2, 3))(weights)
             >>> print(mt)
             [[ 0.          0.          0.          0.        ]
-            [ 0.          0.25       -0.02495835 -0.02495835]
-            [ 0.         -0.02495835  0.01226071  0.01226071]
-            [ 0.         -0.02495835  0.01226071  0.01226071]]
+             [ 0.          0.25       -0.02495835 -0.02495835]
+             [ 0.         -0.02495835  0.01226071  0.01226071]
+             [ 0.         -0.02495835  0.01226071  0.01226071]]
 
         Because the 0-th element of ``weights`` appears second in the QNode and therefore in the
         underlying tape, it is the 1st tape parameter.
diff --git a/pennylane/transforms/optimization/merge_amplitude_embedding.py b/pennylane/transforms/optimization/merge_amplitude_embedding.py
index 76f44a57f..40216c0ee 100644
--- a/pennylane/transforms/optimization/merge_amplitude_embedding.py
+++ b/pennylane/transforms/optimization/merge_amplitude_embedding.py
@@ -55,12 +55,10 @@ def merge_amplitude_embedding(tape: QuantumTape) -> (Sequence[QuantumTape], Call
     >>> optimized_qfunc = qml.transforms.merge_amplitude_embedding(qfunc)
     >>> optimized_qnode = qml.QNode(optimized_qfunc, dev)
     >>> print(qml.draw(optimized_qnode)())
-    0: ─╭●──────────────────────┤  State
-    1: ─╰X──────────────────────┤  State
-    2: ─╭AmplitudeEmbedding(M0)─┤  State
-    3: ─╰AmplitudeEmbedding(M0)─┤  State
-    M0 =
-    [0.+0.j 0.+0.j 0.+0.j 1.+0.j]
+    0: ─╭●───┤  State
+    1: ─╰X───┤  State
+    2: ─╭|Ψ⟩─┤  State
+    3: ─╰|Ψ⟩─┤  State
 
     """
     # Make a working copy of the list to traverse
diff --git a/pennylane/transforms/optimization/merge_rotations.py b/pennylane/transforms/optimization/merge_rotations.py
index 7281825de..ee183a5a5 100644
--- a/pennylane/transforms/optimization/merge_rotations.py
+++ b/pennylane/transforms/optimization/merge_rotations.py
@@ -78,7 +78,7 @@ def merge_rotations(
     By inspection, we can combine the two ``RX`` rotations on the first qubit.
     On the second qubit, we have a cumulative angle of 0, and the gates will cancel.
 
-    >>> optimized_qfunc = merge_rotations()(qfunc)
+    >>> optimized_qfunc = merge_rotations(qfunc)
     >>> optimized_qnode = qml.QNode(optimized_qfunc, dev)
     >>> print(qml.draw(optimized_qnode)(1, 2, 3))
     0: ──RX(3.00)────╭RZ(3.00)─┤  <Z>
@@ -89,7 +89,7 @@ def merge_rotations(
     be merged using the ``include_gates`` argument. For example, if in the above
     circuit we wanted only to merge the "RX" gates, we could do so as follows:
 
-    >>> optimized_qfunc = merge_rotations(include_gates=["RX"])(qfunc)
+    >>> optimized_qfunc = merge_rotations(qfunc, include_gates=["RX"])
     >>> optimized_qnode = qml.QNode(optimized_qfunc, dev)
     >>> print(qml.draw(optimized_qnode)(1, 2, 3))
     0: ──RX(3.00)───────────╭RZ(3.00)────────────┤  <Z>
diff --git a/pennylane/transforms/optimization/pattern_matching.py b/pennylane/transforms/optimization/pattern_matching.py
index 9cb9d728e..77f34d0a6 100644
--- a/pennylane/transforms/optimization/pattern_matching.py
+++ b/pennylane/transforms/optimization/pattern_matching.py
@@ -83,7 +83,7 @@ def pattern_matching_optimization(
 
     >>> dev = qml.device('default.qubit', wires=5)
     >>> qnode = qml.QNode(circuit, dev)
-    >>> optimized_qfunc = pattern_matching_optimization(pattern_tapes=[pattern])(circuit)
+    >>> optimized_qfunc = pattern_matching_optimization(circuit, pattern_tapes=[pattern])
     >>> optimized_qnode = qml.QNode(optimized_qfunc, dev)
 
     >>> print(qml.draw(qnode)())
@@ -101,7 +101,7 @@ def pattern_matching_optimization(
     quantum cost dictionary with a negative cost for ``pennylane.PauliZ``.
 
     >>> my_cost = {"PauliZ": -1 , "S": 1, "Adjoint(S)": 1}
-    >>> optimized_qfunc = pattern_matching_optimization(pattern_tapes=[pattern], custom_quantum_cost=my_cost)(circuit)
+    >>> optimized_qfunc = pattern_matching_optimization(circuit, pattern_tapes=[pattern], custom_quantum_cost=my_cost)
     >>> optimized_qnode = qml.QNode(optimized_qfunc, dev)
 
     >>> print(qml.draw(optimized_qnode)())
@@ -142,7 +142,7 @@ def pattern_matching_optimization(
 
     >>> dev = qml.device('default.qubit', wires=5)
     >>> qnode = qml.QNode(circuit, dev)
-    >>> optimized_qfunc = pattern_matching_optimization(pattern_tapes=[pattern])(circuit)
+    >>> optimized_qfunc = pattern_matching_optimization(circuit, pattern_tapes=[pattern])
     >>> optimized_qnode = qml.QNode(optimized_qfunc, dev)
 
     In our case, it is possible to find three CNOTs and replace this pattern with only two CNOTs and therefore
diff --git a/pennylane/transforms/qcut/montecarlo.py b/pennylane/transforms/qcut/montecarlo.py
index 430105ce9..2ad95bb14 100644
--- a/pennylane/transforms/qcut/montecarlo.py
+++ b/pennylane/transforms/qcut/montecarlo.py
@@ -145,13 +145,13 @@ def cut_circuit_mc(
 
     >>> x = 0.3
     >>> circuit(x)
-    tensor([[1, 1],
-            [0, 1],
-            [0, 1],
-            ...,
-            [0, 1],
-            [0, 1],
-            [0, 1]], requires_grad=True)
+    array([[0., 0.],
+           [0., 1.],
+           [1., 0.],
+           ...,
+           [0., 0.],
+           [0., 0.],
+           [0., 1.]])
 
     Furthermore, the number of shots can be temporarily altered when calling
     the qnode:
diff --git a/pennylane/transforms/qfunc_transforms.py b/pennylane/transforms/qfunc_transforms.py
index 0717e0215..c870cd102 100644
--- a/pennylane/transforms/qfunc_transforms.py
+++ b/pennylane/transforms/qfunc_transforms.py
@@ -325,9 +325,9 @@ def qfunc_transform(tape_transform):
         parameter *and* the transform weights:
 
         >>> circuit(x, y)
-        0.9887793925354269
+        tensor(0.98877939, requires_grad=True)
         >>> qml.grad(circuit)(x, y)
-        (array(-0.02485651), array([-0.02474011, -0.09954244]))
+        (tensor(-0.02485651, requires_grad=True), array([-0.02474011, -0.09954244]))
 
         **Implementation details**
 
diff --git a/pennylane/transforms/qmc.py b/pennylane/transforms/qmc.py
index 7b33c1837..c86a0170c 100644
--- a/pennylane/transforms/qmc.py
+++ b/pennylane/transforms/qmc.py
@@ -305,31 +305,22 @@ def quantum_monte_carlo(fn, wires, target_wire, estimation_wires):
         algorithm
 
         >>> qml.specs(qmc, expansion_strategy="device")()
-        {'resources': Resources(num_wires=12, num_gates=31882, gate_types=defaultdict(<class 'int'>, {'RY': 7747, 'CNOT': 7874, 'Hadamard': 258, 'CZ': 126, 'Adjoint(CNOT)': 7812, 'Adjoint(RY)': 7686, 'PauliX': 252, 'MultiControlledX': 126, 'Adjoint(QFT)': 1}), depth=30610, shots=0),
-        'gate_sizes': defaultdict(int, {1: 15943, 2: 15812, 7: 126, 6: 1}),
-        'gate_types': defaultdict(int,
-                    {'RY': 7747,
-                    'CNOT': 7874,
-                    'Hadamard': 258,
-                    'CZ': 126,
-                    'Adjoint(CNOT)': 7812,
-                    'Adjoint(RY)': 7686,
-                    'PauliX': 252,
-                    'MultiControlledX': 126,
-                    'Adjoint(QFT)': 1}),
-        'num_operations': 31882,
-        'num_observables': 1,
-        'num_diagonalizing_gates': 0,
-        'num_used_wires': 12,
-        'num_trainable_params': 15433,
-        'depth': 30610,
-        'num_device_wires': 12,
-        'device_name': 'default.qubit.autograd',
-        'expansion_strategy': 'gradient',
-        'gradient_options': {},
-        'interface': 'auto',
-        'diff_method': 'best',
-        'gradient_fn': 'backprop'}
+        {'resources': Resources(
+            num_wires=12,
+            num_gates=31882,
+            gate_types=defaultdict(<class 'int'>, {'RY': 7747, 'CNOT': 7874, 'Hadamard': 258, 'CZ': 126, 'Adjoint(CNOT)': 7812, 'Adjoint(RY)': 7686, 'PauliX': 252, 'MultiControlledX': 126, 'Adjoint(QFT)': 1}),
+            gate_sizes=defaultdict(<class 'int'>, {1: 15943, 2: 15812, 7: 126, 6: 1}), depth=30610, shots=Shots(total_shots=None, shot_vector=()),
+         ),
+         'num_observables': 1,
+         'num_diagonalizing_gates': 0,
+         'num_trainable_params': 15433,
+         'num_device_wires': 12,
+         'device_name': 'default.qubit',
+         'expansion_strategy': 'gradient',
+         'gradient_options': {},
+         'interface': 'auto',
+         'diff_method': 'best',
+         'gradient_fn': 'backprop'}
     """
     wires = Wires(wires)
     target_wire = Wires(target_wire)
diff --git a/pennylane/transforms/sign_expand/sign_expand.py b/pennylane/transforms/sign_expand/sign_expand.py
index 312625bec..58d3681dc 100644
--- a/pennylane/transforms/sign_expand/sign_expand.py
+++ b/pennylane/transforms/sign_expand/sign_expand.py
@@ -269,7 +269,7 @@ def sign_expand(  # pylint: disable=too-many-arguments
     one wants to make the circuit approximation of the decomposition:
 
     >>> dev = qml.device("default.qubit", wires=[0,1,2,'Hadamard','Target'])
-    >>> res = dev.batch_execute(tapes)
+    >>> res = dev.execute(tapes)
     >>> fn(res)
     -0.4999999999999999
 
@@ -278,7 +278,7 @@ def sign_expand(  # pylint: disable=too-many-arguments
 
     >>> tapes, fn = qml.transforms.sign_expand(tape, circuit=True, J=20, delta=0)
     >>> dev = qml.device("default.qubit", wires=[0,1,2,'Hadamard','Target'])
-    >>> dev.batch_execute(tapes)
+    >>> dev.execute(tapes)
     >>> fn(res)
     -0.24999999999999994
 
@@ -295,7 +295,7 @@ def sign_expand(  # pylint: disable=too-many-arguments
 
     >>> tapes, fn = qml.transforms.sign_expand(tape, circuit=True, J=20, delta=0)
     >>> dev = qml.device("default.qubit", wires=[0,1,2,'Hadamard','Target'])
-    >>> res = dev.batch_execute(tapes)
+    >>> res = dev.execute(tapes)
     >>> fn(res)
     10.108949481425782
 
diff --git a/pennylane/transforms/specs.py b/pennylane/transforms/specs.py
index ee6c2b6fc..1112e1b9a 100644
--- a/pennylane/transforms/specs.py
+++ b/pennylane/transforms/specs.py
@@ -64,22 +64,17 @@ def specs(qnode, max_expansion=None, expansion_strategy=None):
 
     >>> qml.specs(circuit)(x, add_ry=False)
     {'resources': Resources(num_wires=2, num_gates=2, gate_types=defaultdict(<class 'int'>, {'RX': 1, 'CNOT': 1}),
-    gate_sizes=defaultdict(int, {1: 1, 2: 1}), depth=2, shots=Shots(total_shots=None, shot_vector=())),
-    'gate_sizes': defaultdict(int, {1: 1, 2: 1}),
-    'gate_types': defaultdict(int, {'RX': 1, 'CNOT': 1}),
-    'num_operations': 2,
+    gate_sizes=defaultdict(<class 'int'>, {1: 1, 2: 1}), depth=2, shots=Shots(total_shots=None, shot_vector=())),
     'num_observables': 1,
     'num_diagonalizing_gates': 0,
-    'num_used_wires': 2,
     'num_trainable_params': 1,
-    'depth': 2,
     'num_device_wires': 2,
     'device_name': 'default.qubit',
     'expansion_strategy': 'gradient',
     'gradient_options': {'shifts': 0.7853981633974483},
     'interface': 'auto',
     'diff_method': 'parameter-shift',
-    'gradient_fn': 'pennylane.gradients.parameter_shift.param_shift',
+    'gradient_fn': 'pennylane.transforms.core.transform_dispatcher.param_shift',
     'num_gradient_executions': 2}
 
     """
diff --git a/pennylane/transforms/split_non_commuting.py b/pennylane/transforms/split_non_commuting.py
index ef4c64b32..bd2fe2c9a 100644
--- a/pennylane/transforms/split_non_commuting.py
+++ b/pennylane/transforms/split_non_commuting.py
@@ -103,10 +103,7 @@ def split_non_commuting(tape: qml.tape.QuantumTape) -> (Sequence[qml.tape.Quantu
     :math:`(\langle \sigma_x^0 \rangle, \langle \sigma_z^0 \rangle, \langle \sigma_y^1 \rangle, \langle \sigma_z^0\sigma_z^1 \rangle)`.
 
     >>> circuit0([np.pi/4, np.pi/4])
-    (tensor(0.70710678, requires_grad=True),
-    tensor(0.5, requires_grad=True),
-    tensor(0., requires_grad=True),
-    tensor(0.5, requires_grad=True))
+    [0.7071067811865475, 0.49999999999999994, 0.0, 0.49999999999999994]
 
 
     .. details::
diff --git a/pennylane/transforms/tape_expand.py b/pennylane/transforms/tape_expand.py
index 8b418191f..a349d94d6 100644
--- a/pennylane/transforms/tape_expand.py
+++ b/pennylane/transforms/tape_expand.py
@@ -467,7 +467,7 @@ def set_decomposition(custom_decomps, dev, decomp_depth=10):
     Now let's set up a context where the custom decomposition will be applied:
 
     >>> with qml.transforms.set_decomposition({qml.CNOT : custom_cnot}, dev):
-    ...     print(qml.draw(circuit)())
+    ...     print(qml.draw(circuit, wire_order=[0, 1])())
     0: ────╭●────┤  <Z>
     1: ──H─╰Z──H─┤
 
diff --git a/pennylane/transforms/zx/converter.py b/pennylane/transforms/zx/converter.py
index efbe3d6ed..d005322f9 100644
--- a/pennylane/transforms/zx/converter.py
+++ b/pennylane/transforms/zx/converter.py
@@ -236,7 +236,7 @@ def to_zx(tape, expand_measurements=False):  # pylint: disable=unused-argument
                 return qml.expval(qml.PauliZ(wires=0))
 
         >>> mod_5_4()
-        -0.9999999999999989
+        tensor(1., requires_grad=True)
 
     .. note::
 
@@ -337,7 +337,7 @@ def _to_zx_transform(
             else:
                 expanded_operations.append(op)
 
-        expanded_tape = QuantumScript(expanded_operations, mapped_tape.measurements, [])
+        expanded_tape = QuantumScript(expanded_operations, mapped_tape.measurements)
 
         _add_operations_to_graph(expanded_tape, graph, gate_types, q_mapper, c_mapper)
 
diff --git a/pennylane/vqe/vqe.py b/pennylane/vqe/vqe.py
index bcc27697e..3dca1bcca 100644
--- a/pennylane/vqe/vqe.py
+++ b/pennylane/vqe/vqe.py
@@ -130,10 +130,12 @@ class ExpvalCost:
 
         Grouping these commuting observables leads to fewer device executions:
 
-        >>> cost_opt(params)
-        >>> ex_opt = dev.num_executions
-        >>> cost_no_opt(params)
-        >>> ex_no_opt = dev.num_executions - ex_opt
+        >>> with qml.Tracker(dev) as tracker:
+        ...     cost_opt(params)
+        >>> ex_opt = tracker.totals["executions"]
+        >>> with tracker:
+        ...     cost_no_opt(params)
+        >>> ex_no_opt = tracker.totals["executions"]
         >>> print("Number of executions:", ex_no_opt)
         Number of executions: 2
         >>> print("Number of executions (optimized):", ex_opt)
diff --git a/requirements.txt b/requirements.txt
index 4aab5ba3a..56a611bf6 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -10,7 +10,7 @@ toml~=0.10
 appdirs~=1.4
 semantic_version~=2.10
 dask[delayed]~=2022.4.1
-autoray>=0.3.1
+autoray>=0.6.1
 matplotlib~=3.5
 opt_einsum~=3.3
 requests~=2.31.0
diff --git a/setup.py b/setup.py
index 34021c1c3..15f5a0cfd 100644
--- a/setup.py
+++ b/setup.py
@@ -28,7 +28,7 @@ requirements = [
     "toml",
     "appdirs",
     "semantic-version>=2.7",
-    "autoray>=0.3.1",
+    "autoray>=0.6.1",
     "cachetools",
     "pennylane-lightning>=0.32",
     "requests",


Test patch:
diff --git a/pennylane/devices/tests/__init__.py b/pennylane/devices/tests/__init__.py
index e7b3f68e0..5c508e5e5 100644
--- a/pennylane/devices/tests/__init__.py
+++ b/pennylane/devices/tests/__init__.py
@@ -65,7 +65,7 @@ Finally, the tests can be invoked within a Python session via the :func:`~.test_
 function:
 
 >>> from pennylane.devices.tests import test_device
->>> test_device("default.qubit")
+>>> test_device("default.qubit.legacy")
 
 For more details on the available arguments, see the :func:`~.test_device` documentation.
 
@@ -120,7 +120,7 @@ def test_device(
     **Example**
 
     >>> from pennylane.devices.tests import test_device
-    >>> test_device("default.qubit")
+    >>> test_device("default.qubit.legacy")
     ================================ test session starts =======================================
     platform linux -- Python 3.7.7, pytest-5.4.2, py-1.8.1, pluggy-0.13.1
     rootdir: /home/josh/xanadu/pennylane/pennylane/devices/tests, inifile: pytest.ini
diff --git a/tests/templates/test_subroutines/test_hilbert_schmidt.py b/tests/templates/test_subroutines/test_hilbert_schmidt.py
index 8bb1ad4bd..9d3975f68 100644
--- a/tests/templates/test_subroutines/test_hilbert_schmidt.py
+++ b/tests/templates/test_subroutines/test_hilbert_schmidt.py
@@ -20,7 +20,8 @@ import pennylane as qml
 
 
 # pylint: disable=protected-access
-def test_flatten_unflatten():
+@pytest.mark.parametrize("op_type", (qml.HilbertSchmidt, qml.LocalHilbertSchmidt))
+def test_flatten_unflatten(op_type):
     """Test the flatten and unflatten methods."""
     u_tape = qml.tape.QuantumScript([qml.Hadamard("a"), qml.Identity("b")])
 
@@ -28,7 +29,7 @@ def test_flatten_unflatten():
         qml.RZ(params, wires=1)
 
     v_wires = qml.wires.Wires((0, 1))
-    op = qml.HilbertSchmidt([0.1], v_function=v_circuit, v_wires=v_wires, u_tape=u_tape)
+    op = op_type([0.1], v_function=v_circuit, v_wires=v_wires, u_tape=u_tape)
     data, metadata = op._flatten()
 
     assert data == (0.1,)
@@ -253,9 +254,15 @@ class TestLocalHilbertSchmidt:
         op = qml.LocalHilbertSchmidt([0.1], v_function=v_circuit, v_wires=[1], u_tape=U)
 
         with qml.queuing.AnnotatedQueue() as q_tape_dec:
-            op.decomposition()
+            decomp = op.decomposition()
+
+        unqueued_decomp = op.decomposition()
 
         tape_dec = qml.tape.QuantumScript.from_queue(q_tape_dec)
+
+        assert all(qml.equal(o1, o2) for o1, o2 in zip(decomp, tape_dec))
+        assert all(qml.equal(o1, o2) for o1, o2 in zip(decomp, unqueued_decomp))
+
         expected_operations = [
             qml.Hadamard(wires=[0]),
             qml.CNOT(wires=[0, 1]),
@@ -328,11 +335,46 @@ class TestLocalHilbertSchmidt:
             qml.CNOT(wires=[0, 2]),
             qml.Hadamard(wires=[0]),
         ]
+        assert tape_dec.operations == expected_operations
 
-        for i, j in zip(tape_dec.operations, expected_operations):
-            assert i.name == j.name
-            assert i.data == j.data
-            assert i.wires == j.wires
+    def test_qnode_integration(self):
+        """Test that the local hilbert schmidt template can be used inside a qnode."""
+
+        u_tape = qml.tape.QuantumTape([qml.CZ(wires=(0, 1))])
+
+        def v_function(params):
+            qml.RZ(params[0], wires=2)
+            qml.RZ(params[1], wires=3)
+            qml.CNOT(wires=[2, 3])
+            qml.RZ(params[2], wires=3)
+            qml.CNOT(wires=[2, 3])
+
+        dev = qml.device("default.qubit", wires=4)
+
+        @qml.qnode(dev)
+        def local_hilbert_test(v_params, v_function, v_wires, u_tape):
+            qml.LocalHilbertSchmidt(v_params, v_function=v_function, v_wires=v_wires, u_tape=u_tape)
+            return qml.probs(u_tape.wires + v_wires)
+
+        # pylint: disable=unsubscriptable-object
+        def cost_lhst(parameters, v_function, v_wires, u_tape):
+            return (
+                1
+                - local_hilbert_test(
+                    v_params=parameters, v_function=v_function, v_wires=v_wires, u_tape=u_tape
+                )[0]
+            )
+
+        res = cost_lhst(
+            [3 * qml.numpy.pi / 2, 3 * qml.numpy.pi / 2, qml.numpy.pi / 2],
+            v_function=v_function,
+            v_wires=[2, 3],
+            u_tape=u_tape,
+        )
+
+        assert qml.math.allclose(res, 0.5)
+        # the answer is currently 0.5, and I'm going to assume that's correct. This test will let us know
+        # if the answer changes.
 
     def test_v_not_quantum_function(self):
         """Test that we cannot pass a non quantum function to the HS operation"""
diff --git a/tests/transforms/test_qcut.py b/tests/transforms/test_qcut.py
index 7c96c26d5..a323448ad 100644
--- a/tests/transforms/test_qcut.py
+++ b/tests/transforms/test_qcut.py
@@ -5595,6 +5595,28 @@ class TestCutCircuitWithHamiltonians:
             # each frag should have the device size constraint satisfied.
             assert all(len(set(e[2] for e in f.edges.data("wire"))) <= device_size for f in frags)
 
+    def test_hamiltonian_with_tape(self):
+        """Test that an expand function that generates multiple tapes is applied before the transform and the transform
+        returns correct results."""
+        ops = [qml.Identity(0), qml.PauliZ(0), qml.PauliZ(1), qml.PauliZ(2)]
+        coeffs = [0.4567, 0.25, 0.25, 0.5]
+
+        H = qml.Hamiltonian(observables=ops, coeffs=coeffs)
+
+        dev = qml.device("lightning.qubit", wires=4)
+
+        circuit = [
+            qml.IsingXX(0.1234, wires=[0, 3]),
+            qml.WireCut(0),
+        ]
+
+        tape = qml.tape.QuantumScript(circuit, measurements=[qml.expval(H)])
+        cut_tapes, proc_fn = qml.cut_circuit(tape, device_wires=range(3))
+
+        assert np.allclose(
+            qml.execute([tape], dev, None)[0], proc_fn(qml.execute(cut_tapes, dev, None))
+        )
+
     def test_raise_with_hamiltonian(self):
         """Test that exception is correctly raise when caclulating expectation values of multiple Hamiltonians"""