functionalization: fix x.is_contiguous(channels_last) (pytorch#94195)

Pull Request resolved: pytorch#94195
Approved by: https://github.com/ezyang

Author: bdhirsh

aten/src/ATen/FunctionalTensorWrapper.cpp

@@ -343,7 +343,7 @@ int64_t FunctionalTensorWrapper::numel_custom() const {
   return value_.unsafeGetTensorImpl()->numel();
 }
 bool FunctionalTensorWrapper::is_contiguous_custom(at::MemoryFormat memory_format) const {
-  return value_.unsafeGetTensorImpl()->is_contiguous();
+  return value_.unsafeGetTensorImpl()->is_contiguous(memory_format);
 }
 c10::SymIntArrayRef FunctionalTensorWrapper::sym_sizes_custom() const {
   return value_.unsafeGetTensorImpl()->sym_sizes();

test/test_functionalization.py

@@ -583,6 +583,21 @@ def forward(self, arg0_1):
     return diagonal_scatter
     """)
 
+    def test_channels_last_contiguous(self):
+        def f(x):
+            return x.contiguous(memory_format=torch.channels_last)
+            tmp = torch.ones(2)
+            y = x.diagonal()
+            y.add_(tmp)
+            return x
+        x = torch.randn(4, 8, 8, 3).permute(0, 3, 1, 2)
+        self.assert_functionalization(f, x)
+        logs = self.get_logs(f, x).strip()
+        # There should be no clone in the graph
+        self.assertExpectedInline(logs, """\
+def forward(self, arg0_1):
+    return arg0_1""")
+
     def test_split(self):
         def f(x):
             # test: view ops that return multiple tensors (split)