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_put_along_axis",
     "dpnp_putmask",
     "dpnp_select",
@@ -249,40 +248,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_put_along_axis(dpnp_descriptor arr, dpnp_descriptor indices, dpnp_descriptor values, int axis):
     cdef shape_type_c arr_shape = arr.shape
     cdef DPNPFuncType param1_type = dpnp_dtype_to_DPNPFuncType(arr.dtype)

dpnp/dpnp_iface_indexing.py

@@ -307,31 +307,101 @@ 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``.
 
-    """
+    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 not isinstance(dimensions, (tuple, list)):
-        pass
-    elif len(dimensions) > 2 or len(dimensions) == 0:
-        pass
-    elif dtype != int:
-        pass
-    elif sparse:
-        pass
-    else:
-        return dpnp_indices(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.
 
-    return call_origin(numpy.indices, dimensions, dtype, sparse)
+    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]])
+
+    >>> x = np.arange(20).reshape(5, 4)
+    >>> row, col = np.indices((2, 3))
+    >>> x[row, col]
+    array([[0, 1, 2],
+           [4, 5, 6]])
+
+    >>> 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]])
+
+    """
+
+    dimensions = tuple(dimensions)
+    N = len(dimensions)
+    shape = (1,) * N
+    sycl_queue_normalized = dpnp.get_normalized_queue_device(
+        sycl_queue=sycl_queue, device=device
+    )
+    if sparse:
+        res = ()
+    else:
+        res = dpnp.empty(
+            (N,) + dimensions,
+            dtype=dtype,
+            usm_type=usm_type,
+            sycl_queue=sycl_queue_normalized,
+        )
+    for i, dim in enumerate(dimensions):
+        idx = dpnp.arange(
+            dim,
+            dtype=dtype,
+            usm_type=usm_type,
+            sycl_queue=sycl_queue_normalized,
+        ).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

@@ -179,10 +179,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

@@ -1232,3 +1232,16 @@ def test_take(device):
     result_queue = result.get_array().sycl_queue
 
     assert_sycl_queue_equal(result_queue, expected_queue)
+
+
+@pytest.mark.parametrize(
+    "device",
+    valid_devices,
+    ids=[device.filter_string for device in valid_devices],
+)
+def test_indices(device):
+    dpnp_array = dpnp.indices((2,), device=device)
+    numpy_array = numpy.indices((2,))
+
+    assert_allclose(numpy_array, dpnp_array)
+    assert dpnp_array.sycl_device == device

tests/test_usm_type.py

@@ -480,3 +480,9 @@ def test_take(usm_type_x, usm_type_ind):
     assert x.usm_type == usm_type_x
     assert ind.usm_type == 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)
+def test_indices(usm_type):
+    x = dp.indices((2,), usm_type=usm_type)
+    assert x.usm_type == usm_type

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()
@@ -25,9 +24,10 @@ def test_indices_list1(self, xp, dtype):
     def test_indices_list2(self, xp, dtype):
         return xp.indices((1, 2, 3, 4), dtype)
 
+    @testing.with_requires("numpy>=1.24")
     def test_indices_list3(self):
         for xp in (numpy, cupy):
-            with pytest.raises((ValueError, TypeError)):
+            with pytest.raises(TypeError):
                 xp.indices((1, 2, 3, 4), dtype=xp.bool_)