mindspore.ops.top_k
- mindspore.ops.top_k(input_x, k, sorted=True)[source]
Finds values and indices of the k largest entries along the last dimension.
Warning
If sorted is set to ‘False’, it will use the aicpu operator, the performance may be reduced.
If the input_x is a one-dimensional Tensor, finds the k largest entries in the Tensor, and outputs its value and index as a Tensor. Therefore, values[k] is the k largest item in input_x, and its index is indices [k].
For a multi-dimensional matrix, calculates the first k entries in each row (corresponding vector along the last dimension), therefore:
\[values.shape = indices.shape = input.shape[:-1] + [k].\]If the two compared elements are the same, the one with the smaller index value is returned first.
- Parameters
input_x (Tensor) – Input to be computed, data type must be float16, float32 or int32.
k (int) – The number of top elements to be computed along the last dimension, constant input is needed.
sorted (bool, optional) – If true, the obtained elements will be sorted by the values in descending order. Default: True.
- Returns
Tuple of 2 tensors, the values and the indices.
values (Tensor): The k largest elements in each slice of the last dimension.
indices (Tensor): The indices of values within the last dimension of input.
- Raises
- Supported Platforms:
Ascend
GPU
CPU
Examples
>>> from mindspore import Tensor >>> from mindspore import ops >>> import mindspore >>> input_x = Tensor([1, 2, 3, 4, 5], mindspore.float16) >>> k = 3 >>> values, indices = ops.top_k(input_x, k, sorted=True) >>> print((values, indices)) (Tensor(shape=[3], dtype=Float16, value= [ 5.0000e+00, 4.0000e+00, 3.0000e+00]), Tensor(shape=[3], dtype=Int32, value= [4, 3, 2]))