【complex op】 No.10 add complex support for exp/expm1

paddle/phi/common/complex.h

@@ -466,6 +466,16 @@ HOSTDEVICE inline complex<T> tanh(const complex<T>& a) {
 #endif
 }
 
+template <typename T>
+HOSTDEVICE inline complex<T> exp(const complex<T>& a) {
+#if defined(PADDLE_WITH_CUDA_OR_HIP_COMPLEX) && \
+    (defined(__CUDA_ARCH__) || defined(__HIPCC__))
+  return complex<T>(thrust::exp(thrust::complex<T>(a)));
+#else
+  return complex<T>(std::exp(std::complex<T>(a)));
+#endif
+}
+
 template <typename T>
 HOSTDEVICE inline complex<T> conj(const complex<T>& a) {
 #if defined(PADDLE_WITH_CUDA_OR_HIP_COMPLEX) && \

paddle/phi/kernels/cpu/activation_grad_kernel.cc

@@ -340,15 +340,19 @@ PD_REGISTER_KERNEL(exp_grad,
                    float,
                    double,
                    int,
-                   int64_t) {}
+                   int64_t,
+                   phi::dtype::complex<float>,
+                   phi::dtype::complex<double>) {}
 
 PD_REGISTER_KERNEL(expm1_grad,
                    CPU,
                    ALL_LAYOUT,
                    phi::Expm1GradKernel,
                    float,
                    double,
-                   phi::dtype::float16) {}
+                   phi::dtype::float16,
+                   phi::dtype::complex<float>,
+                   phi::dtype::complex<double>) {}
 
 PD_REGISTER_KERNEL(
     logit_grad, CPU, ALL_LAYOUT, phi::LogitGradKernel, float, double) {}

paddle/phi/kernels/cpu/activation_kernel.cc

@@ -211,7 +211,9 @@ PD_REGISTER_KERNEL(exp,
                    double,
                    int,
                    int64_t,
-                   phi::dtype::float16) {}
+                   phi::dtype::float16,
+                   phi::dtype::complex<float>,
+                   phi::dtype::complex<double>) {}
 
 PD_REGISTER_KERNEL(expm1,
                    CPU,
@@ -221,7 +223,9 @@ PD_REGISTER_KERNEL(expm1,
                    double,
                    int,
                    int64_t,
-                   phi::dtype::float16) {}
+                   phi::dtype::float16,
+                   phi::dtype::complex<float>,
+                   phi::dtype::complex<double>) {}
 
 PD_REGISTER_KERNEL(logit, CPU, ALL_LAYOUT, phi::LogitKernel, float, double) {}
 PD_REGISTER_KERNEL(

paddle/phi/kernels/funcs/activation_functor.h

@@ -1167,6 +1167,30 @@ struct ExpGradFunctor : public BaseActivationFunctor<T> {
   }
 };
 
+template <typename T>
+struct ExpGradFunctor<ComplexType<T>>
+    : public BaseActivationFunctor<ComplexType<T>> {
+  template <typename Device,
+            typename X,
+            typename Out,
+            typename dOut,
+            typename dX>
+  void operator()(Device d, X x UNUSED, Out out, dOut dout, dX dx) const {
+    dx.device(d) = dout * out.unaryExpr(Conj<T>());
+  }
+
+  static constexpr ActBwdOpFwdDeps FwdDeps() {
+    return ActBwdOpFwdDeps::kDepOut;
+  }
+};
+
+template <typename T>
+struct Expm1 {
+  HOSTDEVICE T operator()(const T& val) const {
+    return exp(val) - static_cast<T>(1);
+  }
+};
+
 // expm1(x) = e^x - 1
 template <typename T>
 struct Expm1Functor : public BaseActivationFunctor<T> {
@@ -1178,6 +1202,15 @@ struct Expm1Functor : public BaseActivationFunctor<T> {
   }
 };
 
+template <typename T>
+struct Expm1Functor<ComplexType<T>>
+    : public BaseActivationFunctor<ComplexType<T>> {
+  template <typename Device, typename X, typename Out>
+  void operator()(Device d, X x, Out out) const {
+    out.device(d) = x.unaryExpr(Expm1<ComplexType<T>>()).eval();
+  }
+};
+
 template <typename T>
 struct Expm1GradFunctor : public BaseActivationFunctor<T> {
   template <typename Device,
@@ -1194,6 +1227,21 @@ struct Expm1GradFunctor : public BaseActivationFunctor<T> {
   }
 };
 
+template <typename T>
+struct Expm1GradFunctor<ComplexType<T>>
+    : public BaseActivationFunctor<ComplexType<T>> {
+  template <typename Device,
+            typename X,
+            typename Out,
+            typename dOut,
+            typename dX>
+  void operator()(Device d, X x UNUSED, Out out, dOut dout, dX dx) const {
+    dx.device(d) = dout * out.unaryExpr(Conj<T>()) + dout;
+  }
+
+  static constexpr ActBwdOpFwdDeps FwdDeps() { return kDepOut; }
+};
+
 // relu(x) = max(x, 0)
 template <typename T>
 struct ReluCPUFunctor : public BaseActivationFunctor<T> {
@@ -2790,6 +2838,16 @@ struct CudaExpFunctor<double> : public BaseActivationFunctor<double> {
   }
 };
 
+template <typename T>
+struct CudaExpFunctor<ComplexType<T>>
+    : public BaseActivationFunctor<ComplexType<T>> {
+  // exp(x) = exp(x)
+  __device__ __forceinline__ ComplexType<T> operator()(
+      const ComplexType<T> x) const {
+    return static_cast<ComplexType<T>>(exp(x));
+  }
+};
+
 template <typename T>
 struct CudaSeluFunctor : public BaseActivationFunctor<T> {
   typename BaseActivationFunctor<T>::AttrPair GetAttrs() {
@@ -2866,6 +2924,20 @@ struct CudaExpGradFunctor : public BaseActivationFunctor<T> {
   }
 };
 
+template <typename T>
+struct CudaExpGradFunctor<ComplexType<T>>
+    : public BaseActivationFunctor<ComplexType<T>> {
+  // dx = dout * exp(x)
+  __device__ __forceinline__ ComplexType<T> operator()(
+      const ComplexType<T> dout, const ComplexType<T> out) const {
+    return static_cast<ComplexType<T>>(dout * conj(out));
+  }
+
+  static constexpr ActBwdOpFwdDeps FwdDeps() {
+    return ActBwdOpFwdDeps::kDepOut;
+  }
+};
+
 template <typename T>
 struct CudaReciprocalFunctor : public BaseActivationFunctor<T> {
   using MPType = typename phi::dtype::MPTypeTrait<T>::Type;
@@ -2906,6 +2978,15 @@ struct CudaExpm1Functor<double> : public BaseActivationFunctor<double> {
   }
 };
 
+template <typename T>
+struct CudaExpm1Functor<ComplexType<T>>
+    : public BaseActivationFunctor<ComplexType<T>> {
+  __device__ __forceinline__ ComplexType<T> operator()(
+      const ComplexType<T> x) const {
+    return static_cast<ComplexType<T>>(Expm1<ComplexType<T>>()(x));
+  }
+};
+
 template <typename T>
 struct CudaExpm1GradFunctor : public BaseActivationFunctor<T> {
   // dx = dout * out
@@ -2918,6 +2999,20 @@ struct CudaExpm1GradFunctor : public BaseActivationFunctor<T> {
   }
 };
 
+template <typename T>
+struct CudaExpm1GradFunctor<ComplexType<T>>
+    : public BaseActivationFunctor<ComplexType<T>> {
+  // dx = dout * exp(x)
+  __device__ __forceinline__ ComplexType<T> operator()(
+      const ComplexType<T> dout, const ComplexType<T> out) const {
+    return static_cast<ComplexType<T>>(dout * conj(out) + dout);
+  }
+
+  static constexpr ActBwdOpFwdDeps FwdDeps() {
+    return ActBwdOpFwdDeps::kDepOut;
+  }
+};
+
 template <typename T>
 struct CudaSinFunctor : public BaseActivationFunctor<T> {
   using MPType = typename phi::dtype::MPTypeTrait<T>::Type;

paddle/phi/kernels/gpu/activation_grad_kernel.cu

@@ -398,7 +398,9 @@ PD_REGISTER_KERNEL(exp_grad,
                    int,
                    int64_t,
                    phi::dtype::float16,
-                   phi::dtype::bfloat16) {}
+                   phi::dtype::bfloat16,
+                   phi::dtype::complex<float>,
+                   phi::dtype::complex<double>) {}
 
 PD_REGISTER_ACTIVATION_GRAD_KERNEL(softshrink_grad, SoftShrinkGradKernel)
 PD_REGISTER_ACTIVATION_GRAD_KERNEL(hard_shrink_grad, HardShrinkGradKernel)
@@ -415,7 +417,9 @@ PD_REGISTER_KERNEL(expm1_grad,
                    float,
                    double,
                    phi::dtype::float16,
-                   phi::dtype::bfloat16) {}
+                   phi::dtype::bfloat16,
+                   phi::dtype::complex<float>,
+                   phi::dtype::complex<double>) {}
 
 PD_REGISTER_KERNEL(square_grad,
                    GPU,

paddle/phi/kernels/gpu/activation_kernel.cu

@@ -261,7 +261,10 @@ PD_REGISTER_KERNEL(exp,
                    int,
                    int64_t,
                    phi::dtype::float16,
-                   phi::dtype::bfloat16) {}
+                   phi::dtype::bfloat16,
+                   phi::dtype::complex<float>,
+                   phi::dtype::complex<double>) {}
+
 PD_REGISTER_KERNEL(expm1,
                    GPU,
                    ALL_LAYOUT,
@@ -271,7 +274,10 @@ PD_REGISTER_KERNEL(expm1,
                    int,
                    int64_t,
                    phi::dtype::float16,
-                   phi::dtype::bfloat16) {}
+                   phi::dtype::bfloat16,
+                   phi::dtype::complex<float>,
+                   phi::dtype::complex<double>) {}
+
 PD_REGISTER_KERNEL(square,
                    GPU,
                    ALL_LAYOUT,

python/paddle/tensor/ops.py

@@ -640,7 +640,16 @@ def expm1(x, name=None):
         check_variable_and_dtype(
             x,
             'x',
-            ['float16', 'uint16', 'float32', 'float64', 'int32', 'int64'],
+            [
+                'float16',
+                'uint16',
+                'float32',
+                'float64',
+                'int32',
+                'int64',
+                'complex64',
+                'complex128',
+            ],
             'expm1',
         )
         helper = LayerHelper('expm1', **locals())

test/legacy_test/test_activation_op.py

@@ -149,9 +149,46 @@ def init_dtype(self):
         self.dtype = np.float64
 
 
-class TestExpPrim_ZeroDim(TestExpFp32_Prim):
+class TestExp_Complex64(OpTest):
+    def setUp(self):
+        self.op_type = "exp"
+        self.python_api = paddle.exp
+        self.public_python_api = paddle.exp
+        self.init_dtype()
+        self.init_shape()
+        self.if_enable_cinn()
+        np.random.seed(1024)
+        x = (
+            np.random.uniform(-1, 1, self.shape)
+            + 1j * np.random.uniform(-1, 1, self.shape)
+        ).astype(self.dtype)
+        out = np.exp(x)
+        self.inputs = {'X': OpTest.np_dtype_to_fluid_dtype(x)}
+        self.outputs = {'Out': out}
+        self.convert_input_output()
+
+    def test_check_output(self):
+        self.check_output()
+
+    def test_check_grad(self):
+        self.check_grad(['X'], 'Out', max_relative_error=0.006)
+
+    def init_dtype(self):
+        self.dtype = np.complex64
+
     def init_shape(self):
-        self.shape = []
+        self.shape = [10, 12]
+
+    def if_enable_cinn(self):
+        pass
+
+    def convert_input_output(self):
+        pass
+
+
+class TestExp_Complex128(TestExp_Complex64):
+    def init_dtype(self):
+        self.dtype = np.complex128
 
 
 class Test_Exp_Op_Fp16(unittest.TestCase):
@@ -189,9 +226,13 @@ def setUp(self):
         self.python_api = paddle.expm1
         self.init_dtype()
         self.init_shape()
-
         np.random.seed(2049)
         x = np.random.uniform(0.1, 1, self.shape).astype(self.dtype)
+        if self.dtype == np.complex64 or self.dtype == np.complex128:
+            x = (
+                np.random.uniform(-1, 1, self.shape)
+                + 1j * np.random.uniform(-1, 1, self.shape)
+            ).astype(self.dtype)
         out = np.expm1(x)
 
         self.inputs = {'X': OpTest.np_dtype_to_fluid_dtype(x)}
@@ -205,6 +246,16 @@ def test_check_output(self):
         self.check_output()
 
 
+class TestExpm1_Complex64(TestExpm1):
+    def init_dtype(self):
+        self.dtype = np.complex64
+
+
+class TestExpm1_Complex128(TestExpm1):
+    def init_dtype(self):
+        self.dtype = np.complex128
+
+
 class TestExpm1_ZeroDim(TestExpm1):
     def init_shape(self):
         self.shape = []