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
top_k
def top_k[dtype: DType, out_idx_type: DType, //, largest: Bool = True, target: StringSpan[ImmStaticOrigin] = StringSpan("cpu")](input: TileTensor[dtype, Storage=input.Storage, address_space=input.address_space, linear_idx_type=input.linear_idx_type], max_k: Int, axis: Int, out_vals: TileTensor[dtype, Storage=out_vals.Storage, address_space=out_vals.address_space, linear_idx_type=out_vals.linear_idx_type], out_idxs: TileTensor[out_idx_type, Storage=out_idxs.Storage, address_space=out_idxs.address_space, linear_idx_type=out_idxs.linear_idx_type], sorted: Bool, ctx: DeviceContext, k: Optional[TileTensor[DType.int64, Layout[*(), *()], ImmutAnyOrigin]] = None)
Implementation of the Top K algorithm. Returns the top or bottom K elements and their index along a specified axis.
Parameters:
- dtype (
DType): Data type of the input buffer. - out_idx_type (
DType): The data dtype of the output indices (default == DType.int64). - largest (
Bool): Whether to find the maximum (top k) or minimum value (bottom k). - target (
StringSpan[ImmStaticOrigin]): The target to run on.
Args:
- input (
TileTensor[dtype, Storage=input.Storage, address_space=input.address_space, linear_idx_type=input.linear_idx_type]): The input tensor. - max_k (
Int): The largest number of top elements. - axis (
Int): The axis along which to operate. - out_vals (
TileTensor[dtype, Storage=out_vals.Storage, address_space=out_vals.address_space, linear_idx_type=out_vals.linear_idx_type]): Output values. - out_idxs (
TileTensor[out_idx_type, Storage=out_idxs.Storage, address_space=out_idxs.address_space, linear_idx_type=out_idxs.linear_idx_type]): Output indices. - sorted (
Bool): Indicates if the top/bottom K elements are in (stable) sorted order. - ctx (
DeviceContext): The device call context. - k (
Optional[TileTensor[DType.int64, Layout[*(), *()], ImmutAnyOrigin]]): Per batch element k value.