import torch


def split(input_, indices_or_sections, axis=0):
    """Wrapper of `torch.split`.

    Parameters
    ----------
    input_ : DTensor
        Input dense tensor.
    indices_or_sections : int or DTensor
        If integer, specifing part of splits. Else specifing the positions
        where the splits take place.
    axis : int, optional
        The axis to operate on, by default 0
    """
    return torch.split(input_._data, indices_or_sections, axis)