Support Combined indexing for __getitem__ and __setitem__

paddle/fluid/pybind/eager_method.cc

@@ -927,6 +927,9 @@ static PyObject* tensor__getitem_index_not_tensor(TensorObject* self,
       eager_gil_scoped_release guard;
       out = strided_slice_ad_func(
           self->tensor, slice_axes, slice_starts, slice_ends, slice_strides);
+      if (!decrease_axis_tmp.empty()) {
+        out = squeeze_ad_func(out, decrease_axis_tmp);
+      }
     } else {
       PADDLE_THROW(platform::errors::InvalidArgument(
           "Slice is only support slice and strided_slice, but we got %s which "

python/paddle/fluid/dygraph/tensor_patch_methods.py

@@ -26,7 +26,8 @@
 from ..framework import (
     Variable,
     Parameter,
-    _getitem_impl_,
+    _getitem_static,
+    _setitem_static,
     _setitem_impl_,
     EagerParamBase,
     in_dygraph_mode,
@@ -717,47 +718,34 @@ def contain_tensor(item):
                     return True
         return False
 
-    def __getitem__(self, item):
-        def is_list_tuple(index, contain_type):
-            def _is_list_tuple(item):
-                if isinstance(item, (tuple, list)):
-                    for s in item:
-                        if not _is_list_tuple(s):
-                            return False
-                else:
-                    if type(item) != contain_type:
-                        return False
+    def contain_tensor_or_list(item):
+        if not isinstance(item, tuple):
+            item = (item,)
+
+        for slice_item in item:
+            if isinstance(slice_item, (list, np.ndarray, Variable)):
                 return True
+            elif isinstance(slice_item, slice):
+                if (
+                    isinstance(slice_item.start, Variable)
+                    or isinstance(slice_item.stop, Variable)
+                    or isinstance(slice_item.step, Variable)
+                ):
+                    return True
 
-            if not isinstance(index, (tuple, list)):
-                return False
-            for s in index:
-                if not _is_list_tuple(s):
-                    return False
-            return True
+        return False
 
-        if contain_tensor(item) or is_list_tuple(item, int):
+    def __getitem__(self, item):
+        if contain_tensor_or_list(item):
             # 1. Call _getitem_impl_ when item contains tensor.
             # Why not call a c++ function ? Because item can't be parsed when it contains tensor.
-            return _getitem_impl_(self, item)
+            return _getitem_static(self, item)
 
         else:
             # 2. Call c++ func getitem_index_not_tensor to speedup.
             return self._getitem_index_not_tensor(item)
 
     def __setitem__(self, item, value):
-        def contain_tensor_or_list(item):
-            if not isinstance(item, tuple):
-                item = [item]
-
-            for slice_item in item:
-                if isinstance(slice_item, list):
-                    return True
-                elif isinstance(slice_item, Variable):
-                    return True
-
-            return False
-
         def is_combine_index(item):
             var_type = None
             item_type = None
@@ -779,10 +767,13 @@ def is_combine_index(item):
 
             return False
 
-        if contain_tensor_or_list(item) and not is_combine_index(item):
+        if contain_tensor_or_list(item):
+            if core.is_compiled_with_xpu() and not is_combine_index(item):
+                # (NOTE): Currently, there is no index_put_xpu kernel.
+                return _setitem_impl_(self, item, value)
             # To reuse code with static graph,
-            # Call _setitem_impl_ when item contains tensor or list.
-            return _setitem_impl_(self, item, value)
+            # Call _setitem_static when item contains tensor or list.
+            return _setitem_static(self, item, value)
 
         else:
             return self.__setitem_eager_tensor__(item, value)

python/paddle/fluid/framework.py

@@ -37,7 +37,7 @@
 import paddle.version as fluid_version
 import warnings
 import functools
-from .variable_index import _getitem_impl_, _setitem_impl_
+from .variable_index import _getitem_static, _setitem_static, _setitem_impl_
 import threading
 
 __all__ = [
@@ -2292,13 +2292,16 @@ def _sliceAndConcatVar(self, item, axis):
             raise IndexError("Valid index accept int or slice or tuple")
 
     def __getitem__(self, item):
-        return _getitem_impl_(self, item)
+        return _getitem_static(self, item)
 
     def __setitem__(self, item, value):
         from .dygraph.base import in_declarative_mode
 
         if in_declarative_mode():
-            return _setitem_impl_(self, item, value)
+            if is_compiled_with_xpu():
+                # (NOTE): Currently, there is no index_put_xpu kernel.
+                return _setitem_impl_(self, item, value)
+            return _setitem_static(self, item, value)
         else:
             raise RuntimeError(
                 "In static mode, the __setitem__ (looks like: x[indices] = values) should not be used. Please use x = paddle.static.setitem(x, indices, values)"