torch-ext/fused_split/__init__.py

from typing import Optional

import torch

from ._ops import ops


def fused_split(x: torch.Tensor, out: Optional[torch.Tensor] = None) -> torch.Tensor:
    if out is None:
        out = torch.empty_like(x)
    ops.fused_split(out, x)
    return out