diff --git a/paddle/phi/kernels/cpu/index_select_kernel.cc b/paddle/phi/kernels/cpu/index_select_kernel.cc index 572f7ac3a04814..9f62b567f2a79a 100644 --- a/paddle/phi/kernels/cpu/index_select_kernel.cc +++ b/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(output); return; } auto inputs = x; - if (dim < 0) { - dim += inputs.dims().size(); - } const auto& index_type = index.dtype(); bool index_type_match = diff --git a/paddle/phi/kernels/gpu/index_select_kernel.cu b/paddle/phi/kernels/gpu/index_select_kernel.cu index 3f71cf5810487c..332d9fbb977edf 100644 --- a/paddle/phi/kernels/gpu/index_select_kernel.cu +++ b/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(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]; diff --git a/paddle/phi/kernels/xpu/index_select_kernel.cc b/paddle/phi/kernels/xpu/index_select_kernel.cc index 83c4c65521f2d6..9997f0b4e88279 100644 --- a/paddle/phi/kernels/xpu/index_select_kernel.cc +++ b/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(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 = diff --git a/test/legacy_test/test_index_select_op.py b/test/legacy_test/test_index_select_op.py index e30fb2f2b4b797..beb9e0b669aa6f 100644 --- a/test/legacy_test/test_index_select_op.py +++ b/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( diff --git a/test/xpu/test_index_select_op_xpu.py b/test/xpu/test_index_select_op_xpu.py index 1d6a84657b704b..df386b966ef6e6 100644 --- a/test/xpu/test_index_select_op_xpu.py +++ b/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)