Add output defs for topk kernel (#51233)

paddle/fluid/framework/new_executor/interpreter/interpreter_util.cc

@@ -118,7 +118,6 @@ static std::set<std::string> OpsNeedSetOutputDtypeWhenRegisterPhiKernel = {
     "sgd",
     "svd",
     "sync_batch_norm_grad",
-    "top_k",
     "unique",
     "unique_consecutive_flattened_tensor",
     "unique_raw",

paddle/phi/kernels/cpu/top_k_kernel.cc

@@ -241,4 +241,6 @@ void TopkKernel(const Context& dev_ctx,
 }  // namespace phi
 
 PD_REGISTER_KERNEL(
-    topk, CPU, ALL_LAYOUT, phi::TopkKernel, float, double, int32_t, int64_t) {}
+    topk, CPU, ALL_LAYOUT, phi::TopkKernel, float, double, int32_t, int64_t) {
+  kernel->OutputAt(1).SetDataType(paddle::experimental::DataType::INT64);
+}

paddle/phi/kernels/gpu/top_k_kernel.cu

@@ -348,4 +348,6 @@ PD_REGISTER_KERNEL(topk,
                    double,
                    int,
                    int64_t,
-                   phi::dtype::float16) {}
+                   phi::dtype::float16) {
+  kernel->OutputAt(1).SetDataType(paddle::experimental::DataType::INT64);
+}

paddle/phi/kernels/xpu/top_k_kernel.cc

@@ -187,4 +187,6 @@ void TopkKernel(const Context& dev_ctx,
 }  // namespace phi
 
 PD_REGISTER_KERNEL(
-    topk, XPU, ALL_LAYOUT, phi::TopkKernel, float, phi::dtype::float16) {}
+    topk, XPU, ALL_LAYOUT, phi::TopkKernel, float, phi::dtype::float16) {
+  kernel->OutputAt(1).SetDataType(paddle::experimental::DataType::INT64);
+}