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:
codecov[bot]: ## [Codecov](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4945?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
Attention: `1 lines` in your changes are missing coverage. Please review.
> Comparison is base [(`18608ef`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/18608ef0035b9b566a58f6dabd13fa2bde1c815f?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) 99.66% compared to head [(`cfb1b7a`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4945?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) 99.66%.

| [Files](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4945?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | Patch % | Lines |
|---|---|---|
| [pennylane/gradients/gradient\_transform.py](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4945?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2dyYWRpZW50cy9ncmFkaWVudF90cmFuc2Zvcm0ucHk=) | 92.85% | [1 Missing :warning: ](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4945?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) |

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


```diff
@@            Coverage Diff             @@
##           master    #4945      +/-   ##
==========================================
- Coverage   99.66%   99.66%   -0.01%     
==========================================
  Files         388      388              
  Lines       35379    35121     -258     
==========================================
- Hits        35262    35002     -260     
- Misses        117      119       +2     
```



</details>

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

dwierichs: Thanks @rmoyard. I left the code in to see whether we want to take the small perf hit or keep the code simpler.

I timed the following:

```python3
import pennylane as qml
from pennylane import numpy as pnp

dev = qml.device("default.qubit")

@qml.qnode(dev)
def f(x):
    [qml.RX(_x, 0) for _x in x]
    return qml.expval(qml.PauliZ(0))

n = 100
x = pnp.random.random(n, requires_grad=True)
%timeit qml.gradients.param_shift(f)(x)
```

```python
2.66 s ± 78.2 ms # with "skip_cjac"
2.78 s ± 15.5 ms # without "skip_cjac"
2.77 s ± 64.9 ms # master
```

So there is _some_ slowdown, I don't know how to gauge this against the added code complexity for skipping the contraction?

Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index dc82919ac..4d84dddca 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -388,6 +388,10 @@
 
 <h3>Breaking changes 💔</h3>
 
+* Applying a `gradient_transform` to a QNode directly now gives the same shape and type independent
+  of whether there is classical processing in the node.
+  [(#4945)](https://github.com/PennyLaneAI/pennylane/pull/4945)
+  
 * State measurements preserve `dtype`.
   [(#5547)](https://github.com/PennyLaneAI/pennylane/pull/5547)
 
@@ -496,6 +500,10 @@
 
 <h3>Bug fixes 🐛</h3>
 
+* Fixed a bug where the shape and type of derivatives obtained by applying a gradient transform to
+  a QNode differed, based on whether the QNode uses classical coprocessing.
+  [(#4945)](https://github.com/PennyLaneAI/pennylane/pull/4945)
+
 * `ApproxTimeEvolution`, `CommutingEvolution`, `QDrift`, and `TrotterProduct` 
   now de-queue their input observable.
   [(#5524)](https://github.com/PennyLaneAI/pennylane/pull/5524)
diff --git a/pennylane/gradients/gradient_transform.py b/pennylane/gradients/gradient_transform.py
index 691f28dcc..5c6273012 100644
--- 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
index df8697585..b46cded82 100644
--- 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 @@ class TestGradientTransformIntegration:
         @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 @@ class TestGradientTransformIntegration:
         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 @@ class TestGradientTransformIntegration:
         @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 @@ class TestGradientTransformIntegration:
         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
index 72e4c41ea..78aa01826 100644
--- 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 @@ class TestHadamardGrad:
         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 @@ class TestHadamardGrad:
 
         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 @@ class TestHadamardGrad:
         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 @@ class TestHadamardGradEdgeCases:
         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 @@ class TestJaxArgnums:
             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 @@ class TestJaxArgnums:
         )
         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
index 43dd85617..d135382b3 100644
--- a/tests/gradients/finite_diff/test_finite_difference.py
+++ b/tests/gradients/finite_diff/test_finite_difference.py
@@ -294,35 +294,24 @@ class TestFiniteDiff:
         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 @@ class TestFiniteDiff:
 
         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 @@ class TestFiniteDiff:
         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 @@ class TestFiniteDiff:
 
         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 @@ class TestJaxArgnums:
             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 @@ class TestJaxArgnums:
 
         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
index 48e1e955a..a0c562150 100644
--- 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 @@ class TestFiniteDiff:
         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 @@ class TestFiniteDiff:
         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 @@ class TestFiniteDiff:
 
     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 @@ class TestFiniteDiff:
             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 @@ class TestFiniteDiff:
         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
index 91ea1c898..667a9f00a 100644
--- a/tests/gradients/finite_diff/test_spsa_gradient.py
+++ b/tests/gradients/finite_diff/test_spsa_gradient.py
@@ -387,21 +387,9 @@ class TestSpsaGradient:
 
         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 @@ class TestSpsaGradient:
         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 @@ class TestSpsaGradient:
         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 @@ class TestSpsaGradient:
 
         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
index 7c4bd104f..a88f9c293 100644
--- 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 @@ class TestSpsaGradient:
         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 @@ class TestSpsaGradient:
         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 @@ class TestSpsaGradient:
         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 @@ class TestSpsaGradient:
 
         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 @@ class TestSpsaGradientIntegration:
         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 @@ class TestSpsaGradientIntegration:
         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
index 8a7e9bc67..7a0efd917 100644
--- a/tests/gradients/parameter_shift/test_parameter_shift.py
+++ b/tests/gradients/parameter_shift/test_parameter_shift.py
@@ -2217,7 +2217,7 @@ class TestParameterShiftRule:
 
     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 @@ class TestParameterShiftRule:
     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 @@ class TestParameterShiftRule:
     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 @@ class TestParameterShiftRule:
 
         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 @@ class TestParameterShiftRuleBroadcast:
 
         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 @@ class TestParameterShiftRuleBroadcast:
         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 @@ class TestJaxArgnums:
             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 @@ class TestJaxArgnums:
             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 @@ class TestJaxArgnums:
         )
         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
index 1a391914f..17b6ba6b1 100644
--- 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 @@ class TestParameterShiftRule:
         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 @@ class TestParameterShiftRule:
         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 @@ class TestParameterShiftRule:
         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