make TTOperator picklable

325e579a · Philipp Trunschke · 9c39279d · 325e579a
Commit 325e579a authored Dec 04, 2020 by Philipp Trunschke
Hide whitespace changes
Inline Side-by-side

Showing with 8 additions and 0 deletions

src/xerus/python/ttnetwork.cpp src/xerus/python/ttnetwork.cpp +8 -0

No files found.
--- a/src/xerus/python/ttnetwork.cpp
+++ b/src/xerus/python/ttnetwork.cpp
@@ -79,6 +79,14 @@ void expose_ttnetwork(module& m) {
 	m.def("dyadic_product", static_cast<TTTensor (*)(const std::vector<TTTensor> &)>(&dyadic_product));

 	class_<TTOperator, TensorNetwork>(m, "TTOperator")
+		.def(pickle(
+			[](const TTTensor &_self) { // __getstate__
+				return bytes(misc::serialize(_self));
+			},
+			[](bytes _bytes) { // __setstate__
+				return misc::deserialize<TTTensor>(_bytes);
+			}
+		))
 		.def(init<const Tensor&>())
 		.def(init<const Tensor&, value_t>())
 		.def(init<const Tensor&, value_t, size_t>())