mindspore.Tensor.topk

View Source On Gitee
Tensor.topk(k, dim=- 1, largest=True, sorted=True)

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 self 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 self, 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
  • 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 self.

Raises
Supported Platforms:

Ascend

Examples

>>> import mindspore as ms
>>> from mindspore import Tensor
>>> 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 = Tensor.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 = Tensor.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]]))
Tensor.topk(k, dim=None, largest=True, sorted=True)

For more details, please refer to mindspore.ops.topk().