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 If no keepdim is provided, it will default to False.

Currenly this op only support dimension-specific reduction on last dimension.

Input tensor must have BFLOAT16 data type and ROW_MAJOR layout.

Output tensor will have UINT32 data type.

Equivalent pytorch code:

return torch.argmax(input_tensor, dim=dim, keepdim=keepdim)
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:

List of ttnn.Tensor – the output tensor.