Repository: PennyLaneAI/pennylane
Pull request #3561 (2022-12-19 02:52:02)

Daily rc sync to master 2022-12-19
Automatic sync from the release candidate to master during a feature freeze.

codecov[bot]: # [Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3561?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
> Merging [#3561](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3561?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (d822ce9) into [master](https://codecov.io/gh/PennyLaneAI/pennylane/commit/7ac2c10ce0ce021fc1d42f57b2bfefcc4df92a58?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (7ac2c10) will **not change** coverage.
> The diff coverage is `100.00%`.

```diff
@@           Coverage Diff           @@
##           master    #3561   +/-   ##
=======================================
  Coverage   99.87%   99.87%           
=======================================
  Files         311      311           
  Lines       27903    27903           
=======================================
  Hits        27867    27867           
  Misses         36       36           
```


| [Impacted Files](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3561?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | Coverage Δ | |
|---|---|---|
| [pennylane/operation.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3561/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL29wZXJhdGlvbi5weQ==) | `97.43% <ø> (ø)` | |
| [pennylane/ops/functions/equal.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3561/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL29wcy9mdW5jdGlvbnMvZXF1YWwucHk=) | `99.05% <ø> (ø)` | |
| [pennylane/ops/identity.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3561/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL29wcy9pZGVudGl0eS5weQ==) | `100.00% <ø> (ø)` | |
| [pennylane/ops/op\_math/exp.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3561/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL29wcy9vcF9tYXRoL2V4cC5weQ==) | `100.00% <ø> (ø)` | |
| [pennylane/ops/qubit/matrix\_ops.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3561/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL29wcy9xdWJpdC9tYXRyaXhfb3BzLnB5) | `100.00% <ø> (ø)` | |
| [pennylane/ops/qubit/non\_parametric\_ops.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3561/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL29wcy9xdWJpdC9ub25fcGFyYW1ldHJpY19vcHMucHk=) | `100.00% <ø> (ø)` | |
| [pennylane/ops/qubit/observables.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3561/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL29wcy9xdWJpdC9vYnNlcnZhYmxlcy5weQ==) | `100.00% <ø> (ø)` | |
| [pennylane/ops/qubit/parametric\_ops.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3561/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL29wcy9xdWJpdC9wYXJhbWV0cmljX29wcy5weQ==) | `100.00% <ø> (ø)` | |
| [pennylane/ops/qutrit/non\_parametric\_ops.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3561/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL29wcy9xdXRyaXQvbm9uX3BhcmFtZXRyaWNfb3BzLnB5) | `100.00% <ø> (ø)` | |
| [pennylane/ops/qutrit/observables.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3561/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL29wcy9xdXRyaXQvb2JzZXJ2YWJsZXMucHk=) | `100.00% <ø> (ø)` | |
| ... and [7 more](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3561/diff?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | |

Help us with your feedback. Take ten seconds to tell us [how you rate us](https://about.codecov.io/nps?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI). Have a feature suggestion? [Share it here.](https://app.codecov.io/gh/feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI)


Patch:
diff --git a/doc/releases/changelog-0.28.0.md b/doc/releases/changelog-0.28.0.md
index 95ad139de..2d5bd8bf0 100644
--- a/doc/releases/changelog-0.28.0.md
+++ b/doc/releases/changelog-0.28.0.md
@@ -489,6 +489,9 @@
 * QPE can now accept a target operator instead of a matrix and target wires pair.
   [(#3373)](https://github.com/PennyLaneAI/pennylane/pull/3373)
 
+* The `qml.ops.op_math.Controlled.map_wires` method now uses `base.map_wires` internally instead of the private `_wires` property setter.
+  [(#3405)](https://github.com/PennyLaneAI/pennylane/pull/3405)
+
 * A new function called `qml.tape.make_qscript` has been created for converting a quantum function into a quantum script. This replaces `qml.transforms.make_tape`.
   [(#3429)](https://github.com/PennyLaneAI/pennylane/pull/3429)
 
@@ -573,9 +576,6 @@
 * The `qml.Operation.get_parameter_shift` method has been removed. The `gradients` module should be used for general parameter-shift rules instead.
   [(#3419)](https://github.com/PennyLaneAI/pennylane/pull/3419)
 
-* `_wires` has been removed from the `qml.ops.op_math.Controlled.map_wires` method.
-  [(#3405)](https://github.com/PennyLaneAI/pennylane/pull/3405)
-
 * The signature of the `QubitDevice.statistics` method has been changed from 
 
   ```python
diff --git a/pennylane/interfaces/jax.py b/pennylane/interfaces/jax.py
index 739bab620..49a01c7e9 100644
--- a/pennylane/interfaces/jax.py
+++ b/pennylane/interfaces/jax.py
@@ -535,9 +535,9 @@ def _execute_fwd_new(
         return res, jacs
 
     @execute_wrapper.defjvp
-    def execute_wrapper_jvp(primal, tangents):
+    def execute_wrapper_jvp(primals, tangents):
         """Primals[0] are parameters as Jax tracers and tangents[0] is a list of tangent vectors as Jax tracers."""
-        res, jacs = execute_wrapper(primal[0])
+        res, jacs = execute_wrapper(primals[0])
         multi_measurements = [len(tape.measurements) > 1 for tape in tapes]
 
         jvps = _compute_jvps(jacs, tangents[0], multi_measurements)
diff --git a/pennylane/interfaces/jax_jit_tuple.py b/pennylane/interfaces/jax_jit_tuple.py
index c0555fa19..a4572f6e3 100644
--- a/pennylane/interfaces/jax_jit_tuple.py
+++ b/pennylane/interfaces/jax_jit_tuple.py
@@ -22,7 +22,7 @@ import jax.numpy as jnp
 
 import pennylane as qml
 from pennylane.interfaces.jax import _compute_jvps
-from pennylane.interfaces.jax_jit import _validate_jax_version, _numeric_type_to_dtype
+from pennylane.interfaces.jax_jit import _numeric_type_to_dtype, _validate_jax_version
 
 dtype = jnp.float64
 
@@ -107,9 +107,9 @@ def execute_tuple(tapes, device, execute_fn, gradient_fn, gradient_kwargs, _n=1,
             during the forward pass. This function must return a tuple ``(results, jacobians)``.
             If ``jacobians`` is an empty list, then ``gradient_fn`` is used to
             compute the gradients during the backwards pass.
+        gradient_fn (callable): the gradient function to use to compute quantum gradients
         gradient_kwargs (dict): dictionary of keyword arguments to pass when
             determining the gradients of tapes
-        gradient_fn (callable): the gradient function to use to compute quantum gradients
         _n (int): a positive integer used to track nesting of derivatives, for example
             if the nth-order derivative is requested.
         max_diff (int): If ``gradient_fn`` is a gradient transform, this option specifies
@@ -353,9 +353,9 @@ def _execute_fwd_tuple(
         return res, jacs
 
     @execute_wrapper.defjvp
-    def execute_wrapper_jvp(primal, tangents):
+    def execute_wrapper_jvp(primals, tangents):
         """Primals[0] are parameters as Jax tracers and tangents[0] is a list of tangent vectors as Jax tracers."""
-        res, jacs = execute_wrapper(primal[0])
+        res, jacs = execute_wrapper(primals[0])
         multi_measurements = [len(tape.measurements) > 1 for tape in tapes]
 
         if len(tapes) == 1:
diff --git a/pennylane/operation.py b/pennylane/operation.py
index ec6c20a93..0f511c5bd 100644
--- a/pennylane/operation.py
+++ b/pennylane/operation.py
@@ -656,7 +656,7 @@ class Operator(abc.ABC):
     def compute_eigvals(*params, **hyperparams):
         r"""Eigenvalues of the operator in the computational basis (static method).
 
-        If :attr:`diagonalizing_gates` are specified and implement a unitary :math:`U^{dagger}`,
+        If :attr:`diagonalizing_gates` are specified and implement a unitary :math:`U^{\dagger}`,
         the operator can be reconstructed as
 
         .. math:: O = U \Sigma U^{\dagger},
@@ -679,7 +679,7 @@ class Operator(abc.ABC):
     def eigvals(self):
         r"""Eigenvalues of the operator in the computational basis.
 
-        If :attr:`diagonalizing_gates` are specified and implement a unitary :math:`U^{dagger}`, the operator
+        If :attr:`diagonalizing_gates` are specified and implement a unitary :math:`U^{\dagger}`, the operator
         can be reconstructed as
 
         .. math:: O = U \Sigma U^{\dagger},
diff --git a/pennylane/ops/functions/equal.py b/pennylane/ops/functions/equal.py
index fc04bede1..392043b1e 100644
--- a/pennylane/ops/functions/equal.py
+++ b/pennylane/ops/functions/equal.py
@@ -52,7 +52,7 @@ def equal(
 
         The kwargs ``check_interface`` and ``check_trainability`` can only be set when
         comparing ``Operation`` objects. Comparisons of ``MeasurementProcess``
-        or ``Observable`` objects will use the defualt value of ``True`` for both, regardless
+        or ``Observable`` objects will use the default value of ``True`` for both, regardless
         of what the user specifies when calling the function. For subclasses of ``SymbolicOp``
         with an ``Operation`` as a base, the kwargs will be applied to the base comparison.
 
diff --git a/pennylane/ops/identity.py b/pennylane/ops/identity.py
index 5f2a8038d..6144bc816 100644
--- a/pennylane/ops/identity.py
+++ b/pennylane/ops/identity.py
@@ -58,7 +58,7 @@ class Identity(CVObservable, Operation):
     def compute_eigvals(n_wires=1):  # pylint: disable=arguments-differ
         r"""Eigenvalues of the operator in the computational basis (static method).
 
-        If :attr:`diagonalizing_gates` are specified and implement a unitary :math:`U^{dagger}`,
+        If :attr:`diagonalizing_gates` are specified and implement a unitary :math:`U^{\dagger}`,
         the operator can be reconstructed as
 
         .. math:: O = U \Sigma U^{\dagger},
diff --git a/pennylane/ops/op_math/exp.py b/pennylane/ops/op_math/exp.py
index cd0cce765..ba33aa93f 100644
--- a/pennylane/ops/op_math/exp.py
+++ b/pennylane/ops/op_math/exp.py
@@ -364,16 +364,16 @@ class Evolution(Exp):
 
     Args:
         base (~.operation.Operator): The operator to be used as a generator, G.
-        param (float): The evolution parameter, x. This parameter is not expected to have
+        param (float): The evolution parameter, x. This parameter is expected not to have
             any complex component.
 
     Returns:
-       :class:`Evolution`: A :class`~.operation.Operator` representing an operator exponential of the form :math:`e^{ix\hat{G}}`,
+       :class:`Evolution`: A :class:`~.operation.Operator` representing an operator exponential of the form :math:`e^{ix\hat{G}}`,
        where x is real.
 
     **Usage Details**
 
-    In contrast to the general :class:`~.Exp` class, the Evolution operator :math:`e^{ix\hat{G}}` is constrained to have a single trainable
+    In contrast to the general :class:`~.Exp` class, the ``Evolution`` operator :math:`e^{ix\hat{G}}` is constrained to have a single trainable
     parameter, x. Any parameters contained in the base operator are not trainable. This allows the operator
     to be differentiated with regard to the evolution parameter. Defining a mathematically identical operator
     using the :class:`~.Exp` class will be incompatible with a variety of PennyLane functions that require only a single
@@ -398,7 +398,7 @@ class Evolution(Exp):
 
     >>> @qml.qnode(qml.device('default.qubit', wires=1))
     ... def circuit(x):
-    ...     Evolution(qml.PauliX(0), -0.5 * x)
+    ...     qml.ops.Evolution(qml.PauliX(0), -0.5 * x)
     ...     return qml.expval(qml.PauliZ(0))
     >>> print(qml.draw(circuit)(1.23))
     0: ──Exp(-0.61j X)─┤  <Z>
diff --git a/pennylane/ops/qubit/matrix_ops.py b/pennylane/ops/qubit/matrix_ops.py
index 864b66596..ef54ee8a4 100644
--- a/pennylane/ops/qubit/matrix_ops.py
+++ b/pennylane/ops/qubit/matrix_ops.py
@@ -457,7 +457,7 @@ class DiagonalQubitUnitary(Operation):
     def compute_eigvals(D):  # pylint: disable=arguments-differ
         r"""Eigenvalues of the operator in the computational basis (static method).
 
-        If :attr:`diagonalizing_gates` are specified and implement a unitary :math:`U^{dagger}`,
+        If :attr:`diagonalizing_gates` are specified and implement a unitary :math:`U^{\dagger}`,
         the operator can be reconstructed as
 
         .. math:: O = U \Sigma U^{\dagger},
diff --git a/pennylane/ops/qubit/non_parametric_ops.py b/pennylane/ops/qubit/non_parametric_ops.py
index 3821fb3d9..392e8b8e4 100644
--- a/pennylane/ops/qubit/non_parametric_ops.py
+++ b/pennylane/ops/qubit/non_parametric_ops.py
@@ -89,7 +89,7 @@ class Hadamard(Observable, Operation):
     def compute_eigvals():  # pylint: disable=arguments-differ
         r"""Eigenvalues of the operator in the computational basis (static method).
 
-        If :attr:`diagonalizing_gates` are specified and implement a unitary :math:`U^{dagger}`,
+        If :attr:`diagonalizing_gates` are specified and implement a unitary :math:`U^{\dagger}`,
         the operator can be reconstructed as
 
         .. math:: O = U \Sigma U^{\dagger},
@@ -240,7 +240,7 @@ class PauliX(Observable, Operation):
     def compute_eigvals():  # pylint: disable=arguments-differ
         r"""Eigenvalues of the operator in the computational basis (static method).
 
-        If :attr:`diagonalizing_gates` are specified and implement a unitary :math:`U^{dagger}`,
+        If :attr:`diagonalizing_gates` are specified and implement a unitary :math:`U^{\dagger}`,
         the operator can be reconstructed as
 
         .. math:: O = U \Sigma U^{\dagger},
@@ -394,7 +394,7 @@ class PauliY(Observable, Operation):
     def compute_eigvals():  # pylint: disable=arguments-differ
         r"""Eigenvalues of the operator in the computational basis (static method).
 
-        If :attr:`diagonalizing_gates` are specified and implement a unitary :math:`U^{dagger}`,
+        If :attr:`diagonalizing_gates` are specified and implement a unitary :math:`U^{\dagger}`,
         the operator can be reconstructed as
 
         .. math:: O = U \Sigma U^{\dagger},
@@ -546,7 +546,7 @@ class PauliZ(Observable, Operation):
     def compute_eigvals():  # pylint: disable=arguments-differ
         r"""Eigenvalues of the operator in the computational basis (static method).
 
-        If :attr:`diagonalizing_gates` are specified and implement a unitary :math:`U^{dagger}`,
+        If :attr:`diagonalizing_gates` are specified and implement a unitary :math:`U^{\dagger}`,
         the operator can be reconstructed as
 
         .. math:: O = U \Sigma U^{\dagger},
@@ -688,7 +688,7 @@ class S(Operation):
     def compute_eigvals():  # pylint: disable=arguments-differ
         r"""Eigenvalues of the operator in the computational basis (static method).
 
-        If :attr:`diagonalizing_gates` are specified and implement a unitary :math:`U^{dagger}`,
+        If :attr:`diagonalizing_gates` are specified and implement a unitary :math:`U^{\dagger}`,
         the operator can be reconstructed as
 
         .. math:: O = U \Sigma U^{\dagger},
@@ -797,7 +797,7 @@ class T(Operation):
     def compute_eigvals():  # pylint: disable=arguments-differ
         r"""Eigenvalues of the operator in the computational basis (static method).
 
-        If :attr:`diagonalizing_gates` are specified and implement a unitary :math:`U^{dagger}`,
+        If :attr:`diagonalizing_gates` are specified and implement a unitary :math:`U^{\dagger}`,
         the operator can be reconstructed as
 
         .. math:: O = U \Sigma U^{\dagger},
@@ -906,7 +906,7 @@ class SX(Operation):
     def compute_eigvals():  # pylint: disable=arguments-differ
         r"""Eigenvalues of the operator in the computational basis (static method).
 
-        If :attr:`diagonalizing_gates` are specified and implement a unitary :math:`U^{dagger}`,
+        If :attr:`diagonalizing_gates` are specified and implement a unitary :math:`U^{\dagger}`,
         the operator can be reconstructed as
 
         .. math:: O = U \Sigma U^{\dagger},
@@ -1100,7 +1100,7 @@ class CZ(Operation):
     def compute_eigvals():  # pylint: disable=arguments-differ
         r"""Eigenvalues of the operator in the computational basis (static method).
 
-        If :attr:`diagonalizing_gates` are specified and implement a unitary :math:`U^{dagger}`,
+        If :attr:`diagonalizing_gates` are specified and implement a unitary :math:`U^{\dagger}`,
         the operator can be reconstructed as
 
         .. math:: O = U \Sigma U^{\dagger},
@@ -1244,7 +1244,7 @@ class CH(Operation):
     r"""CH(wires)
     The controlled-Hadamard operator
 
-    .. math:: CY = \begin{bmatrix}
+    .. math:: CH = \begin{bmatrix}
             1 & 0 & 0 & 0 \\
             0 & 1 & 0 & 0 \\
             0 & 0 & \frac{1}{\sqrt{2}} & \frac{1}{\sqrt{2}} \\
@@ -1493,7 +1493,7 @@ class ECR(Operation):
     def compute_eigvals():
         r"""Eigenvalues of the operator in the computational basis (static method).
 
-        If :attr:`diagonalizing_gates` are specified and implement a unitary :math:`U^{dagger}`,
+        If :attr:`diagonalizing_gates` are specified and implement a unitary :math:`U^{\dagger}`,
         the operator can be reconstructed as
 
         .. math:: O = U \Sigma U^{\dagger},
@@ -1611,7 +1611,7 @@ class ISWAP(Operation):
     def compute_eigvals():  # pylint: disable=arguments-differ
         r"""Eigenvalues of the operator in the computational basis (static method).
 
-        If :attr:`diagonalizing_gates` are specified and implement a unitary :math:`U^{dagger}`,
+        If :attr:`diagonalizing_gates` are specified and implement a unitary :math:`U^{\dagger}`,
         the operator can be reconstructed as
 
         .. math:: O = U \Sigma U^{\dagger},
@@ -1734,7 +1734,7 @@ class SISWAP(Operation):
     def compute_eigvals():  # pylint: disable=arguments-differ
         r"""Eigenvalues of the operator in the computational basis (static method).
 
-        If :attr:`diagonalizing_gates` are specified and implement a unitary :math:`U^{dagger}`,
+        If :attr:`diagonalizing_gates` are specified and implement a unitary :math:`U^{\dagger}`,
         the operator can be reconstructed as
 
         .. math:: O = U \Sigma U^{\dagger},
diff --git a/pennylane/ops/qubit/observables.py b/pennylane/ops/qubit/observables.py
index f62b581f4..751501208 100644
--- a/pennylane/ops/qubit/observables.py
+++ b/pennylane/ops/qubit/observables.py
@@ -436,7 +436,7 @@ class Projector(Observable):
     def compute_eigvals(basis_state):  # pylint: disable=arguments-differ
         r"""Eigenvalues of the operator in the computational basis (static method).
 
-        If :attr:`diagonalizing_gates` are specified and implement a unitary :math:`U^{dagger}`,
+        If :attr:`diagonalizing_gates` are specified and implement a unitary :math:`U^{\dagger}`,
         the operator can be reconstructed as
 
         .. math:: O = U \Sigma U^{\dagger},
diff --git a/pennylane/ops/qubit/parametric_ops.py b/pennylane/ops/qubit/parametric_ops.py
index fd0bde632..248abe5b5 100644
--- a/pennylane/ops/qubit/parametric_ops.py
+++ b/pennylane/ops/qubit/parametric_ops.py
@@ -319,7 +319,7 @@ class RZ(Operation):
     def compute_eigvals(theta):  # pylint: disable=arguments-differ
         r"""Eigenvalues of the operator in the computational basis (static method).
 
-        If :attr:`diagonalizing_gates` are specified and implement a unitary :math:`U^{dagger}`,
+        If :attr:`diagonalizing_gates` are specified and implement a unitary :math:`U^{\dagger}`,
         the operator can be reconstructed as
 
         .. math:: O = U \Sigma U^{\dagger},
@@ -462,7 +462,7 @@ class PhaseShift(Operation):
     def compute_eigvals(phi):  # pylint: disable=arguments-differ
         r"""Eigenvalues of the operator in the computational basis (static method).
 
-        If :attr:`diagonalizing_gates` are specified and implement a unitary :math:`U^{dagger}`,
+        If :attr:`diagonalizing_gates` are specified and implement a unitary :math:`U^{\dagger}`,
         the operator can be reconstructed as
 
         .. math:: O = U \Sigma U^{\dagger},
@@ -636,7 +636,7 @@ class ControlledPhaseShift(Operation):
     def compute_eigvals(phi):  # pylint: disable=arguments-differ
         r"""Eigenvalues of the operator in the computational basis (static method).
 
-        If :attr:`diagonalizing_gates` are specified and implement a unitary :math:`U^{dagger}`,
+        If :attr:`diagonalizing_gates` are specified and implement a unitary :math:`U^{\dagger}`,
         the operator can be reconstructed as
 
         .. math:: O = U \Sigma U^{\dagger},
@@ -1533,7 +1533,7 @@ class MultiRZ(Operation):
     def compute_eigvals(theta, num_wires):  # pylint: disable=arguments-differ
         r"""Eigenvalues of the operator in the computational basis (static method).
 
-        If :attr:`diagonalizing_gates` are specified and implement a unitary :math:`U^{dagger}`,
+        If :attr:`diagonalizing_gates` are specified and implement a unitary :math:`U^{\dagger}`,
         the operator can be reconstructed as
 
         .. math:: O = U \Sigma U^{\dagger},
@@ -1829,7 +1829,7 @@ class PauliRot(Operation):
     def compute_eigvals(theta, pauli_word):  # pylint: disable=arguments-differ
         r"""Eigenvalues of the operator in the computational basis (static method).
 
-        If :attr:`diagonalizing_gates` are specified and implement a unitary :math:`U^{dagger}`,
+        If :attr:`diagonalizing_gates` are specified and implement a unitary :math:`U^{\dagger}`,
         the operator can be reconstructed as
 
         .. math:: O = U \Sigma U^{\dagger},
@@ -2360,7 +2360,7 @@ class CRZ(Operation):
     def compute_eigvals(theta):  # pylint: disable=arguments-differ
         r"""Eigenvalues of the operator in the computational basis (static method).
 
-        If :attr:`diagonalizing_gates` are specified and implement a unitary :math:`U^{dagger}`,
+        If :attr:`diagonalizing_gates` are specified and implement a unitary :math:`U^{\dagger}`,
         the operator can be reconstructed as
 
         .. math:: O = U \Sigma U^{\dagger},
@@ -3486,7 +3486,7 @@ class IsingZZ(Operation):
     def compute_eigvals(phi):  # pylint: disable=arguments-differ
         r"""Eigenvalues of the operator in the computational basis (static method).
 
-        If :attr:`diagonalizing_gates` are specified and implement a unitary :math:`U^{dagger}`,
+        If :attr:`diagonalizing_gates` are specified and implement a unitary :math:`U^{\dagger}`,
         the operator can be reconstructed as
 
         .. math:: O = U \Sigma U^{\dagger},
@@ -3685,7 +3685,7 @@ class IsingXY(Operation):
     def compute_eigvals(phi):  # pylint: disable=arguments-differ
         r"""Eigenvalues of the operator in the computational basis (static method).
 
-        If :attr:`diagonalizing_gates` are specified and implement a unitary :math:`U^{dagger}`,
+        If :attr:`diagonalizing_gates` are specified and implement a unitary :math:`U^{\dagger}`,
         the operator can be reconstructed as
 
         .. math:: O = U \Sigma U^{\dagger},
@@ -3841,7 +3841,7 @@ class PSWAP(Operation):
     def compute_eigvals(phi):  # pylint: disable=arguments-differ
         r"""Eigenvalues of the operator in the computational basis (static method).
 
-        If :attr:`diagonalizing_gates` are specified and implement a unitary :math:`U^{dagger}`,
+        If :attr:`diagonalizing_gates` are specified and implement a unitary :math:`U^{\dagger}`,
         the operator can be reconstructed as
 
         .. math:: O = U \Sigma U^{\dagger},
diff --git a/pennylane/ops/qutrit/non_parametric_ops.py b/pennylane/ops/qutrit/non_parametric_ops.py
index 311235098..da9722517 100644
--- a/pennylane/ops/qutrit/non_parametric_ops.py
+++ b/pennylane/ops/qutrit/non_parametric_ops.py
@@ -78,7 +78,7 @@ class TShift(Operation):
     def compute_eigvals():
         r"""Eigenvalues of the operator in the computational basis (static method).
 
-        If :attr:`diagonalizing_gates` are specified and implement a unitary :math:`U^{dagger}`,
+        If :attr:`diagonalizing_gates` are specified and implement a unitary :math:`U^{\dagger}`,
         the operator can be reconstructed as
 
         .. math:: O = U \Sigma U^{\dagger},
@@ -169,7 +169,7 @@ class TClock(Operation):
     def compute_eigvals():
         r"""Eigenvalues of the operator in the computational basis (static method).
 
-        If :attr:`diagonalizing_gates` are specified and implement a unitary :math:`U^{dagger}`,
+        If :attr:`diagonalizing_gates` are specified and implement a unitary :math:`U^{\dagger}`,
         the operator can be reconstructed as
 
         .. math:: O = U \Sigma U^{\dagger},
@@ -286,7 +286,7 @@ class TAdd(Operation):
     def compute_eigvals():
         r"""Eigenvalues of the operator in the computational basis (static method).
 
-        If :attr:`diagonalizing_gates` are specified and implement a unitary :math:`U^{dagger}`,
+        If :attr:`diagonalizing_gates` are specified and implement a unitary :math:`U^{\dagger}`,
         the operator can be reconstructed as
 
         .. math:: O = U \Sigma U^{\dagger},
@@ -403,7 +403,7 @@ class TSWAP(Operation):
     def compute_eigvals():
         r"""Eigenvalues of the operator in the computational basis (static method).
 
-        If :attr:`diagonalizing_gates` are specified and implement a unitary :math:`U^{dagger}`,
+        If :attr:`diagonalizing_gates` are specified and implement a unitary :math:`U^{\dagger}`,
         the operator can be reconstructed as
 
         .. math:: O = U \Sigma U^{\dagger},
diff --git a/pennylane/ops/qutrit/observables.py b/pennylane/ops/qutrit/observables.py
index 7a2c7870a..e83c840b4 100644
--- a/pennylane/ops/qutrit/observables.py
+++ b/pennylane/ops/qutrit/observables.py
@@ -286,7 +286,7 @@ class GellMann(Observable):
     def compute_eigvals(index):  # pylint: disable=arguments-differ
         r"""Eigenvalues of the operator in the computational basis (static method).
 
-        If :attr:`diagonalizing_gates` are specified and implement a unitary :math:`U^{dagger}`,
+        If :attr:`diagonalizing_gates` are specified and implement a unitary :math:`U^{\dagger}`,
         the operator can be reconstructed as
 
         .. math:: O = U \Sigma U^{\dagger},
diff --git a/pennylane/pauli/conversion.py b/pennylane/pauli/conversion.py
index 8e054f9ad..b427c0f52 100644
--- a/pennylane/pauli/conversion.py
+++ b/pennylane/pauli/conversion.py
@@ -14,18 +14,18 @@
 """
 Utility functions to convert between ``~.PauliSentence`` and other PennyLane operators.
 """
-from typing import Union
-from operator import matmul
-from itertools import product
 from functools import reduce, singledispatch
+from itertools import product
+from operator import matmul
+from typing import Union
 
 import numpy as np
 
 from pennylane.operation import Tensor
 from pennylane.ops import Hamiltonian, Identity, PauliX, PauliY, PauliZ, Prod, SProd, Sum
 
+from .pauli_arithmetic import I, PauliSentence, PauliWord, X, Y, Z, mat_map, op_map
 from .utils import is_pauli_word
-from .pauli_arithmetic import PauliWord, PauliSentence, I, X, Y, Z, mat_map, op_map
 
 
 def pauli_decompose(
@@ -52,22 +52,22 @@ def pauli_decompose(
 
     >>> A = np.array(
     ... [[-2, -2+1j, -2, -2], [-2-1j,  0,  0, -1], [-2,  0, -2, -1], [-2, -1, -1,  0]])
-    >>> H = pauli_decompose(A)
+    >>> H = qml.pauli_decompose(A)
     >>> print(H)
-    (-1.0) [I0 I1]
-    + (-1.5) [X1]
-    + (-0.5) [Y1]
-    + (-1.0) [Z1]
-    + (-1.5) [X0]
+    (-1.5) [I0 X1]
+    + (-1.5) [X0 I1]
+    + (-1.0) [I0 I1]
+    + (-1.0) [I0 Z1]
     + (-1.0) [X0 X1]
+    + (-0.5) [I0 Y1]
     + (-0.5) [X0 Z1]
-    + (1.0) [Y0 Y1]
     + (-0.5) [Z0 X1]
     + (-0.5) [Z0 Y1]
+    + (1.0) [Y0 Y1]
 
     We can return a ``~.PauliSentence`` instance by using the keyword argument ``pauli=True``:
 
-    >>> ps = pauli_decompose(A, pauli=True)
+    >>> ps = qml.pauli_decompose(A, pauli=True)
     >>> print(ps)
     -1.0 * I
     + -1.5 * X(1)
@@ -82,7 +82,7 @@ def pauli_decompose(
 
     We can also set custom wires using the ``wire_order`` argument:
 
-    >>> ps = pauli_decompose(A, pauli=True, wire_order=['a', 'b'])
+    >>> ps = qml.pauli_decompose(A, pauli=True, wire_order=['a', 'b'])
     >>> print(ps)
     -1.0 * I
     + -1.5 * X(b)
@@ -153,7 +153,8 @@ def pauli_sentence(op):
     Raises:
         ValueError: Op must be a linear combination of Pauli operators
 
-    Returns ~.PauliSentence
+    Returns:
+        .PauliSentence: the PauliSentence representation of an arithmetic operator or Hamiltonian
     """
     raise ValueError(f"Op must be a linear combination of Pauli operators only, got: {op}")
 
diff --git a/pennylane/pauli/pauli_arithmetic.py b/pennylane/pauli/pauli_arithmetic.py
index a24cb7275..789d66714 100644
--- a/pennylane/pauli/pauli_arithmetic.py
+++ b/pennylane/pauli/pauli_arithmetic.py
@@ -210,7 +210,7 @@ class PauliWord(dict):
         return reduce(kron, (matrix_map[self[w]] for w in wire_order))
 
     def operation(self, wire_order=None):
-        """Returns a native PennyLane``~.Operator`` representing the PauliWord."""
+        """Returns a native PennyLane :class:`~pennylane.operation.Operation` representing the PauliWord."""
         if len(self) == 0:
             if wire_order in (None, [], wires.Wires([])):
                 raise ValueError("Can't get the operation for an empty PauliWord.")
@@ -220,7 +220,7 @@ class PauliWord(dict):
         return factors[0] if len(factors) == 1 else prod(*factors)
 
     def hamiltonian(self, wire_order=None):
-        """Return ~.Hamiltonian representing the PauliWord"""
+        """Return :class:`~pennylane.Hamiltonian` representing the PauliWord"""
         if len(self) == 0:
             if wire_order in (None, [], wires.Wires([])):
                 raise ValueError("Can't get the Hamiltonian for an empty PauliWord.")
@@ -352,7 +352,7 @@ class PauliSentence(dict):
         return math.expand_matrix(reduced_mat, result_wire_order, wire_order=wire_order)
 
     def operation(self, wire_order=None):
-        """Returns a native PennyLane``~.Operator`` representing the PauliSentence."""
+        """Returns a native PennyLane :class:`~pennylane.operation.Operation` representing the PauliSentence."""
         if len(self) == 0:
             if wire_order in (None, [], wires.Wires([])):
                 raise ValueError("Can't get the operation for an empty PauliSentence.")
@@ -364,7 +364,7 @@ class PauliSentence(dict):
         return summands[0] if len(summands) == 1 else op_sum(*summands)
 
     def hamiltonian(self, wire_order=None):
-        """Returns a native PennyLane ~.Hamiltonian representing the PauliSentence."""
+        """Returns a native PennyLane :class:`~pennylane.Hamiltonian` representing the PauliSentence."""
         if len(self) == 0:
             if wire_order in (None, [], wires.Wires([])):
                 raise ValueError("Can't get the Hamiltonian for an empty PauliSentence.")
diff --git a/pennylane/qchem/structure.py b/pennylane/qchem/structure.py
index e1c8329df..62db33d99 100644
--- a/pennylane/qchem/structure.py
+++ b/pennylane/qchem/structure.py
@@ -435,17 +435,17 @@ def mol_data(identifier, identifier_type="name"):
 
     >>> mol_data("BeH2")
     (['Be', 'H', 'H'],
-     array([[ 4.79405604,  0.29290815,  0.        ],
-            [ 3.77946   , -0.29290815,  0.        ],
-            [ 5.80884105, -0.29290815,  0.        ]]))
+    tensor([[ 4.79404621,  0.29290755,  0.        ],
+            [ 3.77945225, -0.29290755,  0.        ],
+            [ 5.80882913, -0.29290755,  0.        ]], requires_grad=True))
 
     >>> mol_data(223, "CID")
     (['N', 'H', 'H', 'H', 'H'],
-    array([[ 4.79404621,  0.        ,  0.        ],
-            [ 5.80882913,  0.5858151 ,  0.        ],
-            [ 3.77945225, -0.5858151 ,  0.        ],
-            [ 4.20823111,  1.01459396,  0.        ],
-            [ 5.3798613 , -1.01459396,  0.        ]]))
+    tensor([[ 0.        ,  0.        ,  0.        ],
+            [ 1.82264085,  0.52836742,  0.40402345],
+            [ 0.01417295, -1.67429735, -0.98038991],
+            [-0.98927163, -0.22714508,  1.65369933],
+            [-0.84773114,  1.373075  , -1.07733286]], requires_grad=True))
 
     .. details::
 
@@ -453,35 +453,35 @@ def mol_data(identifier, identifier_type="name"):
 
         >>> mol_data("74-82-8", "CAS")
         (['C', 'H', 'H', 'H', 'H'],
-        array([[ 0.        ,  0.        ,  0.        ],
+        tensor([[ 0.        ,  0.        ,  0.        ],
                 [ 1.04709725,  1.51102501,  0.93824902],
                 [ 1.29124986, -1.53710323, -0.47923455],
                 [-1.47058487, -0.70581271,  1.26460472],
-                [-0.86795121,  0.7320799 , -1.7236192 ]]))
+                [-0.86795121,  0.7320799 , -1.7236192 ]], requires_grad=True))
 
         >>> mol_data("[C]", "SMILES")
         (['C', 'H', 'H', 'H', 'H'],
-        array([[ 0.        ,  0.        ,  0.        ],
+        tensor([[ 0.        ,  0.        ,  0.        ],
                 [ 1.04709725,  1.51102501,  0.93824902],
                 [ 1.29124986, -1.53710323, -0.47923455],
                 [-1.47058487, -0.70581271,  1.26460472],
-                [-0.86795121,  0.7320799 , -1.7236192 ]]))
+                [-0.86795121,  0.7320799 , -1.7236192 ]], requires_grad=True))
 
         >>> mol_data("InChI=1S/CH4/h1H4", "InChI")
         (['C', 'H', 'H', 'H', 'H'],
-        array([[ 0.        ,  0.        ,  0.        ],
+        tensor([[ 0.        ,  0.        ,  0.        ],
                 [ 1.04709725,  1.51102501,  0.93824902],
                 [ 1.29124986, -1.53710323, -0.47923455],
                 [-1.47058487, -0.70581271,  1.26460472],
-                [-0.86795121,  0.7320799 , -1.7236192 ]]))
+                [-0.86795121,  0.7320799 , -1.7236192 ]], requires_grad=True))
 
         >>> mol_data("VNWKTOKETHGBQD-UHFFFAOYSA-N", "InChIKey")
         (['C', 'H', 'H', 'H', 'H'],
-        array([[ 0.        ,  0.        ,  0.        ],
+        tensor([[ 0.        ,  0.        ,  0.        ],
                 [ 1.04709725,  1.51102501,  0.93824902],
                 [ 1.29124986, -1.53710323, -0.47923455],
                 [-1.47058487, -0.70581271,  1.26460472],
-                [-0.86795121,  0.7320799 , -1.7236192 ]]))
+                [-0.86795121,  0.7320799 , -1.7236192 ]], requires_grad=True))
 
     """
 
@@ -508,7 +508,7 @@ def mol_data(identifier, identifier_type="name"):
     if identifier_type.lower() == "cid":
         cid = int(identifier)
     else:
-        if identifier_type.lower() not in identifier_patterns.keys():
+        if identifier_type.lower() not in identifier_patterns:
             raise ValueError(
                 "Specified identifier type is not supported. Supported type are: name, CAS, SMILES, InChI, InChIKey."
             )
diff --git a/pennylane/return_types.py b/pennylane/return_types.py
index c5333c9d9..62a158486 100644
--- a/pennylane/return_types.py
+++ b/pennylane/return_types.py
@@ -181,6 +181,147 @@ def enable_return():
 
             AttributeError: 'tuple' object has no attribute 'shape'
 
+        **JAX interface upgrades: higher-order derivatives and mixing measurements**
+
+        Higher-order derivatives can now be computed with the JAX interface:
+
+        .. code-block:: python
+
+            import jax
+
+            qml.enable_return()
+
+            dev = qml.device("lightning.qubit", wires=2)
+
+            par_0 = jax.numpy.array(0.1)
+            par_1 = jax.numpy.array(0.2)
+
+            @qml.qnode(dev, interface="jax", diff_method="parameter-shift", max_diff=2)
+            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.PauliX(1))
+
+        >>> jax.hessian(circuit, argnums=[0, 1])(par_0, par_1)
+        ((Array(-0.19767681, dtype=float32, weak_type=True),
+          Array(-0.09784342, dtype=float32, weak_type=True)),
+         (Array(-0.09784339, dtype=float32, weak_type=True),
+          Array(-0.19767687, dtype=float32, weak_type=True)))
+
+        The new return types system also unlocks the use of ``probs`` mixed with different measurements with JAX:
+
+        .. code-block:: python
+
+            import jax
+
+            qml.enable_return()
+
+            dev = qml.device("default.qubit", wires=2)
+            qml.enable_return()
+
+            @qml.qnode(dev, interface="jax")
+            def circuit(a):
+              qml.RX(a[0], wires=0)
+              qml.CNOT(wires=(0, 1))
+              qml.RY(a[1], wires=1)
+              qml.RZ(a[2], wires=1)
+              return qml.expval(qml.PauliZ(wires=0)), qml.probs(wires=[0, 1])
+
+            x = jax.numpy.array([0.1, 0.2, 0.3])
+
+        >>> jax.jacobian(circuit)(x)
+        (Array([-9.9833414e-02, -7.4505806e-09,  6.9285655e-10], dtype=float32),
+         Array([[-4.9419206e-02, -9.9086545e-02,  3.4938008e-09],
+                      [-4.9750542e-04,  9.9086538e-02,  1.2768372e-10],
+                      [ 4.9750548e-04,  2.4812977e-04,  4.8371929e-13],
+                      [ 4.9419202e-02, -2.4812980e-04,  2.6696912e-11]],            dtype=float32))
+
+        where before the following error was raised:
+
+        .. code-block:: python
+
+            ValueError: All input arrays must have the same shape.
+
+        The new return types system also unlocks the use of shot vectors with all the previous features. For example you
+        can take the second derivative and multiple measurement with with JAX:
+
+        .. code-block:: python
+
+            import jax
+
+            qml.enable_return()
+
+            dev = qml.device("default.qubit", wires=2, shots=(1, 10000))
+
+            params = jax.numpy.array([0.1, 0.2])
+
+            @qml.qnode(dev, interface="jax", diff_method="parameter-shift", max_diff=2)
+            def circuit(x):
+                qml.RX(x[0], wires=[0])
+                qml.RY(x[1], wires=[1])
+                qml.CNOT(wires=[0, 1])
+                return qml.var(qml.PauliZ(0) @ qml.PauliX(1)), qml.probs(wires=[0])
+
+        >>> jax.hessian(circuit)(params)
+        ((Array([[ 0.,  0.],
+                        [ 2., -3.]], dtype=float32),
+          Array([[[-0.5,  0. ],
+                       [ 0. ,  0. ]],
+                      [[ 0.5,  0. ],
+                       [ 0. ,  0. ]]], dtype=float32)),
+         (Array([[ 0.07677898,  0.0563341 ],
+                       [ 0.07238522, -1.830669  ]], dtype=float32),
+          Array([[[-4.9707499e-01,  2.9999996e-04],
+                        [-6.2500127e-04,  1.2500001e-04]],
+                       [[ 4.9707499e-01, -2.9999996e-04],
+                        [ 6.2500127e-04, -1.2500001e-04]]], dtype=float32)))
+
+        Note that failing to set the ``max-diff`` with jitting will raise a somewhat unrelated error:
+
+        .. code-block:: python
+
+            import jax
+            import pennylane as qml
+
+            jax.config.update("jax_enable_x64", True)
+            qml.enable_return()
+
+            dev = qml.device("default.qubit", wires=2, shots=10000)
+
+            params = jax.numpy.array([0.1, 0.2])
+
+            @jax.jit
+            @qml.qnode(dev, interface="jax", diff_method="parameter-shift") # Note: max_diff=2 should be passed here
+            def circuit(x):
+                qml.RX(x[0], wires=[0])
+                qml.RY(x[1], wires=[1])
+                qml.CNOT(wires=[0, 1])
+                return qml.var(qml.PauliZ(0) @ qml.PauliX(1)), qml.probs(wires=[0])
+
+        .. code-block:: python
+
+            >>> jax.hessian(circuit)(params)
+            ~/anaconda3/lib/python3.8/site-packages/jax/_src/callback.py in pure_callback_jvp_rule(***failed resolving arguments***)
+                 53 def pure_callback_jvp_rule(*args, **kwargs):
+                 54   del args, kwargs
+            ---> 55   raise ValueError(
+                 56       "Pure callbacks do not support JVP. "
+                 57       "Please use `jax.custom_jvp` to use callbacks while taking gradients.")
+            ValueError: Pure callbacks do not support JVP. Please use `jax.custom_jvp` to use callbacks while taking gradients.
+
+        Correctly specifying ``max_diff=2`` as a QNode argument helps compute the Hessian:
+
+        .. code-block:: python
+
+            >>> jax.hessian(circuit)(params)
+            (Array([[ 0.06617428,  0.07165382],
+                    [ 0.07165382, -1.8348092 ]], dtype=float64),
+             Array([[[-4.974e-01, -2.025e-03],
+                     [-2.025e-03,  1.000e-04]],
+                    [[ 4.974e-01,  2.025e-03],
+                     [ 2.025e-03, -1.000e-04]]], dtype=float64))
+
     .. details::
         :title: Autograd Usage Details
 
@@ -491,148 +632,6 @@ def enable_return():
         >>> torch.autograd.functional.jacobian(jac_fn, (a, b))
         ((tensor([-0.9950, -0.4925, -0.0050,  0.0050,  0.4925]), tensor([ 0.0000,  0.0050, -0.0050,  0.0050, -0.0050])),
         (tensor([ 0.0000,  0.0050, -0.0050,  0.0050, -0.0050]), tensor([ 0.0000, -0.4888,  0.4888,  0.0012, -0.0012])))
-
-    .. details::
-        :title: JAX Usage Details
-
-        Higher-order derivatives can now be computed with the JAX interface:
-
-        .. code-block:: python
-
-            import jax
-
-            qml.enable_return()
-
-            dev = qml.device("lightning.qubit", wires=2)
-
-            par_0 = jax.numpy.array(0.1)
-            par_1 = jax.numpy.array(0.2)
-
-            @qml.qnode(dev, interface="jax", diff_method="parameter-shift", max_diff=2)
-            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.PauliX(1))
-
-        >>> jax.hessian(circuit, argnums=[0, 1])(par_0, par_1)
-        ((Array(-0.19767681, dtype=float32, weak_type=True),
-          Array(-0.09784342, dtype=float32, weak_type=True)),
-         (Array(-0.09784339, dtype=float32, weak_type=True),
-          Array(-0.19767687, dtype=float32, weak_type=True)))
-
-        The new return types system also unlocks the use of ``probs`` mixed with different measurements with JAX:
-
-        .. code-block:: python
-
-            import jax
-
-            qml.enable_return()
-
-            dev = qml.device("default.qubit", wires=2)
-            qml.enable_return()
-
-            @qml.qnode(dev, interface="jax")
-            def circuit(a):
-              qml.RX(a[0], wires=0)
-              qml.CNOT(wires=(0, 1))
-              qml.RY(a[1], wires=1)
-              qml.RZ(a[2], wires=1)
-              return qml.expval(qml.PauliZ(wires=0)), qml.probs(wires=[0, 1])
-
-            x = jax.numpy.array([0.1, 0.2, 0.3])
-
-        >>> jax.jacobian(circuit)(x)
-        (Array([-9.9833414e-02, -7.4505806e-09,  6.9285655e-10], dtype=float32),
-         Array([[-4.9419206e-02, -9.9086545e-02,  3.4938008e-09],
-                [-4.9750542e-04,  9.9086538e-02,  1.2768372e-10],
-                [ 4.9750548e-04,  2.4812977e-04,  4.8371929e-13],
-                [ 4.9419202e-02, -2.4812980e-04,  2.6696912e-11]],            dtype=float32))
-
-        where before the following error was raised:
-
-        .. code-block:: python
-
-            ValueError: All input arrays must have the same shape.
-
-        The new return types system also unlocks the use of shot vectors with all the previous features. For example you
-        can take the second derivative and multiple measurement with with JAX:
-
-        .. code-block:: python
-
-            import jax
-
-            qml.enable_return()
-
-            dev = qml.device("default.qubit", wires=2, shots=(1, 10000))
-
-            params = jax.numpy.array([0.1, 0.2])
-
-            @qml.qnode(dev, interface="jax", diff_method="parameter-shift", max_diff=2)
-            def circuit(x):
-                qml.RX(x[0], wires=[0])
-                qml.RY(x[1], wires=[1])
-                qml.CNOT(wires=[0, 1])
-                return qml.var(qml.PauliZ(0) @ qml.PauliX(1)), qml.probs(wires=[0])
-
-        >>> jax.hessian(circuit)(params)
-        ((Array([[ 0.,  0.],
-                [ 2., -3.]], dtype=float32),
-          Array([[[-0.5,  0. ],
-                   [ 0. ,  0. ]],
-                  [[ 0.5,  0. ],
-                   [ 0. ,  0. ]]], dtype=float32)),
-         (Array([[ 0.07677898,  0.0563341 ],
-                   [ 0.07238522, -1.830669  ]], dtype=float32),
-          Array([[[-4.9707499e-01,  2.9999996e-04],
-                    [-6.2500127e-04,  1.2500001e-04]],
-                   [[ 4.9707499e-01, -2.9999996e-04],
-                    [ 6.2500127e-04, -1.2500001e-04]]], dtype=float32)))
-
-        Note that failing to set the ``max-diff`` with jitting will raise a somewhat unrelated error:
-
-        .. code-block:: python
-
-            import jax
-            import pennylane as qml
-
-            jax.config.update("jax_enable_x64", True)
-            qml.enable_return()
-
-            dev = qml.device("default.qubit", wires=2, shots=10000)
-
-            params = jax.numpy.array([0.1, 0.2])
-
-            @jax.jit
-            @qml.qnode(dev, interface="jax", diff_method="parameter-shift") # Note: max_diff=2 should be passed here
-            def circuit(x):
-                qml.RX(x[0], wires=[0])
-                qml.RY(x[1], wires=[1])
-                qml.CNOT(wires=[0, 1])
-                return qml.var(qml.PauliZ(0) @ qml.PauliX(1)), qml.probs(wires=[0])
-
-        .. code-block:: python
-
-            >>> jax.hessian(circuit)(params)
-            ~/anaconda3/lib/python3.8/site-packages/jax/_src/callback.py in pure_callback_jvp_rule(***failed resolving arguments***)
-                 53 def pure_callback_jvp_rule(*args, **kwargs):
-                 54   del args, kwargs
-            ---> 55   raise ValueError(
-                 56       "Pure callbacks do not support JVP. "
-                 57       "Please use `jax.custom_jvp` to use callbacks while taking gradients.")
-            ValueError: Pure callbacks do not support JVP. Please use `jax.custom_jvp` to use callbacks while taking gradients.
-
-        Correctly specifying ``max_diff=2`` as a QNode argument helps compute the Hessian:
-
-        .. code-block:: python
-
-            >>> jax.hessian(circuit)(params)
-            (Array([[ 0.06617428,  0.07165382],
-                    [ 0.07165382, -1.8348092 ]], dtype=float64),
-             Array([[[-4.974e-01, -2.025e-03],
-                     [-2.025e-03,  1.000e-04]],
-                    [[ 4.974e-01,  2.025e-03],
-                     [ 2.025e-03, -1.000e-04]]], dtype=float64))
     """
 
     global __activated
diff --git a/pennylane/transforms/hamiltonian_expand.py b/pennylane/transforms/hamiltonian_expand.py
index d4b42ead4..40745aa70 100644
--- a/pennylane/transforms/hamiltonian_expand.py
+++ b/pennylane/transforms/hamiltonian_expand.py
@@ -247,15 +247,21 @@ def sum_expand(tape: QuantumTape, group=True):
     >>> tapes, fn = qml.transforms.sum_expand(tape, group=False)
     >>> for tape in tapes:
     ...     print(tape.measurements)
-    [expval(PauliZ(wires=[0]))]
     [expval(PauliY(wires=[2]) @ PauliZ(wires=[1]))]
     [expval(PauliZ(wires=[2]))]
     [expval(PauliZ(wires=[1]))]
+    [expval(PauliZ(wires=[0]))]
     [expval(PauliX(wires=[1]))]
 
     Five tapes are generated: the first three contain the summands of the `Sum` operator,
     and the last two contain the remaining observables. Note that the scalars of the scalar products
-    have been removed. These values will be multiplied by the result obtained from the execution.
+    have been removed. In the processing function, these values will be multiplied by the result obtained
+    from executing the tapes.
+
+    Additionally, the observable expval(PauliZ(wires=[2])) occurs twice in the original tape, but only once
+    in the transformed tapes. When there are multipe identical measurements in the circuit, the measurement
+    is performed once and the outcome is copied when obtaining the final result. This will also be resolved
+    when the processing function is applied.
 
     We can evaluate these tapes on a device:
 


Test patch:
diff --git a/pennylane/devices/tests/test_measurements.py b/pennylane/devices/tests/test_measurements.py
index 756c7666f..15d723a80 100644
--- a/pennylane/devices/tests/test_measurements.py
+++ b/pennylane/devices/tests/test_measurements.py
@@ -1571,9 +1571,6 @@ class TestStateMeasurement:
         dev = device(2)
         _skip_test_for_braket(dev)
 
-        if dev.shots is not None:
-            pytest.skip("Some plugins don't update state information when shots is not None.")
-
         if dev.shots is not None:
             pytest.skip("Some plugins don't update state information when shots is not None.")
 
diff --git a/tests/ops/qubit/test_non_parametric_ops.py b/tests/ops/qubit/test_non_parametric_ops.py
index 761b7158b..c0316852f 100644
--- a/tests/ops/qubit/test_non_parametric_ops.py
+++ b/tests/ops/qubit/test_non_parametric_ops.py
@@ -1323,7 +1323,7 @@ class TestControlledMethod:
         assert qml.equal(out, qml.CZ(("a", 0)))
 
     def test_Hadamard(self):
-        """Test the PauliZ _controlled method."""
+        """Test the Hadamard _controlled method."""
         out = qml.Hadamard(0)._controlled("a")
         assert qml.equal(out, qml.CH(("a", 0)))