Repository: IntelPython/dpnp

Issue:
Improper handling of key as a tuple in dpnp_array.__getitem__()
The below code caused a fault in dpnp:

```
# import dpnp

# x = dpnp.ones((4, 5))
# mask = dpnp.full(x.shape[1], True)
# x[:, mask]
---------------------------------------------------------------------------
IndexError                                Traceback (most recent call last)
Cell In[4], line 1
----> 1 x[:, mask]

File /localdisk/work/antonvol/code/dpnp_dev/dpnp/dpnp/dpnp_array.py:182, in dpnp_array.__getitem__(self, key)
    179 if isinstance(key, dpnp_array):
    180     key = key.get_array()
--> 182 item = self._array_obj.__getitem__(key)
    183 if not isinstance(item, dpt.usm_ndarray):
    184     raise RuntimeError(
    185         "Expected dpctl.tensor.usm_ndarray, got {}"
    186         "".format(type(item)))

File dpctl/tensor/_usmarray.pyx:716, in dpctl.tensor._usmarray.usm_ndarray.__getitem__()

File dpctl/tensor/_slicing.pxi:171, in dpctl.tensor._usmarray._basic_slice_meta()

File dpctl/tensor/_slicing.pxi:63, in dpctl.tensor._usmarray._is_integral()

File /localdisk/work/antonvol/code/dpnp_dev/dpnp/dpnp/dpnp_array.py:217, in dpnp_array.__index__(self)
    216 def __index__(self):
--> 217     return self._array_obj.__index__()

File dpctl/tensor/_usmarray.pyx:888, in dpctl.tensor._usmarray.usm_ndarray.__index__()

IndexError: only integer arrays are valid indices

```

It looks dpnp wrongly handles indexing use cases when key is a tuple.


Discussion:

Resolving patch:
diff --git a/dpnp/dpnp_array.py b/dpnp/dpnp_array.py
--- a/dpnp/dpnp_array.py
+++ b/dpnp/dpnp_array.py
@@ -29,6 +29,23 @@
 
 import dpnp
 
+
+def _get_unwrapped_index_key(key):
+    """
+    Return a key where each nested instance of DPNP array is unwrapped into USM ndarray
+    for futher processing in DPCTL advanced indexing functions.
+
+    """
+
+    if isinstance(key, tuple):
+        if any(isinstance(x, dpnp_array) for x in key):
+            # create a new tuple from the input key with unwrapped DPNP arrays
+            return tuple(x.get_array() if isinstance(x, dpnp_array) else x for x in key)
+    elif isinstance(key, dpnp_array):
+        return key.get_array()
+    return key
+
+
 class dpnp_array:
     """
     Multi-dimensional array object.
@@ -176,8 +193,7 @@ def __ge__(self, other):
  # '__getattribute__',
 
     def __getitem__(self, key):
-        if isinstance(key, dpnp_array):
-            key = key.get_array()
+        key = _get_unwrapped_index_key(key)
 
         item = self._array_obj.__getitem__(key)
         if not isinstance(item, dpt.usm_ndarray):
@@ -340,8 +356,8 @@ def __rxor__(self, other):
  # '__setattr__',
 
     def __setitem__(self, key, val):
-        if isinstance(key, dpnp_array):
-            key = key.get_array()
+        key = _get_unwrapped_index_key(key)
+
         if isinstance(val, dpnp_array):
             val = val.get_array()
 


Test patch:
diff --git a/tests/test_indexing.py b/tests/test_indexing.py
--- a/tests/test_indexing.py
+++ b/tests/test_indexing.py
@@ -6,10 +6,70 @@
 
 import numpy
 from numpy.testing import (
-    assert_array_equal
+    assert_,
+    assert_array_equal,
+    assert_equal
 )
 
 
+class TestIndexing:
+    def test_ellipsis_index(self):
+        a = dpnp.array([[1, 2, 3],
+                        [4, 5, 6],
+                        [7, 8, 9]])
+        assert_(a[...] is not a)
+        assert_equal(a[...], a)
+
+        # test that slicing with ellipsis doesn't skip an arbitrary number of dimensions
+        assert_equal(a[0, ...], a[0])
+        assert_equal(a[0, ...], a[0,:])
+        assert_equal(a[..., 0], a[:, 0])
+
+        # test that slicing with ellipsis always results in an array
+        assert_equal(a[0, ..., 1], dpnp.array(2))
+
+        # assignment with `(Ellipsis,)` on 0-d arrays
+        b = dpnp.array(1)
+        b[(Ellipsis,)] = 2
+        assert_equal(b, 2)
+
+    def test_boolean_indexing_list(self):
+        a = dpnp.array([1, 2, 3])
+        b = dpnp.array([True, False, True])
+
+        assert_equal(a[b], [1, 3])
+        assert_equal(a[None, b], [[1, 3]])
+
+    def test_indexing_array_weird_strides(self):
+        np_x = numpy.ones(10)
+        dp_x = dpnp.ones(10)
+
+        np_ind = numpy.arange(10)[:, None, None, None]
+        np_ind = numpy.broadcast_to(np_ind, (10, 55, 4, 4))
+
+        dp_ind = dpnp.arange(10)[:, None, None, None]
+        dp_ind = dpnp.broadcast_to(dp_ind, (10, 55, 4, 4))
+
+        # single advanced index case
+        assert_array_equal(dp_x[dp_ind], np_x[np_ind])
+
+        np_x2 = numpy.ones((10, 2))
+        dp_x2 = dpnp.ones((10, 2))
+
+        np_zind = numpy.zeros(4, dtype=np_ind.dtype)
+        dp_zind = dpnp.zeros(4, dtype=dp_ind.dtype)
+
+        # higher dimensional advanced index
+        assert_array_equal(dp_x2[dp_ind, dp_zind], np_x2[np_ind, np_zind])
+
+    def test_indexing_array_negative_strides(self):
+        arr = dpnp.zeros((4, 4))[::-1, ::-1]
+
+        slices = (slice(None), dpnp.array([0, 1, 2, 3]))
+        arr[slices] = 10
+        assert_array_equal(arr, 10.)
+
+
 @pytest.mark.usefixtures("allow_fall_back_on_numpy")
 def test_choose():
     a = numpy.r_[:4]