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.

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
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
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
411
412
413
414
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
416
417
418
419
420
421
422
423
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
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
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
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
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
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
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
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
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
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
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
 994
 995
 996
 997
 998
 999
1000
1001
1002
1003
1004
1005
1006
1007
1008
1009
1010
1011
1012
1013
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
1073
1074
1075
1076
1077
1078
1079
1080
1081
1082
1083
1084
1085
1086
1087
1088
1089
1090
1091
1092
1093
1094
1095
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
1148
1149
1150
1151
1152
1153
1154
1155
1156
1157
1158
1159
1160
1161
1162
1163
1164
1165
1166
1167
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
1209
1210
1211
1212
1213
1214
1215
1216
1217
1218
1219
1220
1221
1222
1223
1224
1225
1226
1227
1228
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
1276
1277
1278
1279
1280
1281
1282
1283
1284
1285
1286
1287
1288
1289
1290
1291
1292
1293
1294
1295
1296
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
1388
1389
1390
1391
1392
1393
1394
1395
1396
1397
1398
1399
1400
1401
1402
1403
1404
1405
1406
1407
1408
1409
1410
1411
1412
1413
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
1485
1486
1487
1488
1489
1490
1491
1492
1493
1494
1495
1496
1497
1498
1499
1500
1501
1502
1503
1504
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
1598
1599
1600
1601
1602
1603
1604
1605
1606
1607
1608
1609
1610
1611
1612
1613
1614
1615
1616
1617
1618
1619
1620
1621
1622
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
1640
1641
1642
1643
1644
1645
1646
1647
1648
1649
1650
1651
1652
1653
1654
1655
1656
1657
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
1659
1660
1661
1662
1663
1664
1665
1666
1667
1668
1669
1670
1671
1672
1673
1674
1675
1676
1677
1678
1679
1680
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
1700
1701
1702
1703
1704
1705
1706
1707
1708
1709
1710
1711
1712
1713
1714
1715
1716
1717
1718
1719
1720
1721
1722
1723
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
1745
1746
1747
1748
1749
1750
1751
1752
1753
1754
1755
1756
1757
1758
1759
1760
1761
1762
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
1764
1765
1766
1767
1768
1769
1770
1771
1772
1773
1774
1775
1776
1777
1778
1779
1780
1781
1782
1783
1784
1785
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
1819
1820
1821
1822
1823
1824
1825
1826
1827
1828
1829
1830
1831
1832
1833
1834
1835
1836
1837
1838
1839
1840
1841
1842
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
1886
1887
1888
1889
1890
1891
1892
1893
1894
1895
1896
1897
1898
1899
1900
1901
1902
1903
1904
1905
1906
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

ConditionalOrderedLogistic

ConditionalOrderedLogistic(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)).

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. The codomain must be a finite set.

Source code in src/quivers/continuous/families.py
1942
1943
1944
1945
1946
1947
1948
1949
1950
1951
1952
1953
1954
1955
1956
1957
1958
1959
1960
1961
1962
1963
1964
1965
1966
1967
1968
1969
1970
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(
            "ConditionalOrderedLogistic requires a FinSet codomain, "
            f"got {codomain!r}"
        )
    if codomain.size < 2:
        raise ValueError(
            "ConditionalOrderedLogistic 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,
    )

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
2005
2006
2007
2008
2009
2010
2011
2012
2013
2014
2015
2016
2017
2018
2019
2020
2021
2022
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
2053
2054
2055
2056
2057
2058
2059
2060
2061
2062
2063
2064
2065
2066
2067
2068
2069
2070
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
2102
2103
2104
2105
2106
2107
2108
2109
2110
2111
2112
2113
2114
2115
2116
2117
2118
2119
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
2159
2160
2161
2162
2163
2164
2165
2166
2167
2168
2169
2170
2171
2172
2173
2174
2175
2176
2177
2178
2179
2180
2181
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
2234
2235
2236
2237
2238
2239
2240
2241
2242
2243
2244
2245
2246
2247
2248
2249
2250
2251
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
2316
2317
2318
2319
2320
2321
2322
2323
2324
2325
2326
2327
2328
2329
2330
2331
2332
2333
2334
2335
2336
2337
2338
2339
2340
2341
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
2406
2407
2408
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
2454
2455
2456
2457
2458
2459
2460
2461
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
2516
2517
2518
2519
2520
2521
2522
2523
2524
2525
2526
2527
2528
2529
2530
2531
2532
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
2631
2632
2633
2634
2635
2636
2637
2638
2639
2640
2641
2642
2643
2644
2645
2646
2647
2648
2649
2650
2651
2652
2653
2654
2655
2656
2657
2658
2659
2660
2661
2662
2663
2664
2665
2666
2667
2668
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
2841
2842
2843
2844
2845
2846
2847
2848
2849
2850
2851
2852
2853
2854
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.

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
2944
2945
2946
2947
2948
2949
2950
2951
2952
2953
2954
2955
2956
2957
2958
2959
2960
2961
2962
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
3343
3344
3345
3346
3347
3348
3349
3350
3351
3352
3353
3354
3355
3356
3357
3358
3359
3360
3361
3362
3363
3364
3365
3366
3367
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
3447
3448
3449
3450
3451
3452
3453
3454
3455
3456
3457
3458
3459
3460
3461
3462
3463
3464
3465
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
3525
3526
3527
3528
3529
3530
3531
3532
3533
3534
3535
3536
3537
3538
3539
3540
3541
3542
3543
3544
3545
3546
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
3629
3630
3631
3632
3633
3634
3635
3636
3637
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)

log_prob

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

Log-density of the LKJ prior at the Cholesky factor y.

Up to a normalizing constant that doesn't depend on :math:L, :math:\log p(L) = \sum_{k=2}^{K} (K-k+2(\eta-1)) \log L_{kk}. The diagonal entries are extracted from the flattened representation.

Source code in src/quivers/continuous/families.py
3679
3680
3681
3682
3683
3684
3685
3686
3687
3688
3689
3690
3691
3692
3693
3694
3695
3696
3697
3698
def log_prob(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
    """Log-density of the LKJ prior at the Cholesky factor ``y``.

    Up to a normalizing constant that doesn't depend on
    :math:`L`, :math:`\\log p(L) = \\sum_{k=2}^{K} (K-k+2(\\eta-1))
    \\log L_{kk}`. The diagonal entries are extracted from the
    flattened representation.
    """
    batch = y.shape[0]
    K = self._dim
    L = y.reshape(batch, K, K)
    diag = torch.diagonal(L, dim1=-2, dim2=-1)  # (batch, K)
    # Coefficients per diagonal entry (Stan's lkj_corr_cholesky_lpdf):
    # log_jac_term[k] = (K - k + 2*(eta - 1)) * log(L_kk)  for k = 2..K
    # Pre-K-indexed: power[0..K-1] where power[k] = (K-1-k) + 2*(eta-1).
    # The first diagonal is fixed at 1 so log(1)=0 contributes nothing.
    ks = torch.arange(K, device=y.device, dtype=y.dtype)
    powers = (K - 1 - ks) + 2.0 * (self._eta - 1.0)
    log_diag = torch.log(diag.clamp(min=1e-30))
    return (powers * log_diag).sum(dim=-1)

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
3751
3752
3753
3754
3755
3756
3757
3758
3759
3760
3761
3762
3763
3764
3765
3766
3767
3768
3769
3770
3771
3772
3773
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