【Paddle Tensor 规范化第二期】elementwise_pow support complex (#72518)

fangfangssj · web-flow · commit dd177ffcd281 · 2025-05-12T17:26:36.000+08:00
* add complex

* add test

* fix ci

* ci test

* ci test

* ci test

* fix test

* fix test
diff --git a/paddle/phi/kernels/cpu/elementwise_grad_kernel.cc b/paddle/phi/kernels/cpu/elementwise_grad_kernel.cc
@@ -200,7 +200,9 @@ PD_REGISTER_KERNEL(elementwise_pow_grad,
                    double,
                    int,
                    int64_t,
-                   phi::dtype::bfloat16) {}
+                   phi::dtype::bfloat16,
+                   phi::dtype::complex<float>,
+                   phi::dtype::complex<double>) {}
 
 PD_REGISTER_KERNEL(copysign_grad,
                    CPU,
diff --git a/paddle/phi/kernels/cpu/elementwise_kernel.cc b/paddle/phi/kernels/cpu/elementwise_kernel.cc
@@ -156,7 +156,9 @@ PD_REGISTER_KERNEL(elementwise_pow,
                    double,
                    int,
                    int64_t,
-                   phi::dtype::bfloat16) {}
+                   phi::dtype::bfloat16,
+                   phi::dtype::complex<float>,
+                   phi::dtype::complex<double>) {}
 PD_REGISTER_KERNEL(heaviside,
                    CPU,
                    ALL_LAYOUT,
diff --git a/paddle/phi/kernels/funcs/elementwise_functor.h b/paddle/phi/kernels/funcs/elementwise_functor.h
@@ -1075,6 +1075,32 @@ struct ElementwiseInversePowFunctor {
   }
 };
 
+template <typename T>
+struct ElementwisePowFunctor<ComplexType<T>> {
+  inline HOSTDEVICE ComplexType<T> operator()(const ComplexType<T> a,
+                                              const ComplexType<T> b) const {
+#if defined(__CUDA_ARCH__) || defined(__HIPCC__)
+    return pow(a, b);
+#else
+    return std::pow(static_cast<std::complex<T>>(a),
+                    static_cast<std::complex<T>>(b));
+#endif
+  }
+};
+
+template <typename T>
+struct ElementwiseInversePowFunctor<ComplexType<T>> {
+  inline HOSTDEVICE ComplexType<T> operator()(const ComplexType<T> a,
+                                              const ComplexType<T> b) const {
+#if defined(__CUDA_ARCH__) || defined(__HIPCC__)
+    return pow(a, b);
+#else
+    return std::pow(static_cast<std::complex<T>>(b),
+                    static_cast<std::complex<T>>(a));
+#endif
+  }
+};
+
 // copysign forward and grad functors
 template <typename T>
 inline HOSTDEVICE auto copysign_func(const T& a, const T& b) {
diff --git a/paddle/phi/kernels/gpu/elementwise_grad_kernel.cu b/paddle/phi/kernels/gpu/elementwise_grad_kernel.cu
@@ -396,7 +396,9 @@ PD_REGISTER_KERNEL(elementwise_pow_grad,
                    int,
                    phi::dtype::float16,
                    phi::dtype::bfloat16,
-                   int64_t) {}
+                   int64_t,
+                   phi::dtype::complex<float>,
+                   phi::dtype::complex<double>) {}
 
 PD_REGISTER_KERNEL(add_grad,
                    GPU,
diff --git a/paddle/phi/kernels/impl/elementwise_grad_kernel_impl.h b/paddle/phi/kernels/impl/elementwise_grad_kernel_impl.h
@@ -1379,30 +1379,29 @@ void HeavisideGradKernel(const Context& dev_ctx,
 template <typename T, typename MPType>
 HOSTDEVICE typename std::enable_if<std::is_integral<T>::value, T>::type
 compute_pow_grad_dx(T x, T y, T out, T dout) {
-  return dout * y *
-         std::pow(static_cast<double>(x), static_cast<double>(y - 1));
+  return dout * y * pow(static_cast<double>(x), static_cast<double>(y - 1));
 }
 template <typename T, typename MPType>
 HOSTDEVICE typename std::enable_if<!std::is_integral<T>::value, T>::type
 compute_pow_grad_dx(T x, T y, T out, T dout) {
   MPType x_val = static_cast<MPType>(x);
   MPType y_val = static_cast<MPType>(y);
   return static_cast<T>(static_cast<MPType>(dout) * y_val *
-                        std::pow(x_val, y_val - 1));
+                        pow(x_val, y_val - 1));
 }
 template <typename T, typename MPType>
 HOSTDEVICE typename std::enable_if<std::is_integral<T>::value, T>::type
 compute_pow_grad_dy(T x, T y, T out, T dout) {
-  return dout * std::log(static_cast<double>(x)) *
-         std::pow(static_cast<double>(x), static_cast<double>(y));
+  return dout * log(static_cast<double>(x)) *
+         pow(static_cast<double>(x), static_cast<double>(y));
 }
 template <typename T, typename MPType>
 HOSTDEVICE typename std::enable_if<!std::is_integral<T>::value, T>::type
 compute_pow_grad_dy(T x, T y, T out, T dout) {
   MPType x_val = static_cast<MPType>(x);
   MPType y_val = static_cast<MPType>(y);
-  return static_cast<T>(static_cast<MPType>(dout) * std::log(x_val) *
-                        std::pow(x_val, y_val));
+  return static_cast<T>(static_cast<MPType>(dout) * log(x_val) *
+                        pow(x_val, y_val));
 }
 #else
 template <typename T, typename MPType>
@@ -1437,6 +1436,43 @@ struct PowGradDY {
   }
 };
 
+template <typename T>
+struct PowGradDX<phi::dtype::complex<T>> {
+  HOSTDEVICE phi::dtype::complex<T> operator()(
+      phi::dtype::complex<T> x,
+      phi::dtype::complex<T> y,
+      phi::dtype::complex<T> out,
+      phi::dtype::complex<T> dout) const {
+#if defined(__CUDA_ARCH__) || defined(__HIPCC__)
+    return conj(dout * y * pow(x, y - phi::dtype::complex<T>(1, 0)));
+#else
+    return conj(
+        dout * y *
+        static_cast<phi::dtype::complex<T>>(std::pow(
+            static_cast<std::complex<T>>(x),
+            static_cast<std::complex<T>>(y - phi::dtype::complex<T>(1, 0)))));
+#endif
+  }
+};
+
+template <typename T>
+struct PowGradDY<phi::dtype::complex<T>> {
+  HOSTDEVICE phi::dtype::complex<T> operator()(
+      phi::dtype::complex<T> x,
+      phi::dtype::complex<T> y,
+      phi::dtype::complex<T> out,
+      phi::dtype::complex<T> dout) const {
+#if defined(__CUDA_ARCH__) || defined(__HIPCC__)
+    return conj(dout * log(x) * pow(x, y));
+#else
+    return conj(dout * static_cast<phi::dtype::complex<T>>(
+                           std::log(static_cast<std::complex<T>>(x)) *
+                           std::pow(static_cast<std::complex<T>>(x),
+                                    static_cast<std::complex<T>>(y))));
+#endif
+  }
+};
+
 template <typename T, typename Context>
 void ElementwisePowGradKernel(const Context& dev_ctx,
                               const DenseTensor& x,
diff --git a/paddle/phi/kernels/kps/elementwise_kernel.cu b/paddle/phi/kernels/kps/elementwise_kernel.cu
@@ -252,7 +252,9 @@ PD_REGISTER_KERNEL(elementwise_pow,
                    int,
                    int64_t,
                    phi::dtype::float16,
-                   phi::dtype::bfloat16) {}
+                   phi::dtype::bfloat16,
+                   phi::dtype::complex<float>,
+                   phi::dtype::complex<double>) {}
 PD_REGISTER_KERNEL(copysign,
                    GPU,
                    ALL_LAYOUT,
diff --git a/paddle/phi/kernels/legacy/cpu/elementwise_kernel.cc b/paddle/phi/kernels/legacy/cpu/elementwise_kernel.cc
@@ -159,4 +159,6 @@ PD_REGISTER_KERNEL(elementwise_pow_raw,
                    double,
                    int,
                    int64_t,
-                   phi::dtype::bfloat16) {}
+                   phi::dtype::bfloat16,
+                   phi::dtype::complex<float>,
+                   phi::dtype::complex<double>) {}
diff --git a/paddle/phi/kernels/legacy/kps/elementwise_kernel.cu b/paddle/phi/kernels/legacy/kps/elementwise_kernel.cu
@@ -171,5 +171,7 @@ PD_REGISTER_KERNEL(elementwise_pow_raw,
                    int,
                    float16,
                    int64_t,
-                   bfloat16) {}
+                   bfloat16,
+                   complex64,
+                   complex128) {}
 #endif
diff --git a/python/paddle/tensor/math.py b/python/paddle/tensor/math.py
@@ -531,7 +531,7 @@ def pow(x: Tensor, y: float | Tensor, name: str | None = None) -> Tensor:
 
 
     Args:
-        x (Tensor): An N-D Tensor, the data type is bfloat16, float16, float32, float64, int32 or int64.
+        x (Tensor): An N-D Tensor, the data type is bfloat16, float16, float32, float64, int32, int64, complex64 or complex128.
         y (float|int|Tensor): If it is an N-D Tensor, its data type should be the same as `x`.
         name (str|None, optional): Name for the operation (optional, default is None). For more information, please refer to :ref:`api_guide_Name`.
 
@@ -638,6 +638,7 @@ def _elementwise_op(helper):
         "elementwise_mul",
         "elementwise_div",
         "elementwise_max",
+        "elementwise_pow",
     ]
     if original_op_type in bf16_and_complex_supported_ops:
         data_type = [
diff --git a/test/legacy_test/test_elementwise_pow_op.py b/test/legacy_test/test_elementwise_pow_op.py
@@ -249,6 +249,149 @@ def test_grad(self):
                 np.testing.assert_array_equal(y.gradient(), self.grad_y)
 
 
+@unittest.skipIf(
+    core.is_compiled_with_xpu(),
+    "Skip XPU for complex dtype is not fully supported",
+)
+class TestElementwisePowComplexOp(OpTest):
+    def setUp(self):
+        self.op_type = "elementwise_pow"
+        self.python_api = paddle.pow
+        self.public_python_api = paddle.pow
+        self.prim_op_type = "prim"
+        self._check_cinn = True
+
+        self.inputs = {
+            'X': np.array([1 + 2j, 3 + 4j, 5 + 6j], dtype=np.complex128),
+            'Y': np.array([2.0, 3.0, 4.0], dtype=np.float64),
+        }
+        self.outputs = {'Out': np.power(self.inputs['X'], self.inputs['Y'])}
+
+    def _get_places(self):
+        places = [base.CPUPlace()]
+        if core.is_compiled_with_cuda():
+            places.append(base.CUDAPlace(0))
+        return places
+
+    def test_check_output(self):
+
+        self.check_output(check_pir=True, check_symbol_infer=False)
+
+    def test_check_grad_normal(self):
+        self.check_grad(
+            ['X', 'Y'],
+            'Out',
+            check_pir=True,
+        )
+
+
+@unittest.skipIf(
+    core.is_compiled_with_xpu(),
+    "Skip XPU for complex dtype is not fully supported",
+)
+class TestElementwisePowComplexOp1(TestElementwisePowComplexOp):
+    def setUp(self):
+        self.op_type = "elementwise_pow"
+        self.python_api = paddle.pow
+        self.public_python_api = paddle.pow
+        self.prim_op_type = "prim"
+        self._check_cinn = True
+
+        real_part = np.random.uniform(-5, 5, size=(3, 4))
+        imag_part = np.random.uniform(-5, 5, size=(3, 4))
+        self.inputs = {
+            'X': real_part + 1j * imag_part,
+            'Y': np.random.uniform(1, 5, size=(3, 4)),
+        }
+        self.outputs = {'Out': np.power(self.inputs['X'], self.inputs['Y'])}
+
+
+@unittest.skipIf(
+    core.is_compiled_with_xpu(),
+    "Skip XPU for complex dtype is not fully supported",
+)
+class TestElementwisePowComplexOp2(TestElementwisePowComplexOp):
+    def setUp(self):
+        self.op_type = "elementwise_pow"
+        self.python_api = paddle.pow
+        self.public_python_api = paddle.pow
+        self.prim_op_type = "prim"
+        self._check_cinn = True
+
+        real_part = np.random.uniform(-5, 5, size=(20, 50))
+        imag_part = np.random.uniform(-5, 5, size=(20, 50))
+        self.inputs = {
+            'X': real_part + 1j * imag_part,
+            'Y': np.random.uniform(1, 5, size=(20, 50)),
+        }
+        self.outputs = {'Out': np.power(self.inputs['X'], self.inputs['Y'])}
+
+
+@unittest.skipIf(
+    core.is_compiled_with_xpu(),
+    "Skip XPU for complex dtype is not fully supported",
+)
+class TestElementwisePowComplexOp3(TestElementwisePowComplexOp):
+    def setUp(self):
+        self.op_type = "elementwise_pow"
+        self.python_api = paddle.pow
+        self.public_python_api = paddle.pow
+        self.prim_op_type = "prim"
+        self._check_cinn = True
+
+        real_part = np.random.uniform(-5, 5, size=(3, 5, 3))
+        imag_part = np.random.uniform(-5, 5, size=(3, 5, 3))
+        self.inputs = {
+            'X': real_part + 1j * imag_part,
+            'Y': np.random.uniform(1, 5, size=(3, 5, 3)),
+        }
+        self.outputs = {'Out': np.power(self.inputs['X'], self.inputs['Y'])}
+
+
+@unittest.skipIf(
+    core.is_compiled_with_xpu(),
+    "Skip XPU for complex dtype is not fully supported",
+)
+class TestElementwisePowComplexOp4(TestElementwisePowComplexOp):
+    def setUp(self):
+        self.op_type = "elementwise_pow"
+        self.python_api = paddle.pow
+        self.public_python_api = paddle.pow
+        self.prim_op_type = "prim"
+        self._check_cinn = True
+
+        real_part = np.random.uniform(-5, 5, size=(3, 5, 3))
+        imag_part = np.random.uniform(-5, 5, size=(3, 5, 3))
+        self.inputs = {
+            'X': real_part + 1j * imag_part,
+            'Y': real_part + 1j * imag_part,
+        }
+        self.outputs = {'Out': np.power(self.inputs['X'], self.inputs['Y'])}
+
+
+@unittest.skipIf(
+    core.is_compiled_with_xpu(),
+    "Skip XPU for complex dtype is not fully supported",
+)
+class TestElementwisePowComplexOp5(TestElementwisePowComplexOp):
+    def setUp(self):
+        self.op_type = "elementwise_pow"
+        self.python_api = paddle.pow
+        self.public_python_api = paddle.pow
+        self.prim_op_type = "prim"
+        self._check_cinn = True
+
+        x_real_part = np.random.uniform(-5, 5, size=(5, 3))
+        x_imag_part = np.random.uniform(-5, 5, size=(5, 3))
+        y_real_part = np.random.uniform(-5, 5, size=(3, 5, 3))
+        y_imag_part = np.random.uniform(-5, 5, size=(3, 5, 3))
+        self.inputs = {
+            'X': x_real_part + 1j * x_imag_part,
+            'Y': y_real_part + 1j * y_imag_part,
+        }
+        self.outputs = {'Out': np.power(self.inputs['X'], self.inputs['Y'])}
+
+
 class TestElementwisePowOpFP16(OpTest):
     def setUp(self):
         self.op_type = "elementwise_pow"
