Repository: PennyLaneAI/pennylane
Pull request #2171 (2022-02-04 20:34:30)

Fix `dev.num_executions` bug with QNode caching
**Context:**
See linked issue https://github.com/PennyLaneAI/pennylane/issues/2170. The issue arises when swapping the underlying device in the QNode (e.g., when having `diff_method="backprop"`).

**Description of the Change:**

Instead of incrementing the number of executions, we set the number of executions using the number of executions from the device used.

**Benefits:**

`dev.num_executions` is recorded correctly.

**Possible Drawbacks:**
N/A

**Related GitHub Issues:**
Closes https://github.com/PennyLaneAI/pennylane/issues/2170

codecov[bot]: # [Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2171?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
> :exclamation: No coverage uploaded for pull request base (`v0.21.0-rc0@2a419f4`). [Click here to learn what that means](https://docs.codecov.io/docs/error-reference?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#section-missing-base-commit).
> The diff coverage is `n/a`.

[![Impacted file tree graph](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2171/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/2171?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI)

```diff
@@              Coverage Diff               @@
##             v0.21.0-rc0    #2171   +/-   ##
==============================================
  Coverage               ?   96.38%           
==============================================
  Files                  ?      230           
  Lines                  ?    17670           
  Branches               ?        0           
==============================================
  Hits                   ?    17031           
  Misses                 ?      639           
  Partials               ?        0           
```



------

[Continue to review full report at Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2171?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/2171?src=pr&el=footer&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI). Last update [2a419f4...7ed3fdd](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2171?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).


Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 06e6332ee..19164addb 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -525,6 +525,10 @@
 
 <h3>Bug fixes</h3>
 
+* Fixes a bug where an incorrect number of executions are recorded by
+  a QNode using a custom cache with `diff_method="backprop"`.
+  [(#2171)](https://github.com/PennyLaneAI/pennylane/pull/2171)
+
 * Fixes a bug where the `default.qubit.jax` device can't be used with `diff_method=None` and jitting.
   [(#2136)](https://github.com/PennyLaneAI/pennylane/pull/2136)
 
diff --git a/pennylane/qnode.py b/pennylane/qnode.py
index bdfd1983c..0a0a7da27 100644
--- a/pennylane/qnode.py
+++ b/pennylane/qnode.py
@@ -217,6 +217,7 @@ class QNode:
         self._original_device = device
         self.gradient_fn = None
         self.gradient_kwargs = None
+        self._tape_cached = False
 
         self._update_gradient_fn()
         functools.update_wrapper(self, func)
@@ -265,7 +266,9 @@ class QNode:
         # of the user's device before and after executing the tape.
 
         if self.device is not self._original_device:
-            self._original_device._num_executions += 1  # pylint: disable=protected-access
+
+            if not self._tape_cached:
+                self._original_device._num_executions += 1  # pylint: disable=protected-access
 
             # Update for state vector simulators that have the _pre_rotated_state attribute
             if hasattr(self._original_device, "_pre_rotated_state"):
@@ -546,6 +549,14 @@ class QNode:
         # construct the tape
         self.construct(args, kwargs)
 
+        cache = self.execute_kwargs.get("cache", False)
+        using_custom_cache = (
+            hasattr(cache, "__getitem__")
+            and hasattr(cache, "__setitem__")
+            and hasattr(cache, "__delitem__")
+        )
+        self._tape_cached = using_custom_cache and self.tape.hash in cache
+
         res = qml.execute(
             [self.tape],
             device=self.device,


Test patch:
diff --git a/tests/test_qnode.py b/tests/test_qnode.py
index f52784875..36340e0b6 100644
--- a/tests/test_qnode.py
+++ b/tests/test_qnode.py
@@ -711,6 +711,56 @@ class TestIntegration:
 
         assert dev.num_executions == 6
 
+    def test_num_exec_caching_device_swap(self):
+        """Tests that if we swapped the original device (e.g., when
+        diff_method='backprop') then the number of executions recorded is
+        correct."""
+        dev = qml.device("default.qubit", wires=2)
+
+        cache = {}
+
+        @qml.qnode(dev, diff_method="backprop", cache=cache)
+        def circuit():
+            qml.RY(0.345, wires=0)
+            return qml.expval(qml.PauliZ(0))
+
+        for _ in range(15):
+            circuit()
+
+        # Although we've evaluated the QNode more than once, due to caching,
+        # there was one device execution recorded
+        assert dev.num_executions == 1
+        assert cache != {}
+
+    def test_num_exec_caching_device_swap_two_exec(self):
+        """Tests that if we swapped the original device (e.g., when
+        diff_method='backprop') then the number of executions recorded is
+        correct even with multiple QNode evaluations."""
+        dev = qml.device("default.qubit", wires=2)
+
+        cache = {}
+
+        @qml.qnode(dev, diff_method="backprop", cache=cache)
+        def circuit():
+            qml.RY(0.345, wires=0)
+            return qml.expval(qml.PauliZ(0))
+
+        for _ in range(15):
+            circuit()
+
+        @qml.qnode(dev, diff_method="backprop", cache=cache)
+        def circuit():
+            qml.RZ(0.345, wires=0)
+            return qml.expval(qml.PauliZ(0))
+
+        for _ in range(15):
+            circuit()
+
+        # Although we've evaluated the QNode several times, due to caching,
+        # there were two device executions recorded
+        assert dev.num_executions == 2
+        assert cache != {}
+
     @pytest.mark.parametrize("diff_method", ["parameter-shift", "finite-diff"])
     def test_single_expectation_value_with_argnum_one(self, diff_method, tol):
         """Tests correct output shape and evaluation for a QNode