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 module
topk
Structs
Functions
-
apply_gumbel_noise_kernel: -
fused_token_sampling_cpu: Generalized implementation of the Top K algorithm with sampling. Returns the sampled index from the innermost dimension of the input tensor for each row/subvolume. -
fused_token_sampling_gpu: Top K algorithm with fused sampling. Returns the sampled indices from the Top-K of the innermost dimension of the input tensor for each row/subvolume. -
gumbel_sampling_gpu: Gumbel sampling using the Gumbel-max trick for categorical distributions. -
top_k: Implementation of the Top K algorithm. Returns the top or bottom K elements and their index along a specified axis. -
top_k_shape_impl: Compute the output shape of a top/bottom k operation. -
topk_gpu: Generalized implementation of the Top K algorithm with/without sampling. Returns the sampled index from the innermost dimension of the input tensor for each row/subvolume or the top K values and indices across the tensor.
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!