Compute psi-statistics for ss-linear kernel

def psicomputations(variance, Z, variational_posterior, return_psi2_n=False):
    """
    Compute psi-statistics for ss-linear kernel
    """
    # here are the "statistics" for psi0, psi1 and psi2
    # Produced intermediate results:
    # psi0    N
    # psi1    NxM
    # psi2    MxM
    mu = variational_posterior.mean
    S = variational_posterior.variance

    psi0 = (variance*(np.square(mu)+S)).sum(axis=1)
    Zv = variance * Z
    psi1 = np.dot(mu,Zv.T)
    if return_psi2_n:
        psi2 = psi1[:,:,None] * psi1[:,None,:] + np.dot(S[:,None,:] * Zv[None,:,:], Zv.T)
    else:
        psi2 = np.dot(S.sum(axis=0) * Zv, Zv.T) + tdot(psi1.T)

    return psi0, psi1, psi2