Add support for complex valued keys for dict in TS (pytorch#51472)

Summary: Pull Request resolved: pytorch#51472

Test Plan: Imported from OSS

Reviewed By: ezyang

Differential Revision: D26177963

Pulled By: anjali411

fbshipit-source-id: 5841159c36b07290b1d88d4df27a0bf8c17d9df8

Author: anjali411

Diff for: aten/src/ATen/core/Dict.h

@@ -21,6 +21,7 @@ using valid_dict_key_types = guts::typelist::typelist<
   int64_t,
   std::string,
   double,
+  c10::complex<double>,
   bool,
   at::Tensor
 >;

Diff for: aten/src/ATen/core/Dict_inl.h

@@ -1,6 +1,7 @@
 #pragma once
 
 #include <ATen/core/ivalue.h>
+#include <c10/util/hash.h>
 
 namespace c10 {
 namespace detail {
@@ -43,6 +44,8 @@ inline size_t DictKeyHash::operator()(const IValue& ivalue) const {
     return std::hash<std::string>()(ivalue.toStringRef());
   } else if (ivalue.isDouble()) {
     return std::hash<double>()(ivalue.toDouble());
+  } else if (ivalue.isComplexDouble()) {
+    return c10::hash<c10::complex<double>>()(ivalue.toComplexDouble());
   } else if (ivalue.isBool()) {
     return std::hash<bool>()(ivalue.toBool());
   } else if (ivalue.isTensor()) {

Diff for: aten/src/ATen/core/jit_type.h

@@ -738,14 +738,15 @@ struct TORCH_API DictType : public Type {
       case TypeKind::IntType:
       case TypeKind::BoolType:
       case TypeKind::FloatType:
+      case TypeKind::ComplexDoubleType:
       case TypeKind::StringType:
       case TypeKind::TensorType:
         return DictTypePtr(new DictType(key, value));
       default:
         AT_ERROR(
             "Cannot create dict for key type '",
             key->str(),
-            "', only int, float, Tensor and string keys are supported");
+            "', only int, float, complex, Tensor and string keys are supported");
     }
   }
 

Diff for: aten/src/ATen/test/ivalue_test.cpp

@@ -2,6 +2,7 @@
 #include <gtest/gtest.h>
 #include <torch/torch.h>
 #include <c10/util/intrusive_ptr.h>
+#include <ATen/core/Dict.h>
 
 namespace c10 {
 
@@ -77,6 +78,18 @@ TEST(IValueTest, Basic) {
   ASSERT_EQ(complex_tuple.toTuple()->elements()[1], foo1);
 }
 
+TEST(IValueTest, ComplexDict) {
+  typedef c10::complex<double> c_type;
+  c10::Dict<c_type, c_type> m;
+  auto num1 = c_type(2.3, -3.5);
+  auto num2 = c_type(0, 5);
+  m.insert(num1, 2 * num1);
+  m.insert(num2, 2 * num2);
+  IValue dict(std::move(m));
+  auto m_ = dict.toGenericDict();
+  ASSERT_EQ(m_.at(num1), 2 * num1);
+  ASSERT_EQ(m_.at(num2), 2 * num2);
+}
 static std::array<IValue, 5> makeSampleIValues() {
   return { at::rand({3, 4}), "hello", 42, true, 1.5 };
 }

Diff for: test/jit/test_complex.py

@@ -2,7 +2,7 @@
 import os
 import sys
 from torch.testing._internal.jit_utils import JitTestCase
-from typing import List
+from typing import List, Dict
 
 # Make the helper files in test/ importable
 pytorch_test_dir = os.path.dirname(os.path.dirname(os.path.realpath(__file__)))
@@ -22,16 +22,25 @@ def fn(a: List[complex], idx: int):
         input = [1j, 2, 3 + 4j, -5, -7j]
         self.checkScript(fn, (input, 2))
 
+    def test_complexdict(self):
+        def fn(a: Dict[complex, complex], key: complex) -> complex:
+            return a[key]
+
+        input = {2 + 3j : 2 - 3j, -4.3 - 2j: 3j}
+        self.checkScript(fn, (input, -4.3 - 2j))
+
     def test_pickle(self):
         class ComplexModule(torch.jit.ScriptModule):
             def __init__(self):
                 super().__init__()
                 self.a = 3 + 5j
                 self.b = [2 + 3j, 3 + 4j, 0 - 3j, -4 + 0j]
+                self.c = {2 + 3j : 2 - 3j, -4.3 - 2j: 3j}
 
             def forward(self, b: int):
                 return b
 
         loaded = self.getExportImportCopy(ComplexModule())
         self.assertEqual(loaded.a, 3 + 5j)
         self.assertEqual(loaded.b, [2 + 3j, 3 + 4j, -3j, -4])
+        self.assertEqual(loaded.c, {2 + 3j : 2 - 3j, -4.3 - 2j: 3j})

Diff for: torch/csrc/jit/runtime/register_prim_ops.cpp

@@ -1356,6 +1356,7 @@ RegisterOperators reg_dict_ops({
     CREATE_DICT_OPS("int"),
     CREATE_DICT_OPS("bool"),
     CREATE_DICT_OPS("float"),
+    CREATE_DICT_OPS("complex"),
     CREATE_DICT_OPS("Tensor"),
 });