From 52672ea5eeee47e57d1389528c48e2aae3c159ae Mon Sep 17 00:00:00 2001
From: RedContritio <RedContritio@qq.com>
Date: Tue, 31 Jan 2023 10:52:06 +0800
Subject: [PATCH] Fix Python IndexError of case17:
 paddle.nn.functional.interpolate (#49992)

* add dimension check for interpolate

* modify dimension check for interpolate

* add unittest to size check for interpolate

* fix incorrect shape check for interpolate

* split size check and add unittests
---
 .../tests/unittests/test_bicubic_interp_op.py   |  7 +++++++
 .../unittests/test_bicubic_interp_v2_op.py      | 17 +++++++++++++++++
 python/paddle/nn/functional/common.py           | 17 +++++++++++++++++
 3 files changed, 41 insertions(+)

diff --git a/python/paddle/fluid/tests/unittests/test_bicubic_interp_op.py b/python/paddle/fluid/tests/unittests/test_bicubic_interp_op.py
index b5d1a7d0df..6482a5fddf 100644
--- a/python/paddle/fluid/tests/unittests/test_bicubic_interp_op.py
+++ b/python/paddle/fluid/tests/unittests/test_bicubic_interp_op.py
@@ -390,6 +390,12 @@ class TestBicubicOpError(unittest.TestCase):
                     x, size=[12, 12], mode='BICUBIC', align_corners=False
                 )
 
+            def test_size_shape():
+                x = fluid.data(name="x", shape=[2, 3, 6, 6], dtype="float32")
+                out = interpolate(
+                    x, size=[12], mode='BICUBIC', align_corners=False
+                )
+
             def test_align_corcers():
                 x = fluid.data(name="x", shape=[2, 3, 6, 6], dtype="float32")
                 interpolate(x, size=[12, 12], mode='BICUBIC', align_corners=3)
@@ -481,6 +487,7 @@ class TestBicubicOpError(unittest.TestCase):
 
             self.assertRaises(ValueError, test_mode_type)
             self.assertRaises(ValueError, test_input_shape)
+            self.assertRaises(ValueError, test_size_shape)
             self.assertRaises(TypeError, test_align_corcers)
             self.assertRaises(ValueError, test_attr_data_format)
             self.assertRaises(TypeError, test_actual_shape)
diff --git a/python/paddle/fluid/tests/unittests/test_bicubic_interp_v2_op.py b/python/paddle/fluid/tests/unittests/test_bicubic_interp_v2_op.py
index fed25ad18d..a52a5b3f36 100644
--- a/python/paddle/fluid/tests/unittests/test_bicubic_interp_v2_op.py
+++ b/python/paddle/fluid/tests/unittests/test_bicubic_interp_v2_op.py
@@ -610,6 +610,20 @@ class TestBicubicOpError(unittest.TestCase):
                 x, size={2, 2}, mode='bicubic', align_corners=False
             )
 
+        def test_size_length():
+            x = fluid.data(name="x", shape=[2, 3, 6, 6], dtype="float32")
+            out = interpolate(x, size=[2], mode='bicubic', align_corners=False)
+
+        def test_size_tensor_ndim():
+            x = fluid.data(name="x", shape=[2, 3, 6, 6], dtype="float32")
+            size = paddle.to_tensor(np.array([[2, 2]]))
+            out = interpolate(x, size=size, mode='bicubic', align_corners=False)
+
+        def test_size_tensor_length():
+            x = fluid.data(name="x", shape=[2, 3, 6, 6], dtype="float32")
+            size = paddle.to_tensor(np.array([2]))
+            out = interpolate(x, size=size, mode='bicubic', align_corners=False)
+
         def test_input_shape_1():
             x = fluid.data(name="x", shape=[2, 1, 0, 0], dtype="float32")
             out = interpolate(
@@ -633,6 +647,9 @@ class TestBicubicOpError(unittest.TestCase):
         self.assertRaises(ValueError, test_size_and_scale)
         self.assertRaises(ValueError, test_size_and_scale2)
         self.assertRaises(TypeError, test_size_type)
+        self.assertRaises(ValueError, test_size_length)
+        self.assertRaises(ValueError, test_size_tensor_ndim)
+        self.assertRaises(ValueError, test_size_tensor_length)
         self.assertRaises(ValueError, test_input_shape_1)
 
     def test_errors(self):
diff --git a/python/paddle/nn/functional/common.py b/python/paddle/nn/functional/common.py
index 6631682d0e..d9f5b0b160 100644
--- a/python/paddle/nn/functional/common.py
+++ b/python/paddle/nn/functional/common.py
@@ -397,6 +397,23 @@ def interpolate(
     if size is None and scale_factor is None:
         raise ValueError("One of size and scale_factor must not be None.")
 
+    if (isinstance(size, list) or isinstance(size, tuple)) and len(
+        size
+    ) != x.ndim - 2:
+        raise ValueError(
+            'The x and size should satisfy rank(x) - 2 == len(size).'
+        )
+
+    if isinstance(size, Variable):
+        if size.ndim != 1:
+            raise ValueError(
+                f"If size is a tensor, it's rank must be 1, but received {size.ndim}."
+            )
+        if size.shape[0] != x.ndim - 2:
+            raise ValueError(
+                'The x and size should satisfy rank(x) - 2 == size.shape[0].'
+            )
+
     if not isinstance(align_corners, bool):
         raise TypeError("Attr align_corners should be a bool value")
 
-- 
GitLab