deepCopy also copies type information of lists (#23271)

smessmer · facebook-github-bot · commit 0c0ffccbb6b3 · 2019-07-24T11:53:01.000-07:00
Summary: Pull Request resolved: #23271 ghstack-source-id: 87088503 Differential Revision: D16449220 fbshipit-source-id: 551b7cef8f6d0d2d5a56b24ddbe2e0bb2c0c3dbe
diff --git a/aten/src/ATen/core/List.h b/aten/src/ATen/core/List.h
@@ -435,6 +435,10 @@ class List final {
   // TODO Test use_count
   size_t use_count() const;
 
+  // private API for now because the return type will change to TypePtr
+  // instead of optional<TypePtr> once types are mandatory.
+  optional<TypePtr> _elementType() const;
+
 private:
   explicit List(c10::intrusive_ptr<detail::ListImpl<StorageT>>&& elements);
   friend struct IValue;
diff --git a/aten/src/ATen/core/List_inl.h b/aten/src/ATen/core/List_inl.h
@@ -289,4 +289,9 @@ size_t List<T>::use_count() const {
   return impl_.use_count();
 }
 
+template<class T>
+optional<TypePtr> List<T>::_elementType() const {
+  return impl_->elementType;
+}
+
 }
diff --git a/torch/csrc/jit/passes/utils/check_alias_annotation.cpp b/torch/csrc/jit/passes/utils/check_alias_annotation.cpp
@@ -25,9 +25,12 @@ IValue deepCopy(const IValue& self) {
 
   // Lists of ivalues should recursively deep copy their contents
   if (self.isGenericList()) {
-    auto newList = c10::impl::GenericList(c10::impl::deprecatedUntypedList());
-    newList.reserve(self.toGenericListRef().size());
-    for (const IValue& value : self.toGenericListRef()) {
+    auto source = std::move(self).toGenericList();
+    auto newList = source._elementType().has_value()
+        ? c10::impl::GenericList(*source._elementType())
+        : c10::impl::GenericList(c10::impl::deprecatedUntypedList());
+    newList.reserve(source.size());
+    for (const IValue& value : source) {
       newList.push_back(deepCopy(value));
     }
     return newList;

Original file line number	Diff line number	Diff line change
`@@ -289,4 +289,9 @@ size_t List<T>::use_count() const {`
`289`	`289`	`return impl_.use_count();`
`290`	`290`	`}`
`291`	`291`
	`292`	`+template<class T>`
	`293`	`+optional<TypePtr> List<T>::_elementType() const {`
	`294`	`+ return impl_->elementType;`
	`295`	`+}`
	`296`	`+`
`292`	`297`	`}`