PaddlePaddle
diff --git a/‎paddle/fluid/prim/api/composite_backward/composite_backward_api.h
Copy file name to clipboard
+12-28Lines changed: 12 additions & 28 deletions b/‎paddle/fluid/prim/api/composite_backward/composite_backward_api.h
Copy file name to clipboard
+12-28Lines changed: 12 additions & 28 deletions
diff --git a/‎paddle/fluid/primitive/decomp_rule/decomp_vjp/details.h
Copy file name to clipboard
+1-92Lines changed: 1 addition & 92 deletions b/‎paddle/fluid/primitive/decomp_rule/decomp_vjp/details.h
Copy file name to clipboard
+1-92Lines changed: 1 addition & 92 deletions
diff --git a/‎paddle/phi/kernels/cpu/reduce_max_grad_kernel.cc
Copy file name to clipboard
+18-1Lines changed: 18 additions & 1 deletion b/‎paddle/phi/kernels/cpu/reduce_max_grad_kernel.cc
Copy file name to clipboard
+18-1Lines changed: 18 additions & 1 deletion
diff --git a/‎paddle/phi/kernels/cpu/reduce_min_grad_kernel.cc
Copy file name to clipboard
+18-1Lines changed: 18 additions & 1 deletion b/‎paddle/phi/kernels/cpu/reduce_min_grad_kernel.cc
Copy file name to clipboard
+18-1Lines changed: 18 additions & 1 deletion
diff --git a/‎paddle/phi/kernels/gpu/reduce_kernel.cu
Copy file name to clipboard
+26-118Lines changed: 26 additions & 118 deletions b/‎paddle/phi/kernels/gpu/reduce_kernel.cu
Copy file name to clipboard
+26-118Lines changed: 26 additions & 118 deletions
@@ -1308,13 +1308,9 @@ void max_grad(const Tensor& x,
     auto out_grad_tmp = out_grad.expand(IntArray(x_dim));
     auto out_tmp = out.expand(IntArray(x_dim));
     auto mask = equal<T>(x, out_tmp);
-    if (axis_size == 0) {
-      auto mask_sum = sum<T>(mask, axis, x.dtype(), keepdim = true);
-      auto grad_tmp = out_grad_tmp / mask_sum;
-      x_grad_tmp = where<T>(mask, grad_tmp, zero_tensor);
-    } else {
-      x_grad_tmp = where<T>(mask, out_grad_tmp, zero_tensor);
-    }
+    auto mask_sum = sum<T>(mask, axis, x.dtype(), keepdim = true);
+    auto grad_tmp = out_grad_tmp / mask_sum;
+    x_grad_tmp = where<T>(mask, grad_tmp, zero_tensor);
   } else {
     auto axis_ = std::vector<int64_t>();
     if (reduce_all) {
@@ -1335,13 +1331,9 @@ void max_grad(const Tensor& x,
     auto out_grad_tmp = out_grad_.expand(IntArray(x_dim));
     auto out_tmp = out_.expand(IntArray(x_dim));
     auto mask = equal<T>(x, out_tmp);
-    if (axis_size == 0) {
-      auto mask_sum = sum<T>(mask, axis_, x.dtype(), keepdim = true);
-      auto grad_tmp = out_grad_tmp / mask_sum;
-      x_grad_tmp = where<T>(mask, grad_tmp, zero_tensor);
-    } else {
-      x_grad_tmp = where<T>(mask, out_grad_tmp, zero_tensor);
-    }
+    auto mask_sum = sum<T>(mask, axis_, x.dtype(), keepdim = true);
+    auto grad_tmp = out_grad_tmp / mask_sum;
+    x_grad_tmp = where<T>(mask, grad_tmp, zero_tensor);
   }
   set_output<T>(x_grad_tmp, x_grad);
 }
@@ -1373,13 +1365,9 @@ void min_grad(const Tensor& x,
     auto out_grad_tmp = out_grad.expand(IntArray(x_dim));
     auto out_tmp = out.expand(IntArray(x_dim));
     auto mask = equal<T>(x, out_tmp);
-    if (axis_size == 0) {
-      auto mask_sum = sum<T>(mask, axis, x.dtype(), keepdim = true);
-      auto grad_tmp = out_grad_tmp / mask_sum;
-      x_grad_tmp = where<T>(mask, grad_tmp, zero_tensor);
-    } else {
-      x_grad_tmp = where<T>(mask, out_grad_tmp, zero_tensor);
-    }
+    auto mask_sum = sum<T>(mask, axis, x.dtype(), keepdim = true);
+    auto grad_tmp = out_grad_tmp / mask_sum;
+    x_grad_tmp = where<T>(mask, grad_tmp, zero_tensor);
   } else {
     auto axis_ = std::vector<int64_t>();
     if (reduce_all) {
@@ -1400,13 +1388,9 @@ void min_grad(const Tensor& x,
     auto out_grad_tmp = out_grad_.expand(IntArray(x_dim));
     auto out_tmp = out_.expand(IntArray(x_dim));
     auto mask = equal<T>(x, out_tmp);
-    if (axis_size == 0) {
-      auto mask_sum = sum<T>(mask, axis_, x.dtype(), keepdim = true);
-      auto grad_tmp = out_grad_tmp / mask_sum;
-      x_grad_tmp = where<T>(mask, grad_tmp, zero_tensor);
-    } else {
-      x_grad_tmp = where<T>(mask, out_grad_tmp, zero_tensor);
-    }
+    auto mask_sum = sum<T>(mask, axis_, x.dtype(), keepdim = true);
+    auto grad_tmp = out_grad_tmp / mask_sum;
+    x_grad_tmp = where<T>(mask, grad_tmp, zero_tensor);
   }
   set_output<T>(x_grad_tmp, x_grad);
 }
 
@@ -3409,99 +3409,8 @@ void max_grad(const Tensor& x,
     return;
   }
 
-  if (axis.size() == 0) {
-    Tensor x_grad_tmp;
-    amax_grad<T>(x, out, out_grad, axis, keepdim, reduce_all, &x_grad_tmp);
-    set_output<T>(x_grad_tmp, x_grad);
-    return;
-  }
-
   Tensor x_grad_tmp;
-  if (has_dynamic_shape(x.shape())) {
-    const Tensor x_shape = shape64<T>(x);
-    const Tensor zero_tensor =
-        backend::full_with_tensor<T>(x_shape, 0.0, x.dtype(), x.place());
-    const int64_t axis_size = axis.size();
-    const int64_t x_dim_size = x.dims().size();
-
-    reduce_all = false;
-    if (reduce_all || axis_size == 0 || axis_size == x_dim_size) {
-      reduce_all = true;
-    } else {
-      reduce_all = false;
-    }
-
-    if (x_dim_size == 0 || x_dim_size == 1 || keepdim) {
-      auto out_grad_tmp = backend::expand<T>(out_grad, x_shape);
-      auto out_tmp = backend::expand<T>(out, x_shape);
-      auto mask = equal<T>(x, out_tmp);
-      x_grad_tmp = where<T>(mask, out_grad_tmp, zero_tensor);
-    } else {
-      const Tensor out_grad_shape = shape64<T>(out_grad);
-      auto axis_ = std::vector<int64_t>();
-
-      if (reduce_all) {
-        for (int64_t i = 0; i < x_dim_size; i++) {
-          axis_.push_back(i);
-        }
-      } else {
-        axis_ = axis.GetData();
-        for (int64_t i = 0; i < axis_size; i++) {
-          if (axis[i] < 0) {
-            axis_[i] = axis[i] + x_dim_size;
-          }
-        }
-      }
-      const Tensor out_grad_shape_extend =
-          get_unsqueeze_dims<T>(out_grad_shape, axis_);
-      auto out_grad_ = backend::reshape<T>(out_grad, out_grad_shape_extend);
-      auto out_ = backend::reshape<T>(out, out_grad_shape_extend);
-      auto out_grad_tmp = backend::expand<T>(out_grad_, x_shape);
-      auto out_tmp = backend::expand<T>(out_, x_shape);
-      auto mask = equal<T>(x, out_tmp);
-      x_grad_tmp = where<T>(mask, out_grad_tmp, zero_tensor);
-    }
-  } else {
-    auto zero_tensor =
-        full<T>(common::vectorize(x.dims()), 0.0, x.dtype(), x.place());
-    std::vector<int64_t> x_dim = common::vectorize<int64_t>(x.dims());
-    int64_t axis_size = axis.size();
-    int64_t x_dim_size = x_dim.size();
-    reduce_all = false;
-    if (reduce_all || axis_size == 0 || axis_size == x_dim_size) {
-      reduce_all = true;
-    } else {
-      reduce_all = false;
-    }
-
-    if (x_dim_size == 0 || x_dim_size == 1 || keepdim) {
-      auto out_grad_tmp = out_grad.expand(IntArray(x_dim));
-      auto out_tmp = out.expand(IntArray(x_dim));
-      auto mask = equal<T>(x, out_tmp);
-      x_grad_tmp = where<T>(mask, out_grad_tmp, zero_tensor);
-    } 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 {
-        axis_ = axis.GetData();
-        for (int64_t i = 0; i < axis_size; i++) {
-          if (axis[i] < 0) {
-            axis_[i] = 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);
-      auto out_ = reshape<T>(out, out_grad_shape);
-      auto out_grad_tmp = out_grad_.expand(IntArray(x_dim));
-      auto out_tmp = out_.expand(IntArray(x_dim));
-      auto mask = equal<T>(x, out_tmp);
-      x_grad_tmp = where<T>(mask, out_grad_tmp, zero_tensor);
-    }
-  }
+  amax_grad<T>(x, out, out_grad, axis, keepdim, reduce_all, &x_grad_tmp);
   set_output<T>(x_grad_tmp, x_grad);
 }
 
 
@@ -16,7 +16,24 @@
 
 #include "paddle/phi/backends/cpu/cpu_context.h"
 #include "paddle/phi/core/kernel_registry.h"
-#include "paddle/phi/kernels/impl/reduce_max_grad_kernel_impl.h"
+#include "paddle/phi/kernels/reduce_amax_grad_kernel.h"
+
+namespace phi {
+
+template <typename T, typename Context>
+void ReduceMaxGradKernel(const Context& dev_ctx,
+                         const DenseTensor& x,
+                         const DenseTensor& out,
+                         const DenseTensor& out_grad,
+                         const IntArray& dims,
+                         bool keep_dim,
+                         bool reduce_all,
+                         DenseTensor* x_grad) {
+  ReduceAMaxGradKernel<T, Context>(
+      dev_ctx, x, out, out_grad, dims.GetData(), keep_dim, reduce_all, x_grad);
+}
+
+}  // namespace phi
 
 PD_REGISTER_KERNEL(max_grad,
                    CPU,
 
@@ -16,7 +16,24 @@
 
 #include "paddle/phi/backends/cpu/cpu_context.h"
 #include "paddle/phi/core/kernel_registry.h"
-#include "paddle/phi/kernels/impl/reduce_min_grad_kernel_impl.h"
+#include "paddle/phi/kernels/reduce_amin_grad_kernel.h"
+
+namespace phi {
+
+template <typename T, typename Context>
+void ReduceMinGradKernel(const Context& dev_ctx,
+                         const DenseTensor& x,
+                         const DenseTensor& out,
+                         const DenseTensor& out_grad,
+                         const IntArray& dims,
+                         bool keep_dim,
+                         bool reduce_all,
+                         DenseTensor* x_grad) {
+  ReduceAMinGradKernel<T, Context>(
+      dev_ctx, x, out, out_grad, dims.GetData(), keep_dim, reduce_all, x_grad);
+}
+
+}  // namespace phi
 
 PD_REGISTER_KERNEL(min_grad,
                    CPU,
 
@@ -73,65 +73,6 @@ void ReduceSumGradKernel(const Context& dev_ctx,
       kps::IdentityFunctor<T, MPType>());
 }
 
-template <typename T, typename Context>
-void ReduceMinGradKernel(const Context& dev_ctx,
-                         const DenseTensor& x,
-                         const DenseTensor& out,
-                         const DenseTensor& out_grad,
-                         const IntArray& dims,
-                         bool keep_dim,
-                         bool reduce_all,
-                         DenseTensor* x_grad) {
-  if (dims.size() == 0) {
-    reduce_all = recompute_reduce_all(x, dims, reduce_all);
-    ReduceCudaAMaxAMinGrad<T, Context>(dev_ctx,
-                                       x,
-                                       out,
-                                       out_grad,
-                                       dims.GetData(),
-                                       keep_dim,
-                                       reduce_all,
-                                       x_grad);
-  } else {
-    dev_ctx.Alloc(x_grad, x.dtype());
-    reduce_all = recompute_reduce_all(x, dims, reduce_all);
-
-    // get reduce_dim
-    int dim_size = x.dims().size();
-    auto reduce_dims =
-        funcs::details::GetReduceDim(dims.GetData(), dim_size, reduce_all);
-    auto update_dims = common::vectorize(x.dims());
-    for (auto i : reduce_dims) {
-      update_dims[i] = 1;
-    }
-
-    // make new tensor of out and out_grad
-    phi::DenseTensor new_out(out.type());
-    new_out.ShareDataWith(out);
-    new_out.Resize(common::make_ddim(update_dims));
-
-    phi::DenseTensor new_out_grad(out_grad.type());
-    new_out_grad.ShareDataWith(out_grad);
-    new_out_grad.Resize(common::make_ddim(update_dims));
-
-    // make equal_out
-    phi::DenseTensor* equal_out = new phi::DenseTensor();
-    equal_out->Resize(x.dims());
-    dev_ctx.template Alloc<T>(equal_out);
-
-    // compute
-    // 1. equal_out = Equal(x, y)
-    std::vector<const phi::DenseTensor*> equal_inputs = {&new_out, &x};
-    std::vector<phi::DenseTensor*> equal_outputs = {equal_out};
-    funcs::BroadcastKernel<T>(
-        dev_ctx, equal_inputs, &equal_outputs, funcs::EqualFunctor<T>(), 0);
-
-    // 2. dx = dout * 1
-    phi::MultiplyKernel<T, Context>(dev_ctx, new_out_grad, *equal_out, x_grad);
-    delete equal_out;
-  }
-}
-
 template <typename T, typename Context>
 void ReduceMeanGradKernel(const Context& dev_ctx,
                           const DenseTensor& x,
@@ -173,65 +114,6 @@ void ReduceMeanGradKernel(const Context& dev_ctx,
       dev_ctx, inputs, &outputs, kps::DivideFunctor<T, MPType>(reduce_num), 0);
 }
 
-template <typename T, typename Context>
-void ReduceMaxGradKernel(const Context& dev_ctx,
-                         const DenseTensor& x,
-                         const DenseTensor& out,
-                         const DenseTensor& out_grad,
-                         const IntArray& dims,
-                         bool keep_dim,
-                         bool reduce_all,
-                         DenseTensor* x_grad) {
-  if (dims.size() == 0) {
-    reduce_all = recompute_reduce_all(x, dims, reduce_all);
-    ReduceCudaAMaxAMinGrad<T, Context>(dev_ctx,
-                                       x,
-                                       out,
-                                       out_grad,
-                                       dims.GetData(),
-                                       keep_dim,
-                                       reduce_all,
-                                       x_grad);
-  } else {
-    dev_ctx.Alloc(x_grad, x.dtype());
-    reduce_all = recompute_reduce_all(x, dims, reduce_all);
-
-    // get reduce_dim
-    int dim_size = x.dims().size();
-    auto reduce_dims =
-        funcs::details::GetReduceDim(dims.GetData(), dim_size, reduce_all);
-    auto update_dims = common::vectorize(x.dims());
-    for (auto i : reduce_dims) {
-      update_dims[i] = 1;
-    }
-
-    // make new tensor of out and out_grad
-    phi::DenseTensor new_out(out.type());
-    new_out.ShareDataWith(out);
-    new_out.Resize(common::make_ddim(update_dims));
-
-    phi::DenseTensor new_out_grad(out_grad.type());
-    new_out_grad.ShareDataWith(out_grad);
-    new_out_grad.Resize(common::make_ddim(update_dims));
-
-    // make equal_out
-    phi::DenseTensor* equal_out = new phi::DenseTensor();
-    equal_out->Resize(x.dims());
-    dev_ctx.template Alloc<T>(equal_out);
-
-    // compute
-    // 1. equal_out = Equal(x, y)
-    std::vector<const phi::DenseTensor*> equal_inputs = {&new_out, &x};
-    std::vector<phi::DenseTensor*> equal_outputs = {equal_out};
-    funcs::BroadcastKernel<T>(
-        dev_ctx, equal_inputs, &equal_outputs, funcs::EqualFunctor<T>(), 0);
-
-    // 2. dx = dout * 1
-    phi::MultiplyKernel<T, Context>(dev_ctx, new_out_grad, *equal_out, x_grad);
-    delete equal_out;
-  }
-}
-
 template <typename T, typename Context>
 void ReduceAMinGradKernel(const Context& dev_ctx,
                           const DenseTensor& x,
@@ -246,6 +128,19 @@ void ReduceAMinGradKernel(const Context& dev_ctx,
       dev_ctx, x, out, out_grad, dims, keep_dim, reduce_all, x_grad);
 }
 
+template <typename T, typename Context>
+void ReduceMinGradKernel(const Context& dev_ctx,
+                         const DenseTensor& x,
+                         const DenseTensor& out,
+                         const DenseTensor& out_grad,
+                         const IntArray& dims,
+                         bool keep_dim,
+                         bool reduce_all,
+                         DenseTensor* x_grad) {
+  ReduceAMinGradKernel<T, Context>(
+      dev_ctx, x, out, out_grad, dims.GetData(), keep_dim, reduce_all, x_grad);
+}
+
 template <typename T, typename Context>
 void ReduceAMaxGradKernel(const Context& dev_ctx,
                           const DenseTensor& x,
@@ -260,6 +155,19 @@ void ReduceAMaxGradKernel(const Context& dev_ctx,
       dev_ctx, x, out, out_grad, dims, keep_dim, reduce_all, x_grad);
 }
 
+template <typename T, typename Context>
+void ReduceMaxGradKernel(const Context& dev_ctx,
+                         const DenseTensor& x,
+                         const DenseTensor& out,
+                         const DenseTensor& out_grad,
+                         const IntArray& dims,
+                         bool keep_dim,
+                         bool reduce_all,
+                         DenseTensor* x_grad) {
+  ReduceAMaxGradKernel<T, Context>(
+      dev_ctx, x, out, out_grad, dims.GetData(), keep_dim, reduce_all, x_grad);
+}
+
 template <typename T, typename Context>
 void ReduceKernel(const Context& dev_ctx,
                   const DenseTensor& x,