Repository: PennyLaneAI/pennylane
Pull request #5035 (2024-01-10 16:54:28)

Refactor gradient method analysis and validation
**Context:**

[`_gradient_analysis`](https://github.com/PennyLaneAI/pennylane/blob/f6a7375ee6e6b2efbe4a9d0a76c3570a04731cf2/pennylane/gradients/gradient_transform.py#L117) mutates the input `tape` object so that it can pass information to [`_grad_method_validation`](https://github.com/PennyLaneAI/pennylane/blob/f6a7375ee6e6b2efbe4a9d0a76c3570a04731cf2/pennylane/gradients/gradient_transform.py#L152). Instead it should directly pass the information, not relying on side effects and dynamically mutating a private property of a core data structure.

**Description of the Change:**

1. Instead of returning a list, `gradient_analysis_and_validation` (renamed to `find_and_validate_gradient_methods`) now returns a dictionary of gradient methods for each parameter index, and no longer mutates the tape.
2. Instead of validating all trainable parameters before calling `choose_grad_methods` to select a subset of parameters based on `argnums`, gradient functions now calls a new function `choose_trainable_params` with `argnums`, then calls `find_and_validate_gradient_methods` on the subset of parameters.

**Benefits:**

Crucial to move toward a more functional architecture with immutable data types.

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5035?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
Attention: `2 lines` in your changes are missing coverage. Please review.
> Comparison is base [(`9a6419c`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/9a6419c6c017e0c104a4c110b1ad956806044e6d?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) 99.67% compared to head [(`9e14681`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5035?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) 99.66%.

| [Files](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5035?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | Patch % | Lines |
|---|---|---|
| [pennylane/gradients/gradient\_transform.py](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5035?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2dyYWRpZW50cy9ncmFkaWVudF90cmFuc2Zvcm0ucHk=) | 92.85% | [2 Missing :warning: ](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5035?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    #5035      +/-   ##
==========================================
- Coverage   99.67%   99.66%   -0.02%     
==========================================
  Files         394      394              
  Lines       35582    35340     -242     
==========================================
- Hits        35467    35222     -245     
- Misses        115      118       +3     
```



</details>

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


Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 571dc53d3..84b7b11eb 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -44,6 +44,9 @@
 
 <h3>Breaking changes 💔</h3>
 
+* `gradient_analysis_and_validation` is now renamed to `find_and_validate_gradient_methods`. Instead of returning a list, it now returns a dictionary of gradient methods for each parameter index, and no longer mutates the tape.
+  [(#5035)](https://github.com/PennyLaneAI/pennylane/pull/5035)
+
 * Passing additional arguments to a transform that decorates a QNode must be done through the use
   of `functools.partial`.
   [(#5046)](https://github.com/PennyLaneAI/pennylane/pull/5046)
@@ -71,19 +74,27 @@
 * A typo in a code example in the `qml.transforms` API has been fixed.
   [(#5014)](https://github.com/PennyLaneAI/pennylane/pull/5014)
 
+* Clarification for the definition of `argnum` added to gradient methods
+  [(#5035)](https://github.com/PennyLaneAI/pennylane/pull/5035)
+
 * A typo in the code example for `qml.qchem.dipole_of` has been fixed.
   [(#5036)](https://github.com/PennyLaneAI/pennylane/pull/5036) 
 
 <h3>Bug fixes 🐛</h3>
 
+* If `argnum` is provided to a gradient transform, only the parameters specified in `argnum` will have their gradient methods validated.
+  [(#5035)](https://github.com/PennyLaneAI/pennylane/pull/5035)
+
 * `StatePrep` operations expanded onto more wires are now compatible with backprop.
   [(#5028)](https://github.com/PennyLaneAI/pennylane/pull/5028)
 
+
 <h3>Contributors ✍️</h3>
 
 This release contains contributions from (in alphabetical order):
 
 Abhishek Abhishek,
+Astral Cai,
 Pablo Antonio Moreno Casares,
 Christina Lee,
 Isaac De Vlugt,
diff --git a/pennylane/gradients/finite_difference.py b/pennylane/gradients/finite_difference.py
index e3084f45e..3d6fb2561 100644
--- a/pennylane/gradients/finite_difference.py
+++ b/pennylane/gradients/finite_difference.py
@@ -31,13 +31,12 @@ from pennylane import transform
 from pennylane.transforms.tape_expand import expand_invalid_trainable
 from pennylane.gradients.gradient_transform import _contract_qjac_with_cjac
 
-
 from .general_shift_rules import generate_shifted_tapes
 from .gradient_transform import (
     _all_zero_grad,
     assert_no_tape_batching,
-    choose_grad_methods,
-    gradient_analysis_and_validation,
+    choose_trainable_params,
+    find_and_validate_gradient_methods,
     _no_trainable_grad,
 )
 
@@ -221,7 +220,8 @@ def finite_diff(
         tape (QNode or QuantumTape): quantum circuit to differentiate
         argnum (int or list[int] or None): Trainable parameter indices to differentiate
             with respect to. If not provided, the derivatives with respect to all
-            trainable parameters are returned.
+            trainable parameters are returned. Note that the indices are with respect to
+            the list of trainable parameters.
         h (float): finite difference method step size
         approx_order (int): The approximation order of the finite-difference method to use.
         n (int): compute the :math:`n`-th derivative
@@ -286,9 +286,10 @@ def finite_diff(
     .. details::
         :title: Usage Details
 
-        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 preferred machine learning framework.
+        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
+        preferred machine learning framework.
 
         >>> @qml.qnode(dev)
         ... def circuit(params):
@@ -323,8 +324,19 @@ def finite_diff(
         This can be useful if the underlying circuits representing the gradient
         computation need to be analyzed.
 
-        The output tapes can then be evaluated and post-processed to retrieve
-        the gradient:
+        Note that ``argnum`` refers to the index of a parameter within the list of trainable
+        parameters. For example, if we have:
+
+        >>> tape = qml.tape.QuantumScript(
+        ...     [qml.RX(1.2, wires=0), qml.RY(2.3, wires=0), qml.RZ(3.4, wires=0)],
+        ...     [qml.expval(qml.PauliZ(0))],
+        ...     trainable_params = [1, 2]
+        ... )
+        >>> qml.gradients.finite_diff(tape, argnum=1)
+
+        The code above will differentiate the third parameter rather than the second.
+
+        The output tapes can then be evaluated and post-processed to retrieve the gradient:
 
         >>> dev = qml.device("default.qubit", wires=2)
         >>> fn(qml.execute(gradient_tapes, dev, None))
@@ -369,14 +381,14 @@ def finite_diff(
     if argnum is None and not tape.trainable_params:
         return _no_trainable_grad(tape)
 
-    if validate_params:
-        diff_methods = gradient_analysis_and_validation(
-            tape, "numeric", grad_fn=finite_diff, overwrite=False
-        )
-    else:
-        diff_methods = ["F" for i in tape.trainable_params]
+    trainable_params = choose_trainable_params(tape, argnum)
+    diff_methods = (
+        find_and_validate_gradient_methods(tape, "numeric", trainable_params)
+        if validate_params
+        else {idx: "F" for idx in trainable_params}
+    )
 
-    if all(g == "0" for g in diff_methods):
+    if all(g == "0" for g in diff_methods.values()):
         return _all_zero_grad(tape)
 
     gradient_tapes = []
@@ -399,10 +411,8 @@ def finite_diff(
         shifts = shifts[1:]
         coeffs = coeffs[1:]
 
-    method_map = choose_grad_methods(diff_methods, argnum)
-
     for i, _ in enumerate(tape.trainable_params):
-        if i not in method_map or method_map[i] == "0":
+        if i not in diff_methods or diff_methods[i] == "0":
             # parameter has zero gradient
             shapes.append(0)
             continue
diff --git a/pennylane/gradients/gradient_transform.py b/pennylane/gradients/gradient_transform.py
index 3bebfd364..f7d604f33 100644
--- a/pennylane/gradients/gradient_transform.py
+++ b/pennylane/gradients/gradient_transform.py
@@ -114,70 +114,94 @@ def assert_no_tape_batching(tape, transform_name):
         )
 
 
-def _gradient_analysis(tape, use_graph=True, grad_fn=None):
-    """Update the parameter information dictionary of the tape with
-    gradient information of each parameter.
+def choose_trainable_params(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
+    returns a list of indices with respect to the list of trainable parameters. If argnum
+    is not provided, all trainable parameters are considered.
+
+    Args:
+        tape (`~.QuantumScript`): the tape to analyze
+        argnum (int, list(int), None): Indices for trainable parameters(s)
+            to compute the Jacobian for.
+
+    Returns:
+        list: list of the trainable parameter indices
+
+    """
+
+    if argnum is None:
+        return [idx for idx, _ in enumerate(tape.trainable_params)]
+
+    if isinstance(argnum, int):
+        argnum = [argnum]
+
+    if len(argnum) == 0:
+        warnings.warn(
+            "No trainable parameters were specified for computing the Jacobian.",
+            UserWarning,
+        )
+
+    return argnum
+
+
+def _try_zero_grad_from_graph_or_get_grad_method(tape, param_index, use_graph=True):
+    """Gets the gradient method of a parameter. If use_graph=True, analyze the
+    circuit graph to find if the parameter has zero gradient.
+
+    Args:
+        tape (`~.QuantumScript`): the tape to analyze
+        param_index (int): the index of the parameter to analyze
+        use_graph (bool): whether to use the circuit graph to find if
+            a parameter has zero gradient
+
     """
+
     # pylint:disable=protected-access
-    if grad_fn is not None:
-        if getattr(tape, "_gradient_fn", None) is grad_fn:
-            # gradient analysis has already been performed on this tape
-            return
+    par_info = tape._par_info[param_index]
 
-        tape._gradient_fn = grad_fn
+    if use_graph:
+        op_or_mp = tape[par_info["op_idx"]]
+        if not any(tape.graph.has_path(op_or_mp, mp) for mp in tape.measurements):
+            # there is no influence of this operation on any of the observables
+            return "0"
 
-    for idx, info in enumerate(tape._par_info):
-        if idx not in tape.trainable_params:
-            # non-trainable parameters do not require a grad_method
-            info["grad_method"] = None
-        else:
-            par_info = tape._par_info[idx]
-            op = par_info["op"]
+    return par_info["op"].grad_method
 
-            if not qml.operation.has_grad_method(op):
-                # no differentiation method is registered for this operation
-                info["grad_method"] = None
 
-            elif (tape._graph is not None) or use_graph:
-                # if `op` is an observable, we want the associated MP that's in the graph
-                op_or_mp = tape[par_info["op_idx"]]
-                if not any(tape.graph.has_path(op_or_mp, mp) for mp in tape.measurements):
-                    # there is no influence of this operation on any of the observables
-                    info["grad_method"] = "0"
-                    continue
+def _find_gradient_methods(tape, trainable_param_indices, use_graph=True):
+    """Returns a dictionary with gradient information of each trainable parameter."""
 
-            info["grad_method"] = op.grad_method
+    return {
+        idx: _try_zero_grad_from_graph_or_get_grad_method(
+            tape, tape.trainable_params[idx], use_graph
+        )
+        for idx in trainable_param_indices
+    }
 
 
-def _grad_method_validation(method, tape):
+def _validate_gradient_methods(tape, method, diff_methods):
     """Validates if the gradient method requested is supported by the trainable
     parameters of a tape, and returns the allowed parameter gradient methods."""
-    diff_methods = {
-        idx: info["grad_method"]
-        for idx, info in enumerate(tape._par_info)  # pylint: disable=protected-access
-        if idx in tape.trainable_params
-    }
 
     # check and raise an error if any parameters are non-differentiable
-    if nondiff_params := {idx for idx, g in diff_methods.items() if g is None}:
+    nondiff_params = [tape.trainable_params[idx] for idx, m in diff_methods.items() if m is None]
+    if nondiff_params:
         raise ValueError(f"Cannot differentiate with respect to parameter(s) {nondiff_params}")
 
-    numeric_params = {idx for idx, g in diff_methods.items() if g == "F"}
-
-    # If explicitly using analytic mode, ensure that all parameters
-    # support analytic differentiation.
+    # If explicitly using analytic mode, ensure that all
+    # parameters support analytic differentiation.
+    numeric_params = [tape.trainable_params[idx] for idx, m in diff_methods.items() if m == "F"]
     if method == "analytic" and numeric_params:
         raise ValueError(
             f"The analytic gradient method cannot be used with the parameter(s) {numeric_params}."
         )
 
-    return tuple(diff_methods.values())
-
 
-def gradient_analysis_and_validation(tape, method, use_graph=True, grad_fn=None, overwrite=True):
-    """Update the parameter information dictionary of the tape with gradient information of
-    each parameter. Then validates if the gradient method requested is supported by the trainable
-    parameters of a tape, and returns the allowed parameter gradient methods.
+def find_and_validate_gradient_methods(tape, method, trainable_param_indices, use_graph=True):
+    """Returns a dictionary of gradient methods for each trainable parameter after
+    validating if the gradient method requested is supported by the trainable parameters
 
     Parameter gradient methods include:
 
@@ -189,69 +213,26 @@ def gradient_analysis_and_validation(tape, method, use_graph=True, grad_fn=None,
     In addition, the operator might define its own grad method
     via :attr:`.Operator.grad_method`.
 
-    .. note::
-
-        Note that this function modifies the input tape in-place.
-
     Args:
-        tape (.QuantumTape): the quantum tape to analyze
-        method (str): the overall Jacobian differentiation method
-        use_graph (bool): whether to use a directed-acyclic graph to determine
-            if the parameter has a gradient of 0
-        grad_fn (None or callable): The gradient transform performing the analysis.
-            This is an optional argument; if provided, and the tape has already
-            been analyzed for the gradient information by the same gradient transform,
-            the cached gradient analysis will be used.
-        overwrite (bool): Whether to overwrite existing parameter gradient methods during the
-            first part of the function.
+        tape (`~.QuantumScript`): the tape to analyze
+        method (str): the gradient method to use
+        trainable_param_indices (list[int]): the indices of the trainable parameters
+            to compute the Jacobian for
+        use_graph (bool): whether to use the circuit graph to find if
+            a parameter has zero gradient
+
+    Returns:
+        dict: dictionary of the gradient methods for each trainable parameter
 
     Raises:
         ValueError: If there exist non-differentiable trainable parameters on the tape.
         ValueError: If the Jacobian method is ``"analytic"`` but there exist some trainable
             parameters on the tape that only support numeric differentiation.
 
-    Returns:
-        tuple[str, None]: the allowed parameter gradient methods for each trainable parameter
     """
-    if overwrite or "grad_method" not in tape._par_info[0]:  # pylint: disable=protected-access
-        _gradient_analysis(tape, use_graph=use_graph, grad_fn=grad_fn)
-    return _grad_method_validation(method, tape)
-
-
-def choose_grad_methods(diff_methods, argnum):
-    """Chooses the trainable parameters to use for computing the Jacobian
-    by returning a map of their indices and differentiation methods.
-
-    When there are fewer parameters specified than the total number of
-    trainable parameters, the Jacobian is estimated by using the parameters
-    specified using the ``argnum`` keyword argument.
-
-    Args:
-        diff_methods (list): the ordered list of differentiation methods
-            for each parameter
-        argnum (int, list(int), None): Indices for argument(s) with respect
-            to which to compute the Jacobian.
-
-    Returns:
-        dict: map of the trainable parameter indices and
-        differentiation methods
-    """
-    if argnum is None:
-        return dict(enumerate(diff_methods))
-
-    if isinstance(argnum, int):
-        argnum = [argnum]
-
-    num_params = len(argnum)
-
-    if num_params == 0:
-        warnings.warn(
-            "No trainable parameters were specified for computing the Jacobian.",
-            UserWarning,
-        )
-        return {}
-
-    return {idx: diff_methods[idx] for idx in argnum}
+    diff_methods = _find_gradient_methods(tape, trainable_param_indices, use_graph=use_graph)
+    _validate_gradient_methods(tape, method, diff_methods)
+    return diff_methods
 
 
 def _all_zero_grad(tape):
diff --git a/pennylane/gradients/hadamard_gradient.py b/pennylane/gradients/hadamard_gradient.py
index 9648d3e44..3596b283c 100644
--- a/pennylane/gradients/hadamard_gradient.py
+++ b/pennylane/gradients/hadamard_gradient.py
@@ -30,8 +30,8 @@ from .gradient_transform import (
     assert_no_state_returns,
     assert_no_tape_batching,
     assert_no_variance,
-    choose_grad_methods,
-    gradient_analysis_and_validation,
+    choose_trainable_params,
+    find_and_validate_gradient_methods,
     _no_trainable_grad,
 )
 
@@ -66,29 +66,30 @@ def hadamard_grad(
     aux_wire=None,
     device_wires=None,
 ) -> (Sequence[qml.tape.QuantumTape], Callable):
-    r"""Transform a circuit to compute the Hadamard test gradient of all gates with respect to their inputs.
+    r"""Transform a circuit to compute the Hadamard test gradient of all gates
+    with respect to their inputs.
 
     Args:
         tape (QNode or QuantumTape): quantum circuit to differentiate
         argnum (int or list[int] or None): Trainable tape parameter indices to differentiate
             with respect to. If not provided, the derivatives with respect to all
-            trainable parameters are returned.
-        aux_wire (pennylane.wires.Wires): Auxiliary wire to be used for the Hadamard tests. If ``None`` (the default),
-            a suitable wire is inferred from the wires used in the original circuit and ``device_wires``.
+            trainable parameters are returned. Note that the indices are with respect to
+            the list of trainable parameters.
+        aux_wire (pennylane.wires.Wires): Auxiliary wire to be used for the Hadamard tests.
+            If ``None`` (the default), a suitable wire is inferred from the wires used in
+            the original circuit and ``device_wires``.
         device_wires (pennylane.wires.Wires): Wires of the device that are going to be used for the
-            gradient. Facilitates finding a default for ``aux_wire`` if ``aux_wire``
-            is ``None``.
+            gradient. Facilitates finding a default for ``aux_wire`` if ``aux_wire`` is ``None``.
 
     Returns:
         qnode (QNode) or tuple[List[QuantumTape], function]:
 
-        The transformed circuit as described in :func:`qml.transform <pennylane.transform>`. Executing this circuit
-        will provide the Jacobian in the form of a tensor, a tuple, or a nested tuple depending upon the nesting
-        structure of measurements in the original circuit.
+        The transformed circuit as described in :func:`qml.transform <pennylane.transform>`.
+        Executing this circuit will provide the Jacobian in the form of a tensor, a tuple, or a
+        nested tuple depending upon the nesting structure of measurements in the original circuit.
 
-    For a variational evolution :math:`U(\mathbf{p}) \vert 0\rangle` with
-    :math:`N` parameters :math:`\mathbf{p}`,
-    consider the expectation value of an observable :math:`O`:
+    For a variational evolution :math:`U(\mathbf{p}) \vert 0\rangle` with :math:`N` parameters
+    :math:`\mathbf{p}`, consider the expectation value of an observable :math:`O`:
 
     .. math::
 
@@ -127,9 +128,10 @@ def hadamard_grad(
     .. details::
         :title: Usage Details
 
-        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 preferred machine learning framework.
+        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
+        preferred machine learning framework.
 
         >>> dev = qml.device("default.qubit", wires=2)
         >>> @qml.qnode(dev)
@@ -161,8 +163,19 @@ def hadamard_grad(
         This can be useful if the underlying circuits representing the gradient
         computation need to be analyzed.
 
-        The output tapes can then be evaluated and post-processed to retrieve
-        the gradient:
+        Note that ``argnum`` refers to the index of a parameter within the list of trainable
+        parameters. For example, if we have:
+
+        >>> tape = qml.tape.QuantumScript(
+        ...     [qml.RX(1.2, wires=0), qml.RY(2.3, wires=0), qml.RZ(3.4, wires=0)],
+        ...     [qml.expval(qml.PauliZ(0))],
+        ...     trainable_params = [1, 2]
+        ... )
+        >>> qml.gradients.hadamard_grad(tape, argnum=1)
+
+        The code above will differentiate the third parameter rather than the second.
+
+        The output tapes can then be evaluated and post-processed to retrieve the gradient:
 
         >>> dev = qml.device("default.qubit", wires=2)
         >>> fn(qml.execute(gradient_tapes, dev, None))
@@ -230,14 +243,13 @@ def hadamard_grad(
     if argnum is None and not tape.trainable_params:
         return _no_trainable_grad(tape)
 
-    diff_methods = gradient_analysis_and_validation(tape, "analytic", grad_fn=hadamard_grad)
+    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):
+    if all(g == "0" for g in diff_methods.values()):
         return _all_zero_grad(tape)
 
-    method_map = choose_grad_methods(diff_methods, argnum)
-
-    argnum = [i for i, dm in method_map.items() if dm == "A"]
+    argnum = [i for i, dm in diff_methods.items() if dm == "A"]
 
     # Validate or get default for aux_wire
     aux_wire = _get_aux_wire(aux_wire, tape, device_wires)
@@ -263,13 +275,13 @@ def _expval_hadamard_grad(tape, argnum, aux_wire):
         for idx, m in enumerate(tape.measurements)
         if isinstance(m, qml.measurements.ProbabilityMP)
     ]
-    for id_argnum, _ in enumerate(tape.trainable_params):
-        if id_argnum not in argnums:
+    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(id_argnum)
+        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 :]
diff --git a/pennylane/gradients/parameter_shift.py b/pennylane/gradients/parameter_shift.py
index d0a333f9c..215a01e04 100644
--- a/pennylane/gradients/parameter_shift.py
+++ b/pennylane/gradients/parameter_shift.py
@@ -39,8 +39,8 @@ from .gradient_transform import (
     assert_no_state_returns,
     assert_no_tape_batching,
     assert_multimeasure_not_broadcasted,
-    choose_grad_methods,
-    gradient_analysis_and_validation,
+    choose_trainable_params,
+    find_and_validate_gradient_methods,
     _no_trainable_grad,
     reorder_grads,
 )
@@ -340,7 +340,8 @@ def expval_param_shift(
             tape (.QuantumTape): quantum tape to differentiate
             argnum (int or list[int] or None): Trainable parameter indices to differentiate
                 with respect to. If not provided, the derivatives with respect to all
-                trainable indices are returned.
+                trainable indices are returned. Note that the indices are with respect to
+            the list of trainable parameters.
             shifts (list[tuple[int or float]]): List containing tuples of shift values.
                 If provided, one tuple of shifts should be given per trainable parameter
                 and the tuple should match the number of frequencies for that parameter.
@@ -968,6 +969,18 @@ def param_shift(
         This can be useful if the underlying circuits representing the gradient
         computation need to be analyzed.
 
+        Note that ``argnum`` refers to the index of a parameter within the list of trainable
+        parameters. For example, if we have:
+
+        >>> tape = qml.tape.QuantumScript(
+        ...     [qml.RX(1.2, wires=0), qml.RY(2.3, wires=0), qml.RZ(3.4, wires=0)],
+        ...     [qml.expval(qml.PauliZ(0))],
+        ...     trainable_params = [1, 2]
+        ... )
+        >>> qml.gradients.param_shift(tape, argnum=1)
+
+        The code above will differentiate the third parameter rather than the second.
+
         The output tapes can then be evaluated and post-processed to retrieve
         the gradient:
 
@@ -1050,16 +1063,16 @@ def param_shift(
         return _no_trainable_grad(tape)
 
     method = "analytic" if fallback_fn is None else "best"
-    diff_methods = gradient_analysis_and_validation(tape, method, grad_fn=param_shift)
 
-    if all(g == "0" for g in diff_methods):
-        return _all_zero_grad(tape)
+    trainable_params = choose_trainable_params(tape, argnum)
+    diff_methods = find_and_validate_gradient_methods(tape, method, trainable_params)
 
-    method_map = choose_grad_methods(diff_methods, argnum)
+    if all(g == "0" for g in diff_methods.values()):
+        return _all_zero_grad(tape)
 
     # If there are unsupported operations, call the fallback gradient function
-    unsupported_params = {idx for idx, g in method_map.items() if g == "F"}
-    argnum = [i for i, dm in method_map.items() if dm == "A"]
+    unsupported_params = {idx for idx, g in diff_methods.items() if g == "F"}
+    argnum = [i for i, dm in diff_methods.items() if dm == "A"]
     gradient_tapes = []
 
     if unsupported_params:
@@ -1070,11 +1083,7 @@ def param_shift(
         gradient_tapes.extend(g_tapes)
         fallback_len = len(g_tapes)
 
-        # remove finite difference parameters from the method map
-        method_map = {t_idx: dm for t_idx, dm in method_map.items() if dm != "F"}
-
     # Generate parameter-shift gradient tapes
-
     if gradient_recipes is None:
         gradient_recipes = [None] * len(argnum)
 
diff --git a/pennylane/gradients/parameter_shift_cv.py b/pennylane/gradients/parameter_shift_cv.py
index bd5349da6..b44da47f0 100644
--- a/pennylane/gradients/parameter_shift_cv.py
+++ b/pennylane/gradients/parameter_shift_cv.py
@@ -33,19 +33,18 @@ from pennylane.measurements import (
 )
 from pennylane import transform
 from pennylane.transforms.tape_expand import expand_invalid_trainable
-from pennylane.gradients.gradient_transform import _contract_qjac_with_cjac
-
+from pennylane.gradients.gradient_transform import (
+    _contract_qjac_with_cjac,
+    choose_trainable_params,
+    _validate_gradient_methods,
+)
 from .finite_difference import finite_diff
 from .general_shift_rules import generate_shifted_tapes, process_shifts
-from .gradient_transform import (
-    choose_grad_methods,
-    _grad_method_validation,
-    _no_trainable_grad,
-)
+from .gradient_transform import _no_trainable_grad
 from .parameter_shift import _get_operation_recipe, expval_param_shift
 
 
-def _grad_method(tape, idx):
+def _grad_method_cv(tape, idx):
     """Determine the best CV parameter-shift gradient recipe for a given
     parameter index of a tape.
 
@@ -131,25 +130,19 @@ def _grad_method(tape, idx):
     return "A"
 
 
-def _gradient_analysis_cv(tape):
-    """Update the parameter information dictionary of the tape with
-    gradient information of each parameter."""
-    if getattr(tape, "_gradient_fn", None) is param_shift_cv:
-        # gradient analysis for param_shift_cv has already been performed on this tape
-        return
+def _find_gradient_methods_cv(tape, trainable_param_indices):
+    """Find the best gradient methods for each parameter."""
+    return {
+        idx: _grad_method_cv(tape, tape.trainable_params[idx]) for idx in trainable_param_indices
+    }
 
-    tape._gradient_fn = param_shift_cv
 
-    for idx, info in enumerate(tape._par_info):
-        info["grad_method"] = _grad_method(tape, idx)
-
-
-def _gradient_analysis_and_validation_cv(tape, method):
-    """Update the parameter information dictionary of the tape with
-    gradient information of each parameter. Subsequently validate the
-    gradient methods and return diff_methods."""
-    _gradient_analysis_cv(tape)
-    return _grad_method_validation(method, tape)
+def _gradient_analysis_and_validation_cv(tape, method, trainable_param_indices):
+    """Find the best gradient methods for each parameter. Subsequently, validate
+    the gradient methods and return diff_methods."""
+    diff_methods = _find_gradient_methods_cv(tape, trainable_param_indices)
+    _validate_gradient_methods(tape, method, diff_methods)
+    return diff_methods
 
 
 def _transform_observable(obs, Z, device_wires):
@@ -701,7 +694,9 @@ def param_shift_cv(
         )
 
     method = "analytic" if fallback_fn is None else "best"
-    diff_methods = _gradient_analysis_and_validation_cv(tape, method)
+
+    trainable_params = choose_trainable_params(tape, argnum)
+    method_map = _gradient_analysis_and_validation_cv(tape, method, trainable_params)
 
     if argnum is None and not tape.trainable_params:
         return _no_trainable_grad(tape)
@@ -717,10 +712,9 @@ def param_shift_cv(
         shapes.append(len(data[0]))
         fns.append(data[1])
 
-    if all(g == "0" for g in diff_methods):
+    if all(g == "0" for g in method_map.values()):
         return [], lambda _: np.zeros([tape.output_dim, len(tape.trainable_params)])
 
-    method_map = choose_grad_methods(diff_methods, argnum)
     var_present = any(isinstance(m, VarianceMP) for m in tape.measurements)
 
     unsupported_params = []
diff --git a/pennylane/gradients/parameter_shift_hessian.py b/pennylane/gradients/parameter_shift_hessian.py
index 11a4c9a48..48c5d0ac7 100644
--- a/pennylane/gradients/parameter_shift_hessian.py
+++ b/pennylane/gradients/parameter_shift_hessian.py
@@ -30,7 +30,7 @@ from .general_shift_rules import (
     generate_multishifted_tapes,
     generate_shifted_tapes,
 )
-from .gradient_transform import gradient_analysis_and_validation
+from .gradient_transform import find_and_validate_gradient_methods
 from .parameter_shift import _get_operation_recipe
 from .hessian_transform import _process_jacs
 
@@ -88,7 +88,7 @@ def _process_argnum(argnum, tape):
     return argnum
 
 
-def _collect_recipes(tape, argnum, diff_methods, diagonal_shifts, off_diagonal_shifts):
+def _collect_recipes(tape, argnum, method_map, diagonal_shifts, off_diagonal_shifts):
     r"""Extract second order recipes for the tape operations for the diagonal of the Hessian
     as well as the first-order derivative recipes for the off-diagonal entries.
     """
@@ -99,7 +99,7 @@ def _collect_recipes(tape, argnum, diff_methods, diagonal_shifts, off_diagonal_s
     partial_offdiag_recipes = []
     diag_shifts_idx = offdiag_shifts_idx = 0
     for i, (d, od) in enumerate(zip(diag_argnum, offdiag_argnum)):
-        if not d or diff_methods[i] == "0":
+        if not d or method_map[i] == "0":
             # hessian will be set to 0 for this row/column
             diag_recipes.append(None)
         else:
@@ -108,7 +108,7 @@ def _collect_recipes(tape, argnum, diff_methods, diagonal_shifts, off_diagonal_s
             diag_recipes.append(_get_operation_recipe(tape, i, diag_shifts, order=2))
             diag_shifts_idx += 1
 
-        if not od or diff_methods[i] == "0":
+        if not od or method_map[i] == "0":
             # hessian will be set to 0 for this row/column
             partial_offdiag_recipes.append((None, None, None))
         else:
@@ -217,9 +217,7 @@ def _all_zero_hessian(tape):
     return [], lambda _: tuple(zeros_list)
 
 
-def expval_hessian_param_shift(
-    tape, argnum, diff_methods, diagonal_shifts, off_diagonal_shifts, f0
-):
+def expval_hessian_param_shift(tape, argnum, method_map, diagonal_shifts, off_diagonal_shifts, f0):
     r"""Generate the Hessian tapes that are used in the computation of the second derivative of a
     quantum tape, using analytical parameter-shift rules to do so exactly. Also define a
     post-processing function to combine the results of evaluating the tapes into the Hessian.
@@ -228,9 +226,9 @@ def expval_hessian_param_shift(
         tape (.QuantumTape): quantum tape to differentiate
         argnum (array_like[bool]): Parameter indices to differentiate
             with respect to, in form of a two-dimensional boolean ``array_like`` mask.
-        diff_methods (list[string]): The differentiation method to use for each trainable parameter.
-            Can be "A" or "0", where "A" is the analytical parameter shift rule and "0" indicates
-            a 0 derivative (that is the parameter does not affect the tape's output).
+        method_map (dict[int, string]): The differentiation method to use for each trainable
+            parameter. Can be "A" or "0", where "A" is the analytical parameter shift rule
+            and "0" indicates a 0 derivative (the parameter does not affect the tape's output).
         diagonal_shifts (list[tuple[int or float]]): List containing tuples of shift values
             for the Hessian diagonal.
             If provided, one tuple of shifts should be given per trainable parameter
@@ -262,7 +260,7 @@ def expval_hessian_param_shift(
     # Assemble all univariate recipes for the diagonal and as partial components for the
     # off-diagonal entries.
     diag_recipes, partial_offdiag_recipes = _collect_recipes(
-        tape, argnum, diff_methods, diagonal_shifts, off_diagonal_shifts
+        tape, argnum, method_map, diagonal_shifts, off_diagonal_shifts
     )
 
     hessian_tapes = []
@@ -574,19 +572,18 @@ def param_shift_hessian(
     # finite-difference as method. If they are among the requested argnum, we catch this
     # further below (as no fallback function in analogy to `param_shift` is used currently).
     method = "analytic" if argnum is None else "best"
-    diff_methods = gradient_analysis_and_validation(tape, method, grad_fn=qml.gradients.param_shift)
+    trainable_params = qml.math.where(qml.math.any(bool_argnum, axis=0))[0]
+    diff_methods = find_and_validate_gradient_methods(tape, method, list(trainable_params))
 
-    for i, g in enumerate(diff_methods):
+    for i, g in diff_methods.items():
         if g == "0":
             bool_argnum[i] = bool_argnum[:, i] = False
     if qml.math.all(~bool_argnum):  # pylint: disable=invalid-unary-operand-type
         return _all_zero_hessian(tape)
 
-    # Find all argument indices that appear in at least one derivative that was requested
-    choose_argnum = qml.math.where(qml.math.any(bool_argnum, axis=0))[0]
     # If any of these argument indices correspond to a finite difference
     # derivative (diff_methods[idx]="F"), raise an error.
-    unsupported_params = {idx for idx in choose_argnum if diff_methods[idx] == "F"}
+    unsupported_params = {i for i, m in diff_methods.items() if m == "F"}
     if unsupported_params:
         raise ValueError(
             "The parameter-shift Hessian currently does not support the operations "
diff --git a/pennylane/gradients/pulse_gradient.py b/pennylane/gradients/pulse_gradient.py
index f6a20ecc8..3777d49b0 100644
--- a/pennylane/gradients/pulse_gradient.py
+++ b/pennylane/gradients/pulse_gradient.py
@@ -31,8 +31,8 @@ from .gradient_transform import (
     assert_no_state_returns,
     assert_no_tape_batching,
     assert_no_variance,
-    choose_grad_methods,
-    gradient_analysis_and_validation,
+    choose_trainable_params,
+    find_and_validate_gradient_methods,
     _no_trainable_grad,
     reorder_grads,
 )
@@ -337,7 +337,8 @@ def stoch_pulse_grad(
         tape (QuantumTape): quantum circuit to differentiate
         argnum (int or list[int] or None): Trainable tape parameter indices to differentiate
             with respect to. If not provided, the derivatives with respect to all
-            trainable parameters are returned.
+            trainable parameters are returned. Note that the indices are with respect to
+            the list of trainable parameters.
         num_split_times (int): number of time samples to use in the stochastic parameter-shift
             rule underlying the differentiation; also see details
         sample_seed (int): randomness seed to be used for the time samples in the stochastic
@@ -621,14 +622,13 @@ 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.")
 
-    diff_methods = gradient_analysis_and_validation(tape, "analytic", grad_fn=stoch_pulse_grad)
+    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):
+    if all(g == "0" for g in diff_methods.values()):
         return _all_zero_grad(tape)
 
-    method_map = choose_grad_methods(diff_methods, argnum)
-
-    argnum = [i for i, dm in method_map.items() if dm == "A"]
+    argnum = [i for i, dm in diff_methods.items() if dm == "A"]
 
     sampler_seed = sampler_seed or np.random.randint(18421)
     key = jax.random.PRNGKey(sampler_seed)
diff --git a/pennylane/gradients/pulse_gradient_odegen.py b/pennylane/gradients/pulse_gradient_odegen.py
index 6c873463f..0ea7d317a 100644
--- a/pennylane/gradients/pulse_gradient_odegen.py
+++ b/pennylane/gradients/pulse_gradient_odegen.py
@@ -32,8 +32,8 @@ from .gradient_transform import (
     assert_no_state_returns,
     assert_no_tape_batching,
     assert_no_variance,
-    choose_grad_methods,
-    gradient_analysis_and_validation,
+    choose_trainable_params,
+    find_and_validate_gradient_methods,
     _no_trainable_grad,
     reorder_grads,
 )
@@ -289,7 +289,8 @@ def _expval_pulse_odegen(tape, argnum, atol):
             parameter-shift rule.
         argnum (int or list[int] or None): Trainable tape parameter indices to differentiate
             with respect to. If not provided, the derivatives with respect to all
-            trainable parameters are returned.
+            trainable parameters are returned. Note that the indices are with respect to
+            the list of trainable parameters.
         atol (float): absolute tolerance used to determine vanishing contributions.
 
     Returns:
@@ -430,7 +431,8 @@ def pulse_odegen(
         tape (QuantumTape): quantum circuit to differentiate
         argnum (int or list[int] or None): Trainable tape parameter indices to differentiate
             with respect to. If not provided, the derivatives with respect to all
-            trainable parameters are returned.
+            trainable parameters are returned. Note that the indices are with respect to
+            the list of trainable parameters.
         atol (float): Precision parameter used to truncate the Pauli basis coefficients
             of the effective generators. Coefficients ``x`` satisfying
             ``qml.math.isclose(x, 0., atol=atol, rtol=0) == True`` are neglected.
@@ -684,14 +686,13 @@ def pulse_odegen(
     if argnum is None and not tape.trainable_params:
         return _no_trainable_grad(tape)
 
-    diff_methods = gradient_analysis_and_validation(tape, "analytic", grad_fn=pulse_odegen)
+    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):
+    if all(g == "0" for g in diff_methods.values()):
         return _all_zero_grad(tape)
 
-    method_map = choose_grad_methods(diff_methods, argnum)
-
-    argnum = [i for i, dm in method_map.items() if dm == "A"]
+    argnum = [i for i, dm in diff_methods.items() if dm == "A"]
 
     return _expval_pulse_odegen(tape, argnum, atol)
 
diff --git a/pennylane/gradients/spsa_gradient.py b/pennylane/gradients/spsa_gradient.py
index 18938ccdd..6e2c508f3 100644
--- a/pennylane/gradients/spsa_gradient.py
+++ b/pennylane/gradients/spsa_gradient.py
@@ -30,8 +30,8 @@ from .finite_difference import _processing_fn, finite_diff_coeffs
 from .gradient_transform import (
     _all_zero_grad,
     assert_no_tape_batching,
-    choose_grad_methods,
-    gradient_analysis_and_validation,
+    choose_trainable_params,
+    find_and_validate_gradient_methods,
     _no_trainable_grad,
 )
 from .general_shift_rules import generate_multishifted_tapes
@@ -112,7 +112,8 @@ def spsa_grad(
         tape (QNode or QuantumTape): quantum circuit to differentiate
         argnum (int or list[int] or None): Trainable parameter indices to differentiate
             with respect to. If not provided, the derivatives with respect to all
-            trainable parameters are returned.
+            trainable parameters are returned. Note that the indices are with respect to
+            the list of trainable parameters.
         h (float or tensor_like[float]): Step size for the finite-difference method
             underlying the SPSA. Can be a tensor-like object
             with as many entries as differentiated *gate* parameters
@@ -246,6 +247,18 @@ def spsa_grad(
         computation need to be analyzed. Here we see that for ``num_directions=1``,
         the default, we obtain two tapes.
 
+        Note that ``argnum`` refers to the index of a parameter within the list of trainable
+        parameters. For example, if we have:
+
+        >>> tape = qml.tape.QuantumScript(
+        ...     [qml.RX(1.2, wires=0), qml.RY(2.3, wires=0), qml.RZ(3.4, wires=0)],
+        ...     [qml.expval(qml.PauliZ(0))],
+        ...     trainable_params = [1, 2]
+        ... )
+        >>> qml.gradients.spsa_grad(tape, argnum=1)
+
+        The code above will differentiate the third parameter rather than the second.
+
         The output tapes can then be evaluated and post-processed to retrieve
         the gradient:
 
@@ -284,12 +297,14 @@ def spsa_grad(
     if argnum is None and not tape.trainable_params:
         return _no_trainable_grad(tape)
 
-    if validate_params:
-        diff_methods = gradient_analysis_and_validation(tape, "numeric", grad_fn=spsa_grad)
-    else:
-        diff_methods = ["F" for i in tape.trainable_params]
+    trainable_params = choose_trainable_params(tape, argnum)
+    diff_methods = (
+        find_and_validate_gradient_methods(tape, "numeric", trainable_params)
+        if validate_params
+        else {idx: "F" for idx in trainable_params}
+    )
 
-    if all(g == "0" for g in diff_methods):
+    if all(g == "0" for g in diff_methods.values()):
         return _all_zero_grad(tape)
 
     gradient_tapes = []
@@ -315,10 +330,10 @@ def spsa_grad(
 
     sampler_rng = np.random.default_rng(sampler_rng)
 
-    method_map = choose_grad_methods(diff_methods, argnum)
-
     num_trainable_params = len(tape.trainable_params)
-    indices = [i for i in range(num_trainable_params) if (i in method_map and method_map[i] != "0")]
+    indices = [
+        i for i in range(num_trainable_params) if (i in diff_methods and diff_methods[i] != "0")
+    ]
 
     tapes_per_grad = len(shifts)
     all_coeffs = []


Test patch:
diff --git a/tests/gradients/core/test_gradient_transform.py b/tests/gradients/core/test_gradient_transform.py
index 4eb735fad..df8697585 100644
--- a/tests/gradients/core/test_gradient_transform.py
+++ b/tests/gradients/core/test_gradient_transform.py
@@ -17,9 +17,9 @@ import pytest
 import pennylane as qml
 from pennylane import numpy as np
 from pennylane.gradients.gradient_transform import (
-    _gradient_analysis,
-    choose_grad_methods,
-    _grad_method_validation,
+    _find_gradient_methods,
+    _validate_gradient_methods,
+    choose_trainable_params,
 )
 
 
@@ -37,25 +37,7 @@ class TestGradAnalysis:
 
     def test_non_differentiable(self):
         """Test that a non-differentiable parameter is correctly marked"""
-        psi = np.array([1, 0, 1, 0]) / np.sqrt(2)
-
-        with qml.queuing.AnnotatedQueue() as q:
-            qml.StatePrep(psi, wires=[0, 1])
-            qml.RX(0.543, wires=[0])
-            qml.RY(-0.654, wires=[1])
-            qml.CNOT(wires=[0, 1])
-            qml.probs(wires=[0, 1])
-
-        tape = qml.tape.QuantumScript.from_queue(q)
-        _gradient_analysis(tape)
-
-        assert tape._par_info[0]["grad_method"] is None
-        assert tape._par_info[1]["grad_method"] == "A"
-        assert tape._par_info[2]["grad_method"] == "A"
 
-    def test_analysis_caching(self, mocker):
-        """Test that the gradient analysis is only executed once per tape
-        if grad_fn is set and unchanged."""
         psi = np.array([1, 0, 1, 0]) / np.sqrt(2)
 
         with qml.queuing.AnnotatedQueue() as q:
@@ -66,17 +48,12 @@ class TestGradAnalysis:
             qml.probs(wires=[0, 1])
 
         tape = qml.tape.QuantumScript.from_queue(q)
-        spy = mocker.spy(qml.operation, "has_grad_method")
-        _gradient_analysis(tape, grad_fn=5)
-        spy.assert_called()
-
-        assert tape._par_info[0]["grad_method"] is None
-        assert tape._par_info[1]["grad_method"] == "A"
-        assert tape._par_info[2]["grad_method"] == "A"
+        trainable_params = choose_trainable_params(tape, None)
+        diff_methods = _find_gradient_methods(tape, trainable_params)
 
-        spy = mocker.spy(qml.operation, "has_grad_method")
-        _gradient_analysis(tape, grad_fn=5)
-        spy.assert_not_called()
+        assert diff_methods[0] is None
+        assert diff_methods[1] == "A"
+        assert diff_methods[2] == "A"
 
     def test_independent(self):
         """Test that an independent variable is properly marked
@@ -88,10 +65,11 @@ class TestGradAnalysis:
             qml.expval(qml.PauliY(0))
 
         tape = qml.tape.QuantumScript.from_queue(q)
-        _gradient_analysis(tape)
+        trainable_params = choose_trainable_params(tape, None)
+        diff_methods = _find_gradient_methods(tape, trainable_params)
 
-        assert tape._par_info[0]["grad_method"] == "A"
-        assert tape._par_info[1]["grad_method"] == "0"
+        assert diff_methods[0] == "A"
+        assert diff_methods[1] == "0"
 
     def test_independent_no_graph_mode(self):
         """In non-graph mode, it is impossible to determine
@@ -103,13 +81,15 @@ class TestGradAnalysis:
             qml.expval(qml.PauliY(0))
 
         tape = qml.tape.QuantumScript.from_queue(q)
-        _gradient_analysis(tape, use_graph=False)
+        trainable_params = choose_trainable_params(tape, None)
+        diff_methods = _find_gradient_methods(tape, trainable_params, use_graph=False)
 
-        assert tape._par_info[0]["grad_method"] == "A"
-        assert tape._par_info[1]["grad_method"] == "A"
+        assert diff_methods[0] == "A"
+        assert diff_methods[1] == "A"
 
     def test_finite_diff(self, monkeypatch):
         """If an op has grad_method=F, this should be respected"""
+
         monkeypatch.setattr(qml.RX, "grad_method", "F")
 
         psi = np.array([1, 0, 1, 0]) / np.sqrt(2)
@@ -122,11 +102,12 @@ class TestGradAnalysis:
             qml.probs(wires=[0, 1])
 
         tape = qml.tape.QuantumScript.from_queue(q)
-        _gradient_analysis(tape)
+        trainable_params = choose_trainable_params(tape, None)
+        diff_methods = _find_gradient_methods(tape, trainable_params)
 
-        assert tape._par_info[0]["grad_method"] is None
-        assert tape._par_info[1]["grad_method"] == "F"
-        assert tape._par_info[2]["grad_method"] == "A"
+        assert diff_methods[0] is None
+        assert diff_methods[1] == "F"
+        assert diff_methods[2] == "A"
 
 
 class TestGradMethodValidation:
@@ -143,10 +124,9 @@ class TestGradMethodValidation:
             qml.RX(np.array(0.1, requires_grad=True), wires=0)
             qml.expval(qml.PauliZ(0))
         tape = qml.tape.QuantumScript.from_queue(q)
-        tape._par_info[0]["grad_method"] = "A"
-        tape._par_info[1]["grad_method"] = None
+        diff_methods = {0: "A", 1: None}
         with pytest.raises(ValueError, match="Cannot differentiate with respect"):
-            _grad_method_validation(method, tape)
+            _validate_gradient_methods(tape, method, diff_methods)
 
     def test_with_numdiff_parameters_and_analytic(self):
         """Test that trainable parameters with numerical grad_method ``"F"``
@@ -156,46 +136,49 @@ class TestGradMethodValidation:
             qml.RX(np.array(0.1, requires_grad=True), wires=0)
             qml.expval(qml.PauliZ(0))
         tape = qml.tape.QuantumScript.from_queue(q)
-        tape._par_info[0]["grad_method"] = "A"
-        tape._par_info[1]["grad_method"] = "F"
+        diff_methods = {
+            0: "A",
+            1: "F",
+        }
         with pytest.raises(ValueError, match="The analytic gradient method cannot be used"):
-            _grad_method_validation("analytic", tape)
-
-
-class TestChooseGradMethods:
-    """Test the helper function choose_grad_methods"""
-
-    all_diff_methods = [
-        ["A"] * 2,
-        [None, "A", "F", "A"],
-        ["F"],
-        ["0", "A"],
-    ]
-
-    @pytest.mark.parametrize("diff_methods", all_diff_methods)
-    def test_without_argnum(self, diff_methods):
-        """Test that the method returns all diff_methods when
-        used with ``argnum=None``."""
-        chosen = choose_grad_methods(diff_methods, None)
-        assert chosen == dict(enumerate(diff_methods))
-
-    @pytest.mark.parametrize(
-        "diff_methods, argnum, expected",
-        zip(all_diff_methods, [1, 2, 0, 0], [{1: "A"}, {2: "F"}, {0: "F"}, {0: "0"}]),
-    )
-    def test_with_integer_argnum(self, diff_methods, argnum, expected):
-        """Test that the method returns the correct single diff_method when
-        used with an integer ``argnum``."""
-        chosen = choose_grad_methods(diff_methods, argnum)
-        assert chosen == expected
+            _validate_gradient_methods(tape, "analytic", diff_methods)
 
-    @pytest.mark.parametrize("diff_methods", all_diff_methods[:2] + [[]])
-    def test_warning_with_empty_argnum(self, diff_methods):
+
+class TestChooseParams:
+    """Test the helper function choose_trainable_params."""
+
+    def test_without_argnum(self):
+        """Test that the method returns all params when used with ``argnum=None``."""
+        tape = qml.tape.QuantumScript(
+            [qml.RX(1.2, wires=0), qml.RY(2.3, wires=0), qml.RZ(3.4, wires=0)],
+            [qml.expval(qml.PauliZ(0))],
+            trainable_params=[1, 2],
+        )
+        chosen = choose_trainable_params(tape, None)
+        assert chosen == [0, 1]
+
+    def test_with_integer_argnum(self):
+        """Test that the method returns the correct single param when
+        used with an integer ``argnum``."""
+        tape = qml.tape.QuantumScript(
+            [qml.RX(1.2, wires=0), qml.RY(2.3, wires=0), qml.RZ(3.4, wires=0)],
+            [qml.expval(qml.PauliZ(0))],
+            trainable_params=[1, 2],
+        )
+        chosen = choose_trainable_params(tape, argnum=1)
+        assert chosen == [1]
+
+    def test_warning_with_empty_argnum(self):
         """Test that the method raises a warning when an empty iterable
         is passed as ``argnum``."""
+        tape = qml.tape.QuantumScript(
+            [qml.RX(1.2, wires=0), qml.RY(2.3, wires=0), qml.RZ(3.4, wires=0)],
+            [qml.expval(qml.PauliZ(0))],
+            trainable_params=[1, 2],
+        )
         with pytest.warns(UserWarning, match="No trainable parameters were specified"):
-            chosen = choose_grad_methods(diff_methods, [])
-        assert chosen == {}
+            chosen = choose_trainable_params(tape, [])
+        assert chosen == []
 
 
 class TestGradientTransformIntegration:
diff --git a/tests/gradients/parameter_shift/test_parameter_shift_cv.py b/tests/gradients/parameter_shift/test_parameter_shift_cv.py
index d1029b504..2e569d25b 100644
--- a/tests/gradients/parameter_shift/test_parameter_shift_cv.py
+++ b/tests/gradients/parameter_shift/test_parameter_shift_cv.py
@@ -14,18 +14,19 @@
 """Tests for the gradients.parameter_shift_cv module."""
 # pylint: disable=protected-access, no-self-use, not-callable, no-value-for-parameter
 
+import unittest.mock as mock
 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.parameter_shift_cv import (
-    _grad_method,
-    _gradient_analysis_cv,
+    _grad_method_cv,
+    _find_gradient_methods_cv,
     _transform_observable,
 )
 
-
 hbar = 2
 
 
@@ -43,21 +44,20 @@ class TestGradAnalysis:
             qml.expval(qml.QuadX(wires=[0]))
 
         tape = qml.tape.QuantumScript.from_queue(q)
-        assert _grad_method(tape, 0) is None
-        assert _grad_method(tape, 1) == "A"
-        assert _grad_method(tape, 2) == "A"
-        assert _grad_method(tape, 3) == "A"
-        assert _grad_method(tape, 4) == "A"
-
-        _gradient_analysis_cv(tape)
+        assert _grad_method_cv(tape, 0) is None
+        assert _grad_method_cv(tape, 1) == "A"
+        assert _grad_method_cv(tape, 2) == "A"
+        assert _grad_method_cv(tape, 3) == "A"
+        assert _grad_method_cv(tape, 4) == "A"
 
-        assert tape._par_info[0]["grad_method"] is None
-        assert tape._par_info[1]["grad_method"] == "A"
-        assert tape._par_info[2]["grad_method"] == "A"
-        assert tape._par_info[3]["grad_method"] == "A"
-        assert tape._par_info[4]["grad_method"] == "A"
+        trainable_params = choose_trainable_params(tape, None)
+        diff_methods = _find_gradient_methods_cv(tape, trainable_params)
 
-        _gradient_analysis_cv(tape)
+        assert diff_methods[0] is None
+        assert diff_methods[1] == "A"
+        assert diff_methods[2] == "A"
+        assert diff_methods[3] == "A"
+        assert diff_methods[4] == "A"
 
     def test_independent(self):
         """Test that an independent variable is properly marked
@@ -69,13 +69,14 @@ class TestGradAnalysis:
             qml.expval(qml.QuadP(0))
 
         tape = qml.tape.QuantumScript.from_queue(q)
-        assert _grad_method(tape, 0) == "A"
-        assert _grad_method(tape, 1) == "0"
+        assert _grad_method_cv(tape, 0) == "A"
+        assert _grad_method_cv(tape, 1) == "0"
 
-        _gradient_analysis_cv(tape)
+        trainable_params = choose_trainable_params(tape, None)
+        diff_methods = _find_gradient_methods_cv(tape, trainable_params)
 
-        assert tape._par_info[0]["grad_method"] == "A"
-        assert tape._par_info[1]["grad_method"] == "0"
+        assert diff_methods[0] == "A"
+        assert diff_methods[1] == "0"
 
     def test_finite_diff(self, monkeypatch):
         """If an op has grad_method=F, this should be respected
@@ -88,9 +89,9 @@ class TestGradAnalysis:
             qml.expval(qml.QuadP(0))
 
         tape = qml.tape.QuantumScript.from_queue(q)
-        assert _grad_method(tape, 0) == "F"
-        assert _grad_method(tape, 1) == "A"
-        assert _grad_method(tape, 2) == "A"
+        assert _grad_method_cv(tape, 0) == "F"
+        assert _grad_method_cv(tape, 1) == "A"
+        assert _grad_method_cv(tape, 2) == "A"
 
     def test_non_gaussian_operation(self):
         """Test that a non-Gaussian operation succeeding
@@ -107,11 +108,11 @@ class TestGradAnalysis:
 
         tape = qml.tape.QuantumScript.from_queue(q)
         # First rotation gate has no succeeding non-Gaussian operation
-        assert _grad_method(tape, 0) == "A"
+        assert _grad_method_cv(tape, 0) == "A"
         # Second rotation gate does no succeeding non-Gaussian operation
-        assert _grad_method(tape, 1) == "F"
+        assert _grad_method_cv(tape, 1) == "F"
         # Kerr gate does not support the parameter-shift rule
-        assert _grad_method(tape, 2) == "F"
+        assert _grad_method_cv(tape, 2) == "F"
 
         with qml.queuing.AnnotatedQueue() as q:
             qml.Rotation(1.0, wires=[0])
@@ -126,9 +127,9 @@ class TestGradAnalysis:
         tape = qml.tape.QuantumScript.from_queue(q)
         # After entangling the modes, the Kerr gate now succeeds
         # both initial rotations
-        assert _grad_method(tape, 0) == "F"
-        assert _grad_method(tape, 1) == "F"
-        assert _grad_method(tape, 2) == "F"
+        assert _grad_method_cv(tape, 0) == "F"
+        assert _grad_method_cv(tape, 1) == "F"
+        assert _grad_method_cv(tape, 2) == "F"
 
     def test_probability(self):
         """Probability is the expectation value of a
@@ -140,9 +141,9 @@ class TestGradAnalysis:
             qml.probs(wires=0)
 
         tape = qml.tape.QuantumScript.from_queue(q)
-        assert _grad_method(tape, 0) == "F"
-        assert _grad_method(tape, 1) == "F"
-        assert _grad_method(tape, 2) == "F"
+        assert _grad_method_cv(tape, 0) == "F"
+        assert _grad_method_cv(tape, 1) == "F"
+        assert _grad_method_cv(tape, 2) == "F"
 
     def test_variance(self):
         """If the variance of the observable is first order, then
@@ -154,14 +155,14 @@ class TestGradAnalysis:
             qml.var(qml.QuadP(0))  # first order
 
         tape = qml.tape.QuantumScript.from_queue(q)
-        assert _grad_method(tape, 0) == "A"
+        assert _grad_method_cv(tape, 0) == "A"
 
         with qml.queuing.AnnotatedQueue() as q:
             qml.Rotation(1.0, wires=[0])
             qml.var(qml.NumberOperator(0))  # second order
 
         tape = qml.tape.QuantumScript.from_queue(q)
-        assert _grad_method(tape, 0) == "F"
+        assert _grad_method_cv(tape, 0) == "F"
 
         with qml.queuing.AnnotatedQueue() as q:
             qml.Rotation(1.0, wires=[0])
@@ -171,10 +172,10 @@ class TestGradAnalysis:
             qml.expval(qml.NumberOperator(1))
 
         tape = qml.tape.QuantumScript.from_queue(q)
-        assert _grad_method(tape, 0) == "F"
-        assert _grad_method(tape, 1) == "F"
-        assert _grad_method(tape, 2) == "F"
-        assert _grad_method(tape, 3) == "F"
+        assert _grad_method_cv(tape, 0) == "F"
+        assert _grad_method_cv(tape, 1) == "F"
+        assert _grad_method_cv(tape, 2) == "F"
+        assert _grad_method_cv(tape, 3) == "F"
 
     def test_second_order_expectation(self):
         """Test that the expectation of a second-order observable forces
@@ -185,7 +186,7 @@ class TestGradAnalysis:
             qml.expval(qml.NumberOperator(0))  # second order
 
         tape = qml.tape.QuantumScript.from_queue(q)
-        assert _grad_method(tape, 0) == "A2"
+        assert _grad_method_cv(tape, 0) == "A2"
 
     def test_unknown_op_grad_method(self, monkeypatch):
         """Test that an exception is raised if an operator has a
@@ -199,7 +200,7 @@ class TestGradAnalysis:
 
         tape = qml.tape.QuantumScript.from_queue(q)
         with pytest.raises(ValueError, match="unknown gradient method"):
-            _grad_method(tape, 0)
+            _grad_method_cv(tape, 0)
 
 
 class TestTransformObservable:
@@ -706,12 +707,13 @@ class TestExpectationQuantumGradients:
         tape.trainable_params = {0}
 
         spy = mocker.spy(qml.gradients.parameter_shift_cv, "second_order_param_shift")
+        spy2 = mocker.spy(qml.gradients.parameter_shift_cv, "_find_gradient_methods_cv")
 
         tapes, fn = param_shift_cv(tape, dev)
         grad = fn(dev.batch_execute(tapes))
-        assert tape._par_info[0]["grad_method"] == "F"
 
         spy.assert_not_called()
+        assert spy2.spy_return == {0: "F"}
 
         # (d/dr) |<2|S(r)>|^2 = 0.5 tanh(r)^3 (2 csch(r)^2 - 1) sech(r)
         expected = 0.5 * np.tanh(r) ** 3 * (2 / (np.sinh(r) ** 2) - 1) / np.cosh(r)
@@ -775,7 +777,7 @@ class TestExpectationQuantumGradients:
     @pytest.mark.parametrize(
         "op", [qml.Displacement(0.1, 0.2, wires=0), qml.TwoModeSqueezing(0.1, 0.2, wires=[0, 1])]
     )
-    def test_gradients_gaussian_circuit(self, op, obs, tol):
+    def test_gradients_gaussian_circuit(self, mocker, op, obs, tol):
         """Tests that the gradients of circuits of gaussian gates match between the
         finite difference and analytic methods."""
         tol = 1e-2
@@ -794,15 +796,16 @@ class TestExpectationQuantumGradients:
 
         tape.trainable_params = set(range(2, 2 + op.num_params))
 
+        spy = mocker.spy(qml.gradients.parameter_shift_cv, "_find_gradient_methods_cv")
+
         tapes, fn = qml.gradients.finite_diff(tape)
         grad_F = fn(dev.batch_execute(tapes))
 
         tapes, fn = param_shift_cv(tape, dev, force_order2=True)
         grad_A2 = fn(dev.batch_execute(tapes))
 
-        # check that every parameter is analytic
-        for i in range(op.num_params):
-            assert tape._par_info[2 + i]["grad_method"][0] == "A"
+        for method in spy.spy_return.values():
+            assert method == "A"
 
         assert np.allclose(grad_A2, grad_F, atol=tol, rtol=0)
 
@@ -812,7 +815,7 @@ class TestExpectationQuantumGradients:
             assert np.allclose(grad_A, grad_F, atol=tol, rtol=0)
 
     @pytest.mark.parametrize("t", [0, 1])
-    def test_interferometer_unitary(self, t, tol):
+    def test_interferometer_unitary(self, mocker, t, tol):
         """An integration test for CV gates that support analytic differentiation
         if succeeding the gate to be differentiated, but cannot be differentiated
         themselves (for example, they may be Gaussian but accept no parameters,
@@ -870,6 +873,8 @@ class TestExpectationQuantumGradients:
 
         dev = qml.device("default.gaussian", wires=2)
 
+        spy = mocker.spy(qml.gradients.parameter_shift_cv, "_find_gradient_methods_cv")
+
         tapes, fn = qml.gradients.finite_diff(tape)
         grad_F = fn(dev.batch_execute(tapes))
 
@@ -879,8 +884,7 @@ class TestExpectationQuantumGradients:
         tapes, fn = param_shift_cv(tape, dev, force_order2=True)
         grad_A2 = fn(dev.batch_execute(tapes))
 
-        assert tape._par_info[0]["grad_method"] == "A"
-        assert tape._par_info[1]["grad_method"] == "A"
+        assert spy.spy_return[0] == "A"
 
         # the different methods agree
         assert np.allclose(grad_A, grad_F, atol=tol, rtol=0)
@@ -992,10 +996,11 @@ class TestVarianceQuantumGradients:
         tape = qml.tape.QuantumScript.from_queue(q)
         tape.trainable_params = {0}
 
-        with pytest.raises(ValueError, match=r"cannot be used with the parameter\(s\) \{0\}"):
+        with pytest.raises(ValueError, match=r"cannot be used with the parameter\(s\) \[0\]"):
             param_shift_cv(tape, dev, fallback_fn=None)
 
-    def test_error_unsupported_grad_recipe(self):
+    @mock.patch("pennylane.gradients.parameter_shift_cv._grad_method_cv", return_value="A")
+    def test_error_unsupported_grad_recipe(self, _):
         """Test exception raised if attempting to use the second order rule for
         computing the gradient analytically of an expectation value that
         contains an operation with more than two terms in the gradient recipe"""
@@ -1018,9 +1023,6 @@ class TestVarianceQuantumGradients:
             qml.expval(qml.QuadX(0))
 
         tape = qml.tape.QuantumScript.from_queue(q)
-        tape._gradient_fn = param_shift_cv
-        tape._par_info[0]["grad_method"] = "A"
-        tape.trainable_params = {0}
 
         with pytest.raises(
             NotImplementedError, match=r"analytic gradient for order-2 operators is unsupported"
@@ -1031,7 +1033,7 @@ class TestVarianceQuantumGradients:
     @pytest.mark.parametrize(
         "op", [qml.Squeezing(0.1, 0.2, wires=0), qml.Beamsplitter(0.1, 0.2, wires=[0, 1])]
     )
-    def test_gradients_gaussian_circuit(self, op, obs, tol):
+    def test_gradients_gaussian_circuit(self, mocker, op, obs, tol):
         """Tests that the gradients of circuits of selected gaussian gates match between the
         finite difference and analytic methods."""
         tol = 1e-2
@@ -1050,6 +1052,8 @@ class TestVarianceQuantumGradients:
 
         tape.trainable_params = set(range(2, 2 + op.num_params))
 
+        spy = mocker.spy(qml.gradients.parameter_shift_cv, "_find_gradient_methods_cv")
+
         # jacobians must match
         tapes, fn = qml.gradients.finite_diff(tape)
         grad_F = fn(dev.batch_execute(tapes))
@@ -1063,10 +1067,8 @@ class TestVarianceQuantumGradients:
         assert np.allclose(grad_A2, grad_F, atol=tol, rtol=0)
         assert np.allclose(grad_A, grad_F, atol=tol, rtol=0)
 
-        # check that every parameter is analytic
         if obs != qml.NumberOperator:
-            for i in range(op.num_params):
-                assert tape._par_info[2 + i]["grad_method"][0] == "A"
+            assert all(m == "A" for m in spy.spy_return.values())
 
     def test_squeezed_mean_photon_variance(self, tol):
         """Test gradient of the photon variance of a displaced thermal state"""