# user interface

import torch
from ..jit.core import (
    compile_ops,
)


@compile_ops("module_topk_plain")
def topk_plain(
    x: torch.Tensor,
    topk_ids: torch.Tensor,
    topk_out: torch.Tensor,
    topk: int,
    largest: bool = True,
    rowStarts: torch.Tensor = None, # 变长序列中每个批次的起始索引，形状为[batch_size]
    rowEnds: torch.Tensor = None,   # 变长序列中每个批次的结束索引，形状为[batch_size]。每个批次的实际长度：rowEnds[batch_id]-rowStarts[batch_id]
    stride0: int = -1,
    stride1: int = 1,
) -> None:
    pass
