[Operator Mechanism] Fix index_select empty-axis validation

paddle/phi/kernels/cpu/index_select_kernel.cc

@@ -26,14 +26,19 @@ void IndexSelectKernel(const Context& dev_ctx,
                        const DenseTensor& index,
                        int dim,
                        DenseTensor* output) {
+  auto input_dim = x.dims();
+  dim = dim >= 0 ? dim : dim + input_dim.size();
+  if (input_dim[dim] == 0 && index.numel() > 0) {
+    PADDLE_THROW(common::errors::InvalidArgument(
+        "The dimension of Input(X) on the select axis in OP(index_select) "
+        "must be greater than 0 when Input(Index) is not empty."));
+  }
+
   if (output && output->numel() == 0) {
     dev_ctx.template Alloc<T>(output);
     return;
   }
   auto inputs = x;
-  if (dim < 0) {
-    dim += inputs.dims().size();
-  }
   const auto& index_type = index.dtype();
 
   bool index_type_match =

paddle/phi/kernels/gpu/index_select_kernel.cu

@@ -30,13 +30,20 @@ void IndexSelectKernel(const Context& dev_ctx,
                        const DenseTensor& index,
                        int dim,
                        DenseTensor* output) {
+  auto input_dim = x.dims();
+  dim = dim >= 0 ? dim : dim + input_dim.size();
+  if (input_dim[dim] == 0 && index.numel() > 0) {
+    PADDLE_THROW(common::errors::InvalidArgument(
+        "The dimension of Input(X) on the select axis in OP(index_select) "
+        "must be greater than 0 when Input(Index) is not empty."));
+  }
+
   if (output && output->numel() == 0) {
     dev_ctx.template Alloc<T>(output);
     return;
   }
-  auto input_dim = x.dims();
+
   auto output_dim = output->dims();
-  dim = dim >= 0 ? dim : dim + input_dim.size();
   auto stride_dim = common::stride(input_dim);
   int64_t stride = stride_dim[dim];
   int64_t size = output_dim[dim];

paddle/phi/kernels/xpu/index_select_kernel.cc

@@ -27,12 +27,18 @@ void IndexSelectKernel(const Context& dev_ctx,
                        const DenseTensor& index,
                        int dim,
                        DenseTensor* output) {
+  auto input_dim = x.dims();
+  dim = dim >= 0 ? dim : dim + input_dim.size();
+  if (input_dim[dim] == 0 && index.numel() > 0) {
+    PADDLE_THROW(common::errors::InvalidArgument(
+        "The dimension of Input(X) on the select axis in OP(index_select) "
+        "must be greater than 0 when Input(Index) is not empty."));
+  }
+
   if (output && output->numel() == 0) {
     dev_ctx.template Alloc<T>(output);
     return;
   }
-  auto input_dim = x.dims();
-  dim = dim >= 0 ? dim : dim + input_dim.size();
   const auto& index_type = index.dtype();
 
   bool index_type_match =

test/legacy_test/test_index_select_op.py

@@ -246,6 +246,56 @@ def init_dtype_type(self):
         self.index_size = 10
 
 
+class TestIndexSelectInvalidIndexCPU(unittest.TestCase):
+    def test_index_select_zero_dim_oob(self):
+        paddle.disable_static()
+        paddle.set_device('cpu')
+        try:
+            for index_dtype in ['int32', 'int64']:
+                for shape, axis in [([1, 0], -1), ([0, 0], 0)]:
+                    with self.subTest(
+                        index_dtype=index_dtype, shape=shape, axis=axis
+                    ):
+                        x = paddle.empty(shape, dtype='float32')
+                        index = paddle.to_tensor([0], dtype=index_dtype)
+                        with self.assertRaisesRegex(
+                            Exception,
+                            r'select axis in OP\(index_select\).*Input\(Index\) is not empty',
+                        ):
+                            paddle.index_select(x, index, axis=axis)
+
+            x = paddle.empty([0, 0], dtype='float32')
+            index = paddle.empty([0], dtype='int64')
+            out = paddle.index_select(x, index, axis=0)
+            self.assertEqual(out.shape, [0, 0])
+        finally:
+            paddle.enable_static()
+
+
+class TestIndexSelectInvalidIndex(unittest.TestCase):
+    def test_index_select_zero_dim_oob(self):
+        paddle.disable_static()
+        device = 'gpu' if paddle.is_compiled_with_cuda() else 'cpu'
+        paddle.set_device(device)
+        try:
+            for shape, axis in [([1, 2048, 0], -1), ([0, 0], 0)]:
+                with self.subTest(shape=shape, axis=axis):
+                    x = paddle.empty(shape, dtype='float32')
+                    index = paddle.to_tensor([0], dtype='int64')
+                    with self.assertRaisesRegex(
+                        Exception,
+                        r'select axis in OP\(index_select\).*Input\(Index\) is not empty',
+                    ):
+                        paddle.index_select(x, index, axis=axis)
+
+            x = paddle.empty([0, 0], dtype='float32')
+            index = paddle.empty([0], dtype='int64')
+            out = paddle.index_select(x, index, axis=0)
+            self.assertEqual(out.shape, [0, 0])
+        finally:
+            paddle.enable_static()
+
+
 class TestIndexSelectAPI(unittest.TestCase):
     def input_data(self):
         self.data_x = np.array(

test/xpu/test_index_select_op_xpu.py

@@ -153,6 +153,25 @@ def test_dygraph_api(self):
         np.testing.assert_allclose(expect_out, np_z, rtol=1e-05)
 
 
+class TestIndexSelectInvalidInput(unittest.TestCase):
+    def test_zero_select_dim(self):
+        with base.dygraph.guard(paddle.XPUPlace(0)):
+            for shape, axis in [([1, 0], -1), ([0, 0], 0)]:
+                with self.subTest(shape=shape, axis=axis):
+                    x = paddle.empty(shape, dtype='float32')
+                    index = paddle.to_tensor([0], dtype='int64')
+                    with self.assertRaisesRegex(
+                        Exception,
+                        r'select axis in OP\(index_select\).*Input\(Index\) is not empty',
+                    ):
+                        paddle.index_select(x, index, axis=axis)
+
+            x = paddle.empty([0, 0], dtype='float32')
+            index = paddle.empty([0], dtype='int64')
+            out = paddle.index_select(x, index, axis=0)
+            self.assertEqual(out.shape, [0, 0])
+
+
 support_types = get_xpu_op_support_types('index_select')
 for stype in support_types:
     create_test_class(globals(), XPUTestIndexSelect, stype)