[Prim] Modify the reduce_as op

paddle/fluid/primitive/rule/vjp/details.h

@@ -366,32 +366,28 @@ void reduce_as_grad(const Tensor& x,
   std::vector<int64_t> axis = common::vectorize<int64_t>(
       get_reduce_dims_from_out(x.dims(), target.dims()));
   int64_t axis_size = axis.size();
-  int64_t x_dim_size = x_dim.size();
-  bool reduce_all = false;
-  if (reduce_all || axis_size == 0 || axis_size == x_dim_size) {
-    reduce_all = true;
-  } else {
-    reduce_all = false;
+  if (axis_size == 0) {
+    by_pass<T>(out_grad, x_grad);
+    return;
   }
+  int64_t x_dim_size = x_dim.size();
+
   auto x_grad_tmp = Tensor();
   if (x_dim_size == 1) {
     x_grad_tmp = expand<T>(out_grad, IntArray(x_dim));
   } else {
     auto axis_ = std::vector<int64_t>();
-    if (reduce_all) {
-      for (int64_t i = 0; i < x_dim_size; i++) {
-        axis_.push_back(i);
-      }
-    } else {
-      for (int64_t i = 0; i < axis_size; i++) {
-        axis_.push_back(axis[i]);
-        if (axis[i] < 0) {
-          axis_[i] += x_dim_size;
-        }
+    for (int64_t i = 0; i < axis_size; i++) {
+      axis_.push_back(axis[i]);
+      if (axis[i] < 0) {
+        axis_[i] += x_dim_size;
       }
     }
-    auto out_grad_shape = get_unsqueeze_dims(out_grad, axis_);
-    auto out_grad_ = reshape<T>(out_grad, out_grad_shape);
+    Tensor out_grad_ = out_grad;
+    if (out_grad.shape().size() != x.shape().size()) {
+      auto out_grad_shape = get_unsqueeze_dims(out_grad, axis_);
+      out_grad_ = reshape<T>(out_grad, out_grad_shape);
+    }
     x_grad_tmp = expand<T>(out_grad_, IntArray(x_dim));
   }
 

paddle/phi/kernels/cpu/reduce_as_grad_kernel.cc

@@ -29,15 +29,13 @@ void ReduceAsGradKernel(const Context& dev_ctx,
                         const DenseTensor& out_grad,
                         DenseTensor* x_grad) {
   auto reduce_dim = phi::funcs::GetReduceDims(x, target);
-  bool reduce_all = recompute_reduce_all(x, reduce_dim);
-  ReduceGradKernel<Context, T, funcs::SumGradFunctor, true>(dev_ctx,
-                                                            x,
-                                                            paddle::none,
-                                                            out_grad,
-                                                            reduce_dim,
-                                                            false,
-                                                            reduce_all,
-                                                            x_grad);
+  if (reduce_dim.size() != 0) {
+    ReduceGradKernel<Context, T, funcs::SumGradFunctor, true>(
+        dev_ctx, x, paddle::none, out_grad, reduce_dim, false, false, x_grad);
+  } else {
+    dev_ctx.template Alloc<T>(x_grad);
+    phi::Copy(dev_ctx, out_grad, dev_ctx.GetPlace(), false, x_grad);
+  }
 }
 
 }  // namespace phi

paddle/phi/kernels/cpu/reduce_as_kernel.cc

@@ -18,7 +18,7 @@
 #include "paddle/phi/core/kernel_registry.h"
 #include "paddle/phi/kernels/cpu/reduce.h"
 #include "paddle/phi/kernels/funcs/common_shape.h"
-#include "paddle/phi/kernels/funcs/reduce_functor.h"
+#include "paddle/phi/kernels/reduce_sum_kernel.h"
 
 namespace phi {
 
@@ -28,9 +28,16 @@ void ReduceAsKernel(const Context& dev_ctx,
                     const DenseTensor& target,
                     DenseTensor* out) {
   auto reduce_dim = phi::funcs::GetReduceDims(x, target);
-  bool reduce_all = recompute_reduce_all(x, reduce_dim);
-  phi::Reduce<CPUContext, T, phi::funcs::SumFunctor>(
-      dev_ctx, x, reduce_all, reduce_dim, false, out->type(), out);
+  if (reduce_dim.size() != 0) {
+    MetaTensor meta_out(out);
+    SumInferMeta(x, reduce_dim, out->dtype(), false, &meta_out);
+    phi::SumKernel<T, Context>(
+        dev_ctx, x, reduce_dim, out->dtype(), false, out);
+    out->Resize(target.dims());
+  } else {
+    dev_ctx.template Alloc<T>(out);
+    phi::Copy(dev_ctx, x, dev_ctx.GetPlace(), false, out);
+  }
 }
 
 }  // namespace phi

paddle/phi/kernels/gpu/reduce_as_grad_kernel.cu

@@ -29,7 +29,12 @@ void ReduceAsGradKernel(const Context& dev_ctx,
                         const DenseTensor& out_grad,
                         DenseTensor* x_grad) {
   auto reduce_dim = phi::funcs::GetReduceDims(x, target);
-  bool reduce_all = recompute_reduce_all(x, reduce_dim);
+  dev_ctx.Alloc(x_grad, x.dtype());
+
+  if (reduce_dim.size() == 0) {
+    phi::Copy(dev_ctx, out_grad, dev_ctx.GetPlace(), false, x_grad);
+    return;
+  }
   auto update_dims = common::vectorize(x.dims());
   for (auto i : reduce_dim) {
     update_dims[i] = 1;
@@ -39,7 +44,6 @@ void ReduceAsGradKernel(const Context& dev_ctx,
   new_out_grad.ShareDataWith(out_grad);
   new_out_grad.Resize(common::make_ddim(update_dims));
 
-  dev_ctx.Alloc(x_grad, x.dtype());
   using MPType = typename phi::dtype::MPTypeTrait<T>::Type;
   phi::ReduceGrad<phi::kps::IdentityFunctor<T, MPType>>(
       dev_ctx,

paddle/phi/kernels/gpu/reduce_as_kernel.cu

@@ -28,7 +28,11 @@ void ReduceAsKernel(const Context& dev_ctx,
                     DenseTensor* out) {
   auto reduce_dim = phi::funcs::GetReduceDims(x, target);
   dev_ctx.template Alloc<T>(out);
-  phi::SumKernel<T, Context>(dev_ctx, x, reduce_dim, out->type(), false, out);
+  if (reduce_dim.size() != 0) {
+    phi::SumKernel<T, Context>(dev_ctx, x, reduce_dim, out->type(), false, out);
+  } else {
+    phi::Copy(dev_ctx, x, dev_ctx.GetPlace(), false, out);
+  }
 }
 
 }  // namespace phi

test/deprecated/legacy_test/test_reduce_as_op.py

@@ -78,7 +78,15 @@ def if_enable_cinn(self):
         pass
 
     def calc_output(self):
-        self.out = self.x.sum(axis=tuple(self.attrs['dim']))
+        if len(self.attrs['dim']) != 0:
+            if 1 in self.shape_y:
+                self.out = self.x.sum(
+                    axis=tuple(self.attrs['dim']), keepdims=True
+                )
+            else:
+                self.out = self.x.sum(axis=tuple(self.attrs['dim']))
+        else:
+            self.out = self.x
 
     def test_check_output(self):
         self.check_output(check_pir=True)
@@ -151,6 +159,24 @@ def init_attrs(self):
         self.attrs = {'dim': [0, 1]}
 
 
+class TestReduceAsOp14(TestReduceAsOp):
+    def init_shape(self):
+        self.shape_x = [10, 10, 6]
+        self.shape_y = [10, 10, 6]
+
+    def init_attrs(self):
+        self.attrs = {'dim': []}
+
+
+class TestReduceAsOp15(TestReduceAsOp):
+    def init_shape(self):
+        self.shape_x = [10, 10, 6, 6]
+        self.shape_y = [1, 10, 1, 1]
+
+    def init_attrs(self):
+        self.attrs = {'dim': [0, 2, 3]}
+
+
 class TestReduceAsDynamicShape(unittest.TestCase):
     def setUp(self):
         np.random.seed(2023)