[Fix] adaptive max pool2d in static (#58482)

megemini · web-flow · commit f1144128ef23 · 2023-11-03T16:22:04.000+08:00
diff --git a/python/paddle/nn/functional/pooling.py b/python/paddle/nn/functional/pooling.py
@@ -1998,8 +1998,7 @@ def adaptive_max_pool2d(x, output_size, return_mask=False, name=None):
                 "adaptive": True,
             },
         )
-        # return (pool_out, mask) if return_mask else pool_out
-        return pool_out
+        return (pool_out, mask) if return_mask else pool_out
 
 
 def adaptive_max_pool3d(x, output_size, return_mask=False, name=None):
diff --git a/test/legacy_test/test_adaptive_max_pool2d.py b/test/legacy_test/test_adaptive_max_pool2d.py
@@ -159,6 +159,61 @@ def test_static_graph(self):
 
             np.testing.assert_allclose(res_5, self.res_5_np)
 
+    def test_static_graph_return_mask(self):
+        for use_cuda in (
+            [False, True] if core.is_compiled_with_cuda() else [False]
+        ):
+            place = paddle.CUDAPlace(0) if use_cuda else paddle.CPUPlace()
+            paddle.enable_static()
+            x = paddle.static.data(
+                name="x", shape=[2, 3, 7, 7], dtype="float32"
+            )
+
+            out_1 = paddle.nn.functional.adaptive_max_pool2d(
+                x=x, output_size=[3, 3], return_mask=True
+            )
+
+            out_2 = paddle.nn.functional.adaptive_max_pool2d(
+                x=x, output_size=5, return_mask=True
+            )
+
+            out_3 = paddle.nn.functional.adaptive_max_pool2d(
+                x=x, output_size=[2, 5], return_mask=True
+            )
+
+            # out_4 = paddle.nn.functional.adaptive_max_pool2d(
+            #    x=x, output_size=[3, 3], data_format="NHWC"), return_mask=True
+
+            out_5 = paddle.nn.functional.adaptive_max_pool2d(
+                x=x, output_size=[None, 3], return_mask=True
+            )
+
+            exe = paddle.static.Executor(place=place)
+            [
+                res_1,
+                mask_1,
+                res_2,
+                mask_2,
+                res_3,
+                mask_3,
+                res_5,
+                mask_5,
+            ] = exe.run(
+                base.default_main_program(),
+                feed={"x": self.x_np},
+                fetch_list=[out_1, out_2, out_3, out_5],
+            )
+
+            self.assertEqual(res_1.shape, mask_1.shape)
+
+            self.assertEqual(res_2.shape, mask_2.shape)
+
+            self.assertEqual(res_3.shape, mask_3.shape)
+
+            # self.assertEqual(res_4.shape, mask_4.shape)
+
+            self.assertEqual(res_5.shape, mask_5.shape)
+
     def test_dynamic_graph(self):
         for use_cuda in (
             [False, True] if core.is_compiled_with_cuda() else [False]

Original file line number	Diff line number	Diff line change
`@@ -1998,8 +1998,7 @@ def adaptive_max_pool2d(x, output_size, return_mask=False, name=None):`
`1998`	`1998`	`"adaptive": True,`
`1999`	`1999`	`},`
`2000`	`2000`	`)`
`2001`		`- # return (pool_out, mask) if return_mask else pool_out`
`2002`		`- return pool_out`
	`2001`	`+ return (pool_out, mask) if return_mask else pool_out`
`2003`	`2002`
`2004`	`2003`
`2005`	`2004`	`def adaptive_max_pool3d(x, output_size, return_mask=False, name=None):`