【Paddle Tensor 规范化第二期】为isnan, isfinite, isinf添加复数类型支持 (#70935)

fangfangssj · web-flow · commit da56dd3e000c · 2025-02-13T17:00:45.000+08:00
* add complex

* fix DCU

* fix win
diff --git a/paddle/phi/kernels/cpu/isfinite_kernel.cc b/paddle/phi/kernels/cpu/isfinite_kernel.cc
@@ -15,6 +15,7 @@
 #include "paddle/phi/kernels/isfinite_kernel.h"
 
 #include "paddle/phi/backends/cpu/cpu_context.h"
+#include "paddle/phi/common/complex.h"
 #include "paddle/phi/core/kernel_registry.h"
 #include "paddle/phi/kernels/impl/isfinite_kernel_impl.h"
 
@@ -30,7 +31,9 @@ PD_REGISTER_KERNEL(isinf,
                    int64_t,
                    int16_t,
                    int8_t,
-                   uint8_t) {
+                   uint8_t,
+                   phi::dtype::complex<float>,
+                   phi::dtype::complex<double>) {
   kernel->OutputAt(0).SetDataType(phi::DataType::BOOL);
 }
 
@@ -43,7 +46,9 @@ PD_REGISTER_KERNEL(isnan,
                    phi::dtype::float16,
                    phi::dtype::bfloat16,
                    int,
-                   int64_t) {
+                   int64_t,
+                   phi::dtype::complex<float>,
+                   phi::dtype::complex<double>) {
   kernel->OutputAt(0).SetDataType(phi::DataType::BOOL);
 }
 
@@ -56,6 +61,8 @@ PD_REGISTER_KERNEL(isfinite,
                    phi::dtype::float16,
                    phi::dtype::bfloat16,
                    int,
-                   int64_t) {
+                   int64_t,
+                   phi::dtype::complex<float>,
+                   phi::dtype::complex<double>) {
   kernel->OutputAt(0).SetDataType(phi::DataType::BOOL);
 }
diff --git a/paddle/phi/kernels/gpu/isfinite_kernel.cu b/paddle/phi/kernels/gpu/isfinite_kernel.cu
@@ -15,6 +15,7 @@
 #include "paddle/phi/kernels/isfinite_kernel.h"
 
 #include "paddle/phi/backends/gpu/gpu_context.h"
+#include "paddle/phi/common/complex.h"
 #include "paddle/phi/core/kernel_registry.h"
 #include "paddle/phi/kernels/impl/isfinite_kernel_impl.h"
 
@@ -30,7 +31,9 @@ PD_REGISTER_KERNEL(isinf,
                    int64_t,
                    int16_t,
                    int8_t,
-                   uint8_t) {
+                   uint8_t,
+                   phi::dtype::complex<float>,
+                   phi::dtype::complex<double>) {
   kernel->OutputAt(0).SetDataType(phi::DataType::BOOL);
 }
 
@@ -43,7 +46,9 @@ PD_REGISTER_KERNEL(isnan,
                    phi::dtype::float16,
                    phi::dtype::bfloat16,
                    int,
-                   int64_t) {
+                   int64_t,
+                   phi::dtype::complex<float>,
+                   phi::dtype::complex<double>) {
   kernel->OutputAt(0).SetDataType(phi::DataType::BOOL);
 }
 
@@ -56,6 +61,8 @@ PD_REGISTER_KERNEL(isfinite,
                    phi::dtype::float16,
                    phi::dtype::bfloat16,
                    int,
-                   int64_t) {
+                   int64_t,
+                   phi::dtype::complex<float>,
+                   phi::dtype::complex<double>) {
   kernel->OutputAt(0).SetDataType(phi::DataType::BOOL);
 }
diff --git a/paddle/phi/kernels/impl/isfinite_kernel_impl.h b/paddle/phi/kernels/impl/isfinite_kernel_impl.h
@@ -39,6 +39,14 @@ struct is_other_float
                              std::is_floating_point<T>::value &&
                                  !is_float_or_double<T>::value> {};
 
+// check if complex type
+template <typename T>
+struct is_complex64_or_complex128
+    : std::integral_constant<
+          bool,
+          std::is_same<T, phi::dtype::complex<float>>::value ||
+              std::is_same<T, phi::dtype::complex<double>>::value> {};
+
 namespace phi {
 using Tensor = DenseTensor;
 
@@ -49,6 +57,7 @@ Codes for isfinite/isinf/isnan as constructed as below:
 3. partial specialization for special floating-point numbers(float16/bfloat16
 and other special float),
 4. partial specialization for non-floating-point (integer) types.
+5. partial specialization for complex types.
 */
 
 /* IsfiniteFunctor */
@@ -63,7 +72,8 @@ template <typename T>
 struct IsfiniteFunctor<
     phi::CPUContext,
     T,
-    typename std::enable_if<!std::is_floating_point<T>::value>::type> {
+    typename std::enable_if<!std::is_floating_point<T>::value &&
+                            !is_complex64_or_complex128<T>::value>::type> {
   void operator()(const phi::CPUContext& ctx,
                   const DenseTensor& in,
                   DenseTensor* output) {
@@ -111,6 +121,24 @@ struct IsfiniteFunctor<
   }
 };
 
+template <typename T>
+struct IsfiniteFunctor<
+    phi::CPUContext,
+    T,
+    typename std::enable_if<is_complex64_or_complex128<T>::value>::type> {
+  void operator()(const phi::CPUContext& ctx,
+                  const DenseTensor& in,
+                  DenseTensor* output) {
+    auto* in_a = in.data<T>();
+    auto* out_data = ctx.template Alloc<bool>(output);
+    auto num = in.numel();
+    for (int i = 0; i < num; i++) {
+      const T& a = in_a[i];
+      out_data[i] = std::isfinite(a.real) && std::isfinite(a.imag);
+    }
+  }
+};
+
 /* IsnanFunctor */
 template <typename DeviceContext, typename T, typename Enable = void>
 struct IsnanFunctor {
@@ -123,7 +151,8 @@ template <typename T>
 struct IsnanFunctor<
     phi::CPUContext,
     T,
-    typename std::enable_if<!std::is_floating_point<T>::value>::type> {
+    typename std::enable_if<!std::is_floating_point<T>::value &&
+                            !is_complex64_or_complex128<T>::value>::type> {
   void operator()(const phi::CPUContext& ctx,
                   const DenseTensor& in,
                   DenseTensor* output) {
@@ -170,6 +199,24 @@ struct IsnanFunctor<phi::CPUContext,
   }
 };
 
+template <typename T>
+struct IsnanFunctor<
+    phi::CPUContext,
+    T,
+    typename std::enable_if<is_complex64_or_complex128<T>::value>::type> {
+  void operator()(const phi::CPUContext& ctx,
+                  const DenseTensor& in,
+                  DenseTensor* output) {
+    auto* in_a = in.data<T>();
+    auto* out_data = ctx.template Alloc<bool>(output);
+    auto num = in.numel();
+    for (int i = 0; i < num; i++) {
+      const T& a = in_a[i];
+      out_data[i] = std::isnan(a.real) || std::isnan(a.imag);
+    }
+  }
+};
+
 /* IsinfFunctor */
 template <typename DeviceContext, typename T, typename Enable = void>
 struct IsinfFunctor {
@@ -182,7 +229,8 @@ template <typename T>
 struct IsinfFunctor<
     phi::CPUContext,
     T,
-    typename std::enable_if<!std::is_floating_point<T>::value>::type> {
+    typename std::enable_if<!std::is_floating_point<T>::value &&
+                            !is_complex64_or_complex128<T>::value>::type> {
   void operator()(const phi::CPUContext& ctx,
                   const DenseTensor& in,
                   DenseTensor* output) {
@@ -229,6 +277,24 @@ struct IsinfFunctor<phi::CPUContext,
   }
 };
 
+template <typename T>
+struct IsinfFunctor<
+    phi::CPUContext,
+    T,
+    typename std::enable_if<is_complex64_or_complex128<T>::value>::type> {
+  void operator()(const phi::CPUContext& ctx,
+                  const DenseTensor& in,
+                  DenseTensor* output) {
+    auto* in_a = in.data<T>();
+    auto* out_data = ctx.template Alloc<bool>(output);
+    auto num = in.numel();
+    for (int i = 0; i < num; i++) {
+      const T& a = in_a[i];
+      out_data[i] = std::isinf(a.real) || std::isinf(a.imag);
+    }
+  }
+};
+
 #if defined(__NVCC__) || defined(__HIPCC__)
 /* IsfiniteFunctor */
 template <typename T>
@@ -256,6 +322,19 @@ __global__ void IsfiniteCUDAKernel(
   }
 }
 
+template <typename T>
+__global__ void IsfiniteCUDAKernel(
+    const T* in_data,
+    int num,
+    bool* out_data,
+    typename std::enable_if<is_complex64_or_complex128<T>::value>::type* = 0) {
+  unsigned int idx = threadIdx.x + blockIdx.x * blockDim.x;
+  for (int i = idx; i < num; i += blockDim.x * gridDim.x) {
+    const T& a = in_data[i];
+    out_data[i] = isfinite(a.real) && isfinite(a.imag);
+  }
+}
+
 /* IsnanFunctor */
 template <typename T>
 __global__ void IsnanCUDAKernel(
@@ -282,6 +361,19 @@ __global__ void IsnanCUDAKernel(
   }
 }
 
+template <typename T>
+__global__ void IsnanCUDAKernel(
+    const T* in_data,
+    int num,
+    bool* out_data,
+    typename std::enable_if<is_complex64_or_complex128<T>::value>::type* = 0) {
+  unsigned int idx = threadIdx.x + blockIdx.x * blockDim.x;
+  for (int i = idx; i < num; i += blockDim.x * gridDim.x) {
+    const T& a = in_data[i];
+    out_data[i] = isnan(a.real) || isnan(a.imag);
+  }
+}
+
 /* IsinfFunctor */
 template <typename T>
 __global__ void IsinfCUDAKernel(
@@ -308,6 +400,19 @@ __global__ void IsinfCUDAKernel(
   }
 }
 
+template <typename T>
+__global__ void IsinfCUDAKernel(
+    const T* in_data,
+    int num,
+    bool* out_data,
+    typename std::enable_if<is_complex64_or_complex128<T>::value>::type* = 0) {
+  unsigned int idx = threadIdx.x + blockIdx.x * blockDim.x;
+  for (int i = idx; i < num; i += blockDim.x * gridDim.x) {
+    const T& a = in_data[i];
+    out_data[i] = isinf(a.real) || isinf(a.imag);
+  }
+}
+
 template <typename T>
 struct IsfiniteFunctor<phi::GPUContext, T> {
   void operator()(const phi::GPUContext& dev_ctx,
diff --git a/python/paddle/tensor/math.py b/python/paddle/tensor/math.py
@@ -4771,7 +4771,7 @@ def isfinite(x: Tensor, name: str | None = None) -> Tensor:
     Return whether every element of input tensor is finite number or not.
 
     Args:
-        x (Tensor): The input tensor, it's data type should be float16, float32, float64, int32, int64.
+        x (Tensor): The input tensor, it's data type should be float16, float32, float64, int32, int64, complex64, complex128.
         name (str|None, optional): Name for the operation (optional, default is None). For more information, please refer to :ref:`api_guide_Name`.
 
     Returns:
@@ -4802,6 +4802,8 @@ def isfinite(x: Tensor, name: str | None = None) -> Tensor:
                 'int32',
                 'int64',
                 'uint16',
+                'complex64',
+                'complex128',
             ],
             'isfinite',
         )
@@ -4818,7 +4820,7 @@ def isinf(x: Tensor, name: str | None = None) -> Tensor:
     Return whether every element of input tensor is `+/-INF` or not.
 
     Args:
-        x (Tensor): The input tensor, it's data type should be float16, float32, float64, uint8, int8, int16, int32, int64.
+        x (Tensor): The input tensor, it's data type should be float16, float32, float64, uint8, int8, int16, int32, int64, complex64, complex128.
         name (str|None, optional): Name for the operation (optional, default is None). For more information, please refer to :ref:`api_guide_Name`.
 
     Returns:
@@ -4852,6 +4854,8 @@ def isinf(x: Tensor, name: str | None = None) -> Tensor:
                 'int64',
                 'uint8',
                 'uint16',
+                'complex64',
+                'complex128',
             ],
             'isinf',
         )
@@ -4866,7 +4870,7 @@ def isnan(x: Tensor, name: str | None = None) -> Tensor:
     Return whether every element of input tensor is `NaN` or not.
 
     Args:
-        x (Tensor): The input tensor, it's data type should be float16, float32, float64, int32, int64.
+        x (Tensor): The input tensor, it's data type should be float16, float32, float64, int32, int64, complex64, complex128.
         name (str|None, optional): Name for the operation (optional, default is None). For more information, please refer to :ref:`api_guide_Name`.
 
     Returns:
@@ -4897,6 +4901,8 @@ def isnan(x: Tensor, name: str | None = None) -> Tensor:
                 'int32',
                 'int64',
                 'uint16',
+                'complex64',
+                'complex128',
             ],
             'isnan',
         )
@@ -4916,7 +4922,7 @@ def prod(
     Compute the product of tensor elements over the given axis.
 
     Args:
-        x (Tensor): The input tensor, its data type should be bfloat16, float16, float32, float64, int32, int64.
+        x (Tensor): The input tensor, its data type should be bfloat16, float16, float32, float64, int32, int64, complex64, complex128.
         axis (int|list|tuple|None, optional): The axis along which the product is computed. If :attr:`None`,
             multiply all elements of `x` and return a Tensor with a single element,
             otherwise must be in the range :math:`[-x.ndim, x.ndim)`. If :math:`axis[i]<0`,
diff --git a/test/legacy_test/test_isfinite_v2_op.py b/test/legacy_test/test_isfinite_v2_op.py
