Update dpnp.indices function

dpnp/dpnp_algo/dpnp_algo_indexing.pxi

@@ -40,7 +40,6 @@ __all__ += [
     "dpnp_diag_indices",
     "dpnp_diagonal",
     "dpnp_fill_diagonal",
-    "dpnp_indices",
     "dpnp_putmask",
     "dpnp_select",
     "dpnp_tril_indices",
@@ -237,40 +236,6 @@ cpdef dpnp_fill_diagonal(dpnp_descriptor x1, val):
     c_dpctl.DPCTLEvent_Delete(event_ref)
 
 
-cpdef object dpnp_indices(dimensions):
-    len_dimensions = len(dimensions)
-    res_shape = []
-    res_shape.append(len_dimensions)
-    for i in range(len_dimensions):
-        res_shape.append(dimensions[i])
-
-    result = []
-    if len_dimensions == 1:
-        res = []
-        for i in range(dimensions[0]):
-            res.append(i)
-        result.append(res)
-    else:
-        res1 = []
-        for i in range(dimensions[0]):
-            res = []
-            for j in range(dimensions[1]):
-                res.append(i)
-            res1.append(res)
-        result.append(res1)
-
-        res2 = []
-        for i in range(dimensions[0]):
-            res = []
-            for j in range(dimensions[1]):
-                res.append(j)
-            res2.append(res)
-        result.append(res2)
-
-    dpnp_result = dpnp.array(result)
-    return dpnp_result
-
-
 cpdef dpnp_putmask(utils.dpnp_descriptor arr, utils.dpnp_descriptor mask, utils.dpnp_descriptor values):
     cdef int values_size = values.size
 

dpnp/dpnp_iface_indexing.py

@@ -355,31 +355,117 @@ def fill_diagonal(x1, val, wrap=False):
     return call_origin(numpy.fill_diagonal, x1, val, wrap, dpnp_inplace=True)
 
 
-def indices(dimensions, dtype=int, sparse=False):
+def indices(
+    dimensions,
+    dtype=int,
+    sparse=False,
+    device=None,
+    usm_type="device",
+    sycl_queue=None,
+):
     """
     Return an array representing the indices of a grid.
 
+    Compute an array where the subarrays contain index values 0, 1, …
+    varying only along the corresponding axis.
+
     For full documentation refer to :obj:`numpy.indices`.
 
-    Limitations
-    -----------
-    Parameters `dtype` and `sparse` are supported only with default values.
-    Parameter `dimensions` is supported with len <=2.
+    Parameters
+    ----------
+    dimensions : sequence of ints
+        The shape of the grid.
+    dtype : dtype, optional
+        Data type of the result.
+    sparse : boolean, optional
+        Return a sparse representation of the grid instead of a dense representation.
+        Default is ``False``.
+    device : {None, string, SyclDevice, SyclQueue}, optional
+        An array API concept of device where the output array is created.
+        The `device` can be ``None`` (the default), an OneAPI filter selector string,
+        an instance of :class:`dpctl.SyclDevice` corresponding to a non-partitioned SYCL device,
+        an instance of :class:`dpctl.SyclQueue`, or a `Device` object returned by
+        :obj:`dpnp.dpnp_array.dpnp_array.device` property.
+    usm_type : {"device", "shared", "host"}, optional
+        The type of SYCL USM allocation for the output array.
+    sycl_queue : {None, SyclQueue}, optional
+        A SYCL queue to use for output array allocation and copying.
+
+    Returns
+    -------
+    out : one dpnp.ndarray or tuple of dpnp.ndarray
+        If sparse is ``False``:
+        Returns one array of grid indices, grid.shape = (len(dimensions),) + tuple(dimensions).
+
+        If sparse is ``True``:
+        Returns a tuple of arrays, with grid[i].shape = (1, ..., 1, dimensions[i], 1, ..., 1)
+        with dimensions[i] in the ith place.
+
+    Examples
+    --------
+    >>> import dpnp as np
+    >>> grid = np.indices((2, 3))
+    >>> grid.shape
+    (2, 2, 3)
+    >>> grid[0]
+    array([[0, 0, 0],
+           [1, 1, 1]])
+    >>> grid[1]
+    array([[0, 1, 2],
+           [0, 1, 2]])
+
+    The indices can be used as an index into an array.
+
+    >>> x = np.arange(20).reshape(5, 4)
+    >>> row, col = np.indices((2, 3))
+    >>> x[row, col]
+    array([[0, 1, 2],
+           [4, 5, 6]])
+
+    Note that it would be more straightforward in the above example to
+    extract the required elements directly with ``x[:2, :3]``.
+    If sparse is set to ``True``, the grid will be returned in a sparse
+    representation.
+
+    >>> i, j = np.indices((2, 3), sparse=True)
+    >>> i.shape
+    (2, 1)
+    >>> j.shape
+    (1, 3)
+    >>> i
+    array([[0],
+           [1]])
+    >>> j
+    array([[0, 1, 2]])
 
     """
 
-    if not isinstance(dimensions, (tuple, list)):
-        pass
-    elif len(dimensions) > 2 or len(dimensions) == 0:
-        pass
-    elif dtype != int:
-        pass
-    elif sparse:
-        pass
+    dimensions = tuple(dimensions)
+    N = len(dimensions)
+    shape = (1,) * N
+    if sparse:
+        res = ()
     else:
-        return dpnp_indices(dimensions)
-
-    return call_origin(numpy.indices, dimensions, dtype, sparse)
+        res = dpnp.empty(
+            (N,) + dimensions,
+            dtype=dtype,
+            device=device,
+            usm_type=usm_type,
+            sycl_queue=sycl_queue,
+        )
+    for i, dim in enumerate(dimensions):
+        idx = dpnp.arange(
+            dim,
+            dtype=dtype,
+            device=device,
+            usm_type=usm_type,
+            sycl_queue=sycl_queue,
+        ).reshape(shape[:i] + (dim,) + shape[i + 1 :])
+        if sparse:
+            res = res + (idx,)
+        else:
+            res[i] = idx
+    return res
 
 
 def nonzero(x, /):

tests/test_indexing.py

@@ -363,10 +363,13 @@ def test_fill_diagonal(array, val):
         "[3, 2]",
     ],
 )
-def test_indices(dimension):
-    expected = numpy.indices(dimension)
-    result = dpnp.indices(dimension)
-    assert_array_equal(expected, result)
+@pytest.mark.parametrize("dtype", get_all_dtypes(no_bool=True))
+@pytest.mark.parametrize("sparse", [True, False], ids=["True", "False"])
+def test_indices(dimension, dtype, sparse):
+    expected = numpy.indices(dimension, dtype=dtype, sparse=sparse)
+    result = dpnp.indices(dimension, dtype=dtype, sparse=sparse)
+    for Xnp, X in zip(expected, result):
+        assert_array_equal(Xnp, X)
 
 
 @pytest.mark.parametrize(

tests/test_sycl_queue.py

@@ -1418,6 +1418,19 @@ def test_take(func, device):
     assert_sycl_queue_equal(result_queue, expected_queue)
 
 
+@pytest.mark.parametrize(
+    "device",
+    valid_devices,
+    ids=[device.filter_string for device in valid_devices],
+)
+@pytest.mark.parametrize("sparse", [True, False], ids=["True", "False"])
+def test_indices(device, sparse):
+    sycl_queue = dpctl.SyclQueue(device)
+    grid = dpnp.indices((2, 3), sparse=sparse, sycl_queue=sycl_queue)
+    for dpnp_array in grid:
+        assert_sycl_queue_equal(dpnp_array.sycl_queue, sycl_queue)
+
+
 @pytest.mark.parametrize(
     "device",
     valid_devices,

tests/test_usm_type.py

@@ -521,6 +521,14 @@ def test_take(func, usm_type_x, usm_type_ind):
     assert z.usm_type == du.get_coerced_usm_type([usm_type_x, usm_type_ind])
 
 
+@pytest.mark.parametrize("usm_type", list_of_usm_types, ids=list_of_usm_types)
+@pytest.mark.parametrize("sparse", [True, False], ids=["True", "False"])
+def test_indices(usm_type, sparse):
+    x = dp.indices((2, 3), sparse=sparse, usm_type=usm_type)
+    for i in x:
+        assert i.usm_type == usm_type
+
+
 @pytest.mark.parametrize(
     "usm_type_matrix", list_of_usm_types, ids=list_of_usm_types
 )

tests/third_party/cupy/indexing_tests/test_generate.py

@@ -7,7 +7,6 @@
 from tests.third_party.cupy import testing
 
 
-@pytest.mark.usefixtures("allow_fall_back_on_numpy")
 @testing.gpu
 class TestIndices(unittest.TestCase):
     @testing.for_all_dtypes()