import torch


def gaussian_mixture_moments(mus, sigma_sqs):
    """Estimate moments of a gaussian mixture model

    B - number of observations/samples
    N - number of components in mixture

    Args:
        mus torch.tensor((B, N)): Collection of mu-values
        sigma_sqs torch.tensor((B, N)): Collection of sigma_sq-values
    """

    with torch.no_grad():
        mu = torch.mean(mus, dim=1)
        sigma_sq = torch.mean(sigma_sqs + mus**2, dim=1) - mu**2

    return mu, sigma_sq