Update dpnp.indices function (#1622)

* Update dpnp.indices function

* Clean up dpnp_indices

* address comments

* address comments

* cancel changes for test_indices_list3

* Added description of all parameters

* Implement dpnp.mgrid and dpnp.ogrid function (#1633)

* Implement dpnp.mgrid and dpnp.ogrid function

* address comments

* address comments

* Skip dtype check for Iris Xe

* Update test_ranges.py

* Added description of all parameters

* Resolve pre-commit issue

---------

Co-authored-by: Anton <100830759+antonwolfy@users.noreply.github.com>
Co-authored-by: Anton Volkov <antonwolfy@gmail.com>

Author: npolina4

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_algo/dpnp_arraycreation.py

@@ -1,5 +1,7 @@
+import math
 import operator
 
+import dpctl.utils as dpu
 import numpy
 
 import dpnp
@@ -10,6 +12,7 @@
     "dpnp_geomspace",
     "dpnp_linspace",
     "dpnp_logspace",
+    "dpnp_nd_grid",
 ]
 
 
@@ -256,3 +259,134 @@ def dpnp_logspace(
     if dtype is None:
         return dpnp.power(base, res)
     return dpnp.power(base, res).astype(dtype, copy=False)
+
+
+class dpnp_nd_grid:
+    """
+    Construct a multi-dimensional "meshgrid".
+
+    ``grid = dpnp_nd_grid()`` creates an instance which will return a mesh-grid
+    when indexed. The dimension and number of the output arrays are equal
+    to the number of indexing dimensions.  If the step length is not a
+    complex number, then the stop is not inclusive.
+
+    However, if the step length is a complex number (e.g. 5j), then the
+    integer part of its magnitude is interpreted as specifying the
+    number of points to create between the start and stop values, where
+    the stop value is inclusive.
+
+    If instantiated with an argument of ``sparse=True``, the mesh-grid is
+    open (or not fleshed out) so that only one-dimension of each returned
+    argument is greater than 1.
+
+    Parameters
+    ----------
+    sparse : bool, optional
+        Whether the grid is sparse or not. Default is False.
+
+    """
+
+    def __init__(
+        self, sparse=False, device=None, usm_type="device", sycl_queue=None
+    ):
+        dpu.validate_usm_type(usm_type, allow_none=False)
+        self.sparse = sparse
+        self.usm_type = usm_type
+        self.sycl_queue_normalized = dpnp.get_normalized_queue_device(
+            sycl_queue=sycl_queue, device=device
+        )
+
+    def __getitem__(self, key):
+        if isinstance(key, slice):
+            step = key.step
+            stop = key.stop
+            start = key.start
+            if start is None:
+                start = 0
+            if isinstance(step, complex):
+                step = abs(step)
+                length = int(step)
+                if step != 1:
+                    step = (stop - start) / float(step - 1)
+                stop = stop + step
+                return (
+                    dpnp.arange(
+                        0,
+                        length,
+                        1,
+                        dtype=dpnp.default_float_type(),
+                        usm_type=self.usm_type,
+                        sycl_queue=self.sycl_queue_normalized,
+                    )
+                    * step
+                    + start
+                )
+            else:
+                return dpnp.arange(
+                    start,
+                    stop,
+                    step,
+                    usm_type=self.usm_type,
+                    sycl_queue=self.sycl_queue_normalized,
+                )
+
+        size = []
+        dtype = int
+        for k in range(len(key)):
+            step = key[k].step
+            start = key[k].start
+            stop = key[k].stop
+            if start is None:
+                start = 0
+            if step is None:
+                step = 1
+            if isinstance(step, complex):
+                size.append(int(abs(step)))
+                dtype = dpnp.default_float_type()
+            else:
+                size.append(
+                    int(math.ceil((key[k].stop - start) / (step * 1.0)))
+                )
+            if (
+                isinstance(step, float)
+                or isinstance(start, float)
+                or isinstance(stop, float)
+            ):
+                dtype = dpnp.default_float_type()
+        if self.sparse:
+            nn = [
+                dpnp.arange(
+                    _x,
+                    dtype=_t,
+                    usm_type=self.usm_type,
+                    sycl_queue=self.sycl_queue_normalized,
+                )
+                for _x, _t in zip(size, (dtype,) * len(size))
+            ]
+        else:
+            nn = dpnp.indices(
+                size,
+                dtype,
+                usm_type=self.usm_type,
+                sycl_queue=self.sycl_queue_normalized,
+            )
+        for k in range(len(size)):
+            step = key[k].step
+            start = key[k].start
+            stop = key[k].stop
+            if start is None:
+                start = 0
+            if step is None:
+                step = 1
+            if isinstance(step, complex):
+                step = int(abs(step))
+                if step != 1:
+                    step = (stop - start) / float(step - 1)
+            nn[k] = nn[k] * step + start
+        if self.sparse:
+            slobj = [dpnp.newaxis] * len(size)
+            for k in range(len(size)):
+                slobj[k] = slice(None, None)
+                nn[k] = nn[k][tuple(slobj)]
+                slobj[k] = dpnp.newaxis
+        return nn

dpnp/dpnp_iface_arraycreation.py

@@ -52,6 +52,7 @@
     dpnp_geomspace,
     dpnp_linspace,
     dpnp_logspace,
+    dpnp_nd_grid,
 )
 
 __all__ = [
@@ -1452,6 +1453,24 @@ class MGridClass:
 
     For full documentation refer to :obj:`numpy.mgrid`.
 
+    Parameters
+    ----------
+    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
+        Returns one array of grid indices, grid.shape = (len(dimensions),) + tuple(dimensions).
+
     Examples
     --------
     >>> import dpnp as np
@@ -1466,13 +1485,31 @@ class MGridClass:
             [0, 1, 2, 3, 4],
             [0, 1, 2, 3, 4],
             [0, 1, 2, 3, 4]]])
-    >>> np.mgrid[-1:1:5j]
+
+    >>> x = np.mgrid[-1:1:5j]
+    >>> x
     array([-1. , -0.5,  0. ,  0.5,  1. ])
+    >>> x.usm_type
+    'device'
+
+    >>> y = np.mgrid(usm_type="host")[-1:1:5j]
+    >>> y
+    array([-1. , -0.5,  0. ,  0.5,  1. ])
+    >>> x.usm_type
+    'host'
 
     """
 
     def __getitem__(self, key):
-        return dpnp.array(numpy.mgrid[key])
+        return dpnp_nd_grid(sparse=False)[key]
+
+    def __call__(self, device=None, usm_type="device", sycl_queue=None):
+        return dpnp_nd_grid(
+            sparse=False,
+            device=device,
+            usm_type=usm_type,
+            sycl_queue=sycl_queue,
+        )
 
 
 mgrid = MGridClass()
@@ -1484,23 +1521,56 @@ class OGridClass:
 
     For full documentation refer to :obj:`numpy.ogrid`.
 
+    Parameters
+    ----------
+    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
+        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
-    >>> from numpy import ogrid
-    >>> ogrid[-1:1:5j]
-    array([-1. , -0.5,  0. ,  0.5,  1. ])
-    >>> ogrid[0:5,0:5]
+    >>> np.ogrid[0:5, 0:5]
     [array([[0],
             [1],
             [2],
             [3],
             [4]]), array([[0, 1, 2, 3, 4]])]
 
+    >>> x = np.ogrid[-1:1:5j]
+    >>> x
+    array([-1. , -0.5,  0. ,  0.5,  1. ])
+    >>> x.usm_type
+    'device'
+
+    >>> y = np.ogrid(usm_type="host")[-1:1:5j]
+    >>> y
+    array([-1. , -0.5,  0. ,  0.5,  1. ])
+    >>> x.usm_type
+    'host'
+
     """
 
     def __getitem__(self, key):
-        return dpnp.array(numpy.ogrid[key])
+        return dpnp_nd_grid(sparse=True)[key]
+
+    def __call__(self, device=None, usm_type="device", sycl_queue=None):
+        return dpnp_nd_grid(
+            sparse=True, device=device, usm_type=usm_type, sycl_queue=sycl_queue
+        )
 
 
 ogrid = OGridClass()