From cb814c0b961369a7ab154c58856c730cafaa2307 Mon Sep 17 00:00:00 2001
From: "Wu, Chunyuan" <chunyuan.wu@intel.com>
Date: Thu, 23 Jan 2025 01:04:47 -0500
Subject: [PATCH] Align CPU behavior with CUDA for `ConvTranspose` when
 `out_channels=0` (#142859)

Fixes https://github.com/pytorch/pytorch/issues/142466.
Remove the `weight.numel() != 0` check to align the behavior with CUDA for `ConvTranspose` when `out_channels=0`. After removing this check, the existing code is already able to give an empty output in such case.

Test plan:
```
python -u test/nn/test_convolution.py -k test_ConvTranspose_output_channels_0_cpu_float32
python -u test/nn/test_convolution.py -k test_ConvTranspose_output_channels_0_cuda_float32
```

Pull Request resolved: https://github.com/pytorch/pytorch/pull/142859
Approved by: https://github.com/mingfeima, https://github.com/malfet
---
 .../native/NaiveConvolutionTranspose2d.cpp    |  4 ++--
 .../native/NaiveConvolutionTranspose3d.cpp    |  4 ++--
 test/nn/test_convolution.py                   | 24 +++++++++++++++++++
 3 files changed, 28 insertions(+), 4 deletions(-)

diff --git a/aten/src/ATen/native/NaiveConvolutionTranspose2d.cpp b/aten/src/ATen/native/NaiveConvolutionTranspose2d.cpp
index 799b5ffa2cd..0748194d6a0 100644
--- a/aten/src/ATen/native/NaiveConvolutionTranspose2d.cpp
+++ b/aten/src/ATen/native/NaiveConvolutionTranspose2d.cpp
@@ -78,8 +78,8 @@ static inline void slow_conv_transpose2d_shape_check(
 
   if (weight.defined()) {
     TORCH_CHECK(
-        weight.numel() != 0 && (weight.dim() == 2 || weight.dim() == 4),
-        "non-empty 2D or 4D weight tensor expected, but got: ",
+        (weight.dim() == 2 || weight.dim() == 4),
+        "2D or 4D weight tensor expected, but got: ",
         weight.sizes());
     if (bias.defined()) {
       check_dim_size(bias, 1, 0, weight.size(1));
diff --git a/aten/src/ATen/native/NaiveConvolutionTranspose3d.cpp b/aten/src/ATen/native/NaiveConvolutionTranspose3d.cpp
index 773eb2542ee..b6edcfd6563 100644
--- a/aten/src/ATen/native/NaiveConvolutionTranspose3d.cpp
+++ b/aten/src/ATen/native/NaiveConvolutionTranspose3d.cpp
@@ -98,8 +98,8 @@ static inline void slow_conv_transpose3d_shape_check(
   if (weight.defined()) {
     /* TODO: TORCH_CHECK just have 2 args: condition and message */
     TORCH_CHECK(
-        weight.numel() != 0 && weight.dim() == 5,
-        "non-empty 5D (n_output_plane x n_input_plane x kernel_depth",
+        weight.dim() == 5,
+        "5D (n_output_plane x n_input_plane x kernel_depth",
         " x kernel_height x kernel_width) tensor ",
         "expected for weight, but got: ",
         weight.sizes());
diff --git a/test/nn/test_convolution.py b/test/nn/test_convolution.py
index fbe84a71300..d920b053c25 100644
--- a/test/nn/test_convolution.py
+++ b/test/nn/test_convolution.py
@@ -40,6 +40,7 @@ from torch.testing._internal.common_device_type import (
     skipCUDAIfRocmVersionLessThan,
     skipMeta,
     skipMPS,
+    skipXLA,
 )
 from torch.testing._internal.common_dtype import (
     floating_and_complex_types_and,
@@ -1749,6 +1750,29 @@ class TestConvolutionNNDeviceType(NNTestCase):
         actual = F.conv2d(x, y, padding="same", dilation=3)
         self.assertEqual(expect, actual, rtol=rtol, atol=atol)
 
+    @dtypes(torch.float)
+    # aten/src/ATen/native/mps/OperationUtils.mm: TORCH_INTERNAL_ASSERT([srcBuf length] > 0, "Placeholder tensor is empty!"); on MPS
+    @expectedFailureMPS
+    @skipXLA
+    def test_ConvTranspose_output_channels_0(self, device, dtype):
+        class Model(nn.Module):
+            def __init__(self, operator, dim):
+                super().__init__()
+                self.op = eval(
+                    f"torch.nn.{operator}{dim}d(in_channels=1, out_channels=0, kernel_size={tuple([1] * dim)})"
+                )
+
+            def forward(self, x):
+                x = self.op(x)
+                return x
+
+        for dim in [1, 2, 3]:
+            x = torch.randn([1] * (dim + 1), device=device, dtype=dtype)
+            model = Model("ConvTranspose", dim).to(device).to(dtype=dtype)
+            y = model(x)
+            self.assertEqual(y.numel(), 0)
+            self.assertEqual(x.shape[1:], y.shape[1:])
+
     @dtypes(torch.float, torch.cfloat)
     def test_conv3d_same_padding(self, device, dtype):
         if dtype is torch.cfloat: