From ddb68e88db0fc25e04e928cf05cd7ca1554aea52 Mon Sep 17 00:00:00 2001 From: Nikita Grigorian Date: Mon, 3 Oct 2022 12:49:10 -0700 Subject: [PATCH 1/3] Implemented dpctl.tensor.meshgrid and tests --- dpctl/tensor/__init__.py | 2 ++ dpctl/tensor/_ctors.py | 50 ++++++++++++++++++++++++++++ dpctl/tests/test_usm_ndarray_ctor.py | 26 +++++++++++++++ 3 files changed, 78 insertions(+) diff --git a/dpctl/tensor/__init__.py b/dpctl/tensor/__init__.py index 5a31a05b61..05778d7436 100644 --- a/dpctl/tensor/__init__.py +++ b/dpctl/tensor/__init__.py @@ -31,6 +31,7 @@ full, full_like, linspace, + meshgrid, ones, ones_like, tril, @@ -87,4 +88,5 @@ "from_dlpack", "tril", "triu", + "meshgrid", ] diff --git a/dpctl/tensor/_ctors.py b/dpctl/tensor/_ctors.py index 0e50f38f05..6d5b6bc146 100644 --- a/dpctl/tensor/_ctors.py +++ b/dpctl/tensor/_ctors.py @@ -1198,3 +1198,53 @@ def triu(X, k=0): hev.wait() return res + + +def meshgrid(*arrays, indexing="xy"): + + """ + meshgrid(*arrays, indexing="xy") -> list[usm_ndarray] + + Creates list of `usm_ndarray` coordinate matrices from vectors. + + Args: + arrays: arbitrary number of one-dimensional `USM_ndarray` objects. + If vectors are not of the same data type, + or are not one-dimensional, raises `ValueError.` + indexing: Cartesian (`xy`) or matrix (`ij`) indexing of output. + For a set of `n` vectors with lengths X0, X1, X2, ... + Cartesian indexing results in arrays of shape + (X1, X0, X2, ...) + matrix indexing results in arrays of shape + (X0, X1, X2, ...) + Default: `xy`. + """ + for array in arrays: + if not isinstance(array, dpt.usm_ndarray): + raise TypeError( + f"Expected instance of dpt.usm_ndarray, got {type(array)}." + ) + if array.ndim != 1: + raise ValueError("All arrays must be one-dimensional.") + if len(set([array.dtype for array in arrays])) > 1: + raise ValueError("All arrays must be of the same numeric data type.") + if indexing not in ["xy", "ij"]: + raise ValueError( + "Unrecognized indexing keyword value, expecting 'xy' or 'ij.'" + ) + n = len(arrays) + sh = (-1,) + (1,) * (n - 1) + + res = [] + if n > 1 and indexing == "xy": + res.append(dpt.reshape(arrays[0], (1, -1) + sh[2:], copy=True)) + res.append(dpt.reshape(arrays[1], sh, copy=True)) + arrays, sh = arrays[2:], sh[-2:] + sh[:-2] + + for array in arrays: + res.append(dpt.reshape(array, sh, copy=True)) + sh = sh[-1:] + sh[:-1] + + output = dpt.broadcast_arrays(*res) + + return output diff --git a/dpctl/tests/test_usm_ndarray_ctor.py b/dpctl/tests/test_usm_ndarray_ctor.py index feabecb1ff..770a6fe9eb 100644 --- a/dpctl/tests/test_usm_ndarray_ctor.py +++ b/dpctl/tests/test_usm_ndarray_ctor.py @@ -1427,6 +1427,30 @@ def test_tril_order_k(order, k): assert np.array_equal(Ynp, dpt.asnumpy(Y)) +def test_meshgrid(): + try: + q = dpctl.SyclQueue() + except dpctl.SyclQueueCreationError: + pytest.skip("Queue could not be created") + X = dpt.arange(5, sycl_queue=q) + Y = dpt.arange(3, sycl_queue=q) + Z = dpt.meshgrid(X, Y) + Znp = np.meshgrid(dpt.asnumpy(X), dpt.asnumpy(Y)) + n = len(Z) + assert n == len(Znp) + for i in range(n): + assert np.array_equal(dpt.asnumpy(Z[i]), Znp[i]) + # dimension > 1 must raise ValueError + with pytest.raises(ValueError): + dpt.meshgrid(dpt.usm_ndarray((4, 4))) + # unknown indexing kwarg must raise ValueError + with pytest.raises(ValueError): + dpt.meshgrid(X, indexing="ji") + # input arrays with different data types must raise ValueError + with pytest.raises(ValueError): + dpt.meshgrid(X, dpt.asarray(Y, dtype="b1")) + + def test_common_arg_validation(): order = "I" # invalid order must raise ValueError @@ -1463,3 +1487,5 @@ def test_common_arg_validation(): dpt.tril(X) with pytest.raises(TypeError): dpt.triu(X) + with pytest.raises(TypeError): + dpt.meshgrid(X) From 95e52ff3b887458abc63062e2da6003268023fe2 Mon Sep 17 00:00:00 2001 From: Oleksandr Pavlyk Date: Sun, 9 Oct 2022 14:06:56 -0500 Subject: [PATCH 2/3] Applied PR suggestions --- dpctl/tensor/_ctors.py | 18 +++++++++++++----- 1 file changed, 13 insertions(+), 5 deletions(-) diff --git a/dpctl/tensor/_ctors.py b/dpctl/tensor/_ctors.py index 6d5b6bc146..5d2ca8e303 100644 --- a/dpctl/tensor/_ctors.py +++ b/dpctl/tensor/_ctors.py @@ -1212,13 +1212,15 @@ def meshgrid(*arrays, indexing="xy"): If vectors are not of the same data type, or are not one-dimensional, raises `ValueError.` indexing: Cartesian (`xy`) or matrix (`ij`) indexing of output. - For a set of `n` vectors with lengths X0, X1, X2, ... + For a set of `n` vectors with lengths N0, N1, N2, ... Cartesian indexing results in arrays of shape - (X1, X0, X2, ...) + (N1, N0, N2, ...) matrix indexing results in arrays of shape - (X0, X1, X2, ...) + (n0, N1, N2, ...) Default: `xy`. """ + ref_dt = None + ref_unset = True for array in arrays: if not isinstance(array, dpt.usm_ndarray): raise TypeError( @@ -1226,8 +1228,14 @@ def meshgrid(*arrays, indexing="xy"): ) if array.ndim != 1: raise ValueError("All arrays must be one-dimensional.") - if len(set([array.dtype for array in arrays])) > 1: - raise ValueError("All arrays must be of the same numeric data type.") + if ref_unset: + ref_unset = False + ref_dt = array.dtype + else: + if not ref_dt == array.dtype: + raise ValueError( + "All arrays must be of the same numeric data type." + ) if indexing not in ["xy", "ij"]: raise ValueError( "Unrecognized indexing keyword value, expecting 'xy' or 'ij.'" From 61a7244b4cb70d05aedb35f254a730135f6c7f4a Mon Sep 17 00:00:00 2001 From: Oleksandr Pavlyk Date: Sun, 9 Oct 2022 14:07:34 -0500 Subject: [PATCH 3/3] Added test_meshgrid2 to test different queues --- dpctl/tests/test_usm_ndarray_ctor.py | 27 +++++++++++++++++++++++++++ 1 file changed, 27 insertions(+) diff --git a/dpctl/tests/test_usm_ndarray_ctor.py b/dpctl/tests/test_usm_ndarray_ctor.py index 770a6fe9eb..e4d41a11e3 100644 --- a/dpctl/tests/test_usm_ndarray_ctor.py +++ b/dpctl/tests/test_usm_ndarray_ctor.py @@ -1451,6 +1451,33 @@ def test_meshgrid(): dpt.meshgrid(X, dpt.asarray(Y, dtype="b1")) +def test_meshgrid2(): + try: + q1 = dpctl.SyclQueue() + q2 = dpctl.SyclQueue() + q3 = dpctl.SyclQueue() + except dpctl.SyclQueueCreationError: + pytest.skip("Queue could not be created") + x1 = dpt.arange(0, 2, dtype="int16", sycl_queue=q1) + x2 = dpt.arange(3, 6, dtype="int16", sycl_queue=q2) + x3 = dpt.arange(6, 10, dtype="int16", sycl_queue=q3) + y1, y2, y3 = dpt.meshgrid(x1, x2, x3, indexing="xy") + z1, z2, z3 = dpt.meshgrid(x1, x2, x3, indexing="ij") + assert all( + x.sycl_queue == y.sycl_queue for x, y in zip((x1, x2, x3), (y1, y2, y3)) + ) + assert all( + x.sycl_queue == z.sycl_queue for x, z in zip((x1, x2, x3), (z1, z2, z3)) + ) + assert y1.shape == y2.shape and y2.shape == y3.shape + assert z1.shape == z2.shape and z2.shape == z3.shape + assert y1.shape == (len(x2), len(x1), len(x3)) + assert z1.shape == (len(x1), len(x2), len(x3)) + # FIXME: uncomment out once gh-921 is merged + # assert all(z.flags["C"] for z in (z1, z2, z3)) + # assert all(y.flags["C"] for y in (y1, y2, y3)) + + def test_common_arg_validation(): order = "I" # invalid order must raise ValueError