Repository: PennyLaneAI/pennylane

Issue:
[BUG] Custom gradient recipe with multiple 0 shifts for same gate fails with `param_shift` transform
### Expected behavior

Any custom gradient recipe executes without errors.

### Actual behavior

An indexing error is raised for a gate with the custom recipe of `[[-1e7, 1, 0], [1e7, 1, 0]]`.

### Additional information

_No response_

### Source code

```shell
import pennylane as qml
import numpy as np

dev = qml.device("default.qubit", wires=2)
x = [0.543, -0.654]

ops_with_custom_recipe = [0]

with qml.tape.QuantumTape() as tape:
    qml.RX(x[0], wires=[0])
    qml.RX(x[1], wires=[0])
    qml.expval(qml.PauliZ(0))

gradient_recipes = tuple(
    [[-1e7, 1, 0], [1e7, 1, 0]] if i in ops_with_custom_recipe else None
    for i in range(2)
)
tapes, fn = qml.gradients.param_shift(tape, gradient_recipes=gradient_recipes)

# two tapes per parameter that doesn't use a custom recipe,
# one tape per parameter that uses custom recipe,
# plus one global call if at least one uses the custom recipe
num_ops_standard_recipe = tape.num_params - len(ops_with_custom_recipe)
assert len(tapes) == 2 * num_ops_standard_recipe + len(ops_with_custom_recipe) + 1
# Test that executing the tapes and the postprocessing function works
grad = fn(qml.execute(tapes, dev, None))
assert qml.math.allclose(grad, -np.sin(x[0] + x[1]), atol=1e-5)
```


### Tracebacks

```
~/pennylane/gradients/parameter_shift.py in param_shift(tape, argnum, shifts, gradient_recipes, fallback_fn, f0, broadcast)
   1572         g_tapes, fn = var_param_shift(tape, argnum, shifts, gradient_recipes, f0, broadcast)
   1573     else:
-> 1574         g_tapes, fn = expval_param_shift(tape, argnum, shifts, gradient_recipes, f0, broadcast)
   1575 
   1576     gradient_tapes.extend(g_tapes)

~/pennylane/gradients/parameter_shift.py in expval_param_shift(tape, argnum, shifts, gradient_recipes, f0, broadcast)
    489         # If broadcast=True, g_tapes only contains one tape. If broadcast=False, all returned
    490         # tapes will have the same batch_size=None. Thus we only use g_tapes[0].batch_size here.
--> 491         gradient_data.append((len(g_tapes), coeffs, None, unshifted_coeff, g_tapes[0].batch_size))
    492 
    493     def processing_fn(results):

IndexError: tuple index out of range
```

### System information

```shell
Name: PennyLane
Version: 0.27.0.dev0
Summary: PennyLane is a Python quantum machine learning library by Xanadu Inc.
Home-page: https://github.com/XanaduAI/pennylane
Author: 
Author-email: 
License: Apache License 2.0
Location: /xanadu/pennylane
Requires: appdirs, autograd, autoray, cachetools, networkx, numpy, pennylane-lightning, retworkx, scipy, semantic-version, toml
Required-by: amazon-braket-pennylane-plugin, PennyLane-AQT, PennyLane-Cirq, PennyLane-Forest, PennyLane-Honeywell, PennyLane-IonQ, PennyLane-Lightning, PennyLane-Orquestra, PennyLane-PQ, PennyLane-Qchem, PennyLane-qiskit, PennyLane-qsharp, pennylane-qulacs, PennyLane-SF

Platform info:           Linux-5.15.0-50-generic-x86_64-with-glibc2.10
Python version:          3.8.5
Numpy version:           1.23.0
Scipy version:           1.8.1
Installed devices:
- braket.aws.qubit (amazon-braket-pennylane-plugin-1.7.0)
- braket.local.qubit (amazon-braket-pennylane-plugin-1.7.0)
- lightning.qubit (PennyLane-Lightning-0.26.0)
- forest.numpy_wavefunction (PennyLane-Forest-0.14.0)
- forest.qvm (PennyLane-Forest-0.14.0)
- forest.wavefunction (PennyLane-Forest-0.14.0)
- orquestra.forest (PennyLane-Orquestra-0.15.0)
- orquestra.ibmq (PennyLane-Orquestra-0.15.0)
- orquestra.qiskit (PennyLane-Orquestra-0.15.0)
- orquestra.qulacs (PennyLane-Orquestra-0.15.0)
- qulacs.simulator (pennylane-qulacs-0.17.0.dev0)
- aqt.noisy_sim (PennyLane-AQT-0.18.0)
- aqt.sim (PennyLane-AQT-0.18.0)
- projectq.classical (PennyLane-PQ-0.18.0.dev0)
- projectq.ibm (PennyLane-PQ-0.18.0.dev0)
- projectq.simulator (PennyLane-PQ-0.18.0.dev0)
- microsoft.QuantumSimulator (PennyLane-qsharp-0.19.0)
- strawberryfields.fock (PennyLane-SF-0.21.0.dev0)
- strawberryfields.gaussian (PennyLane-SF-0.21.0.dev0)
- strawberryfields.gbs (PennyLane-SF-0.21.0.dev0)
- strawberryfields.remote (PennyLane-SF-0.21.0.dev0)
- strawberryfields.tf (PennyLane-SF-0.21.0.dev0)
- ionq.qpu (PennyLane-IonQ-0.24.0.dev0)
- ionq.simulator (PennyLane-IonQ-0.24.0.dev0)
- honeywell.hqs (PennyLane-Honeywell-0.19.0.dev0)
- qiskit.aer (PennyLane-qiskit-0.23.0.dev0)
- qiskit.basicaer (PennyLane-qiskit-0.23.0.dev0)
- qiskit.ibmq (PennyLane-qiskit-0.23.0.dev0)
- qiskit.ibmq.circuit_runner (PennyLane-qiskit-0.23.0.dev0)
- qiskit.ibmq.sampler (PennyLane-qiskit-0.23.0.dev0)
- cirq.mixedsimulator (PennyLane-Cirq-0.23.0.dev0)
- cirq.pasqal (PennyLane-Cirq-0.23.0.dev0)
- cirq.qsim (PennyLane-Cirq-0.23.0.dev0)
- cirq.qsimh (PennyLane-Cirq-0.23.0.dev0)
- cirq.simulator (PennyLane-Cirq-0.23.0.dev0)
- default.gaussian (PennyLane-0.27.0.dev0)
- default.mixed (PennyLane-0.27.0.dev0)
- default.qubit (PennyLane-0.27.0.dev0)
- default.qubit.autograd (PennyLane-0.27.0.dev0)
- default.qubit.jax (PennyLane-0.27.0.dev0)
- default.qubit.tf (PennyLane-0.27.0.dev0)
- default.qubit.torch (PennyLane-0.27.0.dev0)
- default.qutrit (PennyLane-0.27.0.dev0)
- null.qubit (PennyLane-0.27.0.dev0)
```


### Existing GitHub issues

- [X] I have searched existing GitHub issues to make sure the issue does not already exist.


Discussion:
@dwierichs any thoughts on the reason for the issue and/or if it could be a potentially significant use case? :thinking: 
The problem here is not that there are multiple zero-shift terms, it is that there are _no nonzero-shift_ terms :) Terms with the same shift (and multiplier) are actually grouped together before creating the tapes.

This bug happens because the code inherently assumes that a non-zero shift will be needed to obtain the derivative (this assumption is not used for Hamiltonian parameters, just for operations).
For all operations I know this makes sense, but maybe there are some for which it doesn't.
It would not be hard to change this, should I do so? Or would we prefer a warning? I would not see any silent bug that would result from the former, i.e. just allowing derivatives with non-zero shifts only.
In particular, this would also allow "dummy" recipes without any terms at all, for whatever reason one might want to use those :grin: 

Frankly, I think any type of handling of the case (either changing the behaviour and/or having a warning) could be nice. I suspect it's not going to be a case widely stumbled upon, though it could confuse some users because it raises an error.
Just making sure: The number of tapes will be `2*num_ops_standard_recipe+int(len(ops_with_custom_recipe)>0)`, right? 2 shifted evaluations per "normal" recipe operation and one global one if at least one op has the custom recipe.

Also: This recipe doesn't work, so the last assertion should not go through, correct? :)

Resolving patch:
diff --git a/pennylane/gradients/parameter_shift.py b/pennylane/gradients/parameter_shift.py
--- a/pennylane/gradients/parameter_shift.py
+++ b/pennylane/gradients/parameter_shift.py
@@ -143,11 +143,20 @@ def _extract_unshifted(recipe, at_least_one_unshifted, f0, gradient_tapes, tape)
 
 
 def _single_meas_grad(result, coeffs, unshifted_coeff, r0):
-    """Compute the gradient for a single measurement by taking the linear combination of the coefficients and the
-    measurement result.
+    """Compute the gradient for a single measurement by taking the linear combination of
+    the coefficients and the measurement result.
 
     If an unshifted term exists, its contribution is added to the gradient.
     """
+    if isinstance(result, list) and result == []:
+        if unshifted_coeff is None:
+            raise ValueError(
+                "This gradient component neither has a shifted nor an unshifted component. "
+                "It should have been identified to have a vanishing gradient earlier on."
+            )  # pragma: no cover
+        # return the unshifted term, which is the only contribution
+        return qml.math.array(unshifted_coeff * r0)
+
     result = qml.math.stack(result)
     coeffs = qml.math.convert_like(coeffs, result)
     g = qml.math.tensordot(result, coeffs, [[0], [0]])
@@ -159,13 +168,12 @@ def _single_meas_grad(result, coeffs, unshifted_coeff, r0):
 
 
 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."""
+    """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[meas_idx])
@@ -175,8 +183,9 @@ def _multi_meas_grad(res, coeffs, r0, unshifted_coeff, num_measurements):
 
 
 def _evaluate_gradient_new(tape, res, data, r0, shots):
-    """Use shifted tape evaluations and parameter-shift rule coefficients
-    to evaluate a gradient result.
+    """Use shifted tape evaluations and parameter-shift rule coefficients to evaluate
+    a gradient result. If res is an empty list, ``r0`` and ``data[3]``, which is the
+    coefficient for the unshifted term, must be given and not None.
 
     This is a helper function for the new return type system.
     """
@@ -229,11 +238,16 @@ def _evaluate_gradient_new(tape, res, data, r0, shots):
 
 
 def _evaluate_gradient(res, data, broadcast, r0, scalar_qfunc_output):
-    """Use shifted tape evaluations and parameter-shift rule coefficients
-    to evaluate a gradient result."""
+    """Use shifted tape evaluations and parameter-shift rule coefficients to evaluate
+    a gradient result. If res is an empty list, ``r0`` and ``data[3]``, which is the
+    coefficient for the unshifted term, must be given and not None."""
 
     _, coeffs, fn, unshifted_coeff, batch_size = data
 
+    if isinstance(res, list) and len(res) == 0:
+        # No shifted evaluations are present, just the unshifted one.
+        return r0 * unshifted_coeff
+
     # individual post-processing of e.g. Hamiltonian grad tapes
     if fn is not None:
         res = fn(res)
@@ -507,7 +521,9 @@ def _expval_param_shift_tuple(
         gradient_tapes.extend(g_tapes)
         # If broadcast=True, g_tapes only contains one tape. If broadcast=False, all returned
         # tapes will have the same batch_size=None. Thus we only use g_tapes[0].batch_size here.
-        gradient_data.append((len(g_tapes), coeffs, None, unshifted_coeff, g_tapes[0].batch_size))
+        # If no gradient tapes are returned (e.g. only unshifted term in recipe), batch_size=None
+        batch_size = g_tapes[0].batch_size if g_tapes else None
+        gradient_data.append((len(g_tapes), coeffs, None, unshifted_coeff, batch_size))
 
     def processing_fn(results):
         start, r0 = (1, results[0]) if at_least_one_unshifted and f0 is None else (0, f0)
@@ -518,11 +534,16 @@ def processing_fn(results):
         grads = []
         for data in gradient_data:
 
-            num_tapes, *_, batch_size = data
+            num_tapes, *_, unshifted_coeff, batch_size = data
             if num_tapes == 0:
-                # parameter has zero gradient. We don't know the output shape yet, so just memorize
-                # that this gradient will be set to zero, via grad = None
-                grads.append(None)
+                if unshifted_coeff is None:
+                    # parameter has zero gradient. We don't know the output shape yet, so just
+                    # memorize that this gradient will be set to zero, via grad = None
+                    grads.append(None)
+                    continue
+                # The gradient for this parameter is computed from r0 alone.
+                g = _evaluate_gradient_new(tape, [], data, r0, shots)
+                grads.append(g)
                 continue
 
             res = results[start : start + num_tapes] if batch_size is None else results[start]
@@ -531,6 +552,8 @@ def processing_fn(results):
             g = _evaluate_gradient_new(tape, res, data, r0, shots)
             grads.append(g)
 
+        # g will have been defined at least once (because otherwise all gradients would have
+        # been zero), providing a representative for a zero gradient to emulate its type/shape.
         if single_measure and not shot_vector:
             zero_rep = qml.math.zeros_like(g)
         elif single_measure:
@@ -664,7 +687,9 @@ def expval_param_shift(
 
         # If broadcast=True, g_tapes only contains one tape. If broadcast=False, all returned
         # tapes will have the same batch_size=None. Thus we only use g_tapes[0].batch_size here.
-        gradient_data.append((len(g_tapes), coeffs, None, unshifted_coeff, g_tapes[0].batch_size))
+        # If no gradient tapes are returned (e.g. only unshifted term in recipe), batch_size=None
+        batch_size = g_tapes[0].batch_size if g_tapes else None
+        gradient_data.append((len(g_tapes), coeffs, None, unshifted_coeff, batch_size))
 
     def processing_fn(results):
         # Apply the same squeezing as in qml.QNode to make the transform output consistent.
@@ -680,22 +705,29 @@ def processing_fn(results):
 
         for data in gradient_data:
 
-            num_tapes, *_, batch_size = data
+            num_tapes, *_, unshifted_coeff, batch_size = data
             if num_tapes == 0:
                 # parameter has zero gradient. We don't know the output shape yet, so just memorize
                 # that this gradient will be set to zero, via grad = None
-                grads.append(None)
+                if unshifted_coeff is None:
+                    # parameter has zero gradient. We don't know the output shape yet, so just
+                    # memorize that this gradient will be set to zero, via grad = None
+                    grads.append(None)
+                    continue
+                # The gradient for this parameter is computed from r0 alone.
+                g = _evaluate_gradient([], data, broadcast, r0, scalar_qfunc_output)
+                grads.append(g)
                 continue
 
             res = results[start : start + num_tapes] if batch_size is None else results[start]
             start = start + num_tapes
 
             g = _evaluate_gradient(res, data, broadcast, r0, scalar_qfunc_output)
-
             grads.append(g)
-            # This clause will be hit at least once (because otherwise all gradients would have
-            # been zero), providing a representative for a zero gradient to emulate its type/shape.
-            zero_rep = qml.math.zeros_like(g)
+
+        # g will have been defined at least once (because otherwise all gradients would have
+        # been zero), providing a representative for a zero gradient to emulate its type/shape.
+        zero_rep = qml.math.zeros_like(g)
 
         for i, g in enumerate(grads):
             # Fill in zero-valued gradients


Test patch:
diff --git a/tests/gradients/test_parameter_shift.py b/tests/gradients/test_parameter_shift.py
--- a/tests/gradients/test_parameter_shift.py
+++ b/tests/gradients/test_parameter_shift.py
@@ -326,8 +326,41 @@ def test_recycled_unshifted_tape(self, ops_with_custom_recipe):
         assert len(tapes) == 2 * num_ops_standard_recipe + len(ops_with_custom_recipe) + 1
         # Test that executing the tapes and the postprocessing function works
         grad = fn(qml.execute(tapes, dev, None))
+        assert qml.math.shape(grad) == (1, 2)
         assert qml.math.allclose(grad, -np.sin(x[0] + x[1]), atol=1e-5)
 
+    @pytest.mark.parametrize("ops_with_custom_recipe", [[0], [1], [0, 1]])
+    def test_custom_recipe_unshifted_only(self, ops_with_custom_recipe):
+        """Test that if the gradient recipe has a zero-shift component, then
+        the tape is executed only once using the current parameter
+        values."""
+        dev = qml.device("default.qubit", wires=2)
+        x = [0.543, -0.654]
+
+        with qml.tape.QuantumTape() as tape:
+            qml.RX(x[0], wires=[0])
+            qml.RX(x[1], wires=[0])
+            qml.expval(qml.PauliZ(0))
+
+        gradient_recipes = tuple(
+            [[-1e7, 1, 0], [1e7, 1, 0]] if i in ops_with_custom_recipe else None for i in range(2)
+        )
+        tapes, fn = qml.gradients.param_shift(tape, gradient_recipes=gradient_recipes)
+
+        # two tapes per parameter that doesn't use a custom recipe,
+        # plus one global (unshifted) call if at least one uses the custom recipe
+        num_ops_standard_recipe = tape.num_params - len(ops_with_custom_recipe)
+        assert len(tapes) == 2 * num_ops_standard_recipe + int(
+            tape.num_params != num_ops_standard_recipe
+        )
+        # Test that executing the tapes and the postprocessing function works
+        grad = fn(qml.execute(tapes, dev, None))
+        expected = [
+            -np.sin(x[0] + x[1]) if i not in ops_with_custom_recipe else 0 for i in range(2)
+        ]
+        assert qml.math.shape(grad) == (1, 2)
+        assert qml.math.allclose(grad, expected, atol=1e-5)
+
     @pytest.mark.parametrize("y_wire", [0, 1])
     def test_f0_provided(self, y_wire):
         """Test that if the original tape output is provided, then
@@ -383,6 +416,7 @@ class RX(qml.RX):
             assert tape.operations[1].data[0] == x[1] + expected[1]
 
         grad = fn(dev.batch_execute(tapes))
+        assert qml.math.shape(grad) == (1, 2)
         assert np.allclose(
             grad, [-np.sin(x[0] + x[1]), -np.sin(x[0] + x[1]) + 0.2 * np.cos(x[0] + x[1])]
         )
diff --git a/tests/returntypes/test_parameter_shift_new.py b/tests/returntypes/test_parameter_shift_new.py
--- a/tests/returntypes/test_parameter_shift_new.py
+++ b/tests/returntypes/test_parameter_shift_new.py
@@ -460,6 +460,76 @@ def test_recycled_unshifted_tape(self, ops_with_custom_recipe):
         grad = fn(qml.execute(tapes, dev, None))
         assert qml.math.allclose(grad, -np.sin(x[0] + x[1]), atol=1e-5)
 
+    @pytest.mark.parametrize("ops_with_custom_recipe", [[0], [1], [0, 1]])
+    @pytest.mark.parametrize("multi_measure", [False, True])
+    def test_custom_recipe_unshifted_only(self, ops_with_custom_recipe, multi_measure):
+        """Test that if the gradient recipe has a zero-shift component, then
+        the tape is executed only once using the current parameter
+        values."""
+        dev = qml.device("default.qubit", wires=2)
+        x = [0.543, -0.654]
+
+        with qml.tape.QuantumTape() as tape:
+            qml.RX(x[0], wires=[0])
+            qml.RX(x[1], wires=[0])
+            qml.expval(qml.PauliZ(0))
+            if multi_measure:
+                qml.expval(qml.PauliZ(1))
+
+        gradient_recipes = tuple(
+            [[-1e7, 1, 0], [1e7, 1, 0]] if i in ops_with_custom_recipe else None for i in range(2)
+        )
+        tapes, fn = qml.gradients.param_shift(tape, gradient_recipes=gradient_recipes)
+
+        # two tapes per parameter that doesn't use a custom recipe,
+        # plus one global (unshifted) call if at least one uses the custom recipe
+        num_ops_standard_recipe = tape.num_params - len(ops_with_custom_recipe)
+        assert len(tapes) == 2 * num_ops_standard_recipe + int(
+            tape.num_params != num_ops_standard_recipe
+        )
+        # Test that executing the tapes and the postprocessing function works
+        grad = fn(qml.execute(tapes, dev, None))
+        if multi_measure:
+            expected = np.array([[-np.sin(x[0] + x[1])] * 2, [0, 0]])
+            # The custom recipe estimates gradients to be 0
+            for i in ops_with_custom_recipe:
+                expected[0, i] = 0
+        else:
+            expected = [
+                -np.sin(x[0] + x[1]) if i not in ops_with_custom_recipe else 0 for i in range(2)
+            ]
+        assert qml.math.allclose(grad, expected, atol=1e-5)
+
+    @pytest.mark.parametrize("ops_with_custom_recipe", [[0], [1], [0, 1]])
+    def test_custom_recipe_mixing_unshifted_shifted(self, ops_with_custom_recipe):
+        """Test that if the gradient recipe has a zero-shift component, then
+        the tape is executed only once using the current parameter
+        values."""
+        dev = qml.device("default.qubit", wires=2)
+        x = [0.543, -0.654]
+
+        with qml.tape.QuantumTape() as tape:
+            qml.RX(x[0], wires=[0])
+            qml.RX(x[1], wires=[0])
+            qml.expval(qml.PauliZ(0))
+            qml.expval(qml.PauliZ(1))
+
+        gradient_recipes = tuple(
+            [[-1e-7, 1, 0], [1e-7, 1, 0], [-1e5, 1, -5e-6], [1e5, 1, 5e-6]]
+            if i in ops_with_custom_recipe
+            else None
+            for i in range(2)
+        )
+        tapes, fn = qml.gradients.param_shift(tape, gradient_recipes=gradient_recipes)
+
+        # two tapes per parameter, independent of recipe
+        # plus one global (unshifted) call if at least one uses the custom recipe
+        assert len(tapes) == 2 * tape.num_params + int(len(ops_with_custom_recipe) > 0)
+        # Test that executing the tapes and the postprocessing function works
+        grad = fn(qml.execute(tapes, dev, None))
+        assert qml.math.allclose(grad[0], -np.sin(x[0] + x[1]), atol=1e-5)
+        assert qml.math.allclose(grad[1], 0, atol=1e-5)
+
     @pytest.mark.parametrize("y_wire", [0, 1])
     def test_f0_provided(self, y_wire):
         """Test that if the original tape output is provided, then