megengine.functional.topk

topk(inp, k, descending=True, kth_only=False, no_sort=False)[源代码]

按行排序,选出二维矩阵中 Top-K (默认情况下)个最小元素。

参数
  • inp (Tensor) – 输入张量,如果是二维的,则该张量的每行将进行排序。

  • k (int) – 所需元素的个数。

  • descending (bool) – 如果为True,则返回最大的元素。 默认值:True

  • kth_only (bool) – 该值如果为 True,则只返回第k个元素。 默认:False

  • no_sort (bool) – 该值如果为 True,则返回值可能是无序的。 默认:False

返回类型

Tuple[Tensor, Tensor]

返回

由两个张量组成的元组 (topk_tensor, indices_of_int32)

实际案例

>>> import numpy as np
>>> x = Tensor(np.array([2, 4, 6, 8, 7, 5, 3, 1], dtype=np.float32))
>>> top, indices = F.topk(x, 5, descending=False)
>>> print(top.numpy(), indices.numpy())
[1. 2. 3. 4. 5.] [7 0 6 1 5]