Add serialization logic for complex numbers (pytorch#50885)

Summary: Pull Request resolved: pytorch#50885

Test Plan: Imported from OSS

Reviewed By: SplitInfinity

Differential Revision: D26094906

Pulled By: anjali411

fbshipit-source-id: 7b2614f3ee4a30c4b4cf04aaa3432988b38a0721

Author: anjali411

test/test_complex.py

@@ -12,6 +12,18 @@ def fn(a: complex):
 
         self.checkScript(fn, (3 + 5j,))
 
+    def test_pickle(self):
+        class ComplexModule(torch.jit.ScriptModule):
+            def __init__(self):
+                super().__init__()
+                self.a = 3 + 5j
+
+            def forward(self, b: int):
+                return b
+
+        loaded = self.getExportImportCopy(ComplexModule())
+        self.assertEqual(loaded.a, 3 + 5j)
+
 class TestComplexTensor(TestCase):
     @dtypes(*torch.testing.get_all_complex_dtypes())
     def test_to_list(self, device, dtype):

torch/csrc/jit/python/pybind_utils.h

@@ -293,6 +293,8 @@ inline InferredType tryToInferType(py::handle input) {
     return InferredType(IntType::get());
   } else if (py::isinstance<py::float_>(input)) {
     return InferredType(FloatType::get());
+  } else if (PyComplex_CheckExact(input.ptr())) {
+    return InferredType(ComplexDoubleType::get());
   } else if (py::isinstance<py::str>(input)) {
     return InferredType(StringType::get());
   } else if (THPLayout_Check(input.ptr())) {

torch/csrc/jit/serialization/pickler.cpp

@@ -49,6 +49,8 @@ void Pickler::pushIValueImpl(const IValue& ivalue) {
     pushTuple(ivalue);
   } else if (ivalue.isDouble()) {
     pushDouble(ivalue.toDouble());
+  } else if (ivalue.isComplexDouble()) {
+    pushComplexDouble(ivalue);
   } else if (ivalue.isInt()) {
     pushInt(ivalue.toInt());
   } else if (ivalue.isBool()) {
@@ -464,6 +466,14 @@ void Pickler::pushDouble(double value) {
   // Python pickle format is big endian, swap.
   push<double>(swapDouble(value));
 }
+void Pickler::pushComplexDouble(const IValue& value) {
+  c10::complex<double> d = value.toComplexDouble();
+  pushGlobal("builtins", "complex");
+  pushIValue(d.real());
+  pushIValue(d.imag());
+  push<PickleOpCode>(PickleOpCode::TUPLE2);
+  push<PickleOpCode>(PickleOpCode::REDUCE);
+}
 
 void Pickler::pushLong(const std::string& data) {
   uint64_t size = data.size();

torch/csrc/jit/serialization/pickler.h

@@ -160,6 +160,7 @@ class TORCH_API Pickler {
   void endTypeTag(const IValue& value);
   void pushBool(bool value);
   void pushDouble(double value);
+  void pushComplexDouble(const IValue& value);
   void pushGenericList(const IValue& ivalue);
   void pushIntList(const IValue& ivalue);
   void pushList(const IValue& ivalue);

torch/csrc/jit/serialization/unpickler.cpp

@@ -57,6 +57,7 @@ void restoreAccurateTypeTags(const IValue& root, const TypePtr& type_tag) {
       case StorageType::Kind:
       case NumberType::Kind:
       case FloatType::Kind:
+      case ComplexDoubleType::Kind:
       case IntType::Kind:
       case NoneType::Kind:
       case GeneratorType::Kind:
@@ -80,9 +81,6 @@ void restoreAccurateTypeTags(const IValue& root, const TypePtr& type_tag) {
       case AnyEnumType::Kind:
         // no op, there is nothing to tag
         break;
-      // TODO(@anjali411): Implement serialization/deserialization for complex
-      // numbers
-      case ComplexDoubleType::Kind:
       case EnumType::Kind:
         // TODO(gmagogsfm): Implement serialization/deserialization of Enum.
         AT_ASSERT(false);
@@ -543,6 +541,15 @@ void Unpickler::readGlobal(
     // Unpickle a tensor
     bool quantized = class_name == "_rebuild_qtensor";
     rebuildTensor(quantized);
+  } else if (module_name == "builtins" && class_name == "complex") {
+    globals_.emplace_back([this] {
+      auto elems = pop(stack_).toTuple()->elements();
+      AT_ASSERT(elems.size() == 2);
+      auto complex =
+          c10::complex<double>(elems.at(0).toDouble(), elems.at(1).toDouble());
+      stack_.emplace_back(complex);
+    });
+
   } else if (module_name == "collections" && class_name == "OrderedDict") {
     // collections.OrderedDict is used in tensor serialization for a tensor's
     // backward hooks (but they are not actually saved with this Pickler)