Repository: PennyLaneAI/pennylane
Pull request #1341 (2021-05-19 18:00:54)

Adjoint differentation improvements
The [current adjoint differentiation](https://github.com/PennyLaneAI/pennylane/blob/98a787643f4928b6078fede6d4ecb7adab3e1800/pennylane/_qubit_device.py#L785) performs a forward pass of the circuit, even though PennyLane will often have performed one just before calling the adjoint jacobian function.  

This PR intends to make it easier to reduce the number of circuit executions when using adjoint diff.  Instead of using a "one-size-fits-all" caching method, this PR introduces several different ways users could reduce forward pass executions:

1) compute expectation values and the jacobian at the same time inside of `adjoint_jacobian` with the `return_obs` keyword

2) start with the current device state if requested by the `use_device_state` keyword

3) start with the provided state if it is provided to the `starting_state` keyword

A note that `starting_state` takes precedence over `use_device_state`.

Reducing the forward pass can reduce execution time of the device method call itself by a full 25%. 

I am also working on editing the interfaces to take advantage of this speed up.  



codecov[bot]: # [Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1341?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
> Merging [#1341](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1341?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (679160d) into [master](https://codecov.io/gh/PennyLaneAI/pennylane/commit/98a787643f4928b6078fede6d4ecb7adab3e1800?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (98a7876) will **decrease** coverage by `0.94%`.
> The diff coverage is `100.00%`.

[![Impacted file tree graph](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1341/graphs/tree.svg?width=650&height=150&src=pr&token=9S3YzlUgGk&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI)](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1341?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI)

```diff
@@            Coverage Diff             @@
##           master    #1341      +/-   ##
==========================================
- Coverage   99.10%   98.16%   -0.95%     
==========================================
  Files         141      145       +4     
  Lines       10575    11089     +514     
==========================================
+ Hits        10480    10885     +405     
- Misses         95      204     +109     
```


| [Impacted Files](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1341?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | Coverage Δ | |
|---|---|---|
| [pennylane/\_qubit\_device.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1341/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL19xdWJpdF9kZXZpY2UucHk=) | `98.55% <100.00%> (-1.45%)` | :arrow_down: |
| [pennylane/interfaces/autograd.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1341/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2ludGVyZmFjZXMvYXV0b2dyYWQucHk=) | `100.00% <100.00%> (ø)` | |
| [pennylane/tape/jacobian\_tape.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1341/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3RhcGUvamFjb2JpYW5fdGFwZS5weQ==) | `97.90% <100.00%> (-1.04%)` | :arrow_down: |
| [pennylane/tape/operation\_recorder.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1341/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3RhcGUvb3BlcmF0aW9uX3JlY29yZGVyLnB5) | `64.51% <0.00%> (-35.49%)` | :arrow_down: |
| [pennylane/measure.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1341/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL21lYXN1cmUucHk=) | `91.76% <0.00%> (-8.24%)` | :arrow_down: |
| [pennylane/circuit\_drawer/grid.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1341/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2NpcmN1aXRfZHJhd2VyL2dyaWQucHk=) | `88.70% <0.00%> (-7.72%)` | :arrow_down: |
| [pennylane/grouping/optimize\_measurements.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1341/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2dyb3VwaW5nL29wdGltaXplX21lYXN1cmVtZW50cy5weQ==) | `92.30% <0.00%> (-7.70%)` | :arrow_down: |
| [pennylane/vqe/vqe.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1341/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3ZxZS92cWUucHk=) | `92.97% <0.00%> (-7.03%)` | :arrow_down: |
| [pennylane/configuration.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1341/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2NvbmZpZ3VyYXRpb24ucHk=) | `93.10% <0.00%> (-6.90%)` | :arrow_down: |
| [pennylane/qnn/torch.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1341/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3Fubi90b3JjaC5weQ==) | `95.08% <0.00%> (-4.92%)` | :arrow_down: |
| ... and [26 more](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1341/diff?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | |

------

[Continue to review full report at Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1341?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1341?src=pr&el=footer&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI). Last update [98a7876...679160d](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1341?src=pr&el=lastupdated&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).

josh146: 🎉 
albi3ro: As of now, this branch is seeing decent time improvements.  I changed the benchmark repo's `GradientComputation_light` benchmark to use the adjoint method.  

The results for current master:
```
[  0.00%] · For pennylane commit 411c3d52 <master>:
[  0.00%] ·· Benchmarking virtualenv-py3.8-Qulacs-jax-jaxlib-networkx-qsimcirq-tensorflow-torch
[  0.05%] ··· Running (asv.core_suite.GradientComputation_light.time_gradient--).
[  0.09%] ··· asv.core_suite.GradientComputation_light.time_gradient                                                                                       4/16 failed
[  0.09%] ··· ========= ============== ============ ============ ========= ============== ============ ============= =========
              --                                                 n_layers / interface
              --------- ------------------------------------------------------------------------------------------------------
               n_wires   3 / autograd     3 / tf     3 / torch    3 / jax   6 / autograd     6 / tf      6 / torch    6 / jax
              ========= ============== ============ ============ ========= ============== ============ ============= =========
                  2       5.11±0.1ms     10.4±2ms    4.77±0.1ms    failed    7.45±0.3ms    14.2±0.2ms   7.14±0.05ms    failed
                  5       10.2±0.2ms    18.6±0.5ms   9.76±0.3ms    failed    17.2±0.3ms    31.7±0.3ms    16.6±0.6ms    failed
              ========= ============== ============ ============ ========= ============== ============ ============= =========
```

And for the most recent commit to this PR:
```
    [  0.00%] · For pennylane commit 7a26a6ce <adjoint_diff_improvements>:
[  0.00%] ·· Benchmarking virtualenv-py3.8-Qulacs-jax-jaxlib-networkx-qsimcirq-tensorflow-torch
[  0.05%] ··· Running (asv.core_suite.GradientComputation_light.time_gradient--).
[  0.09%] ··· asv.core_suite.GradientComputation_light.time_gradient                                                                                       3/16 failed
[  0.09%] ··· ========= ============== ============ ============ ========= ============== ============ ============ =========
              --                                                 n_layers / interface
              --------- -----------------------------------------------------------------------------------------------------
               n_wires   3 / autograd     3 / tf     3 / torch    3 / jax   6 / autograd     6 / tf     6 / torch    6 / jax
              ========= ============== ============ ============ ========= ============== ============ ============ =========
                  2      4.35±0.05ms    9.00±0.3ms    4.98±1ms    235±2ms    6.76±0.3ms     14.5±3ms    6.74±0.1ms    failed
                  5      8.63±0.03ms    17.0±0.6ms   9.82±0.2ms    failed    15.2±0.5ms    29.7±0.7ms   16.1±0.4ms    failed
              ========= ============== ============ ============ ========= ============== ============ ============ =========
```
albi3ro: @dwierichs want to help review?
albi3ro: @trbromley  For the autograd and jax derivatives, we know that we just performed a forward pass, so the device is in the correct state.

For the tensorflow and torch interfaces, we can perform the forward pass, the go do other stuff, maybe perform a forward pass with different parameters, and *then* perform the backward pass where we perform `dev.adjoint_jacobian(qtape)`.  So we cannot use the current device state, since other things may have occurred on the device since the forward pass.

Caching the device state means keeping around a copy of the statevector, which may come with some memory overhead.  At some point, we may switch to having this as default, but I don't want to do so just yet given the memory considerations.
albi3ro: ```
[  0.00%] · For pennylane commit d8b6011d <adjoint_diff_improvements>:
[  0.00%] ·· Benchmarking virtualenv-py3.8-Qulacs-jax-jaxlib-networkx-qsimcirq-tensorflow-torch
[  0.04%] ··· Running (asv.core_suite.GradientComputation_light.time_gradient--).
[  0.09%] ··· ...uite.GradientComputation_light.time_gradient        4/16 failed
[  0.09%] ··· ========= ========== =========== =============
               n_wires   n_layers   interface
              --------- ---------- ----------- -------------
                  2         3        autograd   4.37±0.08ms
                  2         3           tf       8.90±0.1ms
                  2         3         torch     4.68±0.09ms
                  2         3          jax         failed
                  2         6        autograd    6.54±0.2ms
                  2         6           tf       14.0±0.1ms
                  2         6         torch      6.83±0.1ms
                  2         6          jax         failed
                  5         3        autograd    8.89±0.3ms
                  5         3           tf       17.9±0.4ms
                  5         3         torch      9.45±0.1ms
                  5         3          jax         failed
                  5         6        autograd    15.7±0.2ms
                  5         6           tf       31.6±0.1ms
                  5         6         torch      16.5±0.1ms
                  5         6          jax         failed
              ========= ========== =========== =============


[  0.00%] · For pennylane commit 0dcbed53 <master>:
[  0.00%] ·· Benchmarking virtualenv-py3.8-Qulacs-jax-jaxlib-networkx-qsimcirq-tensorflow-torch
[  0.04%] ··· Running (asv.core_suite.GradientComputation_light.time_gradient--).
[  0.09%] ··· ...uite.GradientComputation_light.time_gradient        3/16 failed
[  0.09%] ··· ========= ========== =========== =============
               n_wires   n_layers   interface
              --------- ---------- ----------- -------------
                  2         3        autograd   4.81±0.05ms
                  2         3           tf       9.14±0.1ms
                  2         3         torch     4.60±0.04ms
                  2         3          jax         failed
                  2         6        autograd    7.58±0.3ms
                  2         6           tf       13.9±0.2ms
                  2         6         torch      6.98±0.3ms
                  2         6          jax         failed
                  5         3        autograd   10.1±0.09ms
                  5         3           tf       18.1±0.3ms
                  5         3         torch      9.81±0.3ms
                  5         3          jax        682±8ms
                  5         6        autograd    17.6±0.3ms
                  5         6           tf       31.7±0.1ms
                  5         6         torch     16.4±0.05ms
                  5         6          jax         failed
              ========= ========== =========== =============
```

Looks like some modest gains in autograd, but not much change for tensorflow and torch.  Jax, as before, is problematic when getting timed. I'm curious why tensorflow and torch don't show a difference.
albi3ro: @trbromley Good catch on realizing that tensorflow and torch would be catching the state even when not using adjoint diff!

Now the state gets cached by default, but only if using the adjoint diff method.

I've renamed the keyword to `adjoint_cache`.

I've also removed `return_obs`.  May be useful in the future at some point, but I can easily add it in when the use case arises.

Patch:
diff --git a/.github/CHANGELOG.md b/.github/CHANGELOG.md
index 7aef99260..2981f8e00 100644
--- a/.github/CHANGELOG.md
+++ b/.github/CHANGELOG.md
@@ -310,12 +310,16 @@ random_mat2 = rng.standard_normal(3, requires_grad=False)
 
 <h3>Improvements</h3>
 
+* The adjoint jacobian differentiation method reuses the state computed on the forward pass.
+  This can be turned off to save memory with the Torch and TensorFlow interfaces by passing
+  `adjoint_cache=False` during QNode creation.
+  [(#1341)](https://github.com/PennyLaneAI/pennylane/pull/1341)
+
 * The `Operator` (and by inheritance, the `Operation` and `Observable` class and their children) 
   now have an `id` attribute, which can mark an operator in a circuit, for example to
   identify it on the tape by a tape transform.
   [(#1377)](https://github.com/PennyLaneAI/pennylane/pull/1377)
 
-  
 * Implement special handling for measuring the variance of Projector observables to improve memory usage.
   [(#1368)](https://github.com/PennyLaneAI/pennylane/pull/1368)
 
diff --git a/pennylane/_qubit_device.py b/pennylane/_qubit_device.py
index 64e45efd2..007a33fed 100644
--- a/pennylane/_qubit_device.py
+++ b/pennylane/_qubit_device.py
@@ -109,7 +109,15 @@ class QubitDevice(Device):
         new_array[indices] = array
         return new_array
 
-    observables = {"PauliX", "PauliY", "PauliZ", "Hadamard", "Hermitian", "Identity", "Projector"}
+    observables = {
+        "PauliX",
+        "PauliY",
+        "PauliZ",
+        "Hadamard",
+        "Hermitian",
+        "Identity",
+        "Projector",
+    }
 
     def __init__(self, wires=1, shots=None, cache=0, analytic=None):
         super().__init__(wires=wires, shots=shots, analytic=analytic)
@@ -798,7 +806,7 @@ class QubitDevice(Device):
 
         return samples.reshape((bin_size, -1))
 
-    def adjoint_jacobian(self, tape):
+    def adjoint_jacobian(self, tape, starting_state=None, use_device_state=False):
         """Implements the adjoint method outlined in
         `Jones and Gacon <https://arxiv.org/abs/2009.02823>`__ to differentiate an input tape.
 
@@ -807,7 +815,7 @@ class QubitDevice(Device):
         method, but has a lower time overhead and a similar memory overhead.
 
         .. note::
-            The adjoint differentation method has the following restrictions:
+            The adjoint differentiation method has the following restrictions:
 
             * As it requires knowledge of the statevector, only statevector simulator devices can be
               used.
@@ -817,6 +825,13 @@ class QubitDevice(Device):
         Args:
             tape (.QuantumTape): circuit that the function takes the gradient of
 
+        Keyword Args:
+            starting_state (tensor_like): post-forward pass state to start execution with. It should be
+                complex-valued. Takes precedence over ``use_device_state``.
+            use_device_state (bool): use current device state to initialize. A forward pass of the same
+                circuit should be the last thing the device has executed. If a ``starting_state`` is
+                provided, that takes precedence.
+
         Returns:
             array: the derivative of the tape with respect to trainable parameters.
             Dimensions are ``(len(observables), len(trainable_params))``.
@@ -825,6 +840,9 @@ class QubitDevice(Device):
             QuantumFunctionError: if the input tape has measurements that are not expectation values
                 or contains a multi-parameter operation aside from :class:`~.Rot`
         """
+        # broadcasted inner product not summing over first dimension of b
+        sum_axes = tuple(range(1, self.num_wires + 1))
+        dot_product_real = lambda b, k: self._real(qmlsum(self._conj(b) * k, axis=sum_axes))
 
         for m in tape.measurements:
             if m.return_type is not qml.operation.Expectation:
@@ -836,16 +854,19 @@ class QubitDevice(Device):
             if not hasattr(m.obs, "base_name"):
                 m.obs.base_name = None  # This is needed for when the observable is a tensor product
 
-        # Perform the forward pass.
-        # Consider using caching and calling lower-level functionality. We just need the device to
-        # be in the post-forward pass state.
-        # https://github.com/PennyLaneAI/pennylane/pull/1032/files#r563441040
-        self.reset()
-        self.execute(tape)
-
-        phi = self._reshape(self.state, [2] * self.num_wires)
+        # Initialization of state
+        if starting_state is not None:
+            ket = self._reshape(starting_state, [2] * self.num_wires)
+        else:
+            if not use_device_state:
+                self.reset()
+                self.execute(tape)
+            ket = self._pre_rotated_state
 
-        lambdas = [self._apply_operation(phi, obs) for obs in tape.observables]
+        n_obs = len(tape.observables)
+        bras = np.empty([n_obs] + [2] * self.num_wires, dtype=np.complex128)
+        for kk in range(n_obs):
+            bras[kk, ...] = self._apply_operation(ket, tape.observables[kk])
 
         expanded_ops = []
         for op in reversed(tape.operations):
@@ -863,7 +884,6 @@ class QubitDevice(Device):
                     expanded_ops.append(op)
 
         jac = np.zeros((len(tape.observables), len(tape.trainable_params)))
-        dot_product_real = lambda a, b: self._real(qmlsum(self._conj(a) * b))
 
         param_number = len(tape._par_info) - 1  # pylint: disable=protected-access
         trainable_param_number = len(tape.trainable_params) - 1
@@ -873,20 +893,22 @@ class QubitDevice(Device):
                 d_op_matrix = operation_derivative(op)
 
             op.inv()
-            phi = self._apply_operation(phi, op)
+            # Ideally use use op.adjoint() here
+            # then we don't have to re-invert the operation at the end
+            ket = self._apply_operation(ket, op)
 
             if op.grad_method is not None:
                 if param_number in tape.trainable_params:
-                    mu = self._apply_unitary(phi, d_op_matrix, op.wires)
 
-                    jac_column = np.array(
-                        [2 * dot_product_real(lambda_, mu) for lambda_ in lambdas]
-                    )
-                    jac[:, trainable_param_number] = jac_column
+                    ket_temp = self._apply_unitary(ket, d_op_matrix, op.wires)
+
+                    jac[:, trainable_param_number] = 2 * dot_product_real(bras, ket_temp)
+
                     trainable_param_number -= 1
                 param_number -= 1
 
-            lambdas = [self._apply_operation(lambda_, op) for lambda_ in lambdas]
+            for kk in range(n_obs):
+                bras[kk, ...] = self._apply_operation(bras[kk, ...], op)
             op.inv()
 
         return jac
diff --git a/pennylane/interfaces/jax.py b/pennylane/interfaces/jax.py
index a40418642..a1fb32c10 100644
--- a/pennylane/interfaces/jax.py
+++ b/pennylane/interfaces/jax.py
@@ -95,7 +95,9 @@ class JAXInterface(AnnotatedQueue):
         def wrapped_exec(params):
             exec_fn = partial(self.execute_device, device=device)
             return host_callback.call(
-                exec_fn, params, result_shape=jax.ShapeDtypeStruct((1,), JAXInterface.dtype)
+                exec_fn,
+                params,
+                result_shape=jax.ShapeDtypeStruct((1,), JAXInterface.dtype),
             )
 
         def wrapped_exec_fwd(params):
diff --git a/pennylane/interfaces/tf.py b/pennylane/interfaces/tf.py
index 3fdfd1143..0688dbda6 100644
--- a/pennylane/interfaces/tf.py
+++ b/pennylane/interfaces/tf.py
@@ -138,6 +138,15 @@ class TFInterface(AnnotatedQueue):
         res = self.execute_device(args, input_kwargs["device"])
         self.set_parameters(all_params, trainable_only=False)
 
+        use_adjoint_cached_state = False
+        # tape might not be a jacobian tape
+        jac_options = getattr(self, "jacobian_options", dict())
+        # cache state for adjoint jacobian computation
+        if jac_options.get("jacobian_method", None) == "adjoint_jacobian":
+            if jac_options.get("adjoint_cache", True):
+                use_adjoint_cached_state = True
+                state = input_kwargs["device"]._pre_rotated_state
+
         # The following dictionary caches the Jacobian and Hessian matrices,
         # so that they can be re-used for different vjp/vhp computations
         # within the same backpropagation call.
@@ -174,6 +183,9 @@ class TFInterface(AnnotatedQueue):
             if grad_matrix_fn in saved_grad_matrices:
                 return saved_grad_matrices[grad_matrix_fn]
 
+            if use_adjoint_cached_state:
+                self.jacobian_options["device_pd_options"] = {"starting_state": state}
+
             self.set_parameters(all_params_unwrapped, trainable_only=False)
             grad_matrix = getattr(self, grad_matrix_fn)(
                 input_kwargs["device"], params=args, **self.jacobian_options
diff --git a/pennylane/interfaces/torch.py b/pennylane/interfaces/torch.py
index 3548db155..0eea3c194 100644
--- a/pennylane/interfaces/torch.py
+++ b/pennylane/interfaces/torch.py
@@ -76,6 +76,15 @@ class _TorchInterface(torch.autograd.Function):
         if hasattr(res, "numpy"):
             res = res.numpy()
 
+        use_adjoint_cached_state = False
+        # tape might not be a jacobian tape
+        jac_options = getattr(tape, "jacobian_options", dict())
+        # cache state for adjoint jacobian computation
+        if jac_options.get("jacobian_method", None) == "adjoint_jacobian":
+            if jac_options.get("adjoint_cache", True):
+                use_adjoint_cached_state = True
+                state = device._pre_rotated_state
+
         ctx.saved_grad_matrices = {}
 
         def _evaluate_grad_matrix(grad_matrix_fn):
@@ -105,6 +114,9 @@ class _TorchInterface(torch.autograd.Function):
             if grad_matrix_fn in ctx.saved_grad_matrices:
                 return ctx.saved_grad_matrices[grad_matrix_fn]
 
+            if use_adjoint_cached_state:
+                tape.jacobian_options["device_pd_options"] = {"starting_state": state}
+
             tape.set_parameters(ctx.all_params_unwrapped, trainable_only=False)
             grad_matrix = getattr(tape, grad_matrix_fn)(
                 device, params=ctx.args, **tape.jacobian_options
diff --git a/pennylane/qnode.py b/pennylane/qnode.py
index 49af83b31..78861cc64 100644
--- a/pennylane/qnode.py
+++ b/pennylane/qnode.py
@@ -123,6 +123,9 @@ class QNode:
         order=1 (int): The order of the finite difference method to use. ``1`` corresponds
             to forward finite differences, ``2`` to centered finite differences.
         shift=pi/2 (float): the size of the shift for two-term parameter-shift gradient computations
+        adjoint_cache=True (bool): for TensorFlow and PyTorch interfaces and adjoint differentiation,
+            this indicates whether to save the device state after the forward pass.  Doing so saves a
+            forward execution. Device state automatically reused with autograd and JAX interfaces.
 
     **Example**
 
@@ -411,11 +414,17 @@ class QNode:
                 f"The {device.short_name} device does not support adjoint differentiation."
             )
 
+        jac_options = {"method": "device", "jacobian_method": "adjoint_jacobian"}
+        # reuse the forward pass
+        # torch and tensorflow can cache the state
+        if interface in {"autograd", "jax"}:
+            jac_options["device_pd_options"] = {"use_device_state": True}
+
         return (
             JacobianTape,
             interface,
             device,
-            {"method": "device", "jacobian_method": "adjoint_jacobian"},
+            jac_options,
         )
 
     @staticmethod
@@ -942,6 +951,9 @@ def qnode(
         h=1e-7 (float): Step size for the finite difference method.
         order=1 (int): The order of the finite difference method to use. ``1`` corresponds
             to forward finite differences, ``2`` to centered finite differences.
+        adjoint_cache=True (bool): for TensorFlow and PyTorch interfaces and adjoint differentiation,
+            this indicates whether to save the device state after the forward pass.  Doing so saves a
+            forward execution. Device state automatically reused with autograd and JAX interfaces.
 
     **Example**
 
diff --git a/pennylane/tape/jacobian_tape.py b/pennylane/tape/jacobian_tape.py
index 5d6f12cb5..dc1e640e8 100644
--- a/pennylane/tape/jacobian_tape.py
+++ b/pennylane/tape/jacobian_tape.py
@@ -359,7 +359,7 @@ class JacobianTape(QuantumTape):
 
         # TODO: modify devices that have device Jacobian methods to
         # accept the quantum tape as an argument
-        jac = jacobian_method(self)
+        jac = jacobian_method(self, **options.get("device_pd_options", dict()))
 
         # restore original parameters
         self.set_parameters(saved_parameters)


Test patch:
diff --git a/tests/interfaces/test_qnode_autograd.py b/tests/interfaces/test_qnode_autograd.py
index ee6485b98..02d5de08f 100644
--- a/tests/interfaces/test_qnode_autograd.py
+++ b/tests/interfaces/test_qnode_autograd.py
@@ -1074,6 +1074,25 @@ class TestQNode:
         assert j_spy.call_count == 5
 
 
+def test_adjoint_reuse_device_state(mocker):
+    """Tests that the autograd interface reuses the device state for adjoint differentiation"""
+    dev = qml.device("default.qubit", wires=1)
+
+    @qml.qnode(dev, diff_method="adjoint")
+    def circ(x):
+        qml.RX(x, wires=0)
+        return qml.expval(qml.PauliZ(0))
+
+    spy = mocker.spy(dev, "adjoint_jacobian")
+
+    grad = qml.grad(circ)(1.0)
+    assert circ.device.num_executions == 1
+
+    spy.assert_called_with(mocker.ANY, use_device_state=True)
+
+    assert circ.qtape.jacobian_options["device_pd_options"]["use_device_state"] == True
+
+
 def qtransform(qnode, a, framework=np):
     """Transforms every RY(y) gate in a circuit to RX(-a*cos(y))"""
 
diff --git a/tests/interfaces/test_qnode_jax.py b/tests/interfaces/test_qnode_jax.py
index ba6373461..0e48da3b4 100644
--- a/tests/interfaces/test_qnode_jax.py
+++ b/tests/interfaces/test_qnode_jax.py
@@ -133,6 +133,25 @@ def test_differentiable_expand(diff_method):
     assert np.allclose(res, expected, atol=tol, rtol=0)
 
 
+def test_adjoint_reuse_device_state(mocker):
+    """Tests that the jax interface reuses the device state for adjoint differentiation"""
+    dev = qml.device("default.qubit", wires=1)
+
+    @qml.qnode(dev, interface="jax", diff_method="adjoint")
+    def circ(x):
+        qml.RX(x, wires=0)
+        return qml.expval(qml.PauliZ(0))
+
+    spy = mocker.spy(dev, "adjoint_jacobian")
+
+    grad = jax.grad(circ)(1.0)
+    assert circ.device.num_executions == 1
+
+    spy.assert_called_with(mocker.ANY, use_device_state=True)
+
+    assert circ.qtape.jacobian_options["device_pd_options"]["use_device_state"] == True
+
+
 def qtransform(qnode, a, framework=jnp):
     """Transforms every RY(y) gate in a circuit to RX(-a*cos(y))"""
 
diff --git a/tests/interfaces/test_qnode_tf.py b/tests/interfaces/test_qnode_tf.py
index 993b84e19..114d1bd1f 100644
--- a/tests/interfaces/test_qnode_tf.py
+++ b/tests/interfaces/test_qnode_tf.py
@@ -545,7 +545,11 @@ class TestQNode:
             res = circuit(x, y)
 
         expected = np.array(
-            [tf.cos(x), (1 + tf.cos(x) * tf.cos(y)) / 2, (1 - tf.cos(x) * tf.cos(y)) / 2]
+            [
+                tf.cos(x),
+                (1 + tf.cos(x) * tf.cos(y)) / 2,
+                (1 - tf.cos(x) * tf.cos(y)) / 2,
+            ]
         )
         assert np.allclose(res, expected, atol=tol, rtol=0)
 
@@ -693,7 +697,10 @@ class TestQNode:
 
         a, b = x * 1.0
 
-        expected_res = [0.5 + 0.5 * tf.cos(a) * tf.cos(b), 0.5 - 0.5 * tf.cos(a) * tf.cos(b)]
+        expected_res = [
+            0.5 + 0.5 * tf.cos(a) * tf.cos(b),
+            0.5 - 0.5 * tf.cos(a) * tf.cos(b),
+        ]
         assert np.allclose(res, expected_res, atol=tol, rtol=0)
 
         expected_g = [
@@ -780,6 +787,104 @@ class TestQNode:
         np.testing.assert_allclose(hess, expected_hess, atol=tol, rtol=0, verbose=True)
 
 
+class Test_adjoint:
+    def test_adjoint_default_save_state(self, mocker):
+        """Tests that the state will be saved by default"""
+
+        dev = qml.device("default.qubit", wires=2)
+
+        @qml.qnode(dev, diff_method="adjoint", interface="tf")
+        def circ(x):
+            qml.RX(x[0], wires=0)
+            qml.RY(x[1], wires=1)
+            qml.CNOT(wires=(0, 1))
+            return qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliX(1))
+
+        expected_grad = lambda x: np.array([-np.sin(x[0]), np.cos(x[1])])
+
+        spy = mocker.spy(dev, "adjoint_jacobian")
+
+        x1 = tf.Variable([0.1, 0.2])
+        x2 = tf.Variable([0.3, 0.4])
+
+        with tf.GradientTape() as tape1:
+            res1 = circ(x1)
+
+        with tf.GradientTape() as tape2:
+            res2 = circ(x2)
+
+        grad1 = tape1.gradient(res1, x1)
+        grad2 = tape2.gradient(res2, x2)
+
+        assert np.allclose(grad1, expected_grad(x1))
+        assert np.allclose(grad2, expected_grad(x2))
+
+        assert circ.device.num_executions == 2
+        spy.assert_called_with(mocker.ANY, starting_state=mocker.ANY)
+
+    def test_adjoint_save_state(self, mocker):
+        """Tests that the tf interface reuses device state when prompted by `cache_state=True`.
+        Also makes sure executing a second circuit before backward pass does not interfere
+        with answer.
+        """
+
+        dev = qml.device("default.qubit", wires=2)
+
+        @qml.qnode(dev, diff_method="adjoint", interface="tf", adjoint_cache=True)
+        def circ(x):
+            qml.RX(x[0], wires=0)
+            qml.RY(x[1], wires=1)
+            qml.CNOT(wires=(0, 1))
+            return qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliX(1))
+
+        expected_grad = lambda x: np.array([-np.sin(x[0]), np.cos(x[1])])
+
+        spy = mocker.spy(dev, "adjoint_jacobian")
+
+        x1 = tf.Variable([0.1, 0.2])
+        x2 = tf.Variable([0.3, 0.4])
+
+        with tf.GradientTape() as tape1:
+            res1 = circ(x1)
+
+        with tf.GradientTape() as tape2:
+            res2 = circ(x2)
+
+        grad1 = tape1.gradient(res1, x1)
+        grad2 = tape2.gradient(res2, x2)
+
+        assert np.allclose(grad1, expected_grad(x1))
+        assert np.allclose(grad2, expected_grad(x2))
+
+        assert circ.device.num_executions == 2
+        spy.assert_called_with(mocker.ANY, starting_state=mocker.ANY)
+
+        assert circ.qtape.jacobian_options["adjoint_cache"] == True
+
+    def test_adjoint_no_save_state(self, mocker):
+        """Tests that with `adjoint_cache=False`, the state is not cached"""
+
+        dev = qml.device("default.qubit", wires=1)
+
+        @qml.qnode(dev, diff_method="adjoint", interface="tf", adjoint_cache=False)
+        def circ(x):
+            qml.RX(x, wires=0)
+            return qml.expval(qml.PauliZ(0))
+
+        spy = mocker.spy(dev, "adjoint_jacobian")
+
+        x = tf.Variable(0.1)
+        with tf.GradientTape() as tape:
+            res = circ(x)
+
+        grad = tape.gradient(res, x)
+
+        assert circ.device.num_executions == 2
+        spy.assert_called_with(mocker.ANY)
+
+        assert circ.qtape.jacobian_options.get("adjoint_cache", False) == False
+
+
 def qtransform(qnode, a, framework=tf):
     """Transforms every RY(y) gate in a circuit to RX(-a*cos(y))"""
 
@@ -817,7 +922,8 @@ def qtransform(qnode, a, framework=tf):
 
 
 @pytest.mark.parametrize(
-    "dev_name,diff_method", [("default.qubit", "finite-diff"), ("default.qubit.tf", "backprop")]
+    "dev_name,diff_method",
+    [("default.qubit", "finite-diff"), ("default.qubit.tf", "backprop")],
 )
 def test_transform(dev_name, diff_method, tol):
     """Test an example transform"""
diff --git a/tests/interfaces/test_qnode_torch.py b/tests/interfaces/test_qnode_torch.py
index fc7600deb..fc111a9c9 100644
--- a/tests/interfaces/test_qnode_torch.py
+++ b/tests/interfaces/test_qnode_torch.py
@@ -197,7 +197,10 @@ class TestQNode:
         loss = torch.sum(res)
 
         loss.backward()
-        expected = [-np.sin(a_val) + np.sin(a_val) * np.sin(b_val), -np.cos(a_val) * np.cos(b_val)]
+        expected = [
+            -np.sin(a_val) + np.sin(a_val) * np.sin(b_val),
+            -np.cos(a_val) * np.cos(b_val),
+        ]
         assert np.allclose(a.grad, expected[0], atol=tol, rtol=0)
         assert np.allclose(b.grad, expected[1], atol=tol, rtol=0)
 
@@ -292,7 +295,10 @@ class TestQNode:
         loss = torch.sum(res)
         loss.backward()
 
-        expected = [-np.sin(a_val) + np.sin(a_val) * np.sin(b_val), -np.cos(a_val) * np.cos(b_val)]
+        expected = [
+            -np.sin(a_val) + np.sin(a_val) * np.sin(b_val),
+            -np.cos(a_val) * np.cos(b_val),
+        ]
         assert np.allclose([a.grad, b.grad], expected, atol=tol, rtol=0)
 
         # JacobianTape.numeric_pd has been called for each argument
@@ -376,7 +382,11 @@ class TestQNode:
             res.backward()
 
     @pytest.mark.parametrize(
-        "U", [torch.tensor([[0, 1], [1, 0]], requires_grad=False), np.array([[0, 1], [1, 0]])]
+        "U",
+        [
+            torch.tensor([[0, 1], [1, 0]], requires_grad=False),
+            np.array([[0, 1], [1, 0]]),
+        ],
     )
     def test_matrix_parameter(self, dev_name, diff_method, U, tol):
         """Test that the Torch interface works correctly
@@ -430,7 +440,10 @@ class TestQNode:
 
         tape_params = [i.detach().numpy() for i in circuit.qtape.get_parameters()]
         assert np.allclose(
-            tape_params, [p_val[2], p_val[0], -p_val[2], p_val[1] + p_val[2]], atol=tol, rtol=0
+            tape_params,
+            [p_val[2], p_val[0], -p_val[2], p_val[1] + p_val[2]],
+            atol=tol,
+            rtol=0,
         )
 
         expected = np.cos(a) * np.cos(p_val[1]) * np.sin(p_val[0]) + np.sin(a) * (
@@ -481,7 +494,10 @@ class TestQNode:
         expected = np.array(
             [
                 [np.cos(x_val / 2) ** 2, np.sin(x_val / 2) ** 2],
-                [(1 + np.cos(x_val) * np.cos(y_val)) / 2, (1 - np.cos(x_val) * np.cos(y_val)) / 2],
+                [
+                    (1 + np.cos(x_val) * np.cos(y_val)) / 2,
+                    (1 - np.cos(x_val) * np.cos(y_val)) / 2,
+                ],
             ]
         )
         assert np.allclose(res.detach().numpy(), expected, atol=tol, rtol=0)
@@ -638,7 +654,10 @@ class TestQNode:
 
         a, b = x.detach().numpy()
 
-        expected_res = [0.5 + 0.5 * np.cos(a) * np.cos(b), 0.5 - 0.5 * np.cos(a) * np.cos(b)]
+        expected_res = [
+            0.5 + 0.5 * np.cos(a) * np.cos(b),
+            0.5 - 0.5 * np.cos(a) * np.cos(b),
+        ]
         assert np.allclose(res.detach(), expected_res, atol=tol, rtol=0)
 
         expected_g = [
@@ -717,6 +736,98 @@ class TestQNode:
         assert np.allclose(hess.detach(), expected_hess, atol=tol, rtol=0)
 
 
+class Test_adjoint:
+    def test_adjoint_default_save_state(self, mocker):
+        """tests that the state will be saved by default"""
+
+        dev = qml.device("default.qubit", wires=2)
+
+        @qml.qnode(dev, diff_method="adjoint", interface="torch")
+        def circ(x):
+            qml.RX(x[0], wires=0)
+            qml.RY(x[1], wires=1)
+            qml.CNOT(wires=(0, 1))
+            return qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliX(1))
+
+        expected_grad = lambda x: torch.tensor([-torch.sin(x[0]), torch.cos(x[1])])
+
+        spy = mocker.spy(dev, "adjoint_jacobian")
+
+        x1 = torch.tensor([0.1, 0.2], requires_grad=True)
+        x2 = torch.tensor([0.3, 0.4], requires_grad=True)
+
+        res1 = circ(x1)
+        res2 = circ(x2)
+
+        res1.backward(torch.Tensor([1, 1]))
+        res2.backward(torch.Tensor([1, 1]))
+
+        assert np.allclose(x1.grad, expected_grad(x1))
+        assert np.allclose(x2.grad, expected_grad(x2))
+
+        assert circ.device.num_executions == 2
+
+        spy.assert_called_with(mocker.ANY, starting_state=mocker.ANY)
+
+    def test_adjoint_save_state(self, mocker):
+        """Tests that the torch interface reuses device state when prompted by `cache_state=True`.
+        Also tests a second execution before backward pass does not alter gradient.
+        """
+
+        dev = qml.device("default.qubit", wires=2)
+
+        @qml.qnode(dev, diff_method="adjoint", interface="torch", adjoint_cache=True)
+        def circ(x):
+            qml.RX(x[0], wires=0)
+            qml.RY(x[1], wires=1)
+            qml.CNOT(wires=(0, 1))
+            return qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliX(1))
+
+        expected_grad = lambda x: torch.tensor([-torch.sin(x[0]), torch.cos(x[1])])
+
+        spy = mocker.spy(dev, "adjoint_jacobian")
+
+        x1 = torch.tensor([0.1, 0.2], requires_grad=True)
+        x2 = torch.tensor([0.3, 0.4], requires_grad=True)
+
+        res1 = circ(x1)
+        res2 = circ(x2)
+
+        res1.backward(torch.Tensor([1, 1]))
+        res2.backward(torch.Tensor([1, 1]))
+
+        assert np.allclose(x1.grad, expected_grad(x1))
+        assert np.allclose(x2.grad, expected_grad(x2))
+
+        assert circ.device.num_executions == 2
+
+        spy.assert_called_with(mocker.ANY, starting_state=mocker.ANY)
+
+        assert circ.qtape.jacobian_options["adjoint_cache"] == True
+
+    def test_adjoint_no_save_state(self, mocker):
+        """Tests that with `adjoint_cache=False`, the state is not cached"""
+
+        dev = qml.device("default.qubit", wires=1)
+
+        @qml.qnode(dev, diff_method="adjoint", interface="torch", adjoint_cache=False)
+        def circ(x):
+            qml.RX(x, wires=0)
+            return qml.expval(qml.PauliZ(0))
+
+        spy = mocker.spy(dev, "adjoint_jacobian")
+
+        x = torch.tensor(0.1, requires_grad=True)
+        res = circ(x)
+        res.backward()
+
+        assert circ.device.num_executions == 2
+
+        spy.assert_called_with(mocker.ANY)
+
+        assert circ.qtape.jacobian_options.get("adjoint_cache", False) == False
+
+
 def qtransform(qnode, a, framework=torch):
     """Transforms every RY(y) gate in a circuit to RX(-a*cos(y))"""
 
diff --git a/tests/test_qubit_device_adjoint_jacobian.py b/tests/test_qubit_device_adjoint_jacobian.py
index 9f205084c..b42109006 100644
--- a/tests/test_qubit_device_adjoint_jacobian.py
+++ b/tests/test_qubit_device_adjoint_jacobian.py
@@ -192,6 +192,45 @@ class TestAdjointJacobian:
         # the different methods agree
         assert np.allclose(grad_D, grad_F, atol=tol, rtol=0)
 
+    def test_use_device_state(self, tol, dev):
+        """Tests that when using the device state, the correct answer is still returned."""
+
+        x, y, z = [0.5, 0.3, -0.7]
+
+        with qml.tape.JacobianTape() as tape:
+            qml.RX(0.4, wires=[0])
+            qml.Rot(x, y, z, wires=[0])
+            qml.RY(-0.2, wires=[0])
+            qml.expval(qml.PauliZ(0))
+
+        tape.trainable_params = {1, 2, 3}
+
+        dM1 = dev.adjoint_jacobian(tape)
+
+        tape.execute(dev)
+        dM2 = dev.adjoint_jacobian(tape, use_device_state=True)
+
+        assert np.allclose(dM1, dM2, atol=tol, rtol=0)
+
+    def test_provide_starting_state(self, tol, dev):
+        """Tests provides correct answer when provided starting state."""
+        x, y, z = [0.5, 0.3, -0.7]
+
+        with qml.tape.JacobianTape() as tape:
+            qml.RX(0.4, wires=[0])
+            qml.Rot(x, y, z, wires=[0])
+            qml.RY(-0.2, wires=[0])
+            qml.expval(qml.PauliZ(0))
+
+        tape.trainable_params = {1, 2, 3}
+
+        dM1 = dev.adjoint_jacobian(tape)
+
+        tape.execute(dev)
+        dM2 = dev.adjoint_jacobian(tape, starting_state=dev._pre_rotated_state)
+
+        assert np.allclose(dM1, dM2, atol=tol, rtol=0)
+
 
 class TestAdjointJacobianQNode:
     """Test QNode integration with the adjoint_jacobian method"""