Update channel shuffle to return alias instead of self as-is (#99745)

soulitzer · pytorchmergebot · commit 5ee5afb82c65 · 2023-04-24T14:02:14.000Z
Partially addresses #99655 Pull Request resolved: #99745 Approved by: https://github.com/albanD
diff --git a/aten/src/ATen/native/ChanelShuffle.cpp b/aten/src/ATen/native/ChanelShuffle.cpp
@@ -47,12 +47,12 @@ Tensor channel_shuffle(const Tensor& self, int64_t groups) {
 #if defined(C10_MOBILE) && defined(USE_XNNPACK)
   if (self.is_contiguous(MemoryFormat::ChannelsLast) &&
       xnnpack::use_channel_shuffle(self, groups)) {
-    auto output = self.numel() == 0 ? self : xnnpack::channel_shuffle(self, groups);
+    auto output = self.numel() == 0 ? self.alias() : xnnpack::channel_shuffle(self, groups);
     return output;
   }
 #endif
 
-  auto output = self.numel() == 0 ? self : at::native_channel_shuffle(self, groups);
+  auto output = self.numel() == 0 ? self.alias() : at::native_channel_shuffle(self, groups);
   return namedinference::propagate_names_if_nonempty(
       output,
       self.has_names() ? self.names() : at::ArrayRef<Dimname>{});
diff --git a/test/test_nn.py b/test/test_nn.py
@@ -6414,8 +6414,8 @@ def test_channel_shuffle(self):
         self.assertEqual(y, y_ref)
 
 
-    def test_channel_shuffle_return_self(self):
-        # gh-76616: nn.ChannelShuffle will return self with an  empty input tensor
+    def test_channel_shuffle_return_alias_of_self(self):
+        # gh-76616: nn.ChannelShuffle will return alias of self with an empty input tensor
         groups = 3
         input_tensor = torch.rand([0, 9, 4, 4])
         output = torch.nn.ChannelShuffle(groups)(input_tensor)
