cherry-pick PaddlePaddle#38026:fix bug of reduce_max/reduce_min (Padd…

…lePaddle#38070)
MingMingShangTian · Dec 13, 2021 · 8146961 · 8146961
1 parent 01e2679
commit 8146961
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 3 deletions.
diff --git a/paddle/fluid/operators/reduce_ops/reduce_op.h b/paddle/fluid/operators/reduce_ops/reduce_op.h
@@ -47,10 +47,17 @@ inline void GetShuffledDim(const DDim& src_dims, DDim* dst_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<int> 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
@@ -86,6 +86,18 @@ def test_imperative_api(self):
         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):