IMPORTANT: To view this page as Markdown, append `.md` to the URL (e.g. /max/get-started.md). For the complete documentation index, see llms.txt.
Skip to main content
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_gpu

def argmax_gpu[dtype: DType, output_type: DType](ctx: DeviceContext, input: TileTensor[dtype, Storage=input.Storage, address_space=input.address_space, linear_idx_type=input.linear_idx_type], output: TileTensor[output_type, Storage=output.Storage, address_space=output.address_space, linear_idx_type=output.linear_idx_type])

Computes the indices of the maximum values along the inner-most dimension of the input tensor on the GPU.

Parameters:

  • dtype (DType): DType - The data dtype of the input tensor.
  • output_type (DType): DType - The data dtype of the output tensor.

Was this page helpful?