From a7c4facba3740a2c68f908a10231f1bcf3c9076d Mon Sep 17 00:00:00 2001
From: haosicheng <47998305+HarperCy@users.noreply.github.com>
Date: Tue, 30 Aug 2022 13:13:25 +0800
Subject: [PATCH] fix reduce mean grad bug *test=kunlun (#45511)

fix missing keep_dim variable

    fix missing grad check in unittest

    add new test case
---
 .../operators/reduce_ops/reduce_mean_op_xpu.cc    |  9 ++++++++-
 .../unittests/xpu/test_reduce_mean_op_xpu.py      | 15 +++++++++++++++
 2 files changed, 23 insertions(+), 1 deletion(-)

diff --git a/paddle/fluid/operators/reduce_ops/reduce_mean_op_xpu.cc b/paddle/fluid/operators/reduce_ops/reduce_mean_op_xpu.cc
index 108bba3b452..c86ebbc20c3 100644
--- a/paddle/fluid/operators/reduce_ops/reduce_mean_op_xpu.cc
+++ b/paddle/fluid/operators/reduce_ops/reduce_mean_op_xpu.cc
@@ -90,6 +90,7 @@ class ReduceMeanGradXPUKernel : public framework::OpKernel<T> {
 
     bool reduce_all = ctx.Attr<bool>("reduce_all");
     auto reduce_dims = ctx.Attr<std::vector<int>>("dim");
+    bool keep_dim = ctx.Attr<bool>("keep_dim");
 
     std::vector<int> xdims;
     for (int i = 0; i < input->dims().size(); i++) {
@@ -112,7 +113,13 @@ class ReduceMeanGradXPUKernel : public framework::OpKernel<T> {
         d = d + xdims.size();
       }
       reduce_numel *= xdims[d];
-      ydims.insert(ydims.begin() + d, 1);
+    }
+
+    if (keep_dim != true) {
+      sort(reduce_dims.begin(), reduce_dims.end());
+      for (auto& d : reduce_dims) {
+        ydims.insert(ydims.begin() + d, 1);
+      }
     }
 
     float val = 1.0f / static_cast<float>(reduce_numel);
diff --git a/python/paddle/fluid/tests/unittests/xpu/test_reduce_mean_op_xpu.py b/python/paddle/fluid/tests/unittests/xpu/test_reduce_mean_op_xpu.py
index 90fe474e09c..73fa3fba3c7 100644
--- a/python/paddle/fluid/tests/unittests/xpu/test_reduce_mean_op_xpu.py
+++ b/python/paddle/fluid/tests/unittests/xpu/test_reduce_mean_op_xpu.py
@@ -103,6 +103,9 @@ class XPUTestReduce(XPUOpTestWrapper):
         # def test_check_grad(self):
         #     self.check_output_with_place(self.place, ['X'], 'Out')
 
+        def test_check_grad(self):
+            self.check_grad_with_place(self.place, ['X'], 'Out')
+
     class Test2DReduce0(Test1DReduce):
 
         def setUp(self):
@@ -161,6 +164,18 @@ class XPUTestReduce(XPUOpTestWrapper):
                 'Out': self.inputs['X'].mean(axis=tuple(self.attrs['dim']))
             }
 
+    class Test6DReduce(Test1DReduce):
+
+        def setUp(self):
+            super().setUp()
+            self.attrs = {'dim': [1, -1], 'use_xpu': True}
+            self.inputs = {
+                'X': np.random.random((5, 6, 7, 8, 9, 10)).astype(self.dtype)
+            }
+            self.outputs = {
+                'Out': self.inputs['X'].mean(axis=tuple(self.attrs['dim']))
+            }
+
     class TestKeepDimReduce(Test1DReduce):
 
         def setUp(self):
-- 
GitLab