Repository: PennyLaneAI/pennylane

Issue:
Parameter-shift rule of variance does not recycle unshifted evaluation
### Feature details

Save repeated evaluations at unshifted parameters when using `param_shift` together with gradient recipes that involve unshifted evaluations. Consider the setup
```python
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.var(qml.PauliZ(0))

gradient_recipes = tuple([[[-1e-7, 1, 0], [1e-7, 1, 0], [-1e5, 1, -5e-6], [1e5, 1, 5e-6]], None])
tapes, fn = qml.gradients.param_shift(tape, gradient_recipes=gradient_recipes)
```
Then we get:
```pycon
>>> for tape in tapes:
>>>     print(tape.get_parameters())
[0.543, -0.654]
[0.543, -0.654]
[0.542995, -0.654]
[0.5430050000000001, -0.654]
[0.543, 0.9167963267948965]
[0.543, -2.2247963267948965]
```
i.e. the unshifted tape is included twice, once for the unshifted evaluation in the product rule of the variance, and once for the shift rule that uses an unshifted term.


### Implementation

This could be saved by passing the information that the unshifted evaluation already is available between `var_param_shift` and the calls to `expval_param_shift`.

### How important would you say this feature is?

1: Not important. Would be nice to have.

### Additional information

This issue occurs with and without the new return type system.


Discussion:

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
@@ -587,6 +587,8 @@ def processing_fn(results):
 
         return tuple(grads)
 
+    processing_fn.first_result_unshifted = at_least_one_unshifted
+
     return gradient_tapes, processing_fn
 
 
@@ -738,6 +740,8 @@ def processing_fn(results):
 
         return qml.math.T(qml.math.stack(grads))
 
+    processing_fn.first_result_unshifted = at_least_one_unshifted
+
     return gradient_tapes, processing_fn
 
 
@@ -893,7 +897,7 @@ def processing_fn(results):
         shot_vector = isinstance(shots, Sequence)
 
         # analytic derivative of <A>
-        pdA = pdA_fn(results[1:tape_boundary])
+        pdA = pdA_fn(results[int(not pdA_fn.first_result_unshifted) : tape_boundary])
 
         # analytic derivative of <A^2>
         pdA2 = _get_pdA2(
@@ -969,8 +973,6 @@ def _var_param_shift_tuple(
     # Get <A>, the expectation value of the tape with unshifted parameters.
     expval_tape = tape.copy(copy_operations=True)
 
-    gradient_tapes = [expval_tape]
-
     # Convert all variance measurements on the tape into expectation values
     for i in var_indices:
         obs = expval_tape._measurements[i].obs
@@ -980,11 +982,12 @@ def _var_param_shift_tuple(
     pdA_tapes, pdA_fn = expval_param_shift(
         expval_tape, argnum, shifts, gradient_recipes, f0, broadcast, shots
     )
+    gradient_tapes = [] if pdA_fn.first_result_unshifted else [expval_tape]
     gradient_tapes.extend(pdA_tapes)
 
     # Store the number of first derivative tapes, so that we know
     # the number of results to post-process later.
-    tape_boundary = len(pdA_tapes) + 1
+    tape_boundary = len(gradient_tapes)
 
     # If there are non-involutory observables A present, we must compute d<A^2>/dp.
     # Get the indices in the measurement queue of all non-involutory
@@ -1020,9 +1023,6 @@ def _var_param_shift_tuple(
         )
         gradient_tapes.extend(pdA2_tapes)
 
-    # Store the number of first derivative tapes, so that we know
-    # the number of results to post-process later.
-    tape_boundary = len(pdA_tapes) + 1
     processing_fn = _create_variance_proc_fn(
         tape, var_mask, var_indices, pdA_fn, pdA2_fn, tape_boundary, non_involutory_indices, shots
     )
@@ -1076,8 +1076,6 @@ def var_param_shift(
     # Get <A>, the expectation value of the tape with unshifted parameters.
     expval_tape = tape.copy(copy_operations=True)
 
-    gradient_tapes = [expval_tape]
-
     # Convert all variance measurements on the tape into expectation values
     for i in var_idx:
         obs = expval_tape._measurements[i].obs
@@ -1087,11 +1085,12 @@ def var_param_shift(
     pdA_tapes, pdA_fn = expval_param_shift(
         expval_tape, argnum, shifts, gradient_recipes, f0, broadcast
     )
+    gradient_tapes = [] if pdA_fn.first_result_unshifted else [expval_tape]
     gradient_tapes.extend(pdA_tapes)
 
     # Store the number of first derivative tapes, so that we know
     # the number of results to post-process later.
-    tape_boundary = len(pdA_tapes) + 1
+    tape_boundary = len(gradient_tapes)
 
     # If there are non-involutory observables A present, we must compute d<A^2>/dp.
     # Get the indices in the measurement queue of all non-involutory
@@ -1153,7 +1152,7 @@ def processing_fn(results):
         f0 = qml.math.expand_dims(res, -1)
         mask = qml.math.convert_like(qml.math.reshape(mask, qml.math.shape(f0)), res)
 
-        pdA = pdA_fn(results[1:tape_boundary])
+        pdA = pdA_fn(results[int(not pdA_fn.first_result_unshifted) : tape_boundary])
         pdA2 = 0
 
         if non_involutory:


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
@@ -1298,6 +1298,37 @@ def test_expval_and_variance(self, tol):
         assert gradA == pytest.approx(expected, abs=tol)
         assert gradF == pytest.approx(expected, abs=tol)
 
+    def test_recycling_unshifted_tape_result(self):
+        """Test that an unshifted term in the used gradient recipe is reused
+        for the chain rule computation within the variance parameter shift rule."""
+        dev = qml.device("default.qubit", wires=2)
+        gradient_recipes = ([[-1e-5, 1, 0], [1e-5, 1, 0], [-1e5, 1, -5e-6], [1e5, 1, 5e-6]], None)
+        x = [0.543, -0.654]
+
+        with qml.queuing.AnnotatedQueue() as q:
+            qml.RX(x[0], wires=[0])
+            qml.RX(x[1], wires=[0])
+            qml.var(qml.PauliZ(0))
+
+        tape = qml.tape.QuantumScript.from_queue(q)
+        tapes, fn = qml.gradients.param_shift(tape, gradient_recipes=gradient_recipes)
+        # 2 operations x 2 shifted positions + 1 unshifted term overall
+        assert len(tapes) == 2 * 2 + 1
+
+        with qml.queuing.AnnotatedQueue() as q:
+            qml.RX(x[0], wires=[0])
+            qml.RX(x[1], wires=[0])
+            qml.var(qml.Projector([1], wires=0))
+
+        tape = qml.tape.QuantumScript.from_queue(q)
+        tape.trainable_params = [0, 1]
+        tapes, fn = qml.gradients.param_shift(tape, gradient_recipes=gradient_recipes)
+        for tape in tapes:
+            print(tape.measurements)
+        # 2 operations x 2 shifted positions + 1 unshifted term overall    <-- <H>
+        # + 2 operations x 2 shifted positions + 1 unshifted term          <-- <H^2>
+        assert len(tapes) == (2 * 2 + 1) + (2 * 2 + 1)
+
     def test_projector_variance(self, tol):
         """Test that the variance of a projector is correctly returned"""
         dev = qml.device("default.qubit", wires=2)
@@ -1311,7 +1342,7 @@ def test_projector_variance(self, tol):
             qml.var(qml.Projector(P, wires=0) @ qml.PauliX(1))
 
         tape = qml.tape.QuantumScript.from_queue(q)
-        tape.trainable_params = {0, 1}
+        tape.trainable_params = [0, 1]
 
         res = dev.execute(tape)
         expected = 0.25 * np.sin(x / 2) ** 2 * (3 + np.cos(2 * y) + 2 * np.cos(x) * np.sin(y) ** 2)
diff --git a/tests/returntypes/paramshift/test_parameter_shift_new.py b/tests/returntypes/paramshift/test_parameter_shift_new.py
--- a/tests/returntypes/paramshift/test_parameter_shift_new.py
+++ b/tests/returntypes/paramshift/test_parameter_shift_new.py
@@ -1993,6 +1993,37 @@ def test_expval_and_variance_multi_param(self, tol):
                 assert np.allclose(a_comp, e_comp, atol=tol, rtol=0)
         assert gradF == pytest.approx(expected, abs=tol)
 
+    def test_recycling_unshifted_tape_result(self):
+        """Test that an unshifted term in the used gradient recipe is reused
+        for the chain rule computation within the variance parameter shift rule."""
+        dev = qml.device("default.qubit", wires=2)
+        gradient_recipes = ([[-1e-5, 1, 0], [1e-5, 1, 0], [-1e5, 1, -5e-6], [1e5, 1, 5e-6]], None)
+        x = [0.543, -0.654]
+
+        with qml.queuing.AnnotatedQueue() as q:
+            qml.RX(x[0], wires=[0])
+            qml.RX(x[1], wires=[0])
+            qml.var(qml.PauliZ(0))
+
+        tape = qml.tape.QuantumScript.from_queue(q)
+        tapes, fn = qml.gradients.param_shift(tape, gradient_recipes=gradient_recipes)
+        # 2 operations x 2 shifted positions + 1 unshifted term overall
+        assert len(tapes) == 2 * 2 + 1
+
+        with qml.queuing.AnnotatedQueue() as q:
+            qml.RX(x[0], wires=[0])
+            qml.RX(x[1], wires=[0])
+            qml.var(qml.Projector([1], wires=0))
+
+        tape = qml.tape.QuantumScript.from_queue(q)
+        tape.trainable_params = [0, 1]
+        tapes, fn = qml.gradients.param_shift(tape, gradient_recipes=gradient_recipes)
+        for tape in tapes:
+            print(tape.measurements)
+        # 2 operations x 2 shifted positions + 1 unshifted term overall    <-- <H>
+        # + 2 operations x 2 shifted positions + 1 unshifted term          <-- <H^2>
+        assert len(tapes) == (2 * 2 + 1) + (2 * 2 + 1)
+
     def test_projector_variance(self, tol):
         """Test that the variance of a projector is correctly returned"""
         dev = qml.device("default.qubit", wires=2)