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

Author: npolina4

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()

tests/skipped_tests.tbl

@@ -177,10 +177,6 @@ tests/third_party/cupy/creation_tests/test_ranges.py::TestMeshgrid_param_7_{copy
 tests/third_party/cupy/creation_tests/test_ranges.py::TestMeshgrid_param_7_{copy=True, indexing='ij', sparse=True}::test_meshgrid1
 tests/third_party/cupy/creation_tests/test_ranges.py::TestMeshgrid_param_7_{copy=True, indexing='ij', sparse=True}::test_meshgrid2
 tests/third_party/cupy/creation_tests/test_ranges.py::TestMeshgrid_param_7_{copy=True, indexing='ij', sparse=True}::test_meshgrid3
-tests/third_party/cupy/creation_tests/test_ranges.py::TestMgrid::test_mgrid3
-tests/third_party/cupy/creation_tests/test_ranges.py::TestOgrid::test_ogrid3
-tests/third_party/cupy/creation_tests/test_ranges.py::TestOgrid::test_ogrid4
-tests/third_party/cupy/creation_tests/test_ranges.py::TestOgrid::test_ogrid5
 tests/third_party/cupy/indexing_tests/test_generate.py::TestAxisConcatenator::test_AxisConcatenator_init1
 tests/third_party/cupy/indexing_tests/test_generate.py::TestAxisConcatenator::test_len
 tests/third_party/cupy/indexing_tests/test_generate.py::TestC_::test_c_1

tests/skipped_tests_gpu.tbl

@@ -255,11 +255,6 @@ tests/third_party/cupy/creation_tests/test_ranges.py::TestMeshgrid_param_7_{copy
 tests/third_party/cupy/creation_tests/test_ranges.py::TestMeshgrid_param_7_{copy=True, indexing='ij', sparse=True}::test_meshgrid1
 tests/third_party/cupy/creation_tests/test_ranges.py::TestMeshgrid_param_7_{copy=True, indexing='ij', sparse=True}::test_meshgrid2
 tests/third_party/cupy/creation_tests/test_ranges.py::TestMeshgrid_param_7_{copy=True, indexing='ij', sparse=True}::test_meshgrid3
-tests/third_party/cupy/creation_tests/test_ranges.py::TestMgrid::test_mgrid3
-tests/third_party/cupy/creation_tests/test_ranges.py::TestMgrid::test_mgrid5
-tests/third_party/cupy/creation_tests/test_ranges.py::TestOgrid::test_ogrid3
-tests/third_party/cupy/creation_tests/test_ranges.py::TestOgrid::test_ogrid4
-tests/third_party/cupy/creation_tests/test_ranges.py::TestOgrid::test_ogrid5
 tests/third_party/cupy/creation_tests/test_ranges.py::TestRanges::test_arange_negative_size
 tests/third_party/cupy/creation_tests/test_ranges.py::TestRanges::test_arange_no_dtype_int
 

tests/test_sycl_queue.py

@@ -1431,6 +1431,18 @@ def test_indices(device, sparse):
         assert_sycl_queue_equal(dpnp_array.sycl_queue, sycl_queue)
 
 
+@pytest.mark.parametrize(
+    "device",
+    valid_devices,
+    ids=[device.filter_string for device in valid_devices],
+)
+@pytest.mark.parametrize("func", ["mgrid", "ogrid"])
+def test_grid(device, func):
+    sycl_queue = dpctl.SyclQueue(device)
+    x = getattr(dpnp, func)(sycl_queue=sycl_queue)[0:4]
+    assert_sycl_queue_equal(x.sycl_queue, sycl_queue)
+
+
 @pytest.mark.parametrize(
     "device",
     valid_devices,

tests/test_usm_type.py

@@ -521,6 +521,18 @@ 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)
+def test_indices(usm_type):
+    x = dp.indices((2,), usm_type=usm_type)
+    assert x.usm_type == usm_type
+
+
+@pytest.mark.parametrize("usm_type", list_of_usm_types, ids=list_of_usm_types)
+@pytest.mark.parametrize("func", ["mgrid", "ogrid"])
+def test_grid(usm_type, func):
+    assert getattr(dp, func)(usm_type=usm_type)[0:4].usm_type == usm_type
+
+
 @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):

tests/third_party/cupy/creation_tests/test_ranges.py

@@ -363,7 +363,7 @@ def test_mgrid1(self, xp):
     def test_mgrid2(self, xp):
         return xp.mgrid[-10:10:10j]
 
-    @testing.numpy_cupy_array_equal()
+    @testing.numpy_cupy_allclose(rtol=1e-4, type_check=has_support_aspect64())
     def test_mgrid3(self, xp):
         x = xp.zeros(10)[:, None]
         y = xp.ones(10)[:, None]
@@ -374,7 +374,7 @@ def test_mgrid4(self, xp):
         # check len(keys) > 1
         return xp.mgrid[-10:10:10j, -10:10:10j]
 
-    @testing.numpy_cupy_array_equal()
+    @testing.numpy_cupy_allclose(rtol=1e-4, type_check=has_support_aspect64())
     def test_mgrid5(self, xp):
         # check len(keys) > 1
         x = xp.zeros(10)[:, None]
@@ -396,18 +396,18 @@ def test_ogrid1(self, xp):
     def test_ogrid2(self, xp):
         return xp.ogrid[-10:10:10j]
 
-    @testing.numpy_cupy_array_equal()
+    @testing.numpy_cupy_allclose(rtol=1e-4, type_check=has_support_aspect64())
     def test_ogrid3(self, xp):
         x = xp.zeros(10)[:, None]
         y = xp.ones(10)[:, None]
         return xp.ogrid[x:y:10j]
 
-    @testing.numpy_cupy_array_equal()
+    @testing.numpy_cupy_allclose(rtol=1e-4, type_check=has_support_aspect64())
     def test_ogrid4(self, xp):
         # check len(keys) > 1
         return xp.ogrid[-10:10:10j, -10:10:10j]
 
-    @testing.numpy_cupy_array_equal()
+    @testing.numpy_cupy_allclose(rtol=1e-4, type_check=has_support_aspect64())
     def test_ogrid5(self, xp):
         # check len(keys) > 1
         x = xp.zeros(10)[:, None]