ttnn.argmax

ttnn.argmax(input_tensor: ttnn.Tensor, *, dim: int | None, keepdim: bool | None, sub_core_grids: CoreRangeSet | None, memory_config: ttnn.MemoryConfig | None, output_tensor: ttnn.Tensor | None) → None

Argmax. Returns indices of maximum values. Output is UINT32, ROW_MAJOR, INTERLEAVED (DRAM or L1).

Parameters:

input_tensor (ttnn.Tensor) – On-device, INTERLEAVED input.

Keyword Arguments:

  • dim (int, optional) – Dim to reduce. None reduces all elements (ROW_MAJOR input only). Default: None.

  • keepdim (bool, optional) – Keep reduced dim. Default: False.

  • sub_core_grids (CoreRangeSet, optional) – Limits execution to a subset of cores. Supported on ROW_MAJOR last-dim reductions (<= 2 ranges) and batch/channel dim reductions. Default: None.

  • memory_config (ttnn.MemoryConfig, optional) – Output memory (INTERLEAVED DRAM/L1). Default: input’s memory_config.

  • output_tensor (ttnn.Tensor, optional) – Preallocated output (must be UINT32, ROW_MAJOR, INTERLEAVED, same device). Default: None.

Supported:

  • dim=None (reduce all elements): - input layout: ROW_MAJOR - dtypes: BFLOAT16/FLOAT32/INT32/UINT32/UINT16

  • dim = rank-1 (last / width): - ROW_MAJOR input: BFLOAT16/FLOAT32/INT32/UINT32/UINT16 (multi-core by default) - TILE input: BFLOAT16/FLOAT32 (single-core)

  • dim = rank-2 (height): - BFLOAT16/FLOAT32 only - ROW_MAJOR inputs are internally tilized; this path runs single-core

  • 0 <= dim < rank-2 (batch/channel dims, rank >= 3): - BFLOAT16/FLOAT32 only (integer dtypes not supported) - input may be ROW_MAJOR or TILE (ROW_MAJOR is converted to TILE internally) - output is produced in TILE internally and converted to ROW_MAJOR - sub_core_grids is supported (pass a single-core CoreRangeSet to run on one core)

Not supported:

  • Sharded tensors (inputs/outputs must be INTERLEAVED)

  • TILE input with dim=None

  • Batch/channel dim reductions with INT/UINT inputs

  • Integer dtypes on batch/channel dim reductions

Example

# Create tensor
tensor_input = 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(tensor_input, dim=-1, keepdim=True)
logger.info(f"Argmax onedim result: {output_onedim}")

# All dim reduction yields shape of []
output_alldim = ttnn.argmax(tensor_input)
logger.info(f"Argmax alldim result: {output_alldim}")