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 module

argmaxmin_gpu

Provides a single-pass GPU argmax/argmin kernel over the inner-most dimension.

Functions

argmax_gpu: Computes the indices of the maximum values along the inner-most dimension of the input tensor on the GPU.
argmaxmin_gpu: Reduces the inner-most dimension to the index of its largest/smallest element with a single streaming pass.
argmin_gpu: Computes the indices of the minimum values along the inner-most dimension of the input tensor on the GPU.

Functions