Inline Distribution Builders

The morphisms the DSL compiler emits for an inline family call such as sample x <- Normal(loc=mu, scale=1.0). A call whose arguments are all literals becomes a FixedDistribution, built by a make_fixed_* factory; a call with at least one bound argument becomes a MixedInlineDistribution, which reads the bound arguments from its stacked input and fills in the literals. Both can be constructed directly to build a program in Python.

inline

Inline distribution morphisms for monadic programs.

These ContinuousMorphisms parameterize distributions using either fixed literal values or direct variable values, without learned neural-network transformations. Used by the DSL compiler for inline draw steps like::

draw x ~ LogitNormal(0.0, 1.0)                 # all fixed
draw b ~ Bernoulli(x)                           # direct variable
draw r ~ TruncatedNormal(mu, sigma, 0.0, 1.0)  # mixed

Terminology

• Fixed: all distribution parameters are float literals known at compile time. The morphism has domain Unit (terminal object) and ignores its input.
• Direct: some or all parameters come from bound variables at runtime. The morphism uses the input tensor directly as distribution parameters (no learned neural-net transformation).

Architecture

Each family is described once by a FamilySpec in quivers.continuous.family_spec. The factories and builders below dispatch on the registered spec rather than maintaining per-family hard-coded functions; specials (truncated, Dirichlet, uniform, transformed) override the generic dispatch by setting fixed_factory_override / mixed_builder_override on their FamilySpec entry.

FixedDistribution

FixedDistribution(codomain: AnySpace, make_dist: Callable, discrete: bool = False, support: Constraint | None = None, *, domain: AnySpace = Unit)

Bases: ContinuousMorphism

A distribution with all parameters fixed at construction time.

The rsample input is used only for batch size and device inference; the distribution parameters themselves are constants.

PARAMETER	DESCRIPTION
codomain	The output space. TYPE: AnySpace
make_dist	(batch_size: int, device: torch.device) -> Distribution. TYPE: callable
discrete	Whether the output is discrete (returns LongTensor). TYPE: bool DEFAULT: False
support	The support of the underlying distribution. Used by variational guides to apply the right bijector for unconstrained→constrained sampling. Defaults to constraints.real; the factory functions for each family supply the correct constraint. TYPE: Constraint DEFAULT: None

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

def __init__(
    self,
    codomain: AnySpace,
    make_dist: Callable,
    discrete: bool = False,
    support: _constraints.Constraint | None = None,
    *,
    domain: AnySpace = Unit,
) -> None:
    super().__init__(domain, codomain)
    self._make_dist_fn = make_dist
    self._discrete = discrete
    self._support = support if support is not None else D.constraints.real

with_domain

with_domain(domain: AnySpace) -> FixedDistribution

Return the same fixed family on a declared input space.

PARAMETER	DESCRIPTION
domain	Input space whose values select no family parameters. TYPE: AnySpace

RETURNS	DESCRIPTION
FixedDistribution	A fixed distribution that keeps the declared kernel domain.

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

def with_domain(self, domain: AnySpace) -> FixedDistribution:
    """Return the same fixed family on a declared input space.

    Parameters
    ----------
    domain : AnySpace
        Input space whose values select no family parameters.

    Returns
    -------
    FixedDistribution
        A fixed distribution that keeps the declared kernel domain.
    """
    return FixedDistribution(
        self.codomain,
        self._make_dist_fn,
        discrete=self._discrete,
        support=self._support,
        domain=domain,
    )

rsample

rsample(x: Tensor, sample_shape: Size = Size()) -> Tensor

Sample from the fixed distribution.

PARAMETER	DESCRIPTION
x	Input (used only for batch size and device). TYPE: Tensor
sample_shape	Additional leading sample dimensions. TYPE: Size DEFAULT: Size()

RETURNS	DESCRIPTION
Tensor	Samples from the distribution.

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

def rsample(
    self, x: torch.Tensor, sample_shape: torch.Size = torch.Size()
) -> torch.Tensor:
    """Sample from the fixed distribution.

    Parameters
    ----------
    x : torch.Tensor
        Input (used only for batch size and device).
    sample_shape : torch.Size
        Additional leading sample dimensions.

    Returns
    -------
    torch.Tensor
        Samples from the distribution.
    """
    batch = x.shape[0]
    dist = self._make_dist_fn(batch, x.device)
    shape = torch.Size(sample_shape)
    if not dist.batch_shape and not dist.event_shape:
        # A closed-form operator measure (the `Restrict` a
        # desugared half-family produces, say) carries no batch of
        # its own: one draw per row, and per coordinate of a
        # continuous codomain, keeps the ``(batch, dim)`` contract
        # of every other fixed distribution.
        rows = (
            torch.Size((batch,))
            if self._discrete
            else torch.Size((batch, int(getattr(self.codomain, "dim", 1))))
        )
        shape = shape + rows
    if self._discrete:
        return dist.sample(shape).long()
    if getattr(dist, "has_rsample", True):
        return dist.rsample(shape)
    # Continuous but non-reparameterizable (e.g. VonMises).
    return dist.sample(shape)

log_prob

log_prob(x: Tensor, y: Tensor) -> Tensor

Log-probability under the fixed distribution.

PARAMETER	DESCRIPTION
x	Input (used only for batch size and device). TYPE: Tensor
y	Output values. TYPE: Tensor

RETURNS	DESCRIPTION
Tensor	Log-probabilities. Shape (batch,).

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

def log_prob(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
    """Log-probability under the fixed distribution.

    Parameters
    ----------
    x : torch.Tensor
        Input (used only for batch size and device).
    y : torch.Tensor
        Output values.

    Returns
    -------
    torch.Tensor
        Log-probabilities. Shape ``(batch,)``.
    """
    value = y.float() if self._discrete else y
    # The distribution's parameters are shared constants, so the
    # density is built once (batch of one) and its constant
    # parameters broadcast against whatever value is scored: a
    # globally-shared prior clamped to a single draw scores once, a
    # per-row batch broadcasts to that batch, and a value carrying
    # the family's own event axis (a Dirichlet simplex, say) keeps
    # that axis as the event rather than mistaking it for a batch.
    # Building at the value's leading axis instead would replicate a
    # shared prior across the program's response count, or read an
    # event coordinate as a batch row.
    device = value.device if value.dim() >= 1 else x.device
    dist = self._make_dist_fn(1, device)
    # A `Measure` (an operator-algebra sub-distribution such as
    # the `Restrict` a desugared half-family produces) carries its
    # residual mass symbolically in ``log_normalizer``; the
    # sample / observe boundary is where the algebra collapses it
    # to a probability density, so we score the renormalised
    # log-density here rather than the raw sub-measure value.
    value = _to_event_shape(value, dist)
    if isinstance(dist, _Measure):
        lp = dist.log_prob_normalized(value)
    else:
        lp = dist.log_prob(value)
    if lp.dim() > 1:
        return lp.sum(dim=-1)
    return lp

MixedInlineDistribution

MixedInlineDistribution(domain: AnySpace, codomain: AnySpace, param_spec: list[tuple[str, int | float | None]], dist_builder: Callable, discrete: bool = False, support: Constraint | None = None, param_event_ranks: tuple[int, ...] | None = None)

Bases: ContinuousMorphism

General inline distribution with arbitrary literal/variable mix.

Handles any combination of literal float values and runtime variable inputs for any distribution family. At construction time, receives a specification of which parameter positions are fixed literals vs. variable inputs, plus a builder function that creates the PyTorch distribution from fully-resolved parameter tensors.

This is the general mechanism underlying all inline distribution patterns in the DSL::

draw x ~ Normal(0.0, 1.0)           # all fixed (0 variable inputs)
draw y ~ Normal(mu, sigma)           # all variable (2 variable inputs)
draw z ~ Normal(mu, 0.5)             # mixed (1 variable input)
draw w ~ TruncatedNormal(mu, 0.5, 0.0, 1.0)  # mixed

PARAMETER	DESCRIPTION
domain	Source space (stacked variable parameters from program env). TYPE: AnySpace
codomain	Target space. TYPE: AnySpace
param_spec	For each distribution parameter position, one of: ('var', dim): a variable read from the input, dim columns wide, or None when its width is to be read off the stacked input (see _variable_widths); ('lit', value): a fixed literal float value. TYPE: list of tuple
dist_builder	(list[torch.Tensor]) -> torch.distributions.Distribution. Receives one 1-D tensor per parameter (all same batch size). TYPE: callable
discrete	Whether the output is discrete (returns LongTensor). TYPE: bool DEFAULT: False
support	Support of the underlying distribution; used by variational guides. Defaults to constraints.real. TYPE: Constraint DEFAULT: None

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

def __init__(
    self,
    domain: AnySpace,
    codomain: AnySpace,
    param_spec: list[tuple[str, int | float | None]],
    dist_builder: Callable,
    discrete: bool = False,
    support: _constraints.Constraint | None = None,
    param_event_ranks: tuple[int, ...] | None = None,
) -> None:
    super().__init__(domain, codomain)
    self._param_spec = param_spec
    self._dist_builder = dist_builder
    self._discrete = discrete
    self._support = support if support is not None else _constraints.real
    # Per-position event rank. Position 0 = per-row scalar; >= 1 =
    # vector-shaped distribution parameter (cutpoints, mixture
    # weights / locations / scales). Each consumes the number of
    # columns its ``param_spec`` width records, so the stacked input
    # is split by offset.
    self._param_event_ranks: tuple[int, ...] = (
        param_event_ranks
        if param_event_ranks is not None
        else tuple(0 for _ in param_spec)
    )
    if len(self._param_event_ranks) != len(param_spec):
        raise ValueError(
            "MixedInlineDistribution: param_event_ranks length "
            f"{len(self._param_event_ranks)} disagrees with "
            f"param_spec length {len(param_spec)}"
        )

rsample

rsample(x: Tensor, sample_shape: Size = Size()) -> Tensor

Sample from the distribution.

PARAMETER	DESCRIPTION
x	Stacked variable parameters. TYPE: Tensor
sample_shape	Additional leading sample dimensions. TYPE: Size DEFAULT: Size()

RETURNS	DESCRIPTION
Tensor	Samples from the distribution.

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

def rsample(
    self, x: torch.Tensor, sample_shape: torch.Size = torch.Size()
) -> torch.Tensor:
    """Sample from the distribution.

    Parameters
    ----------
    x : torch.Tensor
        Stacked variable parameters.
    sample_shape : torch.Size
        Additional leading sample dimensions.

    Returns
    -------
    torch.Tensor
        Samples from the distribution.
    """
    params = self._resolve_params(x)
    dist = self._dist_builder(params)
    if self._discrete:
        # Discrete samples are kept at the distribution's native
        # batch shape; downstream code unpacks as needed.
        return dist.sample(sample_shape).long()
    if getattr(dist, "has_rsample", True):
        result = dist.rsample(sample_shape)
    else:
        # Continuous but non-reparameterizable (e.g. VonMises).
        result = dist.sample(sample_shape)
    if result.dim() == 1:
        result = result.unsqueeze(-1)
    return result

log_prob

log_prob(x: Tensor, y: Tensor) -> Tensor

Log-probability under the distribution.

PARAMETER	DESCRIPTION
x	Stacked variable parameters. TYPE: Tensor
y	Observed values. TYPE: Tensor

RETURNS	DESCRIPTION
Tensor	Log-probabilities. Shape matches the broadcast of the distribution's batch shape with y's shape, with any event dimension summed out.

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

def log_prob(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
    """Log-probability under the distribution.

    Parameters
    ----------
    x : torch.Tensor
        Stacked variable parameters.
    y : torch.Tensor
        Observed values.

    Returns
    -------
    torch.Tensor
        Log-probabilities. Shape matches the broadcast of the
        distribution's batch shape with ``y``'s shape, with any
        event dimension summed out.
    """
    params = self._resolve_params(x)
    dist = self._dist_builder(params)
    # For discrete distributions over an integer support,
    # categorical-valued samples must remain integer-typed so
    # PyTorch's ``IntegerInterval`` validator accepts them.
    # Bernoulli / Boolean families accept float-valued samples,
    # but the integer family suite (Categorical, Binomial,
    # Geometric, NegativeBinomial, Poisson) refuses floats. Cast
    # only when the validator's support is a continuous
    # constraint.
    if self._discrete and y.dtype.is_floating_point:
        sup = self._support
        if isinstance(sup, _constraints._Boolean):
            y_in = y.float()
        else:
            y_in = y.long()
    else:
        y_in = y
    # Reconcile y's shape with the distribution's expected
    # event shape: append a trailing singleton if the
    # distribution carries an event dim that y lacks; drop a
    # trailing singleton from y when the distribution is scalar
    # and y carries an extra unit dim from upstream reshapes.
    y_in = _to_event_shape(y_in, dist)
    event_dim = len(dist.event_shape)
    batch_dim = len(dist.batch_shape)
    target_dim = event_dim + batch_dim
    if not self._discrete:
        while y_in.dim() < target_dim:
            y_in = y_in.unsqueeze(-1)
        while y_in.dim() > target_dim and y_in.shape[-1] == 1:
            y_in = y_in.squeeze(-1)
    lp = dist.log_prob(y_in)
    # Scalar parameters read one value per coordinate make the
    # coordinates independent draws of one value: their densities
    # sum, as a fixed distribution's do, leaving one per row.
    coordinates = self._scalar_coordinates(params)
    if coordinates > 1 and lp.dim() >= 1 and lp.shape[-1] == coordinates:
        return lp.sum(dim=-1)
    return lp

make_fixed_logitnormal

make_fixed_logitnormal(mu: float, sigma: float, codomain: AnySpace) -> FixedDistribution

Create a fixed LogitNormal(mu, sigma) distribution.

PARAMETER	DESCRIPTION
mu	Location parameter of the underlying normal. TYPE: float
sigma	Scale parameter of the underlying normal (positive). TYPE: float
codomain	Output space. TYPE: AnySpace

RETURNS	DESCRIPTION
FixedDistribution	Distribution morphism sampling from LogitNormal(mu, sigma).

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

def make_fixed_logitnormal(
    mu: float, sigma: float, codomain: AnySpace
) -> FixedDistribution:
    """Create a fixed LogitNormal(mu, sigma) distribution.

    Parameters
    ----------
    mu : float
        Location parameter of the underlying normal.
    sigma : float
        Scale parameter of the underlying normal (positive).
    codomain : AnySpace
        Output space.

    Returns
    -------
    FixedDistribution
        Distribution morphism sampling from LogitNormal(mu, sigma).
    """
    d = getattr(codomain, "dim", 1)

    def builder(batch: int, device: torch.device) -> D.Distribution:
        mu_t = torch.full((batch, d), mu, device=device)
        sigma_t = torch.full((batch, d), sigma, device=device)
        base = D.Normal(mu_t, sigma_t)
        return D.TransformedDistribution(base, [D.SigmoidTransform()])

    return FixedDistribution(codomain, builder, support=_constraints.unit_interval)

make_fixed_uniform

make_fixed_uniform(low: float, high: float, codomain: AnySpace) -> FixedDistribution

Create a fixed Uniform(low, high) distribution.

PARAMETER	DESCRIPTION
low	Lower bound. TYPE: float
high	Upper bound. TYPE: float
codomain	Output space. TYPE: AnySpace

RETURNS	DESCRIPTION
FixedDistribution	Distribution morphism sampling from Uniform(low, high).

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

def make_fixed_uniform(
    low: float, high: float, codomain: AnySpace
) -> FixedDistribution:
    """Create a fixed Uniform(low, high) distribution.

    Parameters
    ----------
    low : float
        Lower bound.
    high : float
        Upper bound.
    codomain : AnySpace
        Output space.

    Returns
    -------
    FixedDistribution
        Distribution morphism sampling from Uniform(low, high).
    """
    d = getattr(codomain, "dim", 1)

    def builder(batch: int, device: torch.device) -> D.Distribution:
        low_t = torch.full((batch, d), low, device=device)
        high_t = torch.full((batch, d), high, device=device)
        return D.Uniform(low_t, high_t)

    return FixedDistribution(
        codomain, builder, support=_constraints.interval(low, high)
    )

make_fixed_normal

make_fixed_normal(loc: float, scale: float, codomain: AnySpace) -> FixedDistribution

Create a fixed Normal(loc, scale) distribution.

PARAMETER	DESCRIPTION
loc	Mean. TYPE: float
scale	Standard deviation (positive). TYPE: float
codomain	Output space. TYPE: AnySpace

RETURNS	DESCRIPTION
FixedDistribution	Distribution morphism sampling from Normal(loc, scale).

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

def make_fixed_normal(
    loc: float, scale: float, codomain: AnySpace
) -> FixedDistribution:
    """Create a fixed Normal(loc, scale) distribution.

    Parameters
    ----------
    loc : float
        Mean.
    scale : float
        Standard deviation (positive).
    codomain : AnySpace
        Output space.

    Returns
    -------
    FixedDistribution
        Distribution morphism sampling from Normal(loc, scale).
    """
    d = getattr(codomain, "dim", 1)

    def builder(batch: int, device: torch.device) -> D.Distribution:
        loc_t = torch.full((batch, d), loc, device=device)
        scale_t = torch.full((batch, d), scale, device=device)
        return D.Normal(loc_t, scale_t)

    return FixedDistribution(codomain, builder)

make_fixed_bernoulli

make_fixed_bernoulli(prob: float, codomain: AnySpace) -> FixedDistribution

Create a fixed Bernoulli(prob) distribution.

PARAMETER	DESCRIPTION
prob	Success probability. TYPE: float
codomain	Output space (FinSet of size 2). TYPE: AnySpace

RETURNS	DESCRIPTION
FixedDistribution	Distribution morphism sampling from Bernoulli(prob).

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

def make_fixed_bernoulli(prob: float, codomain: AnySpace) -> FixedDistribution:
    """Create a fixed Bernoulli(prob) distribution.

    Parameters
    ----------
    prob : float
        Success probability.
    codomain : AnySpace
        Output space (FinSet of size 2).

    Returns
    -------
    FixedDistribution
        Distribution morphism sampling from Bernoulli(prob).
    """

    def builder(batch: int, device: torch.device) -> D.Distribution:
        probs_t = torch.full((batch,), prob, device=device)
        return D.Bernoulli(probs=probs_t)

    return FixedDistribution(
        codomain, builder, discrete=True, support=_constraints.boolean
    )

make_fixed_beta

make_fixed_beta(concentration1: float, concentration0: float, codomain: AnySpace) -> FixedDistribution

Create a fixed Beta(concentration1, concentration0) distribution.

PARAMETER	DESCRIPTION
concentration1	Alpha parameter. TYPE: float
concentration0	Beta parameter. TYPE: float
codomain	Output space. TYPE: AnySpace

RETURNS	DESCRIPTION
FixedDistribution	Distribution morphism sampling from Beta(a, b).

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

def make_fixed_beta(
    concentration1: float, concentration0: float, codomain: AnySpace
) -> FixedDistribution:
    """Create a fixed Beta(concentration1, concentration0) distribution.

    Parameters
    ----------
    concentration1 : float
        Alpha parameter.
    concentration0 : float
        Beta parameter.
    codomain : AnySpace
        Output space.

    Returns
    -------
    FixedDistribution
        Distribution morphism sampling from Beta(a, b).
    """
    d = getattr(codomain, "dim", 1)

    def builder(batch: int, device: torch.device) -> D.Distribution:
        a = torch.full((batch, d), concentration1, device=device)
        b = torch.full((batch, d), concentration0, device=device)
        return D.Beta(a, b)

    return FixedDistribution(codomain, builder, support=_constraints.unit_interval)

make_fixed_exponential

make_fixed_exponential(rate: float, codomain: AnySpace) -> FixedDistribution

Create a fixed Exponential(rate) distribution.

PARAMETER	DESCRIPTION
rate	Rate parameter (positive). TYPE: float
codomain	Output space. TYPE: AnySpace

RETURNS	DESCRIPTION
FixedDistribution	Distribution morphism sampling from Exponential(rate).

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

def make_fixed_exponential(rate: float, codomain: AnySpace) -> FixedDistribution:
    """Create a fixed Exponential(rate) distribution.

    Parameters
    ----------
    rate : float
        Rate parameter (positive).
    codomain : AnySpace
        Output space.

    Returns
    -------
    FixedDistribution
        Distribution morphism sampling from Exponential(rate).
    """
    d = getattr(codomain, "dim", 1)

    def builder(batch: int, device: torch.device) -> D.Distribution:
        rate_t = torch.full((batch, d), rate, device=device)
        return D.Exponential(rate_t)

    return FixedDistribution(codomain, builder, support=_constraints.positive)

make_fixed_halfcauchy

make_fixed_halfcauchy(scale: float, codomain: AnySpace) -> FixedDistribution

Create a fixed HalfCauchy(scale) distribution.

PARAMETER	DESCRIPTION
scale	Scale parameter (positive). TYPE: float
codomain	Output space. TYPE: AnySpace

RETURNS	DESCRIPTION
FixedDistribution	Distribution morphism sampling from HalfCauchy(scale).

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

def make_fixed_halfcauchy(scale: float, codomain: AnySpace) -> FixedDistribution:
    """Create a fixed HalfCauchy(scale) distribution.

    Parameters
    ----------
    scale : float
        Scale parameter (positive).
    codomain : AnySpace
        Output space.

    Returns
    -------
    FixedDistribution
        Distribution morphism sampling from HalfCauchy(scale).
    """
    d = getattr(codomain, "dim", 1)

    def builder(batch: int, device: torch.device) -> D.Distribution:
        scale_t = torch.full((batch, d), scale, device=device)
        return D.HalfCauchy(scale_t)

    return FixedDistribution(codomain, builder, support=_constraints.positive)

make_fixed_halfnormal

make_fixed_halfnormal(scale: float, codomain: AnySpace) -> FixedDistribution

Create a fixed HalfNormal(scale) distribution.

PARAMETER	DESCRIPTION
scale	Scale parameter (positive). TYPE: float
codomain	Output space. TYPE: AnySpace

RETURNS	DESCRIPTION
FixedDistribution	Distribution morphism sampling from HalfNormal(scale).

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

def make_fixed_halfnormal(scale: float, codomain: AnySpace) -> FixedDistribution:
    """Create a fixed HalfNormal(scale) distribution.

    Parameters
    ----------
    scale : float
        Scale parameter (positive).
    codomain : AnySpace
        Output space.

    Returns
    -------
    FixedDistribution
        Distribution morphism sampling from HalfNormal(scale).
    """
    d = getattr(codomain, "dim", 1)

    def builder(batch: int, device: torch.device) -> D.Distribution:
        scale_t = torch.full((batch, d), scale, device=device)
        return D.HalfNormal(scale_t)

    return FixedDistribution(codomain, builder, support=_constraints.positive)

make_fixed_lognormal

make_fixed_lognormal(loc: float, scale: float, codomain: AnySpace) -> FixedDistribution

Create a fixed LogNormal(loc, scale) distribution.

PARAMETER	DESCRIPTION
loc	Mean of the underlying normal. TYPE: float
scale	Standard deviation of the underlying normal (positive). TYPE: float
codomain	Output space. TYPE: AnySpace

RETURNS	DESCRIPTION
FixedDistribution	Distribution morphism sampling from LogNormal(loc, scale).

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

def make_fixed_lognormal(
    loc: float, scale: float, codomain: AnySpace
) -> FixedDistribution:
    """Create a fixed LogNormal(loc, scale) distribution.

    Parameters
    ----------
    loc : float
        Mean of the underlying normal.
    scale : float
        Standard deviation of the underlying normal (positive).
    codomain : AnySpace
        Output space.

    Returns
    -------
    FixedDistribution
        Distribution morphism sampling from LogNormal(loc, scale).
    """
    d = getattr(codomain, "dim", 1)

    def builder(batch: int, device: torch.device) -> D.Distribution:
        loc_t = torch.full((batch, d), loc, device=device)
        scale_t = torch.full((batch, d), scale, device=device)
        return D.LogNormal(loc_t, scale_t)

    return FixedDistribution(codomain, builder, support=_constraints.positive)

make_fixed_gamma

make_fixed_gamma(concentration: float, rate: float, codomain: AnySpace) -> FixedDistribution

Create a fixed Gamma(concentration, rate) distribution.

PARAMETER	DESCRIPTION
concentration	Shape parameter (positive). TYPE: float
rate	Rate parameter (positive). TYPE: float
codomain	Output space. TYPE: AnySpace

RETURNS	DESCRIPTION
FixedDistribution	Distribution morphism sampling from Gamma(concentration, rate).

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

def make_fixed_gamma(
    concentration: float, rate: float, codomain: AnySpace
) -> FixedDistribution:
    """Create a fixed Gamma(concentration, rate) distribution.

    Parameters
    ----------
    concentration : float
        Shape parameter (positive).
    rate : float
        Rate parameter (positive).
    codomain : AnySpace
        Output space.

    Returns
    -------
    FixedDistribution
        Distribution morphism sampling from Gamma(concentration, rate).
    """
    d = getattr(codomain, "dim", 1)

    def builder(batch: int, device: torch.device) -> D.Distribution:
        conc_t = torch.full((batch, d), concentration, device=device)
        rate_t = torch.full((batch, d), rate, device=device)
        return D.Gamma(conc_t, rate_t)

    return FixedDistribution(codomain, builder, support=_constraints.positive)

make_fixed_dirichlet

make_fixed_dirichlet(concentration: float | list[float], codomain: AnySpace) -> FixedDistribution

Create a fixed Dirichlet(concentration) distribution on the simplex over the codomain's element axis.

PARAMETER	DESCRIPTION
concentration	Concentration parameter. A scalar is broadcast to all simplex components (symmetric Dirichlet). A list / tuple of length equal to the codomain's dimension specifies per-component concentrations. Each entry must be positive. TYPE: float or list of float
codomain	Output space; its dim attribute gives the simplex dimension (the number of components). TYPE: AnySpace

RETURNS	DESCRIPTION
FixedDistribution	Distribution morphism sampling from the chosen Dirichlet on the (d-1)-simplex embedded in R^d.

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

def make_fixed_dirichlet(
    concentration: float | list[float], codomain: AnySpace
) -> FixedDistribution:
    """Create a fixed Dirichlet(concentration) distribution on the simplex
    over the codomain's element axis.

    Parameters
    ----------
    concentration : float or list of float
        Concentration parameter. A scalar is broadcast to all simplex
        components (symmetric Dirichlet). A list / tuple of length
        equal to the codomain's dimension specifies per-component
        concentrations. Each entry must be positive.
    codomain : AnySpace
        Output space; its ``dim`` attribute gives the simplex
        dimension (the number of components).

    Returns
    -------
    FixedDistribution
        Distribution morphism sampling from the chosen Dirichlet on
        the (d-1)-simplex embedded in ``R^d``.
    """
    d = getattr(codomain, "dim", 1)
    if d < 2:
        raise ValueError(f"Dirichlet codomain must have dim >= 2, got dim={d}")
    if isinstance(concentration, (list, tuple)):
        if len(concentration) == 1:
            # Single-element sequence is treated as a scalar
            # (symmetric Dirichlet); broadcast to the codomain's
            # simplex dimension.
            conc_values = [float(concentration[0])] * d
        elif len(concentration) != d:
            raise ValueError(
                f"Dirichlet concentration vector has length "
                f"{len(concentration)} but codomain has dim={d}"
            )
        else:
            conc_values = [float(c) for c in concentration]
    else:
        conc_values = [float(concentration)] * d
    if any(c <= 0.0 for c in conc_values):
        raise ValueError(
            f"Dirichlet concentration must be positive componentwise, got {conc_values}"
        )
    conc_tuple = tuple(conc_values)

    def builder(batch: int, device: torch.device) -> D.Distribution:
        conc_t = torch.tensor(conc_tuple, device=device).expand(batch, d)
        return D.Dirichlet(conc_t)

    return FixedDistribution(codomain, builder, support=_constraints.simplex)

make_fixed_truncated_normal

make_fixed_truncated_normal(mu: float, sigma: float, low: float, high: float, codomain: AnySpace) -> FixedDistribution

Create a fixed truncated normal with literal parameters.

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

def make_fixed_truncated_normal(
    mu: float,
    sigma: float,
    low: float,
    high: float,
    codomain: AnySpace,
) -> FixedDistribution:
    """Create a fixed truncated normal with literal parameters."""
    d = getattr(codomain, "dim", 1)

    def builder(batch: int, device: torch.device) -> D.Distribution:
        params = [
            torch.full((batch, d), value, device=device)
            for value in (mu, sigma, low, high)
        ]
        return _truncated_normal_builder(params)

    return FixedDistribution(
        codomain,
        builder,
        support=_constraints.interval(low, high),
    )

make_fixed_categorical

make_fixed_categorical(values: float | list[float], codomain: AnySpace, *, parameterization: str) -> FixedDistribution

Create a literal categorical from probabilities or logits.

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

def make_fixed_categorical(
    values: float | list[float],
    codomain: AnySpace,
    *,
    parameterization: str,
) -> FixedDistribution:
    """Create a literal categorical from probabilities or logits."""
    cardinality = int(getattr(codomain, "cardinality", getattr(codomain, "size", 2)))
    if isinstance(values, (list, tuple)):
        vector = [float(value) for value in values]
        if len(vector) != cardinality:
            raise ValueError(
                f"Categorical {parameterization} vector has length "
                f"{len(vector)} but codomain has cardinality={cardinality}"
            )
    else:
        vector = [float(values)] * cardinality

    def builder(batch: int, device: torch.device) -> D.Distribution:
        parameter = torch.tensor(vector, device=device).expand(batch, cardinality)
        if parameterization == "logits":
            return D.Categorical(logits=parameter)
        return D.Categorical(probs=parameter)

    return FixedDistribution(
        codomain,
        builder,
        discrete=True,
        support=_constraints.integer_interval(0, cardinality - 1),
    )

make_fixed_lkj_cholesky

make_fixed_lkj_cholesky(concentration: float, codomain: AnySpace) -> FixedDistribution

Create a fixed LKJ prior over Cholesky factors of correlation matrices.

PARAMETER	DESCRIPTION
concentration	The LKJ shape, positive; one is uniform over correlation matrices and larger values concentrate toward the identity. TYPE: float
codomain	Output space; its cardinality or dim gives the correlation matrix's size, as for the variable-parameter form. TYPE: AnySpace

RETURNS	DESCRIPTION
FixedDistribution	Distribution morphism sampling lower-triangular Cholesky factors.

RAISES	DESCRIPTION
ValueError	If concentration is not positive, or the codomain does not fix a matrix size of at least two.

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

def make_fixed_lkj_cholesky(
    concentration: float, codomain: AnySpace
) -> FixedDistribution:
    """Create a fixed LKJ prior over Cholesky factors of correlation matrices.

    Parameters
    ----------
    concentration : float
        The LKJ shape, positive; one is uniform over correlation
        matrices and larger values concentrate toward the identity.
    codomain : AnySpace
        Output space; its ``cardinality`` or ``dim`` gives the
        correlation matrix's size, as for the variable-parameter form.

    Returns
    -------
    FixedDistribution
        Distribution morphism sampling lower-triangular Cholesky factors.

    Raises
    ------
    ValueError
        If ``concentration`` is not positive, or the codomain does not
        fix a matrix size of at least two.
    """
    if concentration <= 0.0:
        raise ValueError(
            f"LKJCholesky concentration must be positive, got {concentration}"
        )
    dim = _codomain_dim_or_raise("LKJCholesky", codomain)
    if dim < 2:
        raise ValueError(f"LKJCholesky needs a matrix size of at least 2, got {dim}")

    def builder(batch: int, device: torch.device) -> D.Distribution:
        shape = torch.full((batch,), float(concentration), device=device)
        return D.LKJCholesky(dim, shape)

    return FixedDistribution(codomain, builder, support=_constraints.corr_cholesky)

get_inline_param_names

get_inline_param_names(family: str) -> tuple[str, ...] | None

Get the ordered parameter names for an inline family.

PARAMETER	DESCRIPTION
family	Distribution family name. TYPE: str

RETURNS	DESCRIPTION
tuple[str, ...] or None	Parameter names, or None if not an inline family.

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

def get_inline_param_names(family: str) -> tuple[str, ...] | None:
    """Get the ordered parameter names for an inline family.

    Parameters
    ----------
    family : str
        Distribution family name.

    Returns
    -------
    tuple[str, ...] or None
        Parameter names, or None if not an inline family.
    """
    if family in _OPERATOR_PARAM_NAMES:
        return _OPERATOR_PARAM_NAMES[family]
    if family in _FIXED_FACTORIES:
        return _FIXED_FACTORIES[family][0]
    if family in _FAMILY_BUILDERS:
        return _FAMILY_BUILDERS[family][0]
    if family == "TruncatedNormal":
        return ("mu", "sigma", "low", "high")
    return None

get_inline_parameterizations

get_inline_parameterizations(family: str) -> tuple[tuple[str, ...], ...]

Return every complete keyword schema accepted by a family call.

Most families expose one canonical schema. Categorical exposes two one-parameter schemas so probabilities and logits occupy the same source-language level without overloading a positional value.

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

def get_inline_parameterizations(family: str) -> tuple[tuple[str, ...], ...]:
    """Return every complete keyword schema accepted by a family call.

    Most families expose one canonical schema. ``Categorical`` exposes two
    one-parameter schemas so probabilities and logits occupy the same
    source-language level without overloading a positional value.
    """
    canonical = get_inline_param_names(family)
    if canonical is None:
        return ()
    return family_parameterizations(family, canonical)

make_inline_distribution

make_inline_distribution(family: str, args: tuple, codomain: AnySpace, variable_types: dict[str, AnySpace] | None = None) -> tuple[ContinuousMorphism, tuple[StepArgument, ...] | None]

Create an inline distribution from family name and mixed args.

Handles any combination of literal and variable arguments for any registered distribution family. The general mechanism:

1. All literals → FixedDistribution (no variable input)
2. Any variables → MixedInlineDistribution with a param_spec that records which positions are literals vs. variable slices

PARAMETER	DESCRIPTION
family	Distribution family name. TYPE: str
args	Arguments from the DSL. Strings are variable names, floats are literal values. TYPE: tuple of str | float
codomain	The output space for the distribution. TYPE: AnySpace
variable_types	Mapping of variable names to their space types (for domain construction of direct distributions). TYPE: dict or None DEFAULT: None

RETURNS	DESCRIPTION
tuple of (ContinuousMorphism, tuple[str, ...] | None)	The inline distribution morphism, and the variable names to pass as step input (None = use program input).

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

def make_inline_distribution(
    family: str,
    args: tuple,
    codomain: AnySpace,
    variable_types: dict[str, AnySpace] | None = None,
) -> tuple[ContinuousMorphism, tuple[StepArgument, ...] | None]:
    """Create an inline distribution from family name and mixed args.

    Handles any combination of literal and variable arguments for any
    registered distribution family. The general mechanism:

    1. All literals → ``FixedDistribution`` (no variable input)
    2. Any variables → ``MixedInlineDistribution`` with a param_spec
       that records which positions are literals vs. variable slices

    Parameters
    ----------
    family : str
        Distribution family name.
    args : tuple of str | float
        Arguments from the DSL. Strings are variable names,
        floats are literal values.
    codomain : AnySpace
        The output space for the distribution.
    variable_types : dict or None
        Mapping of variable names to their space types (for domain
        construction of direct distributions).

    Returns
    -------
    tuple of (ContinuousMorphism, tuple[str, ...] | None)
        The inline distribution morphism, and the variable names
        to pass as step input (None = use program input).
    """
    args, selected_parameters = _bind_inline_arguments(family, args)
    if family in _OPERATOR_FAMILIES:
        return _make_operator_distribution(
            family,
            args,
            codomain,
            variable_types,
            selected_parameters,
        )
    args = _normalize_inline_args(args)
    # An indexed reference counts as a variable reference for the
    # purposes of picking the fixed vs mixed factory dispatch, alongside
    # bare identifier strings.
    var_names = [a for a in args if isinstance(a, str | Indexed)]
    if not var_names:
        if family == "Categorical":
            values = _eval_draw_arg_value(args[0], variable_types)
            return (
                make_fixed_categorical(
                    values,
                    codomain,
                    parameterization=(selected_parameters or ("probs",))[0],
                ),
                None,
            )
        if family in _FIXED_FACTORIES:
            factory_parameters, factory = _FIXED_FACTORIES[family]
            if (
                selected_parameters is not None
                and selected_parameters != factory_parameters
            ):
                return (
                    _make_fixed_selected_parameters(
                        family,
                        args,
                        selected_parameters,
                        codomain,
                    ),
                    None,
                )
            # Vector-parameter families take a single ``list[float]``
            # / ``tuple[float, ...]`` argument rather than splatting
            # the literals. The parser surfaces ``Dirichlet([1, 2, 3])``
            # as three positional float args (the grammar's draw-arg
            # list flattens any bracket-bounded numeric sequence), so
            # we re-bundle here when the factory's documented contract
            # is a single vector argument.
            if family in _VECTOR_PARAM_FAMILIES:
                value = _eval_draw_arg_value(args[0], variable_types)
                morph = factory(value, codomain)
            else:
                all_floats = [float(a) for a in args]
                morph = factory(*all_floats, codomain)
            return (morph, None)
        raise ValueError(f"no fixed factory for inline family {family!r}")
    if family not in _FAMILY_BUILDERS:
        raise ValueError(
            f"no builder for inline family {family!r} with variable arguments"
        )
    param_names, dist_builder, discrete = _FAMILY_BUILDERS[family]
    if family == "Categorical" and selected_parameters == ("logits",):
        param_names = ("logits",)
        dist_builder = _categorical_logits_builder
    elif selected_parameters is not None and selected_parameters != param_names:
        spec = FAMILY_REGISTRY[family]
        source_parameters = selected_parameters

        def _selected_builder(
            params: list[torch.Tensor],
            _spec: FamilySpec = spec,
            _parameters: tuple[str, ...] = source_parameters,
        ) -> D.Distribution:
            values: dict[str, torch.Tensor] = {}
            for index, (name, value) in enumerate(
                zip(_parameters, params, strict=True)
            ):
                if name == "probs":
                    values[name] = value.clamp(min=EPS, max=1.0 - EPS)
                else:
                    values[name] = _spec.params[index].inline_clamp(value)
            return _spec.dist_class(**values)

        param_names = selected_parameters
        dist_builder = _selected_builder
    if len(args) != len(param_names):
        raise ValueError(
            f"inline {family} expects {len(param_names)} args ({', '.join(param_names)}), got {len(args)}"
        )
    fam_event_ranks = _PARAM_EVENT_RANKS.get(
        family,
        tuple(0 for _ in param_names),
    )
    param_spec: list[tuple[str, int | float | None]] = []
    var_name_order: list[StepArgument] = []
    total_var_dim = 0
    for i, arg in enumerate(args):
        if isinstance(arg, (int, float)):
            param_spec.append(("lit", float(arg)))
        else:
            # A parameter's width comes from its binding's type; without
            # one it is left unresolved, for the stacked width to settle.
            var_dim: int | None = None
            lookup_name = arg.name if isinstance(arg, Indexed) else arg
            if variable_types and lookup_name in variable_types:
                vtype = variable_types[lookup_name]
                if isinstance(arg, Indexed) and isinstance(vtype, Euclidean):
                    # A gather selects whole rows of a plate, so the
                    # argument is one row wide.
                    var_dim = vtype.row_width
                else:
                    var_dim = _param_row_width(vtype, fam_event_ranks[i])
            param_spec.append(("var", var_dim))
            var_name_order.append(arg)
            total_var_dim += 1 if var_dim is None else var_dim
    if family == "MixtureNormal":
        # QIEC can currently retain the component width for an indexed draw
        # while conservatively typing a derived vector-valued ``let`` as
        # width one. MixtureNormal's three vector parameters necessarily
        # share one component axis, so a known width on any slot determines
        # the unresolved derived slots without guessing at runtime after the
        # arguments have been concatenated.
        known_component_dims = {
            int(value)
            for (kind, value), rank in zip(param_spec, fam_event_ranks, strict=True)
            if kind == "var" and rank >= 1 and value is not None and int(value) > 1
        }
        if len(known_component_dims) == 1:
            component_dim = known_component_dims.pop()
            param_spec = [
                (kind, component_dim)
                if kind == "var" and rank >= 1 and (value is None or int(value) == 1)
                else (kind, value)
                for (kind, value), rank in zip(param_spec, fam_event_ranks, strict=True)
            ]
            total_var_dim = sum(
                1 if value is None else int(value)
                for kind, value in param_spec
                if kind == "var"
            )
    if total_var_dim == 0:
        domain = Euclidean(name="_inline_domain", dim=1)
    elif len(var_name_order) == 1 and variable_types:
        vtype = _argument_type(var_name_order[0], variable_types)
        domain = (
            vtype
            if vtype is not None
            else Euclidean(name="_inline_domain", dim=total_var_dim)
        )
    else:
        domain = _infer_domain(var_name_order, variable_types)
    fam_support: _constraints.Constraint = _FAMILY_SUPPORTS.get(
        family, _constraints.real
    )
    # TruncatedNormal / Uniform encode bounded supports via two literal
    # arguments; if both bounds are literal we can specialize the
    # constraint to the matching `interval`. Otherwise we fall
    # back to ``real`` (the closest correct guide-side approximation).
    if family in ("TruncatedNormal", "Uniform"):
        lit_args: list[float] = []
        for kind, value in param_spec:
            if kind == "lit" and value is not None:
                lit_args.append(float(value))
        if family == "Uniform" and len(lit_args) == 2:
            fam_support = _constraints.interval(lit_args[0], lit_args[1])
        elif family == "TruncatedNormal" and len(lit_args) >= 2:
            fam_support = _constraints.interval(lit_args[-2], lit_args[-1])
    # Dimension-dependent families need the codomain's matrix / vector
    # size baked into the builder closure. The registered
    # `_FAMILY_BUILDERS` entry carries a dim-agnostic builder; here we
    # replace it with a closure that captures the codomain dim so the
    # resulting torch distribution has the right event_shape at
    # construction.
    if family in _DIM_DEPENDENT_FAMILIES:
        dist_builder = _dim_dependent_builder(family, codomain)
    morph = MixedInlineDistribution(
        domain,
        codomain,
        param_spec,
        dist_builder,
        discrete,
        support=fam_support,
        param_event_ranks=fam_event_ranks,
    )
    return (morph, tuple(var_name_order))

reload_inline_registry

reload_inline_registry() -> None

Refresh the inline dicts from quivers.continuous.family_spec.FAMILY_REGISTRY. Useful when new families are registered after inline has imported (test plugins, downstream extensions).

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

def reload_inline_registry() -> None:
    """Refresh the inline dicts from
    `quivers.continuous.family_spec.FAMILY_REGISTRY`. Useful
    when new families are registered after `inline` has
    imported (test plugins, downstream extensions)."""
    for spec in FAMILY_REGISTRY.values():
        _auto_register_inline(spec)