Repository: PennyLaneAI/pennylane
Pull request #3280 (2022-11-07 19:17:27)

[BUG] [Return-types] Fix param shift recipe
The following fails:
```py
qml.enable_return()
dev = qml.device("default.qubit", wires=2)
@qml.qnode(dev)
def circuit(x):
    qml.RX(x, 0)
    return qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliX(1))

# This is not a correct gradient recipe, but that's not the issue
recipes = ([[1, 1, 0], [-1e5, 1, -1e5]],)
x = pnp.array(0.4, requires_grad=True)
circuit(x)
tapes, fn = qml.gradients.param_shift(circuit.tape, gradient_recipes=recipes)
print( fn(qml.execute(tapes, dev, None)))
```
The scope should be:
* Tapes with multiple returned measurements together with;
* Gradient recipes that contain unshifted terms.

**Changes**

The issue is because `r0` is not being indexed in an aux function here:
https://github.com/PennyLaneAI/pennylane/blob/21b35b6ad368729b0cb9a10071af8ca4c518538c/pennylane/gradients/parameter_shift.py#L169

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.
dwierichs: I discovered  a bug in how `r0` is determined when it is not used (-> silent bug) that surfaces in the failing test here. Will push a fix shortly.
codecov[bot]: # [Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3280?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
> :exclamation: No coverage uploaded for pull request base (`v0.27.0-rc0@81ee743`). [Click here to learn what that means](https://docs.codecov.io/docs/error-reference?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#section-missing-base-commit).
> The diff coverage is `n/a`.

```diff
@@              Coverage Diff               @@
##             v0.27.0-rc0    #3280   +/-   ##
==============================================
  Coverage               ?   99.70%           
==============================================
  Files                  ?      299           
  Lines                  ?    26123           
  Branches               ?        0           
==============================================
  Hits                   ?    26046           
  Misses                 ?       77           
  Partials               ?        0           
```



Help us with your feedback. Take ten seconds to tell us [how you rate us](https://about.codecov.io/nps?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI). Have a feature suggestion? [Share it here.](https://app.codecov.io/gh/feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI)


Patch:
diff --git a/pennylane/gradients/parameter_shift.py b/pennylane/gradients/parameter_shift.py
index 8f30c661e..2cb4bf586 100644
--- a/pennylane/gradients/parameter_shift.py
+++ b/pennylane/gradients/parameter_shift.py
@@ -162,11 +162,13 @@ def _multi_meas_grad(res, coeffs, r0, unshifted_coeff, num_measurements):
     """Compute the gradient for multiple measurements by taking the linear combination of the coefficients and each
     measurement result."""
     g = []
+    if r0 is None:
+        r0 = [None] * num_measurements
     for meas_idx in range(num_measurements):
 
         # Gather the measurement results
         meas_result = [param_result[meas_idx] for param_result in res]
-        g_component = _single_meas_grad(meas_result, coeffs, unshifted_coeff, r0)
+        g_component = _single_meas_grad(meas_result, coeffs, unshifted_coeff, r0[meas_idx])
         g.append(g_component)
 
     return tuple(g)
@@ -197,6 +199,8 @@ def _evaluate_gradient_new(tape, res, data, r0, shots):
         # Res has order of axes:
         # 1. Number of parameters
         # 2. Shot vector
+        if r0 is None:
+            r0 = [None] * len_shot_vec
         for i in range(len_shot_vec):
             shot_comp_res = [r[i] for r in res]
             shot_comp_res = _single_meas_grad(shot_comp_res, coeffs, unshifted_coeff, r0[i])
@@ -506,9 +510,7 @@ def _expval_param_shift_tuple(
         gradient_data.append((len(g_tapes), coeffs, None, unshifted_coeff, g_tapes[0].batch_size))
 
     def processing_fn(results):
-        start = 1 if at_least_one_unshifted and f0 is None else 0
-        r0 = f0 or results[0]
-
+        start, r0 = (1, results[0]) if at_least_one_unshifted and f0 is None else (0, f0)
         single_measure = len(tape.measurements) == 1
         single_param = len(tape.trainable_params) == 1
         shot_vector = isinstance(shots, Sequence)
@@ -674,8 +676,7 @@ def expval_param_shift(
             results = [qml.math.squeeze(res) for res in results]
 
         grads = []
-        start = 1 if at_least_one_unshifted and f0 is None else 0
-        r0 = f0 or results[0]
+        start, r0 = (1, results[0]) if at_least_one_unshifted and f0 is None else (0, f0)
 
         for data in gradient_data:
 


Test patch:
diff --git a/tests/returntypes/test_parameter_shift_new.py b/tests/returntypes/test_parameter_shift_new.py
index 7b5c85b55..2c827a306 100644
--- a/tests/returntypes/test_parameter_shift_new.py
+++ b/tests/returntypes/test_parameter_shift_new.py
@@ -57,6 +57,22 @@ class TestGetOperationRecipe:
             assert qml.math.allclose(np.sort(s), exp_out_shifts)
             assert qml.math.allclose(np.sort(out_recipe[:, 2]), np.sort(exp_out_shifts))
 
+    def test_qnode_custom_recipe(self):
+        """Test a custom recipe using a QNode."""
+        dev = qml.device("default.qubit", wires=2)
+
+        x = np.array(0.4, requires_grad=True)
+        with qml.tape.QuantumTape() as tape:
+            qml.RX(x, 0)
+            qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliX(1))
+
+        # This is not a correct gradient recipe, but that's not the issue
+        recipes = ([[-1e7, 1, 0], [1e7, 1, 1e7]],)
+        tapes, fn = qml.gradients.param_shift(tape, gradient_recipes=recipes)
+        res = fn(qml.execute(tapes, dev, None))
+        assert len(res) == 2
+        assert isinstance(res, tuple)
+
     @pytest.mark.parametrize(
         "orig_op, frequencies, shifts",
         [