Use meshgrid() function from dpctl.tensor. (#1300)

* Use meshgrid() function from dpctl.tensor.

* Added more tests for meshgrid() funcrion.

Author: npolina4

dpnp/dpnp_algo/dpnp_algo_arraycreation.pyx

@@ -41,7 +41,6 @@ __all__ += [
     "dpnp_identity",
     "dpnp_linspace",
     "dpnp_logspace",
-    "dpnp_meshgrid",
     "dpnp_ptp",
     "dpnp_trace",
     "dpnp_tri",
@@ -222,53 +221,6 @@ cpdef utils.dpnp_descriptor dpnp_logspace(start, stop, num, endpoint, base, dtyp
     return dpnp.get_dpnp_descriptor(dpnp.astype(dpnp.power(base, temp), dtype))
 
 
-cpdef list dpnp_meshgrid(xi, copy, sparse, indexing):
-    input_count = len(xi)
-
-    # simple case
-    if input_count == 0:
-        return []
-
-    # simple case
-    if input_count == 1:
-        return [dpnp_copy(dpnp.get_dpnp_descriptor(xi[0])).get_pyobj()]
-
-    shape_mult = 1
-    for i in range(input_count):
-        shape_mult = shape_mult * xi[i].size
-
-    shape_list = []
-    for i in range(input_count):
-        shape_list.append(xi[i].size)
-    if indexing == "xy":
-        temp = shape_list[0]
-        shape_list[0] = shape_list[1]
-        shape_list[1] = temp
-
-    steps = []
-    for i in range(input_count):
-        shape_mult = shape_mult // shape_list[i]
-        steps.append(shape_mult)
-    if indexing == "xy":
-        temp = steps[0]
-        steps[0] = steps[1]
-        steps[1] = temp
-
-    shape = tuple(shape_list)
-
-    cdef utils.dpnp_descriptor res_item
-    result = []
-    for i in range(input_count):
-        res_item = utils_py.create_output_descriptor_py(shape, xi[i].dtype, None)
-
-        for j in range(res_item.size):
-            res_item.get_pyobj()[j] = xi[i][(j // steps[i]) % xi[i].size]
-
-        result.append(res_item.get_pyobj())
-
-    return result
-
-
 cpdef dpnp_ptp(utils.dpnp_descriptor arr, axis=None):
     cdef shape_type_c shape_arr = arr.shape
     cdef shape_type_c output_shape

dpnp/dpnp_container.py

@@ -180,6 +180,15 @@ def eye(N,
     return dpnp_array(array_obj.shape, buffer=array_obj, order=order)
 
 
+def meshgrid(*xi, indexing="xy"):
+    """Creates list of `dpnp_array` coordinate matrices from vectors."""
+    if len(xi) == 0:
+        return []
+    arrays = tuple(x.get_array() if isinstance(x, dpnp_array) else 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,
          *,
          dtype=None,

dpnp/dpnp_iface_arraycreation.py

@@ -1010,8 +1010,10 @@ def meshgrid(*xi, copy=True, sparse=False, indexing='xy'):
 
     Limitations
     -----------
+    Each array instance from `xi` is supported as either :class:`dpnp.dpnp_array` 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.
 
     Examples
     --------
@@ -1045,17 +1047,16 @@ def meshgrid(*xi, copy=True, sparse=False, indexing='xy'):
 
     """
 
-    if not use_origin_backend():
-        # original limitation
-        if indexing not in ["ij", "xy"]:
-            checker_throw_value_error("meshgrid", "indexing", indexing, "'ij' or 'xy'")
-
-        if copy is not True:
-            checker_throw_value_error("meshgrid", "copy", copy, True)
-        if sparse is not False:
-            checker_throw_value_error("meshgrid", "sparse", sparse, False)
-
-        return dpnp_meshgrid(xi, copy, sparse, indexing)
+    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)
 
     return call_origin(numpy.meshgrid, xi, copy, sparse, indexing)
 

tests/skipped_tests.tbl

@@ -406,16 +406,10 @@ tests/third_party/cupy/creation_tests/test_ranges.py::TestMeshgrid_param_3_{copy
 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_4_{copy=True, indexing='xy', sparse=False}::test_meshgrid1
-tests/third_party/cupy/creation_tests/test_ranges.py::TestMeshgrid_param_4_{copy=True, indexing='xy', sparse=False}::test_meshgrid2
-tests/third_party/cupy/creation_tests/test_ranges.py::TestMeshgrid_param_4_{copy=True, indexing='xy', sparse=False}::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_6_{copy=True, indexing='ij', sparse=False}::test_meshgrid1
-tests/third_party/cupy/creation_tests/test_ranges.py::TestMeshgrid_param_6_{copy=True, indexing='ij', sparse=False}::test_meshgrid2
-tests/third_party/cupy/creation_tests/test_ranges.py::TestMeshgrid_param_6_{copy=True, indexing='ij', sparse=False}::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/skipped_tests_gpu.tbl

@@ -581,16 +581,10 @@ tests/third_party/cupy/creation_tests/test_ranges.py::TestMeshgrid_param_3_{copy
 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_4_{copy=True, indexing='xy', sparse=False}::test_meshgrid1
-tests/third_party/cupy/creation_tests/test_ranges.py::TestMeshgrid_param_4_{copy=True, indexing='xy', sparse=False}::test_meshgrid2
-tests/third_party/cupy/creation_tests/test_ranges.py::TestMeshgrid_param_4_{copy=True, indexing='xy', sparse=False}::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_6_{copy=True, indexing='ij', sparse=False}::test_meshgrid1
-tests/third_party/cupy/creation_tests/test_ranges.py::TestMeshgrid_param_6_{copy=True, indexing='ij', sparse=False}::test_meshgrid2
-tests/third_party/cupy/creation_tests/test_ranges.py::TestMeshgrid_param_6_{copy=True, indexing='ij', sparse=False}::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/test_arraycreation.py

@@ -507,7 +507,21 @@ def test_dpctl_tensor_input(func, args):
     new_args = [eval(val, {'x0' : x0}) for val in args]
     X = getattr(dpt, func)(*new_args)
     Y = getattr(dpnp, func)(*new_args)
-    if func is 'empty_like':
+    if func == 'empty_like':
         assert X.shape == Y.shape
     else:
         assert_array_equal(X, Y)
+
+
+@pytest.mark.parametrize("arrays",
+                         [[], [[1]], [[1, 2, 3], [4, 5, 6]], [[1, 2], [3, 4], [5, 6]]],
+                         ids=['[]', '[[1]]', '[[1, 2, 3], [4, 5, 6]]', '[[1, 2], [3, 4], [5, 6]]'])
+@pytest.mark.parametrize("dtype", get_all_dtypes(no_float16=False))
+@pytest.mark.parametrize("indexing",
+                         ["ij", "xy"],
+                         ids=["ij", "xy"])
+def test_meshgrid(arrays, dtype, indexing):
+    func = lambda xp, xi: xp.meshgrid(*xi, indexing=indexing)
+    a = tuple(numpy.array(array, dtype=dtype) for array in arrays)
+    ia = tuple(dpnp.array(array, dtype=dtype) for array in arrays)
+    assert_array_equal(func(numpy, a), func(dpnp, ia))

tests/test_sycl_queue.py

@@ -177,6 +177,20 @@ def test_tril_triu(func, device):
     assert_sycl_queue_equal(x.sycl_queue, x0.sycl_queue)
 
 
+@pytest.mark.parametrize("device_x",
+                          valid_devices,
+                          ids=[device.filter_string for device in valid_devices])
+@pytest.mark.parametrize("device_y",
+                          valid_devices,
+                          ids=[device.filter_string for device in valid_devices])
+def test_meshgrid(device_x, device_y):
+    x = dpnp.arange(100, device = device_x)
+    y = dpnp.arange(100, device = device_y)
+    z = dpnp.meshgrid(x, y)
+    assert_sycl_queue_equal(z[0].sycl_queue, x.sycl_queue)
+    assert_sycl_queue_equal(z[1].sycl_queue, y.sycl_queue)
+
+
 @pytest.mark.usefixtures("allow_fall_back_on_numpy")
 @pytest.mark.parametrize(
     "func,data",

tests/test_usm_type.py

@@ -117,3 +117,13 @@ def test_coerced_usm_types_logic_op(op, usm_type_x, usm_type_y):
     assert x.usm_type == zx.usm_type == usm_type_x
     assert y.usm_type == zy.usm_type == usm_type_y
     assert z.usm_type == du.get_coerced_usm_type([usm_type_x, usm_type_y])
+
+
+@pytest.mark.parametrize("usm_type_x", list_of_usm_types, ids=list_of_usm_types)
+@pytest.mark.parametrize("usm_type_y", list_of_usm_types, ids=list_of_usm_types)
+def test_meshgrid(usm_type_x, usm_type_y):
+    x = dp.arange(100, usm_type = usm_type_x)
+    y = dp.arange(100, usm_type = usm_type_y)
+    z = dp.meshgrid(x, y)
+    assert z[0].usm_type == usm_type_x
+    assert z[1].usm_type == usm_type_y