Top-k Operators

本页暂无中文版。以下为英文原文。

Every op on this page is used the same way: construct it once, then call it. The constructor takes what the kernel is compiled with; the call takes the tensors. Both are documented under each op — __init__ and forward, where forward is what runs when you call op(...).

Top-k selection

tileops.ops.topk_selector.TopkSelectorFwdOp

__init__

__init__(
    topk,
    tune=False,
)

Build the op. Shapes and dtype are taken from the first call.

Parameters:

• topk (int) –

  Manifest params.topk, int.

• tune (bool, default: False ) –

  Whether to autotune, applied when a kernel is first built.

forward

forward(
    index_score,
    starts,
    ends,
)

Run the op on the inputs the manifest declares.

Parameters:

• index_score –

  Input tensor, dtype float32.

• starts –

  Input tensor, dtype int32.

• ends –

  Input tensor, dtype int32.

Returns:

• Tensor –

  indexes, as the manifest declares.