Named inference rule for copy_

Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/23229 Test Plan: Imported from OSS Differential Revision: D16494413 Pulled By: zou3519 fbshipit-source-id: 4acb85e5a4ad09bf5f7cbb84cc8d4ceac0cd9967

Named inference rule for copy_
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/23229 Test Plan: Imported from OSS Differential Revision: D16494413 Pulled By: zou3519 fbshipit-source-id: 4acb85e5a4ad09bf5f7cbb84cc8d4ceac0cd9967
437a8b3e · Richard Zou · Facebook Github Bot · 16da355b · 437a8b3e · 437a8b3e
--- a/aten/src/ATen/native/Copy.cpp
+++ b/aten/src/ATen/native/Copy.cpp
@@ -6,6 +6,9 @@
 #include <ATen/native/TensorIterator.h>
 #include <ATen/native/quantized/Copy.h>
 #include <ATen/quantized/Quantizer.h>
+#ifdef BUILD_NAMEDTENSOR
+#include <ATen/NamedTensorUtils.h>
+#endif
 namespace {
@@ -69,6 +72,14 @@ void copy_same_type_transpose_(Tensor& self, const Tensor& src) {
      }
    }
  });
+#ifdef BUILD_NAMEDTENSOR
+  auto outnames = unify_from_right(self.names(), src.names());
+  if (outnames.has_value()) {
+    at::internal_set_names_inplace(self, *outnames);
+  } else {
+    at::internal_set_names_inplace(self, nullopt);
+  }
+#endif
 }
 // Devices directly supported by this copy implementation. Other device types

--- a/aten/src/ATen/native/native_functions.yaml
+++ b/aten/src/ATen/native/native_functions.yaml
@@ -508,6 +508,7 @@
 - func: copy_(Tensor(a!) self, Tensor src, bool non_blocking=False) -> Tensor(a!)
  variants: method
  device_guard: False
+  named_guard: False
 - func: _copy_from(Tensor self, Tensor dst, bool non_blocking=False) -> Tensor
  dispatch: {}

--- a/test/test_namedtensor.py
+++ b/test/test_namedtensor.py
@@ -74,6 +74,18 @@ class TestNamedTensor(TestCase):
    def test_empty(self):
        self._test_factory(torch.empty, 'cpu')
+    def test_copy_transpose(self):
+        # This type of copy is special-cased and therefore needs its own test
+        def _test(self_names, other_names, expected_names):
+            x = torch.empty(2, 5, names=self_names)
+            y = torch.empty(5, 2).t().set_names_(other_names)
+            x.copy_(y)
+            self.assertEqual(x.names, expected_names)
+        _test(('N', 'C'), ('N', 'C'), ('N', 'C'))
+        _test(('N', None), ('N', 'C'), ('N', 'C'))
+        _test(None, ('N', 'C'), ('N', 'C'))
    def test_set_names_(self):
        tensor = torch.empty(1, 1, names=('N', 'C'))
        self.assertEqual(tensor.set_names_(None).names, (None, None))
@@ -223,6 +235,7 @@ class TestNamedTensor(TestCase):
        tests = [
            fn_method_and_inplace('mul'),
+            method('copy_'),
        ]
        tests = flatten(tests)