【BUPT】[Paddle Tensor 第二期 API 鲁棒性增强] 修改 paddle.atan2 以减少对 CINN 的影响 (#70468)

nizne9 · web-flow · commit 6e3e729f2d21 · 2024-12-30T10:36:39.000+08:00
* Revert Atan2InferMeta and Atan2Kernel to state before PR#70253 with necessary modifications

* Fix a stupid problem
diff --git a/paddle/phi/infermeta/binary.cc b/paddle/phi/infermeta/binary.cc
@@ -195,29 +195,24 @@ void Atan2InferMeta(const MetaTensor& x, const MetaTensor& y, MetaTensor* out) {
   const auto& x_dims = x.dims();
   const auto& y_dims = y.dims();
 
-  if (x_dims == y_dims) {
-    out->share_meta(x);
-    return;
-  }
-
-  const int max_ndim = std::max(x_dims.size(), y_dims.size());
-  const int axis = std::abs(static_cast<int>(x_dims.size()) -
-                            static_cast<int>(y_dims.size()));
-
-  std::vector<int> x_dims_array(max_ndim, 1);
-  std::vector<int> y_dims_array(max_ndim, 1);
-  std::vector<int> out_dims_array(max_ndim, 1);
-
-  funcs::GetBroadcastDimsArrays(x_dims,
-                                y_dims,
-                                x_dims_array.data(),
-                                y_dims_array.data(),
-                                out_dims_array.data(),
-                                max_ndim,
-                                axis);
+  PADDLE_ENFORCE_EQ(
+      x_dims.size(),
+      y_dims.size(),
+      common::errors::InvalidArgument("The rank (%d) of X shall be same as "
+                                      "rank (%d) of Y.",
+                                      x_dims.size(),
+                                      y_dims.size()));
+
+  if (x_dims.size() > 0)
+    PADDLE_ENFORCE_LE(x_dims[0],
+                      y_dims[0],
+                      common::errors::InvalidArgument(
+                          "The count (%d) of elements of X shall not "
+                          "greater than count (%d) of elements of Y.",
+                          x_dims[0],
+                          y_dims[0]));
 
-  out->set_dims(common::make_ddim(out_dims_array));
-  out->share_lod(x);
+  out->share_meta(x);
   if (x.dtype() == DataType::INT32 || x.dtype() == DataType::INT64 ||
       y.dtype() == DataType::INT32 || y.dtype() == DataType::INT64) {
     out->set_dtype(DataType::FLOAT64);
diff --git a/paddle/phi/kernels/gpu/expand_kernel.cu b/paddle/phi/kernels/gpu/expand_kernel.cu
@@ -43,6 +43,7 @@ void ExpandKernel(const Context& ctx,
               "The expanded size (%d) for non-existing dimensions must be "
               "positive for expand_v2 op.",
               expand_shape[i]));
+      if (expand_shape[i] == 0) has_zero_dim = true;
       out_shape[i] = expand_shape[i];
     } else if (expand_shape[i] == -1) {
       out_shape[i] = vec_in_dims[i];
diff --git a/paddle/phi/kernels/impl/atan2_kernel_impl.h b/paddle/phi/kernels/impl/atan2_kernel_impl.h
@@ -76,30 +76,15 @@ void Atan2Kernel(const Context& ctx,
                  const DenseTensor& y,
                  DenseTensor* out) {
   if (x.numel() == 0 || y.numel() == 0) {
-    std::vector<int64_t> out_dims_array = common::vectorize(out->dims());
-    std::replace(out_dims_array.begin(), out_dims_array.end(), -1, 0);
-    out->Resize(common::make_ddim(out_dims_array));
     ctx.template Alloc<typename Atan2Out<T>::type>(out);
     return;
   }
 
+  const auto numel = out->numel();
   const auto* x_data = x.data<T>();
   const auto* y_data = y.data<T>();
 
-  DenseTensor x_broadcasted, y_broadcasted;
-  if (x.dims() != y.dims()) {
-    std::vector<const DenseTensor*> inputs = {&x, &y};
-    std::vector<DenseTensor*> outputs = {&x_broadcasted, &y_broadcasted};
-    for (auto* tensor : outputs) {
-      tensor->Resize(out->dims());
-    }
-    BroadcastTensorsKernel<T, Context>(ctx, inputs, outputs);
-    x_data = x_broadcasted.data<T>();
-    y_data = y_broadcasted.data<T>();
-  }
-
   auto* out_data = ctx.template Alloc<typename Atan2Out<T>::type>(out);
-  const auto numel = out->numel();
   phi::funcs::ForRange<Context> for_range(ctx, numel);
   phi::Atan2Functor<T> functor(x_data, y_data, out_data, numel);
   for_range(functor);
diff --git a/paddle/phi/kernels/xpu/expand_kernel.cc b/paddle/phi/kernels/xpu/expand_kernel.cc
@@ -41,6 +41,7 @@ void ExpandKernel(const Context& ctx,
               "The expanded size (%d) for non-existing dimensions must be "
               "positive for expand_v2 op.",
               expand_shape[i]));
+      if (expand_shape[i] == 0) has_zero_dim = true;
       final_expand_shape[i] = expand_shape[i];
     } else if (expand_shape[i] == -1) {
       final_expand_shape[i] = vec_in_dims[i];
diff --git a/python/paddle/tensor/math.py b/python/paddle/tensor/math.py
@@ -5734,8 +5734,17 @@ def atan2(x: Tensor, y: Tensor, name: str | None = None) -> Tensor:
 
     """
 
+    x_shape = list(x.shape)
+    y_shape = list(y.shape)
     if in_dynamic_or_pir_mode():
-        return _C_ops.atan2(x, y)
+        broadcast_x = x
+        broadcast_y = y
+        if x_shape != y_shape:
+            broadcast_shape = paddle.broadcast_shape(x_shape, y_shape)
+            broadcast_x = paddle.broadcast_to(broadcast_x, broadcast_shape)
+            broadcast_y = paddle.broadcast_to(broadcast_y, broadcast_shape)
+
+        return _C_ops.atan2(broadcast_x, broadcast_y)
     else:
         check_variable_and_dtype(
             x,
