Repository: PennyLaneAI/pennylane (python)

Issue:
[BUG] First-order derivative with unshifted shift term in gradient recipe breaks `param_shift`
### Expected behavior

The gradient is computed correctly by `qml.gradient.param_shift` even when using an operation that contains an unshifted term in its first-order gradient rule, provided via `gradient_recipe`.

### Actual behavior

A custom operation with the gradient recipe containing the term `[coeff, 1, 0]` breaks `param_shift` (`qnode1` in the code below) unless all operations require such a term (`qnode2` in the code below)

### Additional information

This is due to wrong handling of `unshifted_coeffs` in `expval_param_shift`. Will fix.

### Source code

```shell
dev = qml.device('default.mixed', wires=2)

class OP(qml.operation.Operation):
    num_wires = 1
    ndim_params = (0,)
    num_params = 1
    grad_method = "A"
    #the first term corresponds to the unshifted tape
    grad_recipe = ([[1, 1, 0], [-1, 0, 1]],)
    
    @staticmethod
    def compute_decomposition(x, wires):
        return [qml.RX(x, wires=wires)]

@qml.qnode(dev)
def qnode1(params):
    OP(params[0], wires=0)
    qml.RY(0.1, wires=0)
    qml.RZ(params[1], wires=0)
    return qml.expval(qml.PauliZ(0))

@qml.qnode(dev)
def qnode2(params):
    OP(params[0], wires=0)
    OP(params[1], wires=0)
    return qml.expval(qml.PauliZ(0))

params = np.array([0.2, 0.3])
qnode(params)
qml.gradients.param_shift(qnode)(params)
```


### Tracebacks

```shell
---------------------------------------------------------------------------
IndexError                                Traceback (most recent call last)
Input In [6], in <cell line: 27>()
     25 params = np.array([0.2, 0.3])
     26 qnode(params)
---> 27 qml.gradients.param_shift(qnode)(params)

File ~/venvs/xanadu3.10/lib/python3.10/site-packages/PennyLane-0.25.0.dev0-py3.10.egg/pennylane/gradients/gradient_transform.py:269, in gradient_transform.default_qnode_wrapper.<locals>.jacobian_wrapper(*args, **kwargs)
    262     warnings.warn(
    263         "Attempted to compute the gradient of a QNode with no trainable parameters. "
    264         "If this is unintended, please add trainable parameters in accordance with "
    265         "the chosen auto differentiation framework."
    266     )
    267     return ()
--> 269 qjac = _wrapper(*args, **kwargs)
    271 if not hybrid:
    272     return qjac

File ~/venvs/xanadu3.10/lib/python3.10/site-packages/PennyLane-0.25.0.dev0-py3.10.egg/pennylane/transforms/batch_transform.py:313, in batch_transform.default_qnode_wrapper.<locals>._wrapper(*args, **kwargs)
    300     gradient_fn = None
    302 res = qml.execute(
    303     tapes,
    304     device=qnode.device,
   (...)
    310     **execute_kwargs,
    311 )
--> 313 return processing_fn(res)

File ~/venvs/xanadu3.10/lib/python3.10/site-packages/PennyLane-0.25.0.dev0-py3.10.egg/pennylane/gradients/parameter_shift.py:271, in expval_param_shift.<locals>.processing_fn(results)
    267     g = qml.math.tensordot(res, qml.math.convert_like(gradient_coeffs[i], res), [[0], [0]])
    269     if unshifted_coeffs:
    270         # add on the unshifted term
--> 271         g = g + unshifted_coeffs[i] * r0
    273     grads.append(g)
    275 # The following is for backwards compatibility; currently,
    276 # the device stacks multiple measurement arrays, even if not the same
    277 # size, resulting in a ragged array.
    278 # In the future, we might want to change this so that only tuples
    279 # of arrays are returned.

IndexError: list index out of range
```


### System information

```shell
Name: PennyLane
Version: 0.25.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: /home/david/venvs/xanadu3.10/lib/python3.10/site-packages/PennyLane-0.25.0.dev0-py3.10.egg
Requires: appdirs, autograd, autoray, cachetools, networkx, numpy, pennylane-lightning, retworkx, scipy, semantic-version, toml
Required-by: PennyLane-Lightning

Platform info:           Linux-5.4.0-121-generic-x86_64-with-glibc2.31
Python version:          3.10.5
Numpy version:           1.23.1
Scipy version:           1.9.0rc2
Installed devices:
- lightning.qubit (PennyLane-Lightning-0.24.0)
- default.gaussian (PennyLane-0.25.0.dev0)
- default.mixed (PennyLane-0.25.0.dev0)
- default.qubit (PennyLane-0.25.0.dev0)
- default.qubit.autograd (PennyLane-0.25.0.dev0)
- default.qubit.jax (PennyLane-0.25.0.dev0)
- default.qubit.tf (PennyLane-0.25.0.dev0)
- default.qubit.torch (PennyLane-0.25.0.dev0)
```


### Existing GitHub issues

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

Pull request description:
Restructure `param_shift` and fix bug in unshifted term bookkeeping
This PR restructures `param_shift`. In particular, it introduces additional helper methods, trying to modularize the main for loop of `expval_param_shift`.
In addition, this restructuring fixed #2833, as this bug was caused by flawed bookkeeping of unshifted coefficients (It worked for second order derivatives because usually then all operations require the unshifted tape.

Related issue: #2833 

Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 464f368b6..913f80b4f 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -278,6 +278,11 @@
 
 <h3>Bug fixes</h3>
 
+* Fixes a bug where the parameter-shift gradient breaks when using both
+  custom `grad_recipe`s that contain unshifted terms and recipes that
+  do not contains any unshifted terms.
+  [(#2834)](https://github.com/PennyLaneAI/pennylane/pull/2834)
+
 * Fixes mixed CPU-GPU data-locality issues for Torch interface.
   [(#2830)](https://github.com/PennyLaneAI/pennylane/pull/2830)
 
diff --git a/pennylane/gradients/parameter_shift.py b/pennylane/gradients/parameter_shift.py
index 641eb3e3c..0411fbe70 100644
--- a/pennylane/gradients/parameter_shift.py
+++ b/pennylane/gradients/parameter_shift.py
@@ -73,6 +73,63 @@ def _square_observable(obs):
     return NONINVOLUTORY_OBS[obs.name](obs)
 
 
+def _process_op_recipe(op, p_idx, order):
+    """Process an existing recipe of an operation."""
+    recipe = op.grad_recipe[p_idx]
+    if recipe is None:
+        return None
+
+    recipe = qml.math.array(recipe)
+    if order == 1:
+        return process_shifts(recipe, batch_duplicates=False)
+
+    # Try to obtain the period of the operator frequencies for iteration of custom recipe
+    try:
+        period = frequencies_to_period(op.parameter_frequencies[p_idx])
+    except qml.operation.ParameterFrequenciesUndefinedError:
+        period = None
+
+    # Iterate the custom recipe to obtain the second-order recipe
+    if qml.math.allclose(recipe[:, 1], qml.math.ones_like(recipe[:, 1])):
+        # If the multipliers are ones, we do not include them in the iteration
+        # but keep track of them manually
+        iter_c, iter_s = process_shifts(_iterate_shift_rule(recipe[:, ::2], order, period)).T
+        return qml.math.stack([iter_c, qml.math.ones_like(iter_c), iter_s]).T
+
+    return process_shifts(_iterate_shift_rule(recipe, order, period))
+
+
+def _choose_recipe(argnum, idx, gradient_recipes, shifts, tape):
+    # get the gradient recipe for the trainable parameter
+    arg_idx = argnum.index(idx)
+    recipe = gradient_recipes[arg_idx]
+    if recipe is not None:
+        recipe = process_shifts(np.array(recipe))
+    else:
+        op_shifts = None if shifts is None else shifts[arg_idx]
+        recipe = _get_operation_recipe(tape, idx, shifts=op_shifts)
+    return recipe
+
+
+def _extract_unshifted(recipe, at_least_one_unshifted, f0, gradient_tapes, tape):
+    first_c, first_m, first_s = recipe[0]
+    # Extract zero-shift term if present (if so, it will always be the first)
+    if first_s == 0 and first_m == 1:
+        # Gradient recipe includes a term with zero shift.
+        if not at_least_one_unshifted and f0 is None:
+            # Append the unshifted tape to the gradient tapes, if not already present
+            gradient_tapes.insert(0, tape)
+
+        # Store the unshifted coefficient. It is always the first coefficient due to processing
+        unshifted_coeff = first_c
+        at_least_one_unshifted = True
+        recipe = recipe[1:]
+    else:
+        unshifted_coeff = None
+
+    return recipe, at_least_one_unshifted, unshifted_coeff
+
+
 def _get_operation_recipe(tape, t_idx, shifts, order=1):
     """Utility function to return the parameter-shift rule
     of the operation corresponding to trainable parameter
@@ -107,26 +164,9 @@ def _get_operation_recipe(tape, t_idx, shifts, order=1):
     op, p_idx = tape.get_operation(t_idx)
 
     # Try to use the stored grad_recipe of the operation
-    recipe = op.grad_recipe[p_idx]
-    if recipe is not None:
-        recipe = qml.math.array(recipe)
-        if order == 1:
-            return process_shifts(recipe, batch_duplicates=False)
-
-        # Try to obtain the period of the operator frequencies for iteration of custom recipe
-        try:
-            period = frequencies_to_period(op.parameter_frequencies[p_idx])
-        except qml.operation.ParameterFrequenciesUndefinedError:
-            period = None
-
-        # Iterate the custom recipe to obtain the second-order recipe
-        if qml.math.allclose(recipe[:, 1], qml.math.ones_like(recipe[:, 1])):
-            # If the multipliers are ones, we do not include them in the iteration
-            # but keep track of them manually
-            iter_c, iter_s = process_shifts(_iterate_shift_rule(recipe[:, ::2], order, period)).T
-            return qml.math.stack([iter_c, qml.math.ones_like(iter_c), iter_s]).T
-
-        return process_shifts(_iterate_shift_rule(recipe, order, period))
+    op_recipe = _process_op_recipe(op, p_idx, order)
+    if op_recipe is not None:
+        return op_recipe
 
     # Try to obtain frequencies, either via custom implementation or from generator eigvals
     try:
@@ -174,19 +214,18 @@ def expval_param_shift(tape, argnum=None, shifts=None, gradient_recipes=None, f0
     argnum = argnum or tape.trainable_params
 
     gradient_tapes = []
-    gradient_coeffs = []
-    shapes = []
-    unshifted_coeffs = []
-
-    fns = []
+    # Each entry for gradient_data will be a tuple with content
+    # (num_tapes, coeffs, fn, unshifted_coeff)
+    gradient_data = []
+    # Keep track of whether there is at least one unshifted
+    # term in the parameter-shift rules of all parameters
+    at_least_one_unshifted = False
 
     for idx, _ in enumerate(tape.trainable_params):
 
         if idx not in argnum:
             # parameter has zero gradient
-            shapes.append(0)
-            gradient_coeffs.append([])
-            fns.append(None)
+            gradient_data.append((0, [], None, None))
             continue
 
         op, _ = tape.get_operation(idx)
@@ -200,43 +239,21 @@ def expval_param_shift(tape, argnum=None, shifts=None, gradient_recipes=None, f0
                 )
 
             g_tapes, h_fn = qml.gradients.hamiltonian_grad(tape, idx)
+            # hamiltonian_grad always returns a list with a single tape
             gradient_tapes.extend(g_tapes)
-            shapes.append(1)
-            gradient_coeffs.append(np.array([1.0]))
-            fns.append(h_fn)
+            gradient_data.append((1, np.array([1.0]), h_fn, None))
             continue
 
-        # get the gradient recipe for the trainable parameter
-        arg_idx = argnum.index(idx)
-        recipe = gradient_recipes[arg_idx]
-        if recipe is not None:
-            recipe = process_shifts(np.array(recipe))
-        else:
-            op_shifts = None if shifts is None else shifts[arg_idx]
-            recipe = _get_operation_recipe(tape, idx, shifts=op_shifts)
+        recipe = _choose_recipe(argnum, idx, gradient_recipes, shifts, tape)
+        recipe, at_least_one_unshifted, unshifted_coeff = _extract_unshifted(
+            recipe, at_least_one_unshifted, f0, gradient_tapes, tape
+        )
         coeffs, multipliers, op_shifts = recipe.T
-        fns.append(None)
-
-        # Extract zero-shift term if present (if so, it will always be the first)
-        if op_shifts[0] == 0 and multipliers[0] == 1:
-            # Gradient recipe includes a term with zero shift.
-
-            if not unshifted_coeffs and f0 is None:
-                # Ensure that the unshifted tape is appended
-                # to the gradient tapes, if not already.
-                gradient_tapes.append(tape)
-
-            # Store the unshifted coefficient. We know that
-            # it will always be the first coefficient due to processing.
-            unshifted_coeffs.append(coeffs[0])
-            coeffs, multipliers, op_shifts = recipe[1:].T
 
         # generate the gradient tapes
-        gradient_coeffs.append(coeffs)
         g_tapes = generate_shifted_tapes(tape, idx, op_shifts, multipliers)
-
         gradient_tapes.extend(g_tapes)
-        shapes.append(len(g_tapes))
+        gradient_data.append((len(g_tapes), coeffs, None, unshifted_coeff))
 
     def processing_fn(results):
         # Apply the same squeezing as in qml.QNode to make the transform output consistent.
@@ -245,42 +262,39 @@ def expval_param_shift(tape, argnum=None, shifts=None, gradient_recipes=None, f0
             results = [qml.math.squeeze(res) for res in results]
 
         grads = []
-        start = 1 if unshifted_coeffs and f0 is None else 0
+        start = 1 if at_least_one_unshifted and f0 is None else 0
         r0 = f0 or results[0]
 
-        for i, (s, f) in enumerate(zip(shapes, fns)):
+        for num_tapes, coeffs, fn, unshifted_coeff in gradient_data:
 
-            if s == 0:
+            if num_tapes == 0:
                 # parameter has zero gradient
-                g = qml.math.zeros_like(results[0])
-                grads.append(g)
+                grads.append(qml.math.zeros_like(results[0]))
                 continue
 
-            res = results[start : start + s]
-            start = start + s
+            res = results[start : start + num_tapes]
+            start = start + num_tapes
 
-            if f is not None:
-                res = f(res)
+            # individual post-processing of e.g. Hamiltonian grad tapes
+            if fn is not None:
+                res = fn(res)
 
             # compute the linear combination of results and coefficients
             res = qml.math.stack(res)
-            g = qml.math.tensordot(res, qml.math.convert_like(gradient_coeffs[i], res), [[0], [0]])
+            g = qml.math.tensordot(res, qml.math.convert_like(coeffs, res), [[0], [0]])
 
-            if unshifted_coeffs:
-                # add on the unshifted term
-                g = g + unshifted_coeffs[i] * r0
+            if unshifted_coeff is not None:
+                # add the unshifted term
+                g = g + unshifted_coeff * r0
 
             grads.append(g)
 
-        # The following is for backwards compatibility; currently,
-        # the device stacks multiple measurement arrays, even if not the same
-        # size, resulting in a ragged array.
-        # In the future, we might want to change this so that only tuples
-        # of arrays are returned.
+        # The following is for backwards compatibility; currently, the device stacks multiple
+        # measurement arrays, even if not the same size, resulting in a ragged array.
+        # In the future, we might want to change this so that only tuples of arrays are returned.
         for i, g in enumerate(grads):
-            if hasattr(g, "dtype") and g.dtype is np.dtype("object"):
-                if qml.math.ndim(g) > 0:
-                    grads[i] = qml.math.hstack(g)
+            if getattr(g, "dtype", None) is np.dtype("object") and qml.math.ndim(g) > 0:
+                grads[i] = qml.math.hstack(g)
 
         return qml.math.T(qml.math.stack(grads))
 


Test patch:
diff --git a/tests/gradients/test_parameter_shift.py b/tests/gradients/test_parameter_shift.py
index 30a25d03d..d69d6ad7a 100644
--- a/tests/gradients/test_parameter_shift.py
+++ b/tests/gradients/test_parameter_shift.py
@@ -287,22 +287,33 @@ class TestParamShift:
         assert tapes[3].get_parameters(trainable_only=False) == [1.0, 2.0, 2 * 3.0 + 2, 4.0]
         assert tapes[4].get_parameters(trainable_only=False) == [1.0, 2.0, 3 * 3.0 + 3, 4.0]
 
-    def test_recycled_unshifted_tape(self):
+    @pytest.mark.parametrize("ops_with_custom_recipe", [[0], [1], [0, 1]])
+    def test_recycled_unshifted_tape(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(0.543, wires=[0])
-            qml.RY(-0.654, wires=[0])
+            qml.RX(x[0], wires=[0])
+            qml.RX(x[1], wires=[0])
             qml.expval(qml.PauliZ(0))
 
-        gradient_recipes = ([[-1e7, 1, 0], [1e7, 1, 1e7]],) * 2
+        gradient_recipes = tuple(
+            [[-1e7, 1, 0], [1e7, 1, 1e-7]] 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)
 
-        # one tape per parameter, plus one global call
-        assert len(tapes) == tape.num_params + 1
+        # 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)
 
     def test_f0_provided(self):
         """Test that if the original tape output is provided, then
@@ -324,6 +335,44 @@ class TestParamShift:
 
         fn(dev.batch_execute(tapes))
 
+    def test_op_with_custom_unshifted_term(self):
+        """Test that an operation with a gradient recipe that depends on
+        its instantiated parameter values works correctly within the parameter
+        shift rule. Also tests that grad_recipes supersedes paramter_frequencies.
+        """
+        s = np.pi / 2
+
+        class RX(qml.RX):
+            """RX operation with an additional term in the grad recipe.
+            The grad_recipe no longer yields the derivative, but we account for this.
+            For this test, the presence of the unshifted term (with non-vanishing coefficient)
+            is essential."""
+
+            grad_recipe = ([[0.5, 1, s], [-0.5, 1, -s], [0.2, 1, 0]],)
+
+        x = np.array([-0.361, 0.654], requires_grad=True)
+        dev = qml.device("default.qubit", wires=2)
+
+        with qml.tape.QuantumTape() as tape:
+            qml.RX(x[0], wires=0)
+            RX(x[1], wires=0)
+            qml.expval(qml.PauliZ(0))
+
+        tapes, fn = qml.gradients.param_shift(tape)
+
+        # Unshifted tapes always are first within the tapes created for one operation;
+        # They are not batched together because we trust operation recipes to be condensed already
+        expected_shifts = [[0, 0], [s, 0], [-s, 0], [0, s], [0, -s]]
+        assert len(tapes) == 5
+        for tape, expected in zip(tapes, expected_shifts):
+            assert tape.operations[0].data[0] == x[0] + expected[0]
+            assert tape.operations[1].data[0] == x[1] + expected[1]
+
+        grad = fn(dev.batch_execute(tapes))
+        assert np.allclose(
+            grad, [-np.sin(x[0] + x[1]), -np.sin(x[0] + x[1]) + 0.2 * np.cos(x[0] + x[1])]
+        )
+
     def test_independent_parameters_analytic(self):
         """Test the case where expectation values are independent of some parameters. For those
         parameters, the gradient should be evaluated to zero without executing the device."""
@@ -356,14 +405,9 @@ class TestParamShift:
     def test_grad_recipe_parameter_dependent(self, monkeypatch):
         """Test that an operation with a gradient recipe that depends on
         its instantiated parameter values works correctly within the parameter
-        shift rule. Also tests that grad_recipes supersedes paramter_frequencies.
+        shift rule. Also tests that `grad_recipe` supersedes `parameter_frequencies`.
         """
 
-        def fail(*args, **kwargs):
-            raise qml.operation.ParameterFrequenciesUndefinedError
-
-        monkeypatch.setattr(qml.RX, "parameter_frequencies", fail)
-
         class RX(qml.RX):
             @property
             def grad_recipe(self):