ttnn.argmax

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

Returns the indices of the maximum value of elements in the input_tensor. If no dim is provided, it will return the indices of maximum value of all elements in given input_tensor.

Parameters:

input_tensor (ttnn.Tensor) – the input tensor.

Keyword Arguments:

  • dim (int, optional) – dimension to reduce. Defaults to None.

  • keepdim (bool, optional) – whether to keep the reduced dimension. Defaults to False.

  • 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:

ttnn.Tensor – Output tensor containing the indices of the maximum value.

Note

The input tensor supports the following data types and layouts:

Input Tensor

dtype - layout

FLOAT32 - ROW_MAJOR

BFLOAT16 - ROW_MAJOR

UINT32 - ROW_MAJOR

INT32 - ROW_MAJOR

UINT16 - ROW_MAJOR

The output tensor will be of the following data type and layout:

Output Tensor

dtype - layout

UINT32 - ROW_MAJOR
Limitations:

Currently this op only supports dimension-specific reduction on the last dimension (i.e. dim = -1).

Example

input_tensor = ttnn.rand([1, 1, 32, 64], device=device, layout=ttnn.ROW_MAJOR_LAYOUT)

# Last dim reduction yields shape of [1, 1, 32, 1] output_onedim = ttnn.argmax(input_tensor, dim=-1, keepdim=True)

# All dim reduction yields shape of [] output_alldim = ttnn.argmax(input_tensor)