megengine.functional.topk#

topk(inp, k, descending=False, kth_only=False, no_sort=False)[source]#

Selects the Top-K (by default) smallest elements of 2d matrix by row.

Parameters:

  • inp (Tensor) – input tensor. If input tensor is 2d, each row will be sorted.

  • k (int) – number of elements needed.

  • descending (bool) – if True, return the largest elements instead. Default: False

  • kth_only (bool) – if True, only the k-th element will be returned. Default: False

  • no_sort (bool) – if True, the returned elements can be unordered. Default: False

Return type:

Tuple[Tensor, Tensor]

Returns:

tuple of two tensors (topk_tensor, indices_of_int32)

Examples

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