mindspore.ops.topk

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

Return the top k largest or smallest elements of the input tensor along a specified 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.

Parameters

input (Tensor) – The input tensor.
k (int) – The number elements to be returned.
dim (int, optional) – Specify the dimension for sorting. Default None .
largest (bool, optional) – If True , return largest elements. If False , then return smallest elements. Default True .
sorted (bool, optional) – If True , the elements are returned in descending order. If False , the obtained elements will not be sorted. Default True .

Returns

Tuple(values, indices) of 2 tensors.

Supported Platforms:: Ascend GPU CPU

Examples

>>> import mindspore
>>> input = mindspore.tensor([[8, 2, 1],
...                           [5, 9, 3],
...                           [4, 6, 7]])
>>> # case 1: If dim is not given, the last dimension of the input is chosen.
>>> mindspore.ops.topk(input, 2)
(Tensor(shape=[3, 2], dtype=Int64, value=
 [[8, 2],
  [9, 5],
  [7, 6]]),
 Tensor(shape=[3, 2], dtype=Int32, value=
 [[0, 1],
  [1, 0],
  [2, 1]]))
>>> # case 2: when dim is 0:
>>> mindspore.ops.topk(input, 2, dim=0)
(Tensor(shape=[2, 3], dtype=Int64, value=
[[8, 9, 7],
 [5, 6, 3]]),
Tensor(shape=[2, 3], dtype=Int32, value=
[[0, 1, 2],
 [1, 2, 1]]))
>>> # case 3: when largest is False, return smallest values.
>>> mindspore.ops.topk(input, 2, dim=0, largest=False)
(Tensor(shape=[2, 3], dtype=Int64, value=
 [[4, 2, 1],
  [5, 6, 3]]),
 Tensor(shape=[2, 3], dtype=Int32, value=
 [[2, 0, 0],
  [1, 2, 1]]))