Function split

tensorrt_llm/_torch/utils.py:394–402

(x: torch.Tensor,
          tp_size: int,
          idx: int,
          dim: int = 0)

def split(x: torch.Tensor,
          tp_size: int,
          idx: int,
          dim: int = 0) -> torch.Tensor:
    assert x.shape[dim] % tp_size == 0
    split_size = x.shape[dim] // tp_size
    if tp_size == 1:
        return x
    return torch.split(x, split_size, dim=dim)[idx]


def relu2(x: torch.Tensor) -> torch.Tensor: