diff --git a/dpnp/backend/include/dpnp_iface_fptr.hpp b/dpnp/backend/include/dpnp_iface_fptr.hpp index c56f38ffcb5e..6a174b3b647e 100644 --- a/dpnp/backend/include/dpnp_iface_fptr.hpp +++ b/dpnp/backend/include/dpnp_iface_fptr.hpp @@ -231,21 +231,19 @@ enum class DPNPFuncName : size_t DPNP_FN_PTP, /**< Used in numpy.ptp() impl */ DPNP_FN_PUT, /**< Used in numpy.put() impl */ DPNP_FN_PUT_ALONG_AXIS, /**< Used in numpy.put_along_axis() impl */ - DPNP_FN_PUT_ALONG_AXIS_EXT, /**< Used in numpy.put_along_axis() impl, - requires extra parameters */ - DPNP_FN_QR, /**< Used in numpy.linalg.qr() impl */ - DPNP_FN_QR_EXT, /**< Used in numpy.linalg.qr() impl, requires extra - parameters */ - DPNP_FN_RADIANS, /**< Used in numpy.radians() impl */ - DPNP_FN_RADIANS_EXT, /**< Used in numpy.radians() impl, requires extra - parameters */ - DPNP_FN_REMAINDER, /**< Used in numpy.remainder() impl */ - DPNP_FN_RECIP, /**< Used in numpy.recip() impl */ - DPNP_FN_RECIP_EXT, /**< Used in numpy.recip() impl, requires extra - parameters */ - DPNP_FN_REPEAT, /**< Used in numpy.repeat() impl */ - DPNP_FN_RIGHT_SHIFT, /**< Used in numpy.right_shift() impl */ - DPNP_FN_RNG_BETA, /**< Used in numpy.random.beta() impl */ + DPNP_FN_QR, /**< Used in numpy.linalg.qr() impl */ + DPNP_FN_QR_EXT, /**< Used in numpy.linalg.qr() impl, requires extra + parameters */ + DPNP_FN_RADIANS, /**< Used in numpy.radians() impl */ + DPNP_FN_RADIANS_EXT, /**< Used in numpy.radians() impl, requires extra + parameters */ + DPNP_FN_REMAINDER, /**< Used in numpy.remainder() impl */ + DPNP_FN_RECIP, /**< Used in numpy.recip() impl */ + DPNP_FN_RECIP_EXT, /**< Used in numpy.recip() impl, requires extra + parameters */ + DPNP_FN_REPEAT, /**< Used in numpy.repeat() impl */ + DPNP_FN_RIGHT_SHIFT, /**< Used in numpy.right_shift() impl */ + DPNP_FN_RNG_BETA, /**< Used in numpy.random.beta() impl */ DPNP_FN_RNG_BETA_EXT, /**< Used in numpy.random.beta() impl, requires extra parameters */ DPNP_FN_RNG_BINOMIAL, /**< Used in numpy.random.binomial() impl */ diff --git a/dpnp/backend/kernels/dpnp_krnl_indexing.cpp b/dpnp/backend/kernels/dpnp_krnl_indexing.cpp index 2cca84f9e61f..e9addf36b707 100644 --- a/dpnp/backend/kernels/dpnp_krnl_indexing.cpp +++ b/dpnp/backend/kernels/dpnp_krnl_indexing.cpp @@ -796,19 +796,6 @@ void (*dpnp_put_along_axis_default_c)(void *, size_t) = dpnp_put_along_axis_c<_DataType>; -template -DPCTLSyclEventRef (*dpnp_put_along_axis_ext_c)(DPCTLSyclQueueRef, - void *, - long *, - void *, - size_t, - const shape_elem_type *, - size_t, - size_t, - size_t, - const DPCTLEventVectorRef) = - dpnp_put_along_axis_c<_DataType>; - template class dpnp_take_c_kernel; @@ -1005,15 +992,6 @@ void func_map_init_indexing_func(func_map_t &fmap) fmap[DPNPFuncName::DPNP_FN_PUT_ALONG_AXIS][eft_DBL][eft_DBL] = { eft_DBL, (void *)dpnp_put_along_axis_default_c}; - fmap[DPNPFuncName::DPNP_FN_PUT_ALONG_AXIS_EXT][eft_INT][eft_INT] = { - eft_INT, (void *)dpnp_put_along_axis_ext_c}; - fmap[DPNPFuncName::DPNP_FN_PUT_ALONG_AXIS_EXT][eft_LNG][eft_LNG] = { - eft_LNG, (void *)dpnp_put_along_axis_ext_c}; - fmap[DPNPFuncName::DPNP_FN_PUT_ALONG_AXIS_EXT][eft_FLT][eft_FLT] = { - eft_FLT, (void *)dpnp_put_along_axis_ext_c}; - fmap[DPNPFuncName::DPNP_FN_PUT_ALONG_AXIS_EXT][eft_DBL][eft_DBL] = { - eft_DBL, (void *)dpnp_put_along_axis_ext_c}; - fmap[DPNPFuncName::DPNP_FN_TAKE][eft_BLN][eft_INT] = { eft_BLN, (void *)dpnp_take_default_c}; fmap[DPNPFuncName::DPNP_FN_TAKE][eft_INT][eft_INT] = { diff --git a/dpnp/dpnp_algo/dpnp_algo.pxd b/dpnp/dpnp_algo/dpnp_algo.pxd index 80c6035d7a9f..d49adcf0b7fc 100644 --- a/dpnp/dpnp_algo/dpnp_algo.pxd +++ b/dpnp/dpnp_algo/dpnp_algo.pxd @@ -156,8 +156,6 @@ cdef extern from "dpnp_iface_fptr.hpp" namespace "DPNPFuncName": # need this na DPNP_FN_RNG_POISSON_EXT DPNP_FN_RNG_POWER DPNP_FN_RNG_POWER_EXT - DPNP_FN_PUT_ALONG_AXIS - DPNP_FN_PUT_ALONG_AXIS_EXT DPNP_FN_RNG_RAYLEIGH DPNP_FN_RNG_RAYLEIGH_EXT DPNP_FN_RNG_SHUFFLE diff --git a/dpnp/dpnp_algo/dpnp_algo_indexing.pxi b/dpnp/dpnp_algo/dpnp_algo_indexing.pxi index 36fc7ff8eb91..25cebe84d18b 100644 --- a/dpnp/dpnp_algo/dpnp_algo_indexing.pxi +++ b/dpnp/dpnp_algo/dpnp_algo_indexing.pxi @@ -41,10 +41,8 @@ __all__ += [ "dpnp_diagonal", "dpnp_fill_diagonal", "dpnp_indices", - "dpnp_put_along_axis", "dpnp_putmask", "dpnp_select", - "dpnp_take_along_axis", "dpnp_tril_indices", "dpnp_tril_indices_from", "dpnp_triu_indices", @@ -69,16 +67,6 @@ ctypedef c_dpctl.DPCTLSyclEventRef(*custom_indexing_2in_1out_func_ptr_t_)(c_dpct 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) -ctypedef c_dpctl.DPCTLSyclEventRef(*custom_indexing_3in_with_axis_func_ptr_t)(c_dpctl.DPCTLSyclQueueRef, - void * , - void * , - void * , - const size_t, - shape_elem_type * , - const size_t, - const size_t, - const size_t, - const c_dpctl.DPCTLEventVectorRef) cpdef utils.dpnp_descriptor dpnp_choose(utils.dpnp_descriptor x1, list choices1): @@ -283,35 +271,6 @@ cpdef object dpnp_indices(dimensions): return dpnp_result -cpdef dpnp_put_along_axis(dpnp_descriptor arr, dpnp_descriptor indices, dpnp_descriptor values, int axis): - cdef shape_type_c arr_shape = arr.shape - cdef DPNPFuncType param1_type = dpnp_dtype_to_DPNPFuncType(arr.dtype) - - cdef DPNPFuncData kernel_data = get_dpnp_function_ptr(DPNP_FN_PUT_ALONG_AXIS_EXT, param1_type, param1_type) - - utils.get_common_usm_allocation(arr, indices) # check USM allocation is common - _, _, result_sycl_queue = utils.get_common_usm_allocation(arr, values) - - cdef c_dpctl.SyclQueue q = result_sycl_queue - cdef c_dpctl.DPCTLSyclQueueRef q_ref = q.get_queue_ref() - - cdef custom_indexing_3in_with_axis_func_ptr_t func = kernel_data.ptr - - cdef c_dpctl.DPCTLSyclEventRef event_ref = func(q_ref, - arr.get_data(), - indices.get_data(), - values.get_data(), - axis, - arr_shape.data(), - arr.ndim, - indices.size, - values.size, - 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 @@ -341,94 +300,6 @@ cpdef utils.dpnp_descriptor dpnp_select(list condlist, list choicelist, default) return res_array -cpdef object dpnp_take_along_axis(object arr, object indices, int axis): - cdef long size_arr = arr.size - cdef shape_type_c shape_arr = arr.shape - cdef shape_type_c output_shape - cdef long size_indices = indices.size - res_type = arr.dtype - - if axis != arr.ndim - 1: - res_shape_list = list(shape_arr) - res_shape_list[axis] = 1 - res_shape = tuple(res_shape_list) - - output_shape = (0,) * (len(shape_arr) - 1) - ind = 0 - for id, shape_axis in enumerate(shape_arr): - if id != axis: - output_shape[ind] = shape_axis - ind += 1 - - prod = 1 - for i in range(len(output_shape)): - if output_shape[i] != 0: - prod *= output_shape[i] - - result_array = dpnp.empty((prod, ), dtype=res_type) - ind_array = [None] * prod - arr_shape_offsets = [None] * len(shape_arr) - acc = 1 - - for i in range(len(shape_arr)): - ind = len(shape_arr) - 1 - i - arr_shape_offsets[ind] = acc - acc *= shape_arr[ind] - - output_shape_offsets = [None] * len(shape_arr) - acc = 1 - - for i in range(len(output_shape)): - ind = len(output_shape) - 1 - i - output_shape_offsets[ind] = acc - acc *= output_shape[ind] - result_offsets = arr_shape_offsets[:] # need copy. not a reference - result_offsets[axis] = 0 - - for source_idx in range(size_arr): - - # reconstruct x,y,z from linear source_idx - xyz = [] - remainder = source_idx - for i in arr_shape_offsets: - quotient, remainder = divmod(remainder, i) - xyz.append(quotient) - - # extract result axis - result_axis = [] - for idx, offset in enumerate(xyz): - if idx != axis: - result_axis.append(offset) - - # Construct result offset - result_offset = 0 - for i, result_axis_val in enumerate(result_axis): - result_offset += (output_shape_offsets[i] * result_axis_val) - - arr_elem = arr.item(source_idx) - if ind_array[result_offset] is None: - ind_array[result_offset] = 0 - else: - ind_array[result_offset] += 1 - - if ind_array[result_offset] % size_indices == indices.item(result_offset % size_indices): - result_array[result_offset] = arr_elem - - dpnp_result_array = dpnp.reshape(result_array, res_shape) - return dpnp_result_array - - else: - result_array = utils_py.create_output_descriptor_py(shape_arr, res_type, None).get_pyobj() - - result_array_flatiter = result_array.flat - - for i in range(size_arr): - ind = size_indices * (i // size_indices) + indices.item(i % size_indices) - result_array_flatiter[i] = arr.item(ind) - - return result_array - - cpdef tuple dpnp_tril_indices(n, k=0, m=None): array1 = [] array2 = [] diff --git a/dpnp/dpnp_iface.py b/dpnp/dpnp_iface.py index e91a9b991f89..247264a79c56 100644 --- a/dpnp/dpnp_iface.py +++ b/dpnp/dpnp_iface.py @@ -58,6 +58,7 @@ "array_equal", "asnumpy", "astype", + "check_supported_arrays_type", "convert_single_elem_array_to_scalar", "default_float_type", "dpnp_queue_initialize", @@ -203,6 +204,42 @@ def astype(x1, dtype, order="K", casting="unsafe", copy=True): return dpnp_array._create_from_usm_ndarray(array_obj) +def check_supported_arrays_type(*arrays, scalar_type=False): + """ + Return ``True`` if each array has either type of scalar, + :class:`dpnp.ndarray` or :class:`dpctl.tensor.usm_ndarray`. + But if any array has unsupported type, ``TypeError`` will be raised. + + Parameters + ---------- + arrays : {dpnp_array, usm_ndarray} + Input arrays to check for supported types. + scalar_type : {bool}, optional + A scalar type is also considered as supported if flag is True. + + Returns + ------- + out : bool + ``True`` if each type of input `arrays` is supported type, + ``False`` otherwise. + + Raises + ------ + TypeError + If any input array from `arrays` is of unsupported array type. + + """ + + for a in arrays: + if scalar_type and dpnp.isscalar(a) or is_supported_array_type(a): + continue + + raise TypeError( + "An array must be any of supported type, but got {}".format(type(a)) + ) + return True + + def convert_single_elem_array_to_scalar(obj, keepdims=False): """Convert array with single element to scalar.""" diff --git a/dpnp/dpnp_iface_indexing.py b/dpnp/dpnp_iface_indexing.py index bc04a47efafd..6a61f728e7d2 100644 --- a/dpnp/dpnp_iface_indexing.py +++ b/dpnp/dpnp_iface_indexing.py @@ -41,6 +41,7 @@ import dpctl.tensor as dpt import numpy +from numpy.core.numeric import normalize_axis_index import dpnp from dpnp.dpnp_algo import * @@ -70,6 +71,52 @@ ] +def _build_along_axis_index(a, indices, axis): + """ + Build a fancy index used by a family of `_along_axis` functions. + + The fancy index consists of orthogonal arranges, with the + requested index inserted at the right location. + + The resulting index is going to be used inside `dpnp.put_along_axis` + and `dpnp.take_along_axis` implementations. + + """ + + if not dpnp.issubdtype(indices.dtype, dpnp.integer): + raise IndexError("`indices` must be an integer array") + + # normalize array shape and input axis + if axis is None: + a_shape = (a.size,) + axis = 0 + else: + a_shape = a.shape + axis = normalize_axis_index(axis, a.ndim) + + if len(a_shape) != indices.ndim: + raise ValueError( + "`indices` and `a` must have the same number of dimensions" + ) + + # compute dimensions to iterate over + dest_dims = list(range(axis)) + [None] + list(range(axis + 1, indices.ndim)) + shape_ones = (1,) * indices.ndim + + # build the index + fancy_index = [] + for dim, n in zip(dest_dims, a_shape): + if dim is None: + fancy_index.append(indices) + else: + ind_shape = shape_ones[:dim] + (-1,) + shape_ones[dim + 1 :] + fancy_index.append( + dpnp.arange(n, dtype=indices.dtype).reshape(ind_shape) + ) + + return tuple(fancy_index) + + def choose(x1, choices, out=None, mode="raise"): """ Construct an array from an index array and a set of arrays to choose from. @@ -78,7 +125,8 @@ def choose(x1, choices, out=None, mode="raise"): See also -------- - :obj:`take_along_axis` : Preferable if choices is an array. + :obj:`dpnp.take_along_axis` : Preferable if choices is an array. + """ x1_desc = dpnp.get_dpnp_descriptor(x1, copy_when_nondefault_queue=False) @@ -247,7 +295,7 @@ def extract(condition, x): Returns ------- - y : dpnp.ndarray + out : dpnp.ndarray Rank 1 array of values from `x` where `condition` is True. Limitations @@ -342,7 +390,7 @@ def nonzero(x, /): Returns ------- - y : tuple[dpnp.ndarray] + out : tuple[dpnp.ndarray] Indices of elements that are non-zero. Limitations @@ -496,39 +544,55 @@ def put(a, indices, vals, /, *, axis=None, mode="wrap"): return call_origin(numpy.put, a, indices, vals, mode, dpnp_inplace=True) -def put_along_axis(x1, indices, values, axis): +def put_along_axis(a, indices, values, axis): """ Put values into the destination array by matching 1d index and data slices. For full documentation refer to :obj:`numpy.put_along_axis`. + Limitations + ----------- + Parameters `a` and `indices` are supported either as :class:`dpnp.ndarray` + or :class:`dpctl.tensor.usm_ndarray`. + Parameter `values` is supported either as scalar, :class:`dpnp.ndarray` + or :class:`dpctl.tensor.usm_ndarray`. + Otherwise ``TypeError`` exception will be raised. + See Also -------- - :obj:`take_along_axis` : Take values from the input array by matching 1d index and data slices. + :obj:`dpnp.put` : Put values along an axis, using the same indices for every 1d slice. + :obj:`dpnp.take_along_axis` : Take values from the input array by matching 1d index and data slices. + + Examples + -------- + For this sample array + + >>> import dpnp as np + >>> a = np.array([[10, 30, 20], [60, 40, 50]]) + + We can replace the maximum values with: + + >>> ai = np.argmax(a, axis=1, keepdims=True) + >>> ai + array([[1], + [0]]) + >>> np.put_along_axis(a, ai, 99, axis=1) + >>> a + array([[10, 99, 20], + [99, 40, 50]]) + """ - x1_desc = dpnp.get_dpnp_descriptor(x1, copy_when_nondefault_queue=False) - indices_desc = dpnp.get_dpnp_descriptor( - indices, copy_when_nondefault_queue=False - ) - values_desc = dpnp.get_dpnp_descriptor( - values, copy_when_nondefault_queue=False - ) - if x1_desc and indices_desc and values_desc: - if x1_desc.ndim != indices_desc.ndim: - pass - elif not isinstance(axis, int): - pass - elif axis >= x1_desc.ndim: - pass - elif indices_desc.size != values_desc.size: - pass - else: - return dpnp_put_along_axis(x1_desc, indices_desc, values_desc, axis) + dpnp.check_supported_arrays_type(a, indices) - return call_origin( - numpy.put_along_axis, x1, indices, values, axis, dpnp_inplace=True - ) + # TODO: remove when #1382(dpctl) is resolved + if dpnp.is_supported_array_type(values) and a.dtype != values.dtype: + values = values.astype(a.dtype) + + if axis is None: + a = a.ravel() + + a[_build_along_axis_index(a, indices, axis)] = values def putmask(x1, mask, values): @@ -596,7 +660,7 @@ def take(x, indices, /, *, axis=None, out=None, mode="wrap"): Returns ------- - dpnp.ndarray + out : dpnp.ndarray An array with shape x.shape[:axis] + indices.shape + x.shape[axis + 1:] filled with elements from `x`. @@ -613,7 +677,7 @@ def take(x, indices, /, *, axis=None, out=None, mode="wrap"): See Also -------- :obj:`dpnp.compress` : Take elements using a boolean mask. - :obj:`take_along_axis` : Take elements by matching the array and the index arrays. + :obj:`dpnp.take_along_axis` : Take elements by matching the array and the index arrays. Notes ----- @@ -666,44 +730,83 @@ def take(x, indices, /, *, axis=None, out=None, mode="wrap"): return call_origin(numpy.take, x, indices, axis, out, mode) -def take_along_axis(x1, indices, axis): +def take_along_axis(a, indices, axis): """ Take values from the input array by matching 1d index and data slices. For full documentation refer to :obj:`numpy.take_along_axis`. + Returns + ------- + out : dpnp.ndarray + The indexed result. + + Limitations + ----------- + Parameters `a` and `indices` are supported either as :class:`dpnp.ndarray` + or :class:`dpctl.tensor.usm_ndarray`. + Otherwise ``TypeError`` exception will be raised. + See Also -------- :obj:`dpnp.take` : Take along an axis, using the same indices for every 1d slice. - :obj:`put_along_axis` : Put values into the destination array by matching 1d index and data slices. - """ + :obj:`dpnp.put_along_axis` : Put values into the destination array by matching 1d index and data slices. - x1_desc = dpnp.get_dpnp_descriptor(x1, copy_when_nondefault_queue=False) - indices_desc = dpnp.get_dpnp_descriptor( - indices, copy_when_nondefault_queue=False - ) - if x1_desc and indices_desc: - if x1_desc.ndim != indices_desc.ndim: - pass - elif not isinstance(axis, int): - pass - elif axis >= x1_desc.ndim: - pass - elif x1_desc.ndim == indices_desc.ndim: - val_list = [] - for i in list(indices_desc.shape)[:-1]: - if i == 1: - val_list.append(True) - else: - val_list.append(False) - if not all(val_list): - pass - else: - return dpnp_take_along_axis(x1, indices, axis) - else: - return dpnp_take_along_axis(x1, indices, axis) + Examples + -------- + For this sample array - return call_origin(numpy.take_along_axis, x1, indices, axis) + >>> import dpnp as np + >>> a = np.array([[10, 30, 20], [60, 40, 50]]) + + We can sort either by using sort directly, or argsort and this function + + >>> np.sort(a, axis=1) + array([[10, 20, 30], + [40, 50, 60]]) + >>> ai = np.argsort(a, axis=1) + >>> ai + array([[0, 2, 1], + [1, 2, 0]]) + >>> np.take_along_axis(a, ai, axis=1) + array([[10, 20, 30], + [40, 50, 60]]) + + The same works for max and min, if you maintain the trivial dimension + with ``keepdims``: + + >>> np.max(a, axis=1, keepdims=True) + array([[30], + [60]]) + >>> ai = np.argmax(a, axis=1, keepdims=True) + >>> ai + array([[1], + [0]]) + >>> np.take_along_axis(a, ai, axis=1) + array([[30], + [60]]) + + If we want to get the max and min at the same time, we can stack the + indices first + + >>> ai_min = np.argmin(a, axis=1, keepdims=True) + >>> ai_max = np.argmax(a, axis=1, keepdims=True) + >>> ai = np.concatenate([ai_min, ai_max], axis=1) + >>> ai + array([[0, 1], + [1, 0]]) + >>> np.take_along_axis(a, ai, axis=1) + array([[10, 30], + [40, 60]]) + + """ + + dpnp.check_supported_arrays_type(a, indices) + + if axis is None: + a = a.ravel() + + return a[_build_along_axis_index(a, indices, axis)] def tril_indices(n, k=0, m=None): diff --git a/dpnp/dpnp_iface_manipulation.py b/dpnp/dpnp_iface_manipulation.py index b805a3a906e9..7ee53bca3775 100644 --- a/dpnp/dpnp_iface_manipulation.py +++ b/dpnp/dpnp_iface_manipulation.py @@ -799,10 +799,7 @@ def fliplr(m): """ - if not dpnp.is_supported_array_type(m): - raise TypeError( - "An array must be any of supported type, but got {}".format(type(m)) - ) + dpnp.check_supported_arrays_type(m) if m.ndim < 2: raise ValueError(f"Input must be >= 2-d, but got {m.ndim}") @@ -857,10 +854,7 @@ def flipud(m): """ - if not dpnp.is_supported_array_type(m): - raise TypeError( - "An array must be any of supported type, but got {}".format(type(m)) - ) + dpnp.check_supported_arrays_type(m) if m.ndim < 1: raise ValueError(f"Input must be >= 1-d, but got {m.ndim}") diff --git a/dpnp/dpnp_iface_mathematical.py b/dpnp/dpnp_iface_mathematical.py index cacab84510bc..d619b5662b1b 100644 --- a/dpnp/dpnp_iface_mathematical.py +++ b/dpnp/dpnp_iface_mathematical.py @@ -1785,15 +1785,12 @@ def nanprod( """ - if dpnp.is_supported_array_or_scalar(a): - if issubclass(a.dtype.type, dpnp.inexact): - mask = dpnp.isnan(a) - a = dpnp.array(a, copy=True) - dpnp.copyto(a, 1, where=mask) - else: - raise TypeError( - "An array must be any of supported type, but got {}".format(type(a)) - ) + dpnp.check_supported_arrays_type(a) + + if issubclass(a.dtype.type, dpnp.inexact): + mask = dpnp.isnan(a) + a = dpnp.array(a, copy=True) + dpnp.copyto(a, 1, where=mask) return dpnp.prod( a, @@ -2108,10 +2105,7 @@ def prod( # Product reduction for complex output are known to fail for Gen9 with 2024.0 compiler # TODO: get rid of this temporary work around when OneAPI 2024.1 is released - if not isinstance(a, (dpnp_array, dpt.usm_ndarray)): - raise TypeError( - "An array must be any of supported type, but got {}".format(type(a)) - ) + dpnp.check_supported_arrays_type(a) _dtypes = (a.dtype, dtype) _any_complex = any( dpnp.issubdtype(dt, dpnp.complexfloating) for dt in _dtypes diff --git a/dpnp/linalg/dpnp_iface_linalg.py b/dpnp/linalg/dpnp_iface_linalg.py index b41b96c70525..c7437b30da60 100644 --- a/dpnp/linalg/dpnp_iface_linalg.py +++ b/dpnp/linalg/dpnp_iface_linalg.py @@ -233,14 +233,11 @@ def eigh(a, UPLO="L"): """ + dpnp.check_supported_arrays_type(a) + if UPLO not in ("L", "U"): raise ValueError("UPLO argument must be 'L' or 'U'") - if not dpnp.is_supported_array_type(a): - raise TypeError( - "An array must be any of supported type, but got {}".format(type(a)) - ) - if a.ndim < 2: raise ValueError( "%d-dimensional array given. Array must be " diff --git a/tests/helper.py b/tests/helper.py index 243c61504a50..de4db998a7bf 100644 --- a/tests/helper.py +++ b/tests/helper.py @@ -68,6 +68,14 @@ def assert_dtype_allclose( assert dpnp_arr.dtype == numpy_arr.dtype +def get_integer_dtypes(): + """ + Build a list of integer types supported by DPNP. + """ + + return [dpnp.int32, dpnp.int64] + + def get_complex_dtypes(device=None): """ Build a list of complex types supported by DPNP based on device capabilities. @@ -123,7 +131,7 @@ def get_all_dtypes( dtypes = [dpnp.bool] if not no_bool else [] # add integer types - dtypes.extend([dpnp.int32, dpnp.int64]) + dtypes.extend(get_integer_dtypes()) # add floating types dtypes.extend(get_float_dtypes(no_float16=no_float16, device=dev)) diff --git a/tests/test_indexing.py b/tests/test_indexing.py index 03541dc2d55d..4d8229e53ce7 100644 --- a/tests/test_indexing.py +++ b/tests/test_indexing.py @@ -1,10 +1,32 @@ +import functools + import numpy import pytest -from numpy.testing import assert_, assert_array_equal, assert_equal +from numpy.testing import ( + assert_, + assert_array_equal, + assert_equal, + assert_raises, +) import dpnp -from .helper import get_all_dtypes +from .helper import get_all_dtypes, get_integer_dtypes + + +def _add_keepdims(func): + """ + Hack in keepdims behavior into a function taking an axis. + """ + + @functools.wraps(func) + def wrapped(a, axis, **kwargs): + res = func(a, axis=axis, **kwargs) + if axis is None: + axis = 0 # res is now 0d and we can insert this anywhere + return dpnp.expand_dims(res, axis=axis) + + return wrapped class TestIndexing: @@ -63,6 +85,168 @@ def test_indexing_array_negative_strides(self): assert_array_equal(arr, 10.0) +class TestPutAlongAxis: + @pytest.mark.parametrize( + "arr_dt", get_all_dtypes(no_bool=True, no_none=True) + ) + @pytest.mark.parametrize("axis", list(range(2)) + [None]) + def test_replace_max(self, arr_dt, axis): + a = dpnp.array([[10, 30, 20], [60, 40, 50]], dtype=arr_dt) + + # replace the max with a small value + i_max = _add_keepdims(dpnp.argmax)(a, axis=axis) + dpnp.put_along_axis(a, i_max, -99, axis=axis) + + # find the new minimum, which should max + i_min = _add_keepdims(dpnp.argmin)(a, axis=axis) + assert_array_equal(i_min, i_max) + + @pytest.mark.parametrize( + "arr_dt", get_all_dtypes(no_bool=True, no_none=True) + ) + @pytest.mark.parametrize("idx_dt", get_integer_dtypes()) + @pytest.mark.parametrize("ndim", list(range(1, 4))) + @pytest.mark.parametrize( + "values", + [ + 777, + [100, 200, 300, 400], + (42,), + range(4), + numpy.arange(4), + dpnp.ones(4), + ], + ids=[ + "scalar", + "list", + "tuple", + "range", + "numpy.ndarray", + "dpnp.ndarray", + ], + ) + def test_values(self, arr_dt, idx_dt, ndim, values): + np_a = numpy.arange(4**ndim, dtype=arr_dt).reshape((4,) * ndim) + np_ai = numpy.array([3, 0, 2, 1], dtype=idx_dt).reshape( + (1,) * (ndim - 1) + (4,) + ) + + dp_a = dpnp.array(np_a, dtype=arr_dt) + dp_ai = dpnp.array(np_ai, dtype=idx_dt) + + for axis in range(ndim): + numpy.put_along_axis(np_a, np_ai, values, axis) + dpnp.put_along_axis(dp_a, dp_ai, values, axis) + assert_array_equal(np_a, dp_a) + + @pytest.mark.parametrize("arr_dt", get_all_dtypes()) + @pytest.mark.parametrize("idx_dt", get_integer_dtypes()) + def test_broadcast(self, arr_dt, idx_dt): + np_a = numpy.ones((3, 4, 1), dtype=arr_dt) + np_ai = numpy.arange(10, dtype=idx_dt).reshape((1, 2, 5)) % 4 + + dp_a = dpnp.array(np_a, dtype=arr_dt) + dp_ai = dpnp.array(np_ai, dtype=idx_dt) + + numpy.put_along_axis(np_a, np_ai, 20, axis=1) + dpnp.put_along_axis(dp_a, dp_ai, 20, axis=1) + assert_array_equal(np_a, dp_a) + + +class TestTakeAlongAxis: + # TODO: remove fixture once `dpnp.sort` is fully implemented + @pytest.mark.usefixtures("allow_fall_back_on_numpy") + @pytest.mark.parametrize( + "func, argfunc, kwargs", + [ + pytest.param(dpnp.sort, dpnp.argsort, {}), + pytest.param( + _add_keepdims(dpnp.min), _add_keepdims(dpnp.argmin), {} + ), + pytest.param( + _add_keepdims(dpnp.max), _add_keepdims(dpnp.argmax), {} + ), + # TODO: unmute, once `dpnp.argpartition` is implemented + # pytest.param(dpnp.partition, dpnp.argpartition, {"kth": 2}), + ], + ) + def test_argequivalent(self, func, argfunc, kwargs): + a = dpnp.random.random(size=(3, 4, 5)) + + for axis in list(range(a.ndim)) + [None]: + a_func = func(a, axis=axis, **kwargs) + ai_func = argfunc(a, axis=axis, **kwargs) + assert_array_equal( + a_func, dpnp.take_along_axis(a, ai_func, axis=axis) + ) + + @pytest.mark.parametrize( + "arr_dt", get_all_dtypes(no_bool=True, no_none=True) + ) + @pytest.mark.parametrize("idx_dt", get_integer_dtypes()) + @pytest.mark.parametrize("ndim", list(range(1, 4))) + def test_multi_dimensions(self, arr_dt, idx_dt, ndim): + np_a = numpy.arange(4**ndim, dtype=arr_dt).reshape((4,) * ndim) + np_ai = numpy.array([3, 0, 2, 1], dtype=idx_dt).reshape( + (1,) * (ndim - 1) + (4,) + ) + + dp_a = dpnp.array(np_a, dtype=arr_dt) + dp_ai = dpnp.array(np_ai, dtype=idx_dt) + + for axis in range(ndim): + expected = numpy.take_along_axis(np_a, np_ai, axis) + result = dpnp.take_along_axis(dp_a, dp_ai, axis) + assert_array_equal(expected, result) + + @pytest.mark.parametrize("xp", [numpy, dpnp]) + def test_invalid(self, xp): + a = xp.ones((10, 10)) + ai = xp.ones((10, 2), dtype=xp.intp) + + # not enough indices + assert_raises(ValueError, xp.take_along_axis, a, xp.array(1), axis=1) + + # bool arrays not allowed + assert_raises( + IndexError, xp.take_along_axis, a, ai.astype(bool), axis=1 + ) + + # float arrays not allowed + assert_raises( + IndexError, xp.take_along_axis, a, ai.astype(numpy.float32), axis=1 + ) + + # invalid axis + assert_raises(numpy.AxisError, xp.take_along_axis, a, ai, axis=10) + + @pytest.mark.parametrize("arr_dt", get_all_dtypes()) + @pytest.mark.parametrize("idx_dt", get_integer_dtypes()) + def test_empty(self, arr_dt, idx_dt): + np_a = numpy.ones((3, 4, 5), dtype=arr_dt) + np_ai = numpy.ones((3, 0, 5), dtype=idx_dt) + + dp_a = dpnp.array(np_a, dtype=arr_dt) + dp_ai = dpnp.array(np_ai, dtype=idx_dt) + + expected = numpy.take_along_axis(np_a, np_ai, axis=1) + result = dpnp.take_along_axis(dp_a, dp_ai, axis=1) + assert_array_equal(expected, result) + + @pytest.mark.parametrize("arr_dt", get_all_dtypes()) + @pytest.mark.parametrize("idx_dt", get_integer_dtypes()) + def test_broadcast(self, arr_dt, idx_dt): + np_a = numpy.ones((3, 4, 1), dtype=arr_dt) + np_ai = numpy.ones((1, 2, 5), dtype=idx_dt) + + dp_a = dpnp.array(np_a, dtype=arr_dt) + dp_ai = dpnp.array(np_ai, dtype=idx_dt) + + expected = numpy.take_along_axis(np_a, np_ai, axis=1) + result = dpnp.take_along_axis(dp_a, dp_ai, axis=1) + assert_array_equal(expected, result) + + @pytest.mark.usefixtures("allow_fall_back_on_numpy") def test_choose(): a = numpy.r_[:4] @@ -459,42 +643,6 @@ def test_put_invalid_axis(axis): dpnp.put(a, ind, vals, axis=axis) -@pytest.mark.usefixtures("allow_fall_back_on_numpy") -def test_put_along_axis_val_int(): - a = numpy.arange(16).reshape(4, 4) - ai = dpnp.array(a) - ind_r = numpy.array([[3, 0, 2, 1]]) - ind_r_i = dpnp.array(ind_r) - for axis in range(2): - numpy.put_along_axis(a, ind_r, 777, axis) - dpnp.put_along_axis(ai, ind_r_i, 777, axis) - assert_array_equal(a, ai) - - -@pytest.mark.usefixtures("allow_fall_back_on_numpy") -def test_put_along_axis1(): - a = numpy.arange(64).reshape(4, 4, 4) - ai = dpnp.array(a) - ind_r = numpy.array([[[3, 0, 2, 1]]]) - ind_r_i = dpnp.array(ind_r) - for axis in range(3): - numpy.put_along_axis(a, ind_r, 777, axis) - dpnp.put_along_axis(ai, ind_r_i, 777, axis) - assert_array_equal(a, ai) - - -@pytest.mark.usefixtures("allow_fall_back_on_numpy") -def test_put_along_axis2(): - a = numpy.arange(64).reshape(4, 4, 4) - ai = dpnp.array(a) - ind_r = numpy.array([[[3, 0, 2, 1]]]) - ind_r_i = dpnp.array(ind_r) - for axis in range(3): - numpy.put_along_axis(a, ind_r, [100, 200, 300, 400], axis) - dpnp.put_along_axis(ai, ind_r_i, [100, 200, 300, 400], axis) - assert_array_equal(a, ai) - - @pytest.mark.parametrize("vals", [[100, 200]], ids=["[100, 200]"]) @pytest.mark.parametrize( "mask", @@ -688,28 +836,6 @@ def test_take_over_index(indices, array_type, mode): assert_array_equal(expected, result) -def test_take_along_axis(): - a = numpy.arange(16).reshape(4, 4) - ai = dpnp.array(a) - ind_r = numpy.array([[3, 0, 2, 1]]) - ind_r_i = dpnp.array(ind_r) - for axis in range(2): - expected = numpy.take_along_axis(a, ind_r, axis) - result = dpnp.take_along_axis(ai, ind_r_i, axis) - assert_array_equal(expected, result) - - -def test_take_along_axis1(): - a = numpy.arange(64).reshape(4, 4, 4) - ai = dpnp.array(a) - ind_r = numpy.array([[[3, 0, 2, 1]]]) - ind_r_i = dpnp.array(ind_r) - for axis in range(3): - expected = numpy.take_along_axis(a, ind_r, axis) - result = dpnp.take_along_axis(ai, ind_r_i, axis) - assert_array_equal(expected, result) - - @pytest.mark.parametrize( "m", [None, 0, 1, 2, 3, 4], ids=["None", "0", "1", "2", "3", "4"] ) diff --git a/tests/test_sycl_queue.py b/tests/test_sycl_queue.py index 3618a9bb4c54..fc4dbf9f0d6e 100644 --- a/tests/test_sycl_queue.py +++ b/tests/test_sycl_queue.py @@ -1299,20 +1299,21 @@ def test_asarray(device_x, device_y): assert_sycl_queue_equal(y.sycl_queue, x.to_device(device_y).sycl_queue) +@pytest.mark.parametrize("func", ["take", "take_along_axis"]) @pytest.mark.parametrize( "device", valid_devices, ids=[device.filter_string for device in valid_devices], ) -def test_take(device): +def test_take(func, device): numpy_data = numpy.arange(5) dpnp_data = dpnp.array(numpy_data, device=device) - ind = [0, 2, 4] - dpnp_ind = dpnp.array(ind, device=device) + dpnp_ind = dpnp.array([0, 2, 4], device=device) + np_ind = dpnp_ind.asnumpy() - result = dpnp.take(dpnp_data, dpnp_ind) - expected = numpy.take(numpy_data, ind) + result = getattr(dpnp, func)(dpnp_data, dpnp_ind, axis=None) + expected = getattr(numpy, func)(numpy_data, np_ind, axis=None) assert_allclose(expected, result) expected_queue = dpnp_data.get_array().sycl_queue diff --git a/tests/test_usm_type.py b/tests/test_usm_type.py index 18af427e423f..f82e04a2a566 100644 --- a/tests/test_usm_type.py +++ b/tests/test_usm_type.py @@ -472,14 +472,15 @@ def test_broadcast_to(usm_type): assert x.usm_type == y.usm_type +@pytest.mark.parametrize("func", ["take", "take_along_axis"]) @pytest.mark.parametrize("usm_type_x", list_of_usm_types, ids=list_of_usm_types) @pytest.mark.parametrize( "usm_type_ind", list_of_usm_types, ids=list_of_usm_types ) -def test_take(usm_type_x, usm_type_ind): +def test_take(func, usm_type_x, usm_type_ind): x = dp.arange(5, usm_type=usm_type_x) ind = dp.array([0, 2, 4], usm_type=usm_type_ind) - z = dp.take(x, ind) + z = getattr(dp, func)(x, ind, axis=None) assert x.usm_type == usm_type_x assert ind.usm_type == usm_type_ind diff --git a/tests/third_party/cupy/indexing_tests/test_indexing.py b/tests/third_party/cupy/indexing_tests/test_indexing.py index 9e323990891c..20890056ae05 100644 --- a/tests/third_party/cupy/indexing_tests/test_indexing.py +++ b/tests/third_party/cupy/indexing_tests/test_indexing.py @@ -7,7 +7,6 @@ from tests.third_party.cupy import testing -@testing.gpu class TestIndexing(unittest.TestCase): @pytest.mark.usefixtures("allow_fall_back_on_numpy") @testing.numpy_cupy_array_equal() @@ -51,14 +50,12 @@ def test_take_index_range_overflow(self, xp, dtype): b = xp.array([0], dtype=dtype) return a.take(b) - @pytest.mark.usefixtures("allow_fall_back_on_numpy") @testing.numpy_cupy_array_equal() def test_take_along_axis(self, xp): a = testing.shaped_random((2, 4, 3), xp, dtype="float32") b = testing.shaped_random((2, 6, 3), xp, dtype="int64", scale=4) return xp.take_along_axis(a, b, axis=-2) - @pytest.mark.usefixtures("allow_fall_back_on_numpy") @testing.numpy_cupy_array_equal() def test_take_along_axis_none_axis(self, xp): a = testing.shaped_random((2, 4, 3), xp, dtype="float32")