Measure Algebra

The compositional measure algebra behind the operator families of the DSL (Mixture, Restrict, Pushforward, PointMass, ...). Each combinator builds a measure from measures, and normalize_at_boundary renormalizes a sub-probability measure where it is sampled or scored.

measure

Sub-Giry measures and their primitive transformations.

A Measure is an unnormalized positive measure on a Borel space. Probability measures have total mass one; restrictions and likelihood scores may produce sub-probability measures. The runtime tracks total mass as a log normalizer and normalizes at sample or observation boundaries.

The module provides point masses, restriction, pushforward through a bijector, finite mixtures, event-dimension declarations with Independent, and normalization of positive-mass sub-measures.

Measure

Bases: Distribution, ABC

Abstract base for sub-Giry monad values.

The runtime tracks each Measure as a sub-distribution: it has a log_prob returning the unnormalised log-density and a log_normalizer() returning the log of the total mass. For probability measures (mass one), log_normalizer() returns zero by definition. For sub-measures from Restrict or Mixture with unnormalised weights, log_normalizer() carries the residual mass symbolically until the next Normalize or sample / observe boundary.

log_prob_normalized(value) is a derived helper returning log_prob(value) - log_normalizer() so consumers that want the probability-measure log-density can ask for it explicitly.

log_normalizer abstractmethod

log_normalizer() -> Tensor

Log of the total mass of this measure. For probability measures, zero. For sub-measures, the symbolic normaliser the lazy algebra carries through composition.

Source code in src/quivers/continuous/measure.py

@abstractmethod
def log_normalizer(self) -> Tensor:
    """Log of the total mass of this measure. For probability
    measures, zero. For sub-measures, the symbolic normaliser
    the lazy algebra carries through composition.
    """

log_prob_normalized

log_prob_normalized(value: Tensor) -> Tensor

Log-density of the renormalised probability measure.

Source code in src/quivers/continuous/measure.py

def log_prob_normalized(self, value: Tensor) -> Tensor:
    """Log-density of the renormalised probability measure."""
    return self.log_prob(value) - self.log_normalizer()

PointMass

PointMass(value: Tensor | float, validate_args: bool | None = None)

Bases: Measure

Dirac measure at a deterministic value.

The unit \(\eta\) of the Giry monad: PointMass(x) is the measure that places mass one on the singleton \(\{x\}\). log_prob(y) is zero when y == value and \(-\infty\) otherwise.

In a Mixture alongside a non-degenerate component, PointMass is the structural-zero spike that yields zero-inflation and hurdle factorisations.

Source code in src/quivers/continuous/measure.py

def __init__(
    self,
    value: Tensor | float,
    validate_args: bool | None = None,
) -> None:
    self.value = torch.as_tensor(value)
    super().__init__(
        batch_shape=self.value.shape,
        event_shape=torch.Size(()),
        validate_args=validate_args,
    )

Restrict

Restrict(base: Distribution, low: Tensor | float | None = None, high: Tensor | float | None = None, validate_args: bool | None = None)

Bases: Measure

Restriction of a base distribution to a sub-interval.

Restricts base to the interval [low, high] (with low or high set to \(\pm\infty\) for one-sided restrictions). The log-density is base.log_prob(value) inside the interval and \(-\infty\) outside; the symbolic log-normaliser is the log of the base measure's mass on [low, high].

The compiler does not renormalise eagerly; the sub-measure propagates through Mixture, Pushforward, and other operators until a Normalize or sample / observe boundary collapses it. This is the discipline that makes the algebra close, per Di Lavore-Roman-Sobocinski 2025.

Numerically, the log-normaliser is computed via the base's cdf (and the 1 - cdf complement when only low is finite) with the logsubexp stable variant. Bases without a closed-form CDF raise at construction time.

Source code in src/quivers/continuous/measure.py

def __init__(
    self,
    base: Distribution,
    low: Tensor | float | None = None,
    high: Tensor | float | None = None,
    validate_args: bool | None = None,
) -> None:
    if low is None and high is None:
        raise ValueError(
            "Restrict: at least one of `low` or `high` must be "
            "supplied; an unrestricted base is just the base "
            "distribution itself."
        )
    self.base = base
    self.low = (
        torch.as_tensor(low, dtype=torch.get_default_dtype())
        if low is not None
        else None
    )
    self.high = (
        torch.as_tensor(high, dtype=torch.get_default_dtype())
        if high is not None
        else None
    )
    super().__init__(
        batch_shape=base.batch_shape,
        event_shape=base.event_shape,
        validate_args=validate_args,
    )

Pushforward

Pushforward(base: Distribution, bijector: Bijector, validate_args: bool | None = None)

Bases: Measure

Pushforward of a base distribution through a bijector.

Implements \(f_*\mu\), the measure whose log-density at y is base.log_prob(b.inverse(y)) + b.inverse_log_det_jacobian(y). Functorial in b: composing bijectors composes the pushforward.

Does not commute with Restrict in general; the compiler's rewrite pass handles the cases that do (monotone bijectors with interval restrictions).

Source code in src/quivers/continuous/measure.py

def __init__(
    self,
    base: Distribution,
    bijector: Bijector,
    validate_args: bool | None = None,
) -> None:
    self.base = base
    self.bijector = bijector
    super().__init__(
        batch_shape=base.batch_shape,
        event_shape=base.event_shape,
        validate_args=validate_args,
    )

Mixture

Mixture(weights: Tensor, components: list[Distribution], validate_args: bool | None = None)

Bases: Measure

Finite mixture of arbitrary-family components.

weights is a vector summing to one (or, in the sub-measure regime, any non-negative vector; the total mass is the sum). components is a list of Distribution instances, possibly of different families and possibly including PointMass for inflation / hurdle structure.

The log-density is logsumexp(log_weights + log_components) over the component axis, broadcasting weights against the distribution batch shape. The symbolic log-normaliser is log(sum_k w_k * Z_k) where Z_k is the per-component normaliser (in particular when components are themselves sub-measures from Restrict).

Associativity (the Giry monad's multiplication law) is canonicalised by flatten(), which the compiler's rewrite pass calls before lowering.

PyMC issue #5533 documents how nested-mixture composition breaks in production PPLs; the operator algebra closes the gap.

Source code in src/quivers/continuous/measure.py

def __init__(
    self,
    weights: Tensor,
    components: list[Distribution],
    validate_args: bool | None = None,
) -> None:
    if not components:
        raise ValueError("Mixture: components list must be non-empty")
    if weights.shape[-1] != len(components):
        raise ValueError(
            f"Mixture: weights last-dim {weights.shape[-1]} must "
            f"equal len(components) {len(components)}"
        )
    self.weights = weights
    self.components = list(components)
    self._num_components = len(components)
    batch_shape = weights.shape[:-1]
    for c in components:
        batch_shape = torch.broadcast_shapes(batch_shape, c.batch_shape)
    super().__init__(
        batch_shape=batch_shape,
        event_shape=components[0].event_shape,
        validate_args=validate_args,
    )

cdf

cdf(value: Tensor) -> Tensor

Cumulative distribution function: the weighted sum of component CDFs. Required for Restrict(Mixture, ...) to compute its log-normaliser without falling through to the discrete-base fallback.

Source code in src/quivers/continuous/measure.py

def cdf(self, value: Tensor) -> Tensor:
    """Cumulative distribution function: the weighted sum of
    component CDFs. Required for `Restrict(Mixture, ...)` to
    compute its log-normaliser without falling through to the
    discrete-base fallback.
    """
    per_component = []
    for c in self.components:
        per_component.append(c.cdf(value))
    stacked = torch.stack(per_component, dim=-1)
    return (self.weights * stacked).sum(dim=-1)

flatten

flatten() -> Mixture

Apply the Giry-monad-multiplication associativity rule: any nested Mixture component is replaced by its components with weights scaled by the outer weight. Returns a flat Mixture whose components are not themselves Mixture instances.

Source code in src/quivers/continuous/measure.py

def flatten(self) -> Mixture:
    """Apply the Giry-monad-multiplication associativity rule:
    any nested `Mixture` component is replaced by its components
    with weights scaled by the outer weight. Returns a flat
    `Mixture` whose components are not themselves `Mixture`
    instances.
    """
    flat_weights_parts = []
    flat_components: list[Distribution] = []
    for k, comp in enumerate(self.components):
        w_k = self.weights[..., k : k + 1]
        if isinstance(comp, Mixture):
            inner = comp.flatten()
            flat_weights_parts.append(w_k * inner.weights)
            flat_components.extend(inner.components)
        else:
            flat_weights_parts.append(w_k)
            flat_components.append(comp)
    flat_weights = torch.cat(flat_weights_parts, dim=-1)
    return Mixture(flat_weights, flat_components)

pushforward_inside

pushforward_inside(bijector) -> Mixture

Push a single bijector inside every component:

\[g_*\left(\sum_k \pi_k\,\mu_k\right) = \sum_k \pi_k\,g_*\mu_k.\]

Functoriality of the Giry monad on measurable isomorphisms; the per-component pushforward is the canonical shape the renderer prefers.

Source code in src/quivers/continuous/measure.py

def pushforward_inside(self, bijector) -> Mixture:
    """Push a single bijector inside every component:

    $$g_*\\left(\\sum_k \\pi_k\\,\\mu_k\\right)
      = \\sum_k \\pi_k\\,g_*\\mu_k.$$

    Functoriality of the Giry monad on measurable isomorphisms;
    the per-component pushforward is the canonical shape the
    renderer prefers.
    """
    new_components = [Pushforward(c, bijector) for c in self.components]
    return Mixture(self.weights, new_components)

restrict_to

restrict_to(low=None, high=None) -> Mixture

Restrict every component to the same subset, with the per-component reweighting that the non-commutation of Restrict and Mixture requires:

\[\mathrm{Restrict}\!\left(\sum_k \pi_k\,\mu_k,\; S\right) = \sum_k \frac{\pi_k\,\mu_k(S)}{\sum_j \pi_j\,\mu_j(S)}\, \mathrm{Restrict}(\mu_k, S).\]

Returns a Mixture whose components are each Restrict(component, low, high) and whose weights are the reweighted-by-component-mass form. Reference: the v2 design note's rewrite rules section (Welsh et al. 1996 for the statistical motivation).

Source code in src/quivers/continuous/measure.py

def restrict_to(self, low=None, high=None) -> Mixture:
    """Restrict every component to the same subset, with the
    per-component reweighting that the non-commutation of
    Restrict and Mixture requires:

    $$\\mathrm{Restrict}\\!\\left(\\sum_k \\pi_k\\,\\mu_k,\\; S\\right)
      = \\sum_k \\frac{\\pi_k\\,\\mu_k(S)}{\\sum_j \\pi_j\\,\\mu_j(S)}\\,
        \\mathrm{Restrict}(\\mu_k, S).$$

    Returns a `Mixture` whose components are each
    `Restrict(component, low, high)` and whose weights are the
    reweighted-by-component-mass form. Reference: the v2 design
    note's rewrite rules section (Welsh et al. 1996 for the
    statistical motivation).
    """
    restricted_raw = [Restrict(c, low=low, high=high) for c in self.components]
    log_w = torch.log(
        self.weights.clamp_min(
            torch.finfo(self.weights.dtype).tiny,
        )
    )
    log_masses = [r.log_normalizer() for r in restricted_raw]
    stacked = torch.stack(
        [
            lw + lm
            for lw, lm in zip(
                log_w.unbind(dim=-1),
                log_masses,
                strict=False,
            )
        ],
        dim=-1,
    )
    log_denom = torch.logsumexp(stacked, dim=-1, keepdim=True)
    new_log_weights = stacked - log_denom
    new_weights = torch.exp(new_log_weights)
    # Components are NORMALIZED truncations so the Mixture's
    # weighted sum of per-component densities equals the
    # conditional density of the original mixture restricted to
    # `S`. Using unnormalised `Restrict` here would double-count
    # the per-component truncation mass.
    normalized_components = [Normalize(r) for r in restricted_raw]
    return Mixture(new_weights, normalized_components)

lift_point_masses

lift_point_masses() -> Mixture

Surface the PointMass components as a Bernoulli-style branch with the non-degenerate components in the complementary mixture. Returns a Mixture whose components are pairwise measure-disjoint: each PointMass(x) carries only its spike, and the non-degenerate components live in the complement.

For a single PointMass plus a single non-degenerate component (the ZIP / hurdle canonical shape), the rewrite is a no-op: the current form is already the Bernoulli-branch factorisation the renderers expect. For Mixtures with multiple PointMass components at distinct values, returns a canonical form with the point masses ordered first.

Source code in src/quivers/continuous/measure.py

def lift_point_masses(self) -> Mixture:
    """Surface the `PointMass` components as a Bernoulli-style
    branch with the non-degenerate components in the
    complementary mixture. Returns a `Mixture` whose components
    are pairwise measure-disjoint: each `PointMass(x)` carries
    only its spike, and the non-degenerate components live in the
    complement.

    For a single `PointMass` plus a single non-degenerate
    component (the ZIP / hurdle canonical shape), the rewrite is
    a no-op: the current form is already the Bernoulli-branch
    factorisation the renderers expect. For Mixtures with
    multiple `PointMass` components at distinct values, returns a
    canonical form with the point masses ordered first.
    """
    point_masses = []
    rest = []
    point_mass_weights = []
    rest_weights = []
    for k, comp in enumerate(self.components):
        w_k = self.weights[..., k : k + 1]
        if isinstance(comp, PointMass):
            point_masses.append(comp)
            point_mass_weights.append(w_k)
        else:
            rest.append(comp)
            rest_weights.append(w_k)
    if not point_masses:
        return self
    new_components = point_masses + rest
    new_weights = torch.cat(point_mass_weights + rest_weights, dim=-1)
    return Mixture(new_weights, new_components)

Independent

Independent(base: Distribution, reinterpreted_batch_ndims: int, validate_args: bool | None = None)

Bases: Measure

Reinterpret the last reinterpreted_batch_ndims batch dims of base as event dims. Wraps PyTorch's Independent.

The log-density sums base.log_prob(value) over the reinterpreted dims; the log-normaliser sums in the same direction when base is itself a sub-measure.

Source code in src/quivers/continuous/measure.py

def __init__(
    self,
    base: Distribution,
    reinterpreted_batch_ndims: int,
    validate_args: bool | None = None,
) -> None:
    if reinterpreted_batch_ndims < 1:
        raise ValueError(
            "Independent: reinterpreted_batch_ndims must be >= 1, "
            f"got {reinterpreted_batch_ndims}"
        )
    self.base = base
    self.reinterpreted_batch_ndims = int(reinterpreted_batch_ndims)
    self._inner = torch.distributions.Independent(
        base,
        reinterpreted_batch_ndims,
    )
    super().__init__(
        batch_shape=self._inner.batch_shape,
        event_shape=self._inner.event_shape,
        validate_args=validate_args,
    )

Normalize

Normalize(base: Distribution, validate_args: bool | None = None)

Bases: Measure

Rescale a sub-measure to a probability measure.

log_prob_normalized(value) = log_prob(value) - log_normalizer() by construction; this wrapper makes the normalisation explicit so downstream code can rely on log_normalizer() == 0. Use at the observe / sample boundary or whenever a probability measure is required (e.g. before sampling).

Source code in src/quivers/continuous/measure.py

def __init__(
    self,
    base: Distribution,
    validate_args: bool | None = None,
) -> None:
    self.base = base
    if isinstance(base, Measure):
        self._base_log_normalizer = base.log_normalizer()
    else:
        self._base_log_normalizer = (
            torch.zeros(base.batch_shape) if base.batch_shape else torch.tensor(0.0)
        )
    super().__init__(
        batch_shape=base.batch_shape,
        event_shape=base.event_shape,
        validate_args=validate_args,
    )

normalize_at_boundary

normalize_at_boundary(d: Distribution) -> Distribution

Lift d to a probability measure if it is a sub-measure with non-zero log-normaliser; otherwise return d unchanged.

Called by the inline / observe path to enforce the "renormalise at the observe / sample boundary" invariant the sub-Giry algebra relies on.

Source code in src/quivers/continuous/measure.py

def normalize_at_boundary(d: Distribution) -> Distribution:
    """Lift `d` to a probability measure if it is a sub-measure
    with non-zero log-normaliser; otherwise return `d` unchanged.

    Called by the inline / observe path to enforce the
    "renormalise at the observe / sample boundary" invariant the
    sub-Giry algebra relies on.
    """
    if isinstance(d, Measure):
        norm = d.log_normalizer()
        if isinstance(norm, Tensor):
            if torch.any(norm.abs() > 1e-10):
                return Normalize(d)
        else:
            if abs(float(norm)) > 1e-10:
                return Normalize(d)
    return d