From f2a56c6a0bc5f52f043dd86dcd818593e878dcb8 Mon Sep 17 00:00:00 2001
From: zyfncg <zhangyunfei07@baidu.com>
Date: Mon, 15 Nov 2021 11:10:00 +0800
Subject: [PATCH] fix bug of indexing with ellipsis (#37182)

---
 paddle/fluid/pybind/imperative.cc                    | 9 ++++++++-
 python/paddle/fluid/tests/unittests/test_var_base.py | 5 +++++
 python/paddle/fluid/tests/unittests/test_variable.py | 8 ++++++--
 3 files changed, 19 insertions(+), 3 deletions(-)

diff --git a/paddle/fluid/pybind/imperative.cc b/paddle/fluid/pybind/imperative.cc
index eb62e4903f..a7e3c700bd 100644
--- a/paddle/fluid/pybind/imperative.cc
+++ b/paddle/fluid/pybind/imperative.cc
@@ -549,13 +549,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);
 
@@ -660,7 +667,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 95f7c0aca7..e2a90ed135 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 e9e959266d..2df336fbe8 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]]
-- 
GitLab