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
  • TypeError – If sorted is not a bool.

  • TypeError – If input_x is not a Tensor.

  • TypeError – If k is not an int.

  • TypeError – If dtype of input_x is not one of the following: float16, float32 or int32.

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]))