diff --git a/.github/workflows/conda-package.yml b/.github/workflows/conda-package.yml index 75860abff73..4481223e052 100644 --- a/.github/workflows/conda-package.yml +++ b/.github/workflows/conda-package.yml @@ -54,6 +54,7 @@ env: third_party/cupy/fft_tests third_party/cupy/creation_tests third_party/cupy/indexing_tests/test_indexing.py + third_party/cupy/indexing_tests/test_generate.py third_party/cupy/lib_tests third_party/cupy/linalg_tests third_party/cupy/logic_tests diff --git a/dpnp/dpnp_iface_indexing.py b/dpnp/dpnp_iface_indexing.py index 61a472519f0..3275e149310 100644 --- a/dpnp/dpnp_iface_indexing.py +++ b/dpnp/dpnp_iface_indexing.py @@ -72,6 +72,7 @@ "put", "put_along_axis", "putmask", + "ravel_multi_index", "select", "take", "take_along_axis", @@ -79,6 +80,7 @@ "tril_indices_from", "triu_indices", "triu_indices_from", + "unravel_index", ] @@ -133,6 +135,33 @@ def _build_along_axis_index(a, ind, axis): return tuple(fancy_index) +def _ravel_multi_index_checks(multi_index, dims, order): + dpnp.check_supported_arrays_type(*multi_index) + ndim = len(dims) + if len(multi_index) != ndim: + raise ValueError( + f"parameter multi_index must be a sequence of length {ndim}" + ) + dim_mul = 1.0 + for d in dims: + if not isinstance(d, int): + raise TypeError( + f"{type(d)} object cannot be interpreted as an integer" + ) + dim_mul *= d + + if dim_mul > dpnp.iinfo(dpnp.int64).max: + raise ValueError( + "invalid dims: array size defined by dims is larger than the " + "maximum possible size" + ) + if order not in ("C", "c", "F", "f", None): + raise ValueError( + "Unrecognized `order` keyword value, expecting " + f"'C' or 'F', but got '{order}'" + ) + + def choose(x1, choices, out=None, mode="raise"): """ Construct an array from an index array and a set of arrays to choose from. @@ -1429,6 +1458,112 @@ def putmask(x1, mask, values): return call_origin(numpy.putmask, x1, mask, values, dpnp_inplace=True) +def ravel_multi_index(multi_index, dims, mode="raise", order="C"): + """ + Converts a tuple of index arrays into an array of flat indices, applying + boundary modes to the multi-index. + + For full documentation refer to :obj:`numpy.ravel_multi_index`. + + Parameters + ---------- + multi_index : tuple of {dpnp.ndarray, usm_ndarray} + A tuple of integer arrays, one array for each dimension. + dims : tuple or list of ints + The shape of array into which the indices from ``multi_index`` apply. + mode : {"raise", "wrap" or "clip'}, optional + Specifies how out-of-bounds indices are handled. Can specify either + one mode or a tuple of modes, one mode per index: + - "raise" -- raise an error + - "wrap" -- wrap around + - "clip" -- clip to the range + In "clip" mode, a negative index which would normally wrap will + clip to 0 instead. + Default: ``"raise"``. + order : {None, "C", "F"}, optional + Determines whether the multi-index should be viewed as indexing in + row-major (C-style) or column-major (Fortran-style) order. + Default: ``"C"``. + + Returns + ------- + raveled_indices : dpnp.ndarray + An array of indices into the flattened version of an array of + dimensions ``dims``. + + See Also + -------- + :obj:`dpnp.unravel_index` : Converts array of flat indices into a tuple of + coordinate arrays. + + Examples + -------- + >>> import dpnp as np + >>> arr = np.array([[3, 6, 6], [4, 5, 1]]) + >>> np.ravel_multi_index(arr, (7, 6)) + array([22, 41, 37]) + >>> np.ravel_multi_index(arr, (7, 6), order="F") + array([31, 41, 13]) + >>> np.ravel_multi_index(arr, (4, 6), mode="clip") + array([22, 23, 19]) + >>> np.ravel_multi_index(arr, (4, 4), mode=("clip", "wrap")) + array([12, 13, 13]) + >>> arr = np.array([3, 1, 4, 1]) + >>> np.ravel_multi_index(arr, (6, 7, 8, 9)) + array(1621) + + """ + + _ravel_multi_index_checks(multi_index, dims, order) + + ndim = len(dims) + if isinstance(mode, str): + mode = (mode,) * ndim + + s = 1 + ravel_strides = [1] * ndim + + multi_index = tuple(multi_index) + usm_type_alloc, sycl_queue_alloc = get_usm_allocations(multi_index) + + order = "C" if order is None else order.upper() + if order == "C": + for i in range(ndim - 2, -1, -1): + s = s * dims[i + 1] + ravel_strides[i] = s + else: + for i in range(1, ndim): + s = s * dims[i - 1] + ravel_strides[i] = s + + multi_index = dpnp.broadcast_arrays(*multi_index) + raveled_indices = dpnp.zeros( + multi_index[0].shape, + dtype=dpnp.int64, + usm_type=usm_type_alloc, + sycl_queue=sycl_queue_alloc, + ) + for d, stride, idx, _mode in zip(dims, ravel_strides, multi_index, mode): + if not dpnp.can_cast(idx, dpnp.int64, "same_kind"): + raise TypeError( + f"multi_index entries could not be cast from dtype({idx.dtype})" + f" to dtype({dpnp.int64}) according to the rule 'same_kind'" + ) + idx = idx.astype(dpnp.int64, copy=False) + + if _mode == "raise": + if dpnp.any(dpnp.logical_or(idx >= d, idx < 0)): + raise ValueError("invalid entry in coordinates array") + elif _mode == "clip": + idx = dpnp.clip(idx, 0, d - 1) + elif _mode == "wrap": + idx = idx % d + else: + raise ValueError(f"Unrecognized mode: {_mode}") + raveled_indices += stride * idx + return raveled_indices + + def select(condlist, choicelist, default=0): """ Return an array drawn from elements in `choicelist`, depending on @@ -2177,3 +2312,78 @@ def triu_indices_from(arr, k=0): usm_type=arr.usm_type, sycl_queue=arr.sycl_queue, ) + + +def unravel_index(indices, shape, order="C"): + """Converts array of flat indices into a tuple of coordinate arrays. + + For full documentation refer to :obj:`numpy.unravel_index`. + + Parameters + ---------- + indices : {dpnp.ndarray, usm_ndarray} + An integer array whose elements are indices into the flattened version + of an array of dimensions ``shape``. + shape : tuple or list of ints + The shape of the array to use for unraveling ``indices``. + order : {None, "C", "F"}, optional + Determines whether the indices should be viewed as indexing in + row-major (C-style) or column-major (Fortran-style) order. + Default: ``"C"``. + + Returns + ------- + unraveled_coords : tuple of dpnp.ndarray + Each array in the tuple has the same shape as the indices array. + + See Also + -------- + :obj:`dpnp.ravel_multi_index` : Converts a tuple of index arrays into an + array of flat indices. + + + Examples + -------- + import dpnp as np + >>> np.unravel_index(np.array([22, 41, 37]), (7, 6)) + (array([3, 6, 6]), array([4, 5, 1])) + >>> np.unravel_index(np.array([31, 41, 13]), (7, 6), order="F") + (array([3, 6, 6]), array([4, 5, 1])) + + >>> np.unravel_index(np.array(1621), (6, 7, 8, 9)) + (array(3), array(1), array(4), array(1)) + + """ + + dpnp.check_supported_arrays_type(indices) + + if order not in ("C", "c", "F", "f", None): + raise ValueError( + "Unrecognized `order` keyword value, expecting " + f"'C' or 'F', but got '{order}'" + ) + order = "C" if order is None else order.upper() + if order == "C": + shape = reversed(shape) + + if not dpnp.can_cast(indices, dpnp.int64, "same_kind"): + raise TypeError( + "Iterator operand 0 dtype could not be cast from dtype(" + f"{indices.dtype}) to dtype({dpnp.int64}) according to the rule " + "'same_kind'" + ) + + if (indices < 0).any(): + raise ValueError("invalid entry in index array") + + unraveled_coords = [] + for dim in shape: + unraveled_coords.append(indices % dim) + indices = indices // dim + + if (indices > 0).any(): + raise ValueError("invalid entry in index array") + + if order == "C": + unraveled_coords = reversed(unraveled_coords) + return tuple(unraveled_coords) diff --git a/tests/skipped_tests.tbl b/tests/skipped_tests.tbl index 4491cec1e7b..3261a065864 100644 --- a/tests/skipped_tests.tbl +++ b/tests/skipped_tests.tbl @@ -60,18 +60,6 @@ tests/third_party/cupy/indexing_tests/test_generate.py::TestAxisConcatenator::te tests/third_party/cupy/indexing_tests/test_generate.py::TestC_::test_c_1 tests/third_party/cupy/indexing_tests/test_generate.py::TestC_::test_c_2 tests/third_party/cupy/indexing_tests/test_generate.py::TestC_::test_c_3 -tests/third_party/cupy/indexing_tests/test_generate.py::TestRavelMultiIndex::test_basic -tests/third_party/cupy/indexing_tests/test_generate.py::TestRavelMultiIndex::test_basic_clip -tests/third_party/cupy/indexing_tests/test_generate.py::TestRavelMultiIndex::test_basic_nd_coords -tests/third_party/cupy/indexing_tests/test_generate.py::TestRavelMultiIndex::test_basic_raise -tests/third_party/cupy/indexing_tests/test_generate.py::TestRavelMultiIndex::test_basic_wrap -tests/third_party/cupy/indexing_tests/test_generate.py::TestRavelMultiIndex::test_dims_overflow -tests/third_party/cupy/indexing_tests/test_generate.py::TestRavelMultiIndex::test_invalid_float_dims -tests/third_party/cupy/indexing_tests/test_generate.py::TestRavelMultiIndex::test_invalid_mode -tests/third_party/cupy/indexing_tests/test_generate.py::TestRavelMultiIndex::test_invalid_multi_index_dtype -tests/third_party/cupy/indexing_tests/test_generate.py::TestRavelMultiIndex::test_invalid_multi_index_shape -tests/third_party/cupy/indexing_tests/test_generate.py::TestRavelMultiIndex::test_invalid_order -tests/third_party/cupy/indexing_tests/test_generate.py::TestRavelMultiIndex::test_multi_index_broadcasting tests/third_party/cupy/indexing_tests/test_generate.py::TestR_::test_r_1 tests/third_party/cupy/indexing_tests/test_generate.py::TestR_::test_r_2 tests/third_party/cupy/indexing_tests/test_generate.py::TestR_::test_r_3 @@ -81,10 +69,6 @@ tests/third_party/cupy/indexing_tests/test_generate.py::TestR_::test_r_6 tests/third_party/cupy/indexing_tests/test_generate.py::TestR_::test_r_7 tests/third_party/cupy/indexing_tests/test_generate.py::TestR_::test_r_8 tests/third_party/cupy/indexing_tests/test_generate.py::TestR_::test_r_9 -tests/third_party/cupy/indexing_tests/test_generate.py::TestUnravelIndex::test -tests/third_party/cupy/indexing_tests/test_generate.py::TestUnravelIndex::test_invalid_dtype -tests/third_party/cupy/indexing_tests/test_generate.py::TestUnravelIndex::test_invalid_index -tests/third_party/cupy/indexing_tests/test_generate.py::TestUnravelIndex::test_invalid_order tests/third_party/cupy/indexing_tests/test_insert.py::TestPutmaskDifferentDtypes::test_putmask_differnt_dtypes_raises tests/third_party/cupy/indexing_tests/test_insert.py::TestPutmask::test_putmask_non_equal_shape_raises diff --git a/tests/skipped_tests_gpu.tbl b/tests/skipped_tests_gpu.tbl index 3dc8b97e011..7fba58a586c 100644 --- a/tests/skipped_tests_gpu.tbl +++ b/tests/skipped_tests_gpu.tbl @@ -113,18 +113,6 @@ tests/third_party/cupy/indexing_tests/test_generate.py::TestAxisConcatenator::te tests/third_party/cupy/indexing_tests/test_generate.py::TestC_::test_c_1 tests/third_party/cupy/indexing_tests/test_generate.py::TestC_::test_c_2 tests/third_party/cupy/indexing_tests/test_generate.py::TestC_::test_c_3 -tests/third_party/cupy/indexing_tests/test_generate.py::TestRavelMultiIndex::test_basic -tests/third_party/cupy/indexing_tests/test_generate.py::TestRavelMultiIndex::test_basic_clip -tests/third_party/cupy/indexing_tests/test_generate.py::TestRavelMultiIndex::test_basic_nd_coords -tests/third_party/cupy/indexing_tests/test_generate.py::TestRavelMultiIndex::test_basic_raise -tests/third_party/cupy/indexing_tests/test_generate.py::TestRavelMultiIndex::test_basic_wrap -tests/third_party/cupy/indexing_tests/test_generate.py::TestRavelMultiIndex::test_dims_overflow -tests/third_party/cupy/indexing_tests/test_generate.py::TestRavelMultiIndex::test_invalid_float_dims -tests/third_party/cupy/indexing_tests/test_generate.py::TestRavelMultiIndex::test_invalid_mode -tests/third_party/cupy/indexing_tests/test_generate.py::TestRavelMultiIndex::test_invalid_multi_index_dtype -tests/third_party/cupy/indexing_tests/test_generate.py::TestRavelMultiIndex::test_invalid_multi_index_shape -tests/third_party/cupy/indexing_tests/test_generate.py::TestRavelMultiIndex::test_invalid_order -tests/third_party/cupy/indexing_tests/test_generate.py::TestRavelMultiIndex::test_multi_index_broadcasting tests/third_party/cupy/indexing_tests/test_generate.py::TestR_::test_r_1 tests/third_party/cupy/indexing_tests/test_generate.py::TestR_::test_r_2 tests/third_party/cupy/indexing_tests/test_generate.py::TestR_::test_r_3 @@ -134,10 +122,6 @@ tests/third_party/cupy/indexing_tests/test_generate.py::TestR_::test_r_6 tests/third_party/cupy/indexing_tests/test_generate.py::TestR_::test_r_7 tests/third_party/cupy/indexing_tests/test_generate.py::TestR_::test_r_8 tests/third_party/cupy/indexing_tests/test_generate.py::TestR_::test_r_9 -tests/third_party/cupy/indexing_tests/test_generate.py::TestUnravelIndex::test -tests/third_party/cupy/indexing_tests/test_generate.py::TestUnravelIndex::test_invalid_dtype -tests/third_party/cupy/indexing_tests/test_generate.py::TestUnravelIndex::test_invalid_index -tests/third_party/cupy/indexing_tests/test_generate.py::TestUnravelIndex::test_invalid_order tests/third_party/cupy/indexing_tests/test_insert.py::TestPutmaskDifferentDtypes::test_putmask_differnt_dtypes_raises tests/third_party/cupy/indexing_tests/test_insert.py::TestPutmask::test_putmask_non_equal_shape_raises diff --git a/tests/test_indexing.py b/tests/test_indexing.py index 36046420117..ebaa7f4ac6c 100644 --- a/tests/test_indexing.py +++ b/tests/test_indexing.py @@ -1037,6 +1037,112 @@ def test_fill_diagonal_error(): dpnp.fill_diagonal(arr, 5) +class TestRavelIndex: + def test_basic(self): + expected = numpy.ravel_multi_index(numpy.array([1, 0]), (2, 2)) + result = dpnp.ravel_multi_index(dpnp.array([1, 0]), (2, 2)) + assert_equal(expected, result) + + x_np = numpy.array([[3, 6, 6], [4, 5, 1]]) + x_dp = dpnp.array([[3, 6, 6], [4, 5, 1]]) + + expected = numpy.ravel_multi_index(x_np, (7, 6)) + result = dpnp.ravel_multi_index(x_dp, (7, 6)) + assert_equal(expected, result) + + def test_mode(self): + x_np = numpy.array([[3, 6, 6], [4, 5, 1]]) + x_dp = dpnp.array([[3, 6, 6], [4, 5, 1]]) + + expected = numpy.ravel_multi_index(x_np, (4, 6), mode="clip") + result = dpnp.ravel_multi_index(x_dp, (4, 6), mode="clip") + assert_equal(expected, result) + + expected = numpy.ravel_multi_index(x_np, (4, 4), mode=("clip", "wrap")) + result = dpnp.ravel_multi_index(x_dp, (4, 4), mode=("clip", "wrap")) + assert_equal(expected, result) + + def test_order_f(self): + x_np = numpy.array([[3, 6, 6], [4, 5, 1]]) + x_dp = dpnp.array([[3, 6, 6], [4, 5, 1]]) + expected = numpy.ravel_multi_index(x_np, (7, 6), order="F") + result = dpnp.ravel_multi_index(x_dp, (7, 6), order="F") + assert_equal(expected, result) + + def test_error(self): + assert_raises( + ValueError, dpnp.ravel_multi_index, dpnp.array([2, 1]), (2, 2) + ) + assert_raises( + ValueError, dpnp.ravel_multi_index, dpnp.array([0, -3]), (2, 2) + ) + assert_raises( + ValueError, dpnp.ravel_multi_index, dpnp.array([0, 2]), (2, 2) + ) + assert_raises( + TypeError, dpnp.ravel_multi_index, dpnp.array([0.1, 0.0]), (2, 2) + ) + + def test_empty_indices_error(self): + assert_raises(TypeError, dpnp.ravel_multi_index, ([], []), (10, 3)) + assert_raises(TypeError, dpnp.ravel_multi_index, ([], ["abc"]), (10, 3)) + assert_raises( + TypeError, + dpnp.ravel_multi_index, + (dpnp.array([]), dpnp.array([])), + (5, 3), + ) + + def test_empty_indices(self): + assert_equal( + dpnp.ravel_multi_index( + (dpnp.array([], dtype=int), dpnp.array([], dtype=int)), (5, 3) + ), + [], + ) + assert_equal( + dpnp.ravel_multi_index(dpnp.array([[], []], dtype=int), (5, 3)), [] + ) + + +class TestUnravelIndex: + def test_basic(self): + expected = numpy.unravel_index(numpy.array(2), (2, 2)) + result = dpnp.unravel_index(dpnp.array(2), (2, 2)) + assert_equal(expected, result) + + x_np = numpy.array([22, 41, 37]) + x_dp = dpnp.array([22, 41, 37]) + + expected = numpy.unravel_index(x_np, (7, 6)) + result = dpnp.unravel_index(x_dp, (7, 6)) + assert_equal(expected, result) + + def test_order_f(self): + x_np = numpy.array([31, 41, 13]) + x_dp = dpnp.array([31, 41, 13]) + expected = numpy.unravel_index(x_np, (7, 6), order="F") + result = dpnp.unravel_index(x_dp, (7, 6), order="F") + assert_equal(expected, result) + + def test_new_shape(self): + expected = numpy.unravel_index(numpy.array(2), shape=(2, 2)) + result = dpnp.unravel_index(dpnp.array(2), shape=(2, 2)) + assert_equal(expected, result) + + def test_error(self): + assert_raises(ValueError, dpnp.unravel_index, dpnp.array(-1), (2, 2)) + assert_raises(TypeError, dpnp.unravel_index, dpnp.array(0.5), (2, 2)) + assert_raises(ValueError, dpnp.unravel_index, dpnp.array(4), (2, 2)) + assert_raises(TypeError, dpnp.unravel_index, dpnp.array([]), (10, 3, 5)) + + def test_empty_indices(self): + assert_equal( + dpnp.unravel_index(dpnp.array([], dtype=int), (10, 3, 5)), + [[], [], []], + ) + + class TestIx: @pytest.mark.parametrize( "x0", [[0, 1], [True, True]], ids=["[0, 1]", "[True, True]"] diff --git a/tests/test_sycl_queue.py b/tests/test_sycl_queue.py index eb54065e1df..a41b7d52e77 100644 --- a/tests/test_sycl_queue.py +++ b/tests/test_sycl_queue.py @@ -2567,6 +2567,29 @@ def test_ediff1d(device, to_end, to_begin): assert_sycl_queue_equal(res.sycl_queue, x.sycl_queue) +@pytest.mark.parametrize( + "device", + valid_devices, + ids=[device.filter_string for device in valid_devices], +) +def test_unravel_index(device): + x = dpnp.array(2, device=device) + result = dpnp.unravel_index(x, shape=(2, 2)) + for res in result: + assert_sycl_queue_equal(res.sycl_queue, x.sycl_queue) + + +@pytest.mark.parametrize( + "device", + valid_devices, + ids=[device.filter_string for device in valid_devices], +) +def test_ravel_index(device): + x = dpnp.array([1, 0], device=device) + result = dpnp.ravel_multi_index(x, (2, 2)) + assert_sycl_queue_equal(result.sycl_queue, x.sycl_queue) + + @pytest.mark.parametrize( "device_0", valid_devices, diff --git a/tests/test_usm_type.py b/tests/test_usm_type.py index f83b3df5f02..b76f9f42b33 100644 --- a/tests/test_usm_type.py +++ b/tests/test_usm_type.py @@ -1513,6 +1513,21 @@ def test_ediff1d(usm_type_x, usm_type_args, to_end, to_begin): assert res.usm_type == du.get_coerced_usm_type([usm_type_x, usm_type_args]) +@pytest.mark.parametrize("usm_type", list_of_usm_types, ids=list_of_usm_types) +def test_unravel_index(usm_type): + x = dp.array(2, usm_type=usm_type) + result = dp.unravel_index(x, shape=(2, 2)) + for res in result: + assert res.usm_type == x.usm_type + + +@pytest.mark.parametrize("usm_type", list_of_usm_types, ids=list_of_usm_types) +def test_ravel_index(usm_type): + x = dp.array([1, 0], usm_type=usm_type) + result = dp.ravel_multi_index(x, (2, 2)) + assert result.usm_type == x.usm_type + + @pytest.mark.parametrize("usm_type_0", list_of_usm_types, ids=list_of_usm_types) @pytest.mark.parametrize("usm_type_1", list_of_usm_types, ids=list_of_usm_types) def test_ix(usm_type_0, usm_type_1):