import torch


def gaussian_log_prob(x: torch.Tensor, gm: dict) -> torch.Tensor:
    """calculate the Gaussian log-likelihood of data x

    Args:
        x: tensor of size (N, 1, d)
        gm: mean and cov of sizes (K, d) and (K, d, d)
    
    Return: N-by-K matrix of log-likelihood
    """
    if x.dim() == 2:
        x = torch.unsqueeze(input=x, dim=1)

    if 'scale_tril' in gm.keys():
        mvn = torch.distributions.multivariate_normal.MultivariateNormal(
            loc=gm['mean'],
            scale_tril=gm['scale_tril']
        )
    elif 'cov' in gm.keys():
        mvn = torch.distributions.multivariate_normal.MultivariateNormal(
            loc=gm['mean'],
            covariance_matrix=gm['cov']
        )
    else:
        raise ValueError('Missing scale_tril or cov')

    return mvn.log_prob(value=x)