Repository: PennyLaneAI/pennylane

Issue:
[BUG] `argnum` of `stoch_pulse_grad` applied to a tape does not work as expected/documented
### Expected behavior

The `argnum` kwarg of gradient transforms applied to a tape is supposed to index into the set of _trainable_ parameters, not all tape parameters.
That is, for `tape.trainable_params = [1]` and `argnum=0`, the tape parameters with index `1`, which also is marked as trainable, should be differentiated.


### Actual behavior

`stoch_pulse_grad` uses the `argnum` user input to index into _all_ tape parameters. For the example above, this means that the tape parameter with index `0`, which is not marked as trainable, is targeted.
This is an inconsistency compared to other gradient transforms (see first code example). In addition, it disallows differentiation with respect to some parameters (see second code example).

### Additional information

Will have a fix up shortly.

### Source code

```shell
import pennylane as qml
import jax
from jax import numpy as jnp
jax.config.update("jax_enable_x64", True)
H = qml.pulse.constant * qml.PauliX(0)
op = qml.evolve(H)

dev = qml.device("default.qubit")

@qml.qnode(dev)
def circuit(x, a):
    qml.RY(0.3, 0)
    op((a,), t=0.2)
    op((x,), t=0.3)
    return qml.expval(qml.PauliZ(0))

x = jnp.array(0.6)
a = jnp.array(0.2)

circuit(x, a)
print(circuit.tape.trainable_params)
circuit.tape.trainable_params = [1, 2]

# First example: Inconsistency
tapes, fn = qml.gradients.stoch_pulse_grad(circuit.tape, argnum=1)
print(f"stochastic pulse grad, argnum=1: {fn(dev.execute(tapes))}")
# stochastic pulse grad, argnum=1: (Array(-0.16276621, dtype=float64), Array(0., dtype=float64))
tapes, fn = qml.gradients.pulse_odegen(circuit.tape, argnum=1)
print(f"ODEgen pulse grad, argnum=1: {fn(dev.execute(tapes))}")
# ODEgen pulse grad, argnum=1: (Array(0., dtype=float64), Array(-0.24414931, dtype=float64))

# Second example: It's not possible to get the derivative with respect to the second pulse parameter using `stoch_pulse_grad`:
tapes, fn = qml.gradients.stoch_pulse_grad(circuit.tape, argnum=0)
fn(dev.execute(tapes))  # Crashes
tapes, fn = qml.gradients.stoch_pulse_grad(circuit.tape, argnum=1) # Differentiates w.r.t first parameter
tapes, fn = qml.gradients.stoch_pulse_grad(circuit.tape, argnum=2) # Crashes
```


### Tracebacks

```shell
#####################################################
# For argnum=0
#####################################################
---------------------------------------------------------------------------
UnboundLocalError                         Traceback (most recent call last)
Cell In[9], line 29
     26 print(f"ODEgen pulse grad, argnum=0: {fn(dev.execute(tapes))}")
     28 tapes, fn = qml.gradients.stoch_pulse_grad(circuit.tape, argnum=0)
---> 29 fn(dev.execute(tapes))

File ~/repos/pennylane/pennylane/gradients/pulse_gradient.py:848, in _expval_stoch_pulse_grad.<locals>.processing_fn(results)
    844     grads.append(g)
    846 # g will have been defined at least once (because otherwise all gradients would have
    847 # been zero), providing a representative for a zero gradient to emulate its type/shape.
--> 848 zero_rep = _make_zero_rep(g, single_measure, has_partitioned_shots)
    850 # Fill in zero-valued gradients
    851 grads = [zero_rep if g is None else g for g in grads]

UnboundLocalError: local variable 'g' referenced before assignment
#####################################################
# For argnum=2
#####################################################
---------------------------------------------------------------------------
IndexError                                Traceback (most recent call last)
Cell In[6], line 28
     25 tapes, fn = qml.gradients.pulse_odegen(circuit.tape, argnum=1)
     26 print(f"ODEgen pulse grad, argnum=0: {fn(dev.execute(tapes))}")
---> 28 tapes, fn = qml.gradients.stoch_pulse_grad(circuit.tape, argnum=2)

File ~/repos/pennylane/pennylane/transforms/core/transform_dispatcher.py:114, in TransformDispatcher.__call__(self, *targs, **tkwargs)
    111         return expand_processing(processed_results)
    113 else:
--> 114     transformed_tapes, processing_fn = self._transform(obj, *targs, **tkwargs)
    116 if self.is_informative:
    117     return processing_fn(transformed_tapes)

File ~/repos/pennylane/pennylane/gradients/pulse_gradient.py:626, in stoch_pulse_grad(tape, argnum, num_split_times, sampler_seed, use_broadcasting)
    623     raise ValueError("Broadcasting is not supported for tapes that already are broadcasted.")
    625 trainable_params = choose_trainable_params(tape, argnum)
--> 626 diff_methods = find_and_validate_gradient_methods(tape, "analytic", trainable_params)
    628 if all(g == "0" for g in diff_methods.values()):
    629     return _all_zero_grad(tape)

File ~/repos/pennylane/pennylane/gradients/gradient_transform.py:232, in find_and_validate_gradient_methods(tape, method, trainable_param_indices, use_graph)
    201 def find_and_validate_gradient_methods(tape, method, trainable_param_indices, use_graph=True):
    202     """Returns a dictionary of gradient methods for each trainable parameter after
    203     validating if the gradient method requested is supported by the trainable parameters
    204 
   (...)
    230 
    231     """
--> 232     diff_methods = _find_gradient_methods(tape, trainable_param_indices, use_graph=use_graph)
    233     _validate_gradient_methods(tape, method, diff_methods)
    234     return diff_methods

File ~/repos/pennylane/pennylane/gradients/gradient_transform.py:175, in _find_gradient_methods(tape, trainable_param_indices, use_graph)
    172 def _find_gradient_methods(tape, trainable_param_indices, use_graph=True):
    173     """Returns a dictionary with gradient information of each trainable parameter."""
--> 175     return {
    176         idx: _try_zero_grad_from_graph_or_get_grad_method(
    177             tape, tape.trainable_params[idx], use_graph
    178         )
    179         for idx in trainable_param_indices
    180     }

File ~/repos/pennylane/pennylane/gradients/gradient_transform.py:177, in <dictcomp>(.0)
    172 def _find_gradient_methods(tape, trainable_param_indices, use_graph=True):
    173     """Returns a dictionary with gradient information of each trainable parameter."""
    175     return {
    176         idx: _try_zero_grad_from_graph_or_get_grad_method(
--> 177             tape, tape.trainable_params[idx], use_graph
    178         )
    179         for idx in trainable_param_indices
    180     }

IndexError: list index out of range
```


### System information

```shell
pl dev
```


### Existing GitHub issues

- [X] I have searched existing GitHub issues to make sure the issue does not already exist.


Discussion:

Resolving patch:
diff --git a/pennylane/gradients/pulse_gradient.py b/pennylane/gradients/pulse_gradient.py
--- a/pennylane/gradients/pulse_gradient.py
+++ b/pennylane/gradients/pulse_gradient.py
@@ -787,8 +787,8 @@ def _expval_stoch_pulse_grad(tape, argnum, num_split_times, key, use_broadcastin
     """
     tapes = []
     gradient_data = []
-    for idx, trainable_idx in enumerate(tape.trainable_params):
-        if trainable_idx not in argnum:
+    for idx in range(tape.num_params):
+        if idx not in argnum:
             # Only the number of tapes is needed to indicate a zero gradient entry
             gradient_data.append((0, None, None, None))
             continue


Test patch:
diff --git a/tests/gradients/core/test_pulse_gradient.py b/tests/gradients/core/test_pulse_gradient.py
--- a/tests/gradients/core/test_pulse_gradient.py
+++ b/tests/gradients/core/test_pulse_gradient.py
@@ -913,14 +913,14 @@ def test_constant_ry(self, dev_name, num_split_times, t):
 
         jax.config.update("jax_enable_x64", True)
         params = [jnp.array(0.24)]
-        T = t if isinstance(t, tuple) else (0, t)
+        delta_t = t[-1] - t[0] if isinstance(t, tuple) else t
         ham_single_q_const = qml.pulse.constant * qml.PauliY(0)
         op = qml.evolve(ham_single_q_const)(params, t)
         tape = qml.tape.QuantumScript([op], [qml.expval(qml.PauliZ(0))])
 
         dev = qml.device(dev_name, wires=1)
         # Effective rotation parameter
-        p = params[0] * (delta_t := T[-1] - T[0])
+        p = params[0] * delta_t
         r = qml.execute([tape], dev, None)
         assert qml.math.isclose(r, jnp.cos(2 * p), atol=1e-4)
         tapes, fn = stoch_pulse_grad(tape, num_split_times=num_split_times)
@@ -930,6 +930,35 @@ def test_constant_ry(self, dev_name, num_split_times, t):
         assert qml.math.isclose(res, -2 * jnp.sin(2 * p) * delta_t)
         jax.clear_caches()
 
+    def test_constant_ry_argnum(self, dev_name):
+        """Test that the derivative of a pulse generated by a constant Hamiltonian,
+        which is a Pauli word, is computed correctly if it is not the only
+        operation in a tape but selected via `argnum`."""
+        import jax
+        import jax.numpy as jnp
+
+        jax.config.update("jax_enable_x64", True)
+        params = [jnp.array(0.04)]
+        t = 0.1
+        y = 0.3
+        ham_single_q_const = qml.pulse.constant * qml.PauliY(0)
+        op = qml.evolve(ham_single_q_const)(params, t)
+        tape = qml.tape.QuantumScript([qml.RY(y, 0), op], [qml.expval(qml.PauliZ(0))])
+        tape.trainable_params = [1]
+
+        dev = qml.device(dev_name, wires=1)
+        # Effective rotation parameter
+        p = params[0] * t
+        r = qml.execute([tape], dev, None)
+        assert qml.math.isclose(r, jnp.cos(2 * p + y), atol=1e-4)
+        num_split_times = 1
+        tapes, fn = stoch_pulse_grad(tape, num_split_times=num_split_times, argnum=0)
+        assert len(tapes) == num_split_times * 2
+
+        res = fn(qml.execute(tapes, dev, None))
+        assert qml.math.isclose(res, -2 * jnp.sin(2 * p + y) * t)
+        jax.clear_caches()
+
     @pytest.mark.parametrize("num_split_times", [1, 3])
     @pytest.mark.parametrize("t", [2.0, 3, (0.5, 0.6), (0.1, 0.9, 1.2)])
     def test_constant_ry_rescaled(self, dev_name, num_split_times, t):