topk

paddle. topk ( x: Tensor, k: int | Tensor, axis: int | None = None, largest: bool = True, sorted: bool = True, name: str | None = None, *, out: tuple[Tensor, Tensor] | None = None ) → TopKRetType [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|None, 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.

• name (str|None, optional) – Name for the operation (optional, default is None). For more information, please refer to api_guide_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]])