dpnp.top_k¶

dpnp.top_k(a, k, /, *, axis=-1, mode='largest', sorted=True)[source]¶

Returns the k largest or smallest elements and their indices along an axis.

For full documentation refer to numpy.top_k.

Parameters:

a{dpnp.ndarray, usm_ndarray}

Input array.

kint

The number of largest or smallest elements to return. k must be a non-negative integer not greater than the size of a along axis.

axis{None, int}, optional

Axis along which to find the largest or smallest elements. If None, the flattened array is used.

Default: -1.

mode{"largest", "smallest"}, optional

If "largest", the largest elements are returned. If "smallest", the smallest elements are returned. NaN values (and complex values with a NaN component) are ordered to the end regardless of mode, so they are only returned if they are among the top k values.

Default: "largest".

sortedbool, optional

If True, the top k elements are returned in sorted order: largest-first for mode="largest" and smallest-first for mode="smallest". If False, the order of the elements is unspecified.

Default: True.

Returns:

valuesdpnp.ndarray: The top k values of a, with the same data type as a. The shape is that of a with the dimension along axis replaced by k.
indicesdpnp.ndarray: The indices of values along axis, with the same shape as values and the default array index data type.

See also

dpnp.tensor.top_k: Array API version; returns the elements in an unspecified order.
dpnp.partition: Partial sort.
dpnp.sort: Full sorting.

Notes

The returned indices are not guaranteed to be stable, i.e., the order of the indices of equal values is not guaranteed to match their order in a. This is the case regardless of the value of sorted.

In contrast to numpy.top_k, which returns empty arrays when axis is empty, dpnp.top_k raises ValueError for an empty axis and k > 0, as it does for any k greater than the size of axis.

Examples

>>> import dpnp as np
>>> a = np.array([[1, 2, 3, 4], [4, 3, 2, 1]])
>>> res = np.top_k(a, 2)
>>> res.values
array([[4, 3],
       [4, 3]])
>>> res.indices
array([[3, 2],
       [0, 1]])

>>> values, indices = np.top_k(a, 2, axis=0)
>>> values
array([[4, 3, 3, 4],
       [1, 2, 2, 1]])
>>> indices
array([[1, 1, 0, 0],
       [0, 0, 1, 1]])

>>> values, indices = np.top_k(a, 2, axis=1, mode="smallest")
>>> values
array([[1, 2],
       [1, 2]])
>>> indices
array([[0, 1],
       [3, 2]])

>>> values, indices = np.top_k(np.array([1.0, 2.0, 3.0, np.nan]), 2)
>>> values
array([3., 2.])
>>> indices
array([2, 1])