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 struct
TopK
struct TopK
Registers the mo.top_k graph op with the graph compiler.
Implemented traits
Methods
execute
static def execute[dtype: DType, rank: Int, target: StringSpan[ImmStaticOrigin], _trace_name: StringSpan[ImmStaticOrigin]](values: ManagedTensorSlice[IOSpec[_, _].Output, static_spec=values.static_spec], indices: ManagedTensorSlice[IOSpec[_, _].Output, static_spec=indices.static_spec], input: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=input.static_spec], k: Scalar, axis: Scalar, sorted: Scalar[DType.bool], ctx: DeviceContext)
Executes the mo.top_k graph op.
Parameters:
- dtype (
DType): Element type of the input and output tensors. - rank (
Int): Tensor rank of the input and output tensors. - target (
StringSpan[ImmStaticOrigin]): Compilation target string. - _trace_name (
StringSpan[ImmStaticOrigin]): Name used for tracing and debugging.
Args:
- values (
ManagedTensorSlice[IOSpec[_, _].Output, static_spec=values.static_spec]): See the graph op signature. - indices (
ManagedTensorSlice[IOSpec[_, _].Output, static_spec=indices.static_spec]): See the graph op signature. - input (
ManagedTensorSlice[IOSpec[_, _].Input, static_spec=input.static_spec]): Input tensor to reduce. - k (
Scalar): Number of elements to select. - axis (
Scalar): See the graph op signature. - sorted (
Scalar[DType.bool]): Whether to sort the selected elements. - ctx (
DeviceContext): Device context used to enqueue the kernel.
Raises:
Error: If the operation parameters are invalid.