PaddlePaddle
diff --git a/‎paddle/fluid/prim/api/composite_backward/composite_backward_api.h
Copy file name to clipboard
+28-4Lines changed: 28 additions & 4 deletions b/‎paddle/fluid/prim/api/composite_backward/composite_backward_api.h
Copy file name to clipboard
+28-4Lines changed: 28 additions & 4 deletions
diff --git a/‎paddle/fluid/primitive/decomp_rule/decomp_vjp/details.h
Copy file name to clipboard
+108-101Lines changed: 108 additions & 101 deletions b/‎paddle/fluid/primitive/decomp_rule/decomp_vjp/details.h
Copy file name to clipboard
+108-101Lines changed: 108 additions & 101 deletions
diff --git a/‎paddle/phi/kernels/funcs/compare_functors.h
Copy file name to clipboard
+4-4Lines changed: 4 additions & 4 deletions b/‎paddle/phi/kernels/funcs/compare_functors.h
Copy file name to clipboard
+4-4Lines changed: 4 additions & 4 deletions
diff --git a/‎paddle/phi/kernels/funcs/index_calculator.h
Copy file name to clipboard
+13-12Lines changed: 13 additions & 12 deletions b/‎paddle/phi/kernels/funcs/index_calculator.h
Copy file name to clipboard
+13-12Lines changed: 13 additions & 12 deletions
diff --git a/‎paddle/phi/kernels/funcs/norm_utils.cu.h
Copy file name to clipboard
+4-4Lines changed: 4 additions & 4 deletions b/‎paddle/phi/kernels/funcs/norm_utils.cu.h
Copy file name to clipboard
+4-4Lines changed: 4 additions & 4 deletions
@@ -1308,7 +1308,13 @@ 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);
-    x_grad_tmp = where<T>(mask, out_grad_tmp, zero_tensor);
+    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);
+    }
   } else {
     auto axis_ = std::vector<int64_t>();
     if (reduce_all) {
@@ -1329,7 +1335,13 @@ 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);
-    x_grad_tmp = where<T>(mask, out_grad_tmp, zero_tensor);
+    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);
+    }
   }
   set_output<T>(x_grad_tmp, x_grad);
 }
@@ -1361,7 +1373,13 @@ 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);
-    x_grad_tmp = where<T>(mask, out_grad_tmp, zero_tensor);
+    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);
+    }
   } else {
     auto axis_ = std::vector<int64_t>();
     if (reduce_all) {
@@ -1382,7 +1400,13 @@ 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);
-    x_grad_tmp = where<T>(mask, out_grad_tmp, zero_tensor);
+    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);
+    }
   }
   set_output<T>(x_grad_tmp, x_grad);
 }
 
@@ -1575,107 +1575,6 @@ void pad_grad(const Tensor& input,
   }
 }
 
-template <typename T>
-void max_grad(const Tensor& x,
-              const Tensor& out,
-              const Tensor& out_grad,
-              const IntArray& axis,
-              bool keepdim,
-              bool reduce_all,
-              Tensor* x_grad) {
-  if (!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);
-    }
-  }
-  set_output<T>(x_grad_tmp, x_grad);
-}
-
 template <typename T>
 void slice_grad(const Tensor& input,
                 const Tensor& out_grad,
@@ -3498,6 +3397,114 @@ void amin_grad(const Tensor& x,
   }
 }
 
+template <typename T>
+void max_grad(const Tensor& x,
+              const Tensor& out,
+              const Tensor& out_grad,
+              const IntArray& axis,
+              bool keepdim,
+              bool reduce_all,
+              Tensor* x_grad) {
+  if (!x_grad) {
+    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);
+    }
+  }
+  set_output<T>(x_grad_tmp, x_grad);
+}
+
 template <typename T>
 void p_norm_grad(const Tensor& x,
                  /*output of forward was reserved for efficient backward*/
 
@@ -42,7 +42,7 @@ struct EqualFunctor {
       if (isinf(static_cast<float>(a)) || isinf(static_cast<float>(b))) {
         return static_cast<OutT>(a == b);
       }
-      return static_cast<OutT>(fabs(static_cast<double>(a - b)) < 1e-8);
+      return static_cast<OutT>(fabs(static_cast<double>(a - b)) < 1e-15);
     } else {
       return static_cast<OutT>(a == b);
     }
@@ -65,9 +65,9 @@ struct EqualFunctor<phi::dtype::complex<T>> {
         isinf(static_cast<float>(b.imag))) {
       return static_cast<bool>(a.real == b.real && a.imag == b.imag);
     }
-    return static_cast<bool>(fabs(static_cast<double>(a.real - b.real)) <
-                                 1e-8 &&
-                             fabs(static_cast<double>(a.imag - b.imag)) < 1e-8);
+    return static_cast<bool>(
+        fabs(static_cast<double>(a.real - b.real)) < 1e-15 &&
+        fabs(static_cast<double>(a.imag - b.imag)) < 1e-15);
   }
 };
 
 
@@ -50,15 +50,16 @@ static inline Array<T, ElementCount> VectorToArray(const VectorLikeType& vec) {
   return ret;
 }
 }  // namespace details
+template <typename IndexType>
 struct IndexCalculator {
   IndexCalculator(int dim,
-                  const std::vector<int>& cal_dims,
-                  const std::vector<int>& cal_strides,
-                  const std::vector<int>& full_strides)
+                  const std::vector<int64_t>& cal_dims,
+                  const std::vector<int64_t>& cal_strides,
+                  const std::vector<int64_t>& full_strides)
       : dim(dim) {
-    dims = details::VectorToArray<int, kMaxRank>(cal_dims);
-    strides = details::VectorToArray<int, kMaxRank>(full_strides);
-    reduce_strides = details::VectorToArray<int, kMaxRank>(cal_strides);
+    dims = details::VectorToArray<int64_t, kMaxRank>(cal_dims);
+    strides = details::VectorToArray<int64_t, kMaxRank>(full_strides);
+    reduce_strides = details::VectorToArray<int64_t, kMaxRank>(cal_strides);
 #ifndef PADDLE_WITH_XPU_KP
     std::vector<kps::details::FastDivMod> cal_divmoders;
     // fast divmod
@@ -70,9 +71,9 @@ struct IndexCalculator {
 #endif
   }
 
-  __device__ inline int operator()(int offset) const {
+  __device__ inline IndexType operator()(IndexType offset) const {
 #ifdef PADDLE_WITH_XPU_KP
-    int index = 0;
+    IndexType index = 0;
 #pragma unroll
     for (int i = 0; i < kMaxRank; ++i) {
       if (i == dim) {
@@ -83,7 +84,7 @@ struct IndexCalculator {
     }
     return index;
 #else
-    int index = 0;
+    IndexType index = 0;
 #pragma unroll
     for (int i = 0; i < kMaxRank; ++i) {
       if (i == dim) {
@@ -98,9 +99,9 @@ struct IndexCalculator {
   }
 
   int dim;
-  Array<int, kMaxRank> dims;
-  Array<int, kMaxRank> strides;
-  Array<int, kMaxRank> reduce_strides;
+  Array<int64_t, kMaxRank> dims;
+  Array<int64_t, kMaxRank> strides;
+  Array<int64_t, kMaxRank> reduce_strides;
 #ifndef PADDLE_WITH_XPU_KP
   Array<kps::details::FastDivMod, kMaxRank> divmoders;
 #endif
 
@@ -753,14 +753,14 @@ void SetLaunchConfigInfoForChannelLast(const Context &ctx,
                                        dim3 *block,
                                        dim3 *grid) {
   const int MAX_GRID_SIZE = 128;
-  const int WARP_SIZE = 32;
+  const int64_t WARP_SIZE = 32;
 
   int block_x = std::min(phi::funcs::details::GetLastPow2(C), WARP_SIZE);
   int block_y = std::min(phi::funcs::details::GetLastPow2(N * H * W * D / 16),
-                         block_size / block_x);
+                         static_cast<int64_t>(block_size / block_x));
   if (block_x * block_y != block_size) {
-    block_x =
-        std::min(phi::funcs::details::GetLastPow2(C), block_size / block_y);
+    block_x = std::min(phi::funcs::details::GetLastPow2(C),
+                       static_cast<int64_t>(block_size / block_y));
   }
   int grid_x = (C + block_x - 1) / block_x;
   int grid_y = std::min((N * H * W * D + block_y * 16 - 1) / (block_y * 16),
Original file line number	Diff line number	Diff line change
`@@ -42,7 +42,7 @@ struct EqualFunctor {`
`42`	`42`	`if (isinf(static_cast<float>(a)) \|\| isinf(static_cast<float>(b))) {`
`43`	`43`	`return static_cast<OutT>(a == b);`
`44`	`44`	`}`
`45`		`- return static_cast<OutT>(fabs(static_cast<double>(a - b)) < 1e-8);`
	`45`	`+ return static_cast<OutT>(fabs(static_cast<double>(a - b)) < 1e-15);`
`46`	`46`	`} else {`
`47`	`47`	`return static_cast<OutT>(a == b);`
`48`	`48`	`}`
`@@ -65,9 +65,9 @@ struct EqualFunctor<phi::dtype::complex<T>> {`
`65`	`65`	`isinf(static_cast<float>(b.imag))) {`
`66`	`66`	`return static_cast<bool>(a.real == b.real && a.imag == b.imag);`
`67`	`67`	`}`
`68`		`- return static_cast<bool>(fabs(static_cast<double>(a.real - b.real)) <`
`69`		`- 1e-8 &&`
`70`		`- fabs(static_cast<double>(a.imag - b.imag)) < 1e-8);`
	`68`	`+ return static_cast<bool>(`
	`69`	`+ fabs(static_cast<double>(a.real - b.real)) < 1e-15 &&`
	`70`	`+ fabs(static_cast<double>(a.imag - b.imag)) < 1e-15);`
`71`	`71`	`}`
`72`	`72`	`};`
`73`	`73`