oneflow.topk

oneflow.topk()

Finds the values and indices of the k largest entries at specified axis.

Parameters
  • input (oneflow.Tensor) – Input Tensor

  • k (int) – the k in “top-k”

  • dim (int, optional) – the dimension to sort along. Defaults to the last dim (-1)

  • largest (bool, optional) – controls whether to return largest or smallest elements

  • sorted (bool, optional) – controls whether to return the elements in sorted order (Only Support True Now!)

Returns

A tuple of (values, indices), where the indices are the indices of the elements in the original input tensor.

Return type

Tuple(oneflow.Tensor, oneflow.Tensor(dtype=int32))

For example:

>>> import oneflow as flow
>>> import numpy as np
>>> x = np.array([[1, 3, 8, 7, 2], [1, 9, 4, 3, 2]], dtype=np.float32)
>>> result = flow.topk(flow.Tensor(x), k=3, dim=1)
>>> result.values
tensor([[8., 7., 3.],
        [9., 4., 3.]], dtype=oneflow.float32)
>>> result.indices
tensor([[2, 3, 1],
        [1, 2, 3]], dtype=oneflow.int64)
>>> result.values.shape
oneflow.Size([2, 3])
>>> result.indices.shape
oneflow.Size([2, 3])
>>> result = flow.topk(flow.Tensor(x), k=2, dim=1, largest=False)
>>> result.values
tensor([[1., 2.],
        [1., 2.]], dtype=oneflow.float32)
>>> result.indices
tensor([[0, 4],
        [0, 4]], dtype=oneflow.int64)
>>> result.values.shape
oneflow.Size([2, 2])
>>> result.indices.shape
oneflow.Size([2, 2])