import torch


def combine_variance(avg_a, count_a, var_a, avg_b, count_b, var_b):
    """
    Compute variance of X given mean and variances of A and B, where X = A union B.
    Reference: https://en.wikipedia.org/wiki/Algorithms_for_calculating_variance#cite_note-:0-10
    """
    if count_a + count_b <= 1:
        return torch.zeros(var_a.size()).cuda()
    delta = avg_b - avg_a
    M2_a = var_a * (count_a - 1)
    M2_b = var_b * (count_b - 1)
    M2 = M2_a + M2_b + delta ** 2 * count_a * count_b / (count_a + count_b)
    return M2 / (count_a + count_b - 1)