Implement sparse and copy arguments for dpnp.mesgrid function (#1675)

* Implement sparse and copy arguments for dpnp.mesgrid function

* address comments

* Removed limitation block from th description

* added tests

---------

Co-authored-by: Anton Volkov <[email protected]>
Co-authored-by: Anton <[email protected]>

Author: npolina4

dpnp/dpnp_container.py

@@ -271,18 +271,6 @@ def linspace(
     return dpnp_array(array_obj.shape, buffer=array_obj)
 
 
-def meshgrid(*xi, indexing="xy"):
-    """Creates list of `dpnp_array` coordinate matrices from vectors."""
-    if len(xi) == 0:
-        return []
-    arrays = tuple(dpnp.get_usm_ndarray(x) for x in xi)
-    arrays_obj = dpt.meshgrid(*arrays, indexing=indexing)
-    return [
-        dpnp_array._create_from_usm_ndarray(array_obj)
-        for array_obj in arrays_obj
-    ]
-
-
 def ones(
     shape,
     *,

dpnp/dpnp_iface_arraycreation.py

@@ -1394,12 +1394,28 @@ def meshgrid(*xi, copy=True, sparse=False, indexing="xy"):
 
     For full documentation refer to :obj:`numpy.meshgrid`.
 
-    Limitations
-    -----------
-    Each array instance from `xi` is supported as either :class:`dpnp.dpnp.ndarray` or :class:`dpctl.tensor.usm_ndarray`.
-    Parameter `copy` is supported only with default value ``True``.
-    Parameter `sparse` is supported only with default value ``False``.
-    Otherwise the function will be executed sequentially on CPU.
+    Parameters
+    ----------
+    x1, x2,..., xn : {dpnp.ndarray, usm_ndarray}
+        1-D arrays representing the coordinates of a grid.
+    indexing : {'xy', 'ij'}, optional
+        Cartesian ('xy', default) or matrix ('ij') indexing of output.
+    sparse : bool, optional
+        If True the shape of the returned coordinate array for dimension `i`
+        is reduced from ``(N1, ..., Ni, ... Nn)`` to
+        ``(1, ..., 1, Ni, 1, ..., 1)``. Default is False.
+    copy : bool, optional
+        If False, a view into the original arrays are returned in order to
+        conserve memory.  Default is True.
+
+    Returns
+    -------
+    X1, X2,..., XN : tuple of dpnp.ndarrays
+        For vectors `x1`, `x2`,..., `xn` with lengths ``Ni=len(xi)``,
+        returns ``(N1, N2, N3,..., Nn)`` shaped arrays if indexing='ij'
+        or ``(N2, N1, N3,..., Nn)`` shaped arrays if indexing='xy'
+        with the elements of `xi` repeated to fill the matrix along
+        the first dimension for `x1`, the second for `x2` and so on.
 
     Examples
     --------
@@ -1433,18 +1449,32 @@ def meshgrid(*xi, copy=True, sparse=False, indexing="xy"):
 
     """
 
-    if not all((isinstance(x, (dpnp.ndarray, dpt.usm_ndarray)) for x in xi)):
-        pass
-    elif indexing not in ["ij", "xy"]:
-        pass
-    elif copy is not True:
-        pass
-    elif sparse is not False:
-        pass
-    else:
-        return dpnp_container.meshgrid(*xi, indexing=indexing)
+    if not dpnp.check_supported_arrays_type(*xi):
+        raise TypeError("Each input array must be any of supported type")
+
+    ndim = len(xi)
+
+    if indexing not in ["xy", "ij"]:
+        raise ValueError(
+            "Unrecognized indexing keyword value, expecting 'xy' or 'ij'."
+        )
+
+    s0 = (1,) * ndim
+    output = [
+        dpnp.reshape(x, s0[:i] + (-1,) + s0[i + 1 :]) for i, x in enumerate(xi)
+    ]
+
+    if indexing == "xy" and ndim > 1:
+        output[0] = output[0].reshape((1, -1) + s0[2:])
+        output[1] = output[1].reshape((-1, 1) + s0[2:])
+
+    if not sparse:
+        output = dpnp.broadcast_arrays(*output)
+
+    if copy:
+        output = [x.copy() for x in output]
 
-    return call_origin(numpy.meshgrid, xi, copy, sparse, indexing)
+    return output
 
 
 class MGridClass:

tests/skipped_tests.tbl

@@ -151,30 +151,6 @@ tests/third_party/cupy/creation_tests/test_basic.py::TestBasic::test_ones_like_s
 tests/third_party/cupy/creation_tests/test_basic.py::TestBasic::test_zeros_like_subok
 tests/third_party/cupy/creation_tests/test_basic.py::TestBasic::test_zeros_strides
 
-tests/third_party/cupy/creation_tests/test_ranges.py::TestMeshgrid_param_0_{copy=False, indexing='xy', sparse=False}::test_meshgrid0
-tests/third_party/cupy/creation_tests/test_ranges.py::TestMeshgrid_param_0_{copy=False, indexing='xy', sparse=False}::test_meshgrid1
-tests/third_party/cupy/creation_tests/test_ranges.py::TestMeshgrid_param_0_{copy=False, indexing='xy', sparse=False}::test_meshgrid2
-tests/third_party/cupy/creation_tests/test_ranges.py::TestMeshgrid_param_0_{copy=False, indexing='xy', sparse=False}::test_meshgrid3
-tests/third_party/cupy/creation_tests/test_ranges.py::TestMeshgrid_param_1_{copy=False, indexing='xy', sparse=True}::test_meshgrid0
-tests/third_party/cupy/creation_tests/test_ranges.py::TestMeshgrid_param_1_{copy=False, indexing='xy', sparse=True}::test_meshgrid1
-tests/third_party/cupy/creation_tests/test_ranges.py::TestMeshgrid_param_1_{copy=False, indexing='xy', sparse=True}::test_meshgrid2
-tests/third_party/cupy/creation_tests/test_ranges.py::TestMeshgrid_param_1_{copy=False, indexing='xy', sparse=True}::test_meshgrid3
-tests/third_party/cupy/creation_tests/test_ranges.py::TestMeshgrid_param_2_{copy=False, indexing='ij', sparse=False}::test_meshgrid0
-tests/third_party/cupy/creation_tests/test_ranges.py::TestMeshgrid_param_2_{copy=False, indexing='ij', sparse=False}::test_meshgrid1
-tests/third_party/cupy/creation_tests/test_ranges.py::TestMeshgrid_param_2_{copy=False, indexing='ij', sparse=False}::test_meshgrid2
-tests/third_party/cupy/creation_tests/test_ranges.py::TestMeshgrid_param_2_{copy=False, indexing='ij', sparse=False}::test_meshgrid3
-tests/third_party/cupy/creation_tests/test_ranges.py::TestMeshgrid_param_3_{copy=False, indexing='ij', sparse=True}::test_meshgrid0
-tests/third_party/cupy/creation_tests/test_ranges.py::TestMeshgrid_param_3_{copy=False, indexing='ij', sparse=True}::test_meshgrid1
-tests/third_party/cupy/creation_tests/test_ranges.py::TestMeshgrid_param_3_{copy=False, indexing='ij', sparse=True}::test_meshgrid2
-tests/third_party/cupy/creation_tests/test_ranges.py::TestMeshgrid_param_3_{copy=False, indexing='ij', sparse=True}::test_meshgrid3
-tests/third_party/cupy/creation_tests/test_ranges.py::TestMeshgrid_param_5_{copy=True, indexing='xy', sparse=True}::test_meshgrid0
-tests/third_party/cupy/creation_tests/test_ranges.py::TestMeshgrid_param_5_{copy=True, indexing='xy', sparse=True}::test_meshgrid1
-tests/third_party/cupy/creation_tests/test_ranges.py::TestMeshgrid_param_5_{copy=True, indexing='xy', sparse=True}::test_meshgrid2
-tests/third_party/cupy/creation_tests/test_ranges.py::TestMeshgrid_param_5_{copy=True, indexing='xy', sparse=True}::test_meshgrid3
-tests/third_party/cupy/creation_tests/test_ranges.py::TestMeshgrid_param_7_{copy=True, indexing='ij', sparse=True}::test_meshgrid0
-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/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

@@ -230,30 +230,6 @@ tests/third_party/cupy/creation_tests/test_basic.py::TestBasic::test_ones_like_s
 tests/third_party/cupy/creation_tests/test_basic.py::TestBasic::test_zeros_like_subok
 tests/third_party/cupy/creation_tests/test_basic.py::TestBasic::test_zeros_strides
 
-tests/third_party/cupy/creation_tests/test_ranges.py::TestMeshgrid_param_0_{copy=False, indexing='xy', sparse=False}::test_meshgrid0
-tests/third_party/cupy/creation_tests/test_ranges.py::TestMeshgrid_param_0_{copy=False, indexing='xy', sparse=False}::test_meshgrid1
-tests/third_party/cupy/creation_tests/test_ranges.py::TestMeshgrid_param_0_{copy=False, indexing='xy', sparse=False}::test_meshgrid2
-tests/third_party/cupy/creation_tests/test_ranges.py::TestMeshgrid_param_0_{copy=False, indexing='xy', sparse=False}::test_meshgrid3
-tests/third_party/cupy/creation_tests/test_ranges.py::TestMeshgrid_param_1_{copy=False, indexing='xy', sparse=True}::test_meshgrid0
-tests/third_party/cupy/creation_tests/test_ranges.py::TestMeshgrid_param_1_{copy=False, indexing='xy', sparse=True}::test_meshgrid1
-tests/third_party/cupy/creation_tests/test_ranges.py::TestMeshgrid_param_1_{copy=False, indexing='xy', sparse=True}::test_meshgrid2
-tests/third_party/cupy/creation_tests/test_ranges.py::TestMeshgrid_param_1_{copy=False, indexing='xy', sparse=True}::test_meshgrid3
-tests/third_party/cupy/creation_tests/test_ranges.py::TestMeshgrid_param_2_{copy=False, indexing='ij', sparse=False}::test_meshgrid0
-tests/third_party/cupy/creation_tests/test_ranges.py::TestMeshgrid_param_2_{copy=False, indexing='ij', sparse=False}::test_meshgrid1
-tests/third_party/cupy/creation_tests/test_ranges.py::TestMeshgrid_param_2_{copy=False, indexing='ij', sparse=False}::test_meshgrid2
-tests/third_party/cupy/creation_tests/test_ranges.py::TestMeshgrid_param_2_{copy=False, indexing='ij', sparse=False}::test_meshgrid3
-tests/third_party/cupy/creation_tests/test_ranges.py::TestMeshgrid_param_3_{copy=False, indexing='ij', sparse=True}::test_meshgrid0
-tests/third_party/cupy/creation_tests/test_ranges.py::TestMeshgrid_param_3_{copy=False, indexing='ij', sparse=True}::test_meshgrid1
-tests/third_party/cupy/creation_tests/test_ranges.py::TestMeshgrid_param_3_{copy=False, indexing='ij', sparse=True}::test_meshgrid2
-tests/third_party/cupy/creation_tests/test_ranges.py::TestMeshgrid_param_3_{copy=False, indexing='ij', sparse=True}::test_meshgrid3
-tests/third_party/cupy/creation_tests/test_ranges.py::TestMeshgrid_param_5_{copy=True, indexing='xy', sparse=True}::test_meshgrid0
-tests/third_party/cupy/creation_tests/test_ranges.py::TestMeshgrid_param_5_{copy=True, indexing='xy', sparse=True}::test_meshgrid1
-tests/third_party/cupy/creation_tests/test_ranges.py::TestMeshgrid_param_5_{copy=True, indexing='xy', sparse=True}::test_meshgrid2
-tests/third_party/cupy/creation_tests/test_ranges.py::TestMeshgrid_param_5_{copy=True, indexing='xy', sparse=True}::test_meshgrid3
-tests/third_party/cupy/creation_tests/test_ranges.py::TestMeshgrid_param_7_{copy=True, indexing='ij', sparse=True}::test_meshgrid0
-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::TestRanges::test_arange_negative_size
 tests/third_party/cupy/creation_tests/test_ranges.py::TestRanges::test_arange_no_dtype_int
 

tests/test_arraycreation.py

@@ -878,3 +878,12 @@ def test_logspace_axis(axis):
         [2, 3], [20, 15], num=2, base=[[1, 3], [5, 7]], axis=axis
     )
     assert_dtype_allclose(func(dpnp), func(numpy))
+
+
+def test_meshgrid_raise_error():
+    a = numpy.array([1, 2, 3, 4])
+    with pytest.raises(TypeError):
+        dpnp.meshgrid(a)
+    b = dpnp.array([1, 2, 3, 4])
+    with pytest.raises(ValueError):
+        dpnp.meshgrid(b, indexing="ab")