API Reference¶

class BezierCDF(Distribution):
    r"""A continuous probability distribution parameterized by Bernstein polynomials with custom constraints.

    The idea is that the CDF is represented as a Bezier curve, which is a weighted sum of Bernstein basis polynomials,
    defined by control points (betas) that are derived from the input logits.
    This allows for a smooth, flexible CDF that can capture complex shapes while still being differentiable.
    In fact, this formulation is mathematically equivalent to a mixture of Beta distributions, where the mixture
    weights are given by the deltas (softmax of the logits) and the Beta components are defined by the control points.

    Since we know that any CDF must start at 0 and end at 1, we can enforce these constraints by fixing the first
    control point to 0 and the last control point to 1.

    The spacing of the control points along the domain-axis ("x-axis") is strictly uniform and determined by the
    degree of the Bernstein polynomial, hence, number of input logits.

    Note:
        Bernstein polynomials converge slowly: the worst-case pointwise approximation error is $O(1/n)$ where $n$ is
        the polynomial degree, leading to a standard deviation error of $O(1/\sqrt{n})$. However, for smooth CDFs the
        effective rate is better, and Bernstein density estimators achieve the optimal minimax rate (Babu et al., 2002;
        Petrone, 1999). This slower convergence is an inherent trade-off for the structural guarantees they provide:
        monotonicity, values in $[0, 1]$, non-negative PDF, and an unconstrained parameterization (any real-valued
        logits yield a valid distribution). No other polynomial basis offers all of these simultaneously. In practice,
        the bias matters less when logits are learned end-to-end via gradient descent, as the optimizer can compensate.

        The sharpest peak a degree-n Bernstein polynomial can produce is a single Beta component with
        $std \approx 1/(2\sqrt{n})$ in [0,1]-space. Scaled to support range R, the peak std is $R / (2\sqrt{n})$.
    """

    has_rsample = True

    def __init__(
        self,
        logits: torch.Tensor,
        bound_low: float = -1e3,
        bound_up: float = 1e3,
        normalization_method: Literal["sigmoid", "softmax"] = "softmax",
        validate_args: bool | None = None,
    ) -> None:
        """Initializer.

        Args:
            logits: Raw logits for the probabilities before normalization, of shape (*batch_shape, degree).
                The logits also determine the degree of the Bernstein polynomial $n$.
            bound_low: Lower bound of the distribution support, needs to be finite.
            bound_up: Upper bound of the distribution support, needs to be finite.
            normalization_method: How to normalize the probabilities. Either "sigmoid" or "softmax". With "sigmoid",
                each control point is independently activated, while with "softmax", the control point activations
                influence each other.
            validate_args: Whether to validate arguments. Carried over to keep the interface with the base class.
        """
        self.logits = logits
        self.bound_low = bound_low
        self.bound_up = bound_up
        self.normalization_method = normalization_method

        # Precompute binomial coefficients, and store them on the same device as logits.
        self._binom_coeffs_cdf, self._binom_coeffs_pdf = self._compute_binomial_coefficients()

        # Precompute log-space binomial coefficients for numerically stable log_prob.
        self._log_binom_coeffs_pdf = self._binom_coeffs_pdf.log()

        # Calculate parameters (deltas and betas).
        self._deltas, self._betas, self._log_deltas = self._compute_deltas_and_betas()

        # Determine batch shape based on the logits. The event shape is scalar since this is a univariate distribution.
        super().__init__(batch_shape=logits.shape[:-1], event_shape=torch.Size([]), validate_args=validate_args)

    def __repr__(self) -> str:
        """String representation of the distribution."""
        return (
            f"{self.__class__.__name__}(logits_shape: {self.logits.shape}, bound_low: {self.bound_low}, "
            f"bound_up: {self.bound_up}, normalization_method: {self.normalization_method})"
        )

    def _compute_binomial_coefficients(self) -> tuple[torch.Tensor, torch.Tensor]:
        """Compute the binomial coefficients for the CDF and PDF based on the degree of the Bernstein polynomial.

        comb(n, k) = n! / (k! * (n-k)!) is the binomial coefficient, which counts the number of ways to choose k
        elements from a set of n elements.

        Returns:
            coeffs_cdf: Binomial coefficients for the CDF, of shape (degree + 1,)
            coeffs_pdf: Binomial coefficients for the PDF, of shape (degree,)
        """
        coeffs_cdf = torch.tensor(
            [math.comb(self.degree, i) for i in range(self.degree + 1)],
            device=self.logits.device,
            dtype=self.logits.dtype,
        )

        coeffs_pdf = torch.tensor(
            [math.comb(self.degree - 1, i) for i in range(self.degree)],
            device=self.logits.device,
            dtype=self.logits.dtype,
        )

        # Check if any of the binomial coefficients became infinite.
        if torch.isinf(coeffs_cdf).any() or torch.isinf(coeffs_pdf).any():
            raise ValueError(
                f"Binomial coefficients became infinite for degree {self.degree}. "
                "Consider reducing the (last) dimension of the logits, leading to lower degree polynomial."
            )

        return coeffs_cdf, coeffs_pdf

    def _compute_deltas_and_betas(self) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        r"""Compute the deltas (Beta mixture component weights) and betas (control points) for the Bezier curve based
        on the given logits.

        The deltas are the forward differences of the betas, i.e., $ \Delta_i = \beta_{i + 1} - \beta_i $.

        Returns:
            deltas: Weights of the Beta components in the mixture, of shape (*batch_shape, degree)
            betas: Control points of the Bezier curve, of shape (*batch_shape, degree + 1)
            log_deltas: Log of the deltas, computed in a numerically stable way, of shape (*batch_shape, degree)
        """
        # The deltas are the steps themselves (forward differences of betas).
        if self.normalization_method == "softmax":
            deltas = torch.softmax(self.logits, dim=-1)  # shape: (*batch_shape, degree)
            log_deltas = torch.log_softmax(self.logits, dim=-1)

        elif self.normalization_method == "sigmoid":
            raw_deltas = torch.sigmoid(self.logits)
            sum_deltas = raw_deltas.sum(dim=-1, keepdim=True)

            # Prevent division by zero in the rare case where all logits are massively negative.
            eps = torch.finfo(raw_deltas.dtype).eps
            sum_deltas = sum_deltas.clamp_min(eps)

            deltas = raw_deltas / sum_deltas  # shape: (*batch_shape, degree)

            # log(Delta) = log(sigmoid(x) / sum(sigmoid(x))) = logsigmoid(x) - log(sum(sigmoid(x))).
            log_deltas = torch.nn.functional.logsigmoid(self.logits) - sum_deltas.log()

        else:
            raise ValueError(f"Unknown normalization method: {self.normalization_method}")

        # Pad with zeros and ones to enforce the CDF boundary conditions:
        # betas = [0, beta_1, ..., beta_{n-1}, beta_n = 1]
        zeros = torch.zeros(*deltas.shape[:-1], 1, device=deltas.device, dtype=deltas.dtype)  # shape: (*batch_shape, 1)
        inner_betas = torch.cumsum(deltas, dim=-1)[..., :-1]
        ones = torch.ones(*deltas.shape[:-1], 1, device=deltas.device, dtype=deltas.dtype)
        betas = torch.cat([zeros, inner_betas, ones], dim=-1)

        return deltas, betas, log_deltas

    def _map_to_t_space(self, value: torch.Tensor) -> torch.Tensor:
        r"""Map values from the original $X$ space to the $T$ space $[0, 1]$ using the bounds."""
        return torch.clamp((value - self.bound_low) / self.support_range, 0, 1)

    def _map_to_x_space(self, t: torch.Tensor) -> torch.Tensor:
        r"""Map values from the $T$ space $[0, 1]$ back to the original $X$ space using the bounds."""
        return t * self.support_range + self.bound_low

    @property
    def support(self) -> constraints.Constraint:
        """Support of this distribution."""
        return constraints.interval(self.bound_low, self.bound_up)

    @property
    def support_range(self) -> float:
        """Range of the support, i.e., upper bound - lower bound."""
        return self.bound_up - self.bound_low

    @property
    def arg_constraints(self) -> dict[str, constraints.Constraint]:
        """Constraints that should be satisfied by each argument of this distribution. None for this class."""
        return {"logits": constraints.real}

    @property
    def degree(self) -> int:
        r"""Get the degree $n$ of the Bernstein polynomial based on the number of logits.

        For a Bernstein polynomial of degree $n$, there are $n + 1$ control points (betas) and $n$ weights (deltas).
        """
        return self.logits.shape[-1]

    @property
    def mean(self) -> torch.Tensor:
        r"""Compute mean of the distribution, i.e., the weighted average of the control points.

        We transform the random variable $X$ to $T$ in [0, 1] by scaling and shifting according to the bounds.
        Then, the mean of $T$ can be computed as

        $$ E[T] = \sum_{i=0}^{n-1} \Delta_i \frac{i+1}{n+1} $$

        where $\Delta_i$ is the weight of the $i$-th control point, and $n$ is the degree of the Bernstein polynomial.
        We can then get the mean by rescaling $E[T]$ back to the original support:

        $$ E[X] = (U - L) E[T] + L $$

        where $L$ and $U$ are the lower and upper bounds of the distribution support, respectively.

        Note:
            This method uses the exact Beta mixture formula.

        Returns:
            Tensor of shape (*batch_shape,).
        """
        i = torch.arange(self.degree, device=self._deltas.device, dtype=self._deltas.dtype)  # shape: (degree,)
        e_t = torch.sum(self._deltas * (i + 1) / (self.degree + 1), dim=-1)

        return self._map_to_x_space(e_t)

    @property
    def variance(self) -> torch.Tensor:
        r"""Compute variance of the distribution.

        We transform the random variable $X$ to $T$ in [0, 1] by scaling and shifting according to the bounds.
        Then, the variance of $T$ can be computed as

        $$ Var[T] = E[T^2] - (E[T])^2 $$

        with

        $$ E[T^2] = \sum_{i=0}^{n-1} \Delta_i \frac{(i+1)(i+2)}{(n+1)(n+2)} $$

        where $\Delta_i$ is the weight of the $i$-th control point, and $n$ is the degree of the Bernstein polynomial.
        We can then get the variance by rescaling $Var[T]$ back to the original support:

        $$ Var[X] = (U - L)^2 Var[T] $$

        Note:
            This method uses the exact Beta mixture formula.

        Returns:
            Tensor of shape (*batch_shape,).
        """
        i = torch.arange(self.degree, device=self._deltas.device, dtype=self._deltas.dtype)  # shape: (degree,)
        e_t = torch.sum(self._deltas * (i + 1) / (self.degree + 1), dim=-1)
        e_t2 = torch.sum(self._deltas * ((i + 1) * (i + 2)) / ((self.degree + 1) * (self.degree + 2)), dim=-1)
        var_t = e_t2 - e_t**2

        return self.support_range**2 * var_t

    def _eval_bezier_curve(
        self,
        t: torch.Tensor,
        weights: torch.Tensor,
        binom_coeffs: torch.Tensor,
    ) -> torch.Tensor:
        r"""Evaluates a Bezier curve (a Bernstein polynomial) in the $T \in [0, 1]$ space.

        This method computes the weighted sum of Bernstein basis polynomials. Let $d$ be the degree of the polynomial
        being evaluated (either $n$ or $n+1$). Each basis polynomial is defined as:

        $$ B_{i, d}(t) = \binom{d}{i} t^i (1-t)^{d-i} $$

        The polynomial's value $p(t)$ is computed as:

        $$ p(t) = \sum_{i=0}^{d} w_i B_{i, d}(t) $$

        where $w_i$ are the weights (either betas for the CDF or deltas for the PDF).

        Args:
            t: Normalized input values in [0, 1].
                Expected shape: (*sample_shape, *batch_shape).
            weights: The coefficients for the basis polynomials.
                Expected shape: (*batch_shape, d + 1).
            binom_coeffs: Precomputed binomial coefficients corresponding to the polynomial's degree.
                Expected shape: (d + 1,).

        Returns:
            The evaluated polynomial values.
            Output shape: (*sample_shape, *batch_shape)
        """
        # Get n which can be != self.degree as we use this method for both CDF and PDF which have different degrees.
        nun_coeffs = binom_coeffs.shape[0]

        # Create a tensor of indices matching the number of basis polynomials.
        i = torch.arange(nun_coeffs, device=t.device, dtype=t.dtype)

        # Add an empty dimension to t for broadcasting, resulting in shape: (*sample_shape, *batch_shape, 1).
        t_expanded = t.unsqueeze(-1)

        # Compute the entire basis in one shot.
        # PyTorch broadcasts the shapes to shape (*sample_shape, *batch_shape, degree).
        basis = binom_coeffs * (t_expanded**i) * ((1 - t_expanded) ** (nun_coeffs - 1 - i))

        # Multiply by weights and sum across the final dimension, resulting in shape (*sample_shape, *batch_shape).
        return torch.sum(weights * basis, dim=-1)

    def cdf(self, value: torch.Tensor) -> torch.Tensor:
        """Compute cumulative distribution function at given values.

        Args:
            value: Values at which to compute the CDF. Expected shape: (*sample_shape, *batch_shape).

        Returns:
            CDF values in [0, 1] corresponding to the input values. Output shape: same as `value` argument.
        """
        x = value.to(device=self.logits.device, dtype=self.logits.dtype)

        # Map X in [bound_low, bound_up] to T in [0, 1].
        t = self._map_to_t_space(x)

        # Construct and evaluate the Bezier curve in T space.
        return self._eval_bezier_curve(t, weights=self._betas, binom_coeffs=self._binom_coeffs_cdf)

    def prob(self, value: torch.Tensor) -> torch.Tensor:
        """Compute probability density at given values.

        Args:
            value: Values at which to compute the PDF. Expected shape: (*sample_shape, *batch_shape).

        Returns:
            PDF values corresponding to the input values. Output shape: same as `value` argument.
        """
        x = value.to(device=self.logits.device, dtype=self.logits.dtype)

        # Map X in [bound_low, bound_up] to T in [0, 1].
        t = self._map_to_t_space(x)

        # Construct and evaluate the Bezier curve in T space.
        val = self._eval_bezier_curve(t, weights=self._deltas, binom_coeffs=self._binom_coeffs_pdf)

        # Apply the chain rule: dt/dx = 1 / (U - L).
        pdf_val = val * self.degree / self.support_range

        # Mask out values outside [bound_low, bound_up].
        mask = (value >= self.bound_low) & (value <= self.bound_up)
        return torch.where(mask, pdf_val, torch.zeros_like(pdf_val))

    def log_prob(self, value: torch.Tensor) -> torch.Tensor:
        r"""Compute the log-probability density at given values, entirely in log-space for numerical stability.

        Uses the identity

        $$
        \log p(x) = \log \frac{n}{U - L}
                    + \text{logsumexp}_i\!\Big(\log \Delta_i + \log \binom{n-1}{i}
                    + i \log t + (n-1-i) \log(1-t)\Big)
        $$

        where $t = (x - L) / (U - L)$ is the normalized input. Every term is computed in log-space,
        avoiding the numerically problematic ``log(polynomial + eps)`` path.

        Args:
            value: Values at which to compute the log-PDF. Expected shape: (*sample_shape, *batch_shape).

        Returns:
            Log-PDF values corresponding to the input values. Output shape: same as `value` argument.
        """
        x = value.to(device=self.logits.device, dtype=self.logits.dtype)
        t = self._map_to_t_space(x)

        eps = torch.finfo(t.dtype).eps
        n = self.degree

        # Clamp t away from exact 0/1 to avoid log(0).
        t_safe = t.clamp(min=eps, max=1 - eps)

        # Indices for the Bernstein basis: i = 0, ..., n-1.
        i = torch.arange(n, device=t.device, dtype=t.dtype)

        # Expand t for broadcasting: (*sample_shape, *batch_shape, 1).
        log_t = t_safe.unsqueeze(-1)  # will broadcast with i

        # Log of each Bernstein basis term: log(binom) + i*log(t) + (n-1-i)*log(1-t).
        log_basis = self._log_binom_coeffs_pdf + i * log_t.log() + (n - 1 - i) * (1 - log_t).log()

        # Log of each weighted term: log(delta_i) + log(basis_i).
        # _log_deltas shape: (*batch_shape, n),  log_basis shape: (*sample_shape, *batch_shape, n).
        log_terms = self._log_deltas + log_basis

        # Sum via logsumexp over the last dimension.
        log_bezier = torch.logsumexp(log_terms, dim=-1)

        # Apply the chain rule: log(n / (U - L)) + log(bezier).
        log_pdf = math.log(n / self.support_range) + log_bezier

        # Mask values outside the support.
        mask = (value >= self.bound_low) & (value <= self.bound_up)
        return torch.where(mask, log_pdf, torch.full_like(log_pdf, -math.inf))

    def entropy(self, num_quadrature_points: int = 251) -> torch.Tensor:
        r"""Compute differential entropy of the distribution via numerical quadrature.

        $$ H(X) = -\int_{L}^{U} p(x) \log p(x) \, dx $$

        where $L$ and $U$ are the lower and upper bounds of the distribution support, respectively.

        Args:
            num_quadrature_points: Number of points for the trapezoidal rule approximation.

        Returns:
            Tensor of shape (*batch_shape,).
        """
        # Create quadrature points over the support.
        x = torch.linspace(
            self.bound_low, self.bound_up, num_quadrature_points, device=self.logits.device, dtype=self.logits.dtype
        )

        # For batched distributions, expand quadrature points to shape (num_quadrature_points, *batch_shape)
        # so prob/log_prob receive values with explicit batch dimensions.
        x_eval = x.reshape(num_quadrature_points, *([1] * len(self.batch_shape)))
        x_eval = x_eval.expand(num_quadrature_points, *self.batch_shape)

        # Evaluate PDF at quadrature points.
        pdf_val = self.prob(x_eval)  # shape: (num_quadrature_points, *batch_shape)

        # Compute the integrand: -p(x) * log(p(x)), with epsilon for stability.
        eps = torch.finfo(pdf_val.dtype).eps
        log_pdf = torch.log(pdf_val + 2 * eps)
        integrand = -pdf_val * log_pdf  # shape: (num_quadrature_points, *batch_shape)

        # Integrate using the trapezoidal rule.
        return torch.trapezoid(integrand, x, dim=0)

    def icdf(
        self,
        value: torch.Tensor,
        num_iter: int = 8,
        use_newton: bool = True,
        newton_damping: float = 0.9,
        convergence_eps_factor: float = 20.0,
    ) -> torch.Tensor:
        r"""Compute the inverse CDF, i.e., the quantile function, at the given values.

        Two solvers are available for inverting $ F(x) - q = 0 $:

        **Newton's method** uses the PDF as the exact derivative of the CDF and iterates

        $$ x_{k+1} = x_k - \alpha \frac{F(x_k) - q}{f(x_k)} $$

        where $F(x)$ is the CDF, $f(x)$ is the PDF, $q$ is the target quantile in [0, 1],
        and $\alpha \in (0, 1]$ is a damping factor that shrinks each Newton step to improve robustness.
        A bracket $[L_k, U_k]$ is maintained alongside: whenever $F(x_k) < q$ the lower bound tightens,
        otherwise the upper bound tightens. If the Newton step would leave the bracket, a bisection
        step is used instead, guaranteeing monotonic bracket contraction and preventing oscillation.
        The loop exits early once all elements satisfy $|F(x) - q| < \epsilon$.

        **Bisection** halves the search interval each iteration, gaining ~1 bit of precision per step.

        Args:
            value: Values in [0, 1] at which to compute the inverse CDF. Expected shape: (*sample_shape, *batch_shape).
            num_iter: Maximum number of solver iterations. Newton typically converges undamped in ~6-7 iterations;
                bisection needs ~15-20 for full float32 precision.
            use_newton: If True, use Newton's method. If False, use pure bisection.
            newton_damping: Damping factor in (0, 1] applied to the Newton step. A value of 1.0 gives the
                full Newton step (quadratic convergence), while smaller values improve robustness
                at the cost of slower convergence.
            convergence_eps_factor: The factor multiplied by machine epsilon to determine the convergence criterion.

        Returns:
            Quantiles in [bound_low, bound_up] corresponding to the input CDF values.
            Output shape: same as `value` argument.
        """
        q = value.to(device=self.logits.device, dtype=self.logits.dtype)
        eps = torch.finfo(q.dtype).eps

        # Ensure target probability value is strictly in [0, 1].
        q = torch.clamp(q, 0.0, 1.0)

        # Start from the midpoint of the support.
        mid = torch.full_like(q, (self.bound_low + self.bound_up) / 2)
        low = torch.full_like(q, self.bound_low)
        high = torch.full_like(q, self.bound_up)

        for _ in range(num_iter):
            cdf_mid = self.cdf(mid)

            # Early stop when all elements have converged.
            abs_deviation = (cdf_mid - q).abs().max()
            if abs_deviation < convergence_eps_factor * eps:
                break

            # Tighten the bracket based on CDF evaluation.
            low = torch.where(cdf_mid < q, mid, low)
            high = torch.where(cdf_mid >= q, mid, high)
            bisect_mid = (low + high) / 2

            if use_newton:
                # Newton step: x_{k+1} = x_k - (F(x_k) - q) / f(x_k).
                pdf_mid = self.prob(mid)
                newton_mid = mid - newton_damping * (cdf_mid - q) / pdf_mid.clamp_min(2 * eps)

                # Use Newton step if it stays within the bracket, otherwise fall back to bisection.
                in_bracket = (newton_mid >= low) & (newton_mid <= high)
                mid = torch.where(in_bracket, input=newton_mid, other=bisect_mid)

            else:
                mid = bisect_mid

        return mid

    def rsample(self, sample_shape: torch.Size | list[int] | tuple[int, ...] = _size) -> torch.Tensor:
        """Draws reparameterized samples from the distribution, and allows gradients to flow backawards.

        Args:
            sample_shape: Desired shape of the samples to be drawn. Default is empty, which means one sample per batch element.

        Returns:
            Samples drawn from the distribution, with shape (*sample_shape, *batch_shape).
        """
        # Determine the final shape of the output tensor.
        shape = self._extended_shape(sample_shape)

        # Sample uniform noise, u ~ U(0, 1).
        u = torch.rand(shape, dtype=self.logits.dtype, device=self.logits.device)

        # Find the root (the sample x) without tracking gradients for the loop.
        with torch.no_grad():
            x_root = self.icdf(u)

        # Apply the implicit differentiation trick, i.e., evaluate CDF to connect the parameters to the
        # computational graph.
        cdf_val = self.cdf(x_root)

        # Evaluate PDF and detach it to act as the constant denominator.
        pdf_val = self.prob(x_root).detach()

        # Clamp PDF to avoid division by zero near the boundaries where slope is 0. This limits the gradients.
        eps = torch.finfo(pdf_val.dtype).eps
        pdf_val = pdf_val.clamp_min(2 * eps)

        # Attach the exact reparameterized gradient.
        x = x_root + (u - cdf_val) / pdf_val

        # Clamp to the support to prevent the implicit-differentiation correction from pushing samples
        # slightly past the domain boundaries when the CDF is very flat near the bounds.
        return x.clamp(min=self.bound_low, max=self.bound_up)