Repository: PennyLaneAI/pennylane
Pull request #6928 (2025-02-05 16:19:30)

Improve and tidy implementation of hadamard gradient
**Context:**

We are intending to make some improvements to the hadamard gradient to support variations of the differentiation method. To prepare us for these extensions, I've made some minor reworking.

**Description of the Change:**

* `Rot` operations are now decomposed instead of having bespoke logic.  It's basically the same resources at the end, and makes the code substantially cleaner.

* `hadamard_grad` now works with anything with a generator, instead of a manually managed list of simple operators.  This will set us to apply hadamard grad to more complicated hamiltonian time evolution problems.

* The postprocessing function is extracted to global scope.  This makes it more clear which variables are needed from the preprocessing step.

* The postprocessing is slightly reworked to have fewer branches, and to no longer need to reorder everything at the end. By populating the `grad` structure in the correct shape to begin with, we can avoid a potentially expensive and badly scaling step.

* The postprocessing now returns lists instead of tuples.  Our return-spec is supposed to be agnostic to lists versus tuples.  If lists made sense for creating the structure to begin with, we shouldn't have to go back and cast them.  That's an extra step we don't need, if we just allow the other components to expect a list.  Yes, I do generally prefer tuple's to lists, but in this case, the list makes sense.

**Benefits:**

More generic code that is more versatile and easier to maintain.

**Possible Drawbacks:**

**Related Shortcut Stories:**
[sc-84784]

github-actions[bot]: Hello. You may have forgotten to update the changelog!
Please edit [`doc/releases/changelog-dev.md`](/PennyLaneAI/pennylane/blob/master/doc/releases/changelog-dev.md) with:
* A one-to-two sentence description of the change. You may include a small working example for new features.
* A link back to this PR.
* Your name (or GitHub username) in the contributors section.
codecov[bot]: ## [Codecov](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/6928?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
Attention: Patch coverage is `98.07692%` with `2 lines` in your changes missing coverage. Please review.
> Project coverage is 99.58%. Comparing base [(`53df3dc`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/53df3dc89afea4b4c4a9a2f82534a1d85013629c?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) to head [(`2527025`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/252702516875b6fbe158cedc8269c13520264ba8?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).
> Report is 3 commits behind head on master.

| [Files with missing lines](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/6928?dropdown=coverage&src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | Patch % | Lines |
|---|---|---|
| [pennylane/gradients/hadamard\_gradient.py](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/6928?src=pr&el=tree&filepath=pennylane%2Fgradients%2Fhadamard_gradient.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2dyYWRpZW50cy9oYWRhbWFyZF9ncmFkaWVudC5weQ==) | 97.84% | [2 Missing :warning: ](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/6928?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    #6928      +/-   ##
==========================================
- Coverage   99.59%   99.58%   -0.01%     
==========================================
  Files         480      480              
  Lines       45496    45456      -40     
==========================================
- Hits        45311    45269      -42     
- Misses        185      187       +2     
```

</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/6928?dropdown=coverage&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).

JerryChen97: There's a PL Deprecation Warning for the last example of use details. We can mark this as a reminder for todo in the QA round.
```
PennyLaneDeprecationWarning: Specifying gradient keyword arguments ['aux_wire', 'device_wires'] is deprecated and will be removed in v0.42. Instead, please specify all arguments in the gradient_kwargs argument.
  warnings.warn(
```

Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index dbe8fb476..f4fe1b840 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -42,6 +42,9 @@
 
 <h3>Improvements 🛠</h3>
 
+* `qml.gradients.hadamard_grad` can now differentiate anything with a generator, and can accept circuits with non-commuting measurements.
+[(#6928)](https://github.com/PennyLaneAI/pennylane/pull/6928)
+
 * `Controlled` operators now have a full implementation of `sparse_matrix` that supports `wire_order` configuration.
   [(#6994)](https://github.com/PennyLaneAI/pennylane/pull/6994)
 
@@ -280,6 +283,10 @@
 
 <h3>Breaking changes 💔</h3>
 
+* `qml.gradients.gradient_transform.choose_trainable_params` has been renamed to `choose_trainable_param_indices`
+  to better reflect what it actually does.
+  [(#6928)](https://github.com/PennyLaneAI/pennylane/pull/6928)
+
 * `MultiControlledX` no longer accepts strings as control values.
   [(#6835)](https://github.com/PennyLaneAI/pennylane/pull/6835)
 
diff --git a/pennylane/gradients/finite_difference.py b/pennylane/gradients/finite_difference.py
index 985fb245c..2a1d8437a 100644
--- a/pennylane/gradients/finite_difference.py
+++ b/pennylane/gradients/finite_difference.py
@@ -38,7 +38,7 @@ from .gradient_transform import (
     _all_zero_grad,
     _no_trainable_grad,
     assert_no_trainable_tape_batching,
-    choose_trainable_params,
+    choose_trainable_param_indices,
     find_and_validate_gradient_methods,
 )
 
@@ -479,11 +479,11 @@ def finite_diff(
     if argnum is None and not tape.trainable_params:
         return _no_trainable_grad(tape)
 
-    trainable_params = choose_trainable_params(tape, argnum)
+    trainable_params_indices = choose_trainable_param_indices(tape, argnum)
     diff_methods = (
-        find_and_validate_gradient_methods(tape, "numeric", trainable_params)
+        find_and_validate_gradient_methods(tape, "numeric", trainable_params_indices)
         if validate_params
-        else {idx: "F" for idx in trainable_params}
+        else {idx: "F" for idx in trainable_params_indices}
     )
 
     if all(g == "0" for g in diff_methods.values()):
diff --git a/pennylane/gradients/gradient_transform.py b/pennylane/gradients/gradient_transform.py
index be3595a1d..d939de723 100644
--- a/pennylane/gradients/gradient_transform.py
+++ b/pennylane/gradients/gradient_transform.py
@@ -101,7 +101,7 @@ def assert_no_trainable_tape_batching(tape, transform_name):
             )
 
 
-def choose_trainable_params(tape, argnum=None):
+def choose_trainable_param_indices(tape, argnum=None):
     """Returns a list of trainable parameter indices in the tape.
 
     Chooses the subset of trainable parameters to compute the Jacobian for. The function
@@ -116,6 +116,17 @@ def choose_trainable_params(tape, argnum=None):
     Returns:
         list: list of the trainable parameter indices
 
+    Note that trainable param indices are a **double pointer**.
+
+    >>> tape = qml.tape.QuantumScript([qml.RX(0.0, 0), qml.RY(1.0, 0), qml.RZ(2.0, 0)], trainable_params=[1,2])
+    >>> chose_trainable_param_indices(tape, argnum=[0])
+    [0]
+    >>> tape.get_operation(0)
+    (RY(1.0, wires=[0]), 1, 0)
+
+    In this case ``[0]`` points to the ``RY`` parameter. ``0`` selects into ``tape.trainable_params``,
+    which selects into ``tape.data``.
+
     """
 
     if argnum is None:
@@ -382,13 +393,13 @@ def _contract_qjac_with_cjac(qjac, cjac, tape):
     num_measurements = len(tape.measurements)
     has_partitioned_shots = tape.shots.has_partitioned_shots
 
-    if isinstance(qjac, tuple) and len(qjac) == 1:
+    if isinstance(qjac, (tuple, list)) and len(qjac) == 1:
         qjac = qjac[0]
 
-    if isinstance(cjac, tuple) and len(cjac) == 1:
+    if isinstance(cjac, (tuple, list)) and len(cjac) == 1:
         cjac = cjac[0]
 
-    cjac_is_tuple = isinstance(cjac, tuple)
+    cjac_is_tuple = isinstance(cjac, (tuple, list))
 
     multi_meas = num_measurements > 1
 
@@ -402,7 +413,7 @@ def _contract_qjac_with_cjac(qjac, cjac, tape):
             _qjac = _qjac[0]
         if has_partitioned_shots:
             _qjac = _qjac[0]
-        single_tape_param = not isinstance(_qjac, tuple)
+        single_tape_param = not isinstance(_qjac, (tuple, list))
 
     if single_tape_param:
         # Without dimension (e.g. expval) or with dimension (e.g. probs)
diff --git a/pennylane/gradients/hadamard_gradient.py b/pennylane/gradients/hadamard_gradient.py
index 5335bd0a0..fc125d1b4 100644
--- a/pennylane/gradients/hadamard_gradient.py
+++ b/pennylane/gradients/hadamard_gradient.py
@@ -16,80 +16,44 @@ This module contains functions for computing the Hadamard-test gradient
 of a qubit-based quantum tape.
 """
 from functools import partial
+from itertools import islice
 
 import numpy as np
 
 import pennylane as qml
 from pennylane import transform
-from pennylane.gradients.gradient_transform import _contract_qjac_with_cjac
-from pennylane.gradients.metric_tensor import _get_aux_wire
-from pennylane.operation import has_grad_method, is_measurement, is_trainable, not_tape
 from pennylane.tape import QuantumScript, QuantumScriptBatch
-from pennylane.transforms.tape_expand import create_expand_fn
 from pennylane.typing import PostprocessingFn
 
 from .gradient_transform import (
-    _all_zero_grad,
+    _contract_qjac_with_cjac,
     _no_trainable_grad,
+    _try_zero_grad_from_graph_or_get_grad_method,
     assert_no_state_returns,
     assert_no_trainable_tape_batching,
     assert_no_variance,
-    choose_trainable_params,
-    find_and_validate_gradient_methods,
+    choose_trainable_param_indices,
 )
+from .metric_tensor import _get_aux_wire
 
-# pylint: disable=unused-argument,invalid-unary-operand-type
 
-_expand_invalid_trainable_doc_hadamard = """Expand out a tape so that it supports differentiation
-of requested operations with the Hadamard test gradient.
+def _hadamard_stopping_condition(op) -> bool:
+    if not op.has_decomposition:
+        # let things without decompositions through without error
+        # error will happen when calculating hadamard grad
+        return True
+    if isinstance(op, qml.operation.Operator) and any(qml.math.requires_grad(p) for p in op.data):
+        return op.has_generator
+    return True
 
-This is achieved by decomposing all trainable operations that
-are not in the Hadamard compatible list until all resulting operations
-are in the list up to maximum depth ``depth``. Note that this
-might not be possible, in which case the gradient rule will fail to apply.
 
-Args:
-    tape (.QuantumTape): the input tape to expand
-    depth (int) : the maximum expansion depth
-    **kwargs: additional keyword arguments are ignored
-
-Returns:
-    .QuantumTape: the expanded tape
-"""
-
-
-hadamard_comp_list = [
-    "RX",
-    "RY",
-    "RZ",
-    "Rot",
-    "PhaseShift",
-    "U1",
-    "CRX",
-    "CRY",
-    "CRZ",
-    "IsingXX",
-    "IsingYY",
-    "IsingZZ",
-]
-
-
-@qml.BooleanFn
-def _is_hadamard_grad_compatible(obj):
-    """Check if the operation is compatible with Hadamard gradient transform."""
-    return obj.name in hadamard_comp_list
-
-
-expand_invalid_trainable_hadamard_gradient = create_expand_fn(
-    depth=None,
-    stop_at=not_tape
-    | is_measurement
-    | (~is_trainable)
-    | (_is_hadamard_grad_compatible & has_grad_method),
-    docstring=_expand_invalid_trainable_doc_hadamard,
-)
+def _inplace_set_trainable_params(tape):
+    """Update all the trainable params in place."""
+    params = tape.get_parameters(trainable_only=False)
+    tape.trainable_params = qml.math.get_trainable_indices(params)
 
 
+# pylint: disable=unused-argument
 def _expand_transform_hadamard(
     tape: QuantumScript,
     argnum=None,
@@ -97,15 +61,26 @@ def _expand_transform_hadamard(
     device_wires=None,
 ) -> tuple[QuantumScriptBatch, PostprocessingFn]:
     """Expand function to be applied before hadamard gradient."""
-    expanded_tape = expand_invalid_trainable_hadamard_gradient(tape)
-
-    def null_postprocessing(results):
-        """A postprocesing function returned by a transform that only converts the batch of results
-        into a result for a single ``QuantumTape``.
-        """
-        return results[0]
-
-    return [expanded_tape], null_postprocessing
+    batch, postprocessing = qml.devices.preprocess.decompose(
+        tape,
+        stopping_condition=_hadamard_stopping_condition,
+        skip_initial_state_prep=False,
+        name="hadamard",
+        error=qml.operation.DecompositionUndefinedError,
+    )
+    if any(
+        qml.math.requires_grad(d) for mp in tape.measurements for d in getattr(mp.obs, "data", [])
+    ):
+        try:
+            batch, postprocessing = qml.transforms.split_to_single_terms(batch[0])
+        except RuntimeError as e:
+            raise ValueError(
+                "Can only differentiate Hamiltonian "
+                f"coefficients for expectations, not {tape.measurements}."
+            ) from e
+    if len(batch) > 1 or batch[0] is not tape:  # split to single terms modified the tape
+        _ = [_inplace_set_trainable_params(t) for t in batch]
+    return batch, postprocessing
 
 
 @partial(
@@ -156,10 +131,12 @@ def hadamard_grad(
     .. math::
 
         \frac{\partial f}{\partial \mathbf{p}} = -2 \Im[\bra{0} \hat{O} G \ket{0}] = i \left(\bra{0} \hat{O} G \ket{
-        0} - \bra{0} G\hat{O} \ket{0}\right) = -2 \bra{+}\bra{0} ctrl-G^{\dagger} (\hat{Y} \otimes \hat{O}) ctrl-G
+        0} - \bra{0} G\hat{O} \ket{0}\right) = -2 \bra{+}\bra{0} \texttt{ctrl}\left(G^{\dagger}\right) (\hat{Y} \otimes \hat{O}) \texttt{ctrl}\left(G\right)
         \ket{+}\ket{0}
 
-    Here, :math:`G` is the generator of the unitary :math:`U`.
+    Here, :math:`G` is the generator of the unitary :math:`U`. ``hadamard_grad`` will work on any :math:`U` so long
+    as it has a generator :math:`G` defined (i.e., ``op.has_generator == True``). Otherwise, it will try to decompose
+    into gates where this is satisfied.
 
     **Example**
 
@@ -183,6 +160,33 @@ def hadamard_grad(
     .. details::
         :title: Usage Details
 
+        This gradient method can work with any operator that has a generator:
+
+        .. code-block:: pycon
+
+            >>> dev = qml.device('default.qubit')
+            >>> @qml.qnode(dev)
+            ... def circuit(x):
+            ...     qml.evolve(qml.X(0) @ qml.X(1) + qml.Z(0) @ qml.Z(1) + qml.H(0), x )
+            ...     return qml.expval(qml.Z(0))
+            ...
+            >>> print( qml.draw(qml.gradients.hadamard_grad(circuit))(qml.numpy.array(0.5)) )
+            0: ─╭Exp(-0.50j 𝓗)─╭X────┤ ╭<Z@Y>
+            1: ─╰Exp(-0.50j 𝓗)─│─────┤ │
+            2: ──H─────────────╰●──H─┤ ╰<Z@Y>
+
+            0: ─╭Exp(-0.50j 𝓗)─╭X@X────┤ ╭<Z@Y>
+            1: ─╰Exp(-0.50j 𝓗)─├X@X────┤ │
+            2: ──H─────────────╰●────H─┤ ╰<Z@Y>
+
+            0: ─╭Exp(-0.50j 𝓗)─╭Z────┤ ╭<Z@Y>
+            1: ─╰Exp(-0.50j 𝓗)─│─────┤ │
+            2: ──H─────────────╰●──H─┤ ╰<Z@Y>
+
+            0: ─╭Exp(-0.50j 𝓗)─╭Z@Z────┤ ╭<Z@Y>
+            1: ─╰Exp(-0.50j 𝓗)─├Z@Z────┤ │
+            2: ──H─────────────╰●────H─┤ ╰<Z@Y>
+
         This gradient transform can be applied directly to :class:`QNode <pennylane.QNode>`
         objects. However, for performance reasons, we recommend providing the gradient transform
         as the ``diff_method`` argument of the QNode decorator, and differentiating with your
@@ -254,7 +258,8 @@ def hadamard_grad(
 
         >>> dev_wires = ("a", "c")
         >>> dev = qml.device("default.qubit", wires=dev_wires)
-        >>> @qml.qnode(dev, interface="jax", diff_method="hadamard", aux_wire="c", device_wires=dev_wires)
+        >>> gradient_kwargs = {"aux_wire": "c"}
+        >>> @qml.qnode(dev, interface="jax", diff_method="hadamard", gradient_kwargs=gradient_kwargs)
         >>> def circuit(params):
         ...    qml.RX(params[0], wires="a")
         ...    qml.RY(params[1], wires="a")
@@ -263,27 +268,6 @@ def hadamard_grad(
         >>> params = jax.numpy.array([0.1, 0.2, 0.3])
         >>> jax.jacobian(circuit)(params)
         Array([-0.3875172 , -0.18884787, -0.38355704], dtype=float64)
-
-    .. note::
-
-        ``hadamard_grad`` will decompose the operations that are not in the list of supported operations.
-
-        - :class:`~.pennylane.RX`
-        - :class:`~.pennylane.RY`
-        - :class:`~.pennylane.RZ`
-        - :class:`~.pennylane.Rot`
-        - :class:`~.pennylane.PhaseShift`
-        - :class:`~.pennylane.U1`
-        - :class:`~.pennylane.CRX`
-        - :class:`~.pennylane.CRY`
-        - :class:`~.pennylane.CRZ`
-        - :class:`~.pennylane.IsingXX`
-        - :class:`~.pennylane.IsingYY`
-        - :class:`~.pennylane.IsingZZ`
-
-        The expansion will fail if a suitable decomposition in terms of supported operation is not found.
-        The number of trainable parameters may increase due to the decomposition.
-
     """
 
     transform_name = "Hadamard test"
@@ -298,229 +282,139 @@ def hadamard_grad(
     if argnum is None and not tape.trainable_params:
         return _no_trainable_grad(tape)
 
-    trainable_params = choose_trainable_params(tape, argnum)
-    diff_methods = find_and_validate_gradient_methods(tape, "analytic", trainable_params)
-
-    if all(g == "0" for g in diff_methods.values()):
-        return _all_zero_grad(tape)
-
-    argnum = [i for i, dm in diff_methods.items() if dm == "A"]
+    trainable_param_indices = choose_trainable_param_indices(tape, argnum)
 
     # Validate or get default for aux_wire
     aux_wire = _get_aux_wire(aux_wire, tape, device_wires)
 
-    g_tapes, processing_fn = _expval_hadamard_grad(tape, argnum, aux_wire)
-
-    return g_tapes, processing_fn
-
-
-def _expval_hadamard_grad(tape, argnum, aux_wire):
-    r"""Compute the Hadamard test gradient of a tape that returns an expectation value (probabilities are expectations
-    values) with respect to a given set of all trainable gate parameters.
-    The auxiliary wire is the wire which is used to apply the Hadamard gates and controlled gates.
-    """
-    # pylint: disable=too-many-statements
-    argnums = argnum or tape.trainable_params
     g_tapes = []
     coeffs = []
+    generators_per_parameter = []
 
-    gradient_data = []
     for trainable_param_idx, _ in enumerate(tape.trainable_params):
-        if trainable_param_idx not in argnums:
-            # parameter has zero gradient
-            gradient_data.append(0)
-            continue
-
-        trainable_op, idx, p_idx = tape.get_operation(trainable_param_idx)
-
-        ops_to_trainable_op = tape.operations[: idx + 1]
-        ops_after_trainable_op = tape.operations[idx + 1 :]
-
-        # Get a generator and coefficients
-        sub_coeffs, generators = _get_generators(trainable_op)
-        coeffs.extend(sub_coeffs)
-
-        num_tape = 0
-
-        for gen in generators:
-            if isinstance(trainable_op, qml.Rot):
-                # We only registered PauliZ as generator for Rot, therefore we need to apply some gates
-                # before and after the generator for the first two parameters.
-                if p_idx == 0:
-                    # Move the Rot gate past the generator
-                    op_before_trainable_op = ops_to_trainable_op.pop(-1)
-                    ops_after_trainable_op = [op_before_trainable_op] + ops_after_trainable_op
-                elif p_idx == 1:
-                    # Apply additional rotations that effectively move the generator to the middle of Rot
-                    ops_to_add_before = [
-                        qml.RZ(-trainable_op.data[2], wires=trainable_op.wires),
-                        qml.RX(np.pi / 2, wires=trainable_op.wires),
-                    ]
-                    ops_to_trainable_op.extend(ops_to_add_before)
-
-                    ops_to_add_after = [
-                        qml.RX(-np.pi / 2, wires=trainable_op.wires),
-                        qml.RZ(trainable_op.data[2], wires=trainable_op.wires),
-                    ]
-                    ops_after_trainable_op = ops_to_add_after + ops_after_trainable_op
-
-            ctrl_gen = [qml.ctrl(gen, control=aux_wire)]
-            hadamard = [qml.Hadamard(wires=aux_wire)]
-            ops = ops_to_trainable_op + hadamard + ctrl_gen + hadamard + ops_after_trainable_op
-
-            measurements = []
-            # Add the Y measurement on the aux qubit
-            for m in tape.measurements:
-                if m.obs:
-                    obs_new = [m.obs]
-                else:
-                    m_wires = m.wires if len(m.wires) > 0 else tape.wires
-                    obs_new = [qml.Z(i) for i in m_wires]
-
-                obs_new.append(qml.Y(aux_wire))
-                obs_new = qml.prod(*obs_new)
-
-                if isinstance(m, qml.measurements.ExpectationMP):
-                    measurements.append(qml.expval(op=obs_new))
-                else:
-                    measurements.append(qml.probs(op=obs_new))
-
-            new_tape = qml.tape.QuantumScript(ops=ops, measurements=measurements, shots=tape.shots)
-            _rotations, _measurements = qml.tape.tape.rotations_and_diagonal_measurements(new_tape)
-            new_ops = new_tape.operations + _rotations
-            new_tape = qml.tape.QuantumScript(
-                new_ops,
-                _measurements,
-                shots=new_tape.shots,
-                trainable_params=new_tape.trainable_params,
+        if (
+            trainable_param_idx not in trainable_param_indices
+            or _try_zero_grad_from_graph_or_get_grad_method(
+                tape, tape.trainable_params[trainable_param_idx], True
             )
+            == "0"
+        ):
+            generators_per_parameter.append(0)
+        else:
+            # can dispatch between different algorithms here in the future
+            # hadamard test, direct hadamard test, reversed, reversed direct, and flexible
+            batch, new_coeffs = _hadamard_test(tape, trainable_param_idx, aux_wire)
+            g_tapes += batch
+            coeffs += new_coeffs
+            generators_per_parameter.append(len(batch))
+
+    return g_tapes, partial(
+        processing_fn, coeffs=coeffs, tape=tape, generators_per_parameter=generators_per_parameter
+    )
+
+
+def _hadamard_test(tape, trainable_param_idx, aux_wire) -> tuple[list, list]:
+
+    trainable_op, idx, _ = tape.get_operation(trainable_param_idx)
+
+    ops_to_trainable_op = tape.operations[: idx + 1]
+    ops_after_trainable_op = tape.operations[idx + 1 :]
 
-            num_tape += 1
+    # Get a generator and coefficients
+    sub_coeffs, generators = _get_pauli_generators(trainable_op)
 
-            g_tapes.append(new_tape)
+    measurements = [_new_measurement(mp, aux_wire, tape.wires) for mp in tape.measurements]
 
-        gradient_data.append(num_tape)
+    new_batch = []
+    for gen in generators:
+        ctrl_gen = [qml.ctrl(gen, control=aux_wire)]
+        hadamard = [qml.Hadamard(wires=aux_wire)]
+        ops = ops_to_trainable_op + hadamard + ctrl_gen + hadamard + ops_after_trainable_op
 
-    multi_measurements = len(tape.measurements) > 1
-    multi_params = len(tape.trainable_params) > 1
+        new_tape = qml.tape.QuantumScript(ops, measurements, shots=tape.shots)
+        new_batch.append(new_tape)
+    return new_batch, sub_coeffs
+
+
+def _new_measurement(mp, aux_wire, all_wires: qml.wires.Wires):
+    obs = mp.obs or qml.prod(*(qml.Z(w) for w in mp.wires or all_wires))
+    new_obs = qml.simplify(obs @ qml.Y(aux_wire))
+    return type(mp)(obs=new_obs)
+
+
+def _get_pauli_generators(trainable_op):
+    """From a trainable operation, extract the unitary generators and their coefficients.
+    Any operator with a generator is supported.
+    """
+    generator = trainable_op.generator()
+    if generator.pauli_rep is None:
+        mat = qml.matrix(generator, wire_order=generator.wires)
+        pauli_rep = qml.pauli_decompose(mat, wire_order=generator.wires, pauli=True)
+    else:
+        pauli_rep = generator.pauli_rep
+    id_pw = qml.pauli.PauliWord({})
+    if id_pw in pauli_rep:
+        del pauli_rep[qml.pauli.PauliWord({})]  # remove identity term
+    sum_op = pauli_rep.operation()
+    return sum_op.terms()
+
+
+def _postprocess_probs(res, measurement, tape):
+    projector = np.array([1, -1])
+    projector = qml.math.convert_like(projector, res)
+
+    num_wires_probs = len(measurement.wires)
+    if num_wires_probs == 0:
+        num_wires_probs = tape.num_wires
+    res = qml.math.reshape(res, (2**num_wires_probs, 2))
+    return qml.math.tensordot(res, projector, axes=[[1], [0]])
+
+
+def processing_fn(results: qml.typing.ResultBatch, tape, coeffs, generators_per_parameter):
+    """Post processing function for computing a hadamard gradient."""
+
+    final_res = []
+    for coeff, res in zip(coeffs, results):
+        if not isinstance(res, (tuple, list)):
+            res = [res]  # add singleton dimension back in for one measurement
+        final_res.append([qml.math.convert_like(2 * coeff * r, r) for r in res])
+
+    # Post process for probs
     measurements_probs = [
         idx
         for idx, m in enumerate(tape.measurements)
         if isinstance(m, qml.measurements.ProbabilityMP)
     ]
-
-    def _postprocess_probs(res, measurement, projector):
-        num_wires_probs = len(measurement.wires)
-        if num_wires_probs == 0:
-            num_wires_probs = tape.num_wires
-        res = qml.math.reshape(res, (2**num_wires_probs, 2))
-        return qml.math.tensordot(res, projector, axes=[[1], [0]])
-
-    def processing_fn(results):  # pylint: disable=too-many-branches
-        """Post processing function for computing a hadamard gradient."""
-        final_res = []
-        for coeff, res in zip(coeffs, results):
-            if isinstance(res, tuple):
-                new_val = [qml.math.convert_like(2 * coeff * r, r) for r in res]
-            else:
-                new_val = qml.math.convert_like(2 * coeff * res, res)
-            final_res.append(new_val)
-
-        # Post process for probs
-        if measurements_probs:
-            projector = np.array([1, -1])
-            like = final_res[0][0] if multi_measurements else final_res[0]
-            projector = qml.math.convert_like(projector, like)
-            for idx, res in enumerate(final_res):
-                if multi_measurements:
-                    for prob_idx in measurements_probs:
-                        final_res[idx][prob_idx] = _postprocess_probs(
-                            res[prob_idx], tape.measurements[prob_idx], projector
-                        )
-                else:
-                    prob_idx = measurements_probs[0]
-                    final_res[idx] = _postprocess_probs(res, tape.measurements[prob_idx], projector)
-        grads = []
-        idx = 0
-        for num_tape in gradient_data:
-            if num_tape == 0:
-                grads.append(qml.math.zeros(()))
-            elif num_tape == 1:
-                grads.append(final_res[idx])
-                idx += 1
-            else:
-                result = final_res[idx : idx + num_tape]
-                if multi_measurements:
-                    grads.append(
-                        [qml.math.array(qml.math.sum(res, axis=0)) for res in zip(*result)]
-                    )
-                else:
-                    grads.append(qml.math.array(qml.math.sum(result)))
-                idx += num_tape
-
-        if not multi_measurements and not multi_params:
-            return grads[0]
-
-        if not (multi_params and multi_measurements):
-            if multi_measurements:
-                return tuple(grads[0])
-            return tuple(grads)
-
-        # Reordering to match the right shape for multiple measurements
-        grads_reorder = [[0] * len(tape.trainable_params) for _ in range(len(tape.measurements))]
-
-        for i in range(len(tape.measurements)):
-            for j in range(len(tape.trainable_params)):
-                grads_reorder[i][j] = grads[j][i]
-
-        grads_tuple = tuple(tuple(elem) for elem in grads_reorder)
-
-        return grads_tuple
-
-    return g_tapes, processing_fn
-
-
-def _get_generators(trainable_op):
-    """From a trainable operation, extract the unitary generators and their coefficients. If an operation is added here
-    one needs to also update the list of supported operation in the expand function given to the gradient transform.
-    """
-    # For PhaseShift, we need to separate the generator in two unitaries (Hardware compatibility)
-    if isinstance(trainable_op, (qml.PhaseShift, qml.U1)):
-        generators = [qml.Z(trainable_op.wires)]
-        coeffs = [-0.5]
-    elif isinstance(trainable_op, qml.CRX):
-        generators = [
-            qml.X(trainable_op.wires[1]),
-            qml.prod(qml.Z(trainable_op.wires[0]), qml.X(trainable_op.wires[1])),
-        ]
-        coeffs = [-0.25, 0.25]
-    elif isinstance(trainable_op, qml.CRY):
-        generators = [
-            qml.Y(trainable_op.wires[1]),
-            qml.prod(qml.Z(trainable_op.wires[0]), qml.Y(trainable_op.wires[1])),
-        ]
-        coeffs = [-0.25, 0.25]
-    elif isinstance(trainable_op, qml.CRZ):
-        generators = [
-            qml.Z(trainable_op.wires[1]),
-            qml.prod(qml.Z(trainable_op.wires[0]), qml.Z(trainable_op.wires[1])),
-        ]
-        coeffs = [-0.25, 0.25]
-    elif isinstance(trainable_op, qml.IsingXX):
-        generators = [qml.prod(qml.X(trainable_op.wires[0]), qml.X(trainable_op.wires[1]))]
-        coeffs = [-0.5]
-    elif isinstance(trainable_op, qml.IsingYY):
-        generators = [qml.prod(qml.Y(trainable_op.wires[0]), qml.Y(trainable_op.wires[1]))]
-        coeffs = [-0.5]
-    elif isinstance(trainable_op, qml.IsingZZ):
-        generators = [qml.prod(qml.Z(trainable_op.wires[0]), qml.Z(trainable_op.wires[1]))]
-        coeffs = [-0.5]
-    # For rotation it is possible to only use PauliZ by applying some other rotations in the main function
-    elif isinstance(trainable_op, qml.Rot):
-        generators = [qml.Z(trainable_op.wires)]
-        coeffs = [-0.5]
-    else:
-        generators = trainable_op.generator().ops
-        coeffs = trainable_op.generator().coeffs
-
-    return coeffs, generators
+    if measurements_probs:
+        for idx, res in enumerate(final_res):
+            for prob_idx in measurements_probs:
+                final_res[idx][prob_idx] = _postprocess_probs(
+                    res[prob_idx], tape.measurements[prob_idx], tape
+                )
+
+    mps = (
+        tape.measurements * tape.shots.num_copies
+        if tape.shots.has_partitioned_shots
+        else tape.measurements
+    )
+    # first index = into measurements, second index = into parameters
+    grads = tuple([] for _ in mps)
+    results = iter(final_res)
+    for num_generators in generators_per_parameter:
+        if num_generators == 0:
+            for g_for_parameter, mp in zip(grads, mps):
+                zeros_like_mp = np.zeros(
+                    mp.shape(num_device_wires=len(tape.wires)), dtype=mp.numeric_type
+                )
+                g_for_parameter.append(zeros_like_mp)
+        else:
+            sub_results = islice(results, num_generators)  # take the next number of results
+            # sum over batch, iterate over measurements
+            summed_sub_results = (sum(r) for r in zip(*sub_results))
+
+            for g_for_parameter, r in zip(grads, summed_sub_results, strict=True):
+                g_for_parameter.append(r)
+
+    if len(mps) == 1:
+        return grads[0][0] if len(tape.trainable_params) == 1 else grads[0]
+    return tuple(g[0] for g in grads) if len(tape.trainable_params) == 1 else grads
diff --git a/pennylane/gradients/jvp.py b/pennylane/gradients/jvp.py
index 67428ab5c..2f8bb0aa7 100644
--- a/pennylane/gradients/jvp.py
+++ b/pennylane/gradients/jvp.py
@@ -163,7 +163,7 @@ def compute_jvp_single(tangent, jac):
     """
     if jac is None:
         return None
-    single_param = not isinstance(jac, tuple)
+    single_param = not isinstance(jac, (list, tuple))
     if (single_param and jac.shape == (0,)) or (not single_param and len(jac) == 0):
         # No trainable parameters
         return qml.math.zeros((1, 0))
diff --git a/pennylane/gradients/parameter_shift.py b/pennylane/gradients/parameter_shift.py
index 1bb2b25c7..7ee4c6ac2 100644
--- a/pennylane/gradients/parameter_shift.py
+++ b/pennylane/gradients/parameter_shift.py
@@ -41,7 +41,7 @@ from .gradient_transform import (
     _swap_first_two_axes,
     assert_no_state_returns,
     assert_no_trainable_tape_batching,
-    choose_trainable_params,
+    choose_trainable_param_indices,
     find_and_validate_gradient_methods,
     reorder_grads,
 )
@@ -1127,8 +1127,8 @@ def param_shift(
 
     method = "analytic" if fallback_fn is None else "best"
 
-    trainable_params = choose_trainable_params(tape, argnum)
-    diff_methods = find_and_validate_gradient_methods(tape, method, trainable_params)
+    trainable_params_indices = choose_trainable_param_indices(tape, argnum)
+    diff_methods = find_and_validate_gradient_methods(tape, method, trainable_params_indices)
 
     if all(g == "0" for g in diff_methods.values()):
         return _all_zero_grad(tape)
diff --git a/pennylane/gradients/parameter_shift_cv.py b/pennylane/gradients/parameter_shift_cv.py
index d2332bc71..4e8714e82 100644
--- a/pennylane/gradients/parameter_shift_cv.py
+++ b/pennylane/gradients/parameter_shift_cv.py
@@ -26,7 +26,7 @@ from pennylane import transform
 from pennylane.gradients.gradient_transform import (
     _contract_qjac_with_cjac,
     _validate_gradient_methods,
-    choose_trainable_params,
+    choose_trainable_param_indices,
 )
 from pennylane.measurements import (
     ExpectationMP,
@@ -708,8 +708,8 @@ def param_shift_cv(
 
     method = "analytic" if fallback_fn is None else "best"
 
-    trainable_params = choose_trainable_params(tape, argnum)
-    method_map = _gradient_analysis_and_validation_cv(tape, method, trainable_params)
+    trainable_params_indices = choose_trainable_param_indices(tape, argnum)
+    method_map = _gradient_analysis_and_validation_cv(tape, method, trainable_params_indices)
 
     if argnum is None and not tape.trainable_params:
         return _no_trainable_grad(tape)
diff --git a/pennylane/gradients/pulse_gradient.py b/pennylane/gradients/pulse_gradient.py
index 74eb41d6c..1e65c3adb 100644
--- a/pennylane/gradients/pulse_gradient.py
+++ b/pennylane/gradients/pulse_gradient.py
@@ -33,7 +33,7 @@ from .gradient_transform import (
     assert_no_state_returns,
     assert_no_trainable_tape_batching,
     assert_no_variance,
-    choose_trainable_params,
+    choose_trainable_param_indices,
     find_and_validate_gradient_methods,
     reorder_grads,
 )
@@ -617,8 +617,8 @@ def stoch_pulse_grad(
     if use_broadcasting and tape.batch_size is not None:
         raise ValueError("Broadcasting is not supported for tapes that already are broadcasted.")
 
-    trainable_params = choose_trainable_params(tape, argnum)
-    diff_methods = find_and_validate_gradient_methods(tape, "analytic", trainable_params)
+    trainable_params_indices = choose_trainable_param_indices(tape, argnum)
+    diff_methods = find_and_validate_gradient_methods(tape, "analytic", trainable_params_indices)
 
     if all(g == "0" for g in diff_methods.values()):
         return _all_zero_grad(tape)
diff --git a/pennylane/gradients/pulse_gradient_odegen.py b/pennylane/gradients/pulse_gradient_odegen.py
index db7300741..4c85dbca1 100644
--- a/pennylane/gradients/pulse_gradient_odegen.py
+++ b/pennylane/gradients/pulse_gradient_odegen.py
@@ -32,7 +32,7 @@ from .gradient_transform import (
     assert_no_state_returns,
     assert_no_trainable_tape_batching,
     assert_no_variance,
-    choose_trainable_params,
+    choose_trainable_param_indices,
     find_and_validate_gradient_methods,
     reorder_grads,
 )
@@ -689,8 +689,8 @@ def pulse_odegen(
     if argnum is None and not tape.trainable_params:
         return _no_trainable_grad(tape)
 
-    trainable_params = choose_trainable_params(tape, argnum)
-    diff_methods = find_and_validate_gradient_methods(tape, "analytic", trainable_params)
+    trainable_params_indices = choose_trainable_param_indices(tape, argnum)
+    diff_methods = find_and_validate_gradient_methods(tape, "analytic", trainable_params_indices)
 
     if all(g == "0" for g in diff_methods.values()):
         return _all_zero_grad(tape)
diff --git a/pennylane/gradients/spsa_gradient.py b/pennylane/gradients/spsa_gradient.py
index 91135ff86..19e5b8b01 100644
--- a/pennylane/gradients/spsa_gradient.py
+++ b/pennylane/gradients/spsa_gradient.py
@@ -32,7 +32,7 @@ from .gradient_transform import (
     _all_zero_grad,
     _no_trainable_grad,
     assert_no_trainable_tape_batching,
-    choose_trainable_params,
+    choose_trainable_param_indices,
     find_and_validate_gradient_methods,
 )
 
@@ -289,11 +289,11 @@ def spsa_grad(
     if argnum is None and not tape.trainable_params:
         return _no_trainable_grad(tape)
 
-    trainable_params = choose_trainable_params(tape, argnum)
+    trainable_params_indices = choose_trainable_param_indices(tape, argnum)
     diff_methods = (
-        find_and_validate_gradient_methods(tape, "numeric", trainable_params)
+        find_and_validate_gradient_methods(tape, "numeric", trainable_params_indices)
         if validate_params
-        else {idx: "F" for idx in trainable_params}
+        else {idx: "F" for idx in trainable_params_indices}
     )
 
     if all(g == "0" for g in diff_methods.values()):
diff --git a/pennylane/gradients/vjp.py b/pennylane/gradients/vjp.py
index 614406294..5d51b2fbb 100644
--- a/pennylane/gradients/vjp.py
+++ b/pennylane/gradients/vjp.py
@@ -25,7 +25,7 @@ import pennylane as qml
 
 def _convert(jac, dy_row):
     """Utility to convert and cast the jacobian as dy_row."""
-    if isinstance(jac, tuple):
+    if isinstance(jac, (tuple, list)):
         jac_new = []
         for j in jac:
             j_ = qml.math.convert_like(j, dy_row)
@@ -120,7 +120,7 @@ def compute_vjp_single(dy, jac, num=None):
     # TODO: Excplictly catalogue and update raises for known types.
 
     # Single measurement with a single param
-    if not isinstance(jac, (tuple, autograd.builtins.SequenceBox)):
+    if not isinstance(jac, (tuple, list, autograd.builtins.SequenceBox)):
         # No trainable parameters
         if jac.shape == (0,):
             res = qml.math.zeros((1, 0))
diff --git a/pennylane/operation.py b/pennylane/operation.py
index fde5dd25f..2bab2b7d1 100644
--- a/pennylane/operation.py
+++ b/pennylane/operation.py
@@ -1449,7 +1449,7 @@ class Operator(abc.ABC, metaclass=ABCCaptureMeta):
         """
         return cls.generator != Operator.generator
 
-    def generator(self):  # pylint: disable=no-self-use
+    def generator(self) -> "Operator":  # pylint: disable=no-self-use
         r"""Generator of an operator that is in single-parameter-form.
 
         For example, for operator
@@ -1466,8 +1466,6 @@ class Operator(abc.ABC, metaclass=ABCCaptureMeta):
         The generator may also be provided in the form of a dense or sparse Hamiltonian
         (using :class:`.LinearCombination` and :class:`.SparseHamiltonian` respectively).
 
-        The default value to return is ``None``, indicating that the operation has
-        no defined generator.
         """
         raise GeneratorUndefinedError(f"Operation {self.name} does not have a generator")
 


Test patch:
diff --git a/tests/gradients/core/test_gradient_transform.py b/tests/gradients/core/test_gradient_transform.py
index 74d484b4b..16ef21359 100644
--- a/tests/gradients/core/test_gradient_transform.py
+++ b/tests/gradients/core/test_gradient_transform.py
@@ -22,7 +22,7 @@ from pennylane.gradients.gradient_transform import (
     SUPPORTED_GRADIENT_KWARGS,
     _find_gradient_methods,
     _validate_gradient_methods,
-    choose_trainable_params,
+    choose_trainable_param_indices,
 )
 from pennylane.transforms.core import TransformDispatcher
 
@@ -85,7 +85,7 @@ class TestGradAnalysis:
             qml.probs(wires=[0, 1])
 
         tape = qml.tape.QuantumScript.from_queue(q)
-        trainable_params = choose_trainable_params(tape, None)
+        trainable_params = choose_trainable_param_indices(tape, None)
         diff_methods = _find_gradient_methods(tape, trainable_params)
 
         assert diff_methods[0] is None
@@ -102,7 +102,7 @@ class TestGradAnalysis:
             qml.expval(qml.PauliY(0))
 
         tape = qml.tape.QuantumScript.from_queue(q)
-        trainable_params = choose_trainable_params(tape, None)
+        trainable_params = choose_trainable_param_indices(tape, None)
         diff_methods = _find_gradient_methods(tape, trainable_params)
 
         assert diff_methods[0] == "A"
@@ -118,7 +118,7 @@ class TestGradAnalysis:
             qml.expval(qml.PauliY(0))
 
         tape = qml.tape.QuantumScript.from_queue(q)
-        trainable_params = choose_trainable_params(tape, None)
+        trainable_params = choose_trainable_param_indices(tape, None)
         diff_methods = _find_gradient_methods(tape, trainable_params, use_graph=False)
 
         assert diff_methods[0] == "A"
@@ -139,7 +139,7 @@ class TestGradAnalysis:
             qml.probs(wires=[0, 1])
 
         tape = qml.tape.QuantumScript.from_queue(q)
-        trainable_params = choose_trainable_params(tape, None)
+        trainable_params = choose_trainable_param_indices(tape, None)
         diff_methods = _find_gradient_methods(tape, trainable_params)
 
         assert diff_methods[0] is None
@@ -182,7 +182,7 @@ class TestGradMethodValidation:
 
 
 class TestChooseParams:
-    """Test the helper function choose_trainable_params."""
+    """Test the helper function choose_trainable_param_indices."""
 
     def test_without_argnum(self):
         """Test that the method returns all params when used with ``argnum=None``."""
@@ -191,7 +191,7 @@ class TestChooseParams:
             [qml.expval(qml.PauliZ(0))],
             trainable_params=[1, 2],
         )
-        chosen = choose_trainable_params(tape, None)
+        chosen = choose_trainable_param_indices(tape, None)
         assert chosen == [0, 1]
 
     def test_with_integer_argnum(self):
@@ -202,7 +202,7 @@ class TestChooseParams:
             [qml.expval(qml.PauliZ(0))],
             trainable_params=[1, 2],
         )
-        chosen = choose_trainable_params(tape, argnum=1)
+        chosen = choose_trainable_param_indices(tape, argnum=1)
         assert chosen == [1]
 
     def test_warning_with_empty_argnum(self):
@@ -214,7 +214,7 @@ class TestChooseParams:
             trainable_params=[1, 2],
         )
         with pytest.warns(UserWarning, match="No trainable parameters were specified"):
-            chosen = choose_trainable_params(tape, [])
+            chosen = choose_trainable_param_indices(tape, [])
         assert chosen == []
 
 
diff --git a/tests/gradients/core/test_hadamard_gradient.py b/tests/gradients/core/test_hadamard_gradient.py
index 9d22609f6..2474ff4bf 100644
--- a/tests/gradients/core/test_hadamard_gradient.py
+++ b/tests/gradients/core/test_hadamard_gradient.py
@@ -105,6 +105,7 @@ def cost12(x):
     return qml.probs(op=qml.Hadamard(0) @ qml.Y(1) @ qml.Y(2))
 
 
+# pylint: disable=too-many-public-methods
 class TestHadamardGrad:
     """Unit tests for the hadamard_grad function"""
 
@@ -116,6 +117,14 @@ class TestHadamardGrad:
         with pytest.raises(NotImplementedError, match=_match):
             qml.gradients.hadamard_grad(tape)
 
+    def test_error_trainable_obs_probs(self):
+        """Test that there's an error if the observable is trainable with a probs."""
+
+        H = qml.numpy.array(2.0) * (qml.X(0) + qml.Y(0))
+        tape = qml.tape.QuantumScript([], [qml.probs(op=H)])
+        with pytest.raises(ValueError, match="Can only differentiate Hamiltonian coefficients"):
+            qml.gradients.hadamard_grad(tape)
+
     def test_nontrainable_batched_tape(self):
         """Test that no error is raised for a broadcasted/batched tape if the broadcasted
         parameter is not differentiated, and that the results correspond to the stacked
@@ -165,7 +174,6 @@ class TestHadamardGrad:
 
         assert len(tapes) == 1
 
-        assert isinstance(res_hadamard, np.ndarray)
         assert res_hadamard.shape == ()
 
         res_param_shift, _ = grad_fn(tape, dev, qml.gradients.param_shift)
@@ -190,7 +198,7 @@ class TestHadamardGrad:
         res_hadamard, tapes = grad_fn(tape, dev)
 
         assert len(tapes) == num_params
-        assert isinstance(res_hadamard, tuple)
+        # assert isinstance(res_hadamard, tuple)
         assert len(res_hadamard) == num_params
 
         res_param_shift, _ = grad_fn(tape, dev, qml.gradients.param_shift)
@@ -257,12 +265,12 @@ class TestHadamardGrad:
         res_hadamard, _ = grad_fn(tape, dev)
         res_param_shift, _ = grad_fn(tape, dev, qml.gradients.param_shift)
 
-        assert isinstance(res_hadamard, tuple)
+        assert isinstance(res_hadamard, (list, tuple))
         assert np.allclose(res_hadamard[0], res_param_shift[0], atol=tol, rtol=0)
         assert np.allclose(res_hadamard[1], res_param_shift[1], atol=tol, rtol=0)
 
     @pytest.mark.parametrize("theta", np.linspace(-2 * np.pi, 2 * np.pi, 7))
-    @pytest.mark.parametrize("G", [qml.IsingXX, qml.IsingYY, qml.IsingZZ])
+    @pytest.mark.parametrize("G", [qml.IsingXX, qml.IsingYY, qml.IsingZZ, qml.IsingXY])
     def test_ising_gradient(self, G, theta, tol):
         """Test gradient of Ising coupling gates"""
         dev = qml.device("default.qubit", wires=3)
@@ -310,7 +318,7 @@ class TestHadamardGrad:
             ]
         )
 
-        assert isinstance(grad, tuple)
+        assert isinstance(grad, (tuple, list))
         assert len(grad) == 5
         assert np.allclose(-grad[1] / 2, expected[0], atol=tol, rtol=0)
         assert np.allclose(-grad[2] / 2, expected[1], atol=tol, rtol=0)
@@ -402,6 +410,7 @@ class TestHadamardGrad:
             qml.CNOT(wires=[0, 1])
             qml.expval(qml.PauliZ(0))
             qml.expval(qml.PauliX(1))
+            qml.expval(qml.X(0))
 
         tape = qml.tape.QuantumScript.from_queue(q)
 
@@ -409,11 +418,11 @@ class TestHadamardGrad:
 
         assert len(tapes) == 2
 
-        assert len(res_hadamard) == 2
+        assert len(res_hadamard) == 3
         assert len(res_hadamard[0]) == 2
         assert len(res_hadamard[1]) == 2
 
-        expected = np.array([[-np.sin(x), 0], [0, np.cos(y)]])
+        expected = np.array([[-np.sin(x), 0], [0, np.cos(y)], [0, 0]])
         assert np.allclose(res_hadamard[0], expected[0], atol=tol, rtol=0)
         assert np.allclose(res_hadamard[1], expected[1], atol=tol, rtol=0)
 
@@ -434,7 +443,7 @@ class TestHadamardGrad:
         res_hadamard, tapes = grad_fn(tape, dev)
 
         assert len(tapes) == 2
-        assert isinstance(res_hadamard, tuple)
+        assert isinstance(res_hadamard, (tuple, list))
         assert len(res_hadamard) == 2
 
         assert res_hadamard[0].shape == (2,)
@@ -472,19 +481,19 @@ class TestHadamardGrad:
 
         assert len(tapes) == 2
 
-        assert isinstance(res_hadamard, tuple)
+        assert isinstance(res_hadamard, (tuple, list))
         assert len(res_hadamard) == 3
 
-        assert isinstance(res_hadamard[0], tuple)
+        assert isinstance(res_hadamard[0], (tuple, list))
         assert len(res_hadamard[0]) == 2
 
-        assert isinstance(res_hadamard[0][0], np.ndarray)
+        # assert isinstance(res_hadamard[0][0], np.ndarray)
         assert res_hadamard[0][0].shape == ()
-        assert isinstance(res_hadamard[0][1], np.ndarray)
+        # assert isinstance(res_hadamard[0][1], np.ndarray)
         assert res_hadamard[0][1].shape == ()
 
         for res in res_hadamard[1:]:
-            assert isinstance(res, tuple)
+            assert isinstance(res, (tuple, list))
             assert len(res) == 2
 
             assert isinstance(res[0], np.ndarray)
@@ -615,7 +624,7 @@ class TestHadamardGrad:
                     assert isinstance(r, np.ndarray)
                     assert len(r) == exp_shape[1]
 
-    @pytest.mark.parametrize("shots", [None, 100])
+    @pytest.mark.parametrize("shots", [None, 100, (1000, 1000)])
     def test_shots_attribute(self, shots):
         """Tests that the shots attribute is copied to the new tapes"""
         dev = qml.device("default.qubit", wires=3)
@@ -786,10 +795,9 @@ class TestHadamardGradEdgeCases:
         ):
             qml.gradients.hadamard_grad(tape)
 
-    def test_independent_parameter(self, mocker):
+    def test_independent_parameter(self):
         """Test that an independent parameter is skipped
         during the Jacobian computation."""
-        spy = mocker.spy(qml.gradients.hadamard_gradient, "_expval_hadamard_grad")
 
         with qml.queuing.AnnotatedQueue() as q:
             qml.RX(0.543, wires=[0])
@@ -803,14 +811,11 @@ class TestHadamardGradEdgeCases:
 
         assert len(tapes) == 1
 
-        assert isinstance(res_hadamard, tuple)
+        assert isinstance(res_hadamard, (tuple, list))
         assert len(res_hadamard) == 2
         assert res_hadamard[0].shape == ()
         assert res_hadamard[1].shape == ()
 
-        # only called for parameter 0
-        assert spy.call_args[0][0:2] == (tape, [0])
-
     @pytest.mark.autograd
     def test_no_trainable_params_qnode_autograd(self):
         """Test that the correct ouput and warning is generated in the absence of any trainable
@@ -939,7 +944,7 @@ class TestHadamardGradEdgeCases:
 
         assert tapes == []
 
-        assert isinstance(res_hadamard, tuple)
+        assert isinstance(res_hadamard, (list, tuple))
 
         assert len(res_hadamard) == 3
 
@@ -1024,6 +1029,7 @@ class TestHadamardGradEdgeCases:
         result = qml.gradients.hadamard_grad(circuit)(params)
 
         assert isinstance(result, np.ndarray)
+        print(result)
         assert result.shape == (4, 3)
         assert np.allclose(result, 0)
 
@@ -1127,7 +1133,7 @@ class TestHadamardTestGradDiff:
             tape = qml.tape.QuantumScript.from_queue(q)
             tapes, fn = qml.gradients.hadamard_grad(tape)
 
-            return fn(dev.execute(tapes))
+            return tuple(fn(dev.execute(tapes)))
 
         def cost_p(x):
             with qml.queuing.AnnotatedQueue() as q:
@@ -1139,7 +1145,7 @@ class TestHadamardTestGradDiff:
             tape = qml.tape.QuantumScript.from_queue(q)
             tapes, fn = qml.gradients.param_shift(tape)
 
-            return fn(dev.execute(tapes))
+            return tuple(fn(dev.execute(tapes)))
 
         res_hadamard = torch.autograd.functional.jacobian(cost_h, params)
         res_param_shift = torch.autograd.functional.jacobian(cost_p, params)
diff --git a/tests/gradients/parameter_shift/test_parameter_shift_cv.py b/tests/gradients/parameter_shift/test_parameter_shift_cv.py
index d50defbc5..7214e9289 100644
--- a/tests/gradients/parameter_shift/test_parameter_shift_cv.py
+++ b/tests/gradients/parameter_shift/test_parameter_shift_cv.py
@@ -21,7 +21,7 @@ import pytest
 import pennylane as qml
 from pennylane import numpy as np
 from pennylane.gradients import param_shift_cv
-from pennylane.gradients.gradient_transform import choose_trainable_params
+from pennylane.gradients.gradient_transform import choose_trainable_param_indices
 from pennylane.gradients.parameter_shift_cv import (
     _find_gradient_methods_cv,
     _grad_method_cv,
@@ -51,7 +51,7 @@ class TestGradAnalysis:
         assert _grad_method_cv(tape, 3) == "A"
         assert _grad_method_cv(tape, 4) == "A"
 
-        trainable_params = choose_trainable_params(tape, None)
+        trainable_params = choose_trainable_param_indices(tape, None)
         diff_methods = _find_gradient_methods_cv(tape, trainable_params)
 
         assert diff_methods[0] is None
@@ -73,7 +73,7 @@ class TestGradAnalysis:
         assert _grad_method_cv(tape, 0) == "A"
         assert _grad_method_cv(tape, 1) == "0"
 
-        trainable_params = choose_trainable_params(tape, None)
+        trainable_params = choose_trainable_param_indices(tape, None)
         diff_methods = _find_gradient_methods_cv(tape, trainable_params)
 
         assert diff_methods[0] == "A"
diff --git a/tests/workflow/interfaces/qnode/test_tensorflow_qnode.py b/tests/workflow/interfaces/qnode/test_tensorflow_qnode.py
index 79f19bd2f..25a0c5aaf 100644
--- a/tests/workflow/interfaces/qnode/test_tensorflow_qnode.py
+++ b/tests/workflow/interfaces/qnode/test_tensorflow_qnode.py
@@ -1054,6 +1054,7 @@ class TestTapeExpansion:
 
         class PhaseShift(qml.PhaseShift):
             grad_method = None
+            has_generator = False
 
             def decomposition(self):
                 return [qml.RY(3 * self.data[0], wires=self.wires)]
diff --git a/tests/workflow/interfaces/qnode/test_torch_qnode.py b/tests/workflow/interfaces/qnode/test_torch_qnode.py
index a2922a7ce..379089fac 100644
--- a/tests/workflow/interfaces/qnode/test_torch_qnode.py
+++ b/tests/workflow/interfaces/qnode/test_torch_qnode.py
@@ -1205,6 +1205,7 @@ class TestTapeExpansion:
 
         class PhaseShift(qml.PhaseShift):  # pylint:disable=too-few-public-methods
             grad_method = None
+            has_generator = False
 
             def decomposition(self):
                 return [qml.RY(3 * self.data[0], wires=self.wires)]