topk

paddle. topk ( x, k, axis=None, largest=True, sorted=True, name=None ) [source]

Return values and indices of the k largest or smallest at the optional axis. If the input is a 1-D Tensor, finds the k largest or smallest values and indices. If the input is a Tensor with higher rank, this operator computes the top k values and indices along the axis.

Parameters

  • x (Tensor) – Tensor, an input N-D Tensor with type float32, float64, int32, int64.

  • k (int, Tensor) – The number of top elements to look for along the axis.

  • axis (int, optional) – Axis to compute indices along. The effective range is [-R, R), where R is x.ndim. when axis < 0, it works the same way as axis + R. Default is -1.

  • largest (bool, optional) – largest is a flag, if set to true, algorithm will sort by descending order, otherwise sort by ascending order. Default is True.

  • sorted (bool, optional) – controls whether to return the elements in sorted order, default value is True. In gpu device, it always return the sorted value.

  • name (str, optional) – Name for the operation (optional, default is None). For more information, please refer to Name.

Returns

tuple(Tensor), return the values and indices. The value data type is the same as the input x. The indices data type is int64.

Examples

>>> import paddle

>>> data_1 = paddle.to_tensor([1, 4, 5, 7])
>>> value_1, indices_1 = paddle.topk(data_1, k=1)
>>> print(value_1)
Tensor(shape=[1], dtype=int64, place=Place(cpu), stop_gradient=True,
[7])
>>> print(indices_1)
Tensor(shape=[1], dtype=int64, place=Place(cpu), stop_gradient=True,
[3])

>>> data_2 = paddle.to_tensor([[1, 4, 5, 7], [2, 6, 2, 5]])
>>> value_2, indices_2 = paddle.topk(data_2, k=1)
>>> print(value_2)
Tensor(shape=[2, 1], dtype=int64, place=Place(cpu), stop_gradient=True,
[[7],
 [6]])
>>> print(indices_2)
Tensor(shape=[2, 1], dtype=int64, place=Place(cpu), stop_gradient=True,
[[3],
 [1]])

>>> value_3, indices_3 = paddle.topk(data_2, k=1, axis=-1)
>>> print(value_3)
Tensor(shape=[2, 1], dtype=int64, place=Place(cpu), stop_gradient=True,
[[7],
 [6]])
>>> print(indices_3)
Tensor(shape=[2, 1], dtype=int64, place=Place(cpu), stop_gradient=True,
[[3],
 [1]])

>>> value_4, indices_4 = paddle.topk(data_2, k=1, axis=0)
>>> print(value_4)
Tensor(shape=[1, 4], dtype=int64, place=Place(cpu), stop_gradient=True,
[[2, 6, 5, 7]])
>>> print(indices_4)
Tensor(shape=[1, 4], dtype=int64, place=Place(cpu), stop_gradient=True,
[[1, 1, 0, 0]])