import torch


def pool_stds(ns, means, stds):
    """Pool standard deviations from multiple groups. We cannot simply take a
    weighted average like for means.

    Formulas and some notation is from:
    http://www.obg.cuhk.edu.hk/ResearchSupport/StatTools/CombineMeansSDs_Pgm.php

    Parameters
    ----------
    ns: list[int] or torch.LongTensor
        Size of each group. Must be in same order as other args.
    means: list[float] or torch.FloatTensor
        Mean of each group. Must be in same order as other args.
    stds: list[float] or torch.FloatTensor
        Standard deviation of each group. Must be in same order as other args.

    Returns
    -------
    torch.FloatTensor: pooled standard deviation across all groups.
    """
    ns, means, stds = map(torch.tensor, [ns, means, stds])
    n_total = ns.sum()
    totals = means * ns
    tx = totals.sum()
    txx = torch.sum(stds.pow(2).mul(ns - 1) + totals.pow(2).div(ns))
    return torch.sqrt((txx - tx.pow(2).div(n_total)) / (n_total - 1))