Continuous Families

Families of continuous-valued distributions.

families

Parameterized distribution families as continuous morphisms.

Each family is a ContinuousMorphism whose codomain is a continuous space and whose conditional distribution p(y | x) belongs to a specific parametric family. The parameters are learnable functions of x:

• For discrete domains (FinSet): parameters are looked up from a table.
• For continuous domains (ContinuousSpace): parameters are produced by a small neural network.

This module wraps every reparameterizable distribution in torch.distributions as a conditional morphism, plus custom families (TruncatedNormal, MultivariateNormal, etc.).

Architecture

Most per-dimension-independent distributions are built on a shared generic base _IndependentConditional that handles the parameter source, transform, and torch.distributions plumbing. The _make_family class factory generates named classes from a specification. Distributions that need special handling (MultivariateNormal, Dirichlet, TruncatedNormal, etc.) are implemented as standalone classes.

ConditionalCauchy module-attribute

ConditionalCauchy = _make_family('ConditionalCauchy', D.Cauchy, [('loc', 'id'), ('scale', 'softplus')], 'Conditional Cauchy(loc(x), scale(x)). Heavy-tailed, no finite moments.')

Conditional Cauchy(loc(x), scale(x)). Heavy-tailed, no finite moments.

ConditionalLaplace module-attribute

ConditionalLaplace = _make_family('ConditionalLaplace', D.Laplace, [('loc', 'id'), ('scale', 'softplus')], 'Conditional Laplace(loc(x), scale(x)). Sharp peak, heavier tails than normal.')

Conditional Laplace(loc(x), scale(x)). Sharp peak, heavier tails than normal.

ConditionalGumbel module-attribute

ConditionalGumbel = _make_family('ConditionalGumbel', D.Gumbel, [('loc', 'id'), ('scale', 'softplus')], 'Conditional Gumbel(loc(x), scale(x)). Extreme value distribution (type I).')

Conditional Gumbel(loc(x), scale(x)). Extreme value distribution (type I).

ConditionalLogNormal module-attribute

ConditionalLogNormal = _make_family('ConditionalLogNormal', D.LogNormal, [('loc', 'id'), ('scale', 'softplus')], 'Conditional LogNormal(loc(x), scale(x)). Positive-valued, right-skewed.')

Conditional LogNormal(loc(x), scale(x)). Positive-valued, right-skewed.

ConditionalStudentT module-attribute

ConditionalStudentT = _make_family('ConditionalStudentT', D.StudentT, [('df', 'softplus_shifted'), ('loc', 'id'), ('scale', 'softplus')], 'Conditional StudentT(df(x), loc(x), scale(x)). Heavy-tailed with learnable df.')

Conditional StudentT(df(x), loc(x), scale(x)). Heavy-tailed with learnable df.

ConditionalExponential module-attribute

ConditionalExponential = _make_family('ConditionalExponential', D.Exponential, [('rate', 'softplus')], 'Conditional Exponential(rate(x)). Memoryless, positive-valued.')

Conditional Exponential(rate(x)). Memoryless, positive-valued.

ConditionalGamma module-attribute

ConditionalGamma = _make_family('ConditionalGamma', D.Gamma, [('concentration', 'softplus_shifted'), ('rate', 'softplus')], 'Conditional Gamma(concentration(x), rate(x)). Positive-valued, flexible shape.')

Conditional Gamma(concentration(x), rate(x)). Positive-valued, flexible shape.

ConditionalChi2 module-attribute

ConditionalChi2 = _make_family('ConditionalChi2', D.Chi2, [('df', 'softplus_shifted')], 'Conditional Chi2(df(x)). Positive-valued, sum of squared normals.')

Conditional Chi2(df(x)). Positive-valued, sum of squared normals.

ConditionalHalfCauchy module-attribute

ConditionalHalfCauchy = _make_family('ConditionalHalfCauchy', D.HalfCauchy, [('scale', 'softplus')], 'Conditional HalfCauchy(scale(x)). Heavy-tailed positive prior.')

Conditional HalfCauchy(scale(x)). Heavy-tailed positive prior.

ConditionalHalfNormal module-attribute

ConditionalHalfNormal = _make_family('ConditionalHalfNormal', D.HalfNormal, [('scale', 'softplus')], 'Conditional HalfNormal(scale(x)). Folded normal, positive-valued.')

Conditional HalfNormal(scale(x)). Folded normal, positive-valued.

ConditionalInverseGamma module-attribute

ConditionalInverseGamma = _make_family('ConditionalInverseGamma', D.InverseGamma, [('concentration', 'softplus_shifted'), ('rate', 'softplus')], 'Conditional InverseGamma(concentration(x), rate(x)). Conjugate prior for normal variance.')

Conditional InverseGamma(concentration(x), rate(x)). Conjugate prior for normal variance.

ConditionalWeibull module-attribute

ConditionalWeibull = _make_family('ConditionalWeibull', D.Weibull, [('scale', 'softplus'), ('concentration', 'softplus')], 'Conditional Weibull(scale(x), concentration(x)). Survival analysis, reliability.')

Conditional Weibull(scale(x), concentration(x)). Survival analysis, reliability.

ConditionalPareto module-attribute

ConditionalPareto = _make_family('ConditionalPareto', D.Pareto, [('scale', 'softplus'), ('alpha', 'softplus')], 'Conditional Pareto(scale(x), alpha(x)). Power-law tail.')

Conditional Pareto(scale(x), alpha(x)). Power-law tail.

ConditionalKumaraswamy module-attribute

ConditionalKumaraswamy = _make_family('ConditionalKumaraswamy', D.Kumaraswamy, [('concentration1', 'softplus_shifted'), ('concentration0', 'softplus_shifted')], 'Conditional Kumaraswamy(a(x), b(x)). Beta-like on (0,1), closed-form CDF.')

Conditional Kumaraswamy(a(x), b(x)). Beta-like on (0,1), closed-form CDF.

ConditionalContinuousBernoulli module-attribute

ConditionalContinuousBernoulli = _make_family('ConditionalContinuousBernoulli', D.ContinuousBernoulli, [('logits', 'id')], 'Conditional ContinuousBernoulli(logits(x)). Continuous relaxation of Bernoulli.')

Conditional ContinuousBernoulli(logits(x)). Continuous relaxation of Bernoulli.

ConditionalFisherSnedecor module-attribute

ConditionalFisherSnedecor = _make_family('ConditionalFisherSnedecor', D.FisherSnedecor, [('df1', 'softplus_shifted'), ('df2', 'softplus_shifted')], 'Conditional FisherSnedecor(df1(x), df2(x)). F-distribution, ratio of chi-squared.')

Conditional FisherSnedecor(df1(x), df2(x)). F-distribution, ratio of chi-squared.

ConditionalPoisson module-attribute

ConditionalPoisson = _make_family('ConditionalPoisson', D.Poisson, [('rate', 'softplus')], 'Conditional Poisson(rate(x)). Discrete, non-negative integer counts.')

Conditional Poisson(rate(x)). Discrete, non-negative integer counts.

ConditionalGeometric module-attribute

ConditionalGeometric = _make_family('ConditionalGeometric', D.Geometric, [('probs', 'sigmoid')], 'Conditional Geometric(probs(x)). Number of failures before first success.')

Conditional Geometric(probs(x)). Number of failures before first success.

ConditionalNegativeBinomial module-attribute

ConditionalNegativeBinomial = _make_family('ConditionalNegativeBinomial', D.NegativeBinomial, [('total_count', 'softplus_shifted'), ('probs', 'sigmoid')], 'Conditional NegativeBinomial(total_count(x), probs(x)). Discrete, over-dispersed count distribution.')

Conditional NegativeBinomial(total_count(x), probs(x)). Discrete, over-dispersed count distribution.

ConditionalVonMises module-attribute

ConditionalVonMises = _make_family('ConditionalVonMises', D.VonMises, [('loc', 'id'), ('concentration', 'softplus')], 'Conditional VonMises(loc(x), concentration(x)). Circular analogue of the Normal distribution.', discrete=False)

Conditional VonMises(loc(x), concentration(x)). Circular analogue of the Normal distribution.

ConditionalNormal

ConditionalNormal(domain: AnySpace, codomain: ContinuousSpace, hidden_dim: int | Sequence[int] | None = None, param_source=None, param_source_option: str | None = None)

Bases: ContinuousMorphism

Conditional normal (Gaussian) distribution.

For each input x, produces an independent normal distribution on each dimension of the codomain:

y_i ~ Normal(mu_i(x), sigma_i(x))

Parameters are learnable: mu and log(sigma) are functions of x, implemented as lookup tables (discrete domain) or neural networks (continuous domain).

PARAMETER	DESCRIPTION
domain	Source space. TYPE: SetObject or ContinuousSpace
codomain	Target space. TYPE: Euclidean
hidden_dim	Hidden layer width for neural parameter source. TYPE: int DEFAULT: None

Examples:

>>> from quivers import FinSet
>>> from quivers.continuous.spaces import Euclidean
>>> A = FinSet(name="context", cardinality=5)
>>> Y = Euclidean(name="response", dim=3)
>>> f = ConditionalNormal(A, Y)
>>> x = torch.tensor([0, 1, 2])
>>> samples = f.rsample(x)  # shape (3, 3)

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

def __init__(
    self,
    domain: AnySpace,
    codomain: ContinuousSpace,
    hidden_dim: int | Sequence[int] | None = None,
    param_source=None,
    param_source_option: str | None = None,
) -> None:
    super().__init__(domain, codomain)
    d = codomain.dim
    # param_dim = d (mu) + d (log_sigma)
    self.param_source = _make_source(
        domain,
        2 * d,
        hidden_dim,
        param_source=param_source,
        param_source_option=param_source_option,
    )
    self._d = d

base_dimension

base_dimension(x: Tensor) -> int

One standard-normal coordinate per independent output axis.

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

def base_dimension(self, x: torch.Tensor) -> int:
    """One standard-normal coordinate per independent output axis."""
    mu, _ = self._get_params(dimension_probe(x))
    return _event_size(mu.shape)

push_base

push_base(x: Tensor, base: Tensor) -> Tensor

The location-scale map mu(x) + sigma(x) * base.

The same map rsample applies, with the caller's coordinates in place of a fresh draw.

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

def push_base(self, x: torch.Tensor, base: torch.Tensor) -> torch.Tensor:
    """The location-scale map ``mu(x) + sigma(x) * base``.

    The same map `rsample` applies, with the caller's coordinates
    in place of a fresh draw.
    """
    mu, sigma = self._get_params(x)
    return mu + sigma * base.reshape(mu.shape)

ConditionalLogitNormal

ConditionalLogitNormal(domain: AnySpace, codomain: ContinuousSpace, hidden_dim: int | Sequence[int] | None = None, param_source: ParamSource | None = None, param_source_option: str | None = None)

Bases: ContinuousMorphism

Conditional logit-normal distribution on (0, 1)^d.

If z ~ Normal(mu(x), sigma(x)), then y = sigmoid(z) ~ LogitNormal. Useful for modeling probabilities and bounded quantities.

PARAMETER	DESCRIPTION
domain	Source space. TYPE: SetObject or ContinuousSpace
codomain	Target space (should have bounds [0, 1]). TYPE: Euclidean
hidden_dim	Hidden layer width for neural parameter source. TYPE: int DEFAULT: None

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

def __init__(
    self,
    domain: AnySpace,
    codomain: ContinuousSpace,
    hidden_dim: int | Sequence[int] | None = None,
    param_source: ParamSource | None = None,
    param_source_option: str | None = None,
) -> None:
    super().__init__(domain, codomain)
    d = codomain.dim
    self.param_source = _make_source(
        domain,
        2 * d,
        hidden_dim,
        param_source=param_source,
        param_source_option=param_source_option,
    )
    self._d = d

ConditionalBeta

ConditionalBeta(domain: AnySpace, codomain: ContinuousSpace, hidden_dim: int | Sequence[int] | None = None, param_source: ParamSource | None = None, param_source_option: str | None = None)

Bases: ContinuousMorphism

Conditional beta distribution on (0, 1)^d.

For each input x, produces an independent Beta(alpha_i(x), beta_i(x)) on each dimension of the codomain.

PARAMETER	DESCRIPTION
domain	Source space. TYPE: SetObject or ContinuousSpace
codomain	Target space (should have bounds [0, 1]). TYPE: Euclidean
hidden_dim	Hidden layer width for neural parameter source. TYPE: int DEFAULT: None

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

def __init__(
    self,
    domain: AnySpace,
    codomain: ContinuousSpace,
    hidden_dim: int | Sequence[int] | None = None,
    param_source: ParamSource | None = None,
    param_source_option: str | None = None,
) -> None:
    super().__init__(domain, codomain)
    d = codomain.dim
    self.param_source = _make_source(
        domain,
        2 * d,
        hidden_dim,
        param_source=param_source,
        param_source_option=param_source_option,
    )
    self._d = d

ConditionalTruncatedNormal

ConditionalTruncatedNormal(domain: AnySpace, codomain: Euclidean, hidden_dim: int | Sequence[int] | None = None, param_source: ParamSource | None = None, param_source_option: str | None = None)

Bases: ContinuousMorphism

Conditional truncated normal on [low, high]^d.

A normal distribution restricted to a bounded interval. Uses rejection-free sampling via the inverse CDF (Phi-based) method.

PARAMETER	DESCRIPTION
domain	Source space. TYPE: SetObject or ContinuousSpace
codomain	Target space (must be bounded). TYPE: Euclidean
hidden_dim	Hidden layer width for neural parameter source. TYPE: int DEFAULT: None

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

def __init__(
    self,
    domain: AnySpace,
    codomain: Euclidean,
    hidden_dim: int | Sequence[int] | None = None,
    param_source: ParamSource | None = None,
    param_source_option: str | None = None,
) -> None:
    if codomain.low is None or codomain.high is None:
        raise ValueError("ConditionalTruncatedNormal requires a bounded codomain")

    super().__init__(domain, codomain)
    d = codomain.dim
    self.param_source = _make_source(
        domain,
        2 * d,
        hidden_dim,
        param_source=param_source,
        param_source_option=param_source_option,
    )
    self._d = d
    self._low = codomain.low
    self._high = codomain.high

ConditionalDirichlet

ConditionalDirichlet(domain: AnySpace, codomain: ContinuousSpace, hidden_dim: int | Sequence[int] | None = None, param_source: ParamSource | None = None, param_source_option: str | None = None)

Bases: ContinuousMorphism

Conditional Dirichlet distribution on a probability simplex.

For each input x, produces a Dirichlet(alpha(x)) distribution on the simplex.

PARAMETER	DESCRIPTION
domain	Source space. TYPE: SetObject or ContinuousSpace
codomain	Target simplex. TYPE: Simplex
hidden_dim	Hidden layer width for neural parameter source. TYPE: int DEFAULT: None

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

def __init__(
    self,
    domain: AnySpace,
    codomain: ContinuousSpace,
    hidden_dim: int | Sequence[int] | None = None,
    param_source: ParamSource | None = None,
    param_source_option: str | None = None,
) -> None:
    super().__init__(domain, codomain)
    d = codomain.dim
    self.param_source = _make_source(
        domain,
        d,
        hidden_dim,
        param_source=param_source,
        param_source_option=param_source_option,
    )
    self._d = d

ConditionalUniform

ConditionalUniform(domain: AnySpace, codomain: ContinuousSpace, hidden_dim: int | Sequence[int] | None = None, param_source: ParamSource | None = None, param_source_option: str | None = None)

Bases: ContinuousMorphism

Conditional uniform distribution on a learnable interval.

Parameterized as Uniform(loc - width/2, loc + width/2) where loc is unconstrained and width is positive. This ensures low < high is always satisfied.

PARAMETER	DESCRIPTION
domain	Source space. TYPE: SetObject or ContinuousSpace
codomain	Target space. TYPE: ContinuousSpace
hidden_dim	Hidden layer width for neural parameter source. TYPE: int DEFAULT: None

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

def __init__(
    self,
    domain: AnySpace,
    codomain: ContinuousSpace,
    hidden_dim: int | Sequence[int] | None = None,
    param_source: ParamSource | None = None,
    param_source_option: str | None = None,
) -> None:
    super().__init__(domain, codomain)
    d = codomain.dim
    # param_dim = d (loc) + d (raw_width)
    self.param_source = _make_source(
        domain,
        2 * d,
        hidden_dim,
        param_source=param_source,
        param_source_option=param_source_option,
    )
    self._d = d
    # The bounds are data-dependent, so we cannot pin a single
    # interval at construction time; advertise the codomain's
    # declared bounds when available, otherwise fall back to the
    # real line.
    low, high = getattr(codomain, "low", None), getattr(codomain, "high", None)
    if low is not None and high is not None:
        self._support_cache: _constraints.Constraint = _constraints.interval(
            float(low), float(high)
        )
    else:
        self._support_cache = _constraints.real

ConditionalMultivariateNormal

ConditionalMultivariateNormal(domain: AnySpace, codomain: ContinuousSpace, hidden_dim: int | Sequence[int] | None = None, param_source: ParamSource | None = None, param_source_option: str | None = None)

Bases: ContinuousMorphism

Conditional multivariate normal with full covariance.

Parameterized via Cholesky factor: the parameter source outputs loc (d values) and the lower-triangular entries of L (d*(d+1)/2 values), where Sigma = L @ L^T.

PARAMETER	DESCRIPTION
domain	Source space. TYPE: SetObject or ContinuousSpace
codomain	Target space (d-dimensional). TYPE: ContinuousSpace
hidden_dim	Hidden layer width for neural parameter source. TYPE: int DEFAULT: None

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

def __init__(
    self,
    domain: AnySpace,
    codomain: ContinuousSpace,
    hidden_dim: int | Sequence[int] | None = None,
    param_source: ParamSource | None = None,
    param_source_option: str | None = None,
) -> None:
    super().__init__(domain, codomain)
    d = codomain.dim
    n_tril = d * (d + 1) // 2
    self.param_source = _make_source(
        domain,
        d + n_tril,
        hidden_dim,
        param_source=param_source,
        param_source_option=param_source_option,
    )
    self._d = d
    self._n_tril = n_tril

ConditionalLowRankMVN

ConditionalLowRankMVN(domain: AnySpace, codomain: ContinuousSpace, rank: int = 2, hidden_dim: int | Sequence[int] | None = None, param_source: ParamSource | None = None, param_source_option: str | None = None)

Bases: ContinuousMorphism

Conditional low-rank multivariate normal.

Parameterized as loc + low-rank factor + diagonal: Sigma = W @ W^T + diag(d)

This is more parameter-efficient than full MVN for high dimensions.

PARAMETER	DESCRIPTION
domain	Source space. TYPE: SetObject or ContinuousSpace
codomain	Target space (d-dimensional). TYPE: ContinuousSpace
rank	Rank of the low-rank factor W. TYPE: int DEFAULT: 2
hidden_dim	Hidden layer width for neural parameter source. TYPE: int DEFAULT: None

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

def __init__(
    self,
    domain: AnySpace,
    codomain: ContinuousSpace,
    rank: int = 2,
    hidden_dim: int | Sequence[int] | None = None,
    param_source: ParamSource | None = None,
    param_source_option: str | None = None,
) -> None:
    super().__init__(domain, codomain)
    d = codomain.dim
    self._d = d
    self._rank = rank

    # loc (d) + factor (d * rank) + diag (d)
    total = d + d * rank + d
    self.param_source = _make_source(
        domain,
        total,
        hidden_dim,
        param_source=param_source,
        param_source_option=param_source_option,
    )

ConditionalRelaxedBernoulli

ConditionalRelaxedBernoulli(domain: AnySpace, codomain: ContinuousSpace, temperature: float = 0.5, hidden_dim: int | Sequence[int] | None = None, param_source: ParamSource | None = None, param_source_option: str | None = None)

Bases: ContinuousMorphism

Conditional relaxed Bernoulli (concrete) distribution.

Outputs continuous values in (0, 1) that approximate Bernoulli samples. The temperature controls the relaxation: lower temperature = closer to discrete.

PARAMETER	DESCRIPTION
domain	Source space. TYPE: SetObject or ContinuousSpace
codomain	Target space (should be 1-d per Bernoulli component). TYPE: ContinuousSpace
temperature	Relaxation temperature. TYPE: float DEFAULT: 0.5
hidden_dim	Hidden layer width for neural parameter source. TYPE: int DEFAULT: None

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

def __init__(
    self,
    domain: AnySpace,
    codomain: ContinuousSpace,
    temperature: float = 0.5,
    hidden_dim: int | Sequence[int] | None = None,
    param_source: ParamSource | None = None,
    param_source_option: str | None = None,
) -> None:
    super().__init__(domain, codomain)
    d = codomain.dim
    self.param_source = _make_source(
        domain,
        d,
        hidden_dim,
        param_source=param_source,
        param_source_option=param_source_option,
    )
    self._d = d
    self._temperature = temperature

ConditionalRelaxedOneHotCategorical

ConditionalRelaxedOneHotCategorical(domain: AnySpace, codomain: ContinuousSpace, temperature: float = 0.5, hidden_dim: int | Sequence[int] | None = None, param_source: ParamSource | None = None, param_source_option: str | None = None)

Bases: ContinuousMorphism

Conditional relaxed one-hot categorical (Gumbel-Softmax).

Outputs continuous vectors on the simplex that approximate one-hot categorical samples.

PARAMETER	DESCRIPTION
domain	Source space. TYPE: SetObject or ContinuousSpace
codomain	Target space (simplex or d-dimensional). TYPE: ContinuousSpace
temperature	Relaxation temperature. TYPE: float DEFAULT: 0.5
hidden_dim	Hidden layer width for neural parameter source. TYPE: int DEFAULT: None

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

def __init__(
    self,
    domain: AnySpace,
    codomain: ContinuousSpace,
    temperature: float = 0.5,
    hidden_dim: int | Sequence[int] | None = None,
    param_source: ParamSource | None = None,
    param_source_option: str | None = None,
) -> None:
    super().__init__(domain, codomain)
    d = codomain.dim
    self.param_source = _make_source(
        domain,
        d,
        hidden_dim,
        param_source=param_source,
        param_source_option=param_source_option,
    )
    self._d = d
    self._temperature = temperature

ConditionalWishart

ConditionalWishart(domain: AnySpace, codomain: ContinuousSpace, hidden_dim: int | Sequence[int] | None = None, param_source: ParamSource | None = None, param_source_option: str | None = None)

Bases: ContinuousMorphism

Conditional Wishart distribution over positive-definite matrices.

Produces random d x d positive-definite matrices. Parameterized by degrees of freedom df(x) and a scale matrix V(x).

The codomain dimension is interpreted as d, and outputs are d x d matrices flattened to d^2.

PARAMETER	DESCRIPTION
domain	Source space. TYPE: SetObject or ContinuousSpace
codomain	Target space. dim is the matrix size d (output is d x d). TYPE: ContinuousSpace
hidden_dim	Hidden layer width for neural parameter source. TYPE: int DEFAULT: None

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

def __init__(
    self,
    domain: AnySpace,
    codomain: ContinuousSpace,
    hidden_dim: int | Sequence[int] | None = None,
    param_source: ParamSource | None = None,
    param_source_option: str | None = None,
) -> None:
    super().__init__(domain, codomain)
    d = codomain.dim
    n_tril = d * (d + 1) // 2
    # df (1) + lower-triangular scale (n_tril)
    self.param_source = _make_source(
        domain,
        1 + n_tril,
        hidden_dim,
        param_source=param_source,
        param_source_option=param_source_option,
    )
    self._d = d
    self._n_tril = n_tril

ConditionalMatrixNormal

ConditionalMatrixNormal(domain: AnySpace, codomain: ContinuousSpace, rows: int, cols: int, hidden_dim: int | Sequence[int] | None = None, param_source: ParamSource | None = None, param_source_option: str | None = None)

Bases: ContinuousMorphism

Conditional Matrix-Normal :math:\mathcal{MN}(M, U, V).

The matrix-Normal distribution on :math:\mathbb{R}^{n \times p} factorises with a Kronecker-product covariance: if :math:X \sim \mathcal{MN}(M, U, V) then :math:\mathrm{vec}(X) \sim \mathcal{N}(\mathrm{vec}(M), V \otimes U) with :math:U \in \mathbb{R}^{n \times n} the row covariance and :math:V \in \mathbb{R}^{p \times p} the column covariance.

Categorically, the Kronecker structure is the tensor product of two Gaussians; it is strictly more constrained than the flat :math:np-dim MVN that the same parameter tensor could carry, so the surface keeps the two families distinct (no auto-substitution). Use this when the prior should express independent row and column correlation structure.

The codomain's product factorisation supplies the row and column dimensions. The grammar surface ~ MatrixNormal(loc, row_scale, col_scale) over (dom, cod) binds the first axis listed in over to the row covariance and the second to the column covariance.

PARAMETER	DESCRIPTION
domain	Source space. TYPE: AnySpace
codomain	Target space whose factorisation supplies (n, p). Must carry a product structure of two factors; the first is the row axis, the second the column. TYPE: ContinuousSpace
rows	Row dimension :math:n. TYPE: int
cols	Column dimension :math:p. TYPE: int
hidden_dim	Hidden layer width for the parameter network. TYPE: int DEFAULT: None

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

def __init__(
    self,
    domain: AnySpace,
    codomain: ContinuousSpace,
    rows: int,
    cols: int,
    hidden_dim: int | Sequence[int] | None = None,
    param_source: ParamSource | None = None,
    param_source_option: str | None = None,
) -> None:
    super().__init__(domain, codomain)
    self._rows = int(rows)
    self._cols = int(cols)
    n_loc = self._rows * self._cols
    n_row_tril = self._rows * (self._rows + 1) // 2
    n_col_tril = self._cols * (self._cols + 1) // 2
    self._n_loc = n_loc
    self._n_row_tril = n_row_tril
    self._n_col_tril = n_col_tril
    self.param_source = _make_source(
        domain,
        n_loc + n_row_tril + n_col_tril,
        hidden_dim,
        param_source=param_source,
        param_source_option=param_source_option,
    )

ConditionalInverseWishart

ConditionalInverseWishart(domain: AnySpace, codomain: ContinuousSpace, hidden_dim: int | Sequence[int] | None = None, param_source: ParamSource | None = None, param_source_option: str | None = None)

Bases: ContinuousMorphism

Conditional Inverse-Wishart over positive-definite matrices.

Conjugate prior for the covariance of a multivariate normal. Realised as a deterministic inversion of a Wishart sample so autograd flows; equivalent in distribution to drawing :math:\Sigma^{-1} \sim \mathcal{W}(\nu, V^{-1}) and inverting. See Gelman et al. (2013) §3.6 for the conjugacy statement.

PARAMETER	DESCRIPTION
domain	Source space. TYPE: AnySpace
codomain	Target space whose dim is the matrix size :math:d. TYPE: ContinuousSpace
hidden_dim	Hidden layer width for the parameter network. TYPE: int DEFAULT: None

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

def __init__(
    self,
    domain: AnySpace,
    codomain: ContinuousSpace,
    hidden_dim: int | Sequence[int] | None = None,
    param_source: ParamSource | None = None,
    param_source_option: str | None = None,
) -> None:
    super().__init__(domain, codomain)
    d = codomain.dim
    n_tril = d * (d + 1) // 2
    self._d = d
    self._n_tril = n_tril
    self.param_source = _make_source(
        domain,
        1 + n_tril,
        hidden_dim,
        param_source=param_source,
        param_source_option=param_source_option,
    )

ConditionalBernoulli

ConditionalBernoulli(domain: AnySpace, codomain: AnySpace, hidden_dim: int | Sequence[int] | None = None, param_source: ParamSource | None = None, param_source_option: str | None = None)

Bases: ContinuousMorphism

Conditional Bernoulli: continuous probability -> discrete truth value.

Takes a continuous input x and produces learnable logits that parameterize a Bernoulli distribution. The output is a discrete sample in {0, 1}, returned as a LongTensor.

This is the key bridge used in PDS (Grove & White) for the Bern x pattern, where a LogitNormal draw x in (0,1) parameterizes a Bernoulli over truth values.

The codomain must be a FinSet of size 2 (representing {False, True} or {0, 1}).

Note

Sampling from Bernoulli is NOT reparameterizable. Gradients do not flow through the discrete samples. Use score function estimators (REINFORCE) or the Gumbel-Softmax trick if differentiable samples are needed.

PARAMETER	DESCRIPTION
domain	Source space (typically UnitInterval or a FinSet). TYPE: SetObject or ContinuousSpace
codomain	Target FinSet of size 2. TYPE: SetObject
hidden_dim	Hidden layer width for neural parameter source. TYPE: int DEFAULT: None

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

def __init__(
    self,
    domain: AnySpace,
    codomain: AnySpace,
    hidden_dim: int | Sequence[int] | None = None,
    param_source: ParamSource | None = None,
    param_source_option: str | None = None,
) -> None:
    from quivers.core.objects import SetObject

    if not isinstance(codomain, SetObject) or codomain.size != 2:
        raise ValueError(
            f"ConditionalBernoulli requires a FinSet(2) codomain, got {codomain!r}"
        )

    super().__init__(domain, codomain)

    # one logit per input
    self.param_source = _make_source(
        domain,
        1,
        hidden_dim,
        param_source=param_source,
        param_source_option=param_source_option,
    )

log_prob

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

Log-probability of discrete output y given input x.

PARAMETER	DESCRIPTION
x	Input tensor. TYPE: Tensor
y	Discrete output in {0, 1}. Shape (batch,). TYPE: Tensor

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

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

def log_prob(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
    """Log-probability of discrete output y given input x.

    Parameters
    ----------
    x : torch.Tensor
        Input tensor.
    y : torch.Tensor
        Discrete output in {0, 1}. Shape (batch,).

    Returns
    -------
    torch.Tensor
        Log-probabilities. Shape (batch,).
    """
    probs = self._get_probs(x)
    dist = D.Bernoulli(probs=probs)
    return dist.log_prob(y.float())

rsample

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

Sample from Bernoulli (not reparameterizable).

PARAMETER	DESCRIPTION
x	Input tensor. TYPE: Tensor
sample_shape	Additional leading sample dimensions. TYPE: Size DEFAULT: Size()

RETURNS	DESCRIPTION
Tensor	Discrete samples in {0, 1}. Shape (*sample_shape, batch).

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

def rsample(
    self,
    x: torch.Tensor,
    sample_shape: torch.Size = torch.Size(),
) -> torch.Tensor:
    """Sample from Bernoulli (not reparameterizable).

    Parameters
    ----------
    x : torch.Tensor
        Input tensor.
    sample_shape : torch.Size
        Additional leading sample dimensions.

    Returns
    -------
    torch.Tensor
        Discrete samples in {0, 1}. Shape (*sample_shape, batch).
    """
    probs = self._get_probs(x)
    dist = D.Bernoulli(probs=probs)
    return dist.sample(sample_shape).long()

ConditionalCategorical

ConditionalCategorical(domain: AnySpace, codomain: AnySpace, hidden_dim: int | Sequence[int] | None = None, param_source: ParamSource | None = None, param_source_option: str | None = None)

Bases: ContinuousMorphism

Conditional Categorical: continuous input -> discrete category.

Generalizes ConditionalBernoulli to k > 2 categories. Takes a continuous input and produces learnable logits over k categories. The output is a discrete sample in {0, ..., k-1}.

PARAMETER	DESCRIPTION
domain	Source space. TYPE: SetObject or ContinuousSpace
codomain	Target FinSet of size k. TYPE: SetObject
hidden_dim	Hidden layer width for neural parameter source. TYPE: int DEFAULT: None

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

def __init__(
    self,
    domain: AnySpace,
    codomain: AnySpace,
    hidden_dim: int | Sequence[int] | None = None,
    param_source: ParamSource | None = None,
    param_source_option: str | None = None,
) -> None:
    from quivers.core.objects import SetObject

    if not isinstance(codomain, SetObject):
        raise ValueError(
            f"ConditionalCategorical requires a FinSet codomain, got {codomain!r}"
        )

    super().__init__(domain, codomain)
    self._k = codomain.size
    self.param_source = _make_source(
        domain,
        self._k,
        hidden_dim,
        param_source=param_source,
        param_source_option=param_source_option,
    )

support property

support

Discrete-integer support over {0, …, k-1}.

log_prob

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

Log-probability of discrete output y given input x.

PARAMETER	DESCRIPTION
x	Input tensor. TYPE: Tensor
y	Discrete output in {0, ..., k-1}. Shape (batch,). TYPE: Tensor

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

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

def log_prob(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
    """Log-probability of discrete output y given input x.

    Parameters
    ----------
    x : torch.Tensor
        Input tensor.
    y : torch.Tensor
        Discrete output in {0, ..., k-1}. Shape (batch,).

    Returns
    -------
    torch.Tensor
        Log-probabilities. Shape (batch,).
    """
    logits = self._get_logits(x)
    dist = D.Categorical(logits=logits)
    return dist.log_prob(y.long())

rsample

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

Sample from Categorical (not reparameterizable).

PARAMETER	DESCRIPTION
x	Input tensor. TYPE: Tensor
sample_shape	Additional leading sample dimensions. TYPE: Size DEFAULT: Size()

RETURNS	DESCRIPTION
Tensor	Discrete samples in {0, ..., k-1}. Shape (*sample_shape, batch).

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

def rsample(
    self,
    x: torch.Tensor,
    sample_shape: torch.Size = torch.Size(),
) -> torch.Tensor:
    """Sample from Categorical (not reparameterizable).

    Parameters
    ----------
    x : torch.Tensor
        Input tensor.
    sample_shape : torch.Size
        Additional leading sample dimensions.

    Returns
    -------
    torch.Tensor
        Discrete samples in {0, ..., k-1}. Shape (*sample_shape, batch).
    """
    logits = self._get_logits(x)
    dist = D.Categorical(logits=logits)
    return dist.sample(sample_shape).long()

ConditionalBinomial

ConditionalBinomial(domain: AnySpace, codomain: ContinuousSpace, total_count: int = 1, hidden_dim: int | Sequence[int] | None = None, param_source: ParamSource | None = None, param_source_option: str | None = None)

Bases: ContinuousMorphism

Conditional Binomial(total_count, probs(x)).

The total_count (number of trials) is a fixed hyperparameter set at construction time, typical for binomial likelihoods where n is known per observation. Only the probs parameter is learnable.

Outputs integer counts in {0, 1, ..., total_count}.

PARAMETER	DESCRIPTION
domain	Source space. TYPE: SetObject or ContinuousSpace
codomain	Target space. TYPE: ContinuousSpace
total_count	Number of Bernoulli trials per observation. TYPE: int DEFAULT: 1
hidden_dim	Hidden layer width for the parameter source. TYPE: int DEFAULT: None

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

def __init__(
    self,
    domain: AnySpace,
    codomain: ContinuousSpace,
    total_count: int = 1,
    hidden_dim: int | Sequence[int] | None = None,
    param_source: ParamSource | None = None,
    param_source_option: str | None = None,
) -> None:
    if total_count < 1:
        raise ValueError(
            f"ConditionalBinomial: total_count must be >= 1, got {total_count}"
        )
    super().__init__(domain, codomain)
    d = codomain.dim
    self._d = d
    self._total_count = int(total_count)
    self.param_source = _make_source(
        domain,
        d,
        hidden_dim,
        param_source=param_source,
        param_source_option=param_source_option,
    )

ConditionalLogisticNormal

ConditionalLogisticNormal(domain: AnySpace, codomain: ContinuousSpace, hidden_dim: int | Sequence[int] | None = None, param_source: ParamSource | None = None, param_source_option: str | None = None)

Bases: ContinuousMorphism

Conditional LogisticNormal on the simplex.

Pushes a Normal(loc(x), scale(x)) draw through the softmax transform to produce a simplex-valued sample. Multivariate analogue of ConditionalLogitNormal. Useful as an alternative to ConditionalDirichlet when the underlying simplex distribution should be Gaussian in logit space rather than Beta-shaped.

PARAMETER	DESCRIPTION
domain	Source space. TYPE: SetObject or ContinuousSpace
codomain	Target space; codomain.dim is the simplex dimension. TYPE: ContinuousSpace
hidden_dim	Hidden layer width for the parameter source. TYPE: int DEFAULT: None

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

def __init__(
    self,
    domain: AnySpace,
    codomain: ContinuousSpace,
    hidden_dim: int | Sequence[int] | None = None,
    param_source: ParamSource | None = None,
    param_source_option: str | None = None,
) -> None:
    super().__init__(domain, codomain)
    d = codomain.dim
    # We use a Normal in (d-1)-dim space and the
    # StickBreakingTransform to land on the d-simplex.
    # torch.distributions.LogisticNormal handles this.
    self.param_source = _make_source(
        domain,
        2 * (d - 1),
        hidden_dim,
        param_source=param_source,
        param_source_option=param_source_option,
    )
    self._d = d

ConditionalZeroInflatedPoisson

ConditionalZeroInflatedPoisson(domain: AnySpace, codomain: ContinuousSpace, hidden_dim: int | Sequence[int] | None = None, param_source: ParamSource | None = None, param_source_option: str | None = None)

Bases: ContinuousMorphism

Conditional ZeroInflatedPoisson(zero_prob(x), rate(x)).

The parameter source produces 2 * codomain.dim numbers per input row: the first half feeds a sigmoid to produce the zero-inflation probability, the second half a softplus to produce the Poisson rate. Outputs non-negative integer counts.

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

def __init__(
    self,
    domain: AnySpace,
    codomain: ContinuousSpace,
    hidden_dim: int | Sequence[int] | None = None,
    param_source: ParamSource | None = None,
    param_source_option: str | None = None,
) -> None:
    super().__init__(domain, codomain)
    d = codomain.dim
    self._d = d
    self.param_source = _make_source(
        domain,
        2 * d,
        hidden_dim,
        param_source=param_source,
        param_source_option=param_source_option,
    )

ConditionalHurdlePoisson

ConditionalHurdlePoisson(domain: AnySpace, codomain: ContinuousSpace, hidden_dim: int | Sequence[int] | None = None, param_source: ParamSource | None = None, param_source_option: str | None = None)

Bases: ContinuousMorphism

Conditional HurdlePoisson(zero_prob(x), rate(x)).

Same parameter shape as ConditionalZeroInflatedPoisson but the two-stage hurdle log-density: a Bernoulli for zero vs positive, then a zero-truncated Poisson for the strictly-positive branch.

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

def __init__(
    self,
    domain: AnySpace,
    codomain: ContinuousSpace,
    hidden_dim: int | Sequence[int] | None = None,
    param_source: ParamSource | None = None,
    param_source_option: str | None = None,
) -> None:
    super().__init__(domain, codomain)
    d = codomain.dim
    self._d = d
    self.param_source = _make_source(
        domain,
        2 * d,
        hidden_dim,
        param_source=param_source,
        param_source_option=param_source_option,
    )

ConditionalZeroOneInflatedBeta

ConditionalZeroOneInflatedBeta(domain: AnySpace, codomain: ContinuousSpace, hidden_dim: int | Sequence[int] | None = None, param_source: ParamSource | None = None, param_source_option: str | None = None)

Bases: ContinuousMorphism

Conditional ZeroOneInflatedBeta(mu(x), phi(x), zoi(x), coi(x)).

Parameter source emits 4 * dim numbers per row: logit-mu, log-phi, logit-zoi, logit-coi. Used by the runtime family registry for ~ ZeroOneInflatedBeta declarations; formula frontend prefers the inline observe ... <- ZeroOneInflatedBeta(mu, phi, zoi, coi).

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

def __init__(
    self,
    domain: AnySpace,
    codomain: ContinuousSpace,
    hidden_dim: int | Sequence[int] | None = None,
    param_source: ParamSource | None = None,
    param_source_option: str | None = None,
) -> None:
    super().__init__(domain, codomain)
    d = codomain.dim
    self._d = d
    self.param_source = _make_source(
        domain,
        4 * d,
        hidden_dim,
        param_source=param_source,
        param_source_option=param_source_option,
    )

ConditionalMixtureNormal

ConditionalMixtureNormal(domain: AnySpace, codomain: ContinuousSpace, num_components: int = 2, hidden_dim: int | Sequence[int] | None = None, param_source: ParamSource | None = None, param_source_option: str | None = None)

Bases: ContinuousMorphism

Conditional finite Gaussian mixture with input-driven weights, locations, and scales.

The parameter source produces 3 * K numbers per row, where K is the number of mixture components (fixed at construction). Weights pass through softmax, locations are emitted directly, scales pass through softplus + epsilon.

The codomain is assumed scalar (1-d real); a higher-dimensional extension would replace Normal with Independent(Normal(...), 1) and triple the per-component parameter count.

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

def __init__(
    self,
    domain: AnySpace,
    codomain: ContinuousSpace,
    num_components: int = 2,
    hidden_dim: int | Sequence[int] | None = None,
    param_source: ParamSource | None = None,
    param_source_option: str | None = None,
) -> None:
    if num_components < 2:
        raise ValueError(
            "ConditionalMixtureNormal: num_components must be >= 2, "
            f"got {num_components}"
        )
    super().__init__(domain, codomain)
    self._k = int(num_components)
    self.param_source = _make_source(
        domain,
        3 * self._k,
        hidden_dim,
        param_source=param_source,
        param_source_option=param_source_option,
    )

ConditionalOneHotCategorical

ConditionalOneHotCategorical(domain: AnySpace, codomain: ContinuousSpace, hidden_dim: int | Sequence[int] | None = None, param_source: ParamSource | None = None, param_source_option: str | None = None)

Bases: ContinuousMorphism

Conditional OneHotCategorical(probs(x)).

Generalises ConditionalCategorical to one-hot encoded outputs (vector of zeros with a single one). Useful as a discrete-output observation kernel where downstream code wants a vector rather than an integer index.

PARAMETER	DESCRIPTION
domain	Source space. TYPE: SetObject or ContinuousSpace
codomain	Target space; codomain.dim is the number of categories. TYPE: ContinuousSpace
hidden_dim	Hidden layer width for the parameter source. TYPE: int DEFAULT: None

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

def __init__(
    self,
    domain: AnySpace,
    codomain: ContinuousSpace,
    hidden_dim: int | Sequence[int] | None = None,
    param_source: ParamSource | None = None,
    param_source_option: str | None = None,
) -> None:
    super().__init__(domain, codomain)
    d = codomain.dim
    self._d = d
    self.param_source = _make_source(
        domain,
        d,
        hidden_dim,
        param_source=param_source,
        param_source_option=param_source_option,
    )

ConditionalMixture

ConditionalMixture(domain: AnySpace, codomain: ContinuousSpace, component_class: type, num_components: int = 4, hidden_dim: int | Sequence[int] | None = None, param_source: ParamSource | None = None, param_source_option: str | None = None)

Bases: ContinuousMorphism

K-component mixture of a conditional family.

Wraps a single conditional family class (one of the registered ConditionalX types) and gives it K independent parameterizations plus learnable mixture logits, producing

p(y | x) = sum_k pi_k(x) * f_k(y | x)

where each f_k is an instance of the component class and pi is the softmax of K learnable logits.

Sampling is via ancestral simulation (Categorical pick + the chosen component's rsample). log_prob evaluates the log-sum-exp of the per-component log-densities. The Categorical pick is non-reparameterizable; gradient flow through the weights uses the score-function path (which higher-level objectives like IWAE can route through).

PARAMETER	DESCRIPTION
domain	Source space. TYPE: SetObject or ContinuousSpace
codomain	Target space; matches the component family's codomain. TYPE: ContinuousSpace
component_class	A ConditionalX class accepting (domain, codomain, hidden_dim) constructor args. TYPE: type
num_components	Number of mixture components. TYPE: int DEFAULT: 4
hidden_dim	Hidden width for both the mixture-logit MLP and each component's parameter source. TYPE: int DEFAULT: None

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

def __init__(
    self,
    domain: AnySpace,
    codomain: ContinuousSpace,
    component_class: type,
    num_components: int = 4,
    hidden_dim: int | Sequence[int] | None = None,
    param_source: ParamSource | None = None,
    param_source_option: str | None = None,
) -> None:
    if num_components < 2:
        raise ValueError(
            f"ConditionalMixture: num_components must be >= 2, got {num_components}"
        )
    super().__init__(domain, codomain)
    self._K = int(num_components)
    self._components = torch.nn.ModuleList(
        [component_class(domain, codomain, hidden_dim) for _ in range(self._K)]
    )
    self.mixture_logits = _make_source(
        domain,
        self._K,
        hidden_dim,
        param_source=param_source,
        param_source_option=param_source_option,
    )

ConditionalIndependent

ConditionalIndependent(base: ContinuousMorphism)

Bases: ContinuousMorphism

Reinterpret the trailing batch dimension of a base conditional family as an event dimension.

Equivalent to wrapping the base distribution in torch.distributions.Independent with reinterpreted_batch_ndims = 1. Used to make per-element independence explicit when a downstream guide wants to score a vector-valued observation as a single event.

PARAMETER	DESCRIPTION
base	The base conditional family. The wrapped distribution sums the base's per-element log-probabilities along the last axis to score a vector-valued observation. TYPE: ContinuousMorphism

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

def __init__(self, base: ContinuousMorphism) -> None:
    super().__init__(base.domain, base.codomain)
    self._base = base

ConditionalTransformed

ConditionalTransformed(base: ContinuousMorphism, transforms: list)

Bases: ContinuousMorphism

A base conditional family composed with a chain of bijectors.

Equivalent to torch.distributions.TransformedDistribution lifted to ContinuousMorphism. The transforms are applied in forward order to rsample outputs; log_prob includes the log-determinant Jacobian correction.

PARAMETER	DESCRIPTION
base	Base conditional family. TYPE: ContinuousMorphism
transforms	Bijectors applied in forward order. Each must implement __call__, inv, and log_abs_det_jacobian. TYPE: list of torch.distributions.Transform

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

def __init__(
    self,
    base: ContinuousMorphism,
    transforms: list,
) -> None:
    super().__init__(base.domain, base.codomain)
    self._base = base
    self._transforms = list(transforms)

ConditionalLKJCholesky

ConditionalLKJCholesky(domain: AnySpace, codomain: ContinuousSpace, hidden_dim: int | Sequence[int] | None = None, param_source: ParamSource | None = None, param_source_option: str | None = None)

Bases: ContinuousMorphism

Conditional LKJCholesky(dim, concentration(x)).

Produces lower-triangular Cholesky factors of correlation matrices on the LKJ distribution (Lewandowski-Kurowicka-Joe 2009, doi:10.1016/j.jmva.2009.04.008). The matrix dimension is taken from codomain.dim; only the concentration parameter is learnable.

PARAMETER	DESCRIPTION
domain	Source space. TYPE: SetObject or ContinuousSpace
codomain	Target space; codomain.dim is the correlation-matrix size. TYPE: ContinuousSpace
hidden_dim	Hidden layer width for the parameter source. TYPE: int DEFAULT: None

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

def __init__(
    self,
    domain: AnySpace,
    codomain: ContinuousSpace,
    hidden_dim: int | Sequence[int] | None = None,
    param_source: ParamSource | None = None,
    param_source_option: str | None = None,
) -> None:
    super().__init__(domain, codomain)
    self._matrix_dim = codomain.dim
    self.param_source = _make_source(
        domain,
        1,
        hidden_dim,
        param_source=param_source,
        param_source_option=param_source_option,
    )

ConditionalGaussianProcess

ConditionalGaussianProcess(domain: AnySpace, codomain: ContinuousSpace, kernel: str = 'rbf', length_scale: float = 1.0, amplitude: float = 1.0, jitter: float = 1e-06)

Bases: ContinuousMorphism

Gaussian process prior with mean zero and a chosen covariance kernel.

A Gaussian process is a Markov kernel X^N -> G(R^N) whose value at a finite set of input locations x_1, ..., x_N follows a multivariate Normal with covariance matrix K(x_i, x_j). Unlike the parametric families that derive their distribution parameters from a neural network on the input, the GP's "parameters" are the input locations themselves: the kernel function evaluated on the inputs produces the covariance directly.

Reference: Rasmussen & Williams (2006), Gaussian Processes for Machine Learning.

PARAMETER	DESCRIPTION
domain	Source space. Its dim is the per-location feature dimensionality D of the inputs. TYPE: SetObject or ContinuousSpace
codomain	Target space. Its dim is the number of input locations N at which the GP is evaluated. TYPE: ContinuousSpace
kernel	Covariance kernel. "rbf" is the squared-exponential kernel; "matern52" is the Matern kernel with smoothness nu = 5/2; "linear" is the inner-product kernel. TYPE: ('rbf', 'matern52', 'linear') DEFAULT: "rbf"
length_scale	Initial length scale of the kernel (positive; learnable). Ignored by the linear kernel. TYPE: float DEFAULT: 1.0
amplitude	Initial amplitude (positive; learnable). Multiplies the kernel by amplitude^2. TYPE: float DEFAULT: 1.0
jitter	Diagonal regulariser added to K for numerical positive-definiteness of the Cholesky factorisation. TYPE: float DEFAULT: 1e-06

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

def __init__(
    self,
    domain: AnySpace,
    codomain: ContinuousSpace,
    kernel: str = "rbf",
    length_scale: float = 1.0,
    amplitude: float = 1.0,
    jitter: float = 1e-6,
) -> None:
    super().__init__(domain, codomain)
    if kernel not in _GP_KERNEL_CHOICES:
        raise ValueError(
            f"ConditionalGaussianProcess: unknown kernel {kernel!r}; "
            f"valid choices: {_GP_KERNEL_CHOICES}"
        )
    if length_scale <= 0.0:
        raise ValueError(
            f"ConditionalGaussianProcess: length_scale must be > 0, got {length_scale!r}"
        )
    if amplitude <= 0.0:
        raise ValueError(
            f"ConditionalGaussianProcess: amplitude must be > 0, got {amplitude!r}"
        )
    self._kernel = kernel
    self._jitter = jitter
    self._n = codomain.dim
    self._d = getattr(domain, "dim", None)
    # Store raw (pre-softplus) parameters so the transformed value
    # is strictly positive and the optimiser sees unconstrained
    # variables.
    inv_softplus_ls = math.log(math.expm1(length_scale))
    inv_softplus_amp = math.log(math.expm1(amplitude))
    self._raw_length_scale = torch.nn.Parameter(
        torch.tensor(inv_softplus_ls, dtype=torch.get_default_dtype())
    )
    self._raw_amplitude = torch.nn.Parameter(
        torch.tensor(inv_softplus_amp, dtype=torch.get_default_dtype())
    )

length_scale property

length_scale: Tensor

Current (positive) length scale of the kernel.

amplitude property

amplitude: Tensor

Current (positive) amplitude of the kernel.

ConditionalHorseshoe

ConditionalHorseshoe(domain: AnySpace, codomain: ContinuousSpace, scale: float = 1.0)

Bases: ContinuousMorphism

Carvalho-Polson-Scott horseshoe prior.

The horseshoe prior places a global-local shrinkage structure on each coordinate:

.. code-block:: text

tau ~ HalfCauchy(scale) lambda_d ~ HalfCauchy(1) for d = 1, ..., D beta_d | tau, lambda_d ~ Normal(0, (tau * lambda_d)^2)

The marginal density of beta_d after integrating the local scale lambda_d has no closed form; this implementation uses a 16-point Gauss-Legendre quadrature after mapping the half-line lambda in (0, inf) to t in (0, 1) via the change of variables lambda = tan(pi * t / 2), whose Jacobian is (pi / 2) * sec^2(pi * t / 2).

Reference: Carvalho, Polson & Scott (2010), The horseshoe estimator for sparse signals.

PARAMETER	DESCRIPTION
domain	Source space. The prior is conditionally independent of x; x only carries the batch shape. TYPE: SetObject or ContinuousSpace
codomain	Target space. Its dim is the coordinate count D. TYPE: ContinuousSpace
scale	Initial global shrinkage tau (positive; learnable). TYPE: float DEFAULT: 1.0

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

def __init__(
    self,
    domain: AnySpace,
    codomain: ContinuousSpace,
    scale: float = 1.0,
) -> None:
    super().__init__(domain, codomain)
    if scale <= 0.0:
        raise ValueError(f"ConditionalHorseshoe: scale must be > 0, got {scale!r}")
    self._d = codomain.dim
    inv_softplus_scale = math.log(math.expm1(scale))
    self._raw_scale = torch.nn.Parameter(
        torch.tensor(inv_softplus_scale, dtype=torch.get_default_dtype())
    )

scale property

scale: Tensor

Current (positive) global shrinkage tau.

ConditionalLearnedOrderedLogistic

ConditionalLearnedOrderedLogistic(domain: AnySpace, codomain: AnySpace, hidden_dim: int | Sequence[int] | None = None, param_source: ParamSource | None = None, param_source_option: str | None = None)

Bases: ContinuousMorphism

Conditional OrderedLogistic(predictor(x), cutpoints(x)), cutpoints learned.

This is the morphism a ~ OrderedLogistic declaration compiles to. ConditionalOrderedLogistic is the variant that takes its cutpoints explicitly.

The continuous input drives a parameter source that produces 1 + (K - 1) numbers: one real predictor and a K - 1 cutpoint vector. The cutpoints are passed through a strictly-monotonic transform (first entry free, subsequent entries via cumulative softplus) so the cumulative-link contract c_0 < c_1 < ... < c_{K-2} is satisfied unconditionally.

Outputs integer categories in {0, ..., K - 1} where K = codomain.size.

PARAMETER	DESCRIPTION
domain	Source space. TYPE: AnySpace
codomain	A finite set of at least two ordered categories. TYPE: AnySpace
hidden_dim	Hidden widths, read only by a source that has hidden layers. TYPE: int or sequence of int DEFAULT: None
param_source	Select the parameter source directly, or by the DSL's [param_source=...] text. TYPE: ParamSource | None DEFAULT: None
param_source_option	Select the parameter source directly, or by the DSL's [param_source=...] text. TYPE: ParamSource | None DEFAULT: None

RAISES	DESCRIPTION
ValueError	If the codomain is not a finite set of at least two categories.

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

def __init__(
    self,
    domain: AnySpace,
    codomain: AnySpace,
    hidden_dim: int | Sequence[int] | None = None,
    param_source: ParamSource | None = None,
    param_source_option: str | None = None,
) -> None:
    if not isinstance(codomain, SetObject):
        raise ValueError(
            "ConditionalLearnedOrderedLogistic requires a FinSet codomain, "
            f"got {codomain!r}"
        )
    if codomain.size < 2:
        raise ValueError(
            "ConditionalLearnedOrderedLogistic requires K >= 2 categories, "
            f"got codomain.size = {codomain.size}"
        )
    super().__init__(domain, codomain)
    self._k = codomain.size
    self.param_source = _make_source(
        domain,
        1 + (self._k - 1),
        hidden_dim,
        param_source=param_source,
        param_source_option=param_source_option,
    )

ConditionalGeneralizedPareto

ConditionalGeneralizedPareto(domain: AnySpace, codomain: ContinuousSpace, hidden_dim: int | Sequence[int] | None = None, param_source: ParamSource | None = None, param_source_option: str | None = None)

Bases: ContinuousMorphism

Conditional generalized Pareto distribution.

PARAMETER	DESCRIPTION
domain	Source space. TYPE: SetObject or ContinuousSpace
codomain	Target space. TYPE: ContinuousSpace
hidden_dim	Hidden widths, read only by a source that has hidden layers. TYPE: int or sequence of int DEFAULT: None
param_source	Select the parameter source directly, or by the DSL's [param_source=...] text. TYPE: ParamSource | None DEFAULT: None
param_source_option	Select the parameter source directly, or by the DSL's [param_source=...] text. TYPE: ParamSource | None DEFAULT: None

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

def __init__(
    self,
    domain: AnySpace,
    codomain: ContinuousSpace,
    hidden_dim: int | Sequence[int] | None = None,
    param_source: ParamSource | None = None,
    param_source_option: str | None = None,
) -> None:
    super().__init__(domain, codomain)
    d = codomain.dim
    # loc + scale + concentration
    self.param_source = _make_source(
        domain,
        3 * d,
        hidden_dim,
        param_source=param_source,
        param_source_option=param_source_option,
    )
    self._d = d

ConditionalBetaBinomial

ConditionalBetaBinomial(domain: AnySpace, codomain: ContinuousSpace, total_count: int, hidden_dim: int | Sequence[int] | None = None, param_source: ParamSource | None = None, param_source_option: str | None = None)

Bases: ContinuousMorphism

Conjugate Beta-Binomial likelihood with learnable Beta parameters.

The number of successes y arises by drawing p ~ Beta(concentration1(x), concentration0(x)) and then y ~ Binomial(total_count, p). The Beta is integrated out analytically, yielding the log-probability

.. math::

\log p(y \mid x) = \log \Gamma(\alpha + \beta)
    + \log \Gamma(\alpha + y) + \log \Gamma(\beta + n - y)
    - \log \Gamma(\alpha) - \log \Gamma(\beta)
    - \log \Gamma(\alpha + \beta + n)
    + \log \binom{n}{y}

where :math:\alpha = \mathrm{concentration1}(x), :math:\beta = \mathrm{concentration0}(x), and :math:n is total_count.

Sampling is not reparameterisable; rsample raises and sample draws p from the Beta and y from the resulting Binomial.

PARAMETER	DESCRIPTION
domain	Source space. TYPE: SetObject or ContinuousSpace
codomain	Target space. codomain.dim independent Beta-Binomial observations are produced per input. TYPE: ContinuousSpace
total_count	Number of Bernoulli trials per observation; must be a positive integer. TYPE: int
hidden_dim	Hidden layer width for the neural parameter source. TYPE: int DEFAULT: None

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

def __init__(
    self,
    domain: AnySpace,
    codomain: ContinuousSpace,
    total_count: int,
    hidden_dim: int | Sequence[int] | None = None,
    param_source: ParamSource | None = None,
    param_source_option: str | None = None,
) -> None:
    if total_count < 1:
        raise ValueError(
            f"ConditionalBetaBinomial: total_count must be >= 1, got {total_count}"
        )
    super().__init__(domain, codomain)
    d = codomain.dim
    self._d = d
    self._total_count = int(total_count)
    # Two raw scalars per codomain dim: alpha and beta concentrations.
    self.param_source = _make_source(
        domain,
        2 * d,
        hidden_dim,
        param_source=param_source,
        param_source_option=param_source_option,
    )

ConditionalLogistic

ConditionalLogistic(domain: AnySpace, codomain: ContinuousSpace, hidden_dim: int | Sequence[int] | None = None, param_source: ParamSource | None = None, param_source_option: str | None = None)

Bases: ContinuousMorphism

Conditional logistic distribution on the real line.

Independent per codomain dim with learnable location and scale:

.. math::

\log p(y \mid x) = -z - 2 \log(1 + e^{-z}) - \log \sigma

where :math:z = (y - \mu) / \sigma. Samples are reparameterised via the inverse-CDF transform of a uniform draw, :math:y = \mu + \sigma \log(u / (1 - u)).

PARAMETER	DESCRIPTION
domain	Source space. TYPE: SetObject or ContinuousSpace
codomain	Target space. codomain.dim independent logistic components are produced per input. TYPE: ContinuousSpace
hidden_dim	Hidden layer width for the neural parameter source. TYPE: int DEFAULT: None

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

def __init__(
    self,
    domain: AnySpace,
    codomain: ContinuousSpace,
    hidden_dim: int | Sequence[int] | None = None,
    param_source: ParamSource | None = None,
    param_source_option: str | None = None,
) -> None:
    super().__init__(domain, codomain)
    d = codomain.dim
    self._d = d
    # param_dim = d (loc) + d (raw_scale).
    self.param_source = _make_source(
        domain,
        2 * d,
        hidden_dim,
        param_source=param_source,
        param_source_option=param_source_option,
    )

ConditionalHalfStudentT

ConditionalHalfStudentT(domain: AnySpace, codomain: ContinuousSpace, df: float, hidden_dim: int | Sequence[int] | None = None, param_source: ParamSource | None = None, param_source_option: str | None = None)

Bases: ContinuousMorphism

Conditional half-StudentT on the positive reals.

A StudentT with df fixed at construction time, learnable scale, and location fixed at zero, restricted to nonnegative values via reflection. The folded log-density at y >= 0 is log 2 + log_prob_studentt(0, scale(x), df; y) and -inf elsewhere.

Sampling reflects a base StudentT draw through the origin: y = |z| with z ~ StudentT(df, 0, scale). Reparameterisation flows through the absolute-value operation.

PARAMETER	DESCRIPTION
domain	Source space. TYPE: SetObject or ContinuousSpace
codomain	Target space. codomain.dim independent half-StudentT components are produced per input. TYPE: ContinuousSpace
df	Degrees-of-freedom hyperparameter (positive). Held fixed across observations; only scale is learnable. TYPE: float
hidden_dim	Hidden layer width for the neural parameter source. TYPE: int DEFAULT: None

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

def __init__(
    self,
    domain: AnySpace,
    codomain: ContinuousSpace,
    df: float,
    hidden_dim: int | Sequence[int] | None = None,
    param_source: ParamSource | None = None,
    param_source_option: str | None = None,
) -> None:
    if df <= 0.0:
        raise ValueError(f"ConditionalHalfStudentT: df must be > 0, got {df!r}")
    super().__init__(domain, codomain)
    d = codomain.dim
    self._d = d
    self._df = float(df)
    self.param_source = _make_source(
        domain,
        d,
        hidden_dim,
        param_source=param_source,
        param_source_option=param_source_option,
    )

LKJCorrelationFactor

LKJCorrelationFactor(dim: int, eta: float, domain: AnySpace)

Bases: ContinuousMorphism

LKJ prior on Cholesky factors LKJ(K, η) over CholeskyFactor(K).

Density on the Cholesky factor:

.. math::

p(L) \propto \prod_{k=2}^{K} L_{kk}^{K - k + 2(\eta - 1)}.

A higher concentration :math:\eta > 1 pulls toward the identity correlation; :math:\eta = 1 is uniform on correlations. Sampling uses the onion method of Lewandowski-Kurowicka-Joe 2009: draw row-norm partial correlations from Beta distributions and form :math:L row-by-row.

PARAMETER	DESCRIPTION
dim	Correlation-matrix size :math:K \ge 2. TYPE: int
eta	Concentration :math:\eta > 0. TYPE: float
domain	The morphism's source (parameter conditioning); typically the program's input space. The LKJ prior itself does not consume per-observation conditioning, so the rsample path broadcasts the prior across the batch dimension. TYPE: AnySpace

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

def __init__(self, dim: int, eta: float, domain: AnySpace) -> None:
    if dim < 2:
        raise ValueError(f"LKJ requires dim >= 2; got {dim}")
    if eta <= 0:
        raise ValueError(f"LKJ requires eta > 0; got {eta}")
    codomain = CholeskyFactor(name=f"L({dim})", dim=dim)
    super().__init__(domain, codomain)
    self._dim = dim
    self._eta = float(eta)

rsample

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

Draw one Cholesky factor per input row, by the onion method.

PARAMETER	DESCRIPTION
x	The input; only its row count, device, and dtype are read. TYPE: Tensor
sample_shape	Leading sample dimensions. TYPE: Size DEFAULT: Size()

RETURNS	DESCRIPTION
Tensor	Flat factors, shape (*sample_shape, batch, dim * dim).

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

def rsample(
    self, x: torch.Tensor, sample_shape: torch.Size = torch.Size()
) -> torch.Tensor:
    """Draw one Cholesky factor per input row, by the onion method.

    Parameters
    ----------
    x : torch.Tensor
        The input; only its row count, device, and dtype are read.
    sample_shape : torch.Size
        Leading sample dimensions.

    Returns
    -------
    torch.Tensor
        Flat factors, shape ``(*sample_shape, batch, dim * dim)``.
    """
    batch = x.shape[0]
    factors = self._distribution(x).sample(torch.Size(sample_shape) + (batch,))
    return factors.reshape(*factors.shape[:-2], self._dim * self._dim)

log_prob

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

Normalized log density of the LKJ prior at Cholesky factors.

PARAMETER	DESCRIPTION
x	The input; it carries no parameters. TYPE: Tensor
y	Factors, flat (..., dim * dim) or square (..., dim, dim). TYPE: Tensor

RETURNS	DESCRIPTION
Tensor	One log density per factor.

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

def log_prob(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
    """Normalized log density of the LKJ prior at Cholesky factors.

    Parameters
    ----------
    x : torch.Tensor
        The input; it carries no parameters.
    y : torch.Tensor
        Factors, flat ``(..., dim * dim)`` or square
        ``(..., dim, dim)``.

    Returns
    -------
    torch.Tensor
        One log density per factor.
    """
    K = self._dim
    if y.shape[-2:] != (K, K):
        y = y.reshape(*y.shape[:-1], K, K)
    return self._distribution(x).log_prob(y)

Truncated

Truncated(base: ContinuousMorphism, lower: float | None = None, upper: float | None = None, max_rejection_iterations: int = 64)

Bases: ContinuousMorphism

Truncate a base family to an interval :math:[a, b].

Categorical denotation: given a base family :math:F : \Theta \to \mathcal{G}(\mathbb{R}) and constants :math:a, b \in \bar{\mathbb{R}} with :math:a < b, the truncated family has density

.. math::

p_{F_{|[a,b]}}(x) = \frac{p_F(x)}{F_{\text{cdf}}(b)
- F_{\text{cdf}}(a)} \cdot \mathbb{1}_{[a,b]}(x)

and the morphism :math:F_{|[a,b]} : \Theta \to \mathcal{G}([a,b]). Sampling uses inverse-CDF when base supports it; otherwise rejection sampling.

PARAMETER	DESCRIPTION
base	The base distribution-family morphism. Must expose log_prob and rsample plus an icdf method or a base_distribution torch Distribution for inverse-CDF sampling. Falls back to rejection sampling otherwise. TYPE: ContinuousMorphism
lower	Lower bound :math:a. None means :math:-\infty. TYPE: float or None DEFAULT: None
upper	Upper bound :math:b. None means :math:+\infty. TYPE: float or None DEFAULT: None
max_rejection_iterations	Cap on rejection-sampling attempts before raising. TYPE: int DEFAULT: 64

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

def __init__(
    self,
    base: ContinuousMorphism,
    lower: float | None = None,
    upper: float | None = None,
    max_rejection_iterations: int = 64,
) -> None:
    super().__init__(base.domain, base.codomain)
    if lower is None and upper is None:
        raise ValueError(
            "Truncated requires at least one of lower / upper to be finite; "
            "without truncation, use the base family directly"
        )
    if lower is not None and upper is not None and not (lower < upper):
        raise ValueError(
            f"Truncated requires lower < upper; got lower={lower}, upper={upper}"
        )
    self._base = base
    self._lower = lower
    self._upper = upper
    self._max_iters = max_rejection_iterations
    # Attach so the parent nn.Module tracks parameters.
    self._base_mod = base