"""Operators for pipeline parallelism."""

from typing import List  # noqa: UP035

from tvm import relax
from tvm.relax.frontend.nn import Tensor, op


def pipeline_stage_boundary(*tensors: Tensor) -> List[Tensor]:  # noqa: UP006
    """Pipeline parallelism stage boundary mark operator in MLC.

    Parameters
    ----------
    tensors : Tensor
        The tensors to be passed to the next stage.

    Returns
    -------
    tensors : List[Tensor]
        The list of input tensors passed to the next stage.
    """
    return op.wrap_nested(
        relax.call_pure_packed(
            "mlc.pipeline_parallel_stage_boundary",
            *[tensor._expr for tensor in tensors],
            ty_args=(
                tensors[0]._expr.ty
                if len(tensors) == 1
                else relax.TupleType([tensor._expr.ty for tensor in tensors])
            ),
        ),
        name="pipeline_stage_boundary",
    )