From 636780b5be176c57be7049558621f64393976fa9 Mon Sep 17 00:00:00 2001
From: RedContritio <RedContritio@qq.com>
Date: Thu, 2 Feb 2023 16:16:50 +0800
Subject: [PATCH] Fix Python IndexError of case16:
 paddle.nn.functional.max_unpool2d/max_unpool3d (#49991)

* add input dimension check for max_unpool2d/max_unpool3d

* add unittest
---
 .../fluid/tests/unittests/test_unpool3d_op.py | 24 +++++++++++++++++++
 .../fluid/tests/unittests/test_unpool_op.py   | 24 +++++++++++++++++++
 python/paddle/nn/functional/pooling.py        | 18 ++++++++++++++
 3 files changed, 66 insertions(+)

diff --git a/python/paddle/fluid/tests/unittests/test_unpool3d_op.py b/python/paddle/fluid/tests/unittests/test_unpool3d_op.py
index abbabb43a51..aa4da0b7c10 100644
--- a/python/paddle/fluid/tests/unittests/test_unpool3d_op.py
+++ b/python/paddle/fluid/tests/unittests/test_unpool3d_op.py
@@ -195,6 +195,20 @@ class TestUnpool3DOpException(unittest.TestCase):
             ).astype("int32")
             F.max_unpool3d(data, indices, kernel_size=2, stride=2)
 
+        def x_rank_error():
+            data = paddle.rand(shape=[1, 1, 3, 3])
+            indices = paddle.reshape(
+                paddle.arange(0, 27), shape=[1, 1, 3, 3, 3]
+            ).astype("int32")
+            F.max_unpool3d(data, indices, kernel_size=2, stride=2)
+
+        def indices_rank_error():
+            data = paddle.rand(shape=[1, 1, 3, 3, 3])
+            indices = paddle.reshape(
+                paddle.arange(0, 27), shape=[1, 3, 3, 3]
+            ).astype("int32")
+            F.max_unpool3d(data, indices, kernel_size=2, stride=2)
+
         def indices_value_error():
             data = paddle.rand(shape=[1, 1, 3, 3, 3])
             indices = paddle.reshape(
@@ -238,6 +252,16 @@ class TestUnpool3DOpException(unittest.TestCase):
             r"The dimensions of Input\(X\) must equal to",
             indices_size_error,
         )
+        self.assertRaisesRegex(
+            ValueError,
+            r"The x should have \[N, C, D, H, W\] format",
+            x_rank_error,
+        )
+        self.assertRaisesRegex(
+            ValueError,
+            r"The indices should have \[N, C, D, H, W\] format",
+            indices_rank_error,
+        )
         if not core.is_compiled_with_cuda():
             self.assertRaisesRegex(
                 ValueError,
diff --git a/python/paddle/fluid/tests/unittests/test_unpool_op.py b/python/paddle/fluid/tests/unittests/test_unpool_op.py
index 4963f93c67e..e5eefc067e8 100644
--- a/python/paddle/fluid/tests/unittests/test_unpool_op.py
+++ b/python/paddle/fluid/tests/unittests/test_unpool_op.py
@@ -193,6 +193,20 @@ class TestUnpoolOpException(unittest.TestCase):
             ).astype("int32")
             F.max_unpool2d(data, indices, kernel_size=2, stride=2)
 
+        def x_rank_error():
+            data = paddle.rand(shape=[1, 1, 3])
+            indices = paddle.reshape(
+                paddle.arange(0, 9), shape=[1, 1, 3, 3]
+            ).astype("int32")
+            F.max_unpool2d(data, indices, kernel_size=2, stride=2)
+
+        def indices_rank_error():
+            data = paddle.rand(shape=[1, 1, 3, 3])
+            indices = paddle.reshape(
+                paddle.arange(0, 9), shape=[1, 3, 3]
+            ).astype("int32")
+            F.max_unpool2d(data, indices, kernel_size=2, stride=2)
+
         def indices_value_error():
             data = paddle.rand(shape=[1, 1, 3, 3])
             indices = paddle.reshape(
@@ -232,6 +246,16 @@ class TestUnpoolOpException(unittest.TestCase):
             r"The dimensions of Input\(X\) must equal to",
             indices_size_error,
         )
+        self.assertRaisesRegex(
+            ValueError,
+            r"The x should have \[N, C, H, W\] format",
+            x_rank_error,
+        )
+        self.assertRaisesRegex(
+            ValueError,
+            r"The indices should have \[N, C, H, W\] format",
+            indices_rank_error,
+        )
         if not core.is_compiled_with_cuda():
             self.assertRaisesRegex(
                 ValueError,
diff --git a/python/paddle/nn/functional/pooling.py b/python/paddle/nn/functional/pooling.py
index 6aab78ec116..03e7f202c53 100755
--- a/python/paddle/nn/functional/pooling.py
+++ b/python/paddle/nn/functional/pooling.py
@@ -927,6 +927,15 @@ def max_unpool2d(
             # unpool_out shape: [1, 1, 7, 7]
 
     """
+    if x.ndim != 4:
+        raise ValueError(
+            f'The x should have [N, C, H, W] format, but received {x.shape}.'
+        )
+    if indices.ndim != 4:
+        raise ValueError(
+            f'The indices should have [N, C, H, W] format, but received {indices.shape}.'
+        )
+
     kernel_size = utils.convert_to_list(kernel_size, 2, 'pool_size')
     if stride is None:
         stride = kernel_size
@@ -1061,6 +1070,15 @@ def max_unpool3d(
             # unpool_out shape: [1, 1, 4, 4, 6]
 
     """
+    if x.ndim != 5:
+        raise ValueError(
+            f'The x should have [N, C, D, H, W] format, but received {x.shape}.'
+        )
+    if indices.ndim != 5:
+        raise ValueError(
+            f'The indices should have [N, C, D, H, W] format, but received {indices.shape}.'
+        )
+
     kernel_size = utils.convert_to_list(kernel_size, 3, 'pool_size')
     if stride is None:
         stride = kernel_size
-- 
GitLab