diff --git a/dpnp/backend/include/dpnp_iface_fptr.hpp b/dpnp/backend/include/dpnp_iface_fptr.hpp index c69682b43e85..d8e6f8b26e81 100644 --- a/dpnp/backend/include/dpnp_iface_fptr.hpp +++ b/dpnp/backend/include/dpnp_iface_fptr.hpp @@ -112,11 +112,9 @@ 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_DIVIDE, /**< Used in numpy.divide() impl */ - DPNP_FN_DOT, /**< Used in numpy.dot() impl */ + DPNP_FN_DIAGONAL, /**< Used in numpy.diagonal() impl */ + 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 @@ -140,12 +138,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 b07f5e56e77b..dcbf6ca906c7 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, @@ -188,9 +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]; } } @@ -225,8 +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) { @@ -364,15 +358,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, @@ -897,15 +882,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] = { @@ -924,15 +900,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 73a973fad0e3..a82a96ed0c59 100644 --- a/dpnp/dpnp_algo/dpnp_algo.pxd +++ b/dpnp/dpnp_algo/dpnp_algo.pxd @@ -38,13 +38,11 @@ cdef extern from "dpnp_iface_fptr.hpp" namespace "DPNPFuncName": # need this na DPNP_FN_COPY_EXT DPNP_FN_CORRELATE_EXT DPNP_FN_DEGREES_EXT - DPNP_FN_DIAG_INDICES_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 703f65b504ab..d96ebcd816f2 100644 --- a/dpnp/dpnp_algo/dpnp_algo_indexing.pxi +++ b/dpnp/dpnp_algo/dpnp_algo_indexing.pxi @@ -37,26 +37,13 @@ and the rest of the library __all__ += [ "dpnp_choose", - "dpnp_diag_indices", - "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_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): cdef vector[void * ] choices @@ -105,71 +92,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 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 @@ -197,91 +119,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 6be415401254..0a1c8529c425 100644 --- a/dpnp/dpnp_iface_indexing.py +++ b/dpnp/dpnp_iface_indexing.py @@ -46,14 +46,8 @@ # pylint: disable=no-name-in-module from .dpnp_algo import ( dpnp_choose, - dpnp_diag_indices, - 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 ( @@ -69,6 +63,7 @@ "extract", "fill_diagonal", "indices", + "mask_indices", "nonzero", "place", "put", @@ -181,7 +176,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. @@ -193,6 +188,30 @@ 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. 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. + + Returns + ------- + out : tuple of dpnp.ndarray + The indices to access the main diagonal of an array. + See also -------- :obj:`diag_indices_from` : Return the indices to access the main @@ -237,42 +256,78 @@ 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, + device=device, + usm_type=usm_type, + sycl_queue=sycl_queue, + ) + 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 : {dpnp.ndarray, usm_ndarray} + Array at least 2-D + + Returns + ------- + out : tuple of dpnp.ndarray + The indices to access the main diagonal of an n-dimensional array. + 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 + dpnp.check_supported_arrays_type(arr) - # 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) + if not arr.ndim >= 2: + raise ValueError("input array must be at least 2-d") - return call_origin(numpy.diag_indices_from, x1) + if not numpy.all(numpy.diff(arr.shape) == 0): + raise ValueError("All dimensions of input must be of equal length") + + return diag_indices( + arr.shape[0], + arr.ndim, + usm_type=arr.usm_type, + sycl_queue=arr.sycl_queue, + ) def diagonal(a, offset=0, axis1=0, axis2=1): @@ -473,15 +528,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 : {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, the flattened `val` is + written along the diagonal, repeating if necessary to fill all + diagonal entries. + wrap : bool + It enables the diagonal "wrapped" after N columns. This affects only + tall matrices. Default: ``False``. See Also -------- @@ -490,20 +554,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]]) - 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) + 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]]) + + """ - return call_origin(numpy.fill_diagonal, x1, val, wrap, dpnp_inplace=True) + 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") + 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: + 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 + 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) + a[:] = tmp_a def indices( @@ -622,6 +779,104 @@ def indices( return res +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=k)`` returns a new array with zeros in certain locations + (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 :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 :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 + 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 dpnp.ndarray + 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 + 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, + device=device, + usm_type=usm_type, + sycl_queue=sycl_queue, + ) + a = mask_func(m, k=k) + return nonzero(a != 0) + + def nonzero(a): """ Return the indices of the elements that are non-zero. @@ -1118,123 +1373,392 @@ 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. + For full documentation refer to :obj:`numpy.tril_indices`. + Parameters ---------- 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). - - 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 + 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. + 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 + 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, + device=device, + usm_type=usm_type, + sycl_queue=sycl_queue, + ) - 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, + device=device, + usm_type=usm_type, + sycl_queue=sycl_queue, + ) + ) -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 ---------- - 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 :obj:`dpnp.tril` 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. + + 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. + + >>> 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]) + """ - 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) + dpnp.check_supported_arrays_type(arr) - 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], + usm_type=arr.usm_type, + sycl_queue=arr.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. + For full documentation refer to :obj:`numpy.triu_indices`. + Parameters ---------- 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 :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 + 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`). + + 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 + 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, + device=device, + usm_type=usm_type, + sycl_queue=sycl_queue, + ) - 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, + device=device, + usm_type=usm_type, + sycl_queue=sycl_queue, + ) + ) -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 ---------- - 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 :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 + 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. + + >>> 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]) + """ - 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) + dpnp.check_supported_arrays_type(arr) - 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], + usm_type=arr.usm_type, + sycl_queue=arr.sycl_queue, + ) diff --git a/tests/skipped_tests.tbl b/tests/skipped_tests.tbl index f05a82c20492..a9cb3d095604 100644 --- a/tests/skipped_tests.tbl +++ b/tests/skipped_tests.tbl @@ -161,39 +161,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 18c367822010..fa8d00145d17 100644 --- a/tests/skipped_tests_gpu.tbl +++ b/tests/skipped_tests_gpu.tbl @@ -211,56 +211,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_non_broadcastable tests/third_party/cupy/indexing_tests/test_indexing.py::TestSelect::test_select_type_error_condlist -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/test_indexing.py b/tests/test_indexing.py index 89ea4ce39b3f..f001f994dbd9 100644 --- a/tests/test_indexing.py +++ b/tests/test_indexing.py @@ -928,7 +928,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"] @@ -940,7 +940,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", @@ -963,7 +963,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"] @@ -975,7 +975,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", @@ -992,3 +992,21 @@ 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))) + + +def test_fill_diagonal_error(): + arr = dpnp.ones((1, 2, 3)) + with pytest.raises(ValueError): + dpnp.fill_diagonal(arr, 5) diff --git a/tests/test_sycl_queue.py b/tests/test_sycl_queue.py index b0472c56714b..9286131a65b5 100644 --- a/tests/test_sycl_queue.py +++ b/tests/test_sycl_queue.py @@ -2139,6 +2139,49 @@ def test_histogram(weights, device): 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(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", "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) + + @pytest.mark.parametrize("weights", [None, numpy.arange(7, 12)]) @pytest.mark.parametrize( "device", diff --git a/tests/test_usm_type.py b/tests/test_usm_type.py index 530072f35b8b..a2b38b82e8d0 100644 --- a/tests/test_usm_type.py +++ b/tests/test_usm_type.py @@ -1246,6 +1246,35 @@ def test_histogram(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(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", "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 + + @pytest.mark.parametrize("usm_type_v", list_of_usm_types, ids=list_of_usm_types) @pytest.mark.parametrize("usm_type_w", list_of_usm_types, ids=list_of_usm_types) def test_histogram_bin_edges(usm_type_v, usm_type_w): 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) diff --git a/tests/third_party/cupy/indexing_tests/test_insert.py b/tests/third_party/cupy/indexing_tests/test_insert.py index a79956e2f141..538039bd43c6 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):