Repository: PennyLaneAI/pennylane
Pull request #1856 (2021-11-04 19:00:03)

Fix for issue 1848. Pickling and unpickling a qml.numpy.tensor object does not preserve requires_grad attribute


------------------------------------------------------------------------------------------------------------

**Context:** Upon pickling and unpickling a qml.numpy.tensor object, the attribute `requires_grad` is lost. Aside from possible computation errors that could arise, this also causes an AttributeError when attempting to print the object.

**Description of the Change:** Re-implement the parent class' (np.ndarray) `__reduce__` and `__setstate__` methods, which are called during pickling and unpickling. The re-implementation ensures that `requires_grad` is not lost.

**Benefits:** No loss of information or raised errors after unpickling.

**Possible Drawbacks:** None foreseen.

**Related GitHub Issues:** https://github.com/PennyLaneAI/pennylane/issues/966, only in that this issue also results in a loss of the `requires_grad` attribute. 


codecov[bot]: # [Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1856?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
> Merging [#1856](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1856?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (03d1a5f) into [master](https://codecov.io/gh/PennyLaneAI/pennylane/commit/5443339b145c56b89c14451dca3dc6bfe1a3af46?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (5443339) will **increase** coverage by `0.00%`.
> The diff coverage is `100.00%`.

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

```diff
@@           Coverage Diff           @@
##           master    #1856   +/-   ##
=======================================
  Coverage   98.98%   98.98%           
=======================================
  Files         218      218           
  Lines       16426    16433    +7     
=======================================
+ Hits        16260    16267    +7     
  Misses        166      166           
```


| [Impacted Files](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1856?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | Coverage Δ | |
|---|---|---|
| [pennylane/numpy/tensor.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1856/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL251bXB5L3RlbnNvci5weQ==) | `97.53% <100.00%> (+0.23%)` | :arrow_up: |

------

[Continue to review full report at Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1856?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/1856?src=pr&el=footer&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI). Last update [5443339...03d1a5f](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1856?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 757ac23ff..e680f29aa 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -19,10 +19,14 @@
   returns the default empty wires object.
   [(#1830)](https://github.com/PennyLaneAI/pennylane/pull/1830)
 
+* The `requires_grad` attribute of `qml.numpy.tensor` objects is now
+  preserved when pickling/unpickling the object.
+  [(#1856)](https://github.com/PennyLaneAI/pennylane/pull/1856)
+
 <h3>Documentation</h3>
 
 <h3>Contributors</h3>
 
 This release contains contributions from (in alphabetical order):
 
-Christina Lee, Alejandro Montanez
+Jalani Kanem, Christina Lee, Alejandro Montanez
diff --git a/pennylane/numpy/tensor.py b/pennylane/numpy/tensor.py
index 4803c95d7..03ec40767 100644
--- a/pennylane/numpy/tensor.py
+++ b/pennylane/numpy/tensor.py
@@ -197,6 +197,24 @@ class tensor(_np.ndarray):
 
         raise TypeError("unhashable type: 'numpy.tensor'")
 
+    def __reduce__(self):
+        # Called when pickling the object.
+        # Numpy ndarray uses __reduce__ instead of __getstate__ to prepare an object for
+        # pickling. self.requires_grad needs to be included in the tuple returned by
+        # __reduce__ in order to be preserved in the unpickled object.
+        reduced_obj = super().__reduce__()
+        # The last (2nd) element of this tuple holds the data. Add requires_grad to this:
+        full_reduced_data = reduced_obj[2] + (self.requires_grad,)
+        return (reduced_obj[0], reduced_obj[1], full_reduced_data)
+
+    def __setstate__(self, reduced_obj) -> None:
+        # Called when unpickling the object.
+        # Set self.requires_grad with the last element in the tuple returned by __reduce__:
+        # pylint: disable=attribute-defined-outside-init
+        self.requires_grad = reduced_obj[-1]
+        # And call parent's __setstate__ without this element:
+        super().__setstate__(reduced_obj[:-1])
+
     def unwrap(self):
         """Converts the tensor to a standard, non-differentiable NumPy ndarray or Python scalar if
         the tensor is 0-dimensional.


Test patch:
diff --git a/tests/numpy/test_pickling.py b/tests/numpy/test_pickling.py
new file mode 100644
index 000000000..87ed3b10b
--- /dev/null
+++ b/tests/numpy/test_pickling.py
@@ -0,0 +1,32 @@
+# Copyright 2018-2020 Xanadu Quantum Technologies Inc.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+Unit tests for ensuring that objects that are pickled/unpickled are identical to the original.
+"""
+
+import pickle
+
+import pennylane as qml
+
+
+def test_unpickling_tensor():
+    """Tests whether qml.numpy.tensor objects are pickleable."""
+
+    x = qml.numpy.random.random(15)
+    x_str = pickle.dumps(x)
+    x_reloaded = pickle.loads(x_str)
+
+    assert qml.numpy.allclose(x, x_reloaded)
+    assert x.__dict__ == x_reloaded.__dict__
+    assert hasattr(x_reloaded, "requires_grad")