From 512e43396bdcfce1290631f131ce479b0ca64715 Mon Sep 17 00:00:00 2001
From: Shang Zhizhou <shangzhizhou@baidu.com>
Date: Mon, 13 Dec 2021 10:42:00 +0800
Subject: [PATCH] fix reduce_max bug (#38026)

* fix reduce_max bug

* add unittest
---
 paddle/pten/kernels/hybird/eigen/reduce.h          | 14 +++++++++++---
 python/paddle/fluid/tests/unittests/test_max_op.py | 12 ++++++++++++
 2 files changed, 23 insertions(+), 3 deletions(-)

diff --git a/paddle/pten/kernels/hybird/eigen/reduce.h b/paddle/pten/kernels/hybird/eigen/reduce.h
index 52ea1e68e12..e6ab872928c 100644
--- a/paddle/pten/kernels/hybird/eigen/reduce.h
+++ b/paddle/pten/kernels/hybird/eigen/reduce.h
@@ -82,10 +82,18 @@ inline void GetShuffledDim(const DDim& src_dims,
   std::vector<bool> src_dims_check(src_dims.size(), false);
   size_t src_size = src_dims.size();
   size_t reduce_size = reduced_dims.size();
+  std::vector<int64_t> regular_reduced_dims = reduced_dims;
+  for (size_t i = 0; i < regular_reduced_dims.size(); i++) {
+    if (regular_reduced_dims[i] < 0) {
+      regular_reduced_dims[i] = src_size + regular_reduced_dims[i];
+    }
+  }
+
   for (size_t i = 0; i < reduce_size; ++i) {
-    dst_dims->at(src_size - reduce_size + i) = src_dims[reduced_dims[i]];
-    (*perm_axis)[src_size - reduce_size + i] = reduced_dims[i];
-    src_dims_check[reduced_dims[i]] = true;
+    dst_dims->at(src_size - reduce_size + i) =
+        src_dims[regular_reduced_dims[i]];
+    (*perm_axis)[src_size - reduce_size + i] = regular_reduced_dims[i];
+    src_dims_check[regular_reduced_dims[i]] = true;
   }
 
   size_t offset = 0;
diff --git a/python/paddle/fluid/tests/unittests/test_max_op.py b/python/paddle/fluid/tests/unittests/test_max_op.py
index 3a1dbc8f95f..caee7d9e5c2 100644
--- a/python/paddle/fluid/tests/unittests/test_max_op.py
+++ b/python/paddle/fluid/tests/unittests/test_max_op.py
@@ -86,6 +86,18 @@ class ApiMaxTest(unittest.TestCase):
         z_expected = np.array(np.max(np_x, axis=0))
         self.assertEqual((np_z == z_expected).all(), True)
 
+    def test_big_dimension(self):
+        paddle.disable_static()
+        x = paddle.rand(shape=[2, 2, 2, 2, 2, 2, 2])
+        np_x = x.numpy()
+        z1 = paddle.max(x, axis=-1)
+        z2 = paddle.max(x, axis=6)
+        np_z1 = z1.numpy()
+        np_z2 = z2.numpy()
+        z_expected = np.array(np.max(np_x, axis=6))
+        self.assertEqual((np_z1 == z_expected).all(), True)
+        self.assertEqual((np_z2 == z_expected).all(), True)
+
 
 class TestOutDtype(unittest.TestCase):
     def test_max(self):
-- 
GitLab