import torch


def gaussian_likelihood(x, mu, log_std):
    """
    Args:
        x: Tensor with shape [batch, dim]
        mu: Tensor with shape [batch, dim]
        log_std: Tensor with shape [batch, dim] or [dim]

    Returns:
        Tensor with shape [batch]
    """
    m = torch.distributions.normal.Normal(mu, torch.exp(log_std))
    return m.log_prob(x).sum(axis=-1)