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.

Comments:

Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 33b1663d1..3b7fdf41a 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -9,7 +9,7 @@
 * The `qml.math` module now also contains a submodule for
   fast Fourier transforms, `qml.math.fft`.
   [(#1440)](https://github.com/PennyLaneAI/pennylane/pull/1440)
-  
+
   The submodule in particular provides differentiable
   versions of the following functions, available in all common
   interfaces for PennyLane
@@ -42,11 +42,11 @@
   >>> qml.ops.qubit.special_unitary.pauli_basis_strings(1) # 4**1-1 = 3 Pauli words
   ['X', 'Y', 'Z']
   >>> qml.ops.qubit.special_unitary.pauli_basis_strings(2) # 4**2-1 = 15 Pauli words
-  ['IX', 'IY', 'IZ', 'XI', 'XX', 'XY', 'XZ', 'YI', 'YX', 'YY', 'YZ', 'ZI', 'ZX', 'ZY', 'ZZ'] 
+  ['IX', 'IY', 'IZ', 'XI', 'XX', 'XY', 'XZ', 'YI', 'YX', 'YY', 'YZ', 'ZI', 'ZX', 'ZY', 'ZZ']
   ```
-  
+
   For example, on a single qubit, we may define
-  
+
   ```pycon
   >>> theta = np.array([0.2, 0.1, -0.5])
   >>> U = qml.SpecialUnitary(theta, 0)
@@ -54,7 +54,7 @@
   array([[ 0.8537127 -0.47537233j,  0.09507447+0.19014893j],
          [-0.09507447+0.19014893j,  0.8537127 +0.47537233j]])
   ```
-  
+
   A single non-zero entry in the parameters will create a Pauli rotation:
 
   ```pycon
@@ -65,7 +65,7 @@
   >>> qml.math.allclose(su.matrix(), rx.matrix())
   True
   ```
-  
+
   This operation can be differentiated with hardware-compatible methods like parameter shifts
   and it supports parameter broadcasting/batching, but not both at the same time.
 
@@ -80,7 +80,7 @@
   A `ParametrizedHamiltonian` holds information representing a linear combination of operators
   with parametrized coefficents. The `ParametrizedHamiltonian` can be passed parameters to create the operator for
   the specified parameters.
-  
+
   ```pycon
   f1 = lambda p, t: p * np.sin(t) * (t - 1)
   f2 = lambda p, t: p[0] * np.cos(p[1]* t ** 2)
@@ -170,7 +170,7 @@
   ...     qml.RX(x, 0)
   ...     qml.RX(x, 1)
   ...     return qml.expval(qml.PauliZ(0))
-  >>> jax.jacobian(circuit)(jax.numpy.array(0.5)) 
+  >>> jax.jacobian(circuit)(jax.numpy.array(0.5))
   DeviceArray(-0.4792258, dtype=float32, weak_type=True)
   ```
 
@@ -189,20 +189,20 @@
   import pennylane as qml
   import jax
   from jax import numpy as jnp
-  
+
   jax.config.update("jax_enable_x64", True)
-  
+
   qml.enable_return()
-  
+
   dev = qml.device("lightning.qubit", wires=2)
-  
+
   @jax.jit
   @qml.qnode(dev, interface="jax-jit", diff_method="parameter-shift", max_diff=2)
   def circuit(a, b):
       qml.RY(a, wires=0)
       qml.RX(b, wires=1)
       return qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliZ(1))
-  
+
   a, b = jnp.array(1.0), jnp.array(2.0)
   ```
 
@@ -231,7 +231,7 @@
   import pennylane as qml
   from pennylane import numpy as np
   import jax
-  
+
   symbols = ["H", "H"]
   geometry = np.array([[0.0, 0.0, 0.0], [0.0, 0.0, 1.0]])
 
@@ -352,16 +352,16 @@
 
   ```pycon
   >>> f(params=[1.2, 2.3, 3.4, 4.5], t=3.9)
-  DeviceArray(4.5, dtype=float32) 
-  >>> f(params=[1.2, 2.3, 3.4, 4.5], t=6)  # zero outside the range (2, 4) 
+  DeviceArray(4.5, dtype=float32)
+  >>> f(params=[1.2, 2.3, 3.4, 4.5], t=6)  # zero outside the range (2, 4)
   DeviceArray(0., dtype=float32)
   ```
-  
+
 * Added `pwc_from_function` as a decorator for defining a `ParametrizedHamiltonian`.
   This function can be used to decorate a function and create a piecewise constant
   approximation of it.
   [(#3645)](https://github.com/PennyLaneAI/pennylane/pull/3645)
-  
+
   ```pycon
   >>> @pwc_from_function(t=(2, 4), num_bins=10)
   ... def f1(p, t):
@@ -370,7 +370,7 @@
 
   The resulting function approximates the same of `p**2 * t` on the interval `t=(2, 4)`
   in 10 bins, and returns zero outside the interval.
-  
+
   ```pycon
   # t=2 and t=2.1 are within the same bin
   >>> f1(3, 2), f1(3, 2.1)
@@ -382,7 +382,7 @@
   >>> f1(3, 5)
   DeviceArray(0., dtype=float32)
   ```
-  
+
 *Next generation device API:*
 
 * The `apply_operation` single-dispatch function is added to `devices/qubit` that applies an operation
@@ -398,6 +398,10 @@
 
 <h3>Improvements</h3>
 
+* The parameter-shift derivative of variances saves a redundant evaluation of the
+  corresponding unshifted expectation value tape, if possible
+  [(#3744)](https://github.com/PennyLaneAI/pennylane/pull/3744)
+
 * `qml.purity` is added as a measurement process for purity
   [(#3551)](https://github.com/PennyLaneAI/pennylane/pull/3551)
 
@@ -501,20 +505,20 @@
 * `qml.VQECost` is removed.
   [(#3735)](https://github.com/PennyLaneAI/pennylane/pull/3735)
 
-* The default interface is now `auto`. There is no need to specify the interface anymore! It is automatically 
+* The default interface is now `auto`. There is no need to specify the interface anymore! It is automatically
   determined by checking your `QNode` parameters.
   [(#3677)](https://github.com/PennyLaneAI/pennylane/pull/3677)
-  
+
   ```python
   import jax
   import jax.numpy as jnp
-  
+
   qml.enable_return()
   a = jnp.array(0.1)
   b = jnp.array(0.2)
-  
+
   dev = qml.device("default.qubit", wires=2)
-  
+
   @qml.qnode(dev)
   def circuit(a, b):
       qml.RY(a, wires=0)
@@ -522,18 +526,18 @@
       qml.CNOT(wires=[0, 1])
       return qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliY(1))
   ```
-  
+
   ```pycon
   >>> circuit(a, b)
   (Array(0.9950042, dtype=float32), Array(-0.19767681, dtype=float32))
   >>> jac = jax.jacobian(circuit)(a, b)
   (Array(-0.09983341, dtype=float32, weak_type=True), Array(0.01983384, dtype=float32, weak_type=True))
   ```
-  
-  It comes with the fact that the interface is determined during the `QNode` call instead of the 
-  initialization. It means that the `gradient_fn` and `gradient_kwargs` are only defined on the QNode at the beginning 
-  of the call. As well, without specifying the interface it is not possible to guarantee that the device will not be changed 
-  during the call if you are using backprop(`default.qubit` to `default.qubit,jax`e.g.) whereas before it was happening at 
+
+  It comes with the fact that the interface is determined during the `QNode` call instead of the
+  initialization. It means that the `gradient_fn` and `gradient_kwargs` are only defined on the QNode at the beginning
+  of the call. As well, without specifying the interface it is not possible to guarantee that the device will not be changed
+  during the call if you are using backprop(`default.qubit` to `default.qubit,jax`e.g.) whereas before it was happening at
   initialization, therefore you should not try to track the device without specifying the interface.
 
 * The tape method `get_operation` can also now return the operation index in the tape, and it can be
@@ -543,13 +547,13 @@
 
 * `Operation.inv()` and the `Operation.inverse` setter have been removed. Please use `qml.adjoint` or `qml.pow` instead.
   [(#3618)](https://github.com/PennyLaneAI/pennylane/pull/3618)
-  
+
   For example, instead of
-  
+
   ```pycon
   >>> qml.PauliX(0).inv()
   ```
-  
+
   use
 
   ```pycon
@@ -599,7 +603,7 @@
 
 * Updated the code example in `qml.SparseHamiltonian` with the correct wire range.
   [(#3643)](https://github.com/PennyLaneAI/pennylane/pull/3643)
-  
+
 * A hyperlink has been added in the text for a URL in the `qml.qchem.mol_data` docstring.
   [(#3644)](https://github.com/PennyLaneAI/pennylane/pull/3644)
 
diff --git a/pennylane/gradients/parameter_shift.py b/pennylane/gradients/parameter_shift.py
index f8d35909a..0f34169ca 100644
--- a/pennylane/gradients/parameter_shift.py
+++ b/pennylane/gradients/parameter_shift.py
@@ -587,6 +587,8 @@ def _expval_param_shift_tuple(
 
         return tuple(grads)
 
+    processing_fn.first_result_unshifted = at_least_one_unshifted
+
     return gradient_tapes, processing_fn
 
 
@@ -738,6 +740,8 @@ def expval_param_shift(
 
         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 _create_variance_proc_fn(
         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 var_param_shift(
         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
index 15e30190e..dcf42475d 100644
--- a/tests/gradients/test_parameter_shift.py
+++ b/tests/gradients/test_parameter_shift.py
@@ -1298,6 +1298,37 @@ class TestParameterShiftRule:
         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 @@ class TestParameterShiftRule:
             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
index 73d279da3..7832fe972 100644
--- a/tests/returntypes/paramshift/test_parameter_shift_new.py
+++ b/tests/returntypes/paramshift/test_parameter_shift_new.py
@@ -1993,6 +1993,37 @@ class TestParameterShiftRule:
                 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)