mindspore.ops.topk

mindspore.ops.topk(input, k, dim=None, largest=True, sorted=True)[source]

Finds values and indices of the k largest or smallest entries along a given dimension.

Warning

  • If sorted is set to False, it will use the aicpu operator, the performance may be reduced. In addition, due to different memory layout and traversal methods on different platforms, the display order of calculation results may be inconsistent when sorted is False.

If the input is a one-dimensional Tensor, finds the k largest or smallest entries in the Tensor, and outputs its value and index as a Tensor. values[k] is the k largest item in input, and its index is indices [k].

For a multi-dimensional matrix, calculates the first or last k entries in a given dimension, therefore:

\[values.shape = indices.shape\]

If the two compared elements are the same, the one with the smaller index value is returned first.

Note

Currently, Ascend/CPU supported all common data types except bool and complex type, but GPU only supports float16, float32 currently.

Parameters

  • input (Tensor) – Input to be computed.

  • k (int) – The number of top or bottom elements to be computed along the last dimension, constant input is needed.

  • dim (int, optional) – The dimension to sort along. Default: None .

  • largest (bool, optional) – If largest is False then the k smallest elements are returned. Default: True .

  • sorted (bool, optional) – If True , the obtained elements will be sorted by the values in descending order. If False , the obtained elements will not be sorted. Default: True .

Returns

A tuple consisting of values and indexes.

  • values (Tensor): The k largest or smallest elements in each slice of the given dimension.

  • indices (Tensor): The indices of values within the last dimension of input.

Raises
Supported Platforms:

Ascend GPU CPU

Examples

>>> import mindspore as ms
>>> from mindspore import ops
>>> x = ms.Tensor([[0.5368, 0.2447, 0.4302, 0.9673],
...                [0.4388, 0.6525, 0.4685, 0.1868],
...                [0.3563, 0.5152, 0.9675, 0.8230]], dtype=ms.float32)
>>> output = ops.topk(x, 2, dim=1)
>>> print(output)
(Tensor(shape=[3, 2], dtype=Float32, value=
[[ 9.67299998e-01,  5.36800027e-01],
 [ 6.52499974e-01,  4.68499988e-01],
 [ 9.67499971e-01,  8.23000014e-01]]), Tensor(shape=[3, 2], dtype=Int32, value=
[[3, 0],
 [1, 2],
 [2, 3]]))
>>> output2 = ops.topk(x, 2, dim=1, largest=False)
>>> print(output2)
(Tensor(shape=[3, 2], dtype=Float32, value=
[[ 2.44700000e-01,  4.30200011e-01],
 [ 1.86800003e-01,  4.38800007e-01],
 [ 3.56299996e-01,  5.15200019e-01]]), Tensor(shape=[3, 2], dtype=Int32, value=
[[1, 2],
 [3, 0],
 [0, 1]]))