import torch


def hessian_sym_matmul(hess, grad):
    """Matrix-multiplication for small batches of symmetric matrices.

    `>>> hess @ grad`

    `hess` contains only the diagonal and upper part of the matrix, in
    a flattened array. Elements are ordered as:
     `[(i, i) for i in range(P)] +
      [(i, j) for i in range(P) for j in range(i+1, P)]

    Parameters
    ----------
    hess : (P*(P+1)//2, ...) tensor
    grad : (P, ...) tensor

    Returns
    -------
    mm : (P, ...) tensor

    """
    nb_prm = len(grad)
    if nb_prm == 1:
        return hess * grad
    elif nb_prm == 2:
        mm = torch.empty_like(grad)
        mm[0] = hess[0] * grad[0] + hess[2] * grad[1]
        mm[1] = hess[1] * grad[1] + hess[2] * grad[0]
        return mm
    elif nb_prm == 3:
        mm = torch.empty_like(grad)
        mm[0] = hess[0] * grad[0] + hess[3] * grad[1] + hess[4] * grad[2]
        mm[1] = hess[1] * grad[1] + hess[3] * grad[0] + hess[5] * grad[2]
        mm[2] = hess[2] * grad[2] + hess[4] * grad[0] + hess[5] * grad[1]
        return mm
    elif nb_prm == 4:
        mm = torch.empty_like(grad)
        mm[0] = hess[0] * grad[0]
        mm[0] += hess[4] * grad[1]
        mm[0] += hess[5] * grad[2]
        mm[0] += hess[6] * grad[3]
        mm[1] = hess[1] * grad[1]
        mm[1] += hess[4] * grad[0]
        mm[1] += hess[7] * grad[2]
        mm[1] += hess[8] * grad[3]
        mm[2] = hess[2] * grad[2]
        mm[2] += hess[5] * grad[0]
        mm[2] += hess[7] * grad[1]
        mm[2] += hess[9] * grad[3]
        mm[3] = hess[3] * grad[3]
        mm[3] += hess[6] * grad[0]
        mm[3] += hess[8] * grad[1]
        mm[3] += hess[9] * grad[2]
        return mm
    else:
        mm = torch.empty_like(grad)
        for i in range(nb_prm):
            mm[i] = hess[i] * grad[i]
        c = nb_prm
        for i in range(nb_prm):
            for j in range(i+1, nb_prm):
                mm[i] += hess[c] * grad[j]
                mm[j] += hess[c] * grad[i]
                c += 1
        return mm