Repository: PennyLaneAI/pennylane (python)

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.

Pull request description:
Reuse unshifted evaluation when computing parameter shift derivatives of variances
**Context:**
When computing the derivative of a variance using a gradient recipe with unshifted terms, the unshifted tape is produced and evaluated twice.
See #3178 

**Description of the Change:**
The post-processing function of the expectation value parameter-shift is given an attribute that indicates whether it expects the first result in the results list passed to it to be from the unshifted tape. If this is the case, the variance parameter-shift logic does not add the unshifted expectation value tape $\langle H\rangle$ that is required in the chain rule
$\frac{\partial}{\partial \theta} Var[H] = \frac{\partial}{\partial \theta} \langle H^2\rangle - 2 \langle H\rangle \frac{\partial}{\partial \theta} \langle H\rangle$

**Benefits:**
Saves quantum resources.

**Possible Drawbacks:**
N/A

**Related GitHub Issues:**
closes #3178

This patch is for both return type systems.

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)