bucketize-cdf.py

def abs_cdf(t: Tensor, thresholds: list[float]):
    t = t.abs()
    level = torch.bucketize(t, t.new_tensor(thresholds), out_int32=True) # sum(x > v for v in thresholds)
    return level.flatten().bincount(minlength=len(thresholds) + 1).cumsum(0) / t.numel()

# reference: https://github.com/pytorch/pytorch/issues/69519#issuecomment-2500366519
def histogram(input: Tensor, bins: Tensor, *, weight: Optional[Tensor] = None, density: bool = False):
    bucket_indices = torch.bucketize(input, bins)
    counts = torch.bincount(bucket_indices, weights=weight, minlength=bins.size(0)+1)
    counts = counts[1:-1]
    # Processing the last bin right border
    if input.dtype == torch.int:
        counts[-1] += torch.sum(input == bins[-1].item())
    else:
        counts[-1] += torch.sum(torch.isclose(input, bins[-1].item()))
    width = bins[1:] - bins[:-1]
    if not density:
        out = (counts, bins)
    else:
        density = counts / width / counts.sum()
        out = (density, bins)
    return out