import torch


def sample_mean_and_cov(sample, diag=False):
    """Compute mean and covariance of a sample of vectors.

    Parameters
    ----------
    sample: Tensor
        Tensor of dimensions [batch x N x num_samples].
    diag: bool, optional.
        Flag to indicate if the computation has to assume independent or correlated
        variables.

    Returns
    -------
    mean: Tensor
        Tensor of dimension [batch x N].
    covariance: Tensor
        Tensor of dimension [batch x N x N].

    """
    num_samples = sample.shape[-1]
    mean = torch.mean(sample, dim=-1, keepdim=True)

    if diag:
        covariance = torch.diag_embed(sample.var(-1))
    else:
        sigma = (mean - sample) @ (mean - sample).transpose(-2, -1)
        sigma += 1e-6 * torch.eye(sigma.shape[-1])  # Add some jitter.
        covariance = sigma / num_samples
    mean = mean.squeeze(-1)

    return mean, covariance