dictKeys and dictItems ops on typed dicts return typed lists

smessmer · smessmer · commit 377a7283503c · 2019-07-23T14:48:29.000-07:00
Differential Revision: [D16448942](https://our.internmc.facebook.com/intern/diff/D16448942/) ghstack-source-id: 87040601 Pull Request resolved: #23270
diff --git a/aten/src/ATen/core/Dict.h b/aten/src/ATen/core/Dict.h
@@ -363,6 +363,12 @@ class Dict final {
    * having to reallocate or rehash.
    */
   void reserve(size_type count);
+
+
+  // private API for now because the return type will change to TypePtr
+  // instead of optional<TypePtr> once types are mandatory.
+  optional<TypePtr> _keyType() const;
+  optional<TypePtr> _valueType() const;
 };
 
 namespace impl {
diff --git a/aten/src/ATen/core/Dict_inl.h b/aten/src/ATen/core/Dict_inl.h
@@ -212,4 +212,20 @@ void Dict<Key, Value>::reserve(size_type count) {
   impl_->dict.reserve(count);
 }
 
+template<class Key, class Value>
+optional<TypePtr> Dict<Key, Value>::_keyType() const {
+  if (!impl_->elementTypes.has_value()) {
+    return c10::nullopt;
+  }
+  return impl_->elementTypes->keyType;
+}
+
+template<class Key, class Value>
+optional<TypePtr> Dict<Key, Value>::_valueType() const {
+  if (!impl_->elementTypes.has_value()) {
+    return c10::nullopt;
+  }
+  return impl_->elementTypes->valueType;
+}
+
 }
diff --git a/torch/csrc/jit/register_prim_ops.cpp b/torch/csrc/jit/register_prim_ops.cpp
@@ -1693,12 +1693,22 @@ int dictLen(Stack& stack) {
 
 int dictKeys(Stack& stack) {
   auto dict = pop(stack).toGenericDict();
-  auto keys = c10::impl::GenericList(c10::impl::deprecatedUntypedList());
-  keys.reserve(dict.size());
-  for (auto& item : dict) {
-    keys.push_back(item.key());
+  auto addKeysFromDict = [&] (c10::impl::GenericList& keys) {
+    keys.reserve(dict.size());
+    for (auto& item : dict) {
+      keys.push_back(item.key());
+    }
+  };
+  auto key_type = dict._keyType();
+  if (key_type.has_value()) {
+    auto keys = c10::impl::GenericList(*key_type);
+    addKeysFromDict(keys);
+    push(stack, std::move(keys));
+  } else {
+    auto keys = c10::impl::GenericList(c10::impl::deprecatedUntypedList());
+    addKeysFromDict(keys);
+    push(stack, std::move(keys));
   }
-  push(stack, IValue(keys));
   return 0;
 }
 
@@ -1851,12 +1861,23 @@ int dictUpdate(Stack& stack) {
 
 int dictItems(Stack& stack) {
   auto dict = pop(stack).toGenericDict();
-  auto items = c10::impl::GenericList(c10::impl::deprecatedUntypedList());
-  items.reserve(dict.size());
-  for (const auto& item : iterationOrder(dict)) {
-    items.emplace_back(c10::ivalue::Tuple::create({item.first, item.second}));
+  auto addItemsFromDict = [&] (c10::impl::GenericList& items) {
+    items.reserve(dict.size());
+    for (const auto& item : iterationOrder(dict)) {
+      items.emplace_back(c10::ivalue::Tuple::create({item.first, item.second}));
+    }
+  };
+  auto key_type = dict._keyType();
+  auto value_type = dict._valueType();
+  if (key_type.has_value() && value_type.has_value()) {
+    auto items = c10::impl::GenericList(TupleType::create({*key_type, *value_type}));
+    addItemsFromDict(items);
+    push(stack, std::move(items));
+  } else {
+    auto items = c10::impl::GenericList(c10::impl::deprecatedUntypedList());
+    addItemsFromDict(items);
+    push(stack, std::move(items));
   }
-  push(stack, std::move(items));
   return 0;
 }
 
