mindspore.ops.TopK¶
-
class
mindspore.ops.
TopK
(*args, **kwargs)[source]¶ Finds values and indices of the k largest entries along the last dimension.
- Parameters
sorted (bool) – If true, the obtained elements will be sorted by the values in descending order. Default: False.
- Inputs:
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.
- Outputs:
Tuple of 2 tensors, the values and the indices.
values (Tensor) - The k largest elements in each slice of the last dimensional.
indices (Tensor) - The indices of values within the last dimension of input.
- Raises
- Supported Platforms:
Ascend
GPU
CPU
Examples
>>> topk = ops.TopK(sorted=True) >>> input_x = Tensor([1, 2, 3, 4, 5], mindspore.float16) >>> k = 3 >>> values, indices = topk(input_x, k) >>> 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]))