For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /max/get-started.md).

Mojo function

argmax

def argmax(input: TileTensor[address_space=input.address_space, linear_idx_type=input.linear_idx_type, element_size=input.element_size], axis: Int, output: TileTensor[address_space=output.address_space, linear_idx_type=output.linear_idx_type, element_size=output.element_size], ctx: Optional[DeviceContext] = None)

Finds the indices of the maximum element along the specified axis.

Args:

input (TileTensor[address_space=input.address_space, linear_idx_type=input.linear_idx_type, element_size=input.element_size]): The input tensor.
axis (Int): The axis.
output (TileTensor[address_space=output.address_space, linear_idx_type=output.linear_idx_type, element_size=output.element_size]): The output tensor.
ctx (Optional[DeviceContext]): The context to execute the work on.

def argmax(input: TileTensor[address_space=input.address_space, linear_idx_type=input.linear_idx_type, element_size=input.element_size], axis_buf: TileTensor[address_space=axis_buf.address_space, linear_idx_type=axis_buf.linear_idx_type, element_size=axis_buf.element_size], output: TileTensor[address_space=output.address_space, linear_idx_type=output.linear_idx_type, element_size=output.element_size], ctx: Optional[DeviceContext] = None) where (axis_buf.flat_rank == 1)

Finds the indices of the maximum element along the specified axis.

Args:

input (TileTensor[address_space=input.address_space, linear_idx_type=input.linear_idx_type, element_size=input.element_size]): The input tensor.
axis_buf (TileTensor[address_space=axis_buf.address_space, linear_idx_type=axis_buf.linear_idx_type, element_size=axis_buf.element_size]): The axis tensor.
output (TileTensor[address_space=output.address_space, linear_idx_type=output.linear_idx_type, element_size=output.element_size]): The axis tensor.
ctx (Optional[DeviceContext]): The context to execute the work on.