From 133b96069fb355f2d9ac22ee0091b625cafa8c31 Mon Sep 17 00:00:00 2001 From: Natalia Polina Date: Wed, 8 May 2024 16:18:14 -0500 Subject: [PATCH 1/6] Update indexing functions --- dpnp/backend/include/dpnp_iface_fptr.hpp | 22 +- dpnp/backend/kernels/dpnp_krnl_indexing.cpp | 40 +- dpnp/dpnp_algo/dpnp_algo.pxd | 2 - dpnp/dpnp_algo/dpnp_algo_indexing.pxi | 162 ------ dpnp/dpnp_iface_indexing.py | 486 +++++++++++++++--- tests/skipped_tests.tbl | 33 -- tests/skipped_tests_gpu.tbl | 51 -- .../cupy/indexing_tests/test_insert.py | 2 - 8 files changed, 426 insertions(+), 372 deletions(-) diff --git a/dpnp/backend/include/dpnp_iface_fptr.hpp b/dpnp/backend/include/dpnp_iface_fptr.hpp index b63f0e00431f..e3c63a3ded27 100644 --- a/dpnp/backend/include/dpnp_iface_fptr.hpp +++ b/dpnp/backend/include/dpnp_iface_fptr.hpp @@ -114,13 +114,11 @@ enum class DPNPFuncName : size_t DPNP_FN_DET, /**< Used in numpy.linalg.det() impl */ DPNP_FN_DIAG, /**< Used in numpy.diag() impl */ DPNP_FN_DIAG_INDICES, /**< Used in numpy.diag_indices() impl */ - DPNP_FN_DIAG_INDICES_EXT, /**< Used in numpy.diag_indices() impl, requires - extra parameters */ - DPNP_FN_DIAGONAL, /**< Used in numpy.diagonal() impl */ - DPNP_FN_DIAGONAL_EXT, /**< Used in numpy.diagonal() impl, requires extra - parameters */ - DPNP_FN_DIVIDE, /**< Used in numpy.divide() impl */ - DPNP_FN_DOT, /**< Used in numpy.dot() impl */ + DPNP_FN_DIAGONAL, /**< Used in numpy.diagonal() impl */ + DPNP_FN_DIAGONAL_EXT, /**< Used in numpy.diagonal() impl, requires extra + parameters */ + DPNP_FN_DIVIDE, /**< Used in numpy.divide() impl */ + DPNP_FN_DOT, /**< Used in numpy.dot() impl */ DPNP_FN_DOT_EXT, /**< Used in numpy.dot() impl, requires extra parameters */ DPNP_FN_EDIFF1D, /**< Used in numpy.ediff1d() impl */ DPNP_FN_EDIFF1D_EXT, /**< Used in numpy.ediff1d() impl, requires extra @@ -144,12 +142,10 @@ enum class DPNPFuncName : size_t DPNP_FN_FFT_RFFT_EXT, /**< Used in numpy.fft.rfft() impl, requires extra parameters */ DPNP_FN_FILL_DIAGONAL, /**< Used in numpy.fill_diagonal() impl */ - DPNP_FN_FILL_DIAGONAL_EXT, /**< Used in numpy.fill_diagonal() impl, requires - extra parameters */ - DPNP_FN_FLATTEN, /**< Used in numpy.flatten() impl */ - DPNP_FN_FLOOR, /**< Used in numpy.floor() impl */ - DPNP_FN_FLOOR_DIVIDE, /**< Used in numpy.floor_divide() impl */ - DPNP_FN_FMOD, /**< Used in numpy.fmod() impl */ + DPNP_FN_FLATTEN, /**< Used in numpy.flatten() impl */ + DPNP_FN_FLOOR, /**< Used in numpy.floor() impl */ + DPNP_FN_FLOOR_DIVIDE, /**< Used in numpy.floor_divide() impl */ + DPNP_FN_FMOD, /**< Used in numpy.fmod() impl */ DPNP_FN_FMOD_EXT, /**< Used in numpy.fmod() impl, requires extra parameters */ DPNP_FN_FULL, /**< Used in numpy.full() impl */ diff --git a/dpnp/backend/kernels/dpnp_krnl_indexing.cpp b/dpnp/backend/kernels/dpnp_krnl_indexing.cpp index 0889358d989f..969fa903c1c9 100644 --- a/dpnp/backend/kernels/dpnp_krnl_indexing.cpp +++ b/dpnp/backend/kernels/dpnp_krnl_indexing.cpp @@ -150,13 +150,6 @@ template void (*dpnp_diag_indices_default_c)(void *, size_t) = dpnp_diag_indices_c<_DataType>; -template -DPCTLSyclEventRef (*dpnp_diag_indices_ext_c)(DPCTLSyclQueueRef, - void *, - size_t, - const DPCTLEventVectorRef) = - dpnp_diag_indices_c<_DataType>; - template DPCTLSyclEventRef dpnp_diagonal_c(DPCTLSyclQueueRef q_ref, void *array1_in, @@ -190,7 +183,8 @@ DPCTLSyclEventRef dpnp_diagonal_c(DPCTLSyclQueueRef q_ref, if (res_ndim <= 1) { for (size_t i = 0; i < static_cast(res_shape[res_ndim - 1]); - ++i) { + ++i) + { result[i] = array_1[i * shape[res_ndim] + i + offset]; } } @@ -226,7 +220,8 @@ DPCTLSyclEventRef dpnp_diagonal_c(DPCTLSyclQueueRef q_ref, } for (size_t i = 0; i < static_cast(res_shape[res_ndim - 1]); - i++) { + i++) + { for (size_t j = 0; j < xyz.size(); j++) { std::vector ind_list = xyz[j]; if (ind_list.size() == 0) { @@ -376,15 +371,6 @@ void (*dpnp_fill_diagonal_default_c)(void *, const size_t) = dpnp_fill_diagonal_c<_DataType>; -template -DPCTLSyclEventRef (*dpnp_fill_diagonal_ext_c)(DPCTLSyclQueueRef, - void *, - void *, - shape_elem_type *, - const size_t, - const DPCTLEventVectorRef) = - dpnp_fill_diagonal_c<_DataType>; - template DPCTLSyclEventRef dpnp_nonzero_c(DPCTLSyclQueueRef q_ref, const void *in_array1, @@ -909,15 +895,6 @@ void func_map_init_indexing_func(func_map_t &fmap) fmap[DPNPFuncName::DPNP_FN_DIAG_INDICES][eft_DBL][eft_DBL] = { eft_DBL, (void *)dpnp_diag_indices_default_c}; - fmap[DPNPFuncName::DPNP_FN_DIAG_INDICES_EXT][eft_INT][eft_INT] = { - eft_INT, (void *)dpnp_diag_indices_ext_c}; - fmap[DPNPFuncName::DPNP_FN_DIAG_INDICES_EXT][eft_LNG][eft_LNG] = { - eft_LNG, (void *)dpnp_diag_indices_ext_c}; - fmap[DPNPFuncName::DPNP_FN_DIAG_INDICES_EXT][eft_FLT][eft_FLT] = { - eft_FLT, (void *)dpnp_diag_indices_ext_c}; - fmap[DPNPFuncName::DPNP_FN_DIAG_INDICES_EXT][eft_DBL][eft_DBL] = { - eft_DBL, (void *)dpnp_diag_indices_ext_c}; - fmap[DPNPFuncName::DPNP_FN_DIAGONAL][eft_INT][eft_INT] = { eft_INT, (void *)dpnp_diagonal_default_c}; fmap[DPNPFuncName::DPNP_FN_DIAGONAL][eft_LNG][eft_LNG] = { @@ -949,15 +926,6 @@ void func_map_init_indexing_func(func_map_t &fmap) fmap[DPNPFuncName::DPNP_FN_FILL_DIAGONAL][eft_DBL][eft_DBL] = { eft_DBL, (void *)dpnp_fill_diagonal_default_c}; - fmap[DPNPFuncName::DPNP_FN_FILL_DIAGONAL_EXT][eft_INT][eft_INT] = { - eft_INT, (void *)dpnp_fill_diagonal_ext_c}; - fmap[DPNPFuncName::DPNP_FN_FILL_DIAGONAL_EXT][eft_LNG][eft_LNG] = { - eft_LNG, (void *)dpnp_fill_diagonal_ext_c}; - fmap[DPNPFuncName::DPNP_FN_FILL_DIAGONAL_EXT][eft_FLT][eft_FLT] = { - eft_FLT, (void *)dpnp_fill_diagonal_ext_c}; - fmap[DPNPFuncName::DPNP_FN_FILL_DIAGONAL_EXT][eft_DBL][eft_DBL] = { - eft_DBL, (void *)dpnp_fill_diagonal_ext_c}; - fmap[DPNPFuncName::DPNP_FN_NONZERO][eft_INT][eft_INT] = { eft_INT, (void *)dpnp_nonzero_default_c}; fmap[DPNPFuncName::DPNP_FN_NONZERO][eft_LNG][eft_LNG] = { diff --git a/dpnp/dpnp_algo/dpnp_algo.pxd b/dpnp/dpnp_algo/dpnp_algo.pxd index 5859347aac49..6176a0ab6afa 100644 --- a/dpnp/dpnp_algo/dpnp_algo.pxd +++ b/dpnp/dpnp_algo/dpnp_algo.pxd @@ -39,14 +39,12 @@ cdef extern from "dpnp_iface_fptr.hpp" namespace "DPNPFuncName": # need this na DPNP_FN_CORRELATE_EXT DPNP_FN_CUMPROD_EXT DPNP_FN_DEGREES_EXT - DPNP_FN_DIAG_INDICES_EXT DPNP_FN_DIAGONAL_EXT DPNP_FN_EDIFF1D_EXT DPNP_FN_ERF_EXT DPNP_FN_FABS_EXT DPNP_FN_FFT_FFT_EXT DPNP_FN_FFT_RFFT_EXT - DPNP_FN_FILL_DIAGONAL_EXT DPNP_FN_FMOD_EXT DPNP_FN_MAXIMUM_EXT DPNP_FN_MEDIAN_EXT diff --git a/dpnp/dpnp_algo/dpnp_algo_indexing.pxi b/dpnp/dpnp_algo/dpnp_algo_indexing.pxi index ab8eceeb415f..34e3bb03eed5 100644 --- a/dpnp/dpnp_algo/dpnp_algo_indexing.pxi +++ b/dpnp/dpnp_algo/dpnp_algo_indexing.pxi @@ -37,23 +37,14 @@ and the rest of the library __all__ += [ "dpnp_choose", - "dpnp_diag_indices", "dpnp_diagonal", - "dpnp_fill_diagonal", "dpnp_putmask", "dpnp_select", - "dpnp_tril_indices", - "dpnp_tril_indices_from", - "dpnp_triu_indices", - "dpnp_triu_indices_from" ] ctypedef c_dpctl.DPCTLSyclEventRef(*fptr_dpnp_choose_t)(c_dpctl.DPCTLSyclQueueRef, void *, void * , void ** , size_t, size_t, size_t, const c_dpctl.DPCTLEventVectorRef) -ctypedef c_dpctl.DPCTLSyclEventRef(*fptr_dpnp_diag_indices)(c_dpctl.DPCTLSyclQueueRef, - void * , size_t, - const c_dpctl.DPCTLEventVectorRef) ctypedef c_dpctl.DPCTLSyclEventRef(*custom_indexing_2in_1out_func_ptr_t_)(c_dpctl.DPCTLSyclQueueRef, void * , const size_t, @@ -115,38 +106,6 @@ cpdef utils.dpnp_descriptor dpnp_choose(utils.dpnp_descriptor x1, list choices1) return res_array -cpdef tuple dpnp_diag_indices(n, ndim): - cdef size_t res_size = 0 if n < 0 else n - - cdef DPNPFuncType param1_type = dpnp_dtype_to_DPNPFuncType(dpnp.int64) - - cdef DPNPFuncData kernel_data = get_dpnp_function_ptr(DPNP_FN_DIAG_INDICES_EXT, param1_type, param1_type) - - cdef fptr_dpnp_diag_indices func = kernel_data.ptr - - cdef c_dpctl.SyclQueue q - cdef c_dpctl.DPCTLSyclQueueRef q_ref - cdef c_dpctl.DPCTLSyclEventRef event_ref - - res_list = [] - cdef utils.dpnp_descriptor res_arr - cdef shape_type_c result_shape = utils._object_to_tuple(res_size) - for i in range(ndim): - res_arr = utils.create_output_descriptor(result_shape, kernel_data.return_type, None) - - q = res_arr.get_array().sycl_queue - q_ref = q.get_queue_ref() - - event_ref = func(q_ref, res_arr.get_data(), res_size, NULL) - - with nogil: c_dpctl.DPCTLEvent_WaitAndThrow(event_ref) - c_dpctl.DPCTLEvent_Delete(event_ref) - - res_list.append(res_arr.get_pyobj()) - - return tuple(res_list) - - cpdef utils.dpnp_descriptor dpnp_diagonal(dpnp_descriptor x1, offset=0): cdef shape_type_c x1_shape = x1.shape @@ -203,39 +162,6 @@ cpdef utils.dpnp_descriptor dpnp_diagonal(dpnp_descriptor x1, offset=0): return result -cpdef dpnp_fill_diagonal(dpnp_descriptor x1, val): - x1_obj = x1.get_array() - - cdef shape_type_c x1_shape = x1.shape - cdef utils.dpnp_descriptor val_arr = utils_py.create_output_descriptor_py((1,), - x1.dtype, - None, - device=x1_obj.sycl_device, - usm_type=x1_obj.usm_type, - sycl_queue=x1_obj.sycl_queue) - - val_arr.get_pyobj()[0] = val - - cdef DPNPFuncType param1_type = dpnp_dtype_to_DPNPFuncType(x1.dtype) - - cdef DPNPFuncData kernel_data = get_dpnp_function_ptr(DPNP_FN_FILL_DIAGONAL_EXT, param1_type, param1_type) - - cdef c_dpctl.SyclQueue q = x1_obj.sycl_queue - cdef c_dpctl.DPCTLSyclQueueRef q_ref = q.get_queue_ref() - - cdef custom_indexing_2in_func_ptr_t func = kernel_data.ptr - - cdef c_dpctl.DPCTLSyclEventRef event_ref = func(q_ref, - x1.get_data(), - val_arr.get_data(), - x1_shape.data(), - x1.ndim, - NULL) # dep_events_ref - - with nogil: c_dpctl.DPCTLEvent_WaitAndThrow(event_ref) - c_dpctl.DPCTLEvent_Delete(event_ref) - - cpdef dpnp_putmask(utils.dpnp_descriptor arr, utils.dpnp_descriptor mask, utils.dpnp_descriptor values): cdef int values_size = values.size @@ -263,91 +189,3 @@ cpdef utils.dpnp_descriptor dpnp_select(list condlist, list choicelist, default) res_array.get_pyobj()[ind] = val return res_array - - -cpdef tuple dpnp_tril_indices(n, k=0, m=None): - array1 = [] - array2 = [] - if m is None: - for i in range(n): - for j in range(i + 1 + k): - if j >= n: - continue - else: - array1.append(i) - array2.append(j) - else: - for i in range(n): - for j in range(i + 1 + k): - if j < m: - array1.append(i) - array2.append(j) - - array1 = dpnp.array(array1, dtype=dpnp.int64) - array2 = dpnp.array(array2, dtype=dpnp.int64) - return (array1, array2) - - -cpdef tuple dpnp_tril_indices_from(dpnp_descriptor arr, k=0): - m = arr.shape[0] - n = arr.shape[1] - array1 = [] - array2 = [] - if m is None: - for i in range(n): - for j in range(i + 1 + k): - if j >= n: - continue - else: - array1.append(i) - array2.append(j) - else: - for i in range(n): - for j in range(i + 1 + k): - if j < m: - array1.append(i) - array2.append(j) - - array1 = dpnp.array(array1, dtype=dpnp.int64) - array2 = dpnp.array(array2, dtype=dpnp.int64) - return (array1, array2) - - -cpdef tuple dpnp_triu_indices(n, k=0, m=None): - array1 = [] - array2 = [] - if m is None: - for i in range(n): - for j in range(i + k, n): - array1.append(i) - array2.append(j) - else: - for i in range(n): - for j in range(i + k, m): - array1.append(i) - array2.append(j) - - array1 = dpnp.array(array1, dtype=dpnp.int64) - array2 = dpnp.array(array2, dtype=dpnp.int64) - return (array1, array2) - - -cpdef tuple dpnp_triu_indices_from(dpnp_descriptor arr, k=0): - m = arr.shape[0] - n = arr.shape[1] - array1 = [] - array2 = [] - if m is None: - for i in range(n): - for j in range(i + k, n): - array1.append(i) - array2.append(j) - else: - for i in range(n): - for j in range(i + k, m): - array1.append(i) - array2.append(j) - - array1 = dpnp.array(array1, dtype=dpnp.int64) - array2 = dpnp.array(array2, dtype=dpnp.int64) - return (array1, array2) diff --git a/dpnp/dpnp_iface_indexing.py b/dpnp/dpnp_iface_indexing.py index 1db5e33d42a4..a12283514f2f 100644 --- a/dpnp/dpnp_iface_indexing.py +++ b/dpnp/dpnp_iface_indexing.py @@ -46,15 +46,9 @@ # pylint: disable=no-name-in-module from .dpnp_algo import ( dpnp_choose, - dpnp_diag_indices, dpnp_diagonal, - dpnp_fill_diagonal, dpnp_putmask, dpnp_select, - dpnp_tril_indices, - dpnp_tril_indices_from, - dpnp_triu_indices, - dpnp_triu_indices_from, ) from .dpnp_array import dpnp_array from .dpnp_utils import ( @@ -70,6 +64,7 @@ "extract", "fill_diagonal", "indices", + "mask_indices", "nonzero", "place", "put", @@ -194,6 +189,15 @@ def diag_indices(n, ndim=2): For full documentation refer to :obj:`numpy.diag_indices`. + Parameters + ---------- + n : int + The size, along each dimension, of the arrays for which the returned + indices can be used. + + ndim : int, optional + The number of dimensions. + See also -------- :obj:`diag_indices_from` : Return the indices to access the main @@ -238,42 +242,62 @@ def diag_indices(n, ndim=2): """ - if not use_origin_backend(): - return dpnp_diag_indices(n, ndim) - - return call_origin(numpy.diag_indices, n, ndim) + idx = dpnp.arange(n) + return (idx,) * ndim -def diag_indices_from(x1): +def diag_indices_from(arr): """ Return the indices to access the main diagonal of an n-dimensional array. For full documentation refer to :obj:`numpy.diag_indices_from`. + Parameters + ---------- + arr : array, at least 2-D + See also -------- :obj:`diag_indices` : Return the indices to access the main diagonal of an array. + Examples + -------- + Create a 4 by 4 array. + + >>> import dpnp as np + >>> a = np.arange(16).reshape(4, 4) + >>> a + array([[ 0, 1, 2, 3], + [ 4, 5, 6, 7], + [ 8, 9, 10, 11], + [12, 13, 14, 15]]) + + Get the indices of the diagonal elements. + + >>> di = np.diag_indices_from(a) + >>> di + (array([0, 1, 2, 3]), array([0, 1, 2, 3])) + + >>> a[di] + array([ 0, 5, 10, 15]) + + This is simply syntactic sugar for diag_indices. + + >>> np.diag_indices(a.shape[0]) + (array([0, 1, 2, 3]), array([0, 1, 2, 3])) + """ - x1_desc = dpnp.get_dpnp_descriptor(x1, copy_when_nondefault_queue=False) - if x1_desc: - # original limitation - if not x1_desc.ndim >= 2: - pass + if not arr.ndim >= 2: + raise ValueError("input array must be at least 2-d") - # original limitation - # For more than d=2, the strided formula is only valid for arrays with - # all dimensions equal, so we check first. - elif not numpy.alltrue( - numpy.diff(x1_desc.shape) == 0 - ): # TODO: replace alltrue and diff funcs with dpnp own ones - pass - else: - return dpnp_diag_indices(x1_desc.shape[0], x1_desc.ndim) + sh0 = arr.shape[0] + for sh in arr.shape[1:]: + if sh != sh0: + raise ValueError("All dimensions of input must be of equal length") - return call_origin(numpy.diag_indices_from, x1) + return diag_indices(sh0, arr.ndim) def diagonal(x1, offset=0, axis1=0, axis2=1): @@ -344,15 +368,24 @@ def extract(condition, x): return call_origin(numpy.extract, condition, x) -def fill_diagonal(x1, val, wrap=False): +def fill_diagonal(a, val, wrap=False): """ Fill the main diagonal of the given array of any dimensionality. For full documentation refer to :obj:`numpy.fill_diagonal`. - Limitations - ----------- - Parameter `wrap` is supported only with default values. + Parameters + ---------- + a : array, at least 2-D. + Array whose diagonal is to be filled in-place. + val : scalar or array_like + Value(s) to write on the diagonal. If `val` is scalar, the value is + written along the diagonal. If array-like, the flattened `val` is + written along the diagonal, repeating if necessary to fill all + diagonal entries. + wrap : bool + The diagonal "wrapped" after N columns. You can have this behavior + with this option. This affects only tall matrices. See Also -------- @@ -361,20 +394,113 @@ def fill_diagonal(x1, val, wrap=False): :obj:`dpnp.diag_indices_from` : Return the indices to access the main diagonal of an n-dimensional array. + Examples + -------- + >>> import dpnp as np + >>> a = np.zeros((3, 3), dtype=int) + >>> np.fill_diagonal(a, 5) + >>> a + array([[5, 0, 0], + [0, 5, 0], + [0, 0, 5]]) + + The same function can operate on a 4-D array: + + >>> a = np.zeros((3, 3, 3, 3), dtype=int) + >>> np.fill_diagonal(a, 4) + + We only show a few blocks for clarity: + + >>> a[0, 0] + array([[4, 0, 0], + [0, 0, 0], + [0, 0, 0]]) + >>> a[1, 1] + array([[0, 0, 0], + [0, 4, 0], + [0, 0, 0]]) + >>> a[2, 2] + array([[0, 0, 0], + [0, 0, 0], + [0, 0, 4]]) + + The wrap option affects only tall matrices: + + >>> # tall matrices no wrap + >>> a = np.zeros((5, 3), dtype=int) + >>> np.fill_diagonal(a, 4) + >>> a + array([[4, 0, 0], + [0, 4, 0], + [0, 0, 4], + [0, 0, 0], + [0, 0, 0]]) + + >>> # tall matrices wrap + >>> a = np.zeros((5, 3), dtype=int) + >>> np.fill_diagonal(a, 4, wrap=True) + >>> a + array([[4, 0, 0], + [0, 4, 0], + [0, 0, 4], + [0, 0, 0], + [4, 0, 0]]) + + >>> # wide matrices + >>> a = np.zeros((3, 5), dtype=int) + >>> np.fill_diagonal(a, 4, wrap=True) + >>> a + array([[4, 0, 0, 0, 0], + [0, 4, 0, 0, 0], + [0, 0, 4, 0, 0]]) + + The anti-diagonal can be filled by reversing the order of elements + using either `dpnp.flipud` or `dpnp.fliplr`. + + >>> a = np.zeros((3, 3), dtype=int) + >>> val = np.array([1, 2, 3]) + >>> np.fill_diagonal(np.fliplr(a), val) # Horizontal flip + >>> a + array([[0, 0, 1], + [0, 2, 0], + [3, 0, 0]]) + >>> np.fill_diagonal(np.flipud(a), val) # Vertical flip + >>> a + array([[0, 0, 3], + [0, 2, 0], + [1, 0, 0]]) + """ - x1_desc = dpnp.get_dpnp_descriptor( - x1, copy_when_strides=False, copy_when_nondefault_queue=False - ) - if x1_desc: - if not dpnp.isscalar(val): - pass - elif wrap: - pass - else: - return dpnp_fill_diagonal(x1_desc, val) + dpnp.check_supported_arrays_type(a, scalar_type=True, all_scalars=True) - return call_origin(numpy.fill_diagonal, x1, val, wrap, dpnp_inplace=True) + if a.ndim < 2: + raise ValueError("array must be at least 2-d") + end = a.size + if a.ndim == 2: + step = a.shape[1] + 1 + if not wrap and a.shape[0] > a.shape[1]: + end = a.shape[1] * a.shape[1] + else: + sh0 = a.shape[0] + for sh in a.shape[1:]: + if sh != sh0: + raise ValueError( + "All dimensions of input must be of equal length" + ) + step = sum(sh0**x for x in range(a.ndim)) + + # TODO: implement flatiter for slice key + # a.flat[:end:step] = val + fl = a.flat + if dpnp.isscalar(val): + for idx in range(0, end, step): + fl[idx] = val + else: + flat_val = val.flatten() + for i in range(0, flat_val.size): + for idx in range((step * i), end, step * (1 + i)): + fl[idx] = flat_val[i] def indices( @@ -493,6 +619,70 @@ def indices( return res +def mask_indices(n, mask_func, k=0): + """ + Return the indices to access (n, n) arrays, given a masking function. + + Assume `mask_func` is a function that, for a square array a of size + ``(n, n)`` with a possible offset argument `k`, when called as + ``mask_func(a, k)`` returns a new array with zeros in certain locations + (functions like `triu` or `tril` do precisely this). Then this function + returns the indices where the non-zero values would be located. + + Parameters + ---------- + n : int + The returned indices will be valid to access arrays of shape (n, n). + mask_func : callable + A function whose call signature is similar to that of `triu`, `tril`. + That is, ``mask_func(x, k)`` returns a boolean array, shaped like `x`. + `k` is an optional argument to the function. + k : scalar + An optional argument which is passed through to `mask_func`. Functions + like `triu`, `tril` take a second argument that is interpreted as an + offset. + + Returns + ------- + indices : tuple of arrays. + The `n` arrays of indices corresponding to the locations where + ``mask_func(np.ones((n, n)), k)`` is True. + + Examples + -------- + These are the indices that would allow you to access the upper triangular + part of any 3x3 array: + + >>> import dpnp as np + >>> iu = np.mask_indices(3, np.triu) + + For example, if `a` is a 3x3 array: + + >>> a = np.arange(9).reshape(3, 3) + >>> a + array([[0, 1, 2], + [3, 4, 5], + [6, 7, 8]]) + >>> a[iu] + array([0, 1, 2, 4, 5, 8]) + + An offset can be passed also to the masking function. This gets us the + indices starting on the first diagonal right of the main one: + + >>> iu1 = np.mask_indices(3, np.triu, 1) + + with which we now extract only three elements: + + >>> a[iu1] + array([1, 2, 5]) + + """ + + m = dpnp.ones((n, n), dtype=int) + a = mask_func(m, k) + return nonzero(a != 0) + + def nonzero(a): """ Return the indices of the elements that are non-zero. @@ -986,6 +1176,8 @@ def tril_indices(n, k=0, m=None): """ Return the indices for the lower-triangle of an (n, m) array. + For full documentation refer to :obj:`numpy.tril_indices`. + Parameters ---------- n : int @@ -1006,24 +1198,63 @@ def tril_indices(n, k=0, m=None): The indices for the triangle. The returned tuple contains two arrays, each with the indices along one dimension of the array. + Examples + -------- + Compute two different sets of indices to access 4x4 arrays, one for the + lower triangular part starting at the main diagonal, and one starting two + diagonals further right: + + >>> import dpnp as np + >>> il1 = np.tril_indices(4) + >>> il2 = np.tril_indices(4, 2) + + Here is how they can be used with a sample array: + + >>> a = np.arange(16).reshape(4, 4) + >>> a + array([[ 0, 1, 2, 3], + [ 4, 5, 6, 7], + [ 8, 9, 10, 11], + [12, 13, 14, 15]]) + + Both for indexing: + + >>> a[il1] + array([ 0, 4, 5, ..., 13, 14, 15]) + + And for assigning values: + + >>> a[il1] = -1 + >>> a + array([[-1, 1, 2, 3], + [-1, -1, 6, 7], + [-1, -1, -1, 11], + [-1, -1, -1, -1]]) + + These cover almost the whole array (two diagonals right of the main one): + + >>> a[il2] = -10 + >>> a + array([[-10, -10, -10, 3], + [-10, -10, -10, -10], + [-10, -10, -10, -10], + [-10, -10, -10, -10]]) + """ - if not use_origin_backend(): - if ( - isinstance(n, int) - and isinstance(k, int) - and (isinstance(m, int) or m is None) - ): - return dpnp_tril_indices(n, k, m) + tri_ = dpnp.tri(n, m, k=k, dtype=bool) - return call_origin(numpy.tril_indices, n, k, m) + return tuple( + dpnp.broadcast_to(inds, tri_.shape)[tri_] + for inds in indices(tri_.shape, sparse=True) + ) -def tril_indices_from(x1, k=0): +def tril_indices_from(arr, k=0): """ Return the indices for the lower-triangle of arr. - See `tril_indices` for full details. + For full documentation refer to :obj:`numpy.tril_indices_from`. Parameters ---------- @@ -1033,20 +1264,55 @@ def tril_indices_from(x1, k=0): k : int, optional Diagonal offset (see `tril` for details). - """ - x1_desc = dpnp.get_dpnp_descriptor(x1, copy_when_nondefault_queue=False) - if x1_desc: - if isinstance(k, int): - return dpnp_tril_indices_from(x1_desc, k) + Examples + -------- + Create a 4 by 4 array. + + >>> import dpnp as np + >>> a = np.arange(16).reshape(4, 4) + >>> a + array([[ 0, 1, 2, 3], + [ 4, 5, 6, 7], + [ 8, 9, 10, 11], + [12, 13, 14, 15]]) + + Pass the array to get the indices of the lower triangular elements. + + >>> trili = np.tril_indices_from(a) + >>> trili + (array([0, 1, 1, 2, 2, 2, 3, 3, 3, 3]), + array([0, 0, 1, 0, 1, 2, 0, 1, 2, 3])) + + >>> a[trili] + array([ 0, 4, 5, 8, 9, 10, 12, 13, 14, 15]) + + This is syntactic sugar for tril_indices(). + + >>> np.tril_indices(a.shape[0]) + (array([0, 1, 1, 2, 2, 2, 3, 3, 3, 3]), + array([0, 0, 1, 0, 1, 2, 0, 1, 2, 3])) + + Use the `k` parameter to return the indices for the lower triangular array + up to the k-th diagonal. + + >>> trili1 = np.tril_indices_from(a, k=1) + >>> a[trili1] + array([ 0, 1, 4, 5, 6, 8, 9, 10, 11, 12, 13, 14, 15]) + + """ - return call_origin(numpy.tril_indices_from, x1, k) + if arr.ndim != 2: + raise ValueError("input array must be 2-d") + return tril_indices(arr.shape[-2], k=k, m=arr.shape[-1]) def triu_indices(n, k=0, m=None): """ Return the indices for the upper-triangle of an (n, m) array. + For full documentation refer to :obj:`numpy.triu_indices`. + Parameters ---------- n : int @@ -1067,24 +1333,65 @@ def triu_indices(n, k=0, m=None): The indices for the triangle. The returned tuple contains two arrays, each with the indices along one dimension of the array. Can be used to slice a ndarray of shape(`n`, `n`). + + Examples + -------- + Compute two different sets of indices to access 4x4 arrays, one for the + upper triangular part starting at the main diagonal, and one starting two + diagonals further right: + + >>> import dpnp as np + >>> iu1 = np.triu_indices(4) + >>> iu2 = np.triu_indices(4, 2) + + Here is how they can be used with a sample array: + + >>> a = np.arange(16).reshape(4, 4) + >>> a + array([[ 0, 1, 2, 3], + [ 4, 5, 6, 7], + [ 8, 9, 10, 11], + [12, 13, 14, 15]]) + + Both for indexing: + + >>> a[iu1] + array([ 0, 1, 2, ..., 10, 11, 15]) + + And for assigning values: + + >>> a[iu1] = -1 + >>> a + array([[-1, -1, -1, -1], + [ 4, -1, -1, -1], + [ 8, 9, -1, -1], + [12, 13, 14, -1]]) + + These cover only a small part of the whole array (two diagonals right + of the main one): + + >>> a[iu2] = -10 + >>> a + array([[ -1, -1, -10, -10], + [ 4, -1, -1, -10], + [ 8, 9, -1, -1], + [ 12, 13, 14, -1]]) + """ - if not use_origin_backend(): - if ( - isinstance(n, int) - and isinstance(k, int) - and (isinstance(m, int) or m is None) - ): - return dpnp_triu_indices(n, k, m) + tri_ = ~dpnp.tri(n, m, k=k - 1, dtype=bool) - return call_origin(numpy.triu_indices, n, k, m) + return tuple( + dpnp.broadcast_to(inds, tri_.shape)[tri_] + for inds in indices(tri_.shape, sparse=True) + ) -def triu_indices_from(x1, k=0): +def triu_indices_from(arr, k=0): """ Return the indices for the lower-triangle of arr. - See `tril_indices` for full details. + For full documentation refer to :obj:`numpy.triu_indices_from`. Parameters ---------- @@ -1094,11 +1401,44 @@ def triu_indices_from(x1, k=0): k : int, optional Diagonal offset (see `tril` for details). - """ - x1_desc = dpnp.get_dpnp_descriptor(x1, copy_when_nondefault_queue=False) - if x1_desc: - if isinstance(k, int): - return dpnp_triu_indices_from(x1_desc, k) + Examples + -------- + Create a 4 by 4 array. + + >>> import dpnp as np + >>> a = np.arange(16).reshape(4, 4) + >>> a + array([[ 0, 1, 2, 3], + [ 4, 5, 6, 7], + [ 8, 9, 10, 11], + [12, 13, 14, 15]]) + + Pass the array to get the indices of the upper triangular elements. + + >>> triui = np.triu_indices_from(a) + >>> triui + (array([0, 0, 0, 0, 1, 1, 1, 2, 2, 3]), + array([0, 1, 2, 3, 1, 2, 3, 2, 3, 3])) + + >>> a[triui] + array([ 0, 1, 2, 3, 5, 6, 7, 10, 11, 15]) + + This is syntactic sugar for triu_indices(). + + >>> np.triu_indices(a.shape[0]) + (array([0, 0, 0, 0, 1, 1, 1, 2, 2, 3]), + array([0, 1, 2, 3, 1, 2, 3, 2, 3, 3])) + + Use the `k` parameter to return the indices for the upper triangular array + from the k-th diagonal. + + >>> triuim1 = np.triu_indices_from(a, k=1) + >>> a[triuim1] + array([ 1, 2, 3, 6, 7, 11]) + + """ - return call_origin(numpy.triu_indices_from, x1, k) + if arr.ndim != 2: + raise ValueError("input array must be 2-d") + return triu_indices(arr.shape[-2], k=k, m=arr.shape[-1]) diff --git a/tests/skipped_tests.tbl b/tests/skipped_tests.tbl index d6fd82b9e3e3..781e6710e7f3 100644 --- a/tests/skipped_tests.tbl +++ b/tests/skipped_tests.tbl @@ -169,39 +169,6 @@ tests/third_party/cupy/indexing_tests/test_indexing.py::TestSelect::test_select_ tests/third_party/cupy/indexing_tests/test_indexing.py::TestSelect::test_select_odd_shaped_broadcastable_complex tests/third_party/cupy/indexing_tests/test_indexing.py::TestSelect::test_select_odd_shaped_non_broadcastable -tests/third_party/cupy/indexing_tests/test_insert.py::TestDiagIndicesFromRaises_param_4_{shape=(-1,)}::test_non_equal_dims -tests/third_party/cupy/indexing_tests/test_insert.py::TestFillDiagonal_param_0_{shape=(3, 3), val=1, wrap=True}::test_columnar_slice -tests/third_party/cupy/indexing_tests/test_insert.py::TestFillDiagonal_param_11_{shape=(2, 2, 2), val=0, wrap=False}::test_1darray -tests/third_party/cupy/indexing_tests/test_insert.py::TestFillDiagonal_param_16_{shape=(3, 5), val=1, wrap=True}::test_columnar_slice -tests/third_party/cupy/indexing_tests/test_insert.py::TestFillDiagonal_param_17_{shape=(3, 5), val=1, wrap=False}::test_1darray -tests/third_party/cupy/indexing_tests/test_insert.py::TestFillDiagonal_param_17_{shape=(3, 5), val=1, wrap=False}::test_columnar_slice -tests/third_party/cupy/indexing_tests/test_insert.py::TestFillDiagonal_param_18_{shape=(3, 5), val=0, wrap=True}::test_columnar_slice -tests/third_party/cupy/indexing_tests/test_insert.py::TestFillDiagonal_param_19_{shape=(3, 5), val=0, wrap=False}::test_1darray -tests/third_party/cupy/indexing_tests/test_insert.py::TestFillDiagonal_param_19_{shape=(3, 5), val=0, wrap=False}::test_columnar_slice -tests/third_party/cupy/indexing_tests/test_insert.py::TestFillDiagonal_param_1_{shape=(3, 3), val=1, wrap=False}::test_1darray -tests/third_party/cupy/indexing_tests/test_insert.py::TestFillDiagonal_param_1_{shape=(3, 3), val=1, wrap=False}::test_columnar_slice -tests/third_party/cupy/indexing_tests/test_insert.py::TestFillDiagonal_param_20_{shape=(3, 5), val=(2,), wrap=True}::test_columnar_slice -tests/third_party/cupy/indexing_tests/test_insert.py::TestFillDiagonal_param_21_{shape=(3, 5), val=(2,), wrap=False}::test_columnar_slice -tests/third_party/cupy/indexing_tests/test_insert.py::TestFillDiagonal_param_22_{shape=(3, 5), val=(2, 2), wrap=True}::test_columnar_slice -tests/third_party/cupy/indexing_tests/test_insert.py::TestFillDiagonal_param_23_{shape=(3, 5), val=(2, 2), wrap=False}::test_columnar_slice -tests/third_party/cupy/indexing_tests/test_insert.py::TestFillDiagonal_param_24_{shape=(5, 3), val=1, wrap=True}::test_columnar_slice -tests/third_party/cupy/indexing_tests/test_insert.py::TestFillDiagonal_param_25_{shape=(5, 3), val=1, wrap=False}::test_1darray -tests/third_party/cupy/indexing_tests/test_insert.py::TestFillDiagonal_param_25_{shape=(5, 3), val=1, wrap=False}::test_columnar_slice -tests/third_party/cupy/indexing_tests/test_insert.py::TestFillDiagonal_param_26_{shape=(5, 3), val=0, wrap=True}::test_columnar_slice -tests/third_party/cupy/indexing_tests/test_insert.py::TestFillDiagonal_param_27_{shape=(5, 3), val=0, wrap=False}::test_1darray -tests/third_party/cupy/indexing_tests/test_insert.py::TestFillDiagonal_param_27_{shape=(5, 3), val=0, wrap=False}::test_columnar_slice -tests/third_party/cupy/indexing_tests/test_insert.py::TestFillDiagonal_param_28_{shape=(5, 3), val=(2,), wrap=True}::test_columnar_slice -tests/third_party/cupy/indexing_tests/test_insert.py::TestFillDiagonal_param_29_{shape=(5, 3), val=(2,), wrap=False}::test_columnar_slice -tests/third_party/cupy/indexing_tests/test_insert.py::TestFillDiagonal_param_2_{shape=(3, 3), val=0, wrap=True}::test_columnar_slice -tests/third_party/cupy/indexing_tests/test_insert.py::TestFillDiagonal_param_30_{shape=(5, 3), val=(2, 2), wrap=True}::test_columnar_slice -tests/third_party/cupy/indexing_tests/test_insert.py::TestFillDiagonal_param_31_{shape=(5, 3), val=(2, 2), wrap=False}::test_columnar_slice -tests/third_party/cupy/indexing_tests/test_insert.py::TestFillDiagonal_param_3_{shape=(3, 3), val=0, wrap=False}::test_1darray -tests/third_party/cupy/indexing_tests/test_insert.py::TestFillDiagonal_param_3_{shape=(3, 3), val=0, wrap=False}::test_columnar_slice -tests/third_party/cupy/indexing_tests/test_insert.py::TestFillDiagonal_param_4_{shape=(3, 3), val=(2,), wrap=True}::test_columnar_slice -tests/third_party/cupy/indexing_tests/test_insert.py::TestFillDiagonal_param_5_{shape=(3, 3), val=(2,), wrap=False}::test_columnar_slice -tests/third_party/cupy/indexing_tests/test_insert.py::TestFillDiagonal_param_6_{shape=(3, 3), val=(2, 2), wrap=True}::test_columnar_slice -tests/third_party/cupy/indexing_tests/test_insert.py::TestFillDiagonal_param_7_{shape=(3, 3), val=(2, 2), wrap=False}::test_columnar_slice -tests/third_party/cupy/indexing_tests/test_insert.py::TestFillDiagonal_param_9_{shape=(2, 2, 2), val=1, wrap=False}::test_1darray 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 73087bde65e0..54e3e5fc6bad 100644 --- a/tests/skipped_tests_gpu.tbl +++ b/tests/skipped_tests_gpu.tbl @@ -218,57 +218,6 @@ tests/third_party/cupy/indexing_tests/test_indexing.py::TestIndexing::test_diago tests/third_party/cupy/indexing_tests/test_indexing.py::TestIndexing::test_diagonal_negative3 tests/third_party/cupy/indexing_tests/test_indexing.py::TestIndexing::test_diagonal_negative4 tests/third_party/cupy/indexing_tests/test_indexing.py::TestIndexing::test_diagonal_negative5 - -tests/third_party/cupy/indexing_tests/test_insert.py::TestDiagIndices_param_0_{n=2, ndim=2}::test_diag_indices -tests/third_party/cupy/indexing_tests/test_insert.py::TestDiagIndices_param_1_{n=2, ndim=3}::test_diag_indices -tests/third_party/cupy/indexing_tests/test_insert.py::TestDiagIndices_param_2_{n=2, ndim=1}::test_diag_indices -tests/third_party/cupy/indexing_tests/test_insert.py::TestDiagIndices_param_5_{n=4, ndim=2}::test_diag_indices -tests/third_party/cupy/indexing_tests/test_insert.py::TestDiagIndices_param_6_{n=4, ndim=3}::test_diag_indices -tests/third_party/cupy/indexing_tests/test_insert.py::TestDiagIndices_param_7_{n=4, ndim=1}::test_diag_indices -tests/third_party/cupy/indexing_tests/test_insert.py::TestDiagIndices_param_10_{n=-3, ndim=2}::test_diag_indices -tests/third_party/cupy/indexing_tests/test_insert.py::TestDiagIndices_param_11_{n=-3, ndim=3}::test_diag_indices -tests/third_party/cupy/indexing_tests/test_insert.py::TestDiagIndices_param_12_{n=-3, ndim=1}::test_diag_indices -tests/third_party/cupy/indexing_tests/test_insert.py::TestDiagIndices_param_15_{n=0, ndim=2}::test_diag_indices -tests/third_party/cupy/indexing_tests/test_insert.py::TestDiagIndices_param_16_{n=0, ndim=3}::test_diag_indices -tests/third_party/cupy/indexing_tests/test_insert.py::TestDiagIndices_param_17_{n=0, ndim=1}::test_diag_indices -tests/third_party/cupy/indexing_tests/test_insert.py::TestDiagIndicesInvalidValues_param_0_{n=-3, ndim=1}::test_diag_indices -tests/third_party/cupy/indexing_tests/test_insert.py::TestDiagIndicesInvalidValues_param_3_{n=0, ndim=1}::test_diag_indices -tests/third_party/cupy/indexing_tests/test_insert.py::TestDiagIndicesFrom_param_0_{shape=(3, 3)}::test_diag_indices_from -tests/third_party/cupy/indexing_tests/test_insert.py::TestDiagIndicesFrom_param_1_{shape=(0, 0)}::test_diag_indices_from -tests/third_party/cupy/indexing_tests/test_insert.py::TestDiagIndicesFrom_param_2_{shape=(2, 2, 2)}::test_diag_indices_from - -tests/third_party/cupy/indexing_tests/test_insert.py::TestFillDiagonal_param_0_{shape=(3, 3), val=1, wrap=True}::test_columnar_slice -tests/third_party/cupy/indexing_tests/test_insert.py::TestFillDiagonal_param_11_{shape=(2, 2, 2), val=0, wrap=False}::test_1darray -tests/third_party/cupy/indexing_tests/test_insert.py::TestFillDiagonal_param_16_{shape=(3, 5), val=1, wrap=True}::test_columnar_slice -tests/third_party/cupy/indexing_tests/test_insert.py::TestFillDiagonal_param_17_{shape=(3, 5), val=1, wrap=False}::test_1darray -tests/third_party/cupy/indexing_tests/test_insert.py::TestFillDiagonal_param_17_{shape=(3, 5), val=1, wrap=False}::test_columnar_slice -tests/third_party/cupy/indexing_tests/test_insert.py::TestFillDiagonal_param_18_{shape=(3, 5), val=0, wrap=True}::test_columnar_slice -tests/third_party/cupy/indexing_tests/test_insert.py::TestFillDiagonal_param_19_{shape=(3, 5), val=0, wrap=False}::test_1darray -tests/third_party/cupy/indexing_tests/test_insert.py::TestFillDiagonal_param_19_{shape=(3, 5), val=0, wrap=False}::test_columnar_slice -tests/third_party/cupy/indexing_tests/test_insert.py::TestFillDiagonal_param_1_{shape=(3, 3), val=1, wrap=False}::test_1darray -tests/third_party/cupy/indexing_tests/test_insert.py::TestFillDiagonal_param_1_{shape=(3, 3), val=1, wrap=False}::test_columnar_slice -tests/third_party/cupy/indexing_tests/test_insert.py::TestFillDiagonal_param_20_{shape=(3, 5), val=(2,), wrap=True}::test_columnar_slice -tests/third_party/cupy/indexing_tests/test_insert.py::TestFillDiagonal_param_21_{shape=(3, 5), val=(2,), wrap=False}::test_columnar_slice -tests/third_party/cupy/indexing_tests/test_insert.py::TestFillDiagonal_param_22_{shape=(3, 5), val=(2, 2), wrap=True}::test_columnar_slice -tests/third_party/cupy/indexing_tests/test_insert.py::TestFillDiagonal_param_23_{shape=(3, 5), val=(2, 2), wrap=False}::test_columnar_slice -tests/third_party/cupy/indexing_tests/test_insert.py::TestFillDiagonal_param_24_{shape=(5, 3), val=1, wrap=True}::test_columnar_slice -tests/third_party/cupy/indexing_tests/test_insert.py::TestFillDiagonal_param_25_{shape=(5, 3), val=1, wrap=False}::test_1darray -tests/third_party/cupy/indexing_tests/test_insert.py::TestFillDiagonal_param_25_{shape=(5, 3), val=1, wrap=False}::test_columnar_slice -tests/third_party/cupy/indexing_tests/test_insert.py::TestFillDiagonal_param_26_{shape=(5, 3), val=0, wrap=True}::test_columnar_slice -tests/third_party/cupy/indexing_tests/test_insert.py::TestFillDiagonal_param_27_{shape=(5, 3), val=0, wrap=False}::test_1darray -tests/third_party/cupy/indexing_tests/test_insert.py::TestFillDiagonal_param_27_{shape=(5, 3), val=0, wrap=False}::test_columnar_slice -tests/third_party/cupy/indexing_tests/test_insert.py::TestFillDiagonal_param_28_{shape=(5, 3), val=(2,), wrap=True}::test_columnar_slice -tests/third_party/cupy/indexing_tests/test_insert.py::TestFillDiagonal_param_29_{shape=(5, 3), val=(2,), wrap=False}::test_columnar_slice -tests/third_party/cupy/indexing_tests/test_insert.py::TestFillDiagonal_param_2_{shape=(3, 3), val=0, wrap=True}::test_columnar_slice -tests/third_party/cupy/indexing_tests/test_insert.py::TestFillDiagonal_param_30_{shape=(5, 3), val=(2, 2), wrap=True}::test_columnar_slice -tests/third_party/cupy/indexing_tests/test_insert.py::TestFillDiagonal_param_31_{shape=(5, 3), val=(2, 2), wrap=False}::test_columnar_slice -tests/third_party/cupy/indexing_tests/test_insert.py::TestFillDiagonal_param_3_{shape=(3, 3), val=0, wrap=False}::test_1darray -tests/third_party/cupy/indexing_tests/test_insert.py::TestFillDiagonal_param_3_{shape=(3, 3), val=0, wrap=False}::test_columnar_slice -tests/third_party/cupy/indexing_tests/test_insert.py::TestFillDiagonal_param_4_{shape=(3, 3), val=(2,), wrap=True}::test_columnar_slice -tests/third_party/cupy/indexing_tests/test_insert.py::TestFillDiagonal_param_5_{shape=(3, 3), val=(2,), wrap=False}::test_columnar_slice -tests/third_party/cupy/indexing_tests/test_insert.py::TestFillDiagonal_param_6_{shape=(3, 3), val=(2, 2), wrap=True}::test_columnar_slice -tests/third_party/cupy/indexing_tests/test_insert.py::TestFillDiagonal_param_7_{shape=(3, 3), val=(2, 2), wrap=False}::test_columnar_slice -tests/third_party/cupy/indexing_tests/test_insert.py::TestFillDiagonal_param_9_{shape=(2, 2, 2), val=1, wrap=False}::test_1darray 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/third_party/cupy/indexing_tests/test_insert.py b/tests/third_party/cupy/indexing_tests/test_insert.py index 28d25c487b0c..4c886309e2e7 100644 --- a/tests/third_party/cupy/indexing_tests/test_insert.py +++ b/tests/third_party/cupy/indexing_tests/test_insert.py @@ -255,7 +255,6 @@ def test_putmask_differnt_dtypes_mask(self, xp, dtype): } ) ) -@pytest.mark.usefixtures("allow_fall_back_on_numpy") class TestFillDiagonal(unittest.TestCase): def _compute_val(self, xp): if type(self.val) is int: @@ -341,7 +340,6 @@ def test_diag_indices_from(self, xp): } ) ) -@pytest.mark.usefixtures("allow_fall_back_on_numpy") class TestDiagIndicesFromRaises(unittest.TestCase): def test_non_equal_dims(self): for xp in (numpy, cupy): From 24c3b97bd98eea37c68845c19fbf028785f105a3 Mon Sep 17 00:00:00 2001 From: Natalia Polina Date: Thu, 16 May 2024 15:59:56 -0500 Subject: [PATCH 2/6] Add CFD and tests --- dpnp/backend/kernels/dpnp_krnl_indexing.cpp | 11 +- dpnp/dpnp_algo/dpnp_algo_indexing.pxi | 3 - dpnp/dpnp_iface_indexing.py | 299 +++++++++++++++--- tests/test_indexing.py | 20 +- tests/test_sycl_queue.py | 59 ++++ tests/test_usm_type.py | 40 +++ .../cupy/indexing_tests/test_generate.py | 127 ++++++++ 7 files changed, 499 insertions(+), 60 deletions(-) diff --git a/dpnp/backend/kernels/dpnp_krnl_indexing.cpp b/dpnp/backend/kernels/dpnp_krnl_indexing.cpp index 969fa903c1c9..172a1f78ffbd 100644 --- a/dpnp/backend/kernels/dpnp_krnl_indexing.cpp +++ b/dpnp/backend/kernels/dpnp_krnl_indexing.cpp @@ -181,10 +181,11 @@ DPCTLSyclEventRef dpnp_diagonal_c(DPCTLSyclQueueRef q_ref, _DataType *array_1 = input1_ptr.get_ptr(); _DataType *result = result_ptr.get_ptr(); + const size_t res_shape_ndim_sub_1 = + static_cast(res_shape[res_ndim - 1]); + if (res_ndim <= 1) { - for (size_t i = 0; i < static_cast(res_shape[res_ndim - 1]); - ++i) - { + for (size_t i = 0; i < res_shape_ndim_sub_1; ++i) { result[i] = array_1[i * shape[res_ndim] + i + offset]; } } @@ -219,9 +220,7 @@ DPCTLSyclEventRef dpnp_diagonal_c(DPCTLSyclQueueRef q_ref, index += 1; } - for (size_t i = 0; i < static_cast(res_shape[res_ndim - 1]); - i++) - { + for (size_t i = 0; i < res_shape_ndim_sub_1; i++) { for (size_t j = 0; j < xyz.size(); j++) { std::vector ind_list = xyz[j]; if (ind_list.size() == 0) { diff --git a/dpnp/dpnp_algo/dpnp_algo_indexing.pxi b/dpnp/dpnp_algo/dpnp_algo_indexing.pxi index 34e3bb03eed5..fbdb582f9bbf 100644 --- a/dpnp/dpnp_algo/dpnp_algo_indexing.pxi +++ b/dpnp/dpnp_algo/dpnp_algo_indexing.pxi @@ -54,9 +54,6 @@ ctypedef c_dpctl.DPCTLSyclEventRef(*custom_indexing_2in_1out_func_ptr_t_)(c_dpct shape_elem_type *, const size_t, const c_dpctl.DPCTLEventVectorRef) -ctypedef c_dpctl.DPCTLSyclEventRef(*custom_indexing_2in_func_ptr_t)(c_dpctl.DPCTLSyclQueueRef, - void *, void * , shape_elem_type * , const size_t, - const c_dpctl.DPCTLEventVectorRef) cpdef utils.dpnp_descriptor dpnp_choose(utils.dpnp_descriptor x1, list choices1): diff --git a/dpnp/dpnp_iface_indexing.py b/dpnp/dpnp_iface_indexing.py index a12283514f2f..db5b72cc589c 100644 --- a/dpnp/dpnp_iface_indexing.py +++ b/dpnp/dpnp_iface_indexing.py @@ -177,7 +177,7 @@ def choose(x1, choices, out=None, mode="raise"): return call_origin(numpy.choose, x1, choices, out, mode) -def diag_indices(n, ndim=2): +def diag_indices(n, ndim=2, device=None, usm_type="device", sycl_queue=None): """ Return the indices to access the main diagonal of an array. @@ -192,11 +192,26 @@ def diag_indices(n, ndim=2): Parameters ---------- n : int - The size, along each dimension, of the arrays for which the returned - indices can be used. - + The size, along each dimension, of the arrays for which the returned + indices can be used. ndim : int, optional - The number of dimensions. + The number of dimensions. Default: ``2``. + 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. + + Return + ------ + out : tuple of dpnp.ndarray + The indices to access the main diagonal of an array. See also -------- @@ -242,11 +257,16 @@ def diag_indices(n, ndim=2): """ - idx = dpnp.arange(n) + idx = dpnp.arange( + n, + device=device, + usm_type=usm_type, + sycl_queue=sycl_queue, + ) return (idx,) * ndim -def diag_indices_from(arr): +def diag_indices_from(arr, device=None, usm_type=None, sycl_queue=None): """ Return the indices to access the main diagonal of an n-dimensional array. @@ -254,7 +274,24 @@ def diag_indices_from(arr): Parameters ---------- - arr : array, at least 2-D + arr : dpnp.ndarray + Array at least 2-D + 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. + + Return + ------ + out : tuple of dpnp.ndarray + The indices to access the main diagonal of an n-dimensional array. See also -------- @@ -289,6 +326,8 @@ def diag_indices_from(arr): """ + dpnp.check_supported_arrays_type(arr) + if not arr.ndim >= 2: raise ValueError("input array must be at least 2-d") @@ -297,7 +336,16 @@ def diag_indices_from(arr): if sh != sh0: raise ValueError("All dimensions of input must be of equal length") - return diag_indices(sh0, arr.ndim) + _usm_type = arr.usm_type if usm_type is None else usm_type + _sycl_queue = dpnp.get_normalized_queue_device( + arr, sycl_queue=sycl_queue, device=device + ) + return diag_indices( + sh0, + arr.ndim, + usm_type=_usm_type, + sycl_queue=_sycl_queue, + ) def diagonal(x1, offset=0, axis1=0, axis2=1): @@ -377,15 +425,15 @@ def fill_diagonal(a, val, wrap=False): Parameters ---------- a : array, at least 2-D. - Array whose diagonal is to be filled in-place. - val : scalar or array_like - Value(s) to write on the diagonal. If `val` is scalar, the value is - written along the diagonal. If array-like, the flattened `val` is - written along the diagonal, repeating if necessary to fill all - diagonal entries. + Array whose diagonal is to be filled in-place. + val : {dpnp.ndarray, usm_ndarray, scalar} + Value(s) to write on the diagonal. If `val` is scalar, the value is + written along the diagonal. If array-like, the flattened `val` is + written along the diagonal, repeating if necessary to fill all + diagonal entries. wrap : bool - The diagonal "wrapped" after N columns. You can have this behavior - with this option. This affects only tall matrices. + The diagonal "wrapped" after N columns. You can have this behavior + with this option. This affects only tall matrices. Default: ``False``. See Also -------- @@ -619,13 +667,20 @@ def indices( return res -def mask_indices(n, mask_func, k=0): +def mask_indices( + n, + mask_func, + k=0, + device=None, + usm_type="device", + sycl_queue=None, +): """ Return the indices to access (n, n) arrays, given a masking function. Assume `mask_func` is a function that, for a square array a of size ``(n, n)`` with a possible offset argument `k`, when called as - ``mask_func(a, k)`` returns a new array with zeros in certain locations + ``mask_func(a, k=k)`` returns a new array with zeros in certain locations (functions like `triu` or `tril` do precisely this). Then this function returns the indices where the non-zero values would be located. @@ -635,16 +690,27 @@ def mask_indices(n, mask_func, k=0): The returned indices will be valid to access arrays of shape (n, n). mask_func : callable A function whose call signature is similar to that of `triu`, `tril`. - That is, ``mask_func(x, k)`` returns a boolean array, shaped like `x`. + That is, ``mask_func(x, k=k)`` returns a boolean array, shaped like `x`. `k` is an optional argument to the function. k : scalar An optional argument which is passed through to `mask_func`. Functions like `triu`, `tril` take a second argument that is interpreted as an offset. + 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 ------- - indices : tuple of arrays. + indices : tuple of dpnp.ndarray The `n` arrays of indices corresponding to the locations where ``mask_func(np.ones((n, n)), k)`` is True. @@ -678,8 +744,14 @@ def mask_indices(n, mask_func, k=0): """ - m = dpnp.ones((n, n), dtype=int) - a = mask_func(m, k) + m = dpnp.ones( + (n, n), + dtype=int, + device=device, + usm_type=usm_type, + sycl_queue=sycl_queue, + ) + a = mask_func(m, k=k) return nonzero(a != 0) @@ -1172,7 +1244,14 @@ def take_along_axis(a, indices, axis): return a[_build_along_axis_index(a, indices, axis)] -def tril_indices(n, k=0, m=None): +def tril_indices( + n, + k=0, + m=None, + device=None, + usm_type="device", + sycl_queue=None, +): """ Return the indices for the lower-triangle of an (n, m) array. @@ -1183,18 +1262,27 @@ def tril_indices(n, k=0, m=None): n : int The row dimension of the arrays for which the returned indices will be valid. - k : int, optional - Diagonal offset (see `tril` for details). - + Diagonal offset (see `tril` for details). Default: ``0``. m : int, optional The column dimension of the arrays for which the returned arrays will be valid. By default `m` is taken equal to `n`. + 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 ------- - inds : tuple of arrays + inds : tuple of dpnp.ndarray The indices for the triangle. The returned tuple contains two arrays, each with the indices along one dimension of the array. @@ -1242,15 +1330,35 @@ def tril_indices(n, k=0, m=None): """ - tri_ = dpnp.tri(n, m, k=k, dtype=bool) + tri_ = dpnp.tri( + n, + m, + k=k, + dtype=bool, + device=device, + usm_type=usm_type, + sycl_queue=sycl_queue, + ) return tuple( dpnp.broadcast_to(inds, tri_.shape)[tri_] - for inds in indices(tri_.shape, sparse=True) + for inds in indices( + tri_.shape, + sparse=True, + device=device, + usm_type=usm_type, + sycl_queue=sycl_queue, + ) ) -def tril_indices_from(arr, k=0): +def tril_indices_from( + arr, + k=0, + device=None, + usm_type=None, + sycl_queue=None, +): """ Return the indices for the lower-triangle of arr. @@ -1258,12 +1366,28 @@ def tril_indices_from(arr, k=0): Parameters ---------- - arr : array_like + arr : {dpnp.ndarray, usm_ndarray} The indices will be valid for square arrays whose dimensions are the same as arr. - k : int, optional - Diagonal offset (see `tril` for details). + Diagonal offset (see `tril` for details). Default: ``0``. + 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 + ------- + inds : tuple of dpnp.ndarray + The indices for the triangle. The returned tuple contains two arrays, + each with the indices along one dimension of the array. Examples -------- @@ -1302,12 +1426,33 @@ def tril_indices_from(arr, k=0): """ + dpnp.check_supported_arrays_type(arr) + if arr.ndim != 2: raise ValueError("input array must be 2-d") - return tril_indices(arr.shape[-2], k=k, m=arr.shape[-1]) + _usm_type = arr.usm_type if usm_type is None else usm_type + _sycl_queue = dpnp.get_normalized_queue_device( + arr, sycl_queue=sycl_queue, device=device + ) + + return tril_indices( + arr.shape[-2], + k=k, + m=arr.shape[-1], + usm_type=_usm_type, + sycl_queue=_sycl_queue, + ) -def triu_indices(n, k=0, m=None): + +def triu_indices( + n, + k=0, + m=None, + device=None, + usm_type="device", + sycl_queue=None, +): """ Return the indices for the upper-triangle of an (n, m) array. @@ -1318,18 +1463,27 @@ def triu_indices(n, k=0, m=None): n : int The size of the arrays for which the returned indices will be valid. - k : int, optional - Diagonal offset (see `triu` for details). - + Diagonal offset (see `triu` for details). Default: ``0``. m : int, optional The column dimension of the arrays for which the returned arrays will be valid. By default `m` is taken equal to `n`. + 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 ------- - inds : tuple, shape(2) of ndarrays, shape(`n`) + inds : tuple of dpnp.ndarray The indices for the triangle. The returned tuple contains two arrays, each with the indices along one dimension of the array. Can be used to slice a ndarray of shape(`n`, `n`). @@ -1379,15 +1533,35 @@ def triu_indices(n, k=0, m=None): """ - tri_ = ~dpnp.tri(n, m, k=k - 1, dtype=bool) + tri_ = ~dpnp.tri( + n, + m, + k=k - 1, + dtype=bool, + device=device, + usm_type=usm_type, + sycl_queue=sycl_queue, + ) return tuple( dpnp.broadcast_to(inds, tri_.shape)[tri_] - for inds in indices(tri_.shape, sparse=True) + for inds in indices( + tri_.shape, + sparse=True, + device=device, + usm_type=usm_type, + sycl_queue=sycl_queue, + ) ) -def triu_indices_from(arr, k=0): +def triu_indices_from( + arr, + k=0, + device=None, + usm_type=None, + sycl_queue=None, +): """ Return the indices for the lower-triangle of arr. @@ -1395,12 +1569,29 @@ def triu_indices_from(arr, k=0): Parameters ---------- - arr : array_like + arr : {dpnp.ndarray, usm_ndarray} The indices will be valid for square arrays whose dimensions are the same as arr. - k : int, optional - Diagonal offset (see `tril` for details). + Diagonal offset (see `triu` for details). Default: ``0``. + 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 + ------- + inds : tuple of dpnp.ndarray + The indices for the triangle. The returned tuple contains two arrays, + each with the indices along one dimension of the array. Can be used + to slice a ndarray of shape(`n`, `n`). Examples -------- @@ -1439,6 +1630,20 @@ def triu_indices_from(arr, k=0): """ + dpnp.check_supported_arrays_type(arr) + if arr.ndim != 2: raise ValueError("input array must be 2-d") - return triu_indices(arr.shape[-2], k=k, m=arr.shape[-1]) + + _usm_type = arr.usm_type if usm_type is None else usm_type + _sycl_queue = dpnp.get_normalized_queue_device( + arr, sycl_queue=sycl_queue, device=device + ) + + return triu_indices( + arr.shape[-2], + k=k, + m=arr.shape[-1], + usm_type=_usm_type, + sycl_queue=_sycl_queue, + ) diff --git a/tests/test_indexing.py b/tests/test_indexing.py index e0c217c7ab27..79406eac8a53 100644 --- a/tests/test_indexing.py +++ b/tests/test_indexing.py @@ -871,7 +871,7 @@ def test_take_over_index(indices, array_type, mode): "m", [None, 0, 1, 2, 3, 4], ids=["None", "0", "1", "2", "3", "4"] ) @pytest.mark.parametrize( - "k", [0, 1, 2, 3, 4, 5], ids=["0", "1", "2", "3", "4", "5"] + "k", [-3, -2, -1, 0, 1, 2, 3], ids=["-3", "-2", "-1", "0", "1", "2", "3"] ) @pytest.mark.parametrize( "n", [1, 2, 3, 4, 5, 6], ids=["1", "2", "3", "4", "5", "6"] @@ -883,7 +883,7 @@ def test_tril_indices(n, k, m): @pytest.mark.parametrize( - "k", [0, 1, 2, 3, 4, 5], ids=["0", "1", "2", "3", "4", "5"] + "k", [-3, -2, -1, 0, 1, 2, 3], ids=["-3", "-2", "-1", "0", "1", "2", "3"] ) @pytest.mark.parametrize( "array", @@ -906,7 +906,7 @@ def test_tril_indices_from(array, k): "m", [None, 0, 1, 2, 3, 4], ids=["None", "0", "1", "2", "3", "4"] ) @pytest.mark.parametrize( - "k", [0, 1, 2, 3, 4, 5], ids=["0", "1", "2", "3", "4", "5"] + "k", [-3, -2, -1, 0, 1, 2, 3], ids=["-3", "-2", "-1", "0", "1", "2", "3"] ) @pytest.mark.parametrize( "n", [1, 2, 3, 4, 5, 6], ids=["1", "2", "3", "4", "5", "6"] @@ -918,7 +918,7 @@ def test_triu_indices(n, k, m): @pytest.mark.parametrize( - "k", [0, 1, 2, 3, 4, 5], ids=["0", "1", "2", "3", "4", "5"] + "k", [-3, -2, -1, 0, 1, 2, 3], ids=["-3", "-2", "-1", "0", "1", "2", "3"] ) @pytest.mark.parametrize( "array", @@ -935,3 +935,15 @@ def test_triu_indices_from(array, k): result = dpnp.triu_indices_from(ia, k) expected = numpy.triu_indices_from(a, k) assert_array_equal(expected, result) + + +def test_indices_from_err(): + arr = dpnp.array([1, 2, 3]) + with pytest.raises(ValueError): + dpnp.tril_indices_from(arr) + with pytest.raises(ValueError): + dpnp.triu_indices_from(arr) + with pytest.raises(ValueError): + dpnp.diag_indices_from(arr) + with pytest.raises(ValueError): + dpnp.diag_indices_from(dpnp.ones((2, 3))) diff --git a/tests/test_sycl_queue.py b/tests/test_sycl_queue.py index d79ae8b3337d..3ca17120be1e 100644 --- a/tests/test_sycl_queue.py +++ b/tests/test_sycl_queue.py @@ -2101,3 +2101,62 @@ def test_histogram(weights, device): edges_queue = result_edges.sycl_queue assert_sycl_queue_equal(hist_queue, iv.sycl_queue) assert_sycl_queue_equal(edges_queue, iv.sycl_queue) + + +@pytest.mark.parametrize( + "func", ["tril_indices_from", "triu_indices_from", "diag_indices_from"] +) +@pytest.mark.parametrize( + "device", + valid_devices, + ids=[device.filter_string for device in valid_devices], +) +def test_tri_diag_indices_from_follow(func, device): + arr = dpnp.ones((3, 3), device=device) + res = getattr(dpnp, func)(arr) + for x in res: + assert_sycl_queue_equal(x.sycl_queue, arr.sycl_queue) + + +@pytest.mark.parametrize( + "func", ["tril_indices_from", "triu_indices_from", "diag_indices_from"] +) +@pytest.mark.parametrize( + "device", + valid_devices, + ids=[device.filter_string for device in valid_devices], +) +def test_tri_diag_indices_from(func, device): + sycl_queue = dpctl.SyclQueue(device) + arr = dpnp.ones((3, 3)) + res = getattr(dpnp, func)(arr, sycl_queue=sycl_queue) + for x in res: + assert_sycl_queue_equal(x.sycl_queue, sycl_queue) + + +@pytest.mark.parametrize( + "func", ["tril_indices", "triu_indices", "diag_indices"] +) +@pytest.mark.parametrize( + "device", + valid_devices, + ids=[device.filter_string for device in valid_devices], +) +def test_tri_diag_indices(func, device): + sycl_queue = dpctl.SyclQueue(device) + res = getattr(dpnp, func)(4, sycl_queue=sycl_queue) + for x in res: + assert_sycl_queue_equal(x.sycl_queue, sycl_queue) + + +@pytest.mark.parametrize("mask_func", ["tril", "triu"]) +@pytest.mark.parametrize( + "device", + valid_devices, + ids=[device.filter_string for device in valid_devices], +) +def test_mask_indices(mask_func, device): + sycl_queue = dpctl.SyclQueue(device) + res = dpnp.mask_indices(4, getattr(dpnp, mask_func), sycl_queue=sycl_queue) + for x in res: + assert_sycl_queue_equal(x.sycl_queue, sycl_queue) diff --git a/tests/test_usm_type.py b/tests/test_usm_type.py index 75e70a978bc1..2aeb9ebcad79 100644 --- a/tests/test_usm_type.py +++ b/tests/test_usm_type.py @@ -1228,3 +1228,43 @@ def test_histogram(usm_type_v, usm_type_w): assert w.usm_type == usm_type_w assert hist.usm_type == du.get_coerced_usm_type([usm_type_v, usm_type_w]) assert edges.usm_type == du.get_coerced_usm_type([usm_type_v, usm_type_w]) + + +@pytest.mark.parametrize( + "func", ["tril_indices_from", "triu_indices_from", "diag_indices_from"] +) +@pytest.mark.parametrize("usm_type", list_of_usm_types, ids=list_of_usm_types) +def test_tri_diag_indices_from_follow(func, usm_type): + arr = dp.ones((3, 3), usm_type=usm_type) + res = getattr(dp, func)(arr) + for x in res: + assert x.usm_type == usm_type + + +@pytest.mark.parametrize( + "func", ["tril_indices_from", "triu_indices_from", "diag_indices_from"] +) +@pytest.mark.parametrize("usm_type", list_of_usm_types, ids=list_of_usm_types) +def test_tri_diag_indices_from(func, usm_type): + arr = dp.ones((3, 3)) + res = getattr(dp, func)(arr, usm_type=usm_type) + for x in res: + assert x.usm_type == usm_type + + +@pytest.mark.parametrize( + "func", ["tril_indices", "triu_indices", "diag_indices"] +) +@pytest.mark.parametrize("usm_type", list_of_usm_types, ids=list_of_usm_types) +def test_tri_diag_indices(func, usm_type): + res = getattr(dp, func)(4, usm_type=usm_type) + for x in res: + assert x.usm_type == usm_type + + +@pytest.mark.parametrize("mask_func", ["tril", "triu"]) +@pytest.mark.parametrize("usm_type", list_of_usm_types, ids=list_of_usm_types) +def test_mask_indices(mask_func, usm_type): + res = dp.mask_indices(4, getattr(dp, mask_func), usm_type=usm_type) + for x in res: + assert x.usm_type == usm_type diff --git a/tests/third_party/cupy/indexing_tests/test_generate.py b/tests/third_party/cupy/indexing_tests/test_generate.py index 71aa1aa652e2..5c17f2c723e6 100644 --- a/tests/third_party/cupy/indexing_tests/test_generate.py +++ b/tests/third_party/cupy/indexing_tests/test_generate.py @@ -272,3 +272,130 @@ def test_invalid_mode(self, dtype): a = tuple([xp.arange(min(dims), dtype=dtype) for d in dims]) with pytest.raises(ValueError): xp.ravel_multi_index(a, dims, mode="invalid") + + +class TestMaskIndices: + @testing.numpy_cupy_array_equal() + def test_mask_indices(self, xp): + # arr is a square matrix with 50% density + multiplier = testing.shaped_random((10, 10), xp=xp, dtype=xp.bool_) + arr = testing.shaped_random((10, 10), xp=xp) * multiplier + return xp.mask_indices(10, lambda n, k=None: arr) + + @testing.numpy_cupy_array_equal() + def test_mask_indices_k(self, xp): + return xp.mask_indices(10, xp.triu, k=1) + + @testing.numpy_cupy_array_equal() + def test_empty(self, xp): + return xp.mask_indices(0, xp.triu) + + +class TestTrilIndices: + @testing.numpy_cupy_array_equal() + def test_tril_indices_1(self, xp): + return xp.tril_indices(n=29, k=0) + + @testing.numpy_cupy_array_equal() + def test_tril_indices_2(self, xp): + return xp.tril_indices(n=11, k=4, m=4) + + @testing.numpy_cupy_array_equal() + def test_tril_indices_3(self, xp): + return xp.tril_indices(n=4, k=4, m=3) + + @testing.for_all_dtypes() + def test_tril_indices(self, dtype): + for xp in (numpy, cupy): + arr = testing.shaped_random((10, 10), xp=xp, dtype=dtype) + if xp is numpy: + error = ValueError + else: + error = TypeError + with pytest.raises(error): + xp.tril_indices(arr, k=0) + + +class TestTrilIndicesForm: + @testing.for_all_dtypes() + @testing.numpy_cupy_array_equal() + def test_tril_indices_from_1(self, xp, dtype): + arr = testing.shaped_random((10, 10), xp=xp, dtype=dtype) + return xp.tril_indices_from(arr, k=4) + + @testing.for_all_dtypes() + @testing.numpy_cupy_array_equal() + def test_tril_indices_from_2(self, xp, dtype): + arr = testing.shaped_random((10, 20), xp=xp, dtype=dtype) + return xp.tril_indices_from(arr, k=13) + + @testing.for_all_dtypes() + @testing.numpy_cupy_array_equal() + def test_tril_indices_from_3(self, xp, dtype): + arr = testing.shaped_random((4, 6), xp=xp, dtype=dtype) + return xp.tril_indices_from(arr) + + @testing.for_all_dtypes() + def test_tril_indices_from_4(self, dtype): + for xp in (numpy, cupy): + if xp is numpy: + error = AttributeError + else: + error = TypeError + with pytest.raises(error): + xp.tril_indices_from(4, k=1) + + +class TestTriuIndices: + @testing.numpy_cupy_array_equal() + def test_triu_indices_1(self, xp): + return xp.triu_indices(n=10, k=0) + + @testing.numpy_cupy_array_equal() + def test_triu_indices_2(self, xp): + return xp.triu_indices(n=23, k=3, m=4) + + @testing.numpy_cupy_array_equal() + def test_triu_indices_3(self, xp): + return xp.triu_indices(n=4, k=4, m=4) + + @testing.for_all_dtypes() + def test_triu_indices_4(self, dtype): + for xp in (numpy, cupy): + arr = testing.shaped_random((10, 10), xp=xp, dtype=dtype) + if xp is numpy: + error = ValueError + else: + error = TypeError + with pytest.raises(error): + xp.triu_indices(arr, k=0) + + +class TestTriuIndicesFrom: + @testing.for_all_dtypes() + @testing.numpy_cupy_array_equal() + def test_triu_indices_from_1(self, xp, dtype): + arr = testing.shaped_random((20, 20), xp=xp, dtype=dtype) + return xp.triu_indices_from(arr, k=11) + + @testing.for_all_dtypes() + @testing.numpy_cupy_array_equal() + def test_triu_indices_from_2(self, xp, dtype): + arr = testing.shaped_random((20, 5), xp=xp, dtype=dtype) + return xp.triu_indices_from(arr, k=32) + + @testing.for_all_dtypes() + @testing.numpy_cupy_array_equal() + def test_triu_indices_from_3(self, xp, dtype): + arr = testing.shaped_random((4, 6), xp=xp, dtype=dtype) + return xp.triu_indices_from(arr) + + @testing.for_all_dtypes() + def test_triu_indices_from_4(self, dtype): + for xp in (numpy, cupy): + if xp is numpy: + error = AttributeError + else: + error = TypeError + with pytest.raises(error): + xp.triu_indices_from(4, k=1) From fabf3cfebc67e2bd2ef53d09d14da9c26caaad42 Mon Sep 17 00:00:00 2001 From: Natalia Polina Date: Fri, 17 May 2024 02:13:15 -0500 Subject: [PATCH 3/6] Improve code coverage --- tests/test_indexing.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/tests/test_indexing.py b/tests/test_indexing.py index b94eaf16d077..2401fc029f85 100644 --- a/tests/test_indexing.py +++ b/tests/test_indexing.py @@ -967,3 +967,9 @@ def test_indices_from_err(): dpnp.diag_indices_from(arr) with pytest.raises(ValueError): dpnp.diag_indices_from(dpnp.ones((2, 3))) + + +def test_fill_diagonal_error(): + arr = dpnp.ones((1, 2, 3)) + with pytest.raises(ValueError): + dpnp.fill_diagonal(arr, 5) From fed120e30249ff4d05693077f18d8491f1028d42 Mon Sep 17 00:00:00 2001 From: Natalia Polina Date: Mon, 20 May 2024 15:29:56 -0500 Subject: [PATCH 4/6] address comments --- dpnp/dpnp_iface_indexing.py | 193 ++++++++++++++++-------------------- tests/test_sycl_queue.py | 18 +--- tests/test_usm_type.py | 13 +-- 3 files changed, 86 insertions(+), 138 deletions(-) diff --git a/dpnp/dpnp_iface_indexing.py b/dpnp/dpnp_iface_indexing.py index 2ee10762a840..d3b3225612c1 100644 --- a/dpnp/dpnp_iface_indexing.py +++ b/dpnp/dpnp_iface_indexing.py @@ -207,8 +207,8 @@ def diag_indices(n, ndim=2, device=None, usm_type="device", sycl_queue=None): sycl_queue : {None, SyclQueue}, optional A SYCL queue to use for output array allocation and copying. - Return - ------ + Returns + ------- out : tuple of dpnp.ndarray The indices to access the main diagonal of an array. @@ -265,7 +265,7 @@ def diag_indices(n, ndim=2, device=None, usm_type="device", sycl_queue=None): return (idx,) * ndim -def diag_indices_from(arr, device=None, usm_type=None, sycl_queue=None): +def diag_indices_from(arr): """ Return the indices to access the main diagonal of an n-dimensional array. @@ -273,22 +273,11 @@ def diag_indices_from(arr, device=None, usm_type=None, sycl_queue=None): Parameters ---------- - arr : dpnp.ndarray + arr : {dpnp.ndarray, usm_ndarray} Array at least 2-D - 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. - Return - ------ + Returns + ------- out : tuple of dpnp.ndarray The indices to access the main diagonal of an n-dimensional array. @@ -330,20 +319,14 @@ def diag_indices_from(arr, device=None, usm_type=None, sycl_queue=None): if not arr.ndim >= 2: raise ValueError("input array must be at least 2-d") - sh0 = arr.shape[0] - for sh in arr.shape[1:]: - if sh != sh0: - raise ValueError("All dimensions of input must be of equal length") + if not numpy.all(numpy.diff(arr.shape) == 0): + raise ValueError("All dimensions of input must be of equal length") - _usm_type = arr.usm_type if usm_type is None else usm_type - _sycl_queue = dpnp.get_normalized_queue_device( - arr, sycl_queue=sycl_queue, device=device - ) return diag_indices( - sh0, + arr.shape[0], arr.ndim, - usm_type=_usm_type, - sycl_queue=_sycl_queue, + usm_type=arr.usm_type, + sycl_queue=arr.sycl_queue, ) @@ -545,16 +528,16 @@ def fill_diagonal(a, val, wrap=False): Parameters ---------- - a : array, at least 2-D. - Array whose diagonal is to be filled in-place. + a : {dpnp_array, usm_ndarray} + Array whose diagonal is to be filled in-place. It must be at least 2-D. val : {dpnp.ndarray, usm_ndarray, scalar} Value(s) to write on the diagonal. If `val` is scalar, the value is - written along the diagonal. If array-like, the flattened `val` is + written along the diagonal. If array, the flattened `val` is written along the diagonal, repeating if necessary to fill all diagonal entries. wrap : bool - The diagonal "wrapped" after N columns. You can have this behavior - with this option. This affects only tall matrices. Default: ``False``. + It enables the diagonal "wrapped" after N columns. This affects only + tall matrices. Default: ``False``. See Also -------- @@ -593,7 +576,7 @@ def fill_diagonal(a, val, wrap=False): [0, 0, 0], [0, 0, 4]]) - The wrap option affects only tall matrices: + The `wrap` option affects only tall matrices: >>> # tall matrices no wrap >>> a = np.zeros((5, 3), dtype=int) @@ -641,7 +624,7 @@ def fill_diagonal(a, val, wrap=False): """ - dpnp.check_supported_arrays_type(a, scalar_type=True, all_scalars=True) + dpnp.check_supported_arrays_type(a) if a.ndim < 2: raise ValueError("array must be at least 2-d") @@ -651,13 +634,9 @@ def fill_diagonal(a, val, wrap=False): if not wrap and a.shape[0] > a.shape[1]: end = a.shape[1] * a.shape[1] else: - sh0 = a.shape[0] - for sh in a.shape[1:]: - if sh != sh0: - raise ValueError( - "All dimensions of input must be of equal length" - ) - step = sum(sh0**x for x in range(a.ndim)) + if not numpy.all(numpy.diff(a.shape) == 0): + raise ValueError("All dimensions of input must be of equal length") + step = sum(a.shape[0] ** x for x in range(a.ndim)) # TODO: implement flatiter for slice key # a.flat[:end:step] = val @@ -802,21 +781,22 @@ def mask_indices( Assume `mask_func` is a function that, for a square array a of size ``(n, n)`` with a possible offset argument `k`, when called as ``mask_func(a, k=k)`` returns a new array with zeros in certain locations - (functions like `triu` or `tril` do precisely this). Then this function - returns the indices where the non-zero values would be located. + (functions like :obj:`dpnp.triu` or :obj:`dpnp.tril` do precisely this). + Then this function returns the indices where the non-zero values would be + located. Parameters ---------- n : int The returned indices will be valid to access arrays of shape (n, n). mask_func : callable - A function whose call signature is similar to that of `triu`, `tril`. - That is, ``mask_func(x, k=k)`` returns a boolean array, shaped like `x`. - `k` is an optional argument to the function. + A function whose call signature is similar to that of :obj:`dpnp.triu`, + :obj:`dpnp.tril`. That is, ``mask_func(x, k=k)`` returns a boolean + array, shaped like `x`.`k` is an optional argument to the function. k : scalar An optional argument which is passed through to `mask_func`. Functions - like `triu`, `tril` take a second argument that is interpreted as an - offset. + like :obj:`dpnp.triu`, :obj:`dpnp.tril` take a second argument that is + interpreted as an offset. Default: ``0``. 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 @@ -835,6 +815,15 @@ def mask_indices( The `n` arrays of indices corresponding to the locations where ``mask_func(np.ones((n, n)), k)`` is True. + See Also + -------- + :obj:`dpnp.tril` : Return lower triangle of an array. + :obj:`dpnp.triu` : Return upper triangle of an array. + :obj:`dpnp.triu_indices` : Return the indices for the upper-triangle of an + (n, m) array. + :obj:`dpnp.tril_indices` : Return the indices for the lower-triangle of an + (n, m) array. + Examples -------- These are the indices that would allow you to access the upper triangular @@ -853,7 +842,7 @@ def mask_indices( >>> a[iu] array([0, 1, 2, 4, 5, 8]) - An offset can be passed also to the masking function. This gets us the + An offset can be passed also to the masking function. This gets us the indices starting on the first diagonal right of the main one: >>> iu1 = np.mask_indices(3, np.triu, 1) @@ -1384,11 +1373,11 @@ def tril_indices( The row dimension of the arrays for which the returned indices will be valid. k : int, optional - Diagonal offset (see `tril` for details). Default: ``0``. - m : int, optional + Diagonal offset (see :obj:`dpnp.tril` for details). Default: ``0``. + m : {None, int}, optional The column dimension of the arrays for which the returned arrays will be valid. - By default `m` is taken equal to `n`. + By default `m` is taken equal to `n`. Default: ``None``. 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 @@ -1407,6 +1396,14 @@ def tril_indices( The indices for the triangle. The returned tuple contains two arrays, each with the indices along one dimension of the array. + See Also + -------- + :obj:`dpnp.triu_indices` : similar function, for upper-triangular. + :obj:`dpnp.mask_indices` : generic function accepting an arbitrary mask + function. + :obj:`dpnp.tril` : Return lower triangle of an array. + :obj:`dpnp.triu` : Return upper triangle of an array. + Examples -------- Compute two different sets of indices to access 4x4 arrays, one for the @@ -1473,13 +1470,7 @@ def tril_indices( ) -def tril_indices_from( - arr, - k=0, - device=None, - usm_type=None, - sycl_queue=None, -): +def tril_indices_from(arr, k=0): """ Return the indices for the lower-triangle of arr. @@ -1491,18 +1482,7 @@ def tril_indices_from( The indices will be valid for square arrays whose dimensions are the same as arr. k : int, optional - Diagonal offset (see `tril` for details). Default: ``0``. - 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. + Diagonal offset (see :obj:`dpnp.tril` for details). Default: ``0``. Returns ------- @@ -1510,6 +1490,13 @@ def tril_indices_from( The indices for the triangle. The returned tuple contains two arrays, each with the indices along one dimension of the array. + See Also + -------- + :obj:`dpnp.tril_indices` : Return the indices for the lower-triangle of an + (n, m) array. + :obj:`dpnp.tril` : Return lower triangle of an array. + :obj:`dpnp.triu_indices_from` : similar function, for upper-triangular. + Examples -------- Create a 4 by 4 array. @@ -1552,17 +1539,12 @@ def tril_indices_from( if arr.ndim != 2: raise ValueError("input array must be 2-d") - _usm_type = arr.usm_type if usm_type is None else usm_type - _sycl_queue = dpnp.get_normalized_queue_device( - arr, sycl_queue=sycl_queue, device=device - ) - return tril_indices( arr.shape[-2], k=k, m=arr.shape[-1], - usm_type=_usm_type, - sycl_queue=_sycl_queue, + usm_type=arr.usm_type, + sycl_queue=arr.sycl_queue, ) @@ -1585,11 +1567,11 @@ def triu_indices( The size of the arrays for which the returned indices will be valid. k : int, optional - Diagonal offset (see `triu` for details). Default: ``0``. + Diagonal offset (see :obj:`dpnp.triu` for details). Default: ``0``. m : int, optional The column dimension of the arrays for which the returned arrays will be valid. - By default `m` is taken equal to `n`. + By default `m` is taken equal to `n`. Default: ``None``. 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 @@ -1606,9 +1588,17 @@ def triu_indices( ------- inds : tuple of dpnp.ndarray The indices for the triangle. The returned tuple contains two arrays, - each with the indices along one dimension of the array. Can be used + each with the indices along one dimension of the array. Can be used to slice a ndarray of shape(`n`, `n`). + See Also + -------- + :obj:`dpnp.tril_indices` : similar function, for lower-triangular. + :obj:`dpnp.mask_indices` : generic function accepting an arbitrary mask + function. + :obj:`dpnp.tril` : Return lower triangle of an array. + :obj:`dpnp.triu` : Return upper triangle of an array. + Examples -------- Compute two different sets of indices to access 4x4 arrays, one for the @@ -1676,13 +1666,7 @@ def triu_indices( ) -def triu_indices_from( - arr, - k=0, - device=None, - usm_type=None, - sycl_queue=None, -): +def triu_indices_from(arr, k=0): """ Return the indices for the lower-triangle of arr. @@ -1694,26 +1678,22 @@ def triu_indices_from( The indices will be valid for square arrays whose dimensions are the same as arr. k : int, optional - Diagonal offset (see `triu` for details). Default: ``0``. - 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. + Diagonal offset (see :obj:`dpnp.triu` for details). Default: ``0``. Returns ------- inds : tuple of dpnp.ndarray The indices for the triangle. The returned tuple contains two arrays, - each with the indices along one dimension of the array. Can be used + each with the indices along one dimension of the array. Can be used to slice a ndarray of shape(`n`, `n`). + See Also + -------- + :obj:`dpnp.triu_indices` : Return the indices for the upper-triangle of an + (n, m) array. + :obj:`dpnp.triu` : Return upper triangle of an array. + :obj:`dpnp.tril_indices_from` : similar function, for lower-triangular. + Examples -------- Create a 4 by 4 array. @@ -1756,15 +1736,10 @@ def triu_indices_from( if arr.ndim != 2: raise ValueError("input array must be 2-d") - _usm_type = arr.usm_type if usm_type is None else usm_type - _sycl_queue = dpnp.get_normalized_queue_device( - arr, sycl_queue=sycl_queue, device=device - ) - return triu_indices( arr.shape[-2], k=k, m=arr.shape[-1], - usm_type=_usm_type, - sycl_queue=_sycl_queue, + usm_type=arr.usm_type, + sycl_queue=arr.sycl_queue, ) diff --git a/tests/test_sycl_queue.py b/tests/test_sycl_queue.py index fdbbbdb9141d..d562ad024deb 100644 --- a/tests/test_sycl_queue.py +++ b/tests/test_sycl_queue.py @@ -2144,29 +2144,13 @@ def test_histogram(weights, device): valid_devices, ids=[device.filter_string for device in valid_devices], ) -def test_tri_diag_indices_from_follow(func, device): +def test_tri_diag_indices_from(func, device): arr = dpnp.ones((3, 3), device=device) res = getattr(dpnp, func)(arr) for x in res: assert_sycl_queue_equal(x.sycl_queue, arr.sycl_queue) -@pytest.mark.parametrize( - "func", ["tril_indices_from", "triu_indices_from", "diag_indices_from"] -) -@pytest.mark.parametrize( - "device", - valid_devices, - ids=[device.filter_string for device in valid_devices], -) -def test_tri_diag_indices_from(func, device): - sycl_queue = dpctl.SyclQueue(device) - arr = dpnp.ones((3, 3)) - res = getattr(dpnp, func)(arr, sycl_queue=sycl_queue) - for x in res: - assert_sycl_queue_equal(x.sycl_queue, sycl_queue) - - @pytest.mark.parametrize( "func", ["tril_indices", "triu_indices", "diag_indices"] ) diff --git a/tests/test_usm_type.py b/tests/test_usm_type.py index 80360d8f8343..2ba5d911f15e 100644 --- a/tests/test_usm_type.py +++ b/tests/test_usm_type.py @@ -1247,24 +1247,13 @@ def test_histogram(usm_type_v, usm_type_w): "func", ["tril_indices_from", "triu_indices_from", "diag_indices_from"] ) @pytest.mark.parametrize("usm_type", list_of_usm_types, ids=list_of_usm_types) -def test_tri_diag_indices_from_follow(func, usm_type): +def test_tri_diag_indices_from(func, usm_type): arr = dp.ones((3, 3), usm_type=usm_type) res = getattr(dp, func)(arr) for x in res: assert x.usm_type == usm_type -@pytest.mark.parametrize( - "func", ["tril_indices_from", "triu_indices_from", "diag_indices_from"] -) -@pytest.mark.parametrize("usm_type", list_of_usm_types, ids=list_of_usm_types) -def test_tri_diag_indices_from(func, usm_type): - arr = dp.ones((3, 3)) - res = getattr(dp, func)(arr, usm_type=usm_type) - for x in res: - assert x.usm_type == usm_type - - @pytest.mark.parametrize( "func", ["tril_indices", "triu_indices", "diag_indices"] ) From a457656885d4b6a9027ff8d8e1eaf7c09896a068 Mon Sep 17 00:00:00 2001 From: Natalia Polina Date: Tue, 21 May 2024 13:11:15 -0500 Subject: [PATCH 5/6] address comments --- dpnp/dpnp_iface_indexing.py | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/dpnp/dpnp_iface_indexing.py b/dpnp/dpnp_iface_indexing.py index d3b3225612c1..e3be93849044 100644 --- a/dpnp/dpnp_iface_indexing.py +++ b/dpnp/dpnp_iface_indexing.py @@ -625,6 +625,7 @@ def fill_diagonal(a, val, wrap=False): """ dpnp.check_supported_arrays_type(a) + dpnp.check_supported_arrays_type(val, scalar_type=True, all_scalars=True) if a.ndim < 2: raise ValueError("array must be at least 2-d") @@ -638,17 +639,16 @@ def fill_diagonal(a, val, wrap=False): raise ValueError("All dimensions of input must be of equal length") step = sum(a.shape[0] ** x for x in range(a.ndim)) - # TODO: implement flatiter for slice key - # a.flat[:end:step] = val - fl = a.flat + a_sh = a.shape + tmp_a = dpnp.ravel(a) if dpnp.isscalar(val): - for idx in range(0, end, step): - fl[idx] = val + tmp_a[:end:step] = val else: - flat_val = val.flatten() + flat_val = val.ravel() for i in range(0, flat_val.size): - for idx in range((step * i), end, step * (1 + i)): - fl[idx] = flat_val[i] + tmp_a[step * i : end : step * (i + 1)] = flat_val[i] + tmp_a = dpnp.reshape(tmp_a, a_sh) + a[:] = tmp_a def indices( From e77dc79cfe6dd72e1290bcae0744ffd1d2468b0c Mon Sep 17 00:00:00 2001 From: Natalia Polina Date: Tue, 21 May 2024 15:00:24 -0500 Subject: [PATCH 6/6] Added commit --- dpnp/dpnp_iface_indexing.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/dpnp/dpnp_iface_indexing.py b/dpnp/dpnp_iface_indexing.py index 737c704c32ec..9aaa15a9ce8f 100644 --- a/dpnp/dpnp_iface_indexing.py +++ b/dpnp/dpnp_iface_indexing.py @@ -647,12 +647,16 @@ def fill_diagonal(a, val, wrap=False): raise ValueError("All dimensions of input must be of equal length") step = sum(a.shape[0] ** x for x in range(a.ndim)) + # TODO: implement flatiter for slice key + # a.flat[:end:step] = val a_sh = a.shape tmp_a = dpnp.ravel(a) if dpnp.isscalar(val): tmp_a[:end:step] = val else: flat_val = val.ravel() + # Setitem can work only if index size equal val size. + # Using loop for general case without dependencies of val size. for i in range(0, flat_val.size): tmp_a[step * i : end : step * (i + 1)] = flat_val[i] tmp_a = dpnp.reshape(tmp_a, a_sh)