add 0-d support for paddle.mode

paddle/phi/infermeta/unary.cc

@@ -1760,9 +1760,7 @@ void KthvalueInferMeta(const MetaTensor& x,
                           dim_size,
                           axis));
   }
-  if (axis < 0) {
-    axis = (dim_size > 0) ? axis + dim_size : 0;
-  }
+  if (axis < 0) axis += dim_size;
   PADDLE_ENFORCE_GE(
       k,
       1,
@@ -2034,33 +2032,38 @@ void ModeInferMeta(const MetaTensor& x,
                    MetaTensor* indices) {
   auto input_dims = x.dims();
   const int& dim_size = input_dims.size();
-  PADDLE_ENFORCE_EQ(
-      (axis < dim_size) && (axis >= (-1 * dim_size)),
-      true,
-      errors::InvalidArgument(
-          "the axis of ModeOp must be [-%d, %d), but you set axis is %d",
-          dim_size,
-          dim_size,
-          axis));
+  PADDLE_ENFORCE_LT(axis,
+                    dim_size,
+                    phi::errors::InvalidArgument(
+                        "the axis must be [-%d, %d), but received %d .",
+                        dim_size,
+                        dim_size,
+                        axis));
+  if (dim_size > 0) {
+    PADDLE_ENFORCE_GE(axis,
+                      -dim_size,
+                      phi::errors::InvalidArgument(
+                          "the axis must be [-%d, %d), but received %d .",
+                          dim_size,
+                          dim_size,
+                          axis));
+  }
   PADDLE_ENFORCE_GE(
       input_dims.size(),
-      1,
-      errors::InvalidArgument("input of ModeOp must have >= 1d shape"));
+      0,
+      errors::InvalidArgument("input of ModeOp must have >= 0d shape"));
   if (axis < 0) axis += dim_size;
   std::vector<int64_t> dimvec;
   for (int64_t i = 0; i < axis; i++) {
     dimvec.emplace_back(input_dims[i]);
   }
-  if (keepdim) {
+  if (keepdim && dim_size > 0) {
     dimvec.emplace_back(static_cast<int64_t>(1));
   }
   for (int64_t i = axis + 1; i < dim_size; i++) {
     dimvec.emplace_back(input_dims[i]);
   }
   DDim dims = phi::make_ddim(dimvec);
-  PADDLE_ENFORCE_GE(input_dims.size(),
-                    1,
-                    errors::InvalidArgument("input shape should >= 1d"));
   out->set_dims(dims);
   out->share_lod(x);
   out->set_dtype(x.dtype());

paddle/phi/kernels/cpu/kthvalue_kernel.cc

@@ -81,9 +81,8 @@ void KthvalueKernel(const Context& dev_ctx,
                     DenseTensor* output,
                     DenseTensor* indices) {
   const auto& in_dims = x.dims();
-  if (axis < 0) {
-    axis += in_dims.size();
-  }
+  if (axis < 0) axis += in_dims.size();
+
   T* output_data = dev_ctx.template Alloc<T>(output);
   int64_t* indices_data = dev_ctx.template Alloc<int64_t>(indices);
   // For 0D Tensor

paddle/phi/kernels/cpu/mode_grad_kernel.cc

@@ -17,6 +17,7 @@
 #include "paddle/phi/backends/cpu/cpu_context.h"
 #include "paddle/phi/core/kernel_registry.h"
 #include "paddle/phi/core/tensor_utils.h"
+#include "paddle/phi/kernels/funcs/math_function.h"
 #include "paddle/phi/kernels/funcs/mode.h"
 
 namespace phi {
@@ -32,9 +33,17 @@ void ModeGradKernel(const Context& dev_ctx,
   auto in_dims = x.dims();
   auto out_dims = indices.dims();
 
+  T* x_grad_data = dev_ctx.template Alloc<T>(x_grad);
+
   // axis < 0, get the real axis
   axis = (axis < 0) ? (in_dims.size() + axis) : axis;
 
+  // For 0D Tensor
+  if (in_dims.size() == 0) {
+    phi::funcs::set_constant(dev_ctx, x_grad, 1.0);
+    return;
+  }
+
   if (!keepdim) {
     std::vector<int> tmp_out_shape;
     for (int i = 0; i < axis; i++) {
@@ -46,7 +55,6 @@ void ModeGradKernel(const Context& dev_ctx,
     }
     out_dims = phi::make_ddim(tmp_out_shape);
   }
-  T* x_grad_data = dev_ctx.template Alloc<T>(x_grad);
 
   if (axis == in_dims.size() - 1) {
     // allocate the memory for the input_grad

paddle/phi/kernels/cpu/mode_kernel.cc

@@ -16,6 +16,7 @@
 
 #include "paddle/phi/backends/cpu/cpu_context.h"
 #include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/funcs/math_function.h"
 #include "paddle/phi/kernels/funcs/mode.h"
 
 namespace phi {
@@ -34,6 +35,13 @@ void ModeKernel(const Context& dev_ctx,
 
   T* output_data = dev_ctx.template Alloc<T>(out);
   int64_t* indices_data = dev_ctx.template Alloc<int64_t>(indices);
+
+  if (in_dims.size() == 0) {
+    phi::Copy<Context>(dev_ctx, x, dev_ctx.GetPlace(), false, out);
+    phi::funcs::set_constant(dev_ctx, indices, 0);
+    return;
+  }
+
   // if axis is not the last dim, transpose it to the last dim, do the
   // calculation, then tranpose it back to original axis.
   if (axis == in_dims.size() - 1) {

paddle/phi/kernels/gpu/mode_grad_kernel.cu

@@ -16,6 +16,7 @@
 
 #include "paddle/phi/backends/gpu/gpu_context.h"
 #include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/funcs/math_function.h"
 #include "paddle/phi/kernels/funcs/mode.h"
 
 namespace phi {
@@ -61,6 +62,12 @@ void ModeGradKernel(const Context& dev_ctx,
   const T* out_grad_data = out_grad.data<T>();
   const int64_t* indices_data = indices.data<int64_t>();
 
+  // For 0D Tensor
+  if (in_dims.size() == 0) {
+    phi::funcs::set_constant(dev_ctx, x_grad, 1.0);
+    return;
+  }
+
   int pre, n, post;
   funcs::GetDims(in_dims, axis, &pre, &n, &post);
 

paddle/phi/kernels/gpu/mode_kernel.cu

@@ -16,6 +16,7 @@
 
 #include "paddle/phi/backends/gpu/gpu_context.h"
 #include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/funcs/math_function.h"
 #include "paddle/phi/kernels/funcs/mode.h"
 
 namespace phi {
@@ -38,6 +39,13 @@ void ModeKernel(const Context& dev_ctx,
   T* output_data = dev_ctx.template Alloc<T>(out);
   int64_t* indices_data = dev_ctx.template Alloc<int64_t>(indices);
 
+  // For 0D Tensor
+  if (in_dims.size() == 0) {
+    phi::Copy<Context>(dev_ctx, x, dev_ctx.GetPlace(), false, out);
+    phi::funcs::set_constant(dev_ctx, indices, 0);
+    return;
+  }
+
   if (axis == in_dims.size() - 1) {
     const int64_t& input_height =
         phi::product(phi::slice_ddim(in_dims, 0, in_dims.size() - 1));

python/paddle/fluid/tests/unittests/test_zero_dim_tensor.py

@@ -717,7 +717,16 @@ def test_kthvalue(self):
         x.stop_gradient = False
 
         out = paddle.kthvalue(x, 1)
-        out.backward()
+        out[0].backward()
+        self.assertEqual(out[0].shape, [])
+        self.assertEqual(out[1].shape, [])
+
+    def test_mode(self):
+        x = paddle.randn(())
+        x.stop_gradient = False
+
+        out = paddle.mode(x)
+        out[0].backward()
         self.assertEqual(out[0].shape, [])
         self.assertEqual(out[1].shape, [])
 
@@ -934,6 +943,17 @@ def test_kthvalue(self):
         self.assertEqual(len(res[0].shape), 0)
         self.assertEqual(len(res[0].shape), 0)
 
+    @prog_scope()
+    def test_mode(self):
+        x = paddle.full([], 1, 'float32')
+        out = paddle.mode(x)
+        paddle.static.append_backward(out[0])
+
+        prog = paddle.static.default_main_program()
+        res = self.exe.run(prog, fetch_list=[out])
+        self.assertEqual(len(res[0].shape), 0)
+        self.assertEqual(len(res[0].shape), 0)
+
 
 # Use to test API whose zero-dim input tensors don't have grad and not need to test backward in OpTest.
 class TestNoBackwardAPI(unittest.TestCase):