Repository: PennyLaneAI/pennylane
Pull request #1270 (2021-05-03 23:19:08)

Cache mapped wires in `Device`
**Context:**

Wires used internally in devices get mapped based on the initial map wiring defined. Although doing this wire mapping is not too computationally intensive when done a single time, when running an optimization task where it's done thousands of times, the cost for computing mapped wires is non-negligible in terms of time.

Often devices query the mapping for the same wires several times. As the wire mapping of a device is set only when the device is created, it is safe to assume that it's immutable. Therefore we can assume that the value of mapping wires will also not change over the life of the device.

**Description of the Change:**

* When calling `map_wires`, the `Device` class now caches the key-value pair of the wires and the mapped wires in a dictionary
* The `hash` of a `Wires` object is cached

**Benefits:**
More performant device and `Wires` class usage.

**Possible Drawbacks:**
N/A

**Related GitHub Issues:**
N/A

github-actions[bot]: Hello. You may have forgotten to update the changelog!
Please edit [.github/CHANGELOG.md](/PennyLaneAI/pennylane/blob/master/.github/CHANGELOG.md) with:
* A one-to-two sentence description of the change. You may include a small working example for new features.
* A link back to this PR.
* Your name (or GitHub username) in the contributors section.
codecov[bot]: # [Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1270?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
> Merging [#1270](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1270?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (051f845) into [master](https://codecov.io/gh/PennyLaneAI/pennylane/commit/8fe31ad619222e2992130c5d036571f6ef7ab51c?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (8fe31ad) will **increase** coverage by `0.00%`.
> The diff coverage is `100.00%`.

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

```diff
@@           Coverage Diff           @@
##           master    #1270   +/-   ##
=======================================
  Coverage   98.12%   98.12%           
=======================================
  Files         146      146           
  Lines       11071    11078    +7     
=======================================
+ Hits        10863    10870    +7     
  Misses        208      208           
```


| [Impacted Files](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1270?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | Coverage Δ | |
|---|---|---|
| [pennylane/\_device.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1270/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL19kZXZpY2UucHk=) | `96.74% <100.00%> (+0.05%)` | :arrow_up: |
| [pennylane/wires.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1270/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3dpcmVzLnB5) | `98.71% <100.00%> (+0.02%)` | :arrow_up: |

------

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

antalszava: Some benchmarks:

*GradientComputation_light.time_gradient*

Commit 8fe31ad619222e2992130c5d036571f6ef7ab51c

              --                                        n_layers / interface      
              --------- -----------------------------------------------------------------------------
               n_wires   3 / autograd    3 / tf    3 / torch    6 / autograd    6 / tf    6 / torch  
              ========= ============== ========== ============ ============== ========== ============
                  2      9.21±0.09ms    33.6±6ms   14.7±0.2ms     15.0±2ms     52.0±2ms   40.5±0.9ms 
                  5       20.0±0.1ms    92.0±3ms   78.2±0.9ms     36.5±1ms     172±8ms     328±40ms  
              ========= ============== ========== ============ ============== ========== ============


Commit 051f845aae5ed7ed6f91a0eeac0884014916aa86

              --                                        n_layers / interface                               
              --------- -----------------------------------------------------------------------------------
               n_wires   3 / autograd      3 / tf      3 / torch    6 / autograd     6 / tf     6 / torch  
              ========= ============== ============= ============= ============== ============ ============
                  2      5.13±0.03ms    16.6±0.08ms   8.26±0.04ms    7.91±0.1ms    29.8±0.1ms   22.6±0.2ms 
                  5       11.0±0.2ms      45.2±4ms      43.6±7ms      22.6±2ms      85.3±1ms     154±20ms  
              ========= ============== ============= ============= ============== ============ ============


*ML_light*

Commit 8fe31ad619222e2992130c5d036571f6ef7ab51c

Mem:

                autograd   256M 
                 torch     283M 
                   tf      303M 

Time:

                autograd   10.9±0.4s 
                 torch       failed  
                   tf      23.8±0.4s 

Commit 051f845aae5ed7ed6f91a0eeac0884014916aa86

Mem:

                autograd   256M 
                 torch     284M 
                   tf      303M 

Time:

                autograd   9.55±0.3s 
                 torch       failed  
                   tf      23.2±0.3s 

Patch:
diff --git a/.github/CHANGELOG.md b/.github/CHANGELOG.md
index 0a96ba0d4..3615a5139 100644
--- a/.github/CHANGELOG.md
+++ b/.github/CHANGELOG.md
@@ -105,6 +105,12 @@ random_mat2 = rng.standard_normal(3, requires_grad=False)
 
 <h3>Improvements</h3>
 
+* The `Device` class now uses caching when mapping wires.
+  [(#1270)](https://github.com/PennyLaneAI/pennylane/pull/1270)
+
+* The `Wires` class now uses caching for computing its `hash`.
+  [(#1270)](https://github.com/PennyLaneAI/pennylane/pull/1270)
+
 * Added custom gate application for Toffoli in `default.qubit`.
   [(#1249)](https://github.com/PennyLaneAI/pennylane/pull/1249)
 
diff --git a/pennylane/_device.py b/pennylane/_device.py
index 737dbb8b6..a9e3d601b 100644
--- a/pennylane/_device.py
+++ b/pennylane/_device.py
@@ -18,6 +18,7 @@ This module contains the :class:`Device` abstract base class.
 import abc
 from collections.abc import Iterable, Sequence
 from collections import OrderedDict, namedtuple
+from functools import lru_cache
 
 import numpy as np
 
@@ -331,6 +332,7 @@ class Device(abc.ABC):
         wire_map = zip(wires, consecutive_wires)
         return OrderedDict(wire_map)
 
+    @lru_cache()
     def map_wires(self, wires):
         """Map the wire labels of wires using this device's wire map.
 
diff --git a/pennylane/devices/default_qubit.py b/pennylane/devices/default_qubit.py
index 45e8c8684..494fe11c0 100644
--- a/pennylane/devices/default_qubit.py
+++ b/pennylane/devices/default_qubit.py
@@ -154,6 +154,7 @@ class DefaultQubit(QubitDevice):
             "Toffoli": self._apply_toffoli,
         }
 
+    @functools.lru_cache()
     def map_wires(self, wires):
         # temporarily overwrite this method to bypass
         # wire map that produces Wires objects
diff --git a/pennylane/wires.py b/pennylane/wires.py
index 26512e7d5..d851a854f 100644
--- a/pennylane/wires.py
+++ b/pennylane/wires.py
@@ -99,6 +99,8 @@ class Wires(Sequence):
         else:
             self._labels = _process(wires)
 
+        self._hash = None
+
     def __getitem__(self, idx):
         """Method to support indexing. Returns a Wires object if index is a slice, or a label if index is an integer."""
         if isinstance(idx, slice):
@@ -131,7 +133,9 @@ class Wires(Sequence):
 
     def __hash__(self):
         """Implements the hash function."""
-        return hash(self._labels)
+        if self._hash is None:
+            self._hash = hash(self._labels)
+        return self._hash
 
     def __add__(self, other):
         """Defines the addition to return a Wires object containing all wires of the two terms.


Test patch:
diff --git a/tests/devices/test_default_qubit.py b/tests/devices/test_default_qubit.py
index 74ae91f81..702df94cc 100644
--- a/tests/devices/test_default_qubit.py
+++ b/tests/devices/test_default_qubit.py
@@ -23,7 +23,7 @@ import pytest
 import pennylane as qml
 from pennylane import numpy as np, DeviceError
 from pennylane.devices.default_qubit import _get_slice, DefaultQubit
-from pennylane.wires import WireError
+from pennylane.wires import Wires, WireError
 
 U = np.array(
     [
@@ -1855,6 +1855,33 @@ class TestWiresIntegration:
         with pytest.raises(WireError, match="Did not find some of the wires"):
             dev.execute(tape)
 
+    wires_to_try = [
+        (1, Wires([0]), Wires([0])),
+        (4, Wires([1, 3]), Wires([1, 3])),
+        (["a", 2], Wires([2]), Wires([1])),
+        (["a", 2], Wires([2, "a"]), Wires([1, 0])),
+    ]
+
+    @pytest.mark.parametrize("dev_wires, wires_to_map, res", wires_to_try)
+    def test_map_wires_caches(self, dev_wires, wires_to_map, res, mock_device):
+        """Test that multiple calls to map_wires will use caching."""
+        dev = qml.device("default.qubit", wires=dev_wires)
+
+        original_hits = dev.map_wires.cache_info().hits
+        original_misses = dev.map_wires.cache_info().misses
+
+        # The first call is computed: it's a miss as it didn't come from the cache
+        dev.map_wires(wires_to_map)
+
+        # The number of misses increased
+        assert dev.map_wires.cache_info().misses > original_misses
+
+        # The second call comes from the cache: it's a hit
+        dev.map_wires(wires_to_map)
+
+        # The number of hits increased
+        assert dev.map_wires.cache_info().hits > original_hits
+
 
 class TestGetSlice:
     """Tests for the _get_slice function."""
diff --git a/tests/test_device.py b/tests/test_device.py
index bda6f5576..85518f218 100644
--- a/tests/test_device.py
+++ b/tests/test_device.py
@@ -478,6 +478,32 @@ class TestInternalFunctions:
             node_gauss(0.015, 0.02, 0.005)
         assert dev_gauss.num_executions == num_evals_gauss
 
+    wires_to_try = [
+        (1, Wires([0]), Wires([0])),
+        (4, Wires([1, 3]), Wires([1, 3])),
+        (["a", 2], Wires([2]), Wires([1])),
+        (["a", 2], Wires([2, "a"]), Wires([1, 0])),
+    ]
+
+    @pytest.mark.parametrize("dev_wires, wires_to_map, res", wires_to_try)
+    def test_map_wires_caches(self, dev_wires, wires_to_map, res, mock_device):
+        """Test that multiple calls to map_wires will use caching."""
+        dev = mock_device(dev_wires)
+        original_hits = dev.map_wires.cache_info().hits
+        original_misses = dev.map_wires.cache_info().misses
+
+        # The first call is computed: it's a miss as it didn't come from the cache
+        dev.map_wires(wires_to_map)
+
+        # The number of misses increased
+        assert dev.map_wires.cache_info().misses > original_misses
+
+        # The second call comes from the cache: it's a hit
+        dev.map_wires(wires_to_map)
+
+        # The number of hits increased
+        assert dev.map_wires.cache_info().hits > original_hits
+
 
 class TestClassmethods:
     """Test the classmethods of Device"""
diff --git a/tests/test_wires.py b/tests/test_wires.py
index 09db543cd..271d68e92 100644
--- a/tests/test_wires.py
+++ b/tests/test_wires.py
@@ -355,3 +355,10 @@ class TestWires:
         assert Wires([1, 2, 3]) == (1, 2, 3)
         assert Wires([1, 2, 3]) != (1, 5, 3)
         assert (1, 5, 3) != Wires([1, 2, 3])
+
+    def test_hash_cached(self):
+        """Test that the hash of a Wires object is being cached."""
+        wires = Wires([0, 1])
+        assert wires._hash is None
+        h = hash(wires)
+        assert wires._hash == h