diff --git a/dpnp/dpnp_algo/dpnp_algo_indexing.pxi b/dpnp/dpnp_algo/dpnp_algo_indexing.pxi index 82f14cff52bd..ab8eceeb415f 100644 --- a/dpnp/dpnp_algo/dpnp_algo_indexing.pxi +++ b/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 diff --git a/dpnp/dpnp_algo/dpnp_arraycreation.py b/dpnp/dpnp_algo/dpnp_arraycreation.py index b6d3c6120683..0399deea254a 100644 --- a/dpnp/dpnp_algo/dpnp_arraycreation.py +++ b/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 diff --git a/dpnp/dpnp_iface_arraycreation.py b/dpnp/dpnp_iface_arraycreation.py index cd80c239c336..9dfe2ef13d20 100644 --- a/dpnp/dpnp_iface_arraycreation.py +++ b/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() diff --git a/dpnp/dpnp_iface_indexing.py b/dpnp/dpnp_iface_indexing.py index 202db50cc3d8..20a4a60a5bcd 100644 --- a/dpnp/dpnp_iface_indexing.py +++ b/dpnp/dpnp_iface_indexing.py @@ -355,31 +355,117 @@ def fill_diagonal(x1, val, wrap=False): return call_origin(numpy.fill_diagonal, x1, val, wrap, dpnp_inplace=True) -def indices(dimensions, dtype=int, sparse=False): +def indices( + dimensions, + dtype=int, + sparse=False, + device=None, + usm_type="device", + sycl_queue=None, +): """ Return an array representing the indices of a grid. + Compute an array where the subarrays contain index values 0, 1, … + varying only along the corresponding axis. + For full documentation refer to :obj:`numpy.indices`. - Limitations - ----------- - Parameters `dtype` and `sparse` are supported only with default values. - Parameter `dimensions` is supported with len <=2. + Parameters + ---------- + dimensions : sequence of ints + The shape of the grid. + dtype : dtype, optional + Data type of the result. + sparse : boolean, optional + Return a sparse representation of the grid instead of a dense representation. + Default is ``False``. + 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 + If sparse is ``False``: + Returns one array of grid indices, grid.shape = (len(dimensions),) + tuple(dimensions). + + If sparse is ``True``: + 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 + >>> grid = np.indices((2, 3)) + >>> grid.shape + (2, 2, 3) + >>> grid[0] + array([[0, 0, 0], + [1, 1, 1]]) + >>> grid[1] + array([[0, 1, 2], + [0, 1, 2]]) + + The indices can be used as an index into an array. + + >>> x = np.arange(20).reshape(5, 4) + >>> row, col = np.indices((2, 3)) + >>> x[row, col] + array([[0, 1, 2], + [4, 5, 6]]) + + Note that it would be more straightforward in the above example to + extract the required elements directly with ``x[:2, :3]``. + If sparse is set to ``True``, the grid will be returned in a sparse + representation. + + >>> i, j = np.indices((2, 3), sparse=True) + >>> i.shape + (2, 1) + >>> j.shape + (1, 3) + >>> i + array([[0], + [1]]) + >>> j + array([[0, 1, 2]]) """ - if not isinstance(dimensions, (tuple, list)): - pass - elif len(dimensions) > 2 or len(dimensions) == 0: - pass - elif dtype != int: - pass - elif sparse: - pass + dimensions = tuple(dimensions) + N = len(dimensions) + shape = (1,) * N + if sparse: + res = () else: - return dpnp_indices(dimensions) - - return call_origin(numpy.indices, dimensions, dtype, sparse) + res = dpnp.empty( + (N,) + dimensions, + dtype=dtype, + device=device, + usm_type=usm_type, + sycl_queue=sycl_queue, + ) + for i, dim in enumerate(dimensions): + idx = dpnp.arange( + dim, + dtype=dtype, + device=device, + usm_type=usm_type, + sycl_queue=sycl_queue, + ).reshape(shape[:i] + (dim,) + shape[i + 1 :]) + if sparse: + res = res + (idx,) + else: + res[i] = idx + return res def nonzero(x, /): diff --git a/tests/skipped_tests.tbl b/tests/skipped_tests.tbl index e4c88a7298af..8958c7f8f217 100644 --- a/tests/skipped_tests.tbl +++ b/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 diff --git a/tests/skipped_tests_gpu.tbl b/tests/skipped_tests_gpu.tbl index 1ef5b6132431..a6a2886ad86a 100644 --- a/tests/skipped_tests_gpu.tbl +++ b/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 diff --git a/tests/test_indexing.py b/tests/test_indexing.py index 3cb9d4298d61..01e5c33f8efa 100644 --- a/tests/test_indexing.py +++ b/tests/test_indexing.py @@ -364,10 +364,13 @@ def test_fill_diagonal(array, val): "[3, 2]", ], ) -def test_indices(dimension): - expected = numpy.indices(dimension) - result = dpnp.indices(dimension) - assert_array_equal(expected, result) +@pytest.mark.parametrize("dtype", get_all_dtypes(no_bool=True)) +@pytest.mark.parametrize("sparse", [True, False], ids=["True", "False"]) +def test_indices(dimension, dtype, sparse): + expected = numpy.indices(dimension, dtype=dtype, sparse=sparse) + result = dpnp.indices(dimension, dtype=dtype, sparse=sparse) + for Xnp, X in zip(expected, result): + assert_array_equal(Xnp, X) @pytest.mark.parametrize( diff --git a/tests/test_sycl_queue.py b/tests/test_sycl_queue.py index fb94f1cd4cb0..fd09d08b5beb 100644 --- a/tests/test_sycl_queue.py +++ b/tests/test_sycl_queue.py @@ -1467,6 +1467,31 @@ def test_take(func, device): assert_sycl_queue_equal(result_queue, expected_queue) +@pytest.mark.parametrize( + "device", + valid_devices, + ids=[device.filter_string for device in valid_devices], +) +@pytest.mark.parametrize("sparse", [True, False], ids=["True", "False"]) +def test_indices(device, sparse): + sycl_queue = dpctl.SyclQueue(device) + grid = dpnp.indices((2, 3), sparse=sparse, sycl_queue=sycl_queue) + for dpnp_array in grid: + 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, diff --git a/tests/test_usm_type.py b/tests/test_usm_type.py index 0b094920c01e..fa4a14090b37 100644 --- a/tests/test_usm_type.py +++ b/tests/test_usm_type.py @@ -558,6 +558,26 @@ 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_sparse(usm_type, sparse): + x = dp.indices((2, 3), sparse=sparse, usm_type=usm_type) + for i in x: + assert i.usm_type == usm_type + + @pytest.mark.parametrize("usm_type", list_of_usm_types, ids=list_of_usm_types) def test_clip(usm_type): x = dp.arange(10, usm_type=usm_type) diff --git a/tests/third_party/cupy/creation_tests/test_ranges.py b/tests/third_party/cupy/creation_tests/test_ranges.py index 623adc409b7e..83170cb3e37f 100644 --- a/tests/third_party/cupy/creation_tests/test_ranges.py +++ b/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] diff --git a/tests/third_party/cupy/indexing_tests/test_generate.py b/tests/third_party/cupy/indexing_tests/test_generate.py index 58b976c3bd0d..dbec17358b2d 100644 --- a/tests/third_party/cupy/indexing_tests/test_generate.py +++ b/tests/third_party/cupy/indexing_tests/test_generate.py @@ -7,7 +7,6 @@ from tests.third_party.cupy import testing -@pytest.mark.usefixtures("allow_fall_back_on_numpy") @testing.gpu class TestIndices(unittest.TestCase): @testing.for_all_dtypes()