Repository: PennyLaneAI/pennylane

Issue:
[BUG] `param_shift` with `broadcast=True` does not work properly with shot vectors
### Expected behavior

The example below  should work

### Actual behavior

It raises a shape mismatch error

### Additional information

It works if `len(shots)=2`, which is not a good sign. It means that we're wrongly contracting results across the shots axis, which only works as long as they have the right length.

### Source code

```shell
param = pnp.array(0.5, requires_grad=True)
shots = (10, 100, 10)
dev = qml.device("default.qubit", wires=2, shots=shots)
@qml.qnode(dev)
def circuit(param):
    qml.RX(param, wires=0)
    return qml.expval(qml.Z(0))

qml.gradients.param_shift(circuit, broadcast=True)(param)
```


### Tracebacks

```shell
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Cell In[36], line 11
      8     qml.RX(param, wires=0)
      9     return qml.expval(qml.Z(0))
---> 11 qml.gradients.param_shift(circuit, broadcast=True)(param)

File ~/repos/pennylane/pennylane/workflow/qnode.py:1088, in QNode.__call__(self, *args, **kwargs)
   1085 self._update_gradient_fn(shots=override_shots, tape=self._tape)
   1087 try:
-> 1088     res = self._execution_component(args, kwargs, override_shots=override_shots)
   1089 finally:
   1090     if old_interface == "auto":

File ~/repos/pennylane/pennylane/workflow/qnode.py:1042, in QNode._execution_component(self, args, kwargs, override_shots)
   1039 full_transform_program.prune_dynamic_transform()
   1041 # pylint: disable=unexpected-keyword-arg
-> 1042 res = qml.execute(
   1043     (self._tape,),
   1044     device=self.device,
   1045     gradient_fn=self.gradient_fn,
   1046     interface=self.interface,
   1047     transform_program=full_transform_program,
   1048     config=config,
   1049     gradient_kwargs=self.gradient_kwargs,
   1050     override_shots=override_shots,
   1051     **self.execute_kwargs,
   1052 )
   1053 res = res[0]
   1055 # convert result to the interface in case the qfunc has no parameters

File ~/repos/pennylane/pennylane/workflow/execution.py:798, in execute(tapes, device, gradient_fn, interface, transform_program, config, grad_on_execution, gradient_kwargs, cache, cachesize, max_diff, override_shots, expand_fn, max_expansion, device_batch_transform, device_vjp)
    793 else:
    794     results = ml_boundary_execute(
    795         tapes, device, execute_fn, gradient_fn, gradient_kwargs, _n=1, max_diff=max_diff
    796     )
--> 798 return post_processing(results)

File ~/repos/pennylane/pennylane/transforms/core/transform_program.py:88, in _apply_postprocessing_stack(results, postprocessing_stack)
     65 """Applies the postprocessing and cotransform postprocessing functions in a Last-In-First-Out LIFO manner.
     66 
     67 Args:
   (...)
     85 
     86 """
     87 for postprocessing in reversed(postprocessing_stack):
---> 88     results = postprocessing(results)
     89 return results

File ~/repos/pennylane/pennylane/transforms/core/transform_program.py:58, in _batch_postprocessing(results, individual_fns, slices)
     32 def _batch_postprocessing(
     33     results: ResultBatch, individual_fns: List[PostProcessingFn], slices: List[slice]
     34 ) -> ResultBatch:
     35     """Broadcast individual post processing functions onto their respective tapes.
     36 
     37     Args:
   (...)
     56 
     57     """
---> 58     return tuple(fn(results[sl]) for fn, sl in zip(individual_fns, slices))

File ~/repos/pennylane/pennylane/transforms/core/transform_program.py:58, in <genexpr>(.0)
     32 def _batch_postprocessing(
     33     results: ResultBatch, individual_fns: List[PostProcessingFn], slices: List[slice]
     34 ) -> ResultBatch:
     35     """Broadcast individual post processing functions onto their respective tapes.
     36 
     37     Args:
   (...)
     56 
     57     """
---> 58     return tuple(fn(results[sl]) for fn, sl in zip(individual_fns, slices))

File ~/repos/pennylane/pennylane/gradients/parameter_shift.py:433, in expval_param_shift.<locals>.processing_fn(results)
    430     res = results[start : start + num_tapes] if batch_size is None else results[start]
    431     start = start + num_tapes
--> 433     g = _evaluate_gradient(tape, res, data, r0)
    434     grads.append(g)
    436 # g will have been defined at least once (because otherwise all gradients would have
    437 # been zero), providing a representative for a zero gradient to emulate its type/shape.

File ~/repos/pennylane/pennylane/gradients/parameter_shift.py:213, in _evaluate_gradient(tape, res, data, r0)
    211 for i in range(len_shot_vec):
    212     shot_comp_res = [r[i] for r in res]
--> 213     shot_comp_res = _single_meas_grad(shot_comp_res, coeffs, unshifted_coeff, r0[i])
    214     g.append(shot_comp_res)
    215 return tuple(g)

File ~/repos/pennylane/pennylane/gradients/parameter_shift.py:164, in _single_meas_grad(result, coeffs, unshifted_coeff, r0)
    162 result = qml.math.stack(result)
    163 coeffs = qml.math.convert_like(coeffs, result)
--> 164 g = qml.math.tensordot(result, coeffs, [[0], [0]])
    165 if unshifted_coeff is not None:
    166     # add the unshifted term
    167     g = g + unshifted_coeff * r0

File ~/repos/pennylane/pennylane/math/multi_dispatch.py:151, in multi_dispatch.<locals>.decorator.<locals>.wrapper(*args, **kwargs)
    148 interface = interface or get_interface(*dispatch_args)
    149 kwargs["like"] = interface
--> 151 return fn(*args, **kwargs)

File ~/repos/pennylane/pennylane/math/multi_dispatch.py:398, in tensordot(tensor1, tensor2, axes, like)
    372 """Returns the tensor product of two tensors.
    373 In general ``axes`` specifies either the set of axes for both
    374 tensors that are contracted (with the first/second entry of ``axes``
   (...)
    395     tensor_like: the tensor product of the two input tensors
    396 """
    397 tensor1, tensor2 = np.coerce([tensor1, tensor2], like=like)
--> 398 return np.tensordot(tensor1, tensor2, axes=axes, like=like)

File ~/venvs/dev/lib/python3.10/site-packages/autoray/autoray.py:80, in do(fn, like, *args, **kwargs)
     31 """Do function named ``fn`` on ``(*args, **kwargs)``, peforming single
     32 dispatch to retrieve ``fn`` based on whichever library defines the class of
     33 the ``args[0]``, or the ``like`` keyword argument if specified.
   (...)
     77     <tf.Tensor: id=91, shape=(3, 3), dtype=float32>
     78 """
     79 backend = choose_backend(fn, *args, like=like, **kwargs)
---> 80 return get_lib_fn(backend, fn)(*args, **kwargs)

File ~/venvs/dev/lib/python3.10/site-packages/numpy/core/numeric.py:1099, in tensordot(a, b, axes)
   1097             axes_b[k] += ndb
   1098 if not equal:
-> 1099     raise ValueError("shape-mismatch for sum")
   1101 # Move the axes to sum over to the end of "a"
   1102 # and to the front of "b"
   1103 notin = [k for k in range(nda) if k not in axes_a]

ValueError: shape-mismatch for sum
```


### System information

```shell
pl dev
```


### 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
@@ -48,26 +48,6 @@
 }
 
 
-def assert_multimeasure_not_broadcasted(measurements, broadcast):
-    """Assert that there are not simultaneously multiple measurements and
-    broadcasting activated. Otherwise raises an error."""
-    if broadcast and len(measurements) > 1:
-        raise NotImplementedError(
-            "Broadcasting with multiple measurements is not supported yet. "
-            f"Set broadcast to False instead. The tape measurements are {measurements}."
-        )
-
-
-def assert_shot_vector_not_broadcasted(shots, broadcast):
-    """Assert that there are not simultaneously multiple shot settings (shot vector) and
-    broadcasting activated. Otherwise raises an error."""
-    if broadcast and shots.has_partitioned_shots:
-        raise NotImplementedError(
-            "Broadcasting with shot vectors is not supported yet. "
-            f"Set broadcast to False instead. The tape shots are {shots}."
-        )
-
-
 def assert_no_state_returns(measurements, transform_name):
     """Check whether a set of measurements contains a measurement process that returns the quantum
     state and raise an error if this is the case.
diff --git a/pennylane/gradients/parameter_shift.py b/pennylane/gradients/parameter_shift.py
--- a/pennylane/gradients/parameter_shift.py
+++ b/pennylane/gradients/parameter_shift.py
@@ -39,10 +39,8 @@
     _move_first_axis_to_third_pos,
     _no_trainable_grad,
     _swap_first_two_axes,
-    assert_multimeasure_not_broadcasted,
     assert_no_state_returns,
     assert_no_trainable_tape_batching,
-    assert_shot_vector_not_broadcasted,
     choose_trainable_params,
     find_and_validate_gradient_methods,
     reorder_grads,
@@ -1093,8 +1091,6 @@ def param_shift(
 
     transform_name = "parameter-shift rule"
     assert_no_state_returns(tape.measurements, transform_name)
-    assert_multimeasure_not_broadcasted(tape.measurements, broadcast)
-    assert_shot_vector_not_broadcasted(tape.shots, broadcast)
     assert_no_trainable_tape_batching(tape, transform_name)
 
     if argnum is None and not tape.trainable_params:


Test patch:
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
@@ -838,8 +838,6 @@ def test_custom_recipe_unshifted_only(self, ops_with_custom_recipe, multi_measur
         """Test that if the gradient recipe has a zero-shift component, then
         the tape is executed only once using the current parameter
         values."""
-        if multi_measure and broadcast:
-            pytest.skip("Multiple measurements are not supported with `broadcast=True` yet.")
         dev = qml.device("default.qubit", wires=2)
         x = [0.543, -0.654]
 
@@ -876,8 +874,9 @@ def test_custom_recipe_unshifted_only(self, ops_with_custom_recipe, multi_measur
             ]
         assert qml.math.allclose(grad, expected, atol=1e-5)
 
+    @pytest.mark.parametrize("broadcast", [False, True])
     @pytest.mark.parametrize("ops_with_custom_recipe", [[0], [1], [0, 1]])
-    def test_custom_recipe_mixing_unshifted_shifted(self, ops_with_custom_recipe):
+    def test_custom_recipe_mixing_unshifted_shifted(self, ops_with_custom_recipe, broadcast):
         """Test that if the gradient recipe has a zero-shift component, then
         the tape is executed only once using the current parameter
         values."""
@@ -891,19 +890,18 @@ def test_custom_recipe_mixing_unshifted_shifted(self, ops_with_custom_recipe):
             qml.expval(qml.PauliZ(1))
 
         tape = qml.tape.QuantumScript.from_queue(q)
+        custom_recipe = [[-1e-7, 1, 0], [1e-7, 1, 0], [-1e5, 1, -5e-6], [1e5, 1, 5e-6]]
         gradient_recipes = tuple(
-            (
-                [[-1e-7, 1, 0], [1e-7, 1, 0], [-1e5, 1, -5e-6], [1e5, 1, 5e-6]]
-                if i in ops_with_custom_recipe
-                else None
-            )
-            for i in range(2)
+            custom_recipe if i in ops_with_custom_recipe else None for i in range(2)
         )
-        tapes, fn = qml.gradients.param_shift(tape, gradient_recipes=gradient_recipes)
+        tapes, fn = param_shift(tape, gradient_recipes=gradient_recipes, broadcast=broadcast)
 
         # two tapes per parameter, independent of recipe
         # plus one global (unshifted) call if at least one uses the custom recipe
-        assert len(tapes) == 2 * tape.num_params + int(len(ops_with_custom_recipe) > 0)
+        tapes_per_param = 1 if broadcast else 2
+        num_custom = len(ops_with_custom_recipe)
+        assert len(tapes) == tapes_per_param * tape.num_params + (num_custom > 0)
+
         # Test that executing the tapes and the postprocessing function works
         grad = fn(qml.execute(tapes, dev, None))
         assert qml.math.allclose(grad[0], -np.sin(x[0] + x[1]), atol=1e-5)
@@ -1187,8 +1185,7 @@ class TestParamShiftUsingBroadcasting:
     The tests for `param_shift` on already broadcasted tapes can be found above."""
 
     def test_independent_parameter(self, mocker):
-        """Test that an independent parameter is skipped
-        during the Jacobian computation."""
+        """Test that an independent parameter is skipped during the Jacobian computation."""
         spy = mocker.spy(qml.gradients.parameter_shift, "expval_param_shift")
 
         with qml.queuing.AnnotatedQueue() as q:
@@ -2859,7 +2856,6 @@ def test_gradients_agree_finite_differences(self, tol):
         assert np.allclose(grad_A, grad_F1, atol=tol, rtol=0)
         assert np.allclose(grad_A, grad_F2, atol=tol, rtol=0)
 
-    @pytest.mark.xfail(reason="Broadcasting with multiple measurements is not supported yet")
     def test_variance_gradients_agree_finite_differences(self, tol):
         """Tests that the variance parameter-shift rule agrees with the first and second
         order finite differences"""
@@ -2892,17 +2888,21 @@ def test_variance_gradients_agree_finite_differences(self, tol):
                 assert np.allclose(g, grad_F1[idx1][idx2], atol=tol, rtol=0)
                 assert np.allclose(g, grad_F2[idx1][idx2], atol=tol, rtol=0)
 
-    @pytest.mark.autograd
-    @pytest.mark.parametrize("dev_name", ["default.qubit", "default.qubit.autograd"])
-    def test_fallback(self, dev_name, mocker):
+    @pytest.mark.jax
+    @pytest.mark.parametrize("dev_name", ["default.qubit", "default.qubit.jax"])
+    def test_fallback(self, dev_name, mocker, tol):
         """Test that fallback gradient functions are correctly used"""
+
+        import jax
+        from jax import numpy as jnp
+
         spy = mocker.spy(qml.gradients, "finite_diff")
         dev = qml.device(dev_name, wires=2)
         execute_fn = dev.execute if dev_name == "default.qubit" else dev.batch_execute
         x = 0.543
         y = -0.654
 
-        params = np.array([x, y], requires_grad=True)
+        params = jnp.array([x, y])
 
         def cost_fn(params):
             with qml.queuing.AnnotatedQueue() as q:
@@ -2922,18 +2922,16 @@ def cost_fn(params):
 
             return fn(execute_fn(tapes))
 
-        with pytest.raises(NotImplementedError, match="Broadcasting with multiple measurements"):
-            cost_fn(params)
-        # TODO: Uncomment the following when #2693 is resolved. Add test fixture arg `tol`
-        # res = cost_fn(params)
-        # assert res.shape == (2, 2)
-        # expected = np.array([[-np.sin(x), 0], [0, -2 * np.cos(y) * np.sin(y)]])
-        # assert np.allclose(res, expected, atol=tol, rtol=0)
+        res = cost_fn(params)
+        assert len(res) == 2 and isinstance(res, tuple)
+        assert all(len(r) == 2 and isinstance(r, tuple) for r in res)
+        expected = ((-np.sin(x), 0), (0, -2 * np.cos(y) * np.sin(y)))
+        assert np.allclose(res, expected, atol=tol, rtol=0)
 
         # double check the derivative
-        # jac = qml.jacobian(cost_fn)(params)
-        # assert np.allclose(jac[0, 0, 0], -np.cos(x), atol=tol, rtol=0)
-        # assert np.allclose(jac[1, 1, 1], -2 * np.cos(2 * y), atol=tol, rtol=0)
+        jac = jax.jacobian(cost_fn)(params)
+        assert np.allclose(jac[0][0][0], -np.cos(x), atol=tol, rtol=0)
+        assert np.allclose(jac[1][1][1], -2 * np.cos(2 * y), atol=tol, rtol=0)
 
     @pytest.mark.autograd
     @pytest.mark.parametrize("dev_name", ["default.qubit", "default.qubit.autograd"])
@@ -2998,116 +2996,78 @@ def test_single_expectation_value(self, tol):
         for r, e in zip(res, expected):
             assert np.allclose(r, e, atol=tol, rtol=0)
 
-    def test_multiple_expectation_values(self):
+    def test_multiple_expectation_values(self, tol):
         """Tests correct output shape and evaluation for a tape
         with multiple expval outputs"""
         x = 0.543
         y = -0.654
+        ops = [qml.RX(x, 0), qml.RY(y, 1), qml.CNOT([0, 1])]
+        meas = [qml.expval(qml.Z(0)), qml.expval(qml.X(1))]
+        tape = qml.tape.QuantumScript(ops, meas)
 
-        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.expval(qml.PauliX(1))
-
-        tape = qml.tape.QuantumScript.from_queue(q)
-        with pytest.raises(NotImplementedError, match="Broadcasting with multiple measurements"):
-            qml.gradients.param_shift(tape, broadcast=True)
-        # TODO: Uncomment the following when #2693 is resolved. Add test fixture arg `tol`
-        # dev = qml.device("default.qubit", wires=2)
-        # tapes, fn = qml.gradients.param_shift(tape, broadcast=True)
-        # assert len(tapes) == 2
-        # assert tapes[0].batch_size == tapes[1].batch_size == 2
+        tapes, fn = qml.gradients.param_shift(tape, broadcast=True)
+        assert len(tapes) == 2
+        assert tapes[0].batch_size == tapes[1].batch_size == 2
 
-        # res = fn(dev.execute(tapes))
-        # assert res.shape == (2, 2)
+        dev = qml.device("default.qubit", wires=2)
+        res = fn(dev.execute(tapes))
+        assert len(res) == 2
+        assert all(len(r) == 2 for r in res)
 
-        # expected = np.array([[-np.sin(x), 0], [0, np.cos(y)]])
-        # assert np.allclose(res, expected, atol=tol, rtol=0)
+        expected = np.array([[-np.sin(x), 0], [0, np.cos(y)]])
+        assert np.allclose(res, expected, atol=tol, rtol=0)
 
-    def test_var_expectation_values(self):
+    def test_var_expectation_values(self, tol):
         """Tests correct output shape and evaluation for a tape
         with expval and var outputs"""
         x = 0.543
         y = -0.654
+        ops = [qml.RX(x, 0), qml.RY(y, 1), qml.CNOT([0, 1])]
+        meas = [qml.expval(qml.Z(0)), qml.var(qml.X(1))]
+        tape = qml.tape.QuantumScript(ops, meas)
 
-        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))
-
-        tape = qml.tape.QuantumScript.from_queue(q)
-        with pytest.raises(NotImplementedError, match="Broadcasting with multiple measurements"):
-            qml.gradients.param_shift(tape, broadcast=True)
-        # TODO: Uncomment the following when #2693 is resolved. Add test fixture arg `tol`
-        # dev = qml.device("default.qubit", wires=2)
-        # tapes, fn = qml.gradients.param_shift(tape, broadcast=True)
-        # assert len(tapes) == 3  # One unshifted, two broadcasted shifted tapes
-        # assert tapes[0].batch_size is None
-        # assert tapes[1].batch_size == tapes[2].batch_size == 2
+        tapes, fn = qml.gradients.param_shift(tape, broadcast=True)
+        assert len(tapes) == 3  # One unshifted, two broadcasted shifted tapes
+        assert tapes[0].batch_size is None
+        assert tapes[1].batch_size == tapes[2].batch_size == 2
 
-        # res = fn(dev.execute(tapes))
-        # assert res.shape == (2, 2)
+        dev = qml.device("default.qubit", wires=2)
+        res = fn(dev.execute(tapes))
+        assert len(res) == 2
+        assert all(len(r) == 2 for r in res)
 
-        # expected = np.array([[-np.sin(x), 0], [0, -2 * np.cos(y) * np.sin(y)]])
-        # assert np.allclose(res, expected, atol=tol, rtol=0)
+        expected = np.array([[-np.sin(x), 0], [0, -2 * np.cos(y) * np.sin(y)]])
+        assert np.allclose(res, expected, atol=tol, rtol=0)
 
-    def test_prob_expectation_values(self):
+    def test_prob_expectation_values(self, tol):
         """Tests correct output shape and evaluation for a tape
         with prob and expval outputs"""
-        dev = qml.device("default.qubit", wires=2)
         x = 0.543
         y = -0.654
+        ops = [qml.RX(x, 0), qml.RY(y, 1), qml.CNOT([0, 1])]
+        meas = [qml.expval(qml.Z(0)), qml.probs([0, 1])]
+        tape = qml.tape.QuantumScript(ops, meas)
 
-        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.probs(wires=[0, 1])
+        tapes, fn = qml.gradients.param_shift(tape, broadcast=True)
+        assert len(tapes) == 2
+        assert tapes[0].batch_size == tapes[1].batch_size == 2
 
-        tape = qml.tape.QuantumScript.from_queue(q)
-        dev.execute(tape)
-
-        with pytest.raises(NotImplementedError, match="Broadcasting with multiple measurements"):
-            qml.gradients.param_shift(tape, broadcast=True)
-        # TODO: Uncomment the following when #2693 is resolved. Add test fixture arg `tol`
-        # tapes, fn = qml.gradients.param_shift(tape, broadcast=True)
-        # assert len(tapes) == 2
-        # assert tapes[0].batch_size == tapes[1].batch_size == 2
-
-        # res = fn(dev.execute(tapes))
-        # assert res.shape == (5, 2)
-
-        # expected = (
-        # np.array(
-        # [
-        # [-2 * np.sin(x), 0],
-        # [
-        # -(np.cos(y / 2) ** 2 * np.sin(x)),
-        # -(np.cos(x / 2) ** 2 * np.sin(y)),
-        # ],
-        # [
-        # -(np.sin(x) * np.sin(y / 2) ** 2),
-        # (np.cos(x / 2) ** 2 * np.sin(y)),
-        # ],
-        # [
-        # (np.sin(x) * np.sin(y / 2) ** 2),
-        # (np.sin(x / 2) ** 2 * np.sin(y)),
-        # ],
-        # [
-        # (np.cos(y / 2) ** 2 * np.sin(x)),
-        # -(np.sin(x / 2) ** 2 * np.sin(y)),
-        # ],
-        # ]
-        # )
-        # / 2
-        # )
+        dev = qml.device("default.qubit", wires=2)
+        res = fn(dev.execute(tapes))
+        assert isinstance(res, tuple) and len(res) == 2
+        assert all(isinstance(r, tuple) and len(r) == 2 for r in res)
+        assert all(isinstance(r, np.ndarray) and r.shape == () for r in res[0])
+        assert all(isinstance(r, np.ndarray) and r.shape == (4,) for r in res[1])
+
+        expected_expval = (-np.sin(x), 0)
+        sx, cx, sy, cy = np.sin(x / 2), np.cos(x / 2), np.sin(y / 2), np.cos(y / 2)
+        expected_probs = (
+            np.sin(x) / 2 * np.array([-(cy**2), -(sy**2), sy**2, cy**2]),
+            np.array([-(cx**2), cx**2, sx**2, -(sx**2)]) * np.sin(y) / 2,
+        )
 
-        # assert np.allclose(res, expected, atol=tol, rtol=0)
+        assert np.allclose(res[0], expected_expval, atol=tol, rtol=0)
+        assert np.allclose(res[1], expected_probs, atol=tol, rtol=0)
 
     def test_involutory_variance(self, tol):
         """Tests qubit observables that are involutory"""
@@ -3184,13 +3144,8 @@ def test_involutory_and_noninvolutory_variance(self, tol):
         A = np.array([[4, -1 + 6j], [-1 - 6j, 2]])
         a = 0.54
 
-        with qml.queuing.AnnotatedQueue() as q:
-            qml.RX(a, wires=0)
-            qml.RX(a, wires=1)
-            qml.var(qml.PauliZ(0))
-            qml.var(qml.Hermitian(A, 1))
-
-        tape = qml.tape.QuantumScript.from_queue(q)
+        meas = [qml.var(qml.Z(0)), qml.var(qml.Hermitian(A, 1))]
+        tape = qml.tape.QuantumScript([qml.RX(a, 0), qml.RX(a, 1)], meas)
         tape.trainable_params = {0, 1}
 
         res = dev.execute(tape)
@@ -3198,41 +3153,31 @@ def test_involutory_and_noninvolutory_variance(self, tol):
         assert np.allclose(res, expected, atol=tol, rtol=0)
 
         # circuit jacobians
-        with pytest.raises(NotImplementedError, match="Broadcasting with multiple measurements"):
-            qml.gradients.param_shift(tape, broadcast=True)
-        # TODO: Uncomment the following when #2693 is resolved.
-        # tapes, fn = qml.gradients.param_shift(tape, broadcast=True)
-        # gradA = fn(dev.execute(tapes))
-        # assert len(tapes) == 1 + 2 * 4
+        tapes, fn = qml.gradients.param_shift(tape, broadcast=True)
+        gradA = fn(dev.execute(tapes))
+        # 1 unshifted tape and 4 broadcasted shifted tapes
+        assert len(tapes) == 1 + 4
 
-        # tapes, fn = qml.gradients.finite_diff(tape)
-        # gradF = fn(dev.execute(tapes))
-        # assert len(tapes) == 1 + 2
+        tapes, fn = qml.gradients.finite_diff(tape)
+        gradF = fn(dev.execute(tapes))
+        assert len(tapes) == 1 + 2
 
-        # expected = [2 * np.sin(a) * np.cos(a), -35 * np.sin(2 * a) - 12 * np.cos(2 * a)]
-        # assert np.diag(gradA) == pytest.approx(expected, abs=tol)
-        # assert np.diag(gradF) == pytest.approx(expected, abs=tol)
+        expected = [2 * np.sin(a) * np.cos(a), -35 * np.sin(2 * a) - 12 * np.cos(2 * a)]
+        assert np.diag(gradA) == pytest.approx(expected, abs=tol)
+        assert np.diag(gradF) == pytest.approx(expected, abs=tol)
 
     def test_expval_and_variance(self, tol):
-        """Test that the qnode works for a combination of expectation
+        """Test that the gradient transform works for a combination of expectation
         values and variances"""
         dev = qml.device("default.qubit", wires=3)
 
         a = 0.54
         b = -0.423
         c = 0.123
+        ops = [qml.RX(a, 0), qml.RY(b, 1), qml.CNOT([1, 2]), qml.RX(c, 2), qml.CNOT([0, 1])]
+        meas = [qml.var(qml.Z(0)), qml.expval(qml.Z(1)), qml.var(qml.Z(2))]
 
-        with qml.queuing.AnnotatedQueue() as q:
-            qml.RX(a, wires=0)
-            qml.RY(b, wires=1)
-            qml.CNOT(wires=[1, 2])
-            qml.RX(c, wires=2)
-            qml.CNOT(wires=[0, 1])
-            qml.var(qml.PauliZ(0))
-            qml.expval(qml.PauliZ(1))
-            qml.var(qml.PauliZ(2))
-
-        tape = qml.tape.QuantumScript.from_queue(q)
+        tape = qml.tape.QuantumScript(ops, meas)
         res = dev.execute(tape)
         expected = np.array(
             [
@@ -3244,25 +3189,22 @@ def test_expval_and_variance(self, tol):
         assert np.allclose(res, expected, atol=tol, rtol=0)
 
         # circuit jacobians
-        with pytest.raises(NotImplementedError, match="Broadcasting with multiple measurements"):
-            qml.gradients.param_shift(tape, broadcast=True)
-        # TODO: Uncomment the following when #2693 is resolved.
-        # tapes, fn = qml.gradients.param_shift(tape, broadcast=True)
-        # gradA = fn(dev.execute(tapes))
-
-        # tapes, fn = qml.gradients.finite_diff(tape)
-        # gradF = fn(dev.execute(tapes))
-        # ca, sa, cb, sb = np.cos(a), np.sin(a), np.cos(b), np.sin(b)
-        # c2c, s2c = np.cos(2 * c), np.sin(2 * c)
-        # expected = np.array(
-        # [
-        # [2 * ca * sa, -cb * sa, 0],
-        # [0, -ca * sb, 0.5 * (2 * cb * c2c * sb + s2c)],
-        # [0, 0, cb ** 2 * s2c],
-        # ]
-        # ).T
-        # assert gradA == pytest.approx(expected, abs=tol)
-        # assert gradF == pytest.approx(expected, abs=tol)
+        tapes, fn = qml.gradients.param_shift(tape, broadcast=True)
+        gradA = fn(dev.execute(tapes))
+
+        tapes, fn = qml.gradients.finite_diff(tape)
+        gradF = fn(dev.execute(tapes))
+        ca, sa, cb, sb = np.cos(a), np.sin(a), np.cos(b), np.sin(b)
+        c2c, s2c, s2b = np.cos(2 * c), np.sin(2 * c), np.sin(2 * b)
+        expected = np.array(
+            [
+                [2 * ca * sa, -cb * sa, 0],
+                [0, -ca * sb, 0.5 * (2 * cb * c2c * sb + s2b)],
+                [0, 0, cb**2 * s2c],
+            ]
+        ).T
+        assert gradA == pytest.approx(expected, abs=tol)
+        assert gradF == pytest.approx(expected, abs=tol)
 
     @pytest.mark.parametrize("state", [[1], [0, 1]])  # Basis state and state vector
     def test_projector_variance(self, state, tol):
@@ -3337,19 +3279,15 @@ def cost6(x):
             return [qml.probs([0, 1]), qml.probs([2, 3])]
 
         x = np.random.rand(3)
-        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)]
+        single_measure_circuits = [
+            qml.QNode(cost, dev) for cost in (cost1, cost2, cost4, cost5)
+        ] + [qml.QNode(cost, dev) for cost in (cost3, cost6)]
+        expected_shapes = [(3,), (1, 3), (4, 3), (1, 4, 3), (2, 3), (2, 4, 3)]
 
-        for c, exp_shape in zip(single_measure_circuits, [(3,), (1, 3), (4, 3), (1, 4, 3)]):
+        for c, exp_shape in zip(single_measure_circuits, expected_shapes):
             grad = qml.gradients.param_shift(c, broadcast=True)(x)
             assert qml.math.shape(grad) == exp_shape
 
-        for c in multi_measure_circuits:
-            with pytest.raises(
-                NotImplementedError, match="Broadcasting with multiple measurements"
-            ):
-                qml.gradients.param_shift(c, broadcast=True)(x)
-
 
 @pytest.mark.parametrize(
     "broadcast, expected", [(False, (5, [None] * 5)), (True, (3, [None, 2, 2]))]
@@ -3593,15 +3531,10 @@ def test_multiple_hamiltonians(self, mocker, tol, broadcast):
         expected = [-c * np.sin(x) * np.sin(y) + np.cos(x) * (a + b * np.sin(y)), d * np.cos(x)]
         assert np.allclose(res, expected, atol=tol, rtol=0)
 
-        if broadcast:
-            with pytest.raises(
-                NotImplementedError, match="Broadcasting with multiple measurements"
-            ):
-                tapes, fn = qml.gradients.param_shift(tape, broadcast=broadcast)
-            return
         tapes, fn = qml.gradients.param_shift(tape, broadcast=broadcast)
-        # two shifts per rotation gate, one circuit per trainable H term
-        assert len(tapes) == 2 * 2 + 3
+        # two shifts per rotation gate (in one batched tape if broadcasting),
+        # one circuit per trainable H term
+        assert len(tapes) == 2 * (1 if broadcast else 2) + 3
         spy.assert_called()
 
         res = fn(dev.execute(tapes))
@@ -3674,12 +3607,6 @@ def test_autograd(self, dev_name, tol, broadcast):
         weights = np.array([0.4, 0.5], requires_grad=True)
         dev = qml.device(dev_name, wires=2)
 
-        if broadcast:
-            with pytest.raises(
-                NotImplementedError, match="Broadcasting with multiple measurements"
-            ):
-                res = self.cost_fn(weights, coeffs1, coeffs2, dev, broadcast)
-            return
         res = self.cost_fn(weights, coeffs1, coeffs2, dev, broadcast)
         expected = self.cost_fn_expected(weights, coeffs1, coeffs2)
         assert np.allclose(res, np.array(expected), atol=tol, rtol=0)
@@ -3704,12 +3631,12 @@ def test_tf(self, dev_name, tol, broadcast):
 
         dev = qml.device(dev_name, wires=2)
 
-        if broadcast:
+        # Old op math with old device API does not support broadcasting
+        if broadcast and "tf" in dev_name and not qml.operation.active_new_opmath():
             with pytest.raises(
-                NotImplementedError, match="Broadcasting with multiple measurements"
+                NotImplementedError, match="Hamiltonians .* together with parameter broadcasting"
             ):
-                with tf.GradientTape() as _:
-                    self.cost_fn(weights, coeffs1, coeffs2, dev, broadcast)
+                _ = self.cost_fn(weights, coeffs1, coeffs2, dev, broadcast)
             return
         with tf.GradientTape() as _:
             jac = self.cost_fn(weights, coeffs1, coeffs2, dev, broadcast)
@@ -3739,9 +3666,10 @@ def test_torch(self, dev_name, tol, broadcast):
 
         dev = qml.device(dev_name, wires=2)
 
-        if broadcast:
+        # Old op math with old device API does not support broadcasting
+        if broadcast and "torch" in dev_name and not qml.operation.active_new_opmath():
             with pytest.raises(
-                NotImplementedError, match="Broadcasting with multiple measurements"
+                NotImplementedError, match="Hamiltonians .* together with parameter broadcasting"
             ):
                 res = self.cost_fn(weights, coeffs1, coeffs2, dev, broadcast)
             return
@@ -3752,7 +3680,6 @@ def test_torch(self, dev_name, tol, broadcast):
         res = tuple(tuple(_r.detach() for _r in r) for r in res)
         assert np.allclose(res, expected, atol=tol, rtol=0)
 
-        pytest.xfail("does not work with new return system")
         # second derivative wrt to Hamiltonian coefficients should be zero
         # hess = torch.autograd.functional.jacobian(
         #     lambda *args: self.cost_fn(*args, dev, broadcast), (weights, coeffs1, coeffs2)
@@ -3774,9 +3701,10 @@ def test_jax(self, dev_name, tol, broadcast):
         weights = jnp.array([0.4, 0.5])
         dev = qml.device(dev_name, wires=2)
 
-        if broadcast:
+        # Old op math with old device API does not support broadcasting
+        if broadcast and "jax" in dev_name and not qml.operation.active_new_opmath():
             with pytest.raises(
-                NotImplementedError, match="Broadcasting with multiple measurements"
+                NotImplementedError, match="Hamiltonians .* together with parameter broadcasting"
             ):
                 res = self.cost_fn(weights, coeffs1, coeffs2, dev, broadcast)
             return
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
@@ -118,12 +118,6 @@ def test_no_trainable_params_tape(self, broadcast):
         # TODO: remove once #2155 is resolved
         tape.trainable_params = []
 
-        if broadcast:
-            match_ = "Broadcasting with shot vectors is not supported yet"
-            with pytest.raises(NotImplementedError, match=match_):
-                g_tapes, fn = qml.gradients.param_shift(tape, broadcast=broadcast)
-            return
-
         with pytest.warns(UserWarning, match="gradient of a tape with no trainable parameters"):
             g_tapes, fn = qml.gradients.param_shift(tape, broadcast=broadcast)
         all_res = fn(qml.execute(g_tapes, dev, None))
@@ -256,25 +250,25 @@ def test_all_zero_diff_methods_multiple_returns_tape(self):
             assert result[1][2].shape == (4,)
             assert np.allclose(result[1][2], 0)
 
-    # TODO: uncomment when QNode decorator uses new qml.execute pipeline
-    # @pytest.mark.parametrize("broadcast", [True, False])
-    # def test_all_zero_diff_methods(self, broadcast):
-    #     """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, shots=default_shot_vector)
+    @pytest.mark.parametrize("broadcast", [True, False])
+    def test_all_zero_diff_methods(self, broadcast):
+        """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, shots=default_shot_vector)
 
-    #     @qml.qnode(dev)
-    #     def circuit(params):
-    #         qml.Rot(*params, wires=0)
-    #         return qml.probs([2, 3])
+        @qml.qnode(dev)
+        def circuit(params):
+            qml.Rot(*params, wires=0)
+            return qml.probs([2, 3])
 
-    #     params = np.array([0.5, 0.5, 0.5], requires_grad=True)
+        params = np.array([0.5, 0.5, 0.5], requires_grad=True)
+        circuit.construct((params,), {})
 
-    #     result = qml.gradients.param_shift(circuit)(params)
-    #     assert np.allclose(result, np.zeros((4, 3)), atol=0, rtol=0)
+        result = qml.gradients.param_shift(circuit)(params)
+        assert np.allclose(result, np.zeros((4, 3)), atol=0, rtol=0)
 
-    #     tapes, _ = qml.gradients.param_shift(circuit.tape, broadcast=broadcast)
-    #     assert tapes == []
+        tapes, _ = qml.gradients.param_shift(circuit.tape, broadcast=broadcast)
+        assert tapes == []
 
     @pytest.mark.parametrize("ops_with_custom_recipe", [[0], [1], [0, 1]])
     def test_recycled_unshifted_tape(self, ops_with_custom_recipe):
@@ -498,24 +492,25 @@ class NewOp(qml.operation.Operation):
                 qml.gradients.param_shift(tape)
 
 
-# TODO: add test class for parameter broadcasting - both when using broadcasting in the transform
-# via `broadcast=True`, and when using the transform on an already broadcasted tape, with
-# `broadcast=False` (as double broadcasting is not supported).
+# TODO: add test class for parameter broadcasting - when using the transform on an already broadcasted
+# tape with `broadcast=False` (as double broadcasting is not supported).
 
 
 # The first line of pylint disable is for cost1 through cost6
 # pylint: disable=no-self-argument, not-an-iterable
 # pylint: disable=too-many-public-methods
 @pytest.mark.slow
+@pytest.mark.parametrize("broadcast", [False, True])
 class TestParameterShiftRule:
     """Unit tests for the param_shift function used with a device that has a
     shot vector defined"""
 
     @pytest.mark.parametrize("theta", angles)
-    @pytest.mark.parametrize("shift", [np.pi / 2, 0.3, np.sqrt(2)])
+    @pytest.mark.parametrize("shift", [np.pi / 2, 0.3])
     @pytest.mark.parametrize("G", [qml.RX, qml.RY, qml.RZ, qml.PhaseShift])
-    def test_pauli_rotation_gradient(self, mocker, G, theta, shift):
+    def test_pauli_rotation_gradient(self, mocker, G, theta, shift, broadcast):
         """Tests that the automatic gradients of Pauli rotations are correct."""
+        # pylint: disable=too-many-arguments
 
         np.random.seed(824)
         spy = mocker.spy(qml.gradients.parameter_shift, "_get_operation_recipe")
@@ -530,8 +525,8 @@ def test_pauli_rotation_gradient(self, mocker, G, theta, shift):
         tape = qml.tape.QuantumScript.from_queue(q, shots=shot_vec)
         tape.trainable_params = {1}
 
-        tapes, fn = qml.gradients.param_shift(tape, shifts=[(shift,)])
-        assert len(tapes) == 2
+        tapes, fn = qml.gradients.param_shift(tape, shifts=[(shift,)], broadcast=broadcast)
+        assert len(tapes) == (1 if broadcast else 2)
 
         autograd_val = fn(dev.execute(tapes))
 
@@ -541,7 +536,7 @@ def test_pauli_rotation_gradient(self, mocker, G, theta, shift):
         shot_vec_manual_res = dev.execute([tape_fwd, tape_bwd])
 
         # Parameter axis is the first - reorder the results from execute
-        shot_vec_len = len(many_shots_shot_vector)
+        shot_vec_len = len(shot_vec)
         shot_vec_manual_res = [
             tuple(comp[l] for comp in shot_vec_manual_res) for l in range(shot_vec_len)
         ]
@@ -560,8 +555,8 @@ def test_pauli_rotation_gradient(self, mocker, G, theta, shift):
             assert np.allclose(a_val, n_val, atol=finite_diff_tol, rtol=0)
 
     @pytest.mark.parametrize("theta", angles)
-    @pytest.mark.parametrize("shift", [np.pi / 2, 0.3, np.sqrt(2)])
-    def test_Rot_gradient(self, mocker, theta, shift):
+    @pytest.mark.parametrize("shift", [np.pi / 2, 0.3])
+    def test_Rot_gradient(self, mocker, theta, shift, broadcast):
         """Tests that the automatic gradient of an arbitrary Euler-angle-parameterized gate is correct."""
         spy = mocker.spy(qml.gradients.parameter_shift, "_get_operation_recipe")
 
@@ -577,9 +572,10 @@ def test_Rot_gradient(self, mocker, theta, shift):
         tape = qml.tape.QuantumScript.from_queue(q, shots=shot_vec)
         tape.trainable_params = {1, 2, 3}
 
-        tapes, fn = qml.gradients.param_shift(tape, shifts=[(shift,)] * 3)
+        tapes, fn = qml.gradients.param_shift(tape, shifts=[(shift,)] * 3, broadcast=broadcast)
         num_params = len(tape.trainable_params)
-        assert len(tapes) == 2 * num_params
+        tapes_per_param = 1 if broadcast else 2
+        assert len(tapes) == tapes_per_param * num_params
 
         autograd_val = fn(dev.execute(tapes))
         assert isinstance(autograd_val, tuple)
@@ -619,7 +615,7 @@ def test_Rot_gradient(self, mocker, theta, shift):
             assert np.allclose(a_val, n_val, atol=finite_diff_tol, rtol=0)
 
     @pytest.mark.parametrize("G", [qml.CRX, qml.CRY, qml.CRZ])
-    def test_controlled_rotation_gradient(self, G):
+    def test_controlled_rotation_gradient(self, G, broadcast):
         """Test gradient of controlled rotation gates"""
         shot_vec = many_shots_shot_vector
         dev = qml.device("default.qubit", wires=2, shots=shot_vec)
@@ -636,11 +632,12 @@ def test_controlled_rotation_gradient(self, G):
         res = dev.execute(tape)
         assert np.allclose(res, -np.cos(b / 2), atol=shot_vec_tol, rtol=0)
 
-        tapes, fn = qml.gradients.param_shift(tape)
+        tapes, fn = qml.gradients.param_shift(tape, broadcast=broadcast)
+        assert len(tapes) == (1 if broadcast else 4)
         grad = fn(dev.execute(tapes))
         expected = np.sin(b / 2) / 2
         assert isinstance(grad, tuple)
-        assert len(grad) == len(many_shots_shot_vector)
+        assert len(grad) == len(shot_vec)
         assert np.allclose(grad, expected, atol=shot_vec_tol, rtol=0)
 
         tapes, fn = qml.gradients.finite_diff(tape, h=h_val)
@@ -649,7 +646,7 @@ def test_controlled_rotation_gradient(self, G):
             assert np.allclose(a_val, n_val, atol=finite_diff_tol, rtol=0)
 
     @pytest.mark.parametrize("theta", angles)
-    def test_CRot_gradient(self, theta):
+    def test_CRot_gradient(self, theta, broadcast):
         """Tests that the automatic gradient of an arbitrary controlled Euler-angle-parameterized
         gate is correct."""
         shot_vec = tuple([1000000] * 2)
@@ -668,8 +665,9 @@ def test_CRot_gradient(self, theta):
         expected = -np.cos(b / 2) * np.cos(0.5 * (a + c))
         assert np.allclose(res, expected, atol=shot_vec_tol, rtol=0)
 
-        tapes, fn = qml.gradients.param_shift(tape)
-        assert len(tapes) == 4 * len(tape.trainable_params)
+        tapes, fn = qml.gradients.param_shift(tape, broadcast=broadcast)
+        tapes_per_param = 1 if broadcast else 4
+        assert len(tapes) == tapes_per_param * len(tape.trainable_params)
 
         grad = fn(dev.execute(tapes))
         expected = np.array(
@@ -693,7 +691,7 @@ def test_CRot_gradient(self, theta):
         for a_val, n_val in zip(grad, numeric_val):
             assert np.allclose(a_val, n_val, atol=finite_diff_tol, rtol=0)
 
-    def test_gradients_agree_finite_differences(self):
+    def test_gradients_agree_finite_differences(self, broadcast):
         """Tests that the parameter-shift rule agrees with the first and second
         order finite differences"""
         params = np.array([0.1, -1.6, np.pi / 5])
@@ -722,7 +720,7 @@ def test_gradients_agree_finite_differences(self):
             strategy="center",
             h=h_val,
         )
-        grad_A = grad_fn(tape, dev)
+        grad_A = grad_fn(tape, dev, broadcast=broadcast)
 
         # gradients computed with different methods must agree
         for a_val, n_val in zip(grad_A, grad_F1):
@@ -730,7 +728,7 @@ def test_gradients_agree_finite_differences(self):
         for a_val, n_val in zip(grad_A, grad_F2):
             assert np.allclose(a_val, n_val, atol=finite_diff_tol, rtol=0)
 
-    def test_variance_gradients_agree_finite_differences(self):
+    def test_variance_gradients_agree_finite_differences(self, broadcast):
         """Tests that the variance parameter-shift rule agrees with the first and second
         order finite differences"""
         params = np.array([0.1, -1.6, np.pi / 5])
@@ -760,7 +758,7 @@ def test_variance_gradients_agree_finite_differences(self):
             strategy="center",
             h=h_val,
         )
-        grad_A = grad_fn(tape, dev)
+        grad_A = grad_fn(tape, dev, broadcast=broadcast)
 
         # gradients computed with different methods must agree
         for idx1, _grad_A in enumerate(grad_A):
@@ -770,8 +768,10 @@ def test_variance_gradients_agree_finite_differences(self):
 
     @pytest.mark.autograd
     @pytest.mark.parametrize("dev_name", ["default.qubit", "default.qubit.autograd"])
-    def test_fallback(self, dev_name, mocker):
+    def test_fallback(self, dev_name, mocker, broadcast):
         """Test that fallback gradient functions are correctly used"""
+        if broadcast and dev_name == "default.qubit.autograd":
+            pytest.xfail(reason="Return types + autograd + broadcasting does not work")
         spy = mocker.spy(qml.gradients, "finite_diff")
         dev = qml.device(dev_name, wires=3, shots=fallback_shot_vec)
         execute_fn = dev.execute if dev_name == "default.qubit" else dev.batch_execute
@@ -791,8 +791,9 @@ def cost_fn(params):
 
             tape = qml.tape.QuantumScript.from_queue(q, shots=fallback_shot_vec)
             finite_diff = partial(qml.gradients.finite_diff, h=h_val)
-            tapes, fn = param_shift(tape, fallback_fn=finite_diff)
-            assert len(tapes) == 5
+            tapes, fn = param_shift(tape, fallback_fn=finite_diff, broadcast=broadcast)
+            # Saving one tape in the parameter-shift part of the computation
+            assert len(tapes) == (4 if broadcast else 5)
 
             # check that the fallback method was called for the specified argnums
             spy.assert_called()
@@ -828,8 +829,10 @@ def cost_fn(params):
 
     @pytest.mark.autograd
     @pytest.mark.parametrize("dev_name", ["default.qubit", "default.qubit.autograd"])
-    def test_fallback_single_meas(self, dev_name, mocker):
+    def test_fallback_single_meas(self, dev_name, mocker, broadcast):
         """Test that fallback gradient functions are correctly used for a single measurement."""
+        if broadcast and dev_name == "default.qubit.autograd":
+            pytest.xfail(reason="Return types + autograd + broadcasting does not work")
         spy = mocker.spy(qml.gradients, "finite_diff")
         shot_vec = tuple([1000000] * 4)
         dev = qml.device(dev_name, wires=3, shots=shot_vec)
@@ -847,8 +850,8 @@ def cost_fn(params):
 
             tape = qml.tape.QuantumScript.from_queue(q, shots=shot_vec)
             finite_diff = partial(qml.gradients.finite_diff, h=h_val)
-            tapes, fn = param_shift(tape, fallback_fn=finite_diff)
-            assert len(tapes) == 4
+            tapes, fn = param_shift(tape, fallback_fn=finite_diff, broadcast=broadcast)
+            assert len(tapes) == (3 if broadcast else 4)
 
             # check that the fallback method was called for the specified argnums
             spy.assert_called()
@@ -874,9 +877,13 @@ def cost_fn(params):
     @pytest.mark.parametrize("RX, RY, argnum", [(RX_with_F, qml.RY, 0), (qml.RX, RY_with_F, 1)])
     @pytest.mark.parametrize("dev_name", ["default.qubit", "default.qubit.autograd"])
     def test_fallback_probs(
-        self, dev_name, RX, RY, argnum, mocker
+        self, dev_name, RX, RY, argnum, mocker, broadcast
     ):  # pylint:disable=too-many-arguments
         """Test that fallback gradient functions are correctly used with probs"""
+        if broadcast and dev_name == "default.qubit.autograd":
+            pytest.xfail(
+                reason="Return types + autograd + old device API + broadcasting does not work"
+            )
         spy = mocker.spy(qml.gradients, "finite_diff")
         dev = qml.device(dev_name, wires=3, shots=fallback_shot_vec)
         execute_fn = dev.execute if dev_name == "default.qubit" else dev.batch_execute
@@ -895,8 +902,8 @@ def cost_fn(params):
 
             tape = qml.tape.QuantumScript.from_queue(q, shots=fallback_shot_vec)
             finite_diff = partial(qml.gradients.finite_diff, h=h_val)
-            tapes, fn = param_shift(tape, fallback_fn=finite_diff)
-            assert len(tapes) == 4
+            tapes, fn = param_shift(tape, fallback_fn=finite_diff, broadcast=broadcast)
+            assert len(tapes) == (3 if broadcast else 4)
 
             # check that the fallback method was called for the specified argnums
             spy.assert_called()
@@ -964,9 +971,11 @@ def cost_fn(params):
 
     @pytest.mark.autograd
     @pytest.mark.parametrize("dev_name", ["default.qubit", "default.qubit.autograd"])
-    def test_all_fallback(self, dev_name, mocker):
+    def test_all_fallback(self, dev_name, mocker, broadcast):
         """Test that *only* the fallback logic is called if no parameters
         support the parameter-shift rule"""
+        if broadcast:
+            pytest.skip("No need to test fallback with broadcasting")
         spy_fd = mocker.spy(qml.gradients, "finite_diff")
         spy_ps = mocker.spy(qml.gradients.parameter_shift, "expval_param_shift")
 
@@ -1004,7 +1013,7 @@ def test_all_fallback(self, dev_name, mocker):
             assert np.allclose(res[0], expected[0], atol=fallback_shot_vec, rtol=0)
             assert np.allclose(res[1], expected[1], atol=fallback_shot_vec, rtol=0)
 
-    def test_single_expectation_value(self):
+    def test_single_expectation_value(self, broadcast):
         """Tests correct output shape and evaluation for a tape
         with a single expval output"""
         shot_vec = many_shots_shot_vector
@@ -1019,8 +1028,9 @@ def test_single_expectation_value(self):
             qml.expval(qml.PauliZ(0) @ qml.PauliX(1))
 
         tape = qml.tape.QuantumScript.from_queue(q, shots=shot_vec)
-        tapes, fn = qml.gradients.param_shift(tape)
-        assert len(tapes) == 4
+        tapes, fn = qml.gradients.param_shift(tape, broadcast=broadcast)
+        tapes_per_param = 1 if broadcast else 2
+        assert len(tapes) == tapes_per_param * 2
 
         all_res = fn(dev.execute(tapes))
 
@@ -1036,7 +1046,7 @@ def test_single_expectation_value(self):
             assert np.allclose(res[0], expected[0], atol=shot_vec_tol, rtol=0)
             assert np.allclose(res[1], expected[1], atol=shot_vec_tol, rtol=0)
 
-    def test_multiple_expectation_values(self):
+    def test_multiple_expectation_values(self, broadcast):
         """Tests correct output shape and evaluation for a tape
         with multiple expval outputs"""
         shot_vec = many_shots_shot_vector
@@ -1052,8 +1062,9 @@ def test_multiple_expectation_values(self):
             qml.expval(qml.PauliX(1))
 
         tape = qml.tape.QuantumScript.from_queue(q, shots=shot_vec)
-        tapes, fn = qml.gradients.param_shift(tape)
-        assert len(tapes) == 4
+        tapes, fn = qml.gradients.param_shift(tape, broadcast=broadcast)
+        tapes_per_param = 1 if broadcast else 2
+        assert len(tapes) == tapes_per_param * 2
 
         all_res = fn(dev.execute(tapes))
         assert len(all_res) == len(many_shots_shot_vector)
@@ -1069,7 +1080,7 @@ def test_multiple_expectation_values(self):
             assert np.allclose(res[0], expected[0], atol=shot_vec_tol, rtol=0)
             assert np.allclose(res[1], expected[1], atol=shot_vec_tol, rtol=0)
 
-    def test_var_expectation_values(self):
+    def test_var_expectation_values(self, broadcast):
         """Tests correct output shape and evaluation for a tape
         with expval and var outputs"""
         shot_vec = many_shots_shot_vector
@@ -1085,8 +1096,8 @@ def test_var_expectation_values(self):
             qml.var(qml.PauliX(1))
 
         tape = qml.tape.QuantumScript.from_queue(q, shots=shot_vec)
-        tapes, fn = qml.gradients.param_shift(tape)
-        assert len(tapes) == 5
+        tapes, fn = qml.gradients.param_shift(tape, broadcast=broadcast)
+        assert len(tapes) == (3 if broadcast else 5)
 
         all_res = fn(dev.execute(tapes))
         assert len(all_res) == len(many_shots_shot_vector)
@@ -1102,7 +1113,7 @@ def test_var_expectation_values(self):
             for a, e in zip(res, expected):
                 assert np.allclose(np.squeeze(np.stack(a)), e, atol=shot_vec_tol, rtol=0)
 
-    def test_prob_expectation_values(self):
+    def test_prob_expectation_values(self, broadcast):
         """Tests correct output shape and evaluation for a tape
         with prob and expval outputs"""
 
@@ -1119,8 +1130,9 @@ def test_prob_expectation_values(self):
             qml.probs(wires=[0, 1])
 
         tape = qml.tape.QuantumScript.from_queue(q, shots=shot_vec)
-        tapes, fn = qml.gradients.param_shift(tape)
-        assert len(tapes) == 4
+        tapes, fn = qml.gradients.param_shift(tape, broadcast=broadcast)
+        tapes_per_param = 1 if broadcast else 2
+        assert len(tapes) == 2 * tapes_per_param
 
         res = fn(dev.execute(tapes))
         assert isinstance(res, tuple)
@@ -1191,7 +1203,7 @@ def test_prob_expectation_values(self):
             assert isinstance(r_to_check, np.ndarray)
             assert r_to_check.shape == (4,)
 
-    def test_involutory_variance_single_param(self):
+    def test_involutory_variance_single_param(self, broadcast):
         """Tests qubit observables that are involutory with a single trainable param"""
         shot_vec = many_shots_shot_vector
         dev = qml.device("default.qubit", wires=1, shots=shot_vec)
@@ -1208,13 +1220,14 @@ def test_involutory_variance_single_param(self):
             assert np.allclose(r, expected, atol=shot_vec_tol, rtol=0)
 
         # circuit jacobians
-        tapes, fn = qml.gradients.param_shift(tape)
+        tapes, fn = qml.gradients.param_shift(tape, broadcast=broadcast)
         gradA = fn(dev.execute(tapes))
         for _gA in gradA:
             assert isinstance(_gA, np.ndarray)
             assert _gA.shape == ()
 
-        assert len(tapes) == 1 + 2 * 1
+        tapes_per_param = 1 if broadcast else 2
+        assert len(tapes) == 1 + tapes_per_param * 1
 
         tapes, fn = qml.gradients.finite_diff(tape, h=h_val)
         all_gradF = fn(dev.execute(tapes))
@@ -1228,7 +1241,7 @@ def test_involutory_variance_single_param(self):
         for _gA in gradA:
             assert _gA == pytest.approx(expected, abs=shot_vec_tol)
 
-    def test_involutory_variance_multi_param(self):
+    def test_involutory_variance_multi_param(self, broadcast):
         """Tests qubit observables that are involutory with multiple trainable params"""
         shot_vec = many_shots_shot_vector
         dev = qml.device("default.qubit", wires=1, shots=shot_vec)
@@ -1248,7 +1261,10 @@ def test_involutory_variance_multi_param(self):
         assert np.allclose(res, expected, atol=shot_vec_tol, rtol=0)
 
         # circuit jacobians
-        tapes, fn = qml.gradients.param_shift(tape)
+        tapes, fn = qml.gradients.param_shift(tape, broadcast=broadcast)
+        tapes_per_param = 1 if broadcast else 2
+        assert len(tapes) == 1 + tapes_per_param * 2
+
         all_res = fn(dev.execute(tapes))
         assert len(all_res) == len(many_shots_shot_vector)
         assert isinstance(all_res, tuple)
@@ -1260,12 +1276,11 @@ def test_involutory_variance_multi_param(self):
             assert isinstance(gradA[1], np.ndarray)
             assert gradA[1].shape == ()
 
-            assert len(tapes) == 1 + 2 * 2
-
         tapes, fn = qml.gradients.finite_diff(tape, h=h_val)
+        assert len(tapes) == 3
+
         all_Fres = fn(dev.execute(tapes))
         for gradF, gradA in zip(all_Fres, all_res):
-            assert len(tapes) == 3
 
             expected = 2 * np.sin(a + b) * np.cos(a + b)
             assert gradF[0] == pytest.approx(expected, abs=finite_diff_tol)
@@ -1275,7 +1290,7 @@ def test_involutory_variance_multi_param(self):
             assert gradA[1] == pytest.approx(expected, abs=finite_diff_tol)
 
     @flaky(max_runs=5)
-    def test_non_involutory_variance_single_param(self):
+    def test_non_involutory_variance_single_param(self, broadcast):
         """Tests a qubit Hermitian observable that is not involutory with a single trainable parameter"""
         shot_vec = many_shots_shot_vector
         dev = qml.device("default.qubit", wires=1, shots=shot_vec)
@@ -1295,9 +1310,11 @@ def test_non_involutory_variance_single_param(self):
             assert np.allclose(r, expected, atol=_herm_shot_vec_tol, rtol=0)
 
         # circuit jacobians
-        tapes, fn = qml.gradients.param_shift(tape)
+        tapes, fn = qml.gradients.param_shift(tape, broadcast=broadcast)
+        tapes_per_param = 1 if broadcast else 2
+        # forward pass + tapes_per_param * (expval+expval^2) * num_params
+        assert len(tapes) == 1 + tapes_per_param * 2 * 1
         gradA = fn(dev.execute(tapes))
-        assert len(tapes) == 1 + 4 * 1
 
         tapes, fn = qml.gradients.finite_diff(tape, h=h_val)
         all_gradF = fn(dev.execute(tapes))
@@ -1314,7 +1331,7 @@ def test_non_involutory_variance_single_param(self):
             assert qml.math.allclose(gradF, expected, atol=2 * _herm_shot_vec_tol)
 
     @flaky(max_runs=5)
-    def test_non_involutory_variance_multi_param(self):
+    def test_non_involutory_variance_multi_param(self, broadcast):
         """Tests a qubit Hermitian observable that is not involutory with multiple trainable parameters"""
         shot_vec = many_shots_shot_vector
         dev = qml.device("default.qubit", wires=1, shots=shot_vec)
@@ -1337,7 +1354,11 @@ def test_non_involutory_variance_multi_param(self):
             assert np.allclose(res, expected, atol=herm_shot_vec_tol, rtol=0)
 
         # circuit jacobians
-        tapes, fn = qml.gradients.param_shift(tape)
+        tapes, fn = qml.gradients.param_shift(tape, broadcast=broadcast)
+        tapes_per_param = 1 if broadcast else 2
+        # forward pass + tapes_per_param * (expval+expval^2) * num_params
+        assert len(tapes) == 1 + tapes_per_param * 2 * 2
+
         all_res = fn(dev.execute(tapes))
         assert len(all_res) == len(many_shots_shot_vector)
         assert isinstance(all_res, tuple)
@@ -1351,16 +1372,16 @@ def test_non_involutory_variance_multi_param(self):
 
             assert isinstance(gradA[1], np.ndarray)
             assert gradA[1].shape == ()
-            assert len(tapes) == 1 + 4 * 2
             assert gradA[0] == pytest.approx(expected, abs=herm_shot_vec_tol)
             assert gradA[1] == pytest.approx(expected, abs=herm_shot_vec_tol)
 
         tapes, fn = qml.gradients.finite_diff(tape, h=h_val)
+        assert len(tapes) == 3
+
         all_gradF = fn(dev.execute(tapes))
         assert len(all_gradF) == len(many_shots_shot_vector)
         assert isinstance(all_gradF, tuple)
         for gradF in all_gradF:
-            assert len(tapes) == 3
 
             # Note: the tolerances here are significantly higher than in usual tests
             # due to the stochasticity of the test case
@@ -1368,7 +1389,7 @@ def test_non_involutory_variance_multi_param(self):
             assert qml.math.allclose(gradF[1], expected, atol=1.5)
 
     @flaky(max_runs=8)
-    def test_involutory_and_noninvolutory_variance_single_param(self):
+    def test_involutory_and_noninvolutory_variance_single_param(self, broadcast):
         """Tests a qubit Hermitian observable that is not involutory alongside
         an involutory observable when there's a single trainable parameter."""
         shot_vec = tuple([1000000] * 3)
@@ -1392,9 +1413,11 @@ def test_involutory_and_noninvolutory_variance_single_param(self):
             assert np.allclose(r, expected, atol=_herm_shot_vec_tol, rtol=0)
 
         # circuit jacobians
-        tapes, fn = qml.gradients.param_shift(tape)
+        tapes, fn = qml.gradients.param_shift(tape, broadcast=broadcast)
+        tapes_per_param = 1 if broadcast else 2
+        assert len(tapes) == 1 + 2 * tapes_per_param
+
         gradA = fn(dev.execute(tapes))
-        assert len(tapes) == 1 + 4
 
         tapes, fn = qml.gradients.finite_diff(tape, h=h_val)
         gradF = fn(dev.execute(tapes))
@@ -1420,7 +1443,7 @@ def test_involutory_and_noninvolutory_variance_single_param(self):
             assert shot_vec_result[1] == pytest.approx(expected[1], abs=_herm_shot_vec_tol)
 
     @flaky(max_runs=8)
-    def test_involutory_and_noninvolutory_variance_multi_param(self):
+    def test_involutory_and_noninvolutory_variance_multi_param(self, broadcast):
         """Tests a qubit Hermitian observable that is not involutory alongside
         an involutory observable."""
         shot_vec = many_shots_shot_vector
@@ -1443,7 +1466,10 @@ def test_involutory_and_noninvolutory_variance_multi_param(self):
             assert np.allclose(res_shot_item, expected, atol=_herm_shot_vec_tol, rtol=0)
 
         # circuit jacobians
-        tapes, fn = qml.gradients.param_shift(tape)
+        tapes, fn = qml.gradients.param_shift(tape, broadcast=broadcast)
+        tapes_per_param = 1 if broadcast else 2
+        assert len(tapes) == 1 + 2 * 2 * tapes_per_param
+
         gradA = fn(dev.execute(tapes))
 
         assert isinstance(gradA, tuple)
@@ -1458,8 +1484,6 @@ def test_involutory_and_noninvolutory_variance_multi_param(self):
                     assert isinstance(param_res, np.ndarray)
                     assert param_res.shape == ()
 
-        assert len(tapes) == 1 + 2 * 4
-
         tapes, fn = qml.gradients.finite_diff(tape, h=h_val)
         gradF = fn(dev.execute(tapes))
         assert len(tapes) == 1 + 2
@@ -1496,7 +1520,7 @@ def test_involutory_and_noninvolutory_variance_multi_param(self):
             assert np.allclose(shot_vec_result[1][1], expected[3], atol=1.5)
 
     @pytest.mark.parametrize("ind", [0, 1])
-    def test_var_and_probs_single_param(self, ind):
+    def test_var_and_probs_single_param(self, ind, broadcast):
         """Tests a qubit Hermitian observable that is not involutory alongside an involutory observable and probs when
         there's one trainable parameter."""
         shot_vec = many_shots_shot_vector
@@ -1525,7 +1549,9 @@ def test_var_and_probs_single_param(self, ind):
         tape.trainable_params = {ind}
 
         # circuit jacobians
-        tapes, fn = qml.gradients.param_shift(tape)
+        tapes, fn = qml.gradients.param_shift(tape, broadcast=broadcast)
+        tapes_per_param = 1 if broadcast else 2
+        assert len(tapes) == 1 + 2 * tapes_per_param
 
         all_res = fn(dev.execute(tapes))
         assert len(all_res) == len(many_shots_shot_vector)
@@ -1554,7 +1580,7 @@ def test_var_and_probs_single_param(self, ind):
             assert isinstance(gradA[2], np.ndarray)
             assert np.allclose(gradA[2], 0, atol=shot_vec_tol, rtol=0)
 
-    def test_var_and_probs_multi_params(self):
+    def test_var_and_probs_multi_params(self, broadcast):
         """Tests a qubit Hermitian observable that is not involutory alongside an involutory observable and probs when
         there are more trainable parameters."""
         shot_vec = many_shots_shot_vector
@@ -1583,7 +1609,9 @@ def test_var_and_probs_multi_params(self):
         tape.trainable_params = {0, 1, 2, 3}
 
         # circuit jacobians
-        tapes, fn = qml.gradients.param_shift(tape)
+        tapes, fn = qml.gradients.param_shift(tape, broadcast=broadcast)
+        assert len(tapes) == 1 + 4 * 2 * (1 if broadcast else 2)
+
         all_res = fn(dev.execute(tapes))
         assert len(all_res) == len(many_shots_shot_vector)
         assert isinstance(all_res, tuple)
@@ -1650,7 +1678,7 @@ def test_var_and_probs_multi_params(self):
             assert np.allclose(gradA[2][2], probs_expected[:, 0], atol=shot_vec_tol, rtol=0)
             assert np.allclose(gradA[2][3], probs_expected[:, 1], atol=shot_vec_tol, rtol=0)
 
-    def test_expval_and_variance_single_param(self):
+    def test_expval_and_variance_single_param(self, broadcast):
         """Test an expectation value and the variance of involutory and non-involutory observables work well with a
         single trainable parameter"""
         shot_vec = many_shots_shot_vector
@@ -1686,7 +1714,10 @@ def test_expval_and_variance_single_param(self):
         assert np.allclose(res, expected, atol=shot_vec_tol, rtol=0)
 
         # # circuit jacobians
-        tapes, fn = qml.gradients.param_shift(tape)
+        tapes, fn = qml.gradients.param_shift(tape, broadcast=broadcast)
+        tapes_per_param = 1 if broadcast else 2
+        assert len(tapes) == 1 + tapes_per_param
+
         all_res = fn(dev.execute(tapes))
 
         assert len(all_res) == len(many_shots_shot_vector)
@@ -1708,7 +1739,7 @@ def test_expval_and_variance_single_param(self):
             assert isinstance(gradF, tuple)
             assert gradF == pytest.approx(expected, abs=finite_diff_tol)
 
-    def test_expval_and_variance_multi_param(self):
+    def test_expval_and_variance_multi_param(self, broadcast):
         """Test an expectation value and the variance of involutory and non-involutory observables work well with
         multiple trainable parameters"""
         shot_vec = many_shots_shot_vector
@@ -1742,7 +1773,9 @@ def test_expval_and_variance_multi_param(self):
         assert np.allclose(res, expected, atol=shot_vec_tol, rtol=0)
 
         # # circuit jacobians
-        tapes, fn = qml.gradients.param_shift(tape)
+        tapes, fn = qml.gradients.param_shift(tape, broadcast=broadcast)
+        tapes_per_param = 1 if broadcast else 2
+        assert len(tapes) == 1 + 3 * tapes_per_param
         all_res = fn(dev.execute(tapes))
 
         assert len(all_res) == len(many_shots_shot_vector)
@@ -1773,7 +1806,7 @@ def test_expval_and_variance_multi_param(self):
             assert gradF == pytest.approx(expected, abs=finite_diff_tol)
 
     @pytest.mark.parametrize("state", [[1], [0, 1]])  # Basis state and state vector
-    def test_projector_variance(self, state):
+    def test_projector_variance(self, state, broadcast):
         """Test that the variance of a projector is correctly returned"""
         shot_vec = many_shots_shot_vector
         dev = qml.device("default.qubit", wires=2, shots=shot_vec)
@@ -1797,7 +1830,9 @@ def test_projector_variance(self, state):
             assert np.allclose(r, expected, atol=shot_vec_tol, rtol=0)
 
         # # circuit jacobians
-        tapes, fn = qml.gradients.param_shift(tape)
+        tapes, fn = qml.gradients.param_shift(tape, broadcast=broadcast)
+        tapes_per_param = 1 if broadcast else 2
+        assert len(tapes) == 1 + 2 * 2 * tapes_per_param
         all_res = fn(dev.execute(tapes))
 
         assert len(all_res) == len(many_shots_shot_vector)
@@ -1848,16 +1883,14 @@ 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,), np.ndarray),
+        (cost2, (1, 3), list),
+        (cost3, (2, 3), list),
     ]
 
-    @pytest.mark.parametrize("cost, expected_shape", costs_and_expected_expval)
-    def test_output_shape_matches_qnode_expval(self, cost, expected_shape):
+    @pytest.mark.parametrize("cost, expected_shape, exp_type", costs_and_expected_expval)
+    def test_output_shape_matches_qnode_expval(self, cost, expected_shape, exp_type, broadcast):
         """Test that the transform output shape matches that of the QNode."""
-        if cost.__name__ != "cost1":
-            pytest.xfail(reason="new return shape specification")
 
         shot_vec = many_shots_shot_vector
         dev = qml.device("default.qubit", wires=4, shots=shot_vec)
@@ -1865,45 +1898,44 @@ def test_output_shape_matches_qnode_expval(self, cost, expected_shape):
         x = np.random.rand(3)
         circuit = qml.QNode(cost, dev)
 
-        all_res = qml.gradients.param_shift(circuit)(x)
+        all_res = qml.gradients.param_shift(circuit, broadcast=broadcast)(x)
         assert len(all_res) == len(many_shots_shot_vector)
         assert isinstance(all_res, tuple)
 
         for res in all_res:
-            assert isinstance(res, np.ndarray)
-            assert res.shape == expected_shape
+            assert isinstance(res, exp_type)
+            assert qml.math.shape(res) == expected_shape
 
     costs_and_expected_probs = [
-        (cost4, (4, 3)),
-        (cost5, (4, 3)),
-        (cost6, (2, 4, 3)),
+        (cost4, (4, 3), np.ndarray),
+        (cost5, (1, 4, 3), list),
+        (cost6, (2, 4, 3), list),
     ]
 
-    @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, expected_shape, exp_type", costs_and_expected_probs)
+    def test_output_shape_matches_qnode_probs(self, cost, expected_shape, exp_type, broadcast):
         """Test that the transform output shape matches that of the QNode."""
-        if cost.__name__ != "cost4":
-            pytest.xfail(reason="wrong return shape specification")
-
         shot_vec = many_shots_shot_vector
         dev = qml.device("default.qubit", wires=4, shots=shot_vec)
 
         x = np.random.rand(3)
         circuit = qml.QNode(cost, dev)
 
-        all_res = qml.gradients.param_shift(circuit)(x)
+        all_res = qml.gradients.param_shift(circuit, broadcast=broadcast)(x)
         assert len(all_res) == len(many_shots_shot_vector)
         assert isinstance(all_res, tuple)
 
         for res in all_res:
-            assert isinstance(res, np.ndarray)
-            assert res.shape == expected_shape
+            assert isinstance(res, exp_type)
+            assert qml.math.shape(res) == expected_shape
 
     # TODO: revisit the following test when the Autograd interface supports
     # parameter-shift with the new return type system
-    def test_special_observable_qnode_differentiation(self):
+    def test_special_observable_qnode_differentiation(self, broadcast):
         """Test differentiation of a QNode on a device supporting a
         special observable that returns an object rather than a number."""
+        if broadcast:
+            pytest.skip("SpecialObject is not compatible with non-scalar objects.")
 
         class SpecialObject:
             """SpecialObject
@@ -1958,12 +1990,12 @@ def expval(self, observable, **kwargs):
 
         dev = DeviceSupporingSpecialObservable(wires=1, shots=None)
 
-        @qml.qnode(dev, diff_method="parameter-shift")
+        @qml.qnode(dev, diff_method="parameter-shift", broadcast=broadcast)
         def qnode(x):
             qml.RY(x, wires=0)
             return qml.expval(SpecialObservable(wires=0))
 
-        @qml.qnode(dev, diff_method="parameter-shift")
+        @qml.qnode(dev, diff_method="parameter-shift", broadcast=broadcast)
         def reference_qnode(x):
             qml.RY(x, wires=0)
             return qml.expval(qml.PauliZ(wires=0))
@@ -1972,7 +2004,7 @@ def reference_qnode(x):
         assert np.isclose(qnode(par).item().val, reference_qnode(par))
         assert np.isclose(qml.jacobian(qnode)(par).item().val, qml.jacobian(reference_qnode)(par))
 
-    def test_multi_measure_no_warning(self):
+    def test_multi_measure_no_warning(self, broadcast):
         """Test computing the gradient of a tape that contains multiple
         measurements omits no warnings."""
         import warnings
@@ -1991,7 +2023,7 @@ def test_multi_measure_no_warning(self):
 
         tape = qml.tape.QuantumScript.from_queue(q, shots=shot_vec)
         with warnings.catch_warnings(record=True) as record:
-            tapes, fn = qml.gradients.param_shift(tape)
+            tapes, fn = qml.gradients.param_shift(tape, broadcast=broadcast)
             fn(dev.execute(tapes))
 
         assert len(record) == 0