import torch


def unique(input_, return_inverse=False, return_counts=False, axis=None):
    """Wrapper of `torch.unique`.

    Parameters
    ----------
    input_ : DTensor
        Input dense tensor.
    return_inverse : bool, optional
        If True, return the indices also, by default False
    return_counts : bool, optional
        If True, return the count also, by default False
    axis : int or None, optional
        The axis to operate on, by default None
    """
    return torch.unique(input_._data, return_inverse=return_inverse,
                        return_counts=return_counts, dim=axis)