fix zero bug of case21: paddle.mode (#51091)

PaddlePaddle · Mar 2, 2023 · 25d3ed6 · 25d3ed6
1 parent 2bcd393
commit 25d3ed6
Show file tree

Hide file tree

Showing 3 changed files with 24 additions and 0 deletions.
diff --git a/paddle/phi/kernels/cpu/mode_kernel.cc b/paddle/phi/kernels/cpu/mode_kernel.cc
@@ -29,6 +29,12 @@ void ModeKernel(const Context& dev_ctx,
                 DenseTensor* out,
                 DenseTensor* indices) {
   const auto& in_dims = x.dims();
+  for (int i = 0; i < in_dims.size(); i++) {
+    PADDLE_ENFORCE_LT(0,
+                      in_dims[i],
+                      errors::InvalidArgument(
+                          "The dims of Input(X) should be greater than 0."));
+  }
   auto out_dims = out->dims();
   // axis < 0, cacluate the real axis
   if (axis < 0) axis += in_dims.size();

diff --git a/paddle/phi/kernels/gpu/mode_kernel.cu b/paddle/phi/kernels/gpu/mode_kernel.cu
@@ -30,6 +30,12 @@ void ModeKernel(const Context& dev_ctx,
                 DenseTensor* indices) {
   // get the input dims
   const auto& in_dims = x.dims();
+  for (int i = 0; i < in_dims.size(); i++) {
+    PADDLE_ENFORCE_LT(0,
+                      in_dims[i],
+                      errors::InvalidArgument(
+                          "The dims of Input(X) should be greater than 0."));
+  }
   // calcluate the real axis
   if (axis < 0) axis += in_dims.size();
 

diff --git a/python/paddle/fluid/tests/unittests/test_mode_op.py b/python/paddle/fluid/tests/unittests/test_mode_op.py
@@ -182,5 +182,17 @@ def test_run_static(self):
             np.testing.assert_allclose(paddle_result, expect_value, rtol=1e-05)
 
 
+class TestModeZeroError(unittest.TestCase):
+    def test_errors(self):
+        with paddle.fluid.dygraph.guard():
+
+            def test_0_size():
+                array = np.array([], dtype=np.float32)
+                x = paddle.to_tensor(np.reshape(array, [0, 0]), dtype='float32')
+                paddle.mode(x, axis=0)
+
+            self.assertRaises(ValueError, test_0_size)
+
+
 if __name__ == '__main__':
     unittest.main()