Skip to content

Commit a1ed2eb

Browse files
committed
let stack kernel support the same dtypes as concat kernel
1 parent eff0bce commit a1ed2eb

File tree

2 files changed

+10
-3
lines changed

2 files changed

+10
-3
lines changed

paddle/phi/kernels/cpu/stack_kernel.cc

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -57,8 +57,12 @@ PD_REGISTER_KERNEL(stack,
5757
phi::StackKernel,
5858
float,
5959
double,
60-
int,
60+
bool,
6161
int64_t,
62+
int,
6263
uint8_t,
6364
int8_t,
64-
phi::dtype::bfloat16) {}
65+
phi::dtype::float16,
66+
phi::dtype::bfloat16,
67+
phi::dtype::complex<float>,
68+
phi::dtype::complex<double>) {}

paddle/phi/kernels/gpu/stack_kernel.cu

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -175,9 +175,12 @@ PD_REGISTER_KERNEL(stack,
175175
phi::StackKernel,
176176
float,
177177
double,
178+
bool,
178179
int64_t,
179180
int,
180181
uint8_t,
181182
int8_t,
182183
phi::dtype::float16,
183-
phi::dtype::bfloat16) {}
184+
phi::dtype::bfloat16,
185+
phi::dtype::complex<float>,
186+
phi::dtype::complex<double>) {}

0 commit comments

Comments
 (0)