Repository: PennyLaneAI/pennylane

Issue:
[BUG] output shape and type of `gradient_transform` applied to QNodes depend on classical processing
### Expected behavior

When applying a `gradient_transform` to a QNode, the output shape and type should not differ whether or not there is classical processing.

### Actual behavior

If there is no classical processing (and other conditions apply, such as "there is only one QNode argument"), the output shape and type differs from cases where there is classical processing:

```python
import pennylane as qml
from jax import numpy as np
from pennylane import numpy as np

dev = qml.device("default.qubit")

@qml.qnode(dev)
def circuit(w, factor=1.):
    qml.Rot(*(factor * w), 0)
    return qml.expval(qml.PauliZ(0))

w = np.array([0.612, 0.42, 1.2])

print(qml.gradients.hadamard_grad(circuit)(w, factor=1.))
# (tensor(-0., requires_grad=True), tensor(-0.40776045, requires_grad=True), tensor(-0., requires_grad=True))
print(qml.gradients.hadamard_grad(circuit)(w, factor=2.))
# [ 0.         -1.48928624  0.        ]
```

Another example, this time causing a transpose, is

```python
@qml.qnode(dev)
def circuit(w, factor=1.):
    qml.Rot(*(factor * w), 0)
    return qml.probs([0, 1])

w = np.array([0.612, 0.42, 1.2])

print(qml.gradients.hadamard_grad(circuit)(w, factor=1.))
# (tensor([0., 0., 0., 0.], requires_grad=True), tensor([-0.20388023,  0.        ,  0.20388023,  0.        ], requires_grad=True), tensor([0., 0., 0., 0.], requires_grad=True))
# -> Could be interpreted as shape (3, 4)
print(qml.gradients.hadamard_grad(circuit)(w, factor=2.))
# [[ 0.         -0.74464312  0.        ]
#  [ 0.          0.          0.        ]
#  [ 0.          0.74464312  0.        ]
#  [ 0.          0.          0.        ]]
# Has shape (4, 3)
```

### Additional information

This is due to a special branch in `_contract_qjac_with_cjac` in `gradients/gradient_transform.py`, which is activated whenever the classical Jacobian is the identity.
While it is correct that no numerics have to be performed in this case, the contraction that would be performed without this branch (and yields the correct behaviour) also leads to transposing the output. As the contraction is skipped, the transpose is as well, leading to the observed behaviour.

### Source code

_No response_

### Tracebacks

_No response_

### System information

```shell
pl dev; master
```


### Existing GitHub issues

- [X] I have searched existing GitHub issues to make sure the issue does not already exist.


Discussion:

Resolving patch:
diff --git a/pennylane/gradients/gradient_transform.py b/pennylane/gradients/gradient_transform.py
--- a/pennylane/gradients/gradient_transform.py
+++ b/pennylane/gradients/gradient_transform.py
@@ -379,6 +379,10 @@ def reorder_grads(grads, tape_specs):
     return _move_first_axis_to_third_pos(grads, num_params, shots.num_copies, num_measurements)
 
 
+tdot = partial(qml.math.tensordot, axes=[[0], [0]])
+stack = qml.math.stack
+
+
 # pylint: disable=too-many-return-statements,too-many-branches
 def _contract_qjac_with_cjac(qjac, cjac, tape):
     """Contract a quantum Jacobian with a classical preprocessing Jacobian.
@@ -397,52 +401,58 @@ def _contract_qjac_with_cjac(qjac, cjac, tape):
         cjac = cjac[0]
 
     cjac_is_tuple = isinstance(cjac, tuple)
-    if not cjac_is_tuple:
-        is_square = cjac.ndim == 2 and cjac.shape[0] == cjac.shape[1]
-
-        if not qml.math.is_abstract(cjac) and (
-            is_square and qml.math.allclose(cjac, qml.numpy.eye(cjac.shape[0]))
-        ):
-            # Classical Jacobian is the identity. No classical processing is present in the QNode
-            return qjac
 
     multi_meas = num_measurements > 1
 
+    # This block only figures out whether there is a single tape parameter or not
     if cjac_is_tuple:
-        multi_params = True
+        single_tape_param = False
     else:
+        # Peel out a single measurement's and single shot setting's qjac
         _qjac = qjac
         if multi_meas:
             _qjac = _qjac[0]
         if has_partitioned_shots:
             _qjac = _qjac[0]
-        multi_params = isinstance(_qjac, tuple)
-
-    tdot = partial(qml.math.tensordot, axes=[[0], [0]])
+        single_tape_param = not isinstance(_qjac, tuple)
 
-    if not multi_params:
+    if single_tape_param:
         # Without dimension (e.g. expval) or with dimension (e.g. probs)
         def _reshape(x):
             return qml.math.reshape(x, (1,) if x.shape == () else (1, -1))
 
         if not (multi_meas or has_partitioned_shots):
-            # Single parameter, single measurements
+            # Single parameter, single measurements, no shot vector
             return tdot(_reshape(qjac), cjac)
 
         if not (multi_meas and has_partitioned_shots):
+            # Single parameter, multiple measurements or shot vector, but not both
             return tuple(tdot(_reshape(q), cjac) for q in qjac)
 
-        # Single parameter, multiple measurements
+        # Single parameter, multiple measurements, and shot vector
         return tuple(tuple(tdot(_reshape(_q), cjac) for _q in q) for q in qjac)
 
     if not multi_meas:
         # Multiple parameters, single measurement
-        qjac = qml.math.stack(qjac)
+        qjac = stack(qjac)
         if not cjac_is_tuple:
-            return tdot(qjac, qml.math.stack(cjac))
+            cjac = stack(cjac)
+            if has_partitioned_shots:
+                return tuple(tdot(stack(q), cjac) for q in qjac)
+            return tdot(qjac, cjac)
+        if has_partitioned_shots:
+            return tuple(tuple(tdot(q, c) for c in cjac if c is not None) for q in qjac)
         return tuple(tdot(qjac, c) for c in cjac if c is not None)
 
     # Multiple parameters, multiple measurements
     if not cjac_is_tuple:
-        return tuple(tdot(qml.math.stack(q), qml.math.stack(cjac)) for q in qjac)
-    return tuple(tuple(tdot(qml.math.stack(q), c) for c in cjac if c is not None) for q in qjac)
+        cjac = stack(cjac)
+        if has_partitioned_shots:
+            return tuple(tuple(tdot(stack(_q), cjac) for _q in q) for q in qjac)
+        return tuple(tdot(stack(q), cjac) for q in qjac)
+    if has_partitioned_shots:
+        return tuple(
+            tuple(tuple(tdot(stack(_q), c) for c in cjac if c is not None) for _q in q)
+            for q in qjac
+        )
+    return tuple(tuple(tdot(stack(q), c) for c in cjac if c is not None) for q in qjac)


Test patch:
diff --git a/tests/gradients/core/test_gradient_transform.py b/tests/gradients/core/test_gradient_transform.py
--- a/tests/gradients/core/test_gradient_transform.py
+++ b/tests/gradients/core/test_gradient_transform.py
@@ -186,7 +186,8 @@ class TestGradientTransformIntegration:
 
     @pytest.mark.parametrize("shots, atol", [(None, 1e-6), (1000, 1e-1), ([1000, 500], 3e-1)])
     @pytest.mark.parametrize("slicing", [False, True])
-    def test_acting_on_qnodes_single_param(self, shots, slicing, atol):
+    @pytest.mark.parametrize("prefactor", [1.0, 2.0])
+    def test_acting_on_qnodes_single_param(self, shots, slicing, prefactor, atol):
         """Test that a gradient transform acts on QNodes with a single parameter correctly"""
         np.random.seed(412)
         dev = qml.device("default.qubit", wires=2, shots=shots)
@@ -194,9 +195,9 @@ def test_acting_on_qnodes_single_param(self, shots, slicing, atol):
         @qml.qnode(dev)
         def circuit(weights):
             if slicing:
-                qml.RX(weights[0], wires=[0])
+                qml.RX(prefactor * weights[0], wires=[0])
             else:
-                qml.RX(weights, wires=[0])
+                qml.RX(prefactor * weights, wires=[0])
             return qml.expval(qml.PauliZ(0)), qml.var(qml.PauliX(1))
 
         grad_fn = qml.gradients.param_shift(circuit)
@@ -204,14 +205,18 @@ def circuit(weights):
         w = np.array([0.543] if slicing else 0.543, requires_grad=True)
         res = grad_fn(w)
         assert circuit.interface == "auto"
-        expected = np.array([-np.sin(w[0] if slicing else w), 0])
+
+        # Need to multiply 0 with w to get the right output shape for non-scalar w
+        expected = (-prefactor * np.sin(prefactor * w), w * 0)
+
         if isinstance(shots, list):
             assert all(np.allclose(r, expected, atol=atol, rtol=0) for r in res)
         else:
             assert np.allclose(res, expected, atol=atol, rtol=0)
 
     @pytest.mark.parametrize("shots, atol", [(None, 1e-6), (1000, 1e-1), ([1000, 100], 2e-1)])
-    def test_acting_on_qnodes_multi_param(self, shots, atol):
+    @pytest.mark.parametrize("prefactor", [1.0, 2.0])
+    def test_acting_on_qnodes_multi_param(self, shots, prefactor, atol):
         """Test that a gradient transform acts on QNodes with multiple parameters correctly"""
         np.random.seed(412)
         dev = qml.device("default.qubit", wires=2, shots=shots)
@@ -219,9 +224,9 @@ def test_acting_on_qnodes_multi_param(self, shots, atol):
         @qml.qnode(dev)
         def circuit(weights):
             qml.RX(weights[0], wires=[0])
-            qml.RY(weights[1], wires=[1])
+            qml.RY(prefactor * weights[1], wires=[1])
             qml.CNOT(wires=[0, 1])
-            return qml.expval(qml.PauliZ(0)), qml.var(qml.PauliX(1))
+            return qml.expval(qml.PauliZ(0)), qml.var(qml.PauliZ(1))
 
         grad_fn = qml.gradients.param_shift(circuit)
 
@@ -229,7 +234,16 @@ def circuit(weights):
         res = grad_fn(w)
         assert circuit.interface == "auto"
         x, y = w
-        expected = np.array([[-np.sin(x), 0], [0, -2 * np.cos(y) * np.sin(y)]])
+        y *= prefactor
+        expected = np.array(
+            [
+                [-np.sin(x), 0],
+                [
+                    2 * np.cos(x) * np.sin(x) * np.cos(y) ** 2,
+                    2 * prefactor * np.cos(y) * np.sin(y) * np.cos(x) ** 2,
+                ],
+            ]
+        )
         if isinstance(shots, list):
             assert all(np.allclose(r, expected, atol=atol, rtol=0) for r in res)
         else:
diff --git a/tests/gradients/core/test_hadamard_gradient.py b/tests/gradients/core/test_hadamard_gradient.py
--- a/tests/gradients/core/test_hadamard_gradient.py
+++ b/tests/gradients/core/test_hadamard_gradient.py
@@ -29,7 +29,7 @@ def grad_fn(tape, dev, fn=qml.gradients.hadamard_grad, **kwargs):
 
 def cost1(x):
     """Cost function."""
-    qml.Rot(*x, wires=0)
+    qml.Rot(x[0], 0.3 * x[1], x[2], wires=0)
     return qml.expval(qml.PauliZ(0))
 
 
@@ -42,7 +42,7 @@ def cost2(x):
 def cost3(x):
     """Cost function."""
     qml.Rot(*x, wires=0)
-    return [qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliZ(1))]
+    return (qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliZ(1)))
 
 
 def cost4(x):
@@ -60,7 +60,25 @@ def cost5(x):
 def cost6(x):
     """Cost function."""
     qml.Rot(*x, wires=0)
-    return [qml.probs([0, 1]), qml.probs([2, 3])]
+    return (qml.probs([0, 1]), qml.probs([2, 3]))
+
+
+def cost7(x):
+    """Cost function."""
+    qml.RX(x, 0)
+    return qml.expval(qml.PauliZ(0))
+
+
+def cost8(x):
+    """Cost function."""
+    qml.RX(x, 0)
+    return [qml.expval(qml.PauliZ(0))]
+
+
+def cost9(x):
+    """Cost function."""
+    qml.RX(x, 0)
+    return (qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliZ(1)))
 
 
 class TestHadamardGrad:
@@ -480,15 +498,36 @@ def test_prob_expectation_values(self, tol):
         assert np.allclose(res_hadamard[1][0], probs_expected[:, 0], tol)
         assert np.allclose(res_hadamard[1][1], probs_expected[:, 1], tol)
 
-    costs_and_expected_expval = [
-        (cost1, [3]),
-        (cost2, [3]),
-        (cost3, [2, 3]),
+    costs_and_expected_expval_scalar = [
+        (cost7, (), np.ndarray),
+        (cost8, (1,), list),
+        (cost9, (2,), tuple),
     ]
 
-    @pytest.mark.parametrize("cost, expected_shape", costs_and_expected_expval)
-    def test_output_shape_matches_qnode_expval(self, cost, expected_shape):
-        """Test that the transform output shape matches that of the QNode."""
+    @pytest.mark.parametrize("cost, exp_shape, exp_type", costs_and_expected_expval_scalar)
+    def test_output_shape_matches_qnode_expval_scalar(self, cost, exp_shape, exp_type):
+        """Test that the transform output shape matches that of the QNode for
+        expectation values and a scalar parameter."""
+        dev = qml.device("default.qubit", wires=4)
+
+        x = np.array(0.419)
+        circuit = qml.QNode(cost, dev)
+
+        res_hadamard = qml.gradients.hadamard_grad(circuit)(x)
+
+        assert isinstance(res_hadamard, exp_type)
+        assert np.array(res_hadamard).shape == exp_shape
+
+    costs_and_expected_expval_array = [
+        (cost1, [3], np.ndarray),
+        (cost2, [3], list),
+        (cost3, [2, 3], tuple),
+    ]
+
+    @pytest.mark.parametrize("cost, exp_shape, exp_type", costs_and_expected_expval_array)
+    def test_output_shape_matches_qnode_expval_array(self, cost, exp_shape, exp_type):
+        """Test that the transform output shape matches that of the QNode for
+        expectation values and an array-valued parameter."""
         dev = qml.device("default.qubit", wires=4)
 
         x = np.random.rand(3)
@@ -496,24 +535,24 @@ def test_output_shape_matches_qnode_expval(self, cost, expected_shape):
 
         res_hadamard = qml.gradients.hadamard_grad(circuit)(x)
 
-        assert isinstance(res_hadamard, (tuple, list))
+        assert isinstance(res_hadamard, exp_type)
         if len(res_hadamard) == 1:
             res_hadamard = res_hadamard[0]
-        assert len(res_hadamard) == expected_shape[0]
+        assert len(res_hadamard) == exp_shape[0]
 
-        if len(expected_shape) > 1:
+        if len(exp_shape) > 1:
             for r in res_hadamard:
-                assert isinstance(r, tuple)
-                assert len(r) == expected_shape[1]
+                assert isinstance(r, np.ndarray)
+                assert len(r) == exp_shape[1]
 
     costs_and_expected_probs = [
-        (cost4, [3, 4]),
-        (cost5, [3, 4]),
-        (cost6, [2, 3, 4]),
+        (cost4, [4, 3], np.ndarray),
+        (cost5, [4, 3], list),
+        (cost6, [2, 4, 3], tuple),
     ]
 
-    @pytest.mark.parametrize("cost, expected_shape", costs_and_expected_probs)
-    def test_output_shape_matches_qnode_probs(self, cost, expected_shape):
+    @pytest.mark.parametrize("cost, exp_shape, exp_type", costs_and_expected_probs)
+    def test_output_shape_matches_qnode_probs(self, cost, exp_shape, exp_type):
         """Test that the transform output shape matches that of the QNode."""
         dev = qml.device("default.qubit", wires=5)
 
@@ -521,24 +560,24 @@ def test_output_shape_matches_qnode_probs(self, cost, expected_shape):
         circuit = qml.QNode(cost, dev)
 
         res_hadamard = qml.gradients.hadamard_grad(circuit)(x)
-        assert isinstance(res_hadamard, (tuple, list))
+        assert isinstance(res_hadamard, exp_type)
         if len(res_hadamard) == 1:
             res_hadamard = res_hadamard[0]
-        assert len(res_hadamard) == expected_shape[0]
+        assert len(res_hadamard) == exp_shape[0]
 
-        if len(expected_shape) > 2:
+        if len(exp_shape) > 2:
             for r in res_hadamard:
-                assert isinstance(r, tuple)
-                assert len(r) == expected_shape[1]
+                assert isinstance(r, np.ndarray)
+                assert len(r) == exp_shape[1]
 
                 for r_ in r:
-                    assert isinstance(r_, qml.numpy.ndarray)
-                    assert len(r_) == expected_shape[2]
+                    assert isinstance(r_, np.ndarray)
+                    assert len(r_) == exp_shape[2]
 
-        elif len(expected_shape) > 1:
+        elif len(exp_shape) > 1:
             for r in res_hadamard:
-                assert isinstance(r, qml.numpy.ndarray)
-                assert len(r) == expected_shape[1]
+                assert isinstance(r, np.ndarray)
+                assert len(r) == exp_shape[1]
 
     @pytest.mark.parametrize("shots", [None, 100])
     def test_shots_attribute(self, shots):
@@ -997,35 +1036,24 @@ def test_all_zero_diff_methods_multiple_returns_tape(self):
         assert res_hadamard[1][2].shape == (4,)
         assert np.allclose(res_hadamard[1][2], 0)
 
-    def test_all_zero_diff_methods(self):
+    @pytest.mark.parametrize("prefactor", [1.0, 2.0])
+    def test_all_zero_diff_methods(self, prefactor):
         """Test that the transform works correctly when the diff method for every parameter is
         identified to be 0, and that no tapes were generated."""
         dev = qml.device("default.qubit", wires=4)
 
         @qml.qnode(dev)
         def circuit(params):
-            qml.Rot(*params, wires=0)
+            qml.Rot(*(prefactor * params), wires=0)
             return qml.probs([2, 3])
 
         params = np.array([0.5, 0.5, 0.5], requires_grad=True)
 
         result = qml.gradients.hadamard_grad(circuit)(params)
 
-        assert isinstance(result, tuple)
-
-        assert len(result) == 3
-
-        assert isinstance(result[0], np.ndarray)
-        assert result[0].shape == (4,)
-        assert np.allclose(result[0], 0)
-
-        assert isinstance(result[1], np.ndarray)
-        assert result[1].shape == (4,)
-        assert np.allclose(result[1], 0)
-
-        assert isinstance(result[2], np.ndarray)
-        assert result[2].shape == (4,)
-        assert np.allclose(result[2], 0)
+        assert isinstance(result, np.ndarray)
+        assert result.shape == (4, 3)
+        assert np.allclose(result, 0)
 
 
 class TestHadamardTestGradDiff:
@@ -1303,7 +1331,6 @@ def circuit(x, y):
             qml.RX(x[0], wires=[0])
             qml.RY(x[1], wires=[1])
             qml.RY(y, wires=[1])
-            qml.CNOT(wires=[0, 1])
             return qml.expval(qml.PauliZ(0) @ qml.PauliX(1))
 
         x = jax.numpy.array([0.543, -0.654])
@@ -1314,12 +1341,11 @@ def circuit(x, y):
         )
         res_expected = jax.hessian(circuit, argnums=argnums)(x, y)
 
-        if argnums == [0]:
-            assert np.allclose(res[0][0], res_expected[0][0][0])
-            assert np.allclose(res[1][0], res_expected[0][0][1])
+        if len(argnums) == 1:
+            # jax.hessian produces an additional tuple axis, which we have to index away here
+            assert np.allclose(res, res_expected[0])
         else:
-            if len(argnums) != 1:
-                res = res[0]
-
-            for r, r_e in zip(res, res_expected[0]):
-                assert np.allclose(r, r_e)
+            # The Hessian is a 2x2 nested tuple "matrix" for argnums=[0, 1]
+            for r, r_e in zip(res, res_expected):
+                for r_, r_e_ in zip(r, r_e):
+                    assert np.allclose(r_, r_e_)
diff --git a/tests/gradients/finite_diff/test_finite_difference.py b/tests/gradients/finite_diff/test_finite_difference.py
--- a/tests/gradients/finite_diff/test_finite_difference.py
+++ b/tests/gradients/finite_diff/test_finite_difference.py
@@ -294,35 +294,24 @@ def circuit(weights):
         with pytest.raises(qml.QuantumFunctionError, match="No trainable parameters."):
             qml.gradients.finite_diff(circuit)(weights)
 
-    def test_all_zero_diff_methods(self):
+    @pytest.mark.parametrize("prefactor", [1.0, 2.0])
+    def test_all_zero_diff_methods(self, prefactor):
         """Test that the transform works correctly when the diff method for every parameter is
         identified to be 0, and that no tapes were generated."""
         dev = qml.device("default.qubit", wires=4)
 
         @qml.qnode(dev)
         def circuit(params):
-            qml.Rot(*params, wires=0)
+            qml.Rot(*(prefactor * params), wires=0)
             return qml.probs([2, 3])
 
         params = np.array([0.5, 0.5, 0.5], requires_grad=True)
 
         result = qml.gradients.finite_diff(circuit)(params)
 
-        assert isinstance(result, tuple)
-
-        assert len(result) == 3
-
-        assert isinstance(result[0], numpy.ndarray)
-        assert result[0].shape == (4,)
-        assert np.allclose(result[0], 0)
-
-        assert isinstance(result[1], numpy.ndarray)
-        assert result[1].shape == (4,)
-        assert np.allclose(result[1], 0)
-
-        assert isinstance(result[2], numpy.ndarray)
-        assert result[2].shape == (4,)
-        assert np.allclose(result[2], 0)
+        assert isinstance(result, np.ndarray)
+        assert result.shape == (4, 3)
+        assert np.allclose(result, 0)
 
     def test_all_zero_diff_methods_multiple_returns(self):
         """Test that the transform works correctly when the diff method for every parameter is
@@ -343,35 +332,10 @@ def circuit(params):
 
         assert len(result) == 2
 
-        # First elem
-        assert len(result[0]) == 3
-
-        assert isinstance(result[0][0], numpy.ndarray)
-        assert result[0][0].shape == ()
-        assert np.allclose(result[0][0], 0)
-
-        assert isinstance(result[0][1], numpy.ndarray)
-        assert result[0][1].shape == ()
-        assert np.allclose(result[0][1], 0)
-
-        assert isinstance(result[0][2], numpy.ndarray)
-        assert result[0][2].shape == ()
-        assert np.allclose(result[0][2], 0)
-
-        # Second elem
-        assert len(result[0]) == 3
-
-        assert isinstance(result[1][0], numpy.ndarray)
-        assert result[1][0].shape == (4,)
-        assert np.allclose(result[1][0], 0)
-
-        assert isinstance(result[1][1], numpy.ndarray)
-        assert result[1][1].shape == (4,)
-        assert np.allclose(result[1][1], 0)
-
-        assert isinstance(result[1][2], numpy.ndarray)
-        assert result[1][2].shape == (4,)
-        assert np.allclose(result[1][2], 0)
+        for r, exp_shape in zip(result, [(3,), (4, 3)]):
+            assert isinstance(r, np.ndarray)
+            assert r.shape == exp_shape
+            assert np.allclose(r, 0)
 
     def test_y0(self):
         """Test that if first order finite differences is used, then
@@ -440,7 +404,7 @@ def test_output_shape_matches_qnode(self):
         dev = qml.device("default.qubit", wires=4)
 
         def cost1(x):
-            qml.Rot(*x, wires=0)
+            qml.Rot(x[0], 0.3 * x[1], x[2], wires=0)
             return qml.expval(qml.PauliZ(0))
 
         def cost2(x):
@@ -468,18 +432,15 @@ def cost6(x):
 
         transform = [qml.math.shape(qml.gradients.finite_diff(c)(x)) for c in circuits]
 
-        expected = [
+        expected_shapes = [
             (3,),
-            (
-                1,
-                3,
-            ),
+            (1, 3),
             (2, 3),
-            (3, 4),
-            (1, 3, 4),
-            (2, 3, 4),
+            (4, 3),
+            (1, 4, 3),
+            (2, 4, 3),
         ]
-        assert all(t == q for t, q in zip(transform, expected))
+        assert all(t == q for t, q in zip(transform, expected_shapes))
 
     def test_special_observable_qnode_differentiation(self):
         """Test differentiation of a QNode on a device supporting a
@@ -1209,7 +1170,6 @@ def circuit(x, y):
             qml.RX(x[0], wires=[0])
             qml.RY(x[1], wires=[1])
             qml.RY(y, wires=[1])
-            qml.CNOT(wires=[0, 1])
             return qml.var(qml.PauliZ(0) @ qml.PauliX(1))
 
         x = jax.numpy.array([0.543, -0.654])
@@ -1225,12 +1185,11 @@ def circuit(x, y):
 
         tol = 10e-6
 
-        if argnums == [0]:
-            assert np.allclose(res[0][0], res_expected[0][0][0], atol=tol)
-            assert np.allclose(res[1][0], res_expected[0][0][1], atol=tol)
+        if len(argnums) == 1:
+            # jax.hessian produces an additional tuple axis, which we have to index away here
+            assert np.allclose(res, res_expected[0], atol=tol)
         else:
-            if len(argnums) != 1:
-                res = res[0]
-
-            for r, r_e in zip(res, res_expected[0]):
-                assert np.allclose(r, r_e, atol=tol)
+            # The Hessian is a 2x2 nested tuple "matrix" for argnums=[0, 1]
+            for r, r_e in zip(res, res_expected):
+                for r_, r_e_ in zip(r, r_e):
+                    assert np.allclose(r_, r_e_, atol=tol)
diff --git a/tests/gradients/finite_diff/test_finite_difference_shot_vec.py b/tests/gradients/finite_diff/test_finite_difference_shot_vec.py
--- a/tests/gradients/finite_diff/test_finite_difference_shot_vec.py
+++ b/tests/gradients/finite_diff/test_finite_difference_shot_vec.py
@@ -226,24 +226,14 @@ def circuit(params):
         params = np.array([0.5, 0.5, 0.5], requires_grad=True)
 
         all_result = qml.gradients.finite_diff(circuit, h=h_val)(params)
+
+        assert isinstance(all_result, tuple)
         assert len(all_result) == len(default_shot_vector)
 
         for result in all_result:
-            assert isinstance(result, tuple)
-
-            assert len(result) == 3
-
-            assert isinstance(result[0], numpy.ndarray)
-            assert result[0].shape == (4,)
-            assert np.allclose(result[0], 0)
-
-            assert isinstance(result[1], numpy.ndarray)
-            assert result[1].shape == (4,)
-            assert np.allclose(result[1], 0)
-
-            assert isinstance(result[2], numpy.ndarray)
-            assert result[2].shape == (4,)
-            assert np.allclose(result[2], 0)
+            assert isinstance(result, np.ndarray)
+            assert result.shape == (4, 3)
+            assert np.allclose(result, 0)
 
     def test_all_zero_diff_methods_multiple_returns(self):
         """Test that the transform works correctly when the diff method for every parameter is
@@ -259,42 +249,16 @@ def circuit(params):
         params = np.array([0.5, 0.5, 0.5], requires_grad=True)
 
         all_result = qml.gradients.finite_diff(circuit, h=h_val)(params)
+        assert isinstance(all_result, tuple)
         assert len(all_result) == len(many_shots_shot_vector)
 
-        for result in all_result:
-            assert isinstance(result, tuple)
-
-            assert len(result) == 2
-
-            # First elem
-            assert len(result[0]) == 3
-
-            assert isinstance(result[0][0], numpy.ndarray)
-            assert result[0][0].shape == ()
-            assert np.allclose(result[0][0], 0)
-
-            assert isinstance(result[0][1], numpy.ndarray)
-            assert result[0][1].shape == ()
-            assert np.allclose(result[0][1], 0)
-
-            assert isinstance(result[0][2], numpy.ndarray)
-            assert result[0][2].shape == ()
-            assert np.allclose(result[0][2], 0)
-
-            # Second elem
-            assert len(result[0]) == 3
-
-            assert isinstance(result[1][0], numpy.ndarray)
-            assert result[1][0].shape == (4,)
-            assert np.allclose(result[1][0], 0)
-
-            assert isinstance(result[1][1], numpy.ndarray)
-            assert result[1][1].shape == (4,)
-            assert np.allclose(result[1][1], 0)
-
-            assert isinstance(result[1][2], numpy.ndarray)
-            assert result[1][2].shape == (4,)
-            assert np.allclose(result[1][2], 0)
+        for res in all_result:
+            assert isinstance(res, tuple)
+            assert len(res) == 2
+            for r, exp_shape in zip(res, [(3,), (4, 3)]):
+                assert isinstance(r, np.ndarray)
+                assert r.shape == exp_shape
+                assert np.allclose(r, 0)
 
     def test_y0(self):
         """Test that if first order finite differences is used, then
@@ -366,10 +330,10 @@ def test_independent_parameters(self):
 
     def test_output_shape_matches_qnode(self):
         """Test that the transform output shape matches that of the QNode."""
-        dev = qml.device("default.qubit", wires=4, shots=many_shots_shot_vector)
+        dev = qml.device("default.qubit", wires=4, shots=(1, 1000, 1001))
 
         def cost1(x):
-            qml.Rot(*x, wires=0)
+            qml.Rot(x[0], 0.3 * x[1], x[2], wires=0)
             return qml.expval(qml.PauliZ(0))
 
         def cost2(x):
@@ -385,7 +349,7 @@ def cost4(x):
             return qml.probs([0, 1])
 
         def cost5(x):
-            qml.Rot(*x, wires=0)
+            qml.Rot(x[0], 0.3 * x[1], x[2], wires=0)
             return [qml.probs([0, 1])]
 
         def cost6(x):
@@ -396,7 +360,31 @@ def cost6(x):
         circuits = [qml.QNode(cost, dev) for cost in (cost1, cost2, cost3, cost4, cost5, cost6)]
 
         transform = [qml.math.shape(qml.gradients.finite_diff(c, h=h_val)(x)) for c in circuits]
-        expected = [(3, 3), (3, 1, 3), (3, 2, 3), (3, 3, 4), (3, 1, 3, 4), (3, 2, 3, 4)]
+        expected = [(3, 3), (3, 1, 3), (3, 2, 3), (3, 4, 3), (3, 1, 4, 3), (3, 2, 4, 3)]
+        assert all(t == q for t, q in zip(transform, expected))
+
+    def test_output_shape_matches_qnode_two_args(self):
+        """Test that the transform output shape matches that of a QNode with multiple args."""
+        dev = qml.device("default.qubit", wires=4, shots=(1, 1000, 1001))
+
+        def cost1(x, y, z):
+            qml.Rot(x[0], 2 * y[1], -0.1 * z[0], wires=0)
+            return qml.expval(qml.PauliZ(0))
+
+        def cost2(x, y, z):
+            qml.Rot(x[0], 2 * y[1], -0.1 * z[0], wires=0)
+            return qml.probs([0, 1]), qml.probs([2, 3])
+
+        x, y, z = np.random.rand(3, 2)
+        circuits = [qml.QNode(cost, dev) for cost in (cost1, cost2)]
+
+        transform = [
+            qml.math.shape(qml.gradients.finite_diff(c, h=h_val)(x, y, z)) for c in circuits
+        ]
+        expected = [
+            (3, 3, 2),  # shot vector, params, param shape
+            (3, 2, 3, 4, 2),  # shot vector, measurements, params, measurement shape, param shape
+        ]
         assert all(t == q for t, q in zip(transform, expected))
 
     def test_special_observable_qnode_differentiation(self):
diff --git a/tests/gradients/finite_diff/test_spsa_gradient.py b/tests/gradients/finite_diff/test_spsa_gradient.py
--- a/tests/gradients/finite_diff/test_spsa_gradient.py
+++ b/tests/gradients/finite_diff/test_spsa_gradient.py
@@ -387,21 +387,9 @@ def circuit(params):
 
         result = spsa_grad(circuit)(params)
 
-        assert isinstance(result, tuple)
-
-        assert len(result) == 3
-
-        assert isinstance(result[0], numpy.ndarray)
-        assert result[0].shape == (4,)
-        assert np.allclose(result[0], 0)
-
-        assert isinstance(result[1], numpy.ndarray)
-        assert result[1].shape == (4,)
-        assert np.allclose(result[1], 0)
-
-        assert isinstance(result[2], numpy.ndarray)
-        assert result[2].shape == (4,)
-        assert np.allclose(result[2], 0)
+        assert isinstance(result, np.ndarray)
+        assert result.shape == (4, 3)
+        assert np.allclose(result, 0)
 
     def test_all_zero_diff_methods_multiple_returns(self):
         """Test that the transform works correctly when the diff method for every parameter is
@@ -419,38 +407,12 @@ def circuit(params):
         result = spsa_grad(circuit)(params)
 
         assert isinstance(result, tuple)
-
         assert len(result) == 2
 
-        # First elem
-        assert len(result[0]) == 3
-
-        assert isinstance(result[0][0], numpy.ndarray)
-        assert result[0][0].shape == ()
-        assert np.allclose(result[0][0], 0)
-
-        assert isinstance(result[0][1], numpy.ndarray)
-        assert result[0][1].shape == ()
-        assert np.allclose(result[0][1], 0)
-
-        assert isinstance(result[0][2], numpy.ndarray)
-        assert result[0][2].shape == ()
-        assert np.allclose(result[0][2], 0)
-
-        # Second elem
-        assert len(result[0]) == 3
-
-        assert isinstance(result[1][0], numpy.ndarray)
-        assert result[1][0].shape == (4,)
-        assert np.allclose(result[1][0], 0)
-
-        assert isinstance(result[1][1], numpy.ndarray)
-        assert result[1][1].shape == (4,)
-        assert np.allclose(result[1][1], 0)
-
-        assert isinstance(result[1][2], numpy.ndarray)
-        assert result[1][2].shape == (4,)
-        assert np.allclose(result[1][2], 0)
+        for r, exp_shape in zip(result, [(3,), (4, 3)]):
+            assert isinstance(r, np.ndarray)
+            assert r.shape == exp_shape
+            assert np.allclose(r, 0)
 
     def test_y0(self):
         """Test that if first order finite differences is underlying the SPSA, then
@@ -529,7 +491,7 @@ def test_output_shape_matches_qnode(self):
         dev = qml.device("default.qubit", wires=4)
 
         def cost1(x):
-            qml.Rot(*x, wires=0)
+            qml.Rot(x[0], 0.3 * x[1], x[2], wires=0)
             return qml.expval(qml.PauliZ(0))
 
         def cost2(x):
@@ -559,14 +521,11 @@ def cost6(x):
 
         expected = [
             (3,),
-            (
-                1,
-                3,
-            ),
+            (1, 3),
             (2, 3),
-            (3, 4),
-            (1, 3, 4),
-            (2, 3, 4),
+            (4, 3),
+            (1, 4, 3),
+            (2, 4, 3),
         ]
 
         assert all(t == q for t, q in zip(transform, expected))
diff --git a/tests/gradients/finite_diff/test_spsa_gradient_shot_vec.py b/tests/gradients/finite_diff/test_spsa_gradient_shot_vec.py
--- a/tests/gradients/finite_diff/test_spsa_gradient_shot_vec.py
+++ b/tests/gradients/finite_diff/test_spsa_gradient_shot_vec.py
@@ -252,24 +252,13 @@ def circuit(params):
         grad_fn = spsa_grad(circuit, h=h_val, sampler_rng=rng)
         all_result = grad_fn(params)
 
+        assert isinstance(all_result, tuple)
         assert len(all_result) == len(default_shot_vector)
 
         for result in all_result:
-            assert isinstance(result, tuple)
-
-            assert len(result) == 3
-
-            assert isinstance(result[0], numpy.ndarray)
-            assert result[0].shape == (4,)
-            assert np.allclose(result[0], 0)
-
-            assert isinstance(result[1], numpy.ndarray)
-            assert result[1].shape == (4,)
-            assert np.allclose(result[1], 0)
-
-            assert isinstance(result[2], numpy.ndarray)
-            assert result[2].shape == (4,)
-            assert np.allclose(result[2], 0)
+            assert isinstance(result, np.ndarray)
+            assert result.shape == (4, 3)
+            assert np.allclose(result, 0)
 
     def test_all_zero_diff_methods_multiple_returns(self):
         """Test that the transform works correctly when the diff method for every parameter is
@@ -288,42 +277,16 @@ def circuit(params):
         grad_fn = spsa_grad(circuit, h=h_val, sampler_rng=rng)
         all_result = grad_fn(params)
 
-        assert len(all_result) == len(default_shot_vector)
-
-        for result in all_result:
-            assert isinstance(result, tuple)
-
-            assert len(result) == 2
-
-            # First elem
-            assert len(result[0]) == 3
-
-            assert isinstance(result[0][0], numpy.ndarray)
-            assert result[0][0].shape == ()
-            assert np.allclose(result[0][0], 0)
-
-            assert isinstance(result[0][1], numpy.ndarray)
-            assert result[0][1].shape == ()
-            assert np.allclose(result[0][1], 0)
+        assert isinstance(all_result, tuple)
+        assert len(all_result) == len(many_shots_shot_vector)
 
-            assert isinstance(result[0][2], numpy.ndarray)
-            assert result[0][2].shape == ()
-            assert np.allclose(result[0][2], 0)
-
-            # Second elem
-            assert len(result[0]) == 3
-
-            assert isinstance(result[1][0], numpy.ndarray)
-            assert result[1][0].shape == (4,)
-            assert np.allclose(result[1][0], 0)
-
-            assert isinstance(result[1][1], numpy.ndarray)
-            assert result[1][1].shape == (4,)
-            assert np.allclose(result[1][1], 0)
-
-            assert isinstance(result[1][2], numpy.ndarray)
-            assert result[1][2].shape == (4,)
-            assert np.allclose(result[1][2], 0)
+        for res in all_result:
+            assert isinstance(res, tuple)
+            assert len(res) == 2
+            for r, exp_shape in zip(res, [(3,), (4, 3)]):
+                assert isinstance(r, np.ndarray)
+                assert r.shape == exp_shape
+                assert np.allclose(r, 0)
 
     def test_y0(self):
         """Test that if first order finite differences is underlying the SPSA, then
@@ -433,7 +396,7 @@ def test_output_shape_matches_qnode(self):
         dev = qml.device("default.qubit", wires=4, shots=many_shots_shot_vector)
 
         def cost1(x):
-            qml.Rot(*x, wires=0)
+            qml.Rot(x[0], 0.3 * x[1], x[2], wires=0)
             return qml.expval(qml.PauliZ(0))
 
         def cost2(x):
@@ -461,7 +424,7 @@ def cost6(x):
 
         transform = [qml.math.shape(spsa_grad(c, h=h_val)(x)) for c in circuits]
 
-        expected = [(3, 3), (3, 1, 3), (3, 2, 3), (3, 3, 4), (3, 1, 3, 4), (3, 2, 3, 4)]
+        expected = [(3, 3), (3, 1, 3), (3, 2, 3), (3, 4, 3), (3, 1, 4, 3), (3, 2, 4, 3)]
 
         assert all(t == q for t, q in zip(transform, expected))
 
@@ -852,7 +815,6 @@ def test_var_expectation_values(self, approx_order, strategy, validate):
         with qml.queuing.AnnotatedQueue() as q:
             qml.RX(x, wires=[0])
             qml.RY(y, wires=[1])
-            qml.CNOT(wires=[0, 1])
             qml.expval(qml.PauliZ(0))
             qml.var(qml.PauliX(1))
 
@@ -870,7 +832,7 @@ def test_var_expectation_values(self, approx_order, strategy, validate):
         all_res = fn(dev.execute(tapes))
 
         assert isinstance(all_res, tuple)
-        assert len(all_res) == len(default_shot_vector)
+        assert len(all_res) == len(many_shots_shot_vector)
 
         for res in all_res:
             assert isinstance(res, tuple)
diff --git a/tests/gradients/parameter_shift/test_parameter_shift.py b/tests/gradients/parameter_shift/test_parameter_shift.py
--- a/tests/gradients/parameter_shift/test_parameter_shift.py
+++ b/tests/gradients/parameter_shift/test_parameter_shift.py
@@ -2217,7 +2217,7 @@ def test_projector_variance(self, state, tol):
 
     def cost1(x):
         """Perform rotation and return a scalar expectation value."""
-        qml.Rot(*x, wires=0)
+        qml.Rot(x[0], 0.3 * x[1], x[2], wires=0)
         return qml.expval(qml.PauliZ(0))
 
     def cost2(x):
@@ -2228,7 +2228,7 @@ def cost2(x):
     def cost3(x):
         """Perform rotation and return two expectation value in a 1d array."""
         qml.Rot(*x, wires=0)
-        return [qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliZ(1))]
+        return (qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliZ(1)))
 
     def cost4(x):
         """Perform rotation and return probabilities."""
@@ -2243,17 +2243,53 @@ def cost5(x):
     def cost6(x):
         """Perform rotation and return two sets of probabilities in a 2d object."""
         qml.Rot(*x, wires=0)
-        return [qml.probs([0, 1]), qml.probs([2, 3])]
+        return (qml.probs([0, 1]), qml.probs([2, 3]))
+
+    def cost7(x):
+        """Perform rotation and return a scalar expectation value."""
+        qml.RX(x, wires=0)
+        return qml.expval(qml.PauliZ(0))
+
+    def cost8(x):
+        """Perform rotation and return an expectation value in a 1d array."""
+        qml.RX(x, wires=0)
+        return [qml.expval(qml.PauliZ(0))]
+
+    def cost9(x):
+        """Perform rotation and return two expectation value in a 1d array."""
+        qml.RX(x, wires=0)
+        return (qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliZ(1)))
+
+    costs_and_expected_expval_scalar = [
+        (cost7, (), np.ndarray),
+        (cost8, (1,), list),
+        (cost9, (2,), tuple),
+    ]
+
+    @pytest.mark.parametrize("cost, exp_shape, exp_type", costs_and_expected_expval_scalar)
+    def test_output_shape_matches_qnode_expval_scalar(self, cost, exp_shape, exp_type):
+        """Test that the transform output shape matches that of the QNode for
+        expectation values and a scalar parameter."""
+        dev = qml.device("default.qubit", wires=4)
+
+        x = np.array(0.419)
+        circuit = qml.QNode(cost, dev)
+
+        res_parshift = qml.gradients.param_shift(circuit)(x)
+
+        assert isinstance(res_parshift, exp_type)
+        assert np.array(res_parshift).shape == exp_shape
 
     costs_and_expected_expval = [
-        (cost1, [3]),
-        (cost2, [1, 3]),
-        (cost3, [2, 3]),
+        (cost1, [3], np.ndarray),
+        (cost2, [3], list),
+        (cost3, [2, 3], tuple),
     ]
 
-    @pytest.mark.parametrize("cost, expected_shape", costs_and_expected_expval)
-    def test_output_shape_matches_qnode_expval(self, cost, expected_shape):
-        """Test that the transform output shape matches that of the QNode."""
+    @pytest.mark.parametrize("cost, exp_shape, exp_type", costs_and_expected_expval)
+    def test_output_shape_matches_qnode_expval_array(self, cost, exp_shape, exp_type):
+        """Test that the transform output shape matches that of the QNode for
+        expectation values and an array-valued parameter."""
         dev = qml.device("default.qubit", wires=4)
 
         x = np.random.rand(3)
@@ -2261,44 +2297,40 @@ def test_output_shape_matches_qnode_expval(self, cost, expected_shape):
 
         res = qml.gradients.param_shift(circuit)(x)
 
-        assert len(res) == expected_shape[0]
+        assert isinstance(res, exp_type)
+        if len(res) == 1:
+            res = res[0]
+        assert len(res) == exp_shape[0]
 
-        if len(expected_shape) > 1:
+        if len(exp_shape) > 1:
             for r in res:
-                assert isinstance(r, tuple)
-                assert len(r) == expected_shape[1]
+                assert isinstance(r, np.ndarray)
+                assert len(r) == exp_shape[1]
 
     costs_and_expected_probs = [
-        (cost4, [3, 4]),
-        (cost5, [1, 3, 4]),
-        (cost6, [2, 3, 4]),
+        (cost4, [4, 3], np.ndarray),
+        (cost5, [4, 3], list),
+        (cost6, [2, 4, 3], tuple),
     ]
 
-    @pytest.mark.parametrize("cost, expected_shape", costs_and_expected_probs)
-    def test_output_shape_matches_qnode_probs(self, cost, expected_shape):
+    @pytest.mark.parametrize("cost, exp_shape, exp_type", costs_and_expected_probs)
+    def test_output_shape_matches_qnode_probs(self, cost, exp_shape, exp_type):
         """Test that the transform output shape matches that of the QNode."""
         dev = qml.device("default.qubit", wires=4)
 
         x = np.random.rand(3)
         circuit = qml.QNode(cost, dev)
 
-        res = qml.gradients.param_shift(circuit)(x)
-
-        assert len(res) == expected_shape[0]
+        res_parshift = qml.gradients.param_shift(circuit)(x)
 
-        if len(expected_shape) > 2:
-            for r in res:
-                assert isinstance(r, tuple)
-                assert len(r) == expected_shape[1]
+        # Check data types
+        assert isinstance(res_parshift, exp_type)
+        if len(exp_shape) > 1:
+            for r in res_parshift:
+                assert isinstance(r, np.ndarray)
 
-                for _r in r:
-                    assert isinstance(_r, qml.numpy.ndarray)
-                    assert len(_r) == expected_shape[2]
-
-        elif len(expected_shape) > 1:
-            for r in res:
-                assert isinstance(r, qml.numpy.ndarray)
-                assert len(r) == expected_shape[1]
+        # Check shape, result can be put into a single array by assumption
+        assert np.allclose(np.squeeze(np.array(res_parshift)).shape, exp_shape)
 
     # TODO: revisit the following test when the Autograd interface supports
     # parameter-shift with the new return type system
@@ -3028,7 +3060,7 @@ def test_output_shape_matches_qnode(self):
 
         def cost1(x):
             """Perform rotation and return a scalar expectation value."""
-            qml.Rot(*x, wires=0)
+            qml.Rot(x[0], 0.3 * x[1], x[2], wires=0)
             return qml.expval(qml.PauliZ(0))
 
         def cost2(x):
@@ -3060,7 +3092,7 @@ def cost6(x):
         single_measure_circuits = [qml.QNode(cost, dev) for cost in (cost1, cost2, cost4, cost5)]
         multi_measure_circuits = [qml.QNode(cost, dev) for cost in (cost3, cost6)]
 
-        for c, exp_shape in zip(single_measure_circuits, [(3,), (1, 3), (3, 4), (1, 3, 4)]):
+        for c, exp_shape in zip(single_measure_circuits, [(3,), (1, 3), (4, 3), (1, 4, 3)]):
             grad = qml.gradients.param_shift(c, broadcast=True)(x)
             assert qml.math.shape(grad) == exp_shape
 
@@ -4590,7 +4622,7 @@ def circuit(x, y):
             qml.QuantumFunctionError,
             match="argnum does not work with the Jax interface. You should use argnums instead.",
         ):
-            qml.gradients.hadamard_grad(circuit, argnum=argnums)(x, y)
+            qml.gradients.param_shift(circuit, argnum=argnums)(x, y)
 
     def test_single_expectation_value(self, argnums, interface):
         """Test for single expectation value."""
@@ -4664,7 +4696,6 @@ def circuit(x, y):
             qml.RX(x[0], wires=[0])
             qml.RY(x[1], wires=[1])
             qml.RY(y, wires=[1])
-            qml.CNOT(wires=[0, 1])
             return qml.var(qml.PauliZ(0) @ qml.PauliX(1))
 
         x = jax.numpy.array([0.543, -0.654])
@@ -4675,12 +4706,11 @@ def circuit(x, y):
         )
         res_expected = jax.hessian(circuit, argnums=argnums)(x, y)
 
-        if argnums == [0]:
-            assert np.allclose(res[0][0], res_expected[0][0][0])
-            assert np.allclose(res[1][0], res_expected[0][0][1])
+        if len(argnums) == 1:
+            # jax.hessian produces an additional tuple axis, which we have to index away here
+            assert np.allclose(res, res_expected[0])
         else:
-            if len(argnums) != 1:
-                res = res[0]
-
-            for r, r_e in zip(res, res_expected[0]):
-                assert np.allclose(r, r_e)
+            # The Hessian is a 2x2 nested tuple "matrix" for argnums=[0, 1]
+            for r, r_e in zip(res, res_expected):
+                for r_, r_e_ in zip(r, r_e):
+                    assert np.allclose(r_, r_e_)
diff --git a/tests/gradients/parameter_shift/test_parameter_shift_shot_vec.py b/tests/gradients/parameter_shift/test_parameter_shift_shot_vec.py
--- a/tests/gradients/parameter_shift/test_parameter_shift_shot_vec.py
+++ b/tests/gradients/parameter_shift/test_parameter_shift_shot_vec.py
@@ -1842,9 +1842,9 @@ def cost6(x):
         return [qml.probs([0, 1]), qml.probs([2, 3])]
 
     costs_and_expected_expval = [
-        (cost1, [3]),
-        (cost2, [3]),
-        (cost3, [2, 3]),
+        (cost1, (3,)),
+        (cost2, (3,)),
+        (cost3, (2, 3)),
     ]
 
     @pytest.mark.parametrize("cost, expected_shape", costs_and_expected_expval)
@@ -1864,18 +1864,13 @@ def test_output_shape_matches_qnode_expval(self, cost, expected_shape):
         assert isinstance(all_res, tuple)
 
         for res in all_res:
-            assert isinstance(res, tuple)
-            assert len(res) == expected_shape[0]
-
-            if len(expected_shape) > 1:
-                for r in res:
-                    assert isinstance(r, tuple)
-                    assert len(r) == expected_shape[1]
+            assert isinstance(res, np.ndarray)
+            assert res.shape == expected_shape
 
     costs_and_expected_probs = [
-        (cost4, [3, 4]),
-        (cost5, [3, 4]),
-        (cost6, [2, 3, 4]),
+        (cost4, (4, 3)),
+        (cost5, (4, 3)),
+        (cost6, (2, 4, 3)),
     ]
 
     @pytest.mark.parametrize("cost, expected_shape", costs_and_expected_probs)
@@ -1895,22 +1890,8 @@ def test_output_shape_matches_qnode_probs(self, cost, expected_shape):
         assert isinstance(all_res, tuple)
 
         for res in all_res:
-            assert isinstance(res, tuple)
-            assert len(res) == expected_shape[0]
-
-            if len(expected_shape) > 2:
-                for r in res:
-                    assert isinstance(r, tuple)
-                    assert len(r) == expected_shape[1]
-
-                    for _r in r:
-                        assert isinstance(_r, qml.numpy.ndarray)
-                        assert len(_r) == expected_shape[2]
-
-            elif len(expected_shape) > 1:
-                for r in res:
-                    assert isinstance(r, qml.numpy.ndarray)
-                    assert len(r) == expected_shape[1]
+            assert isinstance(res, np.ndarray)
+            assert res.shape == expected_shape
 
     # TODO: revisit the following test when the Autograd interface supports
     # parameter-shift with the new return type system