mindspore.ops.topk

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

按指定维度返回前 k 个最大或最小元素及其对应索引。

警告

如果 sorted 设置为False，它将使用aicpu运算符，性能可能会降低。此外，由于在不同平台上存在内存排布以及遍历方式不同等问题，sorted 设置为False时计算结果的显示顺序可能会出现不一致的情况。

参数：

input (Tensor) - 输入tensor。
k (int) - 返回的元素数量。
dim (int, 可选) - 指定排序的维度。如果为 None，则会按最后一个维度排序。默认 None 。
largest (bool, 可选) - 如果为 True ，返回最大元素。如果为 False ，返回最小元素。默认 True 。
sorted (bool, 可选) - 如果为 True ，则返回的元素按降序排序。如果为 False ，则不对获取的元素进行排序。默认值： True 。

返回：

两个tensor组成的tuple(values, indices)

支持平台：

Ascend GPU CPU

样例：

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