mindspore.mint.topk

mindspore.mint.topk(input, k, dim=- 1, 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, 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.

Parameters

  • input (Tensor) – Input to be computed.

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

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

  • 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 or ascending order according to largest. If False , the obtained elements will not be sorted. Default: True .

Returns

A tuple consisting of values and indices.

  • 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

Examples

>>> import mindspore as ms
>>> from mindspore import mint
>>> 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 = mint.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 = mint.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]]))