add grad_output shape check for adaptive_max_pool2d_backward and adaptive_max_pool3d_backward (pytorch#141663)

Fix pytorch#141099, pytorch#141100.

Pull Request resolved: pytorch#141663
Approved by: https://github.com/mingfeima, https://github.com/malfet

Author: jiayisunx

aten/src/ATen/native/AdaptiveMaxPooling2d.cpp

@@ -61,8 +61,12 @@ TORCH_META_FUNC(adaptive_max_pool2d_backward)
 
   at::native::adaptive_pool_empty_output_check(grad_output, "adaptive_max_pool2d_backward");
 
+  TORCH_CHECK(input.ndimension() == indices.ndimension(),
+    "expected dimensions ", input.ndimension(), " for `indices` but got dimensions ", indices.ndimension());
   TORCH_CHECK(input.dtype() == grad_output.dtype(),
     "expected dtype ", input.dtype(), " for `grad_output` but got dtype ", grad_output.dtype());
+  TORCH_CHECK(indices.sizes() == grad_output.sizes(),
+    "expected sizes ", indices.sizes(), " for `grad_output` but got sizes ", grad_output.sizes());
 
   set_output_raw_strided(0, input.sizes(), {}, input.options().memory_format(input.suggest_memory_format()));
 }

aten/src/ATen/native/AdaptiveMaxPooling3d.cpp

@@ -66,7 +66,19 @@ TORCH_META_FUNC(adaptive_max_pool3d) (const Tensor& input, IntArrayRef output_si
 
 TORCH_META_FUNC(adaptive_max_pool3d_backward)
 (const Tensor& gradOutput, const Tensor& input, const Tensor& indices) {
+  int64_t ndim = gradOutput.ndimension();
+  TORCH_CHECK(ndim == 4 || ndim == 5,
+    "adaptive_max_pool3d_backward(): Expected 4D or 5D gradOutput, but got: ", gradOutput.sizes());
+
     at::native::adaptive_pool_empty_output_check(gradOutput, "adaptive_max_pool3d_backward");
+
+    TORCH_CHECK(input.ndimension() == indices.ndimension(),
+    "expected dimensions ", input.ndimension(), " for `indices` but got dimensions ", indices.ndimension());
+    TORCH_CHECK(input.dtype() == gradOutput.dtype(),
+      "expected dtype ", input.dtype(), " for `gradOutput` but got dtype ", gradOutput.dtype());
+    TORCH_CHECK(indices.sizes() == gradOutput.sizes(),
+      "expected sizes ", indices.sizes(), " for `gradOutput` but got sizes ", gradOutput.sizes());
+
     set_output_raw_strided(0, input.sizes(), {}, input.options());
 }
 } // namespace meta

test/nn/test_pooling.py

@@ -556,6 +556,20 @@ def test_adaptive_pooling_empty_output_size(self, dtype, device):
             with self.assertRaisesRegex(RuntimeError, error_msg):
                 fn(input2, output_size).sum().backward()
 
+    @onlyNativeDeviceTypes
+    def test_adaptive_pooling_backward_fails(self, device):
+        grad_output = torch.randn(1, 2, 7, 7, device=device)
+        input = torch.randn(1, 2, 7, 7, device=device)
+        indices = torch.ones(1, 2, 3, 3, dtype=torch.long, device=device)
+        with self.assertRaisesRegex(RuntimeError, "expected sizes"):
+            torch.ops.aten.adaptive_max_pool2d_backward(grad_output, input, indices)
+
+        grad_output = torch.randn(1, 2, 7, 7, 7, device=device)
+        input = torch.randn(1, 2, 3, 3, 3, device=device)
+        indices = torch.ones(1, 2, 3, 3, dtype=torch.long, device=device)
+        with self.assertRaisesRegex(RuntimeError, "expected dimensions"):
+            torch.ops.aten.adaptive_max_pool3d_backward(grad_output, input, indices)
+
     @onlyNativeDeviceTypes
     def test_FractionalMaxPool2d_zero_batch(self, device):
         mod = nn.FractionalMaxPool2d(3, output_ratio=(0.5, 0.5))