ttnn.max

ttnn.max = Operation(python_fully_qualified_name='ttnn.max', function=<ttnn._ttnn.operations.reduction.max_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)

ttnn.max(input_tensor: ttnn.Tensor, dim: Optional[int] = None, keepdim: bool = False, memory_config: Optional[ttnn.MemoryConfig] = None, compute_kernel_config: Optional[ttnn.ComputeKernelConfig] = None, scalar: float = 1.0, correction: bool = True) -> ttnn.Tensor

Computes the max of the input tensor input_a along the specified dimension dim. If no dimension is provided, max is computed over all dimensions yielding a single value.

Parameters:

  • input_a (ttnn.Tensor) – the input tensor. Must be on the device.

  • dim (number) – dimension value to reduce over.

  • keepdim (bool, optional) – keep original dimension size. Defaults to False.

Keyword Arguments:

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

  • compute_kernel_config (ttnn.ComputeKernelConfig, optional) – Compute kernel configuration for the operation. Defaults to None.

  • scalar (float, optional) – A scaling factor to be applied to the input tensor. Defaults to 1.0.

  • correction (bool, optional) – Applies only to ttnn.std() - whether to apply Bessel’s correction (i.e. N-1). Defaults to True.

Returns:

ttnn.Tensor – the output tensor.

Note

The input tensor supports the following data types and layouts:

Input Tensor

dtype

layout

FLOAT32

ROW_MAJOR, TILE

BFLOAT16

ROW_MAJOR, TILE

BFLOAT8_B

ROW_MAJOR, TILE

INT32

ROW_MAJOR, TILE

UINT32

ROW_MAJOR, TILE

The output tensor will match the data type and layout of the input tensor.

Memory Support:

  • Interleaved: DRAM and L1

  • Sharded (L1): Width, Height, and ND sharding

  • Output sharding/layout will mirror the input

Example

input_a = ttnn.rand(1, 2), dtype=torch.bfloat16, device=device)
output = ttnn.max(input_a, dim, memory_config)