IMPORTANT: To view this page as Markdown, append `.md` to the URL (e.g. /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. /get-started.md).

Mojo function

argmaxmin_gpu

def argmaxmin_gpu[dtype: DType, output_type: DType, largest: Bool](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])

Reduces the inner-most dimension to the index of its largest/smallest element with a single streaming pass.

Parameters:

  • โ€‹dtype (DType): DType - The data dtype of the input tensor.
  • โ€‹output_type (DType): DType - The data dtype of the output tensor.
  • โ€‹largest (Bool): Bool - Whether to perform argmax or argmin.