Compute Quantile

quantile

FUNCTION	DESCRIPTION
quantile	Compute quantiles along a specified dimension, optionally weighted.

Functions

quantile(X: torch.Tensor, q: torch.Tensor | float, weights: torch.Tensor | None = None, dim: int | None = None, interpolation: Literal['linear', 'nearest', 'higher', 'lower', 'midpoint'] = 'linear', keepdim: bool = False) -> torch.Tensor

Compute quantiles along a specified dimension, optionally weighted.

Uses torch.sort for efficiency and handles multiple quantiles simultaneously. Overcomes limitations of torch.quantile which is restricted to 2^24 elements.

When weights are provided, computes weighted quantiles using midpoint CDF interpolation and torch.searchsorted.

PARAMETER	DESCRIPTION
X	Data tensor. TYPE: Tensor
q	Quantile(s) to compute. Must be in [0, 1]. Scalar values are supported and will squeeze the quantile dimension from the output. TYPE: Tensor | float
weights	Non-negative sample weights, same shape as X (or broadcastable along dim). If None, computes unweighted quantiles. TYPE: torch.Tensor | None, optional, by default None DEFAULT: None
dim	Dimension to reduce. If None, the tensor is flattened. TYPE: int | None, optional, by default None DEFAULT: None
interpolation	Interpolation method, by default "linear". Only "linear" is supported when weights are provided. TYPE: (linear, nearest, higher, lower, midpoint) DEFAULT: "linear"
keepdim	Whether to keep the reduced dimension in the output. TYPE: bool, optional, by default False DEFAULT: False

RETURNS	DESCRIPTION
Tensor	Quantile values. Shape depends on q, dim, and keepdim.

Examples:

>>> import torch
>>> X = torch.randn(100, 5)
>>> q = torch.tensor([0.25, 0.5, 0.75])
>>> quantile(X, q, dim=0).shape
torch.Size([3, 5])

Weighted quantiles:

>>> w = torch.rand(100, 5)
>>> quantile(X, q, weights=w, dim=0).shape
torch.Size([3, 5])

References

.. [1] https://discuss.pytorch.org/t/efficient-quantile-k-largest-value/ .. [2] https://github.com/pytorch/pytorch/issues/157431 .. [3] https://github.com/pytorch/pytorch/issues/64947

Source code in spectre/compute/quantile.py

def quantile(
    X: torch.Tensor,
    q: torch.Tensor | float,
    weights: torch.Tensor | None = None,
    dim: int | None = None,
    interpolation: Literal[
        "linear", "nearest", "higher", "lower", "midpoint"
    ] = "linear",
    keepdim: bool = False,
) -> torch.Tensor:
    """
    Compute quantiles along a specified dimension, optionally weighted.

    Uses `torch.sort` for efficiency and handles multiple quantiles
    simultaneously. Overcomes limitations of `torch.quantile` which is
    restricted to 2^24 elements.

    When `weights` are provided, computes weighted quantiles using midpoint CDF
    interpolation and `torch.searchsorted`.

    Parameters
    ----------
    X : torch.Tensor
        Data tensor.

    q : torch.Tensor | float
        Quantile(s) to compute. Must be in [0, 1]. Scalar values are supported
        and will squeeze the quantile dimension from the output.

    weights : torch.Tensor | None, optional, by default None
        Non-negative sample weights, same shape as `X` (or broadcastable along
        `dim`). If None, computes unweighted quantiles.

    dim : int | None, optional, by default None
        Dimension to reduce. If None, the tensor is flattened.

    interpolation : {"linear", "nearest", "higher", "lower", "midpoint"}
        Interpolation method, by default "linear". Only ``"linear"`` is
        supported when `weights` are provided.

    keepdim : bool, optional, by default False
        Whether to keep the reduced dimension in the output.

    Returns
    -------
    torch.Tensor
        Quantile values. Shape depends on `q`, `dim`, and `keepdim`.

    Examples
    --------
    >>> import torch
    >>> X = torch.randn(100, 5)
    >>> q = torch.tensor([0.25, 0.5, 0.75])
    >>> quantile(X, q, dim=0).shape
    torch.Size([3, 5])

    Weighted quantiles:

    >>> w = torch.rand(100, 5)
    >>> quantile(X, q, weights=w, dim=0).shape
    torch.Size([3, 5])

    References
    ----------
    .. [1] https://discuss.pytorch.org/t/efficient-quantile-k-largest-value/
    .. [2] https://github.com/pytorch/pytorch/issues/157431
    .. [3] https://github.com/pytorch/pytorch/issues/64947
    """
    if weights is not None:
        if interpolation != "linear":
            raise ValueError(
                f"Only 'linear' interpolation is supported with weights, "
                f"got '{interpolation}'."
            )

    if isinstance(q, (int, float)):
        q = torch.tensor([q], dtype=X.dtype, device=X.device)
    else:
        q = q.to(dtype=X.dtype, device=X.device)

    if not torch.all((0 <= q) & (q <= 1)):
        raise ValueError("Quantiles `q` must be in [0, 1].")

    # Sanitize `dim` since we cannot pass `dim=None` to `torch.index_select`.
    if dim_was_none := dim is None:
        dim = 0
        X = X.flatten()
        if weights is not None:
            weights = weights.flatten()

    if weights is None:
        out = _unweighted_quantile(
            X, q, dim=dim, interpolation=interpolation, dim_was_none=dim_was_none
        )
    else:
        out = _weighted_quantile(X, q, weights=weights, dim=dim)

    len_q = len(q)
    if dim_was_none:
        if keepdim:
            out = out.view(len_q, 1)
        else:
            out = out.view(len_q) if len_q > 1 else out.squeeze()
    else:
        if not keepdim and len_q == 1:
            out = out.squeeze(dim)

    return out