fix bugs in mean_grad (#69912)

cubehan3 · web-flow · commit 31ca1d998971 · 2024-12-04T09:41:20.000+08:00
diff --git a/paddle/fluid/primitive/decomp_rule/decomp_vjp/details.h b/paddle/fluid/primitive/decomp_rule/decomp_vjp/details.h
@@ -287,21 +287,23 @@ void mean_grad(const Tensor& x,
         axis_data.push_back(i);
       }
     }
+
+    for (int64_t& idx : axis_data) {
+      if (idx < 0) {
+        idx += x_dim.size();
+      }
+    }
+
     if (has_dynamic_shape(x_dim, axis_data)) {
       auto x_shape = shape64<T>(x);
-      factor_tensor =
-          slice<T>(x_shape, {0}, {axis_data[0]}, {axis_data[0] + 1}, {1}, {0});
-      for (size_t i = 1; i < axis_data.size(); ++i) {
-        factor_tensor =
-            factor_tensor *
-            slice<T>(
-                x_shape, {0}, {axis_data[i]}, {axis_data[i] + 1}, {1}, {0});
+      factor_tensor = full<T>({1}, 1.0, x_shape.dtype(), x_shape.place());
+      for (int64_t idx : axis_data) {
+        factor_tensor = factor_tensor * get_slice<T>(x_shape, idx);
       }
       factor_tensor = cast<T>(factor_tensor, x.dtype());
     } else {
       int64_t factor = 1;
       for (int64_t idx : axis_data) {
-        if (idx < 0) idx += x_dim.size();
         factor *= x_dim[idx];
       }
       factor_tensor =
