[Complex] support paddle.remainder for complex dtype - part1

paddle/phi/kernels/cpu/elementwise_kernel.cc

@@ -166,6 +166,8 @@ PD_REGISTER_KERNEL(remainder,
  float,
  double,
  int,
+ phi::dtype::complex<float>,
+ phi::dtype::complex<double>,
  int64_t) {}
 PD_REGISTER_KERNEL(floor_divide,
  CPU,

paddle/phi/kernels/funcs/elementwise_functor.h

@@ -751,6 +751,70 @@ struct RemainderFunctor<dtype::bfloat16> {
  }
 };
 
+/**
+ * Remainder for complex number rule
+ * Regarding a and b is gaussian integer, then
+ * r = mod(a, b) = a - b * round(a/b)
+ * and a, b is complex number
+ */
+template <typename T>
+struct RemainderFunctor<ComplexType<T>> {
+ inline HOSTDEVICE ComplexType<T> operator()(ComplexType<T> a,
+ ComplexType<T> b) const {
+ // remainder = z1 - q_rounded * z2
+ T a__ = a.real;
+ T b__ = a.imag;
+ T c__ = b.real;
+ T d__ = b.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;
+ }
+ auto q = ComplexType<T>(real_, imag_);
+
+#if defined(__CUDA_ARCH__) || defined(__HIPCC__)
+ const auto& q_rounded = ComplexType<T>(round(q.real), round(q.imag));
+#else
+ const auto& q_rounded =
+ ComplexType<T>(std::round(q.real), std::round(q.imag));
+#endif
+ const auto& a_ = q_rounded.real;
+ const auto& b_ = q_rounded.imag;
+ const auto& c = b.real;
+ const auto& d = b.imag;
+ const auto& t_real_ = a_ * c - b_ * d;
+ const auto& t_imag_ = a_ * d + b_ * c;
+ auto remainder = ComplexType<T>(a.real - t_real_, a.imag - t_imag_);
+ return remainder;
+ }
+};
+
 // RemainderGradXFunctor
 template <typename T>
 struct RemainderGradXFunctor {
@@ -879,6 +943,72 @@ struct InverseRemainderFunctor<
  }
 };
 
+/**
+ * Remainder for complex number rule
+ * Regarding a and b is gaussian integer, then
+ * r = mod(a, b) = a - b * round(a/b)
+ * and a, b is complex number
+ */
+template <typename T>
+struct InverseRemainderFunctor<
+ ComplexType<T>,
+ typename std::enable_if_t<std::is_floating_point<T>::value>> {
+ inline HOSTDEVICE ComplexType<T> operator()(ComplexType<T> b,
+ ComplexType<T> a) const {
+ // remainder = z1 - q_rounded * z2
+ T a__ = a.real;
+ T b__ = a.imag;
+ T c__ = b.real;
+ T d__ = b.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;
+ }
+ auto q = ComplexType<T>(real_, imag_);
+
+#if defined(__CUDA_ARCH__) || defined(__HIPCC__)
+ const auto& q_rounded = ComplexType<T>(round(q.real), round(q.imag));
+#else
+ const auto& q_rounded =
+ ComplexType<T>(std::round(q.real), std::round(q.imag));
+#endif
+ const auto& a_ = q_rounded.real;
+ const auto& b_ = q_rounded.imag;
+ const auto& c = b.real;
+ const auto& d = b.imag;
+ const auto& t_real_ = a_ * c - b_ * d;
+ const auto& t_imag_ = a_ * d + b_ * c;
+ auto remainder = ComplexType<T>(a.real - t_real_, a.imag - t_imag_);
+ return remainder;
+ }
+};
+
 template <typename T>
 struct ElementwiseHeavisideFunctor {
  inline HOSTDEVICE T operator()(const T a, const T b) const {

paddle/phi/kernels/kps/elementwise_kernel.cu

@@ -245,6 +245,8 @@ PD_REGISTER_KERNEL(remainder,
  int,
  int64_t,
  phi::dtype::float16,
+ phi::dtype::complex<float>,
+ phi::dtype::complex<double>,
  phi::dtype::bfloat16) {}
 PD_REGISTER_KERNEL(floor_divide,
  KPS,

paddle/phi/kernels/legacy/cpu/elementwise_kernel.cc

@@ -136,6 +136,8 @@ PD_REGISTER_KERNEL(remainder_raw,
  phi::RemainderRawKernel,
  float,
  double,
+ phi::dtype::complex<float>,
+ phi::dtype::complex<double>,
  int,
  int64_t) {}
 PD_REGISTER_KERNEL(floor_divide_raw,

paddle/phi/kernels/legacy/kps/elementwise_kernel.cu

@@ -147,6 +147,8 @@ PD_REGISTER_KERNEL(remainder_raw,
  double,
  int,
  float16,
+ complex64,
+ complex128,
  int64_t,
  bfloat16) {}
 PD_REGISTER_KERNEL(floor_divide_raw,

test/legacy_test/test_elementwise_mod_op.py

@@ -20,7 +20,7 @@
 from utils import dygraph_guard, static_guard
 
 import paddle
-from paddle import static
+from paddle import base, static
 from paddle.base import core
 
 
@@ -198,6 +198,50 @@ def init_dtype(self):
  self.dtype = np.float64
 
 
+class TestElementwiseModOpComplex64(unittest.TestCase):
+ def test_check_output(self):
+ with dygraph_guard():
+ dtype = "complex64"
+ a = np.array([6 + 4j]).astype(dtype)
+ b = np.array([3 + 5j]).astype(dtype)
+ res = np.array([-2 + 2j]).astype(dtype)
+
+ res_pd = paddle.remainder(paddle.to_tensor(a), paddle.to_tensor(b))
+ np.testing.assert_allclose(res, res_pd.numpy())
+
+ dtype = "complex64"
+ a = np.array([6 + 4j]).astype(dtype)
+ b = np.array([3 + 5j]).astype(dtype)
+ res = np.array([-2 + 2j]).astype(dtype)
+
+ res_pd = paddle.remainder(paddle.to_tensor(a), paddle.to_tensor(b))
+ np.testing.assert_allclose(res, res_pd.numpy())
+
+ with base.device_guard("cpu"):
+ res_pd = paddle.remainder(
+ paddle.to_tensor(a), paddle.to_tensor(b)
+ )
+ np.testing.assert_allclose(res, res_pd.numpy())
+
+
+class TestElementwiseModOpComplex128(unittest.TestCase):
+ def test_check_output(self):
+ with dygraph_guard():
+ dtype = "complex128"
+ a = np.array([6 + 4j]).astype(dtype)
+ b = np.array([3 + 5j]).astype(dtype)
+ res = np.array([-2 + 2j]).astype(dtype)
+
+ res_pd = paddle.remainder(paddle.to_tensor(a), paddle.to_tensor(b))
+ np.testing.assert_allclose(res, res_pd.numpy())
+
+ with base.device_guard("cpu"):
+ res_pd = paddle.remainder(
+ paddle.to_tensor(a), paddle.to_tensor(b)
+ )
+ np.testing.assert_allclose(res, res_pd.numpy())
+
+
 class TestElementwiseDygraph(unittest.TestCase):
  def test_dygraph_same_shape(self):
  with dygraph_guard():