def sample_portion(vec, p=0.5):
    """
    Subsamples a fraction (given by p) of the given batch.  Used by
    Quantile when the data gets very very large.
    """
    bits = torch.bernoulli(
        torch.zeros(vec.shape[0], dtype=torch.uint8, device=vec.device), p
    )
    return vec[bits]

    Subsamples a fraction (given by p) of the given batch.  Used by
    Quantile when the data gets very very large.
    
util/runningstats.py