Skip to content

Commit

Permalink
No.10 add complex support for exp/expm1 (#57116)
Browse files Browse the repository at this point in the history
  • Loading branch information
@Wanglongzhi2001
Wanglongzhi2001 authored Sep 11, 2023
1 parent 8a6da17 commit da40602
Show file tree
Hide file tree
Showing 7 changed files with 191 additions and 11 deletions.
8 changes: 6 additions & 2 deletions paddle/phi/kernels/cpu/activation_grad_kernel.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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) {}
Expand Down
8 changes: 6 additions & 2 deletions paddle/phi/kernels/cpu/activation_kernel.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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(
Expand Down
98 changes: 98 additions & 0 deletions paddle/phi/kernels/funcs/activation_functor.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1167,6 +1167,33 @@ 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 {};

template <typename T>
struct Expm1<ComplexType<T>> {
HOSTDEVICE ComplexType<T> operator()(const ComplexType<T>& val) const {
return exp(val) - static_cast<ComplexType<T>>(1);
}
};

// expm1(x) = e^x - 1
template <typename T>
struct Expm1Functor : public BaseActivationFunctor<T> {
Expand All @@ -1178,6 +1205,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,
Expand All @@ -1194,6 +1230,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> {
Expand Down Expand Up @@ -2831,6 +2882,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() {
Expand Down Expand Up @@ -2907,6 +2968,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;
Expand Down Expand Up @@ -2947,6 +3022,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
Expand All @@ -2959,6 +3043,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;
Expand Down
8 changes: 6 additions & 2 deletions paddle/phi/kernels/gpu/activation_grad_kernel.cu
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -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,
Expand Down
8 changes: 6 additions & 2 deletions paddle/phi/kernels/gpu/activation_kernel.cu
View file
Original file line number Diff line number Diff line change
Expand Up @@ -261,7 +261,9 @@ 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,
Expand All @@ -271,7 +273,9 @@ 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,
Expand Down
15 changes: 12 additions & 3 deletions python/paddle/tensor/ops.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -567,7 +567,7 @@ def exp(x, name=None):
out = e^x
Args:
x (Tensor): Input of Exp operator, an N-D Tensor, with data type int32, int64, float32, float64 or float16.
x (Tensor): Input of Exp operator, an N-D Tensor, with data type int32, int64, float16, float32, float64, complex64 or complex128.
name (str, optional): Name for the operation (optional, default is None). For more information, please refer to :ref:`api_guide_Name`.
Returns:
Expand Down Expand Up @@ -617,7 +617,7 @@ def expm1(x, name=None):
out = e^x - 1
Args:
x (Tensor): Input of Expm1 operator, an N-D Tensor, with data type int32, int64, float32, float64 or float16.
x (Tensor): Input of Expm1 operator, an N-D Tensor, with data type int32, int64, float16, float32, float64, complex64 or complex128.
name (str, optional): Name for the operation (optional, default is None). For more information, please refer to :ref:`api_guide_Name`.
Returns:
Expand All @@ -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())
Expand Down
57 changes: 57 additions & 0 deletions test/legacy_test/test_activation_op.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -154,6 +154,48 @@ def init_shape(self):
self.shape = []


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_base_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 = [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):
def test_api_fp16(self):
with paddle.base.framework._static_guard():
Expand Down Expand Up @@ -192,6 +234,11 @@ def setUp(self):

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_base_dtype(x)}
Expand All @@ -205,6 +252,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 = []
Expand Down

0 comments on commit da40602

Please sign in to comment.