polish

PaddlePaddle · Apr 26, 2022 · 4054f7c · 4054f7c
1 parent 5c3b6bb
commit 4054f7c
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 19 deletions.
diff --git a/paddle/phi/kernels/cpu/cum_kernel.cc b/paddle/phi/kernels/cpu/cum_kernel.cc
@@ -48,14 +48,14 @@ void ComputeImp(Device d,
   }
 }
 
-template <typename T, typename Context, typename Op>
+template <typename T, typename Context, typename Reducer>
 void ScanKernel(const Context& dev_ctx,
                 const DenseTensor& x,
                 int axis,
                 bool flatten,
                 bool exclusive,
                 bool reverse,
-                Op op,
+                Reducer reducer,
                 DenseTensor* out) {
   auto out_dims = out->dims();
 
@@ -98,7 +98,7 @@ void ScanKernel(const Context& dev_ctx,
                  /* axis= */ 0,
                  reverse,
                  exclusive,
-                 op);
+                 reducer);
     } else {
       ComputeImp(place,
                  Eigen::DSizes<IndexT, 2>(mid, post),
@@ -107,7 +107,7 @@ void ScanKernel(const Context& dev_ctx,
                  /* axis= */ 0,
                  reverse,
                  exclusive,
-                 op);
+                 reducer);
     }
   } else {
     if (post == 1) {
@@ -118,7 +118,7 @@ void ScanKernel(const Context& dev_ctx,
                  /* axis= */ 1,
                  reverse,
                  exclusive,
-                 op);
+                 reducer);
     } else {
       ComputeImp(place,
                  Eigen::DSizes<IndexT, 3>(pre, mid, post),
@@ -127,7 +127,7 @@ void ScanKernel(const Context& dev_ctx,
                  /* axis= */ 1,
                  reverse,
                  exclusive,
-                 op);
+                 reducer);
     }
   }
 }
@@ -146,6 +146,7 @@ void CumsumKernel(const Context& dev_ctx,
       dev_ctx, x, axis, flatten, exclusive, reverse, reducer, out);
 }
 
+// Copied from https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/kernels/scan_ops.h
 template <typename T>
 struct LogSumExp {
   EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE T operator()(const T& a,

diff --git a/paddle/phi/kernels/gpu/cum_kernel.cu b/paddle/phi/kernels/gpu/cum_kernel.cu
@@ -141,14 +141,6 @@ struct LogAddExp {
   }
 };
 
-struct Prod {
-  template <typename T>
-  __host__ __device__ __forceinline__ T operator()(const T& a,
-                                                   const T& b) const {
-    return a * b;
-  }
-};
-
 template <typename T, typename op>
 struct Identity;
 
@@ -157,11 +149,6 @@ struct Identity<T, cub::Sum> {
   static constexpr T value = 0;
 };
 
-template <typename T>
-struct Identity<T, Prod> {
-  static constexpr T value = 1;
-};
-
 template <typename T>
 struct Identity<T, LogAddExp> {
   static constexpr T value = std::numeric_limits<T>::lowest();

diff --git a/paddle/phi/kernels/impl/logcumsumexp_grad_impl.h b/paddle/phi/kernels/impl/logcumsumexp_grad_impl.h
@@ -47,6 +47,7 @@ void LogcumsumexpGradKernel(const Context& dev_ctx,
                             bool exclusive,
                             bool reverse,
                             DenseTensor* d_x) {
+  // Reference: https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/ops/math_grad.py
   reverse = !reverse;
   dev_ctx.template Alloc<T>(d_x);