[0-size Tensor No.99、310] Add 0-size Tensor support for kthvalue

paddle/phi/kernels/cpu/kthvalue_grad_kernel.cc

@@ -56,6 +56,9 @@ void KthvalueGradKernel(const Context& dev_ctx,
   auto in_dims = x.dims();
   auto out_dims = indices.dims();
   T* x_grad_data = dev_ctx.template Alloc<T>(d_x);
+  if (d_x && d_x->numel() == 0) {
+    return;
+  }
 
   // For 0D Tensor
   if (in_dims.size() == 0) {

paddle/phi/kernels/cpu/kthvalue_kernel.cc

@@ -16,9 +16,9 @@
 
 #include "paddle/phi/backends/cpu/cpu_context.h"
 #include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/full_kernel.h"
 #include "paddle/phi/kernels/funcs/eigen/common.h"
 #include "paddle/phi/kernels/funcs/math_function.h"
-
 namespace phi {
 template <typename T, typename Type>
 static void getKthvalue(Type input_height,
@@ -80,6 +80,13 @@ void KthvalueKernel(const Context& dev_ctx,
                     bool keepdim,
                     DenseTensor* output,
                     DenseTensor* indices) {
+  if (x.numel() == 0) {
+    phi::Full<T, Context>(
+        dev_ctx, phi::IntArray(common::vectorize(output->dims())), NAN, output);
+    phi::Full<int64_t, Context>(
+        dev_ctx, phi::IntArray(common::vectorize(indices->dims())), 0, indices);
+    return;
+  }
   const auto& in_dims = x.dims();
   if (axis < 0) axis += in_dims.size();
 

paddle/phi/kernels/gpu/kthvalue_grad_kernel.cu

@@ -44,6 +44,10 @@ void KthvalueGradKernel(const Context& dev_ctx,
   const auto& in_dims = x.dims();
   auto out_dims = indices.dims();
   T* x_grad_data = dev_ctx.template Alloc<T>(d_x);
+  if (d_x && d_x->numel() == 0) {
+    return;
+  }
+
   // For 0D Tensor
   if (in_dims.size() == 0) {
     phi::funcs::set_constant(dev_ctx, d_x, static_cast<T>(1.0));

paddle/phi/kernels/gpu/kthvalue_kernel.cu

@@ -16,6 +16,7 @@
 #include "paddle/phi/backends/gpu/gpu_context.h"
 #include "paddle/phi/common/amp_type_traits.h"
 #include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/full_kernel.h"
 #include "paddle/phi/kernels/funcs/eigen/common.h"
 #include "paddle/phi/kernels/funcs/eigen/eigen_function.h"
 #include "paddle/phi/kernels/funcs/math_function.h"
@@ -160,6 +161,14 @@ void KthvalueKernel(const Context& dev_ctx,
                     bool keepdim,
                     DenseTensor* output,
                     DenseTensor* indices) {
+  if (x.numel() == 0) {
+    phi::Full<T, Context>(
+        dev_ctx, phi::IntArray(common::vectorize(output->dims())), NAN, output);
+    phi::Full<int64_t, Context>(
+        dev_ctx, phi::IntArray(common::vectorize(indices->dims())), 0, indices);
+    return;
+  }
+
   const auto& in_dims = x.dims();
   if (axis < 0) axis += in_dims.size();
   auto out_dims = output->dims();

test/legacy_test/test_kthvalue_op.py

@@ -43,13 +43,17 @@ def init_args(self):
     def init_dtype(self):
         self.dtype = np.float64
 
+    def init_shape(self):
+        self.shape = [2, 1, 2, 4, 10]
+
     def setUp(self):
         self.op_type = "kthvalue"
         self.prim_op_type = "prim"
         self.python_api = paddle.kthvalue
         self.public_python_api = paddle.kthvalue
         self.init_dtype()
-        self.input_data = np.random.random([2, 1, 2, 4, 10]).astype(self.dtype)
+        self.init_shape()
+        self.input_data = np.random.random(self.shape).astype(self.dtype)
         self.init_args()
         self.inputs = {'X': self.input_data}
         self.attrs = {'k': self.k, 'axis': self.axis}
@@ -77,6 +81,11 @@ def init_dtype(self):
         self.dtype = np.float16
 
 
+class TestKthvalueOp_ZeroSize(TestKthvalueOp):
+    def init_shape(self):
+        self.shape = [2, 1, 0, 4, 10]
+
+
 class TestKthvalueOpWithKeepdim(OpTest):
     def init_args(self):
         self.k = 2