Merge pull request #125 from firedrakeproject/fix-listtensor-pickle

Fix ListTensor pickling
firedrakeproject · May 30, 2017 · 7cbae93 · 7cbae93
2 parents 361c1b0 + 78f1af1
commit 7cbae93
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 0 deletions.
diff --git a/gem/gem.py b/gem/gem.py
@@ -655,6 +655,9 @@ def children(self):
     def shape(self):
         return self.array.shape
 
+    def __reduce__(self):
+        return type(self), (self.array,)
+
     def reconstruct(self, *args):
         return ListTensor(asarray(args).reshape(self.array.shape))
 

diff --git a/tests/test_pickle_gem.py b/tests/test_pickle_gem.py
@@ -18,6 +18,14 @@ def test_pickle_gem(protocol):
     assert repr(expr) == repr(unpickled)
 
 
+@pytest.mark.parametrize('protocol', range(3))
+def test_listtensor(protocol):
+    expr = gem.ListTensor([gem.Variable('x', ()), gem.Zero()])
+
+    unpickled = pickle.loads(pickle.dumps(expr, protocol))
+    assert expr == unpickled
+
+
 if __name__ == "__main__":
     import os
     import sys