[0-size Tensor No.16] Add 0-size Tensor support for paddle.atan2 API. (#72851)

DanielSun11 · web-flow · commit 5edfa61896d9 · 2025-05-23T10:55:01.000+08:00
* fix 0 size Tensor for expand kernel in onednn

* format

* expand_grad support 0 size Tensor

* fix bug and add  unittest

* modify the unittest for onednn kernel

* modify expand unittest to force enable onednn op

* fix atan2Grad

* modify unittest
diff --git a/paddle/phi/kernels/impl/atan2_grad_kernel_impl.h b/paddle/phi/kernels/impl/atan2_grad_kernel_impl.h
@@ -16,6 +16,7 @@
 
 #include "paddle/phi/core/dense_tensor.h"
 #include "paddle/phi/kernels/atan2_grad_kernel.h"
+#include "paddle/phi/kernels/full_kernel.h"
 #include "paddle/phi/kernels/funcs/for_range.h"
 
 namespace phi {
@@ -88,6 +89,20 @@ void Atan2GradKernel(const Context& ctx,
   auto y_data = y.data<T>();
   auto out_grad_data = out_grad.data<T>();
 
+  if (out_grad.numel() == 0) {
+    ctx.template Alloc<T>(x_grad);
+    ctx.template Alloc<T>(y_grad);
+    if (x_grad && x_grad->numel() != 0) {
+      phi::Full<T, Context>(
+          ctx, phi::IntArray(common::vectorize(x_grad->dims())), 0, x_grad);
+    }
+    if (y_grad && y_grad->numel() != 0) {
+      phi::Full<T, Context>(
+          ctx, phi::IntArray(common::vectorize(y_grad->dims())), 0, y_grad);
+    }
+    return;
+  }
+
   auto* x_grad_data =
       x_grad ? ctx.template Alloc<T>(x_grad, size_t(x.numel() * sizeof(T)))
              : nullptr;
diff --git a/test/legacy_test/test_atan2_op.py b/test/legacy_test/test_atan2_op.py
@@ -161,8 +161,8 @@ def setUp(self):
         self.public_python_api = paddle.atan2
         self.dtype = np.uint16
         self.check_cinn = True
-        x1 = np.random.uniform(-1, -0.1, [15, 17]).astype('float32')
-        x2 = np.random.uniform(0.1, 1, [15, 17]).astype('float32')
+        x1 = np.random.uniform(-1, -0.1, [15, 17]).astype('float64')
+        x2 = np.random.uniform(0.1, 1, [15, 17]).astype('float64')
         out = np.arctan2(x1, x2)
 
         self.inputs = {
@@ -208,8 +208,20 @@ def _test_with_shapes(self, shapes, place=None):
 
         if place is None:  # Dygraph mode
             with paddle.base.dygraph.guard():
-                tensors = [paddle.to_tensor(inp) for inp in inputs]
+                tensors = [
+                    paddle.to_tensor(inp, stop_gradient=False) for inp in inputs
+                ]
                 result = paddle.atan2(tensors[0], tensors[1])
+                loss = paddle.sum(result)
+                loss.backward()
+
+                np.testing.assert_allclose(
+                    tensors[0].shape, tensors[0].grad.shape, rtol=1e-05
+                )
+                np.testing.assert_allclose(
+                    tensors[1].shape, tensors[1].grad.shape, rtol=1e-05
+                )
+
         else:  # Static mode
             with paddle.static.program_guard(paddle.static.Program()):
                 data_tensors = [
