dictKeys and dictItems ops on typed dicts return typed lists (pytorch#23270)

Summary:
Pull Request resolved: pytorch#23270
ghstack-source-id: 87389530

Differential Revision: D16448942

fbshipit-source-id: e6b578f0e97776112259d7ea38e143e4716ec273

Author: smessmer

aten/src/ATen/core/Dict.h

@@ -348,6 +348,12 @@ class Dict final {
    * having to reallocate or rehash.
    */
   void reserve(size_type count) const;
+
+
+  // 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 {

aten/src/ATen/core/Dict_inl.h

@@ -189,4 +189,20 @@ void Dict<Key, Value>::reserve(size_type count) const {
   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;
+}
+
 }

torch/csrc/jit/register_prim_ops.cpp

@@ -1841,7 +1841,15 @@ int dictLen(Stack& stack) {
 
 template <unsigned int Index, typename Elem>
 c10::List<Elem> makeListForDictKeysOrValues(
+    const std::pair<c10::optional<TypePtr>, c10::optional<TypePtr>>& types,
     const std::vector<std::pair<IValue, IValue>>& order) {
+  TORCH_INTERNAL_ASSERT(
+      (!std::get<Index>(types).has_value())
+      || (*std::get<Index>(types) == getTypePtr<Elem>()),
+      "Type mismatch when trying to get a List of keys/values from Dict. ",
+      "Type in Dict is ", toString(*std::get<Index>(types)),
+      ". Type in List is ", toString(getTypePtr<Elem>()),
+      ". Index is ", c10::guts::to_string(Index));
   c10::List<Elem> values;
   values.reserve(order.size());
   for (const auto& item : order) {
@@ -1852,8 +1860,12 @@ c10::List<Elem> makeListForDictKeysOrValues(
 
 template <unsigned int Index>
 c10::impl::GenericList makeGenericListForDictKeysOrValues(
+    const std::pair<c10::optional<TypePtr>, c10::optional<TypePtr>>& types,
     const std::vector<std::pair<IValue, IValue>>& order) {
-  auto values = c10::impl::GenericList(c10::impl::deprecatedUntypedList());
+  auto type = std::get<Index>(types);
+  auto values = type.has_value()
+    ? c10::impl::GenericList(*type)
+    : c10::impl::GenericList(c10::impl::deprecatedUntypedList());
   values.reserve(order.size());
   for (const auto& item : order) {
     values.push_back(std::get<Index>(item));
@@ -1865,17 +1877,19 @@ template <unsigned int Index>
 Operation dictKeysOrValues(const Node* n) {
   auto outputType = n->output()->type()->expect<ListType>();
   return [=](Stack& stack) -> int {
-    const auto& order = iterationOrder(pop(stack).toGenericDict());
+    auto dict = pop(stack).toGenericDict();
+    const auto& order = iterationOrder(dict);
+    const auto types = std::make_pair(dict._keyType(), dict._valueType());
     if (outputType->getElementType()->isSubtypeOf(TensorType::get())) {
-      push(stack, makeListForDictKeysOrValues<Index, at::Tensor>(order));
+      push(stack, makeListForDictKeysOrValues<Index, at::Tensor>(types, order));
     } else if (outputType->getElementType() == IntType::get()) {
-      push(stack, makeListForDictKeysOrValues<Index, int64_t>(order));
+      push(stack, makeListForDictKeysOrValues<Index, int64_t>(types, order));
     } else if (outputType->getElementType() == FloatType::get()) {
-      push(stack, makeListForDictKeysOrValues<Index, double>(order));
+      push(stack, makeListForDictKeysOrValues<Index, double>(types, order));
     } else if (outputType->getElementType() == BoolType::get()) {
-      push(stack, makeListForDictKeysOrValues<Index, bool>(order));
+      push(stack, makeListForDictKeysOrValues<Index, bool>(types, order));
     } else {
-      push(stack, makeGenericListForDictKeysOrValues<Index>(order));
+      push(stack, makeGenericListForDictKeysOrValues<Index>(types, order));
     }
     return 0;
   };
@@ -1999,7 +2013,11 @@ int dictUpdate(Stack& stack) {
 
 int dictItems(Stack& stack) {
   auto dict = pop(stack).toGenericDict();
-  auto items = c10::impl::GenericList(c10::impl::deprecatedUntypedList());
+  auto key_type = dict._keyType();
+  auto value_type = dict._valueType();
+  auto items = (key_type.has_value() && value_type.has_value())
+      ? c10::impl::GenericList(TupleType::create({*key_type, *value_type}))
+      : 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}));