propagates NaNs for paddle.angle (#73239)

ooooo-create · web-flow · commit 067e42fbca95 · 2025-06-18T15:02:55.000+08:00
* fix angle

* update docstring

* fix compile

* refine

* fix tests and nan implement

* fix
diff --git a/paddle/phi/kernels/funcs/complex_functors.h b/paddle/phi/kernels/funcs/complex_functors.h
@@ -377,6 +377,25 @@ struct AngleFunctor<T, phi::funcs::NoComplex<T, dtype::Real<T>>> {
       : input_(input), output_(output), numel_(numel) {}
 
   HOSTDEVICE void operator()(int64_t idx) const {
+    if constexpr (std::is_same_v<T, phi::dtype::bfloat16> ||
+                  std::is_same_v<T, phi::dtype::float16>) {
+      if (phi::dtype::isnan(input_[idx])) {
+        output_[idx] = input_[idx];
+        return;
+      }
+    } else if constexpr (std::is_floating_point_v<T>) {
+#if defined(__CUDACC__) || defined(__HIPCC__)
+      if (::isnan(input_[idx])) {
+        output_[idx] = input_[idx];
+        return;
+      }
+#else
+      if (std::isnan(input_[idx])) {
+        output_[idx] = input_[idx];
+        return;
+      }
+#endif
+    }
     output_[idx] = input_[idx] < static_cast<T>(0) ? M_PI : 0;
   }
 
diff --git a/python/paddle/tensor/math.py b/python/paddle/tensor/math.py
@@ -6808,7 +6808,7 @@ def _diff_handler(x, n=1, axis=-1, prepend=None, append=None, name=None):
 def angle(x: Tensor, name: str | None = None) -> Tensor:
     r"""
     Element-wise angle of complex numbers. For non-negative real numbers, the angle is 0 while
-    for negative real numbers, the angle is :math:`\pi`.
+    for negative real numbers, the angle is :math:`\pi`, and NaNs are propagated..
 
     Equation:
         .. math::
diff --git a/test/legacy_test/test_angle_op.py b/test/legacy_test/test_angle_op.py
@@ -148,6 +148,7 @@ class TestAngleAPI(unittest.TestCase):
     def setUp(self):
         self.x = np.random.randn(2, 3) + 1j * np.random.randn(2, 3)
         self.out = np.angle(self.x)
+        self.dtype = "complex128"
 
     def test_dygraph(self):
         with dygraph.guard():
@@ -158,7 +159,7 @@ def test_dygraph(self):
     def test_static(self):
         mp, sp = static.Program(), static.Program()
         with static.program_guard(mp, sp):
-            x = static.data("x", shape=[2, 3], dtype="complex128")
+            x = static.data("x", shape=[2, 3], dtype=self.dtype)
             out = paddle.angle(x)
 
         exe = static.Executor()
@@ -167,6 +168,13 @@ def test_static(self):
         np.testing.assert_allclose(self.out, out_np, rtol=1e-05)
 
 
+class TestAngleAPIWithNan(TestAngleAPI):
+    def setUp(self):
+        self.x = np.array([np.nan, -1, 1], dtype=np.float64)
+        self.out = np.angle(self.x)
+        self.dtype = "float64"
+
+
 class TestZeroSize(unittest.TestCase):
     def setUp(self):
         self.x = np.random.randn(2, 0) + 1j * np.random.randn(2, 0)
