| from __future__ import annotations |
|
|
| import cupy |
| from cupy import _core |
| from cupy._core import fusion |
| from cupy import _util |
|
|
| from cupy._core import _routines_indexing as _indexing |
| from cupy._core import _routines_statistics as _statistics |
|
|
|
|
| def argmax(a, axis=None, dtype=None, out=None, keepdims=False): |
| """Returns the indices of the maximum along an axis. |
| |
| Args: |
| a (cupy.ndarray): Array to take argmax. |
| axis (int): Along which axis to find the maximum. ``a`` is flattened by |
| default. |
| dtype: Data type specifier. |
| out (cupy.ndarray): Output array. |
| keepdims (bool): If ``True``, the axis ``axis`` is preserved as an axis |
| of length one. |
| |
| Returns: |
| cupy.ndarray: The indices of the maximum of ``a`` along an axis. |
| |
| .. note:: |
| ``dtype`` and ``keepdim`` arguments are specific to CuPy. They are |
| not in NumPy. |
| |
| .. note:: |
| ``axis`` argument accepts a tuple of ints, but this is specific to |
| CuPy. NumPy does not support it. |
| |
| .. seealso:: :func:`numpy.argmax` |
| |
| """ |
| |
| return a.argmax(axis=axis, dtype=dtype, out=out, keepdims=keepdims) |
|
|
|
|
| def nanargmax(a, axis=None, dtype=None, out=None, keepdims=False): |
| """Return the indices of the maximum values in the specified axis ignoring |
| NaNs. For all-NaN slice ``-1`` is returned. |
| Subclass cannot be passed yet, subok=True still unsupported |
| |
| Args: |
| a (cupy.ndarray): Array to take nanargmax. |
| axis (int): Along which axis to find the maximum. ``a`` is flattened by |
| default. |
| |
| Returns: |
| cupy.ndarray: The indices of the maximum of ``a`` |
| along an axis ignoring NaN values. |
| |
| .. note:: For performance reasons, ``cupy.nanargmax`` returns |
| ``out of range values`` for all-NaN slice |
| whereas ``numpy.nanargmax`` raises ``ValueError`` |
| .. seealso:: :func:`numpy.nanargmax` |
| """ |
| if a.dtype.kind in 'biu': |
| return argmax(a, axis=axis, dtype=dtype, out=out, keepdims=keepdims) |
|
|
| return _statistics._nanargmax(a, axis, out, dtype, keepdims) |
|
|
|
|
| def argmin(a, axis=None, dtype=None, out=None, keepdims=False): |
| """Returns the indices of the minimum along an axis. |
| |
| Args: |
| a (cupy.ndarray): Array to take argmin. |
| axis (int): Along which axis to find the minimum. ``a`` is flattened by |
| default. |
| dtype: Data type specifier. |
| out (cupy.ndarray): Output array. |
| keepdims (bool): If ``True``, the axis ``axis`` is preserved as an axis |
| of length one. |
| |
| Returns: |
| cupy.ndarray: The indices of the minimum of ``a`` along an axis. |
| |
| .. note:: |
| ``dtype`` and ``keepdim`` arguments are specific to CuPy. They are |
| not in NumPy. |
| |
| .. note:: |
| ``axis`` argument accepts a tuple of ints, but this is specific to |
| CuPy. NumPy does not support it. |
| |
| .. seealso:: :func:`numpy.argmin` |
| |
| """ |
| |
| return a.argmin(axis=axis, dtype=dtype, out=out, keepdims=keepdims) |
|
|
|
|
| def nanargmin(a, axis=None, dtype=None, out=None, keepdims=False): |
| """Return the indices of the minimum values in the specified axis ignoring |
| NaNs. For all-NaN slice ``-1`` is returned. |
| Subclass cannot be passed yet, subok=True still unsupported |
| |
| Args: |
| a (cupy.ndarray): Array to take nanargmin. |
| axis (int): Along which axis to find the minimum. ``a`` is flattened by |
| default. |
| |
| Returns: |
| cupy.ndarray: The indices of the minimum of ``a`` |
| along an axis ignoring NaN values. |
| |
| .. note:: For performance reasons, ``cupy.nanargmin`` returns |
| ``out of range values`` for all-NaN slice |
| whereas ``numpy.nanargmin`` raises ``ValueError`` |
| .. seealso:: :func:`numpy.nanargmin` |
| """ |
| if a.dtype.kind in 'biu': |
| return argmin(a, axis=axis, dtype=dtype, out=out, keepdims=keepdims) |
|
|
| return _statistics._nanargmin(a, axis, out, dtype, keepdims) |
|
|
|
|
| def nonzero(a): |
| """Return the indices of the elements that are non-zero. |
| |
| Returns a tuple of arrays, one for each dimension of a, |
| containing the indices of the non-zero elements in that dimension. |
| |
| Args: |
| a (cupy.ndarray): array |
| |
| Returns: |
| tuple of arrays: Indices of elements that are non-zero. |
| |
| .. warning:: |
| |
| This function may synchronize the device. |
| |
| .. seealso:: :func:`numpy.nonzero` |
| |
| """ |
| _util.check_array(a, arg_name='a') |
| return a.nonzero() |
|
|
|
|
| def flatnonzero(a): |
| """Return indices that are non-zero in the flattened version of a. |
| |
| This is equivalent to a.ravel().nonzero()[0]. |
| |
| Args: |
| a (cupy.ndarray): input array |
| |
| Returns: |
| cupy.ndarray: Output array, |
| containing the indices of the elements of a.ravel() that are non-zero. |
| |
| .. warning:: |
| |
| This function may synchronize the device. |
| |
| .. seealso:: :func:`numpy.flatnonzero` |
| """ |
| _util.check_array(a, arg_name='a') |
| return a.ravel().nonzero()[0] |
|
|
|
|
| _where_ufunc = _core.create_ufunc( |
| 'cupy_where', |
| ('???->?', '?bb->b', '?BB->B', '?hh->h', '?HH->H', '?ii->i', '?II->I', |
| '?ll->l', '?LL->L', '?qq->q', '?QQ->Q', '?ee->e', '?ff->f', |
| '?dd->d', '?FF->F', '?DD->D'), |
| 'out0 = in0 ? in1 : in2') |
|
|
|
|
| def where(condition, x=None, y=None): |
| """Return elements, either from x or y, depending on condition. |
| |
| If only condition is given, return ``condition.nonzero()``. |
| |
| Args: |
| condition (cupy.ndarray): When True, take x, otherwise take y. |
| x (cupy.ndarray): Values from which to choose on ``True``. |
| y (cupy.ndarray): Values from which to choose on ``False``. |
| |
| Returns: |
| cupy.ndarray: Each element of output contains elements of ``x`` when |
| ``condition`` is ``True``, otherwise elements of ``y``. If only |
| ``condition`` is given, return the tuple ``condition.nonzero()``, |
| the indices where ``condition`` is True. |
| |
| .. warning:: |
| |
| This function may synchronize the device if both ``x`` and ``y`` are |
| omitted. |
| |
| .. seealso:: :func:`numpy.where` |
| |
| """ |
|
|
| missing = (x is None, y is None).count(True) |
|
|
| if missing == 1: |
| raise ValueError('Must provide both \'x\' and \'y\' or neither.') |
| if missing == 2: |
| return nonzero(condition) |
|
|
| if fusion._is_fusing(): |
| return fusion._call_ufunc(_where_ufunc, condition, x, y) |
|
|
| return _where_ufunc(condition.astype('?', copy=False), x, y) |
|
|
|
|
| def argwhere(a): |
| """Return the indices of the elements that are non-zero. |
| |
| Returns a (N, ndim) dimantional array containing the |
| indices of the non-zero elements. Where `N` is number of |
| non-zero elements and `ndim` is dimension of the given array. |
| |
| Args: |
| a (cupy.ndarray): array |
| |
| Returns: |
| cupy.ndarray: Indices of elements that are non-zero. |
| |
| .. seealso:: :func:`numpy.argwhere` |
| |
| """ |
| _util.check_array(a, arg_name='a') |
| return _indexing._ndarray_argwhere(a) |
|
|
|
|
| |
| |
| _preamble = ''' |
| template<typename T> |
| __device__ bool _isnan(T val) { |
| return val != val; |
| } |
| ''' |
|
|
|
|
| _hip_preamble = r''' |
| #ifdef __HIP_DEVICE_COMPILE__ |
| #define no_thread_divergence(do_work, to_return) \ |
| if (!is_done) { \ |
| do_work; \ |
| is_done = true; \ |
| } |
| #else |
| #define no_thread_divergence(do_work, to_return) \ |
| do_work; \ |
| if (to_return) { return; } |
| #endif |
| ''' |
|
|
|
|
| _searchsorted_code = ''' |
| #ifdef __HIP_DEVICE_COMPILE__ |
| bool is_done = false; |
| #endif |
| |
| // Array is assumed to be monotonically |
| // increasing unless a check is requested with the |
| // `assume_increasing = False` parameter. |
| // `digitize` allows increasing and decreasing arrays. |
| bool inc = true; |
| if (!assume_increasing && n_bins >= 2) { |
| // In the case all the bins are nan the array is considered |
| // to be decreasing in numpy |
| inc = (bins[0] <= bins[n_bins-1]) |
| || (!_isnan<T>(bins[0]) && _isnan<T>(bins[n_bins-1])); |
| } |
| |
| if (_isnan<S>(x)) { |
| long long pos = (inc ? n_bins : 0); |
| if (!side_is_right) { |
| if (inc) { |
| while (pos > 0 && _isnan<T>(bins[pos-1])) { |
| --pos; |
| } |
| } else { |
| while (pos < n_bins && _isnan<T>(bins[pos])) { |
| ++pos; |
| } |
| } |
| } |
| no_thread_divergence( y = pos , true ) |
| } |
| |
| bool greater = false; |
| if (side_is_right) { |
| greater = inc && x >= bins[n_bins-1]; |
| } else { |
| greater = (inc ? x > bins[n_bins-1] : x <= bins[n_bins-1]); |
| } |
| if (greater) { |
| no_thread_divergence( y = n_bins , true ) |
| } |
| |
| long long left = 0; |
| // In the case the bins is all NaNs, digitize |
| // needs to place all the valid values to the right |
| if (!inc) { |
| while (_isnan<T>(bins[left]) && left < n_bins) { |
| ++left; |
| } |
| if (left == n_bins) { |
| no_thread_divergence( y = n_bins , true ) |
| } |
| if (side_is_right |
| && !_isnan<T>(bins[n_bins-1]) && !_isnan<S>(x) |
| && bins[n_bins-1] > x) { |
| no_thread_divergence( y = n_bins , true ) |
| } |
| } |
| |
| long long right = n_bins-1; |
| while (left < right) { |
| long long m = left + (right - left) / 2; |
| bool look_right = true; |
| if (side_is_right) { |
| look_right = (inc ? bins[m] <= x : bins[m] > x); |
| } else { |
| look_right = (inc ? bins[m] < x : bins[m] >= x); |
| } |
| if (look_right) { |
| left = m + 1; |
| } else { |
| right = m; |
| } |
| } |
| no_thread_divergence( y = right , false ) |
| ''' |
|
|
|
|
| _searchsorted_kernel = _core.ElementwiseKernel( |
| 'S x, raw T bins, int64 n_bins, bool side_is_right, ' |
| 'bool assume_increasing', |
| 'int64 y', |
| _searchsorted_code, |
| name='cupy_searchsorted_kernel', preamble=_preamble+_hip_preamble) |
|
|
|
|
| _hip_preamble = r''' |
| #ifdef __HIP_DEVICE_COMPILE__ |
| #define no_thread_divergence(do_work, to_return) \ |
| if (!is_done) { \ |
| do_work; \ |
| is_done = true; \ |
| } |
| #else |
| #define no_thread_divergence(do_work, to_return) \ |
| do_work; \ |
| if (to_return) { \ |
| out = (y == n_bins ? false : bins[y] == x); \ |
| if (invert) out = !out; \ |
| return; \ |
| } |
| #endif |
| ''' |
|
|
|
|
| _exists_kernel = _core.ElementwiseKernel( |
| 'S x, raw T bins, int64 n_bins, bool invert', |
| 'bool out', |
| ''' |
| // Handle empty bins array to avoid illegal memory access |
| if (n_bins == 0) { |
| out = (invert ? true : false); |
| return; |
| } |
| |
| const bool assume_increasing = true; |
| const bool side_is_right = false; |
| long long y; |
| ''' |
| + _searchsorted_code + ''' |
| out = (y == n_bins ? false : bins[y] == x); |
| if (invert) out = !out; |
| ''', name='cupy_exists_kernel', preamble=_preamble+_hip_preamble) |
|
|
|
|
| _exists_and_searchsorted_kernel = _core.ElementwiseKernel( |
| 'S x, raw T bins, int64 n_bins, bool invert', |
| 'bool out, int64 y', |
| ''' |
| // Handle empty bins array to avoid illegal memory access |
| if (n_bins == 0) { |
| out = (invert ? true : false); |
| y = 0; |
| return; |
| } |
| |
| const bool assume_increasing = true; |
| const bool side_is_right = false; |
| ''' |
| + _searchsorted_code + ''' |
| out = (y == n_bins ? false : bins[y] == x); |
| if (invert) out = !out; |
| ''', name='cupy_exists_and_searchsorted_kernel', |
| preamble=_preamble+_hip_preamble) |
|
|
|
|
| def searchsorted(a, v, side='left', sorter=None): |
| """Finds indices where elements should be inserted to maintain order. |
| |
| Find the indices into a sorted array ``a`` such that, |
| if the corresponding elements in ``v`` were inserted before the indices, |
| the order of ``a`` would be preserved. |
| |
| Args: |
| a (cupy.ndarray): Input array. If ``sorter`` is ``None``, then |
| it must be sorted in ascending order, |
| otherwise ``sorter`` must be an array of indices that sort it. |
| v (cupy.ndarray): Values to insert into ``a``. |
| side : {'left', 'right'} |
| If ``left``, return the index of the first suitable location found |
| If ``right``, return the last such index. |
| If there is no suitable index, return either 0 or length of ``a``. |
| sorter : 1-D array_like |
| Optional array of integer indices that sort array ``a`` into |
| ascending order. They are typically the result of |
| :func:`~cupy.argsort`. |
| |
| Returns: |
| cupy.ndarray: Array of insertion points with the same shape as ``v``. |
| |
| .. note:: When a is not in ascending order, behavior is undefined. |
| |
| .. seealso:: :func:`numpy.searchsorted` |
| |
| """ |
| return _searchsorted(a, v, side, sorter, True) |
|
|
|
|
| def _searchsorted(a, v, side, sorter, assume_increasing): |
| """`assume_increasing` is used in the kernel to |
| skip monotonically increasing or decreasing verification |
| inside the cuda kernel. |
| """ |
| if not isinstance(a, cupy.ndarray): |
| raise NotImplementedError('Only int or ndarray are supported for a') |
|
|
| if not isinstance(v, cupy.ndarray): |
| if not isinstance(v, int | float | complex): |
| raise NotImplementedError( |
| 'Only python scalars or ndarrays are supported for v') |
| v = cupy.asarray(v) |
|
|
| if a.ndim > 1: |
| raise ValueError('object too deep for desired array') |
| if a.ndim < 1: |
| raise ValueError('object of too small depth for desired array') |
| if a.size == 0: |
| return cupy.zeros(v.shape, dtype=cupy.int64) |
|
|
| a_iscomplex = a.dtype.kind == 'c' |
| v_iscomplex = v.dtype.kind == 'c' |
|
|
| if a_iscomplex and not v_iscomplex: |
| v = v.astype(a.dtype) |
| elif v_iscomplex and not a_iscomplex: |
| a = a.astype(v.dtype) |
|
|
| |
| |
| if sorter is not None: |
| if sorter.dtype.kind not in ('i', 'u'): |
| raise TypeError('sorter must be of integer type') |
| if sorter.size != a.size: |
| raise ValueError('sorter.size must equal a.size') |
| a = a.take(sorter) |
|
|
| y = cupy.zeros(v.shape, dtype=cupy.int64) |
|
|
| _searchsorted_kernel(v, a, a.size, side == 'right', assume_increasing, y) |
| return y |
|
|
|
|
| |
|
|
