diff --git a/paddle/fluid/pybind/imperative.cc b/paddle/fluid/pybind/imperative.cc index 4403eb469723a5afdbc602d3f6ce19966536079c..60b97b76491c80fd83c05fd08606de492928e6e7 100644 --- a/paddle/fluid/pybind/imperative.cc +++ b/paddle/fluid/pybind/imperative.cc @@ -528,13 +528,20 @@ static void ParseIndexingSlice( // specified_dims is the number of dimensions which indexed by Interger, // Slices. int specified_dims = 0; + int ell_count = 0; for (int dim = 0; dim < size; ++dim) { PyObject *slice_item = PyTuple_GetItem(index, dim); if (PyCheckInteger(slice_item) || PySlice_Check(slice_item)) { specified_dims++; + } else if (slice_item == Py_Ellipsis) { + ell_count++; } } + PADDLE_ENFORCE_LE(ell_count, 1, + platform::errors::InvalidArgument( + "An index can only have a single ellipsis ('...')")); + for (int i = 0, dim = 0; i < size; ++i) { PyObject *slice_item = PyTuple_GetItem(index, i); @@ -639,7 +646,7 @@ static void ParseIndexingSlice( } // valid_index is the number of dimensions exclude None index - const int valid_indexs = size - none_axes->size(); + const int valid_indexs = size - none_axes->size() - ell_count; PADDLE_ENFORCE_EQ(valid_indexs <= rank, true, platform::errors::InvalidArgument( "Too many indices (%d) for tensor of dimension %d.", diff --git a/python/paddle/fluid/tests/unittests/test_var_base.py b/python/paddle/fluid/tests/unittests/test_var_base.py index cfaef15c1d335a6f441b57deb204d2e3821c9e78..e4ef7428202d5e9d7a1f70fae34fd2584cdd9770 100644 --- a/python/paddle/fluid/tests/unittests/test_var_base.py +++ b/python/paddle/fluid/tests/unittests/test_var_base.py @@ -702,6 +702,11 @@ class TestVarBase(unittest.TestCase): assert_getitem_ellipsis_index(var_fp32, np_fp32_value) assert_getitem_ellipsis_index(var_int, np_int_value) + # test 1 dim tensor + var_one_dim = paddle.to_tensor([1, 2, 3, 4]) + self.assertTrue( + np.array_equal(var_one_dim[..., 0].numpy(), np.array([1]))) + def _test_none_index(self): shape = (8, 64, 5, 256) np_value = np.random.random(shape).astype('float32') diff --git a/python/paddle/fluid/tests/unittests/test_variable.py b/python/paddle/fluid/tests/unittests/test_variable.py index e9e959266dba8e8f9751e3285eb6c272284e4968..2df336fbe8e77142377ce54eed7ab7b365a01e2e 100644 --- a/python/paddle/fluid/tests/unittests/test_variable.py +++ b/python/paddle/fluid/tests/unittests/test_variable.py @@ -226,19 +226,22 @@ class TestVariable(unittest.TestCase): prog = paddle.static.Program() with paddle.static.program_guard(prog): x = paddle.assign(data) + y = paddle.assign([1, 2, 3, 4]) out1 = x[0:, ..., 1:] out2 = x[0:, ...] out3 = x[..., 1:] out4 = x[...] out5 = x[[1, 0], [0, 0]] out6 = x[([1, 0], [0, 0])] + out7 = y[..., 0] exe = paddle.static.Executor(place) - result = exe.run(prog, fetch_list=[out1, out2, out3, out4, out5, out6]) + result = exe.run(prog, + fetch_list=[out1, out2, out3, out4, out5, out6, out7]) expected = [ data[0:, ..., 1:], data[0:, ...], data[..., 1:], data[...], - data[[1, 0], [0, 0]], data[([1, 0], [0, 0])] + data[[1, 0], [0, 0]], data[([1, 0], [0, 0])], np.array([1]) ] self.assertTrue((result[0] == expected[0]).all()) @@ -247,6 +250,7 @@ class TestVariable(unittest.TestCase): self.assertTrue((result[3] == expected[3]).all()) self.assertTrue((result[4] == expected[4]).all()) self.assertTrue((result[5] == expected[5]).all()) + self.assertTrue((result[6] == expected[6]).all()) with self.assertRaises(IndexError): res = x[[1.2, 0]]