[BUG] argnum of stoch_pulse_grad applied to a tape does not work as expected/documented

[dwierichs]

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

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

#####################################################
# 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

pl dev

Existing GitHub issues

• I have searched existing GitHub issues to make sure the issue does not already exist.