Repository: PennyLaneAI/pennylane
Pull request #2059 (2021-12-22 18:06:01)

Remove the stacking (and transposing) behaviour of `qml.jacobian`
**Context:**
This PR implements the proposal #2056. This is a breaking change for multiple equal-shaped QNode arguments in `qml.jacobian` and in `metric_tensor`, `adjoint_metric_tensor`, and `classical_jacobian` when using the Autograd interface.

**Description of the Change:**
stacking and transposition is removed from `qml.jacobian`.
A lot of tests have to be adapted to this, including those for `metric_tensor`, `adjoint_metric_tensor`, and `classical_jacobian`.
Furthermore, `qml.transforms.metric_tensor._contract_metric_tensor_with_cjac` is change to not replicate the old Jacobian behaviour.

**Benefits:**
- Removes arguably strange behaviour regarding the output shape of `qml.jacobian`, and even more so for second order derivatives like metric tensors and Hessians.
- Makes the output shape more predictable 
- Reduces code complexity
- Makes the behaviour of `qml.jacobian` more similar to `jax.jacobian`, `tf.GradientTape().jacobian` and `torch.autograd.functional.jacobian`.

**Possible Drawbacks:**
- It is a breaking change. 
- For multiple equal-shaped QNode arguments, higher-order derivatives are not computed properly yet,
in particular `QNGOptimizer` does not work for two of its tests any more, because it is written to use a single metric tensor, not one per QNode argument.
- Using a computed Jacobian is easier if the Jacobian w.r.t. different QNode arguments are stacked. On the other hand, the way they were stacked previously is not necessarily how one would expect, and in particular for array-valued functions, the output dimensions are transposed.

**Related GitHub Issues:**
#1884 #1725 #1989 #1992 #2056


github-actions[bot]: Hello. You may have forgotten to update the changelog!
Please edit [`doc/releases/changelog-dev.md`](/PennyLaneAI/pennylane/blob/master/doc/releases/changelog-dev.md) with:
* A one-to-two sentence description of the change. You may include a small working example for new features.
* A link back to this PR.
* Your name (or GitHub username) in the contributors section.
codecov[bot]: # [Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2059?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
> Merging [#2059](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2059?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (da91072) into [master](https://codecov.io/gh/PennyLaneAI/pennylane/commit/deffcee1615687d74a25f4385dd1f097d0d7e9cd?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (deffcee) will **decrease** coverage by `2.12%`.
> The diff coverage is `100.00%`.

[![Impacted file tree graph](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2059/graphs/tree.svg?width=650&height=150&src=pr&token=9S3YzlUgGk&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI)](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2059?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI)

```diff
@@            Coverage Diff             @@
##           master    #2059      +/-   ##
==========================================
- Coverage   99.17%   97.05%   -2.13%     
==========================================
  Files         225      225              
  Lines       17291    17265      -26     
==========================================
- Hits        17149    16756     -393     
- Misses        142      509     +367     
```


| [Impacted Files](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2059?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | Coverage Δ | |
|---|---|---|
| [pennylane/\_grad.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2059/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL19ncmFkLnB5) | `100.00% <100.00%> (ø)` | |
| [pennylane/gradients/gradient\_transform.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2059/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2dyYWRpZW50cy9ncmFkaWVudF90cmFuc2Zvcm0ucHk=) | `79.48% <100.00%> (-20.52%)` | :arrow_down: |
| [pennylane/transforms/adjoint\_metric\_tensor.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2059/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3RyYW5zZm9ybXMvYWRqb2ludF9tZXRyaWNfdGVuc29yLnB5) | `100.00% <100.00%> (ø)` | |
| [pennylane/transforms/classical\_jacobian.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2059/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3RyYW5zZm9ybXMvY2xhc3NpY2FsX2phY29iaWFuLnB5) | `77.08% <100.00%> (-22.92%)` | :arrow_down: |
| [pennylane/transforms/metric\_tensor.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2059/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3RyYW5zZm9ybXMvbWV0cmljX3RlbnNvci5weQ==) | `99.47% <100.00%> (-0.03%)` | :arrow_down: |
| [pennylane/interfaces/batch/jax.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2059/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2ludGVyZmFjZXMvYmF0Y2gvamF4LnB5) | `0.00% <0.00%> (-100.00%)` | :arrow_down: |
| [pennylane/interfaces/jax.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2059/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2ludGVyZmFjZXMvamF4LnB5) | `0.00% <0.00%> (-97.44%)` | :arrow_down: |
| [pennylane/devices/default\_qubit\_jax.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2059/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2RldmljZXMvZGVmYXVsdF9xdWJpdF9qYXgucHk=) | `0.00% <0.00%> (-96.25%)` | :arrow_down: |
| [pennylane/qnn/cost.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2059/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3Fubi9jb3N0LnB5) | `42.85% <0.00%> (-57.15%)` | :arrow_down: |
| ... and [18 more](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2059/diff?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | |

------

[Continue to review full report at Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2059?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2059?src=pr&el=footer&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI). Last update [deffcee...da91072](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2059?src=pr&el=lastupdated&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).

dime10: Hi @dwierichs, thank you so much for working on this, this change is really helpful 🥳 

One question I have is, is it worth considering removing only the transpose and leaving the stacking?

This would still make the behaviour more consistent (the Jacobian is always of shape (num_args, *qnode_output_shape, *arg_shape), where the only difference is whether the first dimension is a tuple or array/tensor), while maybe negating some of the drawbacks you listed (working with a single tensor where possible).
dwierichs: Hey @dime10, that's a good question! On one hand it might come in handy because of the listed possible drawbacks that indeed would be reduced by keeping the `stack`. Personally, I am a bit hesitant with this, because a) we have the opportunity to match the behaviour of the other interfaces (even though it might be a bit less convenient to work with the output in cases where the `stack` succeeds) and b) the output shape of `qml.jacobian` becomes dependent on the argument shapes so that building more advanced functionalities on top of it might become cumbersome?
I find it hard to weigh the consistency between interfaces and stable output format against simpler output formatting when possible. What do you think? Maybe @josh146 also has an opinion on how this should be handled?
josh146: @dime10 @dwierichs I fear by this point you are both more informed than I am on the nuances 😄 

From a higher-level, I am in favour of the approach which brings us more in line with the existing conventions of the autodiff frameworks.

Further, an issue I see with the stacking is that it is only well-defined if the QNode itself returns a single tensor. This is not always the case (and in the future might be even more of a minority --- I can imagine cases where QNodes return tuples of tensors). And, because of the existing `try`-`except` around the stacking, the `qml.jacobian` return type will fluctuate!
dime10: True these are good points 😃 In this case I do think it sensible to remove both the transpose and the stacking.
dwierichs: @josh146 This would be ready for another review, and the mentioned override of CodeCov would be required :)
dime10: I can review this next week :)
josh146: [sc-13026]
josh146: [sc-13047]

Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index db73f3558..550ac85b7 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -216,6 +216,18 @@
 
 <h3>Improvements</h3>
 
+* The function `qml.math.safe_squeeze` is introduced and `gradient_transform` allows
+  for QNode argument axes of size `1`.
+  [(#2080)](https://github.com/PennyLaneAI/pennylane/pull/2080)
+
+  `qml.math.safe_squeeze` wraps `qml.math.squeeze`, with slight modifications:
+
+  - When provided the `axis` keyword argument, axes that do not have size `1` will be
+    ignored, instead of raising an error.
+
+  - The keyword argument `exclude_axis` allows to explicitly exclude axes from the
+    squeezing.
+
 * The `adjoint` transform now raises and error whenever the object it is applied to
   is not callable.
   [(#2060)](https://github.com/PennyLaneAI/pennylane/pull/2060)
@@ -279,6 +291,35 @@
 
 <h3>Breaking changes</h3>
 
+* `qml.metric_tensor`, `qml.adjoint_metric_tensor` and `qml.transforms.classical_jacobian`
+  now follow a different convention regarding their output shape when being used
+  with the Autograd interface
+  [(#2059)](https://github.com/PennyLaneAI/pennylane/pull/2059)
+
+  See the previous entry for details. This breaking change immediately follows from
+  the change in `qml.jacobian` whenever `hybrid=True` is used in the above methods.
+
+* `qml.jacobian` now follows a different convention regarding its output shape.
+  [(#2059)](https://github.com/PennyLaneAI/pennylane/pull/2059)
+
+  Previously, `qml.jacobian` would attempt to stack the Jacobian for multiple
+  QNode arguments, which succeeded whenever the arguments have the same shape.
+  In this case, the stacked Jacobian would also be transposed, leading to the 
+  output shape `(*reverse_QNode_args_shape, *reverse_output_shape, num_QNode_args)`
+
+  If no stacking and transposing occurs, the output shape instead is a `tuple`
+  where each entry corresponds to one QNode argument and has the shape
+  `(*output_shape, *QNode_arg_shape)`.
+
+  This breaking change alters the behaviour in the first case and removes the attempt
+  to stack and transpose, so that the output always has the shape of the second
+  type.
+
+  Note that the behaviour is unchanged --- that is, the Jacobian tuple is unpacked into
+  a single Jacobian --- if `argnum=None` and there is only one QNode argument
+  with respect to which the differentiation takes place, or if an integer
+  is provided as `argnum`.
+
 * The behaviour of `RotosolveOptimizer` has been changed regarding
   its keyword arguments.
   [(#2081)](https://github.com/PennyLaneAI/pennylane/pull/2081)
@@ -301,9 +342,13 @@
   For more details, see the
   [RotosolveOptimizer documentation](https://pennylane.readthedocs.io/en/stable/code/api/pennylane.RotosolveOptimizer.html).
 
-
 <h3>Bug fixes</h3>
 
+* Fixes a bug in `gradient_transform` where the hybrid differentiation
+  of circuits with a single parametrized gate failed and QNode argument
+  axes of size `1` where removed from the output gradient.
+  [(#2080)](https://github.com/PennyLaneAI/pennylane/pull/2080)
+
 * The available `diff_method` options for QNodes has been corrected in both the
   error messages and the documentation.
   [(#2078)](https://github.com/PennyLaneAI/pennylane/pull/2078)
diff --git a/pennylane/_grad.py b/pennylane/_grad.py
index 9ed50655c..f2498281b 100644
--- a/pennylane/_grad.py
+++ b/pennylane/_grad.py
@@ -24,8 +24,6 @@ from autograd.numpy.numpy_boxes import ArrayBox
 from autograd.extend import vspace
 from autograd.wrap_util import unary_to_nary
 
-from pennylane import numpy as np
-
 make_vjp = unary_to_nary(_make_vjp)
 
 
@@ -158,43 +156,147 @@ def jacobian(func, argnum=None):
             must consist of a single NumPy array (if classical) or a tuple of
             expectation values (if a quantum node)
         argnum (int or Sequence[int]): Which argument to take the gradient
-            with respect to. If a sequence is given, the Jacobian matrix
-            corresponding to all input elements and all output elements is returned.
+            with respect to. If a sequence is given, the Jacobian corresponding
+            to all marked inputs and all output elements is returned.
 
     Returns:
         function: the function that returns the Jacobian of the input
         function with respect to the arguments in argnum
-    """
-    # pylint: disable=no-value-for-parameter
 
-    if argnum is not None:
-        # for backwards compatibility with existing code
-        # that manually specifies argnum
-        if isinstance(argnum, int):
-            return _jacobian(func, argnum)
+    For ``argnum=None``, the trainable arguments are inferred dynamically from the arguments
+    passed to the function. The returned function takes the same arguments as the original
+    function and outputs a ``tuple``. The ``i`` th entry of the ``tuple`` has shape
+    ``(*output shape, *shape of args[argnum[i]])``.
 
-        return lambda *args, **kwargs: np.stack(
-            [_jacobian(func, arg)(*args, **kwargs) for arg in argnum]
-        ).T
+    If a single trainable argument is inferred, or if a single integer
+    is provided as ``argnum``, the tuple is unpacked and its only entry is returned instead.
 
-    def _jacobian_function(*args, **kwargs):
-        """Inspect the arguments for differentiability, and
-        compute the autograd gradient function with required argnums
-        dynamically.
+    **Example**
 
-        This wrapper function is returned to the user instead of autograd.jacobian,
-        so that we can take into account cases where the user computes the
-        jacobian function once, but then calls it with arguments that change
-        in differentiability.
-        """
+    Consider the QNode
+
+    .. code-block::
+
+        dev = qml.device("default.qubit", wires=2)
+
+        @qml.qnode(dev)
+        def circuit(weights):
+            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)
+
+        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)``:
+
+    >>> qml.jacobian(circuit)(weights).shape
+    (2, 2, 1, 3)
+
+    On the other hand, consider the following QNode for the same circuit
+    structure:
+
+    .. code-block::
+
+        dev = qml.device("default.qubit", wires=2)
+
+        @qml.qnode(dev)
+        def circuit(x, y, z):
+            qml.RX(x, wires=0)
+            qml.RY(y, wires=1)
+            qml.RZ(z, wires=0)
+            return tuple(qml.expval(qml.PauliZ(w)) for w in dev.wires)
+
+        x = np.array(0.2, requires_grad=True)
+        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,)``:
+
+    >>> jac = qml.jacobian(circuit)(x, y, z)
+    >>> type(jac)
+    tuple
+    >>> for sub_jac in jac:
+    ...     print(sub_jac.shape)
+    (2,)
+    (2,)
+    (2,)
+
+    For a more advanced setting of QNode arguments, consider the QNode
+
+    .. code-block::
+
+        dev = qml.device("default.qubit", wires=3)
+
+        @qml.qnode(dev)
+        def circuit(x, y):
+            qml.RX(x[0], wires=0)
+            qml.RY(y[0, 3], wires=1)
+            qml.RX(x[1], wires=2)
+            return [qml.expval(qml.PauliZ(w)) for w in [0, 1, 2]]
+
+        x = np.array([0.1, 0.5], requires_grad=True)
+        y = np.array([[-0.3, 1.2, 0.1, 0.9], [-0.2, -3.1, 0.5, -0.7]], requires_grad=True)
+
+    If we do not provide ``argnum``, ``qml.jacobian`` will correctly identify both,
+    ``x`` and ``y``, as trainable function arguments:
+
+    >>> jac = qml.jacobian(circuit)(x, y)
+    >>> print(type(jac), len(jac))
+    <class 'tuple'> 2
+    >>> qml.math.shape(jac[0])
+    (3, 2)
+    >>> qml.math.shape(jac[1])
+    (3, 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,)``).
+    Similarily, the shape ``(2, 4)`` of ``y`` leads to a Jacobian shape ``(3, 2, 4)``.
+
+    Instead we may choose the output to contain only one of the two
+    entries by providing an iterable as ``argnum``:
+
+    >>> jac = qml.jacobian(circuit, argnum=[1])(x, y)
+    >>> print(type(jac), len(jac))
+    <class 'tuple'> 1
+    >>> qml.math.shape(jac)
+    (1, 3, 2, 4)
+
+    Here we included the size of the tuple in the shape analysis, corresponding to the
+    first dimension of size ``1``.
+
+    Finally, we may want to receive the single entry above directly, not as a tuple
+    with a single entry. This is done by providing a single integer as ``argnum``
+
+    >>> jac = qml.jacobian(circuit, argnum=1)(x, y)
+    >>> print(type(jac), len(jac))
+    <class 'numpy.ndarray'> 3
+    >>> qml.math.shape(jac)
+    (3, 2, 4)
+
+    As expected, the tuple was unpacked and we directly received the Jacobian of the
+    QNode with respect to ``y``.
+    """
+    # pylint: disable=no-value-for-parameter
+
+    def _get_argnum(args):
+        """Inspect the arguments for differentiability and return the
+        corresponding indices."""
         argnum = []
 
         for idx, arg in enumerate(args):
 
             trainable = getattr(arg, "requires_grad", None)
-            array_box = isinstance(arg, ArrayBox)
+            is_array_box = isinstance(arg, ArrayBox)
 
-            if trainable is None and not array_box:
+            if trainable is None and not is_array_box:
 
                 warnings.warn(
                     "Starting with PennyLane v0.21.0, when using Autograd, inputs "
@@ -210,20 +312,29 @@ def jacobian(func, argnum=None):
             if trainable:
                 argnum.append(idx)
 
-        if not argnum:
-            return tuple()
-
-        if len(argnum) == 1:
-            return _jacobian(func, argnum[0])(*args, **kwargs)
+        return argnum
 
-        jacobians = [_jacobian(func, arg)(*args, **kwargs) for arg in argnum]
+    def _jacobian_function(*args, **kwargs):
+        """Compute the autograd Jacobian.
 
-        try:
-            return np.stack(jacobians).T
-        except ValueError:
-            # The Jacobian of each argument is a different shape and cannot
-            # be stacked; simply return the tuple of argument Jacobians.
-            return tuple(jacobians)
+        This wrapper function is returned to the user instead of autograd.jacobian,
+        so that we can take into account cases where the user computes the
+        jacobian function once, but then calls it with arguments that change
+        in differentiability.
+        """
+        if argnum is None:
+            # Infer which arguments to consider trainable
+            _argnum = _get_argnum(args)
+            # Infer whether to unpack from the infered argnum
+            unpack = len(_argnum) == 1
+        else:
+            # For a single integer as argnum, unpack the Jacobian tuple
+            unpack = isinstance(argnum, int)
+            _argnum = [argnum] if unpack else argnum
+
+        jac = tuple(_jacobian(func, arg)(*args, **kwargs) for arg in _argnum)
+
+        return jac[0] if unpack else jac
 
     return _jacobian_function
 
diff --git a/pennylane/gradients/gradient_transform.py b/pennylane/gradients/gradient_transform.py
index bd9c5ab63..66c286d77 100644
--- a/pennylane/gradients/gradient_transform.py
+++ b/pennylane/gradients/gradient_transform.py
@@ -14,8 +14,6 @@
 """This module contains utilities for defining custom gradient transforms,
 including a decorator for specifying gradient expansions."""
 # pylint: disable=too-few-public-methods
-import warnings
-
 import pennylane as qml
 from pennylane.transforms.tape_expand import expand_invalid_trainable
 
@@ -110,27 +108,6 @@ class gradient_transform(qml.batch_transform):
         self.hybrid = hybrid
         super().__init__(transform_fn, expand_fn=expand_fn, differentiable=differentiable)
 
-    @staticmethod
-    def _jacobian_trainable_args(args, interface):
-        """Return the indices of QNode arguments for which a Jacobian was
-        computed by `qml.transforms.classical_jacobian` with argnum=None.
-        """
-        trainable_args = []
-
-        if interface == "autograd":
-            for idx, arg in enumerate(args):
-                # TODO: make default False once this change is done in qml.jacobian
-                if getattr(arg, "requires_grad", True):
-                    trainable_args.append(idx)
-
-        elif interface == "jax":
-            trainable_args = [0]
-
-        # Torch and Tensorflow interfaces are not considered since `classical_jacobian`
-        # always returns a tuple for them, thus not invoking this function.
-
-        return trainable_args
-
     def default_qnode_wrapper(self, qnode, targs, tkwargs):
         # Here, we overwrite the QNode execution wrapper in order
         # to take into account that classical processing may be present
@@ -143,21 +120,22 @@ class gradient_transform(qml.batch_transform):
             qjac = _wrapper(*args, **kwargs)
 
             if any(m.return_type is qml.operation.Probability for m in qnode.qtape.measurements):
-                qjac = qml.math.squeeze(qjac)
+                qjac = qml.math.safe_squeeze(qjac, exclude_axis=-1)
 
             if not hybrid:
                 return qjac
 
+            qjac_axes = list(range(len(qml.math.shape(qjac)) - 1))
             kwargs.pop("shots", False)
             cjac = cjac_fn(*args, **kwargs)
 
             if isinstance(cjac, tuple):
                 # Classical processing of multiple arguments is present. Return qjac @ cjac.
-                jacs = [
-                    qml.math.squeeze(qml.math.tensordot(c, qjac, [[0], [-1]]))
+                jacs = tuple(
+                    qml.math.safe_squeeze(qml.math.tensordot(qjac, c, [[-1], [0]]), axis=qjac_axes)
                     for c in cjac
                     if c is not None
-                ]
+                )
                 return jacs
 
             is_square = cjac.shape == (1,) or (cjac.ndim == 2 and cjac.shape[0] == cjac.shape[1])
@@ -167,42 +145,8 @@ class gradient_transform(qml.batch_transform):
                 # is present inside the QNode.
                 return qjac
 
-            # Classical processing present of either:
-            #   a) a single argument or
-            #   b) multiple arguments of the same shape
-            # The shape of the classical jacobian returned by qml.jacobian (invoked by
-            # classical_jacobian with autograd) depends on the scenario:
-            #   a) (# gate args, qnode arg shape)
-            #   b) (reverse qnode arg shape, # gate args, # qnode args)
-            # It then needs to be contracted with the quantum jacobian of shape:
-            #      (qnode output shape, # gate args)
-            # The result should have shape:
-            #   a) (qnode ouput shape, qnode arg shape)
-            #   b) (reverse qnode arg shape, qnode output shape, # qnode args)
-            num_gate_args = qml.math.shape(qjac)[-1]
-            # Consider only QNode arguments regarded as trainable by the interfaces.
-            trainable_args_idx = self._jacobian_trainable_args(args, qnode.interface)
-            # Since all arguments have the same shape, obtain shape from the first trainable arg
-            qnode_arg_shape = qml.math.shape(args[trainable_args_idx[0]])
-            num_qnode_args = len(trainable_args_idx)
-
-            if qml.math.shape(cjac) == (num_gate_args, *qnode_arg_shape):
-                # single QNode argument
-                jac = qml.math.tensordot(qjac, cjac, [[-1], [0]])
-            elif qml.math.shape(cjac) == (*qnode_arg_shape[::-1], num_gate_args, num_qnode_args):
-                # multiple QNode arguments with stacking
-                cjac = qml.math.swapaxes(cjac, -1, -2)
-                jac = qml.math.tensordot(cjac, qjac, [[-1], [-1]])
-                jac = qml.math.moveaxis(jac, -1, len(qnode_arg_shape))
-            else:  # pragma: no cover
-                jac = ()
-                warnings.warn(
-                    "Unexpected classical Jacobian encoutered, could not compute the hybrid "
-                    "Jacobian of the QNode. You can still attempt to obtain the quantum Jacobian "
-                    "with the `hybrid=False` parameter.",
-                    UserWarning,
-                )
-
-            return qml.math.squeeze(jac)
+            return qml.math.safe_squeeze(
+                qml.math.tensordot(qjac, cjac, [[-1], [0]]), axis=qjac_axes
+            )
 
         return jacobian_wrapper
diff --git a/pennylane/math/__init__.py b/pennylane/math/__init__.py
index 883f09ba2..6f73182c8 100644
--- a/pennylane/math/__init__.py
+++ b/pennylane/math/__init__.py
@@ -42,6 +42,7 @@ from .multi_dispatch import (
     frobenius_inner_product,
     get_trainable_indices,
     ones_like,
+    safe_squeeze,
     scatter_element_add,
     stack,
     tensordot,
@@ -93,6 +94,7 @@ __all__ = [
     "marginal_prob",
     "ones_like",
     "requires_grad",
+    "safe_squeeze",
     "scatter_element_add",
     "stack",
     "tensordot",
diff --git a/pennylane/math/multi_dispatch.py b/pennylane/math/multi_dispatch.py
index e75a419ec..3de21426f 100644
--- a/pennylane/math/multi_dispatch.py
+++ b/pennylane/math/multi_dispatch.py
@@ -361,6 +361,55 @@ def ones_like(tensor, dtype=None):
     return np.ones_like(tensor)
 
 
+def safe_squeeze(tensor, axis=None, exclude_axis=None):
+    """Squeeze a tensor either along all axes, specified axes or all
+    but a set of excluded axes. For selective squeezing, catch errors
+    and do nothing if the selected axes do not have size 1.
+
+    Args:
+        tensor (tensor_like): input tensor
+        axis (int or Sequence[int]): Axis/axes to squeeze
+        exclude_axis (int or Sequence[int]): Axis/axes not to squeeze
+
+    Return:
+        tensor_like: The input tensor with those axes removed that were specified
+        or not excluded and that have size 1. If no axes are specified or excluded,
+        all axes are attempted to be squeezed.
+    """
+    interface = _multi_dispatch([tensor])
+    if interface == "tensorflow":
+        from tensorflow.python.framework.errors_impl import InvalidArgumentError
+
+        exception = InvalidArgumentError
+    else:
+        exception = ValueError
+
+    num_axes = len(np.shape(tensor))
+    if axis is None:
+        if exclude_axis is None:
+            return np.squeeze(tensor)
+        if np.isscalar(exclude_axis):
+            exclude_axis = [exclude_axis if exclude_axis >= 0 else num_axes + exclude_axis]
+        else:
+            exclude_axis = [(i if i >= 0 else num_axes + i) for i in exclude_axis]
+        axis = [i for i in range(num_axes) if not i in exclude_axis]
+    elif np.isscalar(axis):
+        axis = [axis if axis >= 0 else num_axes + axis]
+    else:
+        axis = [(i if i >= 0 else num_axes + i) for i in axis]
+
+    for ax in sorted(set(axis)):
+        # Modify the axis index to squeeze by the number of squeezes already
+        # performed - this works because we squeeze in increasing axis index order.
+        # If the axis index is negative, no modification is needed.
+        ax -= num_axes - len(np.shape(tensor))
+        try:
+            tensor = np.squeeze(tensor, axis=ax)
+        except exception:
+            pass
+    return tensor
+
+
 def stack(values, axis=0):
     """Stack a sequence of tensors along the specified axis.
 
diff --git a/pennylane/optimize/qng.py b/pennylane/optimize/qng.py
index cba9c047b..de7624919 100644
--- a/pennylane/optimize/qng.py
+++ b/pennylane/optimize/qng.py
@@ -220,10 +220,25 @@ class QNGOptimizer(GradientDescentOptimizer):
         if forward is None:
             forward = qnode(*args, **kwargs)
 
-        # unwrap from list if one argument, cleaner return
-        if len(new_args) == 1:
-            return new_args[0], forward
-        return new_args, forward
+        # Note: for now, we only have single element lists as the new
+        # arguments, but this might change, see TODO below.
+        # Once the other approach is implemented, we need to unwrap from list
+        # if one argument for a cleaner return.
+        # if len(new_args) == 1:
+        return new_args[0], forward
+
+        # TODO: The scenario of the following return statement is not implemented
+        # yet, as currently only a single metric tensor can be processed.
+        # An optimizer refactor is needed to accomodate for this (similar to other
+        # optimizers for which `apply_grad` will have to be patched to allow for
+        # tuple-valued gradients to be processed)
+        #
+        # For multiple QNode arguments, `qml.jacobian` and `qml.metric_tensor`
+        # return a tuple of arrays. Each of the gradient arrays has to be processed
+        # together with the corresponding array in the metric tensor tuple.
+        # This requires modifications of the `GradientDescentOptimizer` base class
+        # as none of the optimizers accomodate for this use case.
+        # return new_args, forward
 
     # pylint: disable=arguments-differ
     def step(
diff --git a/pennylane/optimize/rotosolve.py b/pennylane/optimize/rotosolve.py
index 2f0d58c16..44f74f74b 100644
--- a/pennylane/optimize/rotosolve.py
+++ b/pennylane/optimize/rotosolve.py
@@ -46,6 +46,26 @@ def _shgo_optimizer(fun, **kwargs):
     return opt_res.x[0], opt_res.fun
 
 
+def _validate_inputs(requires_grad, args, nums_frequency, spectra):
+    """Checks that for each trainable argument either the number of
+    frequencies or the frequency spectrum is given."""
+
+    for arg, (arg_name, _requires_grad) in zip(args, requires_grad.items()):
+        if _requires_grad:
+            _nums_frequency = nums_frequency.get(arg_name, {})
+            _spectra = spectra.get(arg_name, {})
+            all_keys = set(_nums_frequency) | set(_spectra)
+
+            shape = qml.math.shape(arg)
+            indices = np.ndindex(shape) if len(shape) > 0 else [()]
+            for par_idx in indices:
+                if par_idx not in all_keys:
+                    raise ValueError(
+                        "Neither the number of frequencies nor the frequency spectrum "
+                        f"was provided for the entry {par_idx} of argument {arg_name}."
+                    )
+
+
 def _restrict_to_univariate(fn, arg_idx, par_idx, args, kwargs):
     r"""Restrict a function to a univariate function for given argument
     and parameter indices.
@@ -300,26 +320,6 @@ class RotosolveOptimizer:
         else:
             self.substep_optimizer = substep_optimizer
 
-    @staticmethod
-    def _validate_inputs(requires_grad, args, nums_frequency, spectra):
-        """Checks that for each trainable argument either the number of
-        frequencies or the frequency spectrum is given."""
-
-        for arg, (arg_name, _requires_grad) in zip(args, requires_grad.items()):
-            if _requires_grad:
-                _nums_frequency = nums_frequency.get(arg_name, {})
-                _spectra = spectra.get(arg_name, {})
-                all_keys = set(_nums_frequency) | set(_spectra)
-
-                shape = qml.math.shape(arg)
-                indices = np.ndindex(shape) if len(shape) > 0 else [()]
-                for par_idx in indices:
-                    if par_idx not in all_keys:
-                        raise ValueError(
-                            "Neither the number of frequencies nor the frequency spectrum "
-                            f"was provided for the entry {par_idx} of argument {arg_name}."
-                        )
-
     def step_and_cost(
         self,
         objective_fn,
@@ -408,7 +408,7 @@ class RotosolveOptimizer:
         }
         nums_frequency = nums_frequency or {}
         spectra = spectra or {}
-        self._validate_inputs(requires_grad, args, nums_frequency, spectra)
+        _validate_inputs(requires_grad, args, nums_frequency, spectra)
 
         # we will single out one arg to change at a time
         # the following hold the arguments not getting updated
diff --git a/pennylane/transforms/adjoint_metric_tensor.py b/pennylane/transforms/adjoint_metric_tensor.py
index 8dfbc22ee..5a25118dd 100644
--- a/pennylane/transforms/adjoint_metric_tensor.py
+++ b/pennylane/transforms/adjoint_metric_tensor.py
@@ -303,6 +303,6 @@ def _adjoint_metric_tensor_qnode(qnode, device, hybrid):
 
         cjac = cjac_fn(*args, **kwargs)
 
-        return _contract_metric_tensor_with_cjac(mt, cjac, args, qnode.interface)
+        return _contract_metric_tensor_with_cjac(mt, cjac)
 
     return wrapper
diff --git a/pennylane/transforms/classical_jacobian.py b/pennylane/transforms/classical_jacobian.py
index aa8b95a57..668202c0b 100644
--- a/pennylane/transforms/classical_jacobian.py
+++ b/pennylane/transforms/classical_jacobian.py
@@ -88,15 +88,15 @@ def classical_jacobian(qnode, argnum=None, expand_fn=None, trainable_only=True):
     The output and its format depend on the backend:
 
     .. list-table:: Output format of ``classical_jacobian``
-       :widths: 25 25 25 25
+       :widths: 15 25 25 35
        :header-rows: 1
 
        * - Interface
          - ``argnum=None``
          - ``type(argnum)=int``
-         - ``argnum=Sequence[int]``
+         - ``type(argnum) = Sequence[int]``
        * - ``'autograd'``
-         - ``tuple(arrays)`` [1]
+         - ``tuple(array)`` [1]
          - ``array``
          - ``tuple(array)``
        * - ``'jax'``
@@ -104,17 +104,16 @@ def classical_jacobian(qnode, argnum=None, expand_fn=None, trainable_only=True):
          - ``array``
          - ``tuple(array)``
        * - ``'tf'``
-         - ``tuple(arrays)``
+         - ``tuple(array)``
          - ``array``
          - ``tuple(array)``
        * - ``'torch'``
-         - ``tuple(arrays)``
+         - ``tuple(array)``
          - ``array``
          - ``tuple(array)``
 
-    [1] If all QNode arguments are of the same shape, the tuple is unpacked and the
-    Jacobian arrays are stacked into one ``array``. If there only is one QNode argument,
-    the tuple is unpacked as well. Both is based on the behaviour of ``qml.jacobian``.
+    [1] If there only is one trainable QNode argument, the tuple is unpacked to a
+    single ``array``, as is the case for :func:`.jacobian`.
 
     [2] For JAX, ``argnum=None`` defaults to ``argnum=0`` in contrast to all other
     interfaces. This means that only the classical Jacobian with respect to the first
@@ -136,7 +135,6 @@ def classical_jacobian(qnode, argnum=None, expand_fn=None, trainable_only=True):
 
     Only the Jacobians with respect to the arguments ``x`` and ``y`` were computed, and
     returned as a tuple of ``arrays``.
-
     """
 
     def classical_preprocessing(*args, **kwargs):
@@ -152,21 +150,7 @@ def classical_jacobian(qnode, argnum=None, expand_fn=None, trainable_only=True):
 
     if qnode.interface == "autograd":
 
-        def _jacobian(*args, **kwargs):
-            if argnum is None:
-                jac = qml.jacobian(classical_preprocessing)(*args, **kwargs)
-            elif np.isscalar(argnum):
-                jac = qml.jacobian(classical_preprocessing, argnum=argnum)(*args, **kwargs)
-            else:
-                jac = tuple(
-                    (
-                        qml.jacobian(classical_preprocessing, argnum=i)(*args, **kwargs)
-                        for i in argnum
-                    )
-                )
-            return jac
-
-        return _jacobian
+        return qml.jacobian(classical_preprocessing, argnum=argnum)
 
     if qnode.interface == "torch":
         import torch
diff --git a/pennylane/transforms/metric_tensor.py b/pennylane/transforms/metric_tensor.py
index 84092ea8e..b9b522cbd 100644
--- a/pennylane/transforms/metric_tensor.py
+++ b/pennylane/transforms/metric_tensor.py
@@ -259,7 +259,7 @@ def metric_tensor(tape, approx=None, allow_nonunitary=True, aux_wire=None, devic
     )
 
 
-def _contract_metric_tensor_with_cjac(mt, cjac, args, interface):
+def _contract_metric_tensor_with_cjac(mt, cjac):
     """Execute the contraction of pre-computed classical Jacobian(s)
     and the metric tensor of a tape in order to obtain the hybrid
     metric tensor of a QNode.
@@ -267,14 +267,12 @@ def _contract_metric_tensor_with_cjac(mt, cjac, args, interface):
     Args:
         mt (array): Metric tensor of a tape (2-dimensional)
         cjac (array or tuple[array]): The classical Jacobian of a QNode
-        args (tuple): QNode arguments
-        interface (str): QNode interface
 
     Returns:
         array or tuple[array]: Hybrid metric tensor(s) of the QNode.
         The number of metric tensors depends on the number of QNode arguments
-        for which the classical Jacobian was computed, their shape on the
-        shape of these QNode arguments.
+        for which the classical Jacobian was computed, the tensor shape(s)
+        depend on the shape of these QNode arguments.
     """
     if isinstance(cjac, tuple):
         # Classical processing of multiple arguments is present. Return cjac.T @ mt @ cjac
@@ -296,36 +294,7 @@ def _contract_metric_tensor_with_cjac(mt, cjac, args, interface):
         # is present inside the QNode.
         return mt
 
-    # TODO: Remove the following behaviour once the stacking behaviour in `qml.jacobian`
-    # has been removed. The additional arguments `args` and `interface` can be removed
-    # accordingly.
-
-    # Get number of gate arguments that were considered trainable
-    num_gate_args = qml.math.shape(mt)[0]
-    # Get trainable args
-    # pylint: disable=protected-access
-    trainable_args_idx = qml.gradients.gradient_transform._jacobian_trainable_args(args, interface)
-    # Since all arguments have the same shape, obtain shape from the first trainable arg
-    qnode_arg_shape = qml.math.shape(args[trainable_args_idx[0]])
-    num_qnode_args = len(trainable_args_idx)
-
-    if qml.math.shape(cjac) == (num_gate_args, *qnode_arg_shape):
-        # single QNode argument
-        cjac_axis = [0]
-    elif qml.math.shape(cjac) == (*qnode_arg_shape[::-1], num_gate_args, num_qnode_args):
-        # multiple QNode arguments with stacking
-        cjac_axis = [-2]
-    else:  # pragma: no cover
-        warnings.warn(
-            "Unexpected classical Jacobian encoutered, could not compute the hybrid "
-            "metric tensor of the QNode. You can still attempt to obtain the quantum "
-            "metric tensor with the `hybrid=False` parameter.",
-            UserWarning,
-        )
-        return ()
-
-    mt = qml.math.tensordot(mt, cjac, [[-1], cjac_axis])
-    mt = qml.math.tensordot(cjac, mt, [cjac_axis, [0]])
+    mt = qml.math.tensordot(cjac, qml.math.tensordot(mt, cjac, axes=[[-1], [0]]), axes=[[0], [0]])
 
     return mt
 
@@ -382,7 +351,7 @@ def qnode_execution_wrapper(self, qnode, targs, tkwargs):
         kwargs.pop("shots", False)
         cjac = cjac_fn(*args, **kwargs)
 
-        return _contract_metric_tensor_with_cjac(mt, cjac, args, qnode.interface)
+        return _contract_metric_tensor_with_cjac(mt, cjac)
 
     return wrapper
 


Test patch:
diff --git a/tests/fourier/test_reconstruct.py b/tests/fourier/test_reconstruct.py
index ac2c4843c..1c9568f7f 100644
--- a/tests/fourier/test_reconstruct.py
+++ b/tests/fourier/test_reconstruct.py
@@ -66,7 +66,6 @@ def fun_close(fun1, fun2, zero=None, tol=1e-5, samples=100):
 
     for x in X:
         if not np.isclose(fun1(x), fun2(x), atol=tol, rtol=0):
-            print(f"{fun1(x)} != {fun2(x)} at {x}")
             return False
     return True
 
diff --git a/tests/gradients/test_gradient_transform.py b/tests/gradients/test_gradient_transform.py
index 5b1d0c571..967b4dea8 100644
--- a/tests/gradients/test_gradient_transform.py
+++ b/tests/gradients/test_gradient_transform.py
@@ -140,6 +140,28 @@ class TestGradientTransformIntegration:
         correctly when multiple tensor QNode arguments are present"""
         dev = qml.device("default.qubit", wires=2)
 
+        @qml.qnode(dev)
+        def circuit(x, y):
+            qml.RX(x[0, 0], wires=0)
+            qml.RY(y[0, 0], wires=0)
+            qml.RZ(x[1, 0], wires=0)
+            return qml.probs(wires=[0, 1])
+
+        x = np.array([[0.1, 0.3], [0.2, -0.1]], requires_grad=True)
+        y = np.array([[0.2, 0.2], [0.3, 0.5]], requires_grad=True)
+
+        expected = qml.jacobian(circuit)(x, y)
+        res = qml.gradients.param_shift(circuit, hybrid=True)(x, y)
+        assert isinstance(res, tuple) and len(res) == 2
+        assert all(np.allclose(_r, _e, atol=tol, rtol=0) for _r, _e in zip(res, expected))
+
+    # TODO: Include the following test once the gradient_transform is fixed regarding the
+    # usage of qml.math.squeeze.
+    def test_multiple_tensor_arguments_old_version(self, tol):
+        """Test that a gradient transform acts on QNodes
+        correctly when multiple tensor QNode arguments are present"""
+        dev = qml.device("default.qubit", wires=2)
+
         @qml.qnode(dev)
         def circuit(x, y):
             qml.RX(x[0, 0], wires=0)
@@ -152,7 +174,61 @@ class TestGradientTransformIntegration:
 
         expected = qml.jacobian(circuit)(x, y)
         res = qml.gradients.param_shift(circuit)(x, y)
+        assert isinstance(res, tuple) and len(res) == 2
+        assert all(np.allclose(_r, _e, atol=tol, rtol=0) for _r, _e in zip(res, expected))
+
+    def test_high_dimensional_single_parameter_arg(self, tol):
+        """Test that a gradient transform acts on QNodes correctly
+        when a single high-dimensional tensor QNode arguments is used"""
+        dev = qml.device("default.qubit", wires=2)
+
+        @qml.qnode(dev)
+        def circuit(x):
+            qml.RX(x[0, 0] / 2, wires=0)
+            qml.RX(x[0, 0] / 2, wires=0)
+            return qml.probs(wires=[0, 1])
+
+        x = np.array([[0.1, 0.1], [0.1, 0.1]], requires_grad=True)
+
+        expected = qml.jacobian(circuit)(x)
+        res = qml.gradients.param_shift(circuit)(x)
+        assert np.allclose(res, expected, atol=tol, rtol=0)
+
+    # TODO: Include the following test once the gradient_transform is fixed regarding the
+    # usage of qml.math.squeeze.
+    def test_high_dimensional_single_parameter_arg_and_single_gate(self, tol):
+        """Test that a gradient transform acts on QNodes correctly
+        when a single high-dimensional tensor QNode arguments is used"""
+        dev = qml.device("default.qubit", wires=2)
+
+        @qml.qnode(dev)
+        def circuit(x):
+            qml.RX(x[0, 0], wires=0)
+            return qml.probs(wires=[0, 1])
+
+        x = np.array([[0.1]], requires_grad=True)
+
+        expected = qml.jacobian(circuit)(x)
+        res = qml.gradients.param_shift(circuit)(x)
+        assert np.allclose(res, expected, atol=tol, rtol=0)
+
+    # TODO: Include the following test once the gradient_transform is fixed regarding the
+    # usage of qml.math.squeeze.
+    def test_single_gate_arg(self, tol):
+        """Test that a gradient transform acts on QNodes correctly
+        when a single QNode argument and gate are present"""
+        dev = qml.device("default.qubit", wires=2)
+
+        @qml.qnode(dev)
+        def circuit(x):
+            qml.RX(x[0], wires=0)
+            return qml.probs(wires=[0, 1])
+
+        x = np.array([0.1, 0.2], requires_grad=True)
 
+        expected = qml.jacobian(circuit)(x)
+        cjac = qml.transforms.classical_jacobian(circuit)(x)
+        res = qml.gradients.param_shift(circuit)(x)
         assert np.allclose(res, expected, atol=tol, rtol=0)
 
     def test_first_non_trainable_argument(self, tol):
diff --git a/tests/interfaces/test_batch_autograd.py b/tests/interfaces/test_batch_autograd.py
index 0b0f72ad1..a8ad8083c 100644
--- a/tests/interfaces/test_batch_autograd.py
+++ b/tests/interfaces/test_batch_autograd.py
@@ -434,10 +434,12 @@ class TestAutogradExecuteIntegration:
         assert np.allclose(res, expected, atol=tol, rtol=0)
 
         res = qml.jacobian(cost)(a, b, device=dev)
-        assert res.shape == (2, 2)
+        assert isinstance(res, tuple) and len(res) == 2
+        assert res[0].shape == (2,)
+        assert res[1].shape == (2,)
 
-        expected = [[-np.sin(a), 0], [np.sin(a) * np.sin(b), -np.cos(a) * np.cos(b)]]
-        assert np.allclose(res, expected, atol=tol, rtol=0)
+        expected = ([-np.sin(a), np.sin(a) * np.sin(b)], [0, -np.cos(a) * np.cos(b)])
+        assert all(np.allclose(_r, _e, atol=tol, rtol=0) for _r, _e in zip(res, expected))
 
     def test_tape_no_parameters(self, execute_kwargs, tol):
         """Test that a tape with no parameters is correctly
@@ -505,11 +507,12 @@ class TestAutogradExecuteIntegration:
 
         jac_fn = qml.jacobian(lambda a, b: cost(2 * a, b))
         jac = jac_fn(a, b)
-        expected = [
-            [-2 * np.sin(2 * a), 0],
-            [2 * np.sin(2 * a) * np.sin(b), -np.cos(2 * a) * np.cos(b)],
-        ]
-        assert np.allclose(jac, expected, atol=tol, rtol=0)
+        expected = (
+            [-2 * np.sin(2 * a), 2 * np.sin(2 * a) * np.sin(b)],
+            [0, -np.cos(2 * a) * np.cos(b)],
+        )
+        assert isinstance(jac, tuple) and len(jac) == 2
+        assert all(np.allclose(_j, _e, atol=tol, rtol=0) for _j, _e in zip(jac, expected))
 
     def test_classical_processing(self, execute_kwargs, tol):
         """Test classical processing within the quantum tape"""
@@ -528,7 +531,10 @@ class TestAutogradExecuteIntegration:
 
         dev = qml.device("default.qubit", wires=2)
         res = qml.jacobian(cost)(a, b, c, device=dev)
-        assert res.shape == (1, 2)
+        # Only two arguments are trainable
+        assert isinstance(res, tuple) and len(res) == 2
+        assert res[0].shape == (1,)
+        assert res[1].shape == (1,)
 
     def test_no_trainable_parameters(self, execute_kwargs, tol):
         """Test evaluation and Jacobian if there are no trainable parameters"""
@@ -662,19 +668,27 @@ class TestAutogradExecuteIntegration:
 
         jac_fn = qml.jacobian(cost)
         res = jac_fn(x, y, device=dev)
-        assert res.shape == (2, 2, 2)
+        assert isinstance(res, tuple) and len(res) == 2
+        assert res[0].shape == (2, 2)
+        assert res[1].shape == (2, 2)
 
-        expected = np.array(
-            [
-                [[-np.sin(x) / 2, 0], [-np.sin(x) * np.cos(y) / 2, -np.cos(x) * np.sin(y) / 2]],
+        expected = (
+            np.array(
                 [
-                    [np.sin(x) / 2, 0],
-                    [np.cos(y) * np.sin(x) / 2, np.cos(x) * np.sin(y) / 2],
-                ],
-            ]
+                    [-np.sin(x) / 2, np.sin(x) / 2],
+                    [-np.sin(x) * np.cos(y) / 2, np.sin(x) * np.cos(y) / 2],
+                ]
+            ),
+            np.array(
+                [
+                    [0, 0],
+                    [-np.cos(x) * np.sin(y) / 2, np.cos(x) * np.sin(y) / 2],
+                ]
+            ),
         )
 
-        assert np.allclose(res, expected, atol=tol, rtol=0)
+        assert np.allclose(res[0], expected[0], atol=tol, rtol=0)
+        assert np.allclose(res[1], expected[1], atol=tol, rtol=0)
 
     def test_ragged_differentiation(self, execute_kwargs, tol):
         """Tests correct output shape and evaluation for a tape
@@ -704,14 +718,16 @@ class TestAutogradExecuteIntegration:
 
         jac_fn = qml.jacobian(cost)
         res = jac_fn(x, y, device=dev)
-        expected = np.array(
-            [
-                [-np.sin(x), 0],
-                [-np.sin(x) * np.cos(y) / 2, -np.cos(x) * np.sin(y) / 2],
-                [np.cos(y) * np.sin(x) / 2, np.cos(x) * np.sin(y) / 2],
-            ]
+        assert isinstance(res, tuple) and len(res) == 2
+        assert res[0].shape == (3,)
+        assert res[1].shape == (3,)
+
+        expected = (
+            np.array([-np.sin(x), -np.sin(x) * np.cos(y) / 2, np.sin(x) * np.cos(y) / 2]),
+            np.array([0, -np.cos(x) * np.sin(y) / 2, np.cos(x) * np.sin(y) / 2]),
         )
-        assert np.allclose(res, expected, atol=tol, rtol=0)
+        assert np.allclose(res[0], expected[0], atol=tol, rtol=0)
+        assert np.allclose(res[1], expected[1], atol=tol, rtol=0)
 
     def test_sampling(self, execute_kwargs):
         """Test sampling works as expected"""
@@ -963,10 +979,14 @@ class TestOverridingShots:
 
         res = qml.jacobian(cost_fn)(a, b, shots=[10000, 10000, 10000])
         assert dev.shots is None
-        assert len(res) == 3
+        assert isinstance(res, tuple) and len(res) == 2
+        assert res[0].shape == (3,)
+        assert res[1].shape == (3,)
 
         expected = [np.sin(a) * np.sin(b), -np.cos(a) * np.cos(b)]
-        assert np.allclose(np.mean(res, axis=0), expected, atol=0.1, rtol=0)
+        assert all(
+            np.allclose(np.mean(r, axis=0), e, atol=0.1, rtol=0) for r, e in zip(res, expected)
+        )
 
 
 execute_kwargs = [
diff --git a/tests/interfaces/test_batch_autograd_qnode.py b/tests/interfaces/test_batch_autograd_qnode.py
index 17e8618e1..eddbe910b 100644
--- a/tests/interfaces/test_batch_autograd_qnode.py
+++ b/tests/interfaces/test_batch_autograd_qnode.py
@@ -164,8 +164,10 @@ class TestQNode:
         assert np.allclose(res, expected, atol=tol, rtol=0)
 
         res = qml.jacobian(circuit)(a, b)
-        expected = [[-np.sin(a), 0], [np.sin(a) * np.sin(b), -np.cos(a) * np.cos(b)]]
-        assert np.allclose(res, expected, atol=tol, rtol=0)
+        assert isinstance(res, tuple) and len(res) == 2
+        expected = ([-np.sin(a), np.sin(a) * np.sin(b)], [0, -np.cos(a) * np.cos(b)])
+        assert np.allclose(res[0], expected[0], atol=tol, rtol=0)
+        assert np.allclose(res[1], expected[1], atol=tol, rtol=0)
 
         if diff_method in ("parameter-shift", "finite-diff"):
             spy.assert_called()
@@ -191,8 +193,9 @@ class TestQNode:
 
         jac_fn = qml.jacobian(circuit)
         res = jac_fn(a, b)
-        expected = [[-np.sin(a), 0], [np.sin(a) * np.sin(b), -np.cos(a) * np.cos(b)]]
-        assert np.allclose(res, expected, atol=tol, rtol=0)
+        expected = ([-np.sin(a), np.sin(a) * np.sin(b)], [0, -np.cos(a) * np.cos(b)])
+        assert np.allclose(res[0], expected[0], atol=tol, rtol=0)
+        assert np.allclose(res[1], expected[1], atol=tol, rtol=0)
 
         if diff_method in ("parameter-shift", "finite-diff"):
             spy.assert_called()
@@ -202,8 +205,10 @@ class TestQNode:
         b = np.array(0.832, requires_grad=True)
 
         res = jac_fn(a, b)
-        expected = [[-np.sin(a), 0], [np.sin(a) * np.sin(b), -np.cos(a) * np.cos(b)]]
-        assert np.allclose(res, expected, atol=tol, rtol=0)
+        expected = ([-np.sin(a), np.sin(a) * np.sin(b)], [0, -np.cos(a) * np.cos(b)])
+        expected = ([-np.sin(a), np.sin(a) * np.sin(b)], [0, -np.cos(a) * np.cos(b)])
+        assert np.allclose(res[0], expected[0], atol=tol, rtol=0)
+        assert np.allclose(res[1], expected[1], atol=tol, rtol=0)
 
     def test_jacobian_options(self, dev_name, diff_method, mode, mocker, tol):
         """Test setting jacobian options"""
@@ -305,7 +310,9 @@ class TestQNode:
             tape_params = np.array(circuit.qtape.get_parameters())
             assert np.all(tape_params == [a * c, c + c ** 2 + np.sin(a)])
 
-        assert res.shape == (2,)
+        assert isinstance(res, tuple) and len(res) == 2
+        assert res[0].shape == ()
+        assert res[1].shape == ()
 
     def test_no_trainable_parameters(self, dev_name, diff_method, mode, tol):
         """Test evaluation and Jacobian if there are no trainable parameters"""
@@ -472,10 +479,13 @@ class TestShotsIntegration:
 
         res = qml.jacobian(cost_fn)(a, b, shots=[10000, 10000, 10000])
         assert dev.shots is None
-        assert len(res) == 3
+        assert isinstance(res, tuple) and len(res) == 2
+        assert all(r.shape == (3,) for r in res)
 
         expected = [np.sin(a) * np.sin(b), -np.cos(a) * np.cos(b)]
-        assert np.allclose(np.mean(res, axis=0), expected, atol=0.1, rtol=0)
+        assert all(
+            np.allclose(np.mean(r, axis=0), e, atol=0.1, rtol=0) for r, e in zip(res, expected)
+        )
 
     def test_update_diff_method(self, mocker, tol):
         """Test that temporarily setting the shots updates the diff method"""
@@ -531,14 +541,13 @@ class TestQubitIntegration:
             return qml.probs(wires=[1])
 
         res = qml.jacobian(circuit)(x, y)
+        assert isinstance(res, tuple) and len(res) == 2
 
-        expected = np.array(
-            [
-                [-np.sin(x) * np.cos(y) / 2, -np.cos(x) * np.sin(y) / 2],
-                [np.cos(y) * np.sin(x) / 2, np.cos(x) * np.sin(y) / 2],
-            ]
+        expected = (
+            np.array([-np.sin(x) * np.cos(y) / 2, np.cos(y) * np.sin(x) / 2]),
+            np.array([-np.cos(x) * np.sin(y) / 2, np.cos(x) * np.sin(y) / 2]),
         )
-        assert np.allclose(res, expected, atol=tol, rtol=0)
+        assert all(np.allclose(r, e, atol=tol, rtol=0) for r, e in zip(res, expected))
 
     def test_multiple_probability_differentiation(self, dev_name, diff_method, mode, tol):
         """Tests correct output shape and evaluation for a tape
@@ -569,17 +578,26 @@ class TestQubitIntegration:
         assert np.allclose(res, expected, atol=tol, rtol=0)
 
         res = qml.jacobian(circuit)(x, y)
-        expected = np.array(
-            [
-                [[-np.sin(x) / 2, 0], [-np.sin(x) * np.cos(y) / 2, -np.cos(x) * np.sin(y) / 2]],
+        assert isinstance(res, tuple) and len(res) == 2
+        assert res[0].shape == (2, 2)
+        assert res[1].shape == (2, 2)
+
+        expected = (
+            np.array(
                 [
-                    [np.sin(x) / 2, 0],
-                    [np.cos(y) * np.sin(x) / 2, np.cos(x) * np.sin(y) / 2],
-                ],
-            ]
+                    [-np.sin(x) / 2, np.sin(x) / 2],
+                    [-np.sin(x) * np.cos(y) / 2, np.sin(x) * np.cos(y) / 2],
+                ]
+            ),
+            np.array(
+                [
+                    [0, 0],
+                    [-np.cos(x) * np.sin(y) / 2, np.cos(x) * np.sin(y) / 2],
+                ]
+            ),
         )
 
-        assert np.allclose(res, expected, atol=tol, rtol=0)
+        assert all(np.allclose(r, e, atol=tol, rtol=0) for r, e in zip(res, expected))
 
     def test_ragged_differentiation(self, dev_name, diff_method, mode, tol):
         """Tests correct output shape and evaluation for a tape
@@ -606,14 +624,16 @@ class TestQubitIntegration:
         assert np.allclose(res, expected, atol=tol, rtol=0)
 
         res = qml.jacobian(circuit)(x, y)
-        expected = np.array(
-            [
-                [-np.sin(x), 0],
-                [-np.sin(x) * np.cos(y) / 2, -np.cos(x) * np.sin(y) / 2],
-                [np.cos(y) * np.sin(x) / 2, np.cos(x) * np.sin(y) / 2],
-            ]
+        assert isinstance(res, tuple) and len(res) == 2
+        assert res[0].shape == (3,)
+        assert res[1].shape == (3,)
+
+        expected = (
+            np.array([-np.sin(x), -np.sin(x) * np.cos(y) / 2, np.sin(x) * np.cos(y) / 2]),
+            np.array([0, -np.cos(x) * np.sin(y) / 2, np.cos(x) * np.sin(y) / 2]),
         )
-        assert np.allclose(res, expected, atol=tol, rtol=0)
+        assert np.allclose(res[0], expected[0], atol=tol, rtol=0)
+        assert np.allclose(res[1], expected[1], atol=tol, rtol=0)
 
     def test_ragged_differentiation_variance(self, dev_name, diff_method, mode, tol):
         """Tests correct output shape and evaluation for a tape
@@ -640,14 +660,14 @@ class TestQubitIntegration:
         assert np.allclose(res, expected, atol=tol, rtol=0)
 
         res = qml.jacobian(circuit)(x, y)
-        expected = np.array(
-            [
-                [2 * np.cos(x) * np.sin(x), 0],
-                [-np.sin(x) * np.cos(y) / 2, -np.cos(x) * np.sin(y) / 2],
-                [np.cos(y) * np.sin(x) / 2, np.cos(x) * np.sin(y) / 2],
-            ]
+        assert isinstance(res, tuple) and len(res) == 2
+
+        expected = (
+            np.array([np.sin(2 * x), -np.sin(x) * np.cos(y) / 2, np.sin(x) * np.cos(y) / 2]),
+            np.array([0, -np.cos(x) * np.sin(y) / 2, np.cos(x) * np.sin(y) / 2]),
         )
-        assert np.allclose(res, expected, atol=tol, rtol=0)
+        assert np.allclose(res[0], expected[0], atol=tol, rtol=0)
+        assert np.allclose(res[1], expected[1], atol=tol, rtol=0)
 
     def test_chained_qnodes(self, dev_name, diff_method, mode):
         """Test that the gradient of chained QNodes works without error"""
@@ -898,13 +918,13 @@ class TestQubitIntegration:
         assert np.allclose(res, expected_res, atol=tol, rtol=0)
 
         jac_fn = qml.jacobian(circuit)
-        g = jac_fn(x)
+        res = jac_fn(x)
 
-        expected_g = [
+        expected_res = [
             [-0.5 * np.sin(a) * np.cos(b), -0.5 * np.cos(a) * np.sin(b)],
             [0.5 * np.sin(a) * np.cos(b), 0.5 * np.cos(a) * np.sin(b)],
         ]
-        assert np.allclose(g, expected_g, atol=tol, rtol=0)
+        assert np.allclose(res, expected_res, atol=tol, rtol=0)
 
         hess = qml.jacobian(jac_fn)(x)
 
@@ -989,32 +1009,39 @@ class TestQubitIntegration:
 
         jac_fn = qml.jacobian(circuit)
         g = jac_fn(a, b)
+        assert isinstance(g, tuple) and len(g) == 2
 
-        expected_g = [
-            [-0.5 * np.sin(a) * np.cos(b), -0.5 * np.cos(a) * np.sin(b)],
-            [0.5 * np.sin(a) * np.cos(b), 0.5 * np.cos(a) * np.sin(b)],
-        ]
-        assert np.allclose(g, expected_g, atol=tol, rtol=0)
+        expected_g = (
+            [-0.5 * np.sin(a) * np.cos(b), 0.5 * np.sin(a) * np.cos(b)],
+            [-0.5 * np.cos(a) * np.sin(b), 0.5 * np.cos(a) * np.sin(b)],
+        )
+        assert np.allclose(g[0], expected_g[0], atol=tol, rtol=0)
+        assert np.allclose(g[1], expected_g[1], atol=tol, rtol=0)
 
         spy = mocker.spy(qml.gradients.param_shift, "transform_fn")
-        hess = qml.jacobian(jac_fn)(a, b)
+        jac_fn_a = lambda *args: jac_fn(*args)[0]
+        jac_fn_b = lambda *args: jac_fn(*args)[1]
+        hess_a = qml.jacobian(jac_fn_a)(a, b)
+        hess_b = qml.jacobian(jac_fn_b)(a, b)
+        assert isinstance(hess_a, tuple) and len(hess_a) == 2
+        assert isinstance(hess_b, tuple) and len(hess_b) == 2
 
         if diff_method == "backprop":
             spy.assert_not_called()
         elif diff_method == "parameter-shift":
             spy.assert_called()
 
-        expected_hess = [
-            [
-                [-0.5 * np.cos(a) * np.cos(b), 0.5 * np.sin(a) * np.sin(b)],
-                [0.5 * np.cos(a) * np.cos(b), -0.5 * np.sin(a) * np.sin(b)],
-            ],
-            [
-                [0.5 * np.sin(a) * np.sin(b), -0.5 * np.cos(a) * np.cos(b)],
-                [-0.5 * np.sin(a) * np.sin(b), 0.5 * np.cos(a) * np.cos(b)],
-            ],
-        ]
-        assert np.allclose(hess, expected_hess, atol=tol, rtol=0)
+        exp_hess_a = (
+            [-0.5 * np.cos(a) * np.cos(b), 0.5 * np.cos(a) * np.cos(b)],
+            [0.5 * np.sin(a) * np.sin(b), -0.5 * np.sin(a) * np.sin(b)],
+        )
+        exp_hess_b = (
+            [0.5 * np.sin(a) * np.sin(b), -0.5 * np.sin(a) * np.sin(b)],
+            [-0.5 * np.cos(a) * np.cos(b), 0.5 * np.cos(a) * np.cos(b)],
+        )
+        for hess, exp_hess in zip([hess_a, hess_b], [exp_hess_a, exp_hess_b]):
+            assert np.allclose(hess[0], exp_hess[0], atol=tol, rtol=0)
+            assert np.allclose(hess[1], exp_hess[1], atol=tol, rtol=0)
 
     def test_hessian_ragged(self, dev_name, diff_method, mode, tol):
         """Test hessian calculation of a ragged QNode"""
diff --git a/tests/interfaces/test_qnode_autograd.py b/tests/interfaces/test_qnode_autograd.py
index fd3ef7305..ce5798fbf 100644
--- a/tests/interfaces/test_qnode_autograd.py
+++ b/tests/interfaces/test_qnode_autograd.py
@@ -194,8 +194,10 @@ class TestQNode:
         assert np.allclose(res, expected, atol=tol, rtol=0)
 
         res = qml.jacobian(circuit)(a, b)
-        expected = [[-np.sin(a), 0], [np.sin(a) * np.sin(b), -np.cos(a) * np.cos(b)]]
-        assert np.allclose(res, expected, atol=tol, rtol=0)
+        assert isinstance(res, tuple) and len(res) == 2
+        expected = ([-np.sin(a), np.sin(a) * np.sin(b)], [0, -np.cos(a) * np.cos(b)])
+        assert np.allclose(res[0], expected[0], atol=tol, rtol=0)
+        assert np.allclose(res[1], expected[1], atol=tol, rtol=0)
 
         if diff_method == "finite-diff":
             spy.assert_called()
@@ -219,8 +221,10 @@ class TestQNode:
 
         jac_fn = qml.jacobian(circuit)
         res = jac_fn(a, b)
-        expected = [[-np.sin(a), 0], [np.sin(a) * np.sin(b), -np.cos(a) * np.cos(b)]]
-        assert np.allclose(res, expected, atol=tol, rtol=0)
+        assert isinstance(res, tuple) and len(res) == 2
+        expected = ([-np.sin(a), np.sin(a) * np.sin(b)], [0, -np.cos(a) * np.cos(b)])
+        assert np.allclose(res[0], expected[0], atol=tol, rtol=0)
+        assert np.allclose(res[1], expected[1], atol=tol, rtol=0)
 
         if diff_method == "finite-diff":
             spy.assert_called()
@@ -232,8 +236,10 @@ class TestQNode:
         b = np.array(0.832, requires_grad=True)
 
         res = jac_fn(a, b)
-        expected = [[-np.sin(a), 0], [np.sin(a) * np.sin(b), -np.cos(a) * np.cos(b)]]
-        assert np.allclose(res, expected, atol=tol, rtol=0)
+        assert isinstance(res, tuple) and len(res) == 2
+        expected = ([-np.sin(a), np.sin(a) * np.sin(b)], [0, -np.cos(a) * np.cos(b)])
+        assert np.allclose(res[0], expected[0], atol=tol, rtol=0)
+        assert np.allclose(res[1], expected[1], atol=tol, rtol=0)
 
     def test_jacobian_options(self, dev_name, diff_method, mocker, tol):
         """Test setting jacobian options"""
@@ -337,7 +343,7 @@ class TestQNode:
             tape_params = np.array(circuit.qtape.get_parameters())
             assert np.all(tape_params == [a * c, c + c ** 2 + np.sin(a)])
 
-        assert res.shape == (2,)
+        assert isinstance(res, tuple) and len(res) == 2
 
     def test_no_trainable_parameters(self, dev_name, diff_method, tol):
         """Test evaluation and Jacobian if there are no trainable parameters"""
@@ -461,14 +467,14 @@ class TestQNode:
             return qml.probs(wires=[1])
 
         res = qml.jacobian(circuit)(x, y)
+        assert isinstance(res, tuple) and len(res) == 2
 
-        expected = np.array(
-            [
-                [-np.sin(x) * np.cos(y) / 2, -np.cos(x) * np.sin(y) / 2],
-                [np.cos(y) * np.sin(x) / 2, np.cos(x) * np.sin(y) / 2],
-            ]
+        expected = (
+            [-np.sin(x) * np.cos(y) / 2, np.cos(y) * np.sin(x) / 2],
+            [-np.cos(x) * np.sin(y) / 2, np.cos(x) * np.sin(y) / 2],
         )
-        assert np.allclose(res, expected, atol=tol, rtol=0)
+        assert np.allclose(res[0], expected[0], atol=tol, rtol=0)
+        assert np.allclose(res[1], expected[1], atol=tol, rtol=0)
 
     def test_multiple_probability_differentiation(self, dev_name, diff_method, tol):
         """Tests correct output shape and evaluation for a tape
@@ -499,17 +505,21 @@ class TestQNode:
         assert np.allclose(res, expected, atol=tol, rtol=0)
 
         res = qml.jacobian(circuit)(x, y)
-        expected = np.array(
+        assert isinstance(res, tuple) and len(res) == 2
+
+        expected = (
             [
-                [[-np.sin(x) / 2, 0], [-np.sin(x) * np.cos(y) / 2, -np.cos(x) * np.sin(y) / 2]],
-                [
-                    [np.sin(x) / 2, 0],
-                    [np.cos(y) * np.sin(x) / 2, np.cos(x) * np.sin(y) / 2],
-                ],
-            ]
+                [-np.sin(x) / 2, np.sin(x) / 2],
+                [-np.sin(x) * np.cos(y) / 2, np.sin(x) * np.cos(y) / 2],
+            ],
+            [
+                [0, 0],
+                [-np.cos(x) * np.sin(y) / 2, np.cos(x) * np.sin(y) / 2],
+            ],
         )
 
-        assert np.allclose(res, expected, atol=tol, rtol=0)
+        assert np.allclose(res[0], expected[0], atol=tol, rtol=0)
+        assert np.allclose(res[1], expected[1], atol=tol, rtol=0)
 
     def test_ragged_differentiation(self, dev_name, diff_method, tol):
         """Tests correct output shape and evaluation for a tape
@@ -536,14 +546,14 @@ class TestQNode:
         assert np.allclose(res, expected, atol=tol, rtol=0)
 
         res = qml.jacobian(circuit)(x, y)
-        expected = np.array(
-            [
-                [-np.sin(x), 0],
-                [-np.sin(x) * np.cos(y) / 2, -np.cos(x) * np.sin(y) / 2],
-                [np.cos(y) * np.sin(x) / 2, np.cos(x) * np.sin(y) / 2],
-            ]
+        assert isinstance(res, tuple) and len(res) == 2
+
+        expected = (
+            [-np.sin(x), -np.sin(x) * np.cos(y) / 2, np.cos(y) * np.sin(x) / 2],
+            [0, -np.cos(x) * np.sin(y) / 2, np.cos(x) * np.sin(y) / 2],
         )
-        assert np.allclose(res, expected, atol=tol, rtol=0)
+        assert np.allclose(res[0], expected[0], atol=tol, rtol=0)
+        assert np.allclose(res[1], expected[1], atol=tol, rtol=0)
 
     def test_ragged_differentiation_variance(self, dev_name, diff_method, tol):
         """Tests correct output shape and evaluation for a tape
@@ -570,14 +580,14 @@ class TestQNode:
         assert np.allclose(res, expected, atol=tol, rtol=0)
 
         res = qml.jacobian(circuit)(x, y)
-        expected = np.array(
-            [
-                [2 * np.cos(x) * np.sin(x), 0],
-                [-np.sin(x) * np.cos(y) / 2, -np.cos(x) * np.sin(y) / 2],
-                [np.cos(y) * np.sin(x) / 2, np.cos(x) * np.sin(y) / 2],
-            ]
+        assert isinstance(res, tuple) and len(res) == 2
+
+        expected = (
+            [np.sin(2 * x), -np.sin(x) * np.cos(y) / 2, np.sin(x) * np.cos(y) / 2],
+            [0, -np.cos(x) * np.sin(y) / 2, np.cos(x) * np.sin(y) / 2],
         )
-        assert np.allclose(res, expected, atol=tol, rtol=0)
+        assert np.allclose(res[0], expected[0], atol=tol, rtol=0)
+        assert np.allclose(res[1], expected[1], atol=tol, rtol=0)
 
     def test_sampling(self, dev_name, diff_method):
         """Test sampling works as expected"""
@@ -977,15 +987,22 @@ class TestQNode:
 
         jac_fn = qml.jacobian(circuit)
         g = jac_fn(a, b)
+        assert isinstance(g, tuple) and len(g) == 2
 
-        expected_g = [
-            [-0.5 * np.sin(a) * np.cos(b), -0.5 * np.cos(a) * np.sin(b)],
-            [0.5 * np.sin(a) * np.cos(b), 0.5 * np.cos(a) * np.sin(b)],
-        ]
-        assert np.allclose(g, expected_g, atol=tol, rtol=0)
+        expected_g = (
+            [-0.5 * np.sin(a) * np.cos(b), 0.5 * np.sin(a) * np.cos(b)],
+            [-0.5 * np.cos(a) * np.sin(b), 0.5 * np.cos(a) * np.sin(b)],
+        )
+        assert np.allclose(g[0], expected_g[0], atol=tol, rtol=0)
+        assert np.allclose(g[1], expected_g[1], atol=tol, rtol=0)
 
         spy = mocker.spy(JacobianTape, "hessian")
-        hess = qml.jacobian(jac_fn)(a, b)
+        jac_fn_a = lambda *args: jac_fn(*args)[0]
+        jac_fn_b = lambda *args: jac_fn(*args)[1]
+        hess_a = qml.jacobian(jac_fn_a)(a, b)
+        hess_b = qml.jacobian(jac_fn_b)(a, b)
+        assert isinstance(hess_a, tuple) and len(hess_a) == 2
+        assert isinstance(hess_b, tuple) and len(hess_b) == 2
 
         if diff_method == "backprop":
             spy.assert_not_called()
@@ -994,17 +1011,17 @@ class TestQNode:
             # since autograd will treat them as separate arguments.
             assert spy.call_count == 4
 
-        expected_hess = [
-            [
-                [-0.5 * np.cos(a) * np.cos(b), 0.5 * np.sin(a) * np.sin(b)],
-                [0.5 * np.cos(a) * np.cos(b), -0.5 * np.sin(a) * np.sin(b)],
-            ],
-            [
-                [0.5 * np.sin(a) * np.sin(b), -0.5 * np.cos(a) * np.cos(b)],
-                [-0.5 * np.sin(a) * np.sin(b), 0.5 * np.cos(a) * np.cos(b)],
-            ],
-        ]
-        assert np.allclose(hess, expected_hess, atol=tol, rtol=0)
+        exp_hess_a = (
+            [-0.5 * np.cos(a) * np.cos(b), 0.5 * np.cos(a) * np.cos(b)],
+            [0.5 * np.sin(a) * np.sin(b), -0.5 * np.sin(a) * np.sin(b)],
+        )
+        exp_hess_b = (
+            [0.5 * np.sin(a) * np.sin(b), -0.5 * np.sin(a) * np.sin(b)],
+            [-0.5 * np.cos(a) * np.cos(b), 0.5 * np.cos(a) * np.cos(b)],
+        )
+        for hess, exp_hess in zip([hess_a, hess_b], [exp_hess_a, exp_hess_b]):
+            assert np.allclose(hess[0], exp_hess[0], atol=tol, rtol=0)
+            assert np.allclose(hess[1], exp_hess[1], atol=tol, rtol=0)
 
     def test_hessian_ragged(self, dev_name, diff_method, mocker, tol):
         """Test hessian calculation of a ragged QNode"""
diff --git a/tests/interfaces/test_tape_autograd.py b/tests/interfaces/test_tape_autograd.py
index 6d850e55e..d7ea9db87 100644
--- a/tests/interfaces/test_tape_autograd.py
+++ b/tests/interfaces/test_tape_autograd.py
@@ -114,10 +114,13 @@ class TestAutogradQuantumTape:
         assert np.allclose(res, expected, atol=tol, rtol=0)
 
         res = qml.jacobian(cost)(a, b, device=dev)
-        assert res.shape == (2, 2)
+        assert isinstance(res, tuple) and len(res) == 2
+        assert res[0].shape == (2,)
+        assert res[1].shape == (2,)
 
-        expected = [[-np.sin(a), 0], [np.sin(a) * np.sin(b), -np.cos(a) * np.cos(b)]]
-        assert np.allclose(res, expected, atol=tol, rtol=0)
+        expected = ([-np.sin(a), np.sin(a) * np.sin(b)], [0, -np.cos(a) * np.cos(b)])
+        assert np.allclose(res[0], expected[0], atol=tol, rtol=0)
+        assert np.allclose(res[1], expected[1], atol=tol, rtol=0)
 
     def test_jacobian_options(self, mocker, tol):
         """Test setting jacobian options"""
@@ -174,11 +177,14 @@ class TestAutogradQuantumTape:
 
         jac_fn = qml.jacobian(lambda a, b: cost(2 * a, b))
         jac = jac_fn(a, b)
-        expected = [
-            [-2 * np.sin(2 * a), 0],
-            [2 * np.sin(2 * a) * np.sin(b), -np.cos(2 * a) * np.cos(b)],
-        ]
-        assert np.allclose(jac, expected, atol=tol, rtol=0)
+        assert isinstance(jac, tuple) and len(jac) == 2
+
+        expected = (
+            [-2 * np.sin(2 * a), 2 * np.sin(2 * a) * np.sin(b)],
+            [0, -np.cos(2 * a) * np.cos(b)],
+        )
+        assert np.allclose(jac[0], expected[0], atol=tol, rtol=0)
+        assert np.allclose(jac[1], expected[1], atol=tol, rtol=0)
 
     def test_classical_processing(self, tol):
         """Test classical processing within the quantum tape"""
@@ -197,7 +203,9 @@ class TestAutogradQuantumTape:
 
         dev = qml.device("default.qubit", wires=2)
         res = qml.jacobian(cost)(a, b, c, device=dev)
-        assert res.shape == (1, 2)
+        assert isinstance(res, tuple) and len(res) == 2
+        assert res[0].shape == (1,)
+        assert res[1].shape == (1,)
 
     def test_no_trainable_parameters(self, tol):
         """Test evaluation and Jacobian if there are no trainable parameters"""
@@ -335,17 +343,21 @@ class TestAutogradQuantumTape:
 
         jac_fn = qml.jacobian(cost)
         res = jac_fn(x, y, device=dev)
-        assert res.shape == (2, 2, 2)
-        expected = np.array(
+        assert isinstance(res, tuple) and len(res) == 2
+        assert res[0].shape == (2, 2)
+        assert res[1].shape == (2, 2)
+        expected = (
             [
-                [[-np.sin(x) / 2, 0], [-np.sin(x) * np.cos(y) / 2, -np.cos(x) * np.sin(y) / 2]],
-                [
-                    [np.sin(x) / 2, 0],
-                    [np.cos(y) * np.sin(x) / 2, np.cos(x) * np.sin(y) / 2],
-                ],
-            ]
+                [-np.sin(x) / 2, np.sin(x) / 2],
+                [-np.sin(x) * np.cos(y) / 2, np.sin(x) * np.cos(y) / 2],
+            ],
+            [
+                [0, 0],
+                [-np.cos(x) * np.sin(y) / 2, np.cos(x) * np.sin(y) / 2],
+            ],
         )
-        assert np.allclose(res, expected, atol=tol, rtol=0)
+        assert np.allclose(res[0], expected[0], atol=tol, rtol=0)
+        assert np.allclose(res[1], expected[1], atol=tol, rtol=0)
 
     def test_ragged_differentiation(self, tol):
         """Tests correct output shape and evaluation for a tape
@@ -373,14 +385,12 @@ class TestAutogradQuantumTape:
 
         jac_fn = qml.jacobian(cost)
         res = jac_fn(x, y, device=dev)
-        expected = np.array(
-            [
-                [-np.sin(x), 0],
-                [-np.sin(x) * np.cos(y) / 2, -np.cos(x) * np.sin(y) / 2],
-                [np.cos(y) * np.sin(x) / 2, np.cos(x) * np.sin(y) / 2],
-            ]
+        expected = (
+            [-np.sin(x), -np.sin(x) * np.cos(y) / 2, np.sin(x) * np.cos(y) / 2],
+            [0, -np.cos(x) * np.sin(y) / 2, np.cos(x) * np.sin(y) / 2],
         )
-        assert np.allclose(res, expected, atol=tol, rtol=0)
+        assert np.allclose(res[0], expected[0], atol=tol, rtol=0)
+        assert np.allclose(res[1], expected[1], atol=tol, rtol=0)
 
     def test_sampling(self):
         """Test sampling works as expected"""
@@ -474,7 +484,9 @@ class TestAutogradPassthru:
         dev = qml.device("default.qubit.autograd", wires=2)
         res = qml.jacobian(cost)(a, b, device=dev)
         spy.assert_not_called()
-        assert res.shape == (2, 2)
+        assert isinstance(res, tuple) and len(res) == 2
+        assert res[0].shape == (2,)
+        assert res[1].shape == (2,)
 
         # compare to standard tape jacobian
         with JacobianTape() as tape:
@@ -485,7 +497,8 @@ class TestAutogradPassthru:
 
         expected = tape.jacobian(dev)
         assert expected.shape == (2, 2)
-        assert np.allclose(res, expected, atol=tol, rtol=0)
+        assert np.allclose(res[0], expected.T[0], atol=tol, rtol=0)
+        assert np.allclose(res[1], expected.T[1], atol=tol, rtol=0)
 
     def test_classical_processing(self, mocker, tol):
         """Test classical processing within the quantum tape"""
@@ -504,7 +517,9 @@ class TestAutogradPassthru:
 
         dev = qml.device("default.qubit.autograd", wires=2)
         res = qml.jacobian(cost)(a, b, c, device=dev)
-        assert res.shape == (1, 2)
+        assert isinstance(res, tuple) and len(res) == 2
+        assert res[0].shape == (1,)
+        assert res[1].shape == (1,)
         spy.assert_not_called()
 
     def test_no_trainable_parameters(self, mocker, tol):
@@ -651,18 +666,22 @@ class TestAutogradPassthru:
 
         jac_fn = qml.jacobian(cost)
         res = jac_fn(x, y, device=dev)
-        assert res.shape == (2, 2, 2)
+        assert isinstance(res, tuple) and len(res) == 2
+        assert res[0].shape == (2, 2)
+        assert res[1].shape == (2, 2)
 
-        expected = np.array(
+        expected = (
             [
-                [[-np.sin(x) / 2, 0], [-np.sin(x) * np.cos(y) / 2, -np.cos(x) * np.sin(y) / 2]],
-                [
-                    [np.sin(x) / 2, 0],
-                    [np.cos(y) * np.sin(x) / 2, np.cos(x) * np.sin(y) / 2],
-                ],
-            ]
+                [-np.sin(x) / 2, np.sin(x) / 2],
+                [-np.sin(x) * np.cos(y) / 2, np.sin(x) * np.cos(y) / 2],
+            ],
+            [
+                [0, 0],
+                [-np.cos(x) * np.sin(y) / 2, np.cos(x) * np.sin(y) / 2],
+            ],
         )
-        assert np.allclose(res, expected, atol=tol, rtol=0)
+        assert np.allclose(res[0], expected[0], atol=tol, rtol=0)
+        assert np.allclose(res[1], expected[1], atol=tol, rtol=0)
         spy.assert_not_called()
 
     def test_ragged_differentiation(self, mocker, monkeypatch, tol):
@@ -702,14 +721,15 @@ class TestAutogradPassthru:
 
         jac_fn = qml.jacobian(cost)
         res = jac_fn(x, y, device=dev)
-        expected = np.array(
-            [
-                [-np.sin(x), 0],
-                [-np.sin(x) * np.cos(y) / 2, -np.cos(x) * np.sin(y) / 2],
-                [np.cos(y) * np.sin(x) / 2, np.cos(x) * np.sin(y) / 2],
-            ]
+        assert isinstance(res, tuple) and len(res) == 2
+        assert res[0].shape == (3,)
+        assert res[1].shape == (3,)
+        expected = (
+            [-np.sin(x), -np.sin(x) * np.cos(y) / 2, np.sin(x) * np.cos(y) / 2],
+            [0, -np.cos(x) * np.sin(y) / 2, np.cos(x) * np.sin(y) / 2],
         )
-        assert np.allclose(res, expected, atol=tol, rtol=0)
+        assert np.allclose(res[0], expected[0], atol=tol, rtol=0)
+        assert np.allclose(res[1], expected[1], atol=tol, rtol=0)
         spy.assert_not_called()
 
     def test_sampling(self):
diff --git a/tests/math/test_functions.py b/tests/math/test_functions.py
index a82625ab5..2be238b7f 100644
--- a/tests/math/test_functions.py
+++ b/tests/math/test_functions.py
@@ -1851,12 +1851,17 @@ def test_block_diag(tensors):
 
 class TestBlockDiagDiffability:
 
-    expected = lambda self, x, y: np.array(
+    expected = lambda self, x, y: (
         [
-            [[-np.sin(x * y) * y, -np.sin(x * y) * x], [0, 0], [0, 0]],
-            [[0, 0], [1.0, 0.0], [0, 1.2]],
-            [[0, 0], [2 * x, -1 / 3], [-1 / y, x / y ** 2]],
-        ]
+            [-np.sin(x * y) * y, 0, 0],
+            [0, 1.0, 0],
+            [0, 2 * x, -1 / y],
+        ],
+        [
+            [-np.sin(x * y) * x, 0, 0],
+            [0, 0.0, 1.2],
+            [0, -1 / 3, x / y ** 2],
+        ],
     )
 
     def test_autograd(self):
@@ -1869,9 +1874,8 @@ class TestBlockDiagDiffability:
         x, y = 0.2, 1.5
         res = qml.jacobian(f)(x, y)
         exp = self.expected(x, y)
-        # Transposes in the following because autograd behaves strangely
-        assert fn.allclose(res[:, :, 0].T, exp[:, :, 0])
-        assert fn.allclose(res[:, :, 1].T, exp[:, :, 1])
+        assert fn.allclose(res[0], exp[0])
+        assert fn.allclose(res[1], exp[1])
 
     def test_jax(self):
         """Tests for differentiating the block diagonal function with JAX."""
@@ -1884,8 +1888,8 @@ class TestBlockDiagDiffability:
         x, y = 0.2, 1.5
         res = jax.jacobian(f, argnums=[0, 1])(x, y)
         exp = self.expected(x, y)
-        assert fn.allclose(exp[:, :, 0], res[0])
-        assert fn.allclose(exp[:, :, 1], res[1])
+        assert fn.allclose(exp[0], res[0])
+        assert fn.allclose(exp[1], res[1])
 
     def test_tf(self):
         """Tests for differentiating the block diagonal function with Tensorflow."""
@@ -2167,3 +2171,126 @@ class TestSortFunction:
         result = fn.sort(input)
 
         assert all(result == test_output)
+
+
+ones_functions = [onp.ones, np.ones, jnp.ones, torch.ones, tf.ones]
+
+
+class TestSafeSqueeze:
+    """Test the safe version of squeeze that allows to specify
+    axes of size bigger than one without raising an error and that
+    allows to exclude axes."""
+
+    input_shapes = [
+        (2, 5, 3),
+        (1, 1, 9, 1),
+        (1, 1),
+        (5, 1),
+        (),
+    ]
+
+    @pytest.mark.parametrize(
+        "shape, expected_shape",
+        list(zip(input_shapes, [(2, 5, 3), (9,), (), (5,), ()])),
+    )
+    @pytest.mark.parametrize("ones_fn", ones_functions)
+    def test_without_kwargs(self, ones_fn, shape, expected_shape):
+        """Test squeezing without specifying axes."""
+
+        tensor = ones_fn(shape)
+        squeezed = fn.safe_squeeze(tensor)
+        assert fn.shape(squeezed) == expected_shape
+
+    # Skipping the first and last input shape
+    @pytest.mark.parametrize(
+        "shape, axis, expected_shape",
+        list(zip(input_shapes[1:4], [3, -2, 1], [(1, 1, 9), (1,), (5,)])),
+    )
+    @pytest.mark.parametrize("ones_fn", ones_functions)
+    def test_with_axis_of_size_one(self, ones_fn, shape, axis, expected_shape):
+        """Test squeezing a single axis with size one."""
+
+        tensor = ones_fn(shape)
+        squeezed = fn.safe_squeeze(tensor, axis=axis)
+        assert fn.shape(squeezed) == expected_shape
+
+    @pytest.mark.parametrize(
+        "shape, axis",
+        list(zip(input_shapes[:2], [1, 2])),
+    )
+    @pytest.mark.parametrize("ones_fn", ones_functions)
+    def test_with_axis_of_any_size(self, ones_fn, shape, axis):
+        """Test squeezing a single axis of any size."""
+
+        tensor = ones_fn(shape)
+        squeezed = fn.safe_squeeze(tensor, axis=axis)
+        assert fn.shape(squeezed) == shape
+
+    # Skipping the first and last input shape
+    @pytest.mark.parametrize(
+        "shape, axis, expected_shape",
+        list(zip(input_shapes[1:4], [(0, 1, -1), [1, -2], {1}], [(9,), (), (5,)])),
+    )
+    @pytest.mark.parametrize("ones_fn", ones_functions)
+    def test_with_multiple_axes_of_size_one(self, ones_fn, shape, axis, expected_shape):
+        """Test squeezing multiple axes with size one."""
+
+        tensor = ones_fn(shape)
+        squeezed = fn.safe_squeeze(tensor, axis=axis)
+        assert fn.shape(squeezed) == expected_shape
+
+    @pytest.mark.parametrize(
+        "shape, axis, expected_shape",
+        list(
+            zip(
+                input_shapes,
+                [(2, 1), (0, 1, 2), [0], {0, 1}, []],
+                [(2, 5, 3), (9, 1), (1,), (5,), ()],
+            )
+        ),
+    )
+    @pytest.mark.parametrize("ones_fn", ones_functions)
+    def test_with_multiple_axes_of_any_size(self, ones_fn, shape, axis, expected_shape):
+        """Test squeezing multiple axes with any size."""
+
+        tensor = ones_fn(shape)
+        squeezed = fn.safe_squeeze(tensor, axis=axis)
+        assert fn.shape(squeezed) == expected_shape
+
+    @pytest.mark.parametrize(
+        "shape, exclude_axis, expected_shape",
+        list(
+            zip(
+                input_shapes[:4],
+                [-1, -3, 0, 1],
+                [(2, 5, 3), (1, 9), (1,), (5, 1)],
+            )
+        ),
+    )
+    @pytest.mark.parametrize("ones_fn", ones_functions)
+    def test_with_single_exclude_axis(self, ones_fn, shape, exclude_axis, expected_shape):
+        """Tests squeezing with the kwarg exclude_axis and a single axis to ignore
+        while squeezing."""
+
+        tensor = ones_fn(shape)
+        squeezed = fn.safe_squeeze(tensor, exclude_axis=exclude_axis)
+        assert fn.shape(squeezed) == expected_shape
+
+    @pytest.mark.parametrize(
+        "shape, exclude_axis, expected_shape",
+        list(
+            zip(
+                input_shapes,
+                [(0,), (-2, 3), [1], {-1, -2}, []],
+                [(2, 5, 3), (9, 1), (1,), (5, 1), ()],
+            )
+        ),
+    )
+    @pytest.mark.parametrize("ones_fn", ones_functions)
+    def test_with_multiple_exclude_axis(self, ones_fn, shape, exclude_axis, expected_shape):
+        """Tests squeezing with the kwarg exclude_axis and multiple axes to ignore
+        while squeezing."""
+
+        tensor = ones_fn(shape)
+        squeezed = fn.safe_squeeze(tensor, exclude_axis=exclude_axis)
+        assert fn.shape(squeezed) == expected_shape
diff --git a/tests/optimize/test_qng.py b/tests/optimize/test_qng.py
index 2e512b178..e4cc6c1fe 100644
--- a/tests/optimize/test_qng.py
+++ b/tests/optimize/test_qng.py
@@ -104,6 +104,7 @@ class TestOptimize:
         assert np.isclose(cost1, expected_cost)
         assert np.isclose(cost2, expected_cost)
 
+    @pytest.mark.skip("QNGOptimizer is not yet implemented for split inputs.")
     def test_step_and_cost_with_grad_fn_split_input(self, tol):
         """Test that the correct cost and update is returned via the step_and_cost
         method for the QNG optimizer when providing an explicit grad_fn.
@@ -278,6 +279,7 @@ class TestOptimize:
         assert np.allclose(cost_fn(theta), -1.41421356, atol=tol, rtol=0)
         assert len(recwarn) == 0
 
+    @pytest.mark.skip("QNGOptimizer is not yet implemented for split inputs.")
     def test_single_qubit_vqe_using_expval_h_multiple_input_params(self, tol, recwarn):
         """Test single-qubit VQE by returning qml.expval(H) in the QNode and
         check for the correct QNG value every step, the correct parameter updates, and
diff --git a/tests/test_classical_gradients.py b/tests/test_classical_gradients.py
index e9a948725..a4a4a6811 100644
--- a/tests/test_classical_gradients.py
+++ b/tests/test_classical_gradients.py
@@ -355,8 +355,11 @@ class TestJacobian:
 
         jac_fn = qml.jacobian(cost_fn, argnum=[0, 1])
         res = jac_fn(x, y)
-        expected = np.array([[np.cos(x) * np.cos(y), -np.sin(x) * np.sin(y)], [y ** 2, 2 * x * y]])
-        assert np.allclose(res, expected, atol=tol, rtol=0)
+        expected = (
+            np.array([np.cos(x) * np.cos(y), y ** 2]),
+            np.array([-np.sin(x) * np.sin(y), 2 * x * y]),
+        )
+        assert all(np.allclose(_r, _e, atol=tol, rtol=0) for _r, _e in zip(res, expected))
 
     def test_no_argnum_jacobian(self, tol):
         """Test the qml.jacobian function for inferred argnums"""
@@ -367,12 +370,14 @@ class TestJacobian:
 
         jac_fn = qml.jacobian(cost_fn)
         res = jac_fn(x, y)
-        expected = np.array([[np.cos(x) * np.cos(y), -np.sin(x) * np.sin(y)], [y ** 2, 2 * x * y]])
-        assert np.allclose(res, expected, atol=tol, rtol=0)
+        expected = (
+            np.array([np.cos(x) * np.cos(y), y ** 2]),
+            np.array([-np.sin(x) * np.sin(y), 2 * x * y]),
+        )
+        assert all(np.allclose(_r, _e, atol=tol, rtol=0) for _r, _e in zip(res, expected))
 
-        x = np.array(0.5, requires_grad=True)
-        y = np.array(0.2, requires_grad=False)
+        x = np.array(0.5, requires_grad=False)
+        y = np.array(0.2, requires_grad=True)
 
         res = jac_fn(x, y)
-        expected = np.array([np.cos(x) * np.cos(y), y ** 2])
-        assert np.allclose(res, expected, atol=tol, rtol=0)
+        assert np.allclose(res, expected[1], atol=tol, rtol=0)
diff --git a/tests/test_prob.py b/tests/test_prob.py
index 0f8c292ee..08fa61883 100644
--- a/tests/test_prob.py
+++ b/tests/test_prob.py
@@ -150,11 +150,13 @@ def test_numerical_analytic_diff_agree(init_state, tol):
     res_A = qml.jacobian(circuit_A)(*params)
 
     # Both jacobians should be of shape (2**prob.wires, num_params)
-    assert res_F.shape == (2 ** 2, 3)
-    assert res_F.shape == (2 ** 2, 3)
+    assert isinstance(res_F, tuple) and len(res_F) == 3
+    assert all(_r.shape == (2 ** 2,) for _r in res_F)
+    assert isinstance(res_A, tuple) and len(res_A) == 3
+    assert all(_r.shape == (2 ** 2,) for _r in res_A)
 
     # Check that they agree up to numeric tolerance
-    assert np.allclose(res_F, res_A, atol=tol, rtol=0)
+    assert all(np.allclose(_rF, _rA, atol=tol, rtol=0) for _rF, _rA in zip(res_F, res_A))
 
 
 @pytest.mark.parametrize("hermitian", [1 / np.sqrt(2) * np.array([[1, 1], [1, -1]])])
diff --git a/tests/test_qnode.py b/tests/test_qnode.py
index f5d8e75cd..3088ce243 100644
--- a/tests/test_qnode.py
+++ b/tests/test_qnode.py
@@ -430,7 +430,9 @@ class TestTapeConstruction:
         assert qn.gradient_kwargs["approx_order"] == 2
 
         jac = qn.gradient_fn(qn)(0.45, 0.1)
-        assert jac.shape == (2, 2, 2)
+        assert isinstance(jac, tuple) and len(jac) == 2
+        assert jac[0].shape == (2, 2)
+        assert jac[1].shape == (2, 2)
 
     def test_returning_non_measurements(self):
         """Test that an exception is raised if a non-measurement
diff --git a/tests/test_quantum_gradients.py b/tests/test_quantum_gradients.py
index 8c749e01b..86fb440b2 100644
--- a/tests/test_quantum_gradients.py
+++ b/tests/test_quantum_gradients.py
@@ -199,8 +199,8 @@ class TestCVGradient:
         grad_A2 = qml.gradients.param_shift_cv(q, dev=gaussian_dev, force_order2=True)(*par)
 
         # the different methods agree
-        assert grad_A == pytest.approx(grad_F, abs=tol)
-        assert grad_A2 == pytest.approx(grad_F, abs=tol)
+        assert qml.math.allclose(grad_A, grad_F, atol=tol, rtol=0)
+        assert qml.math.allclose(grad_A2, grad_F, atol=tol, rtol=0)
 
         # check against the known analytic formula
         r0, phi0, r1, phi1 = par
@@ -214,7 +214,7 @@ class TestCVGradient:
         )
         dn[3] = 0.5 * np.sin(phi0 - phi1) * np.sinh(2 * r0) * np.sinh(2 * r1)
 
-        assert dn[np.newaxis, :] == pytest.approx(grad_F, abs=tol)
+        assert all(qml.math.isclose(dn[i], grad_F[i], atol=tol, rtol=0) for i in range(4))
 
     def test_cv_gradients_repeated_gate_parameters(self, gaussian_dev, tol):
         "Tests that repeated use of a free parameter in a multi-parameter gate yield correct gradients."
@@ -232,8 +232,8 @@ class TestCVGradient:
         grad_A2 = qml.gradients.param_shift_cv(q, dev=gaussian_dev, force_order2=True)(*par)
 
         # the different methods agree
-        assert grad_A == pytest.approx(grad_F, abs=tol)
-        assert grad_A2 == pytest.approx(grad_F, abs=tol)
+        assert qml.math.allclose(grad_A, grad_F, atol=tol, rtol=0)
+        assert qml.math.allclose(grad_A2, grad_F, atol=tol, rtol=0)
 
     def test_cv_gradients_parameters_inside_array(self, gaussian_dev, tol):
         "Tests that free parameters inside an array passed to an Operation yield correct gradients."
@@ -276,8 +276,8 @@ class TestCVGradient:
         grad_A2 = qml.gradients.param_shift_cv(q, dev=gaussian_dev, force_order2=True)(*par)
 
         # the different methods agree
-        assert grad_A == pytest.approx(grad_F, abs=tol)
-        assert grad_A2 == pytest.approx(grad_F, abs=tol)
+        assert qml.math.allclose(grad_A, grad_F, atol=tol, rtol=0)
+        assert qml.math.allclose(grad_A2, grad_F, atol=tol, rtol=0)
 
     def test_CVOperation_with_heisenberg_and_no_params(self, gaussian_dev, tol):
         """An integration test for InterferometerUnitary, a gate that supports analytic differentiation
@@ -457,9 +457,9 @@ class TestQubitGradient:
         grad_auto = grad_fn(*params)
 
         # gradients computed with different methods must agree
-        assert grad_fd1 == pytest.approx(grad_fd2, abs=tol)
-        assert grad_fd1 == pytest.approx(grad_angle, abs=tol)
-        assert np.allclose(grad_fd1, grad_auto, atol=tol, rtol=0)
+        assert qml.math.allclose(grad_fd1, grad_fd2, atol=tol, rtol=0)
+        assert qml.math.allclose(grad_fd1, grad_angle, atol=tol, rtol=0)
+        assert qml.math.allclose(grad_fd1, grad_auto, atol=tol, rtol=0)
 
     def test_hybrid_gradients(self, qubit_device_2_wires, tol):
         "Tests that the various ways of computing the gradient of a hybrid computation all agree."
@@ -533,7 +533,7 @@ class TestQubitGradient:
             "Gradient of classical computed symbolically, can use normal numpy functions."
             val = classical((a, b))
             J = grad_method(quantum)(a, np.log(b))
-            return val * np.array([J[0, 0] + J[1, 0], (J[0, 1] + J[1, 1]) / b])
+            return val * np.array([J[0][0] + J[0][1], (J[1][0] + J[1][1]) / b])
 
         param = np.array([-0.1259, 1.53])
         y0 = classical(param)
@@ -544,9 +544,9 @@ class TestQubitGradient:
         grad_angle = d_classical(*param, qml.gradients.param_shift)
 
         # gradients computed with different methods must agree
-        assert grad_fd1 == pytest.approx(grad_angle, abs=tol)
-        assert grad_fd1 == pytest.approx(grad_auto, abs=tol)
-        assert grad_angle == pytest.approx(grad_auto, abs=tol)
+        assert qml.math.allclose(grad_fd1, grad_angle, atol=tol, rtol=0)
+        assert qml.math.allclose(grad_fd1, grad_auto, atol=tol, rtol=0)
+        assert qml.math.allclose(grad_angle, grad_auto, atol=tol, rtol=0)
 
     def test_qnode_gradient_fanout(self, qubit_device_1_wire, tol):
         "Tests that the correct gradient is computed for qnodes which use the same parameter in multiple gates."
diff --git a/tests/test_qubit_device.py b/tests/test_qubit_device.py
index 830c8fc8e..6a79cbefa 100644
--- a/tests/test_qubit_device.py
+++ b/tests/test_qubit_device.py
@@ -924,7 +924,9 @@ class TestShotList:
 
         # test gradient works
         res = qml.jacobian(circuit)(0.5, 0.1)
-        assert res.shape == (2,) + expected_shape
+        assert isinstance(res, tuple) and len(res) == 2
+        assert res[0].shape == expected_shape
+        assert res[1].shape == expected_shape
 
     shot_data = [
         [[1, 2, 3, 10], [(1, 1), (2, 1), (3, 1), (10, 1)], (4, 4), 16],
diff --git a/tests/transforms/test_adjoint_metric_tensor.py b/tests/transforms/test_adjoint_metric_tensor.py
index 3903dfa1e..0684c9efa 100644
--- a/tests/transforms/test_adjoint_metric_tensor.py
+++ b/tests/transforms/test_adjoint_metric_tensor.py
@@ -403,7 +403,7 @@ fubini_params = [
 ]
 
 
-def autodiff_metric_tensor(ansatz, num_wires, mimic_autograd=False):
+def autodiff_metric_tensor(ansatz, num_wires):
     """Compute the metric tensor by full state vector
     differentiation via autograd."""
     dev = qml.device("default.qubit", wires=num_wires)
@@ -420,13 +420,6 @@ def autodiff_metric_tensor(ansatz, num_wires, mimic_autograd=False):
         rjac = qml.jacobian(rqnode)(*params)
         ijac = qml.jacobian(iqnode)(*params)
 
-        if not mimic_autograd:
-            if len(params) > 1 and all(
-                qml.math.shape(p) == qml.math.shape(params[0]) for p in params[1:]
-            ):
-                rjac = tuple(qml.math.transpose(rjac))
-                ijac = tuple(qml.math.transpose(ijac))
-
         if isinstance(rjac, tuple):
             out = []
             for rc, ic in zip(rjac, ijac):
@@ -440,18 +433,10 @@ def autodiff_metric_tensor(ansatz, num_wires, mimic_autograd=False):
                 )
             return tuple(out)
 
-        if (
-            mimic_autograd
-            and len(params) > 1
-            and all(qml.math.shape(p) == qml.math.shape(params[0]) for p in params[1:])
-        ):
-            jac_contract = [-2]
-        else:
-            jac_contract = [0]
         jac = rjac + 1j * ijac
-        psidpsi = np.tensordot(np.conj(state), jac, axes=([0], jac_contract))
+        psidpsi = np.tensordot(np.conj(state), jac, axes=([0], [0]))
         return np.real(
-            np.tensordot(np.conj(jac), jac, axes=(jac_contract, jac_contract))
+            np.tensordot(np.conj(jac), jac, axes=([0], [0]))
             - np.tensordot(np.conj(psidpsi), psidpsi, axes=0)
         )
 
@@ -469,7 +454,7 @@ class TestAdjointMetricTensorTape:
     def test_correct_output_tape_autograd(self, ansatz, params):
         """Test that the output is correct when using Autograd and
         calling the adjoint metric tensor directly on a tape."""
-        expected = autodiff_metric_tensor(ansatz, self.num_wires, True)(*params)
+        expected = autodiff_metric_tensor(ansatz, self.num_wires)(*params)
         dev = qml.device("default.qubit.autograd", wires=self.num_wires)
 
         @qml.qnode(dev, interface="autograd")
@@ -577,7 +562,7 @@ class TestAdjointMetricTensorQNode:
     def test_correct_output_qnode_autograd(self, ansatz, params):
         """Test that the output is correct when using Autograd and
         calling the adjoint metric tensor on a QNode."""
-        expected = autodiff_metric_tensor(ansatz, self.num_wires, True)(*params)
+        expected = autodiff_metric_tensor(ansatz, self.num_wires)(*params)
         dev = qml.device("default.qubit", wires=self.num_wires)
 
         @qml.qnode(dev, interface="autograd")
@@ -588,8 +573,6 @@ class TestAdjointMetricTensorQNode:
 
         mt = qml.adjoint_metric_tensor(circuit)(*params)
 
-        print(mt)
-        print(expected)
         if isinstance(mt, tuple):
             assert all(qml.math.allclose(_mt, _exp) for _mt, _exp in zip(mt, expected))
         else:
@@ -642,8 +625,6 @@ class TestAdjointMetricTensorQNode:
 
         mt = qml.adjoint_metric_tensor(circuit)(*t_params)
 
-        print(mt)
-        print(expected)
         if isinstance(mt, tuple):
             assert all(qml.math.allclose(_mt, _exp) for _mt, _exp in zip(mt, expected))
         else:
@@ -679,7 +660,7 @@ class TestAdjointMetricTensorQNode:
         ansatz = fubini_ansatz2
         params = fubini_params[2]
 
-        exp_fn = autodiff_metric_tensor(ansatz, self.num_wires, True)
+        exp_fn = autodiff_metric_tensor(ansatz, self.num_wires)
         expected = qml.jacobian(exp_fn)(*params)
         dev = qml.device("default.qubit", wires=self.num_wires)
         dev2 = qml.device("default.qubit.autograd", wires=self.num_wires)
@@ -722,7 +703,7 @@ class TestAdjointMetricTensorDifferentiability:
     def test_autograd(self, ansatz, params):
         """Test that the derivative is correct when using Autograd and
         calling the adjoint metric tensor on a QNode."""
-        exp_fn = autodiff_metric_tensor(ansatz, self.num_wires, True)
+        exp_fn = autodiff_metric_tensor(ansatz, self.num_wires)
         expected = qml.jacobian(exp_fn)(*params)
         dev = qml.device("default.qubit", wires=self.num_wires)
 
diff --git a/tests/transforms/test_classical_jacobian.py b/tests/transforms/test_classical_jacobian.py
index 6ae87e60a..683392128 100644
--- a/tests/transforms/test_classical_jacobian.py
+++ b/tests/transforms/test_classical_jacobian.py
@@ -128,14 +128,10 @@ def test_autograd_without_argnum(circuit, args, expected_jac, diff_method):
     qnode = qml.QNode(circuit, dev, interface="autograd", diff_method=diff_method)
     jac = classical_jacobian(qnode)(*args)
 
-    # NOTE: We use stacking to replicate qml.jacobian behaviour for equal-shaped inputs
     arg_shapes = [qml.math.shape(arg) for arg in args]
     if len(args) == 1:
         # For a single argument, the Jacobian is unpacked
         assert np.allclose(jac, expected_jac[0])
-    elif all(sh == arg_shapes[0] for sh in arg_shapes[1:]):
-        expected_jac = qml.math.stack(expected_jac).T
-        assert np.allclose(jac, expected_jac)
     else:
         assert len(jac) == len(expected_jac)
         for _jac, _expected_jac in zip(jac, expected_jac):
@@ -268,6 +264,7 @@ def test_autograd_with_single_list_argnum(circuit, args, expected_jac, argnum, d
     qnode = qml.QNode(circuit, dev, interface="autograd", diff_method=diff_method)
     jac = classical_jacobian(qnode, argnum=argnum)(*args)
     expected_jac = (expected_jac[argnum[0]],)
+
     assert len(jac) == 1
     assert np.allclose(jac[0], expected_jac[0])
 
diff --git a/tests/transforms/test_decompositions.py b/tests/transforms/test_decompositions.py
index d871943f6..3aa6c7763 100644
--- a/tests/transforms/test_decompositions.py
+++ b/tests/transforms/test_decompositions.py
@@ -116,9 +116,6 @@ class TestQubitUnitaryZYZDecomposition:
         assert obtained_gates[0].wires == Wires("a")
         assert qml.math.allclose(qml.math.unwrap(obtained_gates[0].parameters), expected_params)
 
-        print(qml.math.unwrap(obtained_gates[0].parameters))
-        print(expected_params)
-
         if obtained_gates[0].num_params == 1:
             # With TF and RZ, need to cast since can't just unwrap
             obtained_mat = qml.RZ(obtained_gates[0].parameters[0].numpy(), wires=0).matrix
diff --git a/tests/transforms/test_metric_tensor.py b/tests/transforms/test_metric_tensor.py
index a97d8556f..e8f410e42 100644
--- a/tests/transforms/test_metric_tensor.py
+++ b/tests/transforms/test_metric_tensor.py
@@ -71,7 +71,9 @@ class TestMetricTensor:
 
         params = np.array([0.1, 0.2], requires_grad=True)
         result = qml.metric_tensor(circuit, approx="block-diag")(*params)
-        assert qml.math.shape(result) == (2, 2)
+        assert isinstance(result, tuple) and len(result) == 2
+        assert qml.math.shape(result[0]) == ()
+        assert qml.math.shape(result[1]) == ()
 
     @pytest.mark.parametrize("diff_method", ["parameter-shift", "backprop"])
     def test_parameter_fan_out(self, diff_method):
@@ -207,7 +209,8 @@ class TestMetricTensor:
         block-diagonal and diagonal setting."""
         dev = qml.device("default.qubit", wires=2)
 
-        def circuit(a, b, c):
+        def circuit(abc):
+            a, b, c = abc
             qml.RX(a, wires=0)
             qml.RY(b, wires=0)
             qml.CNOT(wires=[0, 1])
@@ -216,13 +219,12 @@ class TestMetricTensor:
 
         circuit = qml.QNode(circuit, dev)
 
-        a = np.array(0.432, requires_grad=True)
-        b = np.array(0.12, requires_grad=True)
-        c = np.array(-0.432, requires_grad=True)
+        abc = np.array([0.432, 0.12, -0.432], requires_grad=True)
+        a, b, c = abc
 
         # evaluate metric tensor
-        g_diag = qml.metric_tensor(circuit, approx="diag")(a, b, c)
-        g_blockdiag = qml.metric_tensor(circuit, approx="block-diag")(a, b, c)
+        g_diag = qml.metric_tensor(circuit, approx="diag")(abc)
+        g_blockdiag = qml.metric_tensor(circuit, approx="block-diag")(abc)
 
         # check that the metric tensor is correct
         expected = (
@@ -757,7 +759,7 @@ fubini_params = [
 ]
 
 
-def autodiff_metric_tensor(ansatz, num_wires, mimic_autograd=False):
+def autodiff_metric_tensor(ansatz, num_wires):
     """Compute the metric tensor by full state vector
     differentiation via autograd."""
     dev = qml.device("default.qubit", wires=num_wires)
@@ -774,13 +776,6 @@ def autodiff_metric_tensor(ansatz, num_wires, mimic_autograd=False):
         rjac = qml.jacobian(rqnode)(*params)
         ijac = qml.jacobian(iqnode)(*params)
 
-        if not mimic_autograd:
-            if len(params) > 1 and all(
-                qml.math.shape(p) == qml.math.shape(params[0]) for p in params[1:]
-            ):
-                rjac = tuple(qml.math.transpose(rjac))
-                ijac = tuple(qml.math.transpose(ijac))
-
         if isinstance(rjac, tuple):
             out = []
             for rc, ic in zip(rjac, ijac):
@@ -794,18 +789,10 @@ def autodiff_metric_tensor(ansatz, num_wires, mimic_autograd=False):
                 )
             return tuple(out)
 
-        if (
-            mimic_autograd
-            and len(params) > 1
-            and all(qml.math.shape(p) == qml.math.shape(params[0]) for p in params[1:])
-        ):
-            jac_contract = [-2]
-        else:
-            jac_contract = [0]
         jac = rjac + 1j * ijac
-        psidpsi = np.tensordot(np.conj(state), jac, axes=([0], jac_contract))
+        psidpsi = np.tensordot(np.conj(state), jac, axes=([0], [0]))
         return np.real(
-            np.tensordot(np.conj(jac), jac, axes=(jac_contract, jac_contract))
+            np.tensordot(np.conj(jac), jac, axes=([0], [0]))
             - np.tensordot(np.conj(psidpsi), psidpsi, axes=0)
         )
 
@@ -818,7 +805,7 @@ class TestFullMetricTensor:
 
     @pytest.mark.parametrize("ansatz, params", zip(fubini_ansatze, fubini_params))
     def test_correct_output_autograd(self, ansatz, params):
-        expected = autodiff_metric_tensor(ansatz, self.num_wires, True)(*params)
+        expected = autodiff_metric_tensor(ansatz, self.num_wires)(*params)
         dev = qml.device("default.qubit.autograd", wires=self.num_wires + 1)
 
         @qml.qnode(dev, interface="autograd")
@@ -1020,7 +1007,7 @@ class TestDifferentiability:
         def cost_full(*weights):
             return np.array(qml.metric_tensor(qnode, approx=None)(*weights))
 
-        _cost_full = lambda *weights: np.array(autodiff_metric_tensor(ansatz, 3, True)(*weights))
+        _cost_full = lambda *weights: np.array(autodiff_metric_tensor(ansatz, 3)(*weights))
         _c = _cost_full(*weights)
         c = cost_full(*weights)
         assert all(
diff --git a/tests/transforms/test_optimization/test_commute_controlled.py b/tests/transforms/test_optimization/test_commute_controlled.py
index 0c7b2eb94..2f287e915 100644
--- a/tests/transforms/test_optimization/test_commute_controlled.py
+++ b/tests/transforms/test_optimization/test_commute_controlled.py
@@ -53,7 +53,6 @@ class TestCommuteControlled:
         transformed_qfunc = commute_controlled(direction=direction)(qfunc)
 
         ops = qml.transforms.make_tape(transformed_qfunc)().operations
-        print(ops)
 
         names_expected = ["PauliX", "ControlledQubitUnitary", "PauliX"]
         wires_expected = [Wires(2), Wires([0, 2]), Wires(2)]
@@ -71,7 +70,6 @@ class TestCommuteControlled:
         transformed_qfunc = commute_controlled(direction=direction)(qfunc)
 
         ops = qml.transforms.make_tape(transformed_qfunc)().operations
-        print(ops)
 
         names_expected = ["PauliZ", "CNOT", "PauliY"]
         wires_expected = [Wires("b"), Wires([2, "b"]), Wires("b")]
diff --git a/tests/transforms/test_qfunc_transform.py b/tests/transforms/test_qfunc_transform.py
index d8943e370..7fbccb5b8 100644
--- a/tests/transforms/test_qfunc_transform.py
+++ b/tests/transforms/test_qfunc_transform.py
@@ -419,5 +419,4 @@ class TestQFuncTransformGradients:
 
         grad = jax.grad(qnode, argnums=[0, 1, 2])(x, a, b)
         expected = qml.grad(self.expval)(np.array(x), np.array(a), np.array(b))
-        print(grad, expected)
         assert all(np.allclose(g, e) for g, e in zip(grad, expected))
diff --git a/tests/transforms/test_tape_expand.py b/tests/transforms/test_tape_expand.py
index 7f5e1d61d..99631fb3a 100644
--- a/tests/transforms/test_tape_expand.py
+++ b/tests/transforms/test_tape_expand.py
@@ -493,7 +493,6 @@ class TestCreateCustomDecompExpandFn:
         assert np.allclose(original_grad, decomp_grad)
 
         expected_ops = ["Hadamard", "RZ", "PauliX", "RY", "PauliX", "RZ", "Hadamard"]
-        print(decomp_qnode.qtape.operations)
         assert all(
             [op.name == name for op, name in zip(decomp_qnode.qtape.operations, expected_ops)]
         )