import torch
import math


def sym_diag(mat):
    """Diagonal of a symmetric matrix

    Parameters
    ----------
    mat : (..., M * (M+1) // 2) tensor
        A symmetric matrix that is stored in a sparse way.
        Its elements along the last (flat) dimension are the
        diagonal elements followed by the flattened upper-half elements.
        E.g., [a00, a11, aa22, a01, a02, a12]

    Returns
    -------
    diag : (..., M) tensor
        Main diagonal of the matrix

    """
    mat = torch.as_tensor(mat)
    nb_prm = int((math.sqrt(1 + 8 * mat.shape[-1]) - 1)//2)
    return mat[..., :nb_prm]