未验证 提交 ff803bdc 编写于 作者: 张春乔 提交者: GitHub

add REGISTER of as_real (#51313)

上级 59c1ffa8
......@@ -58,7 +58,6 @@ static std::set<std::string> OpsNeedSetOutputDtypeWhenRegisterPhiKernel = {
"angle",
"any_raw",
"arg_sort",
"as_real",
"atan2",
"auc",
"bincount",
......
......@@ -273,6 +273,7 @@ void AsRealInferMeta(const MetaTensor& input, MetaTensor* output) {
auto out_dims = phi::make_ddim(out_dims_v);
output->set_dims(out_dims);
output->share_lod(input);
output->set_dtype(dtype::ToReal(input.dtype()));
}
void AsComplexInferMeta(const MetaTensor& input, MetaTensor* output) {
......
......@@ -23,4 +23,6 @@ using complex64 = ::phi::dtype::complex<float>;
using complex128 = ::phi::dtype::complex<double>;
PD_REGISTER_KERNEL(
as_real, CPU, ALL_LAYOUT, phi::AsRealKernel, complex64, complex128) {}
as_real, CPU, ALL_LAYOUT, phi::AsRealKernel, complex64, complex128) {
kernel->OutputAt(0).SetDataType(paddle::experimental::DataType::UNDEFINED);
}
......@@ -23,4 +23,6 @@ using complex64 = ::phi::dtype::complex<float>;
using complex128 = ::phi::dtype::complex<double>;
PD_REGISTER_KERNEL(
as_real, GPU, ALL_LAYOUT, phi::AsRealKernel, complex64, complex128) {}
as_real, GPU, ALL_LAYOUT, phi::AsRealKernel, complex64, complex128) {
kernel->OutputAt(0).SetDataType(paddle::experimental::DataType::UNDEFINED);
}
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册