fix wrong divide result when meetscomplex dtype (#73331)

HydrogenSulfate · web-flow · commit f363aec252aa · 2025-06-16T11:58:54.000+08:00
diff --git a/paddle/phi/kernels/funcs/elementwise_functor.h b/paddle/phi/kernels/funcs/elementwise_functor.h
@@ -123,6 +123,91 @@ struct InverseDivideFunctor {
 template <typename T>
 using ComplexType = phi::dtype::complex<T>;
 
+// Reference: https://github.com/pytorch/pytorch/pull/92539
+template <typename T>
+struct DivideFunctor<ComplexType<T>> {
+  inline HOSTDEVICE ComplexType<T> operator()(const ComplexType<T> x,
+                                              const ComplexType<T> y) const {
+    T a = x.real;
+    T b = x.imag;
+    T c = y.real;
+    T d = y.imag;
+
+    // (a + bi) / (c + di) = (ac + bd)/(c^2 + d^2) + (bc - ad)/(c^2 + d^2) i
+    // the calculation below follows numpy's complex division
+#if defined(__GNUC__) && !defined(__clang__)
+    // std::abs is already constexpr by gcc
+    auto abs_c = std::abs(c);
+    auto abs_d = std::abs(d);
+#else
+    auto abs_c = c < 0 ? -c : c;
+    auto abs_d = d < 0 ? -d : d;
+#endif
+
+    T real_, imag_;
+    if (abs_c >= abs_d) {
+      if (abs_c == T(0) && abs_d == T(0)) {
+        /* divide by zeros should yield a complex inf or nan */
+        real_ = a / abs_c;
+        imag_ = b / abs_d;
+      } else {
+        auto rat = d / c;
+        auto scl = T(1.0) / (c + d * rat);
+        real_ = (a + b * rat) * scl;
+        imag_ = (b - a * rat) * scl;
+      }
+    } else {
+      auto rat = c / d;
+      auto scl = T(1.0) / (d + c * rat);
+      real_ = (a * rat + b) * scl;
+      imag_ = (b * rat - a) * scl;
+    }
+    return ComplexType<T>(real_, imag_);
+  }
+};
+
+template <typename T>
+struct InverseDivideFunctor<ComplexType<T>> {
+  inline HOSTDEVICE ComplexType<T> operator()(const ComplexType<T> x,
+                                              const ComplexType<T> y) const {
+    T a = y.real;
+    T b = y.imag;
+    T c = x.real;
+    T d = x.imag;
+
+    // (a + bi) / (c + di) = (ac + bd)/(c^2 + d^2) + (bc - ad)/(c^2 + d^2) i
+    // the calculation below follows numpy's complex division
+#if defined(__GNUC__) && !defined(__clang__)
+    // std::abs is already constexpr by gcc
+    auto abs_c = std::abs(c);
+    auto abs_d = std::abs(d);
+#else
+    auto abs_c = c < 0 ? -c : c;
+    auto abs_d = d < 0 ? -d : d;
+#endif
+
+    T real_, imag_;
+    if (abs_c >= abs_d) {
+      if (abs_c == T(0) && abs_d == T(0)) {
+        /* divide by zeros should yield a complex inf or nan */
+        real_ = a / abs_c;
+        imag_ = b / abs_d;
+      } else {
+        auto rat = d / c;
+        auto scl = T(1.0) / (c + d * rat);
+        real_ = (a + b * rat) * scl;
+        imag_ = (b - a * rat) * scl;
+      }
+    } else {
+      auto rat = c / d;
+      auto scl = T(1.0) / (d + c * rat);
+      real_ = (a * rat + b) * scl;
+      imag_ = (b * rat - a) * scl;
+    }
+    return ComplexType<T>(real_, imag_);
+  }
+};
+
 template <typename InT, typename OutT>
 struct DivGradXYFunctor {
   inline HOSTDEVICE phi::Array<OutT, 2> operator()(const InT a,
diff --git a/test/legacy_test/test_elementwise_div_op.py b/test/legacy_test/test_elementwise_div_op.py
@@ -12,10 +12,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import itertools
 import unittest
 
 import numpy as np
 from op_test import OpTest, convert_float_to_uint16, skip_check_grad_ci
+from utils import dygraph_guard
 
 import paddle
 import paddle.static
@@ -782,6 +784,63 @@ def init_data(self):
         self.y_numpy = np.random.rand(3, 0, 1).astype('float32')
 
 
+class TestDivComplexDtype(unittest.TestCase):
+    def test(self):
+        with dygraph_guard():
+            places = ['cpu']
+            if core.is_compiled_with_cuda():
+                places.append('gpu')
+            shapes = [[], [1], [1, 1]]
+            values = [
+                -paddle.inf,
+                paddle.inf,
+                paddle.nan,
+                -np.zeros([]),
+                +np.zeros([]),
+                paddle.nan,
+                -paddle.nan,
+                1e-23,
+                -1e-23,
+            ]
+            dtypes = ["float32", "float64", "complex64", "complex128"]
+
+            for place in places:
+                with base.device_guard(place):
+                    for (
+                        shape_x,
+                        shape_y,
+                        x,
+                        y,
+                        dtype_x,
+                        dtype_y,
+                    ) in itertools.product(
+                        shapes, shapes, values, values, dtypes, dtypes
+                    ):
+                        pd_x = paddle.to_tensor(x, dtype=dtype_x).reshape(
+                            shape_x
+                        )
+                        pd_y = paddle.to_tensor(y, dtype=dtype_y).reshape(
+                            shape_y
+                        )
+                        pd_z = paddle.divide(pd_x, pd_y)
+
+                        np_x = np.asarray(x, dtype=dtype_x).reshape(shape_x)
+                        np_y = np.asarray(y, dtype=dtype_y).reshape(shape_y)
+                        np_z = np.divide(np_x, np_y)
+
+                        err_msg = (
+                            f"\n❌ Mismatch detected!\n"
+                            f"Place: {place}\n"
+                            f"x={x}, y={y}, dtype_x={dtype_x}, dtype_y={dtype_y}\n"
+                            f"Shape_x: {shape_x}, Shape_y: {shape_y}\n"
+                            f"np_x={np_x.item()}, np_y={np_y.item()}, np_z={np_z.item()}\n"
+                            f"pd_x={pd_x.item()}, pd_y={pd_y.item()}, pd_z={pd_z.item()}"
+                        )
+                        np.testing.assert_allclose(
+                            pd_z.item(), np_z, 0.0, 0.0, err_msg=err_msg
+                        )
+
+
 if __name__ == '__main__':
     paddle.enable_static()
     unittest.main()
