[numpy] symbolic advanced indexing

python/mxnet/numpy/multiarray.py

@@ -699,6 +699,14 @@ def __getitem__(self, key):
             elif key.step == 0:
                 raise ValueError("slice step cannot be zero")
 
+
+        all = __builtins__['all']  # `def all` below shadows the all builtin
+        if (isinstance(key, tuple) and all( \
+        (isinstance(arr, NDArray) \
+        and _np.issubdtype(arr.dtype, _np.integer) and arr.ndim > 0) \
+        for arr in key)):
+            return _npi.advanced_indexing_multiple(self, _npi.stack(*key))
+
         # For 0-d boolean indices: A new axis is added,
         # but at the same time no axis is "used". So if we have True,
         # we add a new axis (a bit like with np.newaxis). If it is

python/mxnet/symbol/numpy/_symbol.py

@@ -117,6 +117,8 @@ def __getitem__(self, key): # pylint: disable = too-many-return-statements, inco
                 end = []
                 step = []
                 new_shape = ()
+                result = self
+                is_symbol_tuple = False
                 if len(key) == 0:
                     return self
                 for index in key:
@@ -137,14 +139,27 @@ def __getitem__(self, key): # pylint: disable = too-many-return-statements, inco
                             end.append(index - 1)
                             step.append(-1)
                         new_shape += (-3,)
+                    elif isinstance(index, Symbol):
+                        if new_shape != ():
+                            new_shape += (-4,)
+                            sliced = _npi.slice(result, begin, end, step)
+                            result = _npi.reshape(sliced, new_shape)
+                        if not is_symbol_tuple:
+                            is_symbol_tuple = True
                     else:
-                        raise IndexError('Only integer, slice, or tuple of these types'
+                        raise IndexError('Only integer, slice, symbol or tuple of these types'
                                          ' are supported! Received key={}'.format(key))
+                if is_symbol_tuple:
+                    key = _npi.stack(*[i for i in key])
+                    sliced = _npi.advanced_indexing_multiple(self, key)
+                    return sliced
                 new_shape += (-4,)
                 sliced = _npi.slice(self, begin, end, step)
                 return _npi.reshape(sliced, new_shape)
+            elif isinstance(key, Symbol):
+                return _npi.advanced_indexing(self, key)
             else:
-                raise IndexError('Only integer, slice, or tuple of these types are supported! '
+                raise IndexError('Only integer, slice, tuple or Symbol of these types are supported! '
                                  'Received key={}'.format(key))
 
     def __setitem__(self, key, value):