topk¶

nvtripy.topk(input: Tensor, k: int, dim: int) → Tuple[Tensor, Tensor][source]¶

Returns the k largest values in the tensor and their indices along the specified dimension.

Parameters:

input (Tensor) – [dtype=T1] The input tensor.
k (int) – The number of values to take.
dim (int) – The dimension along which to find the top-k values.

Returns:

[dtype=[‘T1’, ‘T2’]] The top-k values and indices, in sorted order.

Return type:

Tuple[Tensor, Tensor]

DATA TYPE CONSTRAINTS:

T1: float16, bfloat16, float32, int32, int64
T2: int32

Example

inp = tp.iota((1, 5), dim=1) + 2.5
values, indices = tp.topk(inp, k=2, dim=1)

Local Variables¶

>>> inp
tensor(
    [[2.5, 3.5, 4.5, 5.5, 6.5]], 
    dtype=float32, loc=gpu:0, shape=(1, 5))

>>> values
tensor(
    [[6.5, 5.5]], 
    dtype=float32, loc=gpu:0, shape=(1, 2))

>>> indices
tensor(
    [[4, 3]], 
    dtype=int32, loc=gpu:0, shape=(1, 2))