ttnn.topk

ttnn.topk(input_tensor: ttnn.Tensor, k: number, dim: number, largest: bool = False, sorted: bool = False, *, memory_config: ttnn.MemoryConfig | None = None, output_tensor: ttnn.Tensor | None = None, queue_id: int | None = 0) → List of ttnn.Tensor

Returns the k largest or k smallest elements of the given input tensor along a given dimension.

If dim is not provided, the last dimension of the input tensor is used.

If largest is True, the k largest elements are returned. Otherwise, the k smallest elements are returned.

The boolean option sorted if True, will make sure that the returned k elements are sorted.

Input tensor must have BFLOAT8 or BFLOAT16 data type and TILE_LAYOUT layout.

Output value tensor will have the same data type as input tensor and output index tensor will have UINT16 data type.

Equivalent pytorch code:

return torch.topk(input_tensor, k, dim=dim, largest=largest, sorted=sorted, *, out=None)
Parameters:

  • input_tensor (ttnn.Tensor) – the input tensor.

  • k (number) – the number of top elements to look for.

  • dim (number) – the dimension to reduce.

  • largest (bool) – whether to return the largest or the smallest elements. Defaults to False.

  • sorted (bool) – whether to return the elements in sorted order. Defaults to False.

Keyword Arguments:

  • memory_config (ttnn.MemoryConfig, optional) – Memory configuration for the operation. Defaults to None.

  • output_tensor (ttnn.Tensor, optional) – Preallocated output tensor. Defaults to None.

  • queue_id (int, optional) – command queue id. Defaults to 0.

Returns:

List of ttnn.Tensor – the output tensor.