import torch


def diagonal(input_, offset=0, axis1=0, axis2=1):
    """Wrapper of `torch.diagonal`.

    Parameters
    ----------
    input_ : DTensor or STensor
        Input tensor.
    offset : int, optional
        Offset to the main diagonal, by default 0
    axis1 : int, optional
        The first axis of the diagonal, by default 0
    axis2 : int, optional
        The second axis of the diagonal, by default 1
    """
    return torch.diagonal(input_._data, offset=offset, dim1=axis1, dim2=axis2)