ttnn.topk

ttnn.topk = Operation(python_fully_qualified_name='ttnn.topk', function=<ttnn._ttnn.operations.reduction.topk_t object>, preprocess_golden_function_inputs=<function default_preprocess_golden_function_inputs>, golden_function=<function _create_golden_function_topk.<locals>.golden_function>, postprocess_golden_function_outputs=<function default_postprocess_golden_function_outputs>, is_cpp_operation=True, is_experimental=False)

ttnn.Tensor, k: int, dim: int, largest: bool, sorted: bool, out : Optional[ttnn.Tensor] = std::nullopt, memory_config: MemoryConfig = std::nullopt, queue_id : [int] = 0) -> Tuple[ttnn.Tensor, 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. Note that when using BFLOAT8, a different set of elements than in the input may share the same exponent, causing some values to be rounded up or down.

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.

Type:

topk(input_tensor