Variational Guides

Variational guide distributions for approximate inference. The guides live in submodules of quivers.inference.guides and share the Guide ABC and LatentRegistry introspection layer. They are AutoNormalGuide, AutoDeltaGuide, AutoMultivariateNormalGuide, AutoLowRankMultivariateNormalGuide, AutoLaplaceApproximation, AutoNormalizingFlow, AutoIAFGuide, AutoNeuralSplineGuide, AutoMixtureGuide, AutoGuideList, and AutoStructured.

quivers.inference re-exports every guide, and quivers.inference.guides adds the Pyro-style short names (AutoNormal, AutoMultivariateNormal, AutoLowRankMVN, AutoDelta, AutoLaplace, AutoIAFNormal), each bound to the guide class it names, and the type aliases that configure AutoStructured.

guides

Variational guide families.

Public surface (re-exported by the parent quivers.inference package): one ABC (Guide) plus a zoo of concrete Auto*Guide subclasses spanning the standard variational-family ladder from mean-field Normal to normalizing-flow stacks and hierarchical / mixture / structured guides.

Every concrete guide is built against a single quivers.inference.registry.LatentRegistry and obeys the shape contract documented on Guide.

Short-name aliases matching the Pyro AutoGuide naming (AutoNormal, AutoMultivariateNormal, AutoLowRankMVN, AutoDelta, AutoLaplace, AutoIAFNormal) are re-exported for users switching over from Pyro. The longer *Guide names remain the canonical quivers names.

AutoNormal module-attribute

AutoNormal = AutoNormalGuide

Pyro-style short name for AutoNormalGuide.

AutoMultivariateNormal module-attribute

AutoMultivariateNormal = AutoMultivariateNormalGuide

Pyro-style short name for AutoMultivariateNormalGuide.

AutoLowRankMVN module-attribute

AutoLowRankMVN = AutoLowRankMultivariateNormalGuide

Pyro-style short name for AutoLowRankMultivariateNormalGuide.

AutoDelta module-attribute

AutoDelta = AutoDeltaGuide

Pyro-style short name for AutoDeltaGuide.

AutoLaplace module-attribute

AutoLaplace = AutoLaplaceApproximation

Pyro-style short name for AutoLaplaceApproximation.

AutoIAFNormal module-attribute

AutoIAFNormal = AutoIAFGuide

Pyro-style short name for AutoIAFGuide.

ConditionalSpec

ConditionalSpec = str | Mapping[str, str]

The per-site marginal families of an AutoStructured guide: one family name ("delta", "normal", or "mvn") for every site, or a mapping from site name to family name.

DependencyKind

DependencyKind = str | Callable[[Tensor], Tensor]

One dependency edge of an AutoStructured guide: the string "linear" for a learned linear shift, or a callable mapping the parent's unconstrained flat draw to a shift on the child's unconstrained flat location.

DependencySpec

DependencySpec = Mapping[str, Mapping[str, DependencyKind]]

The dependency edges of an AutoStructured guide, keyed first by child site and then by parent site.

AutoGuideList

AutoGuideList(parts: dict[str, Guide])

Bases: Guide

A block-composed variational guide.

PARAMETER	DESCRIPTION
parts	Named subguides, each covering a disjoint set of latent site names. Keys are labels used for parameter-namespacing and error messages; values are any Guide subclass. Every guide's Guide.latent_names must be pairwise disjoint from every other's. TYPE: dict[str, Guide]

Source code in src/quivers/inference/guides/auto_guide_list.py

def __init__(self, parts: dict[str, Guide]) -> None:
    super().__init__()
    if len(parts) == 0:
        raise ValueError("AutoGuideList: parts must be non-empty")

    seen: dict[str, str] = {}
    for label, part in parts.items():
        for name in part.latent_names:
            if name in seen:
                raise ValueError(
                    f"AutoGuideList: latent {name!r} appears in "
                    f"both part {seen[name]!r} and part {label!r}; "
                    f"parts must cover disjoint site sets"
                )
            seen[name] = label
    self._parts_dict: dict[str, Guide] = dict(parts)
    self.parts = nn.ModuleDict(parts)
    self._latent_names: list[str] = list(seen.keys())
    self._site_to_part: dict[str, str] = seen

part_labels property

part_labels: tuple[str, ...]

Labels of the sub-guides, in declaration order.

registry property

registry

The list-guide has no single fused registry; downstream code that needs one should walk part_labels and pull each sub-guide's Guide.registry directly.

part

part(label: str) -> Guide

Look up a sub-guide by label.

Source code in src/quivers/inference/guides/auto_guide_list.py

def part(self, label: str) -> Guide:
    """Look up a sub-guide by label."""
    if label not in self._parts_dict:
        raise KeyError(
            f"AutoGuideList: no part labelled {label!r}; "
            f"available: {tuple(self._parts_dict)!r}"
        )
    return self._parts_dict[label]

rsample

rsample(x: Tensor) -> dict[str, Tensor]

Concatenate per-part reparameterized draws.

Source code in src/quivers/inference/guides/auto_guide_list.py

def rsample(self, x: torch.Tensor) -> dict[str, torch.Tensor]:
    """Concatenate per-part reparameterized draws."""
    result: dict[str, torch.Tensor] = {}
    for part in self._parts_dict.values():
        part_sample = part.rsample(x)
        for name, tensor in part_sample.items():
            result[name] = tensor
    return result

log_prob

log_prob(x: Tensor, sites: dict[str, Tensor]) -> Tensor

Sum per-part log-densities.

Each sub-guide receives only the site keys it owns; a missing key raises KeyError in the sub-guide.

Source code in src/quivers/inference/guides/auto_guide_list.py

def log_prob(
    self,
    x: torch.Tensor,
    sites: dict[str, torch.Tensor],
) -> torch.Tensor:
    """Sum per-part log-densities.

    Each sub-guide receives only the site keys it owns; a
    missing key raises `KeyError` in the sub-guide.
    """
    total: torch.Tensor | None = None
    for label, part in self._parts_dict.items():
        part_sites = {
            name: sites[name] for name in part.latent_names if name in sites
        }
        missing = [n for n in part.latent_names if n not in sites]
        if missing:
            raise KeyError(
                f"AutoGuideList.log_prob: part {label!r} needs "
                f"sites {missing!r} but they are not in the input"
            )
        contribution = part.log_prob(x, part_sites)
        total = contribution if total is None else total + contribution
    assert total is not None
    return total

AutoStructured

AutoStructured(model: MonadicProgram, observed_names: set[str], conditionals: ConditionalSpec, dependencies: DependencySpec | None = None, init_scale: float = 0.1)

Bases: Guide

Per-site / per-edge structured variational guide.

PARAMETER	DESCRIPTION
model	Generative model to build a guide for. TYPE: MonadicProgram
observed_names	Variable names treated as observations. TYPE: set[str]
conditionals	Per-site marginal-family choice. A single string sets every site to the same family. A dict maps site name to family. Allowed families: "delta": Dirac mass at a learnable location (analogous to AutoDeltaGuide). "normal": independent-Normal marginal. "mvn": the site is part of a joint multivariate Normal block. Every "mvn" site shares one Cholesky factor sized to the sum of every "mvn" site's flat length. This matches Pyro's convention. TYPE: str or dict[str, str]
dependencies	Directed edges. dependencies[child][parent] is either the string "linear" (build a _LinearDependency from the parent's unconstrained flat draw to a shift on the child's unconstrained flat location) or a user-supplied callable with the same signature. Every dependency must reference a site strictly upstream in declaration order. TYPE: dict[str, dict[str, DependencyKind]] or None DEFAULT: None
init_scale	Initial scale of every "normal" and "mvn" marginal. Default 0.1. TYPE: float DEFAULT: 0.1

Source code in src/quivers/inference/guides/auto_structured.py

def __init__(
    self,
    model: MonadicProgram,
    observed_names: set[str],
    conditionals: ConditionalSpec,
    dependencies: DependencySpec | None = None,
    init_scale: float = 0.1,
) -> None:
    super().__init__()
    self._registry = self.build_registry(model, observed_names)

    # --- Normalise conditionals ---
    site_names = list(self._registry.names)
    if isinstance(conditionals, str):
        if conditionals not in _ALLOWED_CONDITIONALS:
            raise ValueError(
                f"AutoStructured: conditional {conditionals!r} "
                f"not in {sorted(_ALLOWED_CONDITIONALS)!r}"
            )
        resolved: dict[str, str] = {n: conditionals for n in site_names}
    else:
        resolved = {}
        for n in site_names:
            if n not in conditionals:
                raise ValueError(
                    f"AutoStructured: conditionals dict has no "
                    f"entry for latent site {n!r}"
                )
            choice = conditionals[n]
            if choice not in _ALLOWED_CONDITIONALS:
                raise ValueError(
                    f"AutoStructured: conditional {choice!r} for "
                    f"site {n!r} not in "
                    f"{sorted(_ALLOWED_CONDITIONALS)!r}"
                )
            resolved[n] = choice
    self._conditionals: dict[str, str] = resolved

    # --- Register per-site location + scale parameters ---
    init_log_scale = float(torch.tensor(init_scale).log().item())
    for name in site_names:
        site = self._registry.sites[name]
        self.register_parameter(
            f"loc_{name}",
            nn.Parameter(torch.zeros(site.unconstrained_shape)),
        )
        if resolved[name] == "normal":
            self.register_parameter(
                f"log_scale_{name}",
                nn.Parameter(torch.full(site.unconstrained_shape, init_log_scale)),
            )
        # "delta" has no scale.  "mvn" scale is shared and set below.

    # --- Set up the shared MVN block, if any ---
    mvn_names: list[str] = [n for n in site_names if resolved[n] == "mvn"]
    self._mvn_names: list[str] = mvn_names
    self._mvn_offsets: dict[str, int] = {}
    cursor = 0
    for n in mvn_names:
        self._mvn_offsets[n] = cursor
        cursor += _site_flat_length(self._registry.sites[n])
    self._mvn_total: int = cursor
    if mvn_names:
        init_diag = torch.full((self._mvn_total,), float(init_scale))
        init_diag_raw = torch.log(torch.expm1(init_diag.clamp(min=1e-6)))
        self.mvn_scale_diag_raw = nn.Parameter(init_diag_raw)
        self.mvn_scale_offdiag = nn.Parameter(
            torch.zeros(self._mvn_total, self._mvn_total)
        )

    # --- Validate + register dependencies ---
    deps = dependencies or {}
    name_order: dict[str, int] = {n: i for i, n in enumerate(site_names)}
    dep_modules: dict[str, nn.ModuleDict] = {}
    for child, parents in deps.items():
        if child not in name_order:
            raise ValueError(
                f"AutoStructured: dependency child {child!r} is not a latent site"
            )
        child_flat = _site_flat_length(self._registry.sites[child])
        per_parent = nn.ModuleDict()
        for parent, kind in parents.items():
            if parent not in name_order:
                raise ValueError(
                    f"AutoStructured: dependency parent {parent!r} "
                    f"(for child {child!r}) is not a latent site"
                )
            if name_order[parent] >= name_order[child]:
                raise ValueError(
                    f"AutoStructured: dependency {parent!r} -> "
                    f"{child!r} is not strictly upstream in "
                    f"declaration order"
                )
            parent_flat = _site_flat_length(self._registry.sites[parent])
            if kind == "linear":
                per_parent[parent] = _LinearDependency(
                    parent_flat=parent_flat, child_flat=child_flat
                )
            elif callable(kind):
                per_parent[parent] = _CallableDependency(kind)
            else:
                raise ValueError(
                    f"AutoStructured: dependency kind {kind!r} "
                    f"for edge {parent!r} -> {child!r} must be "
                    f"either the string 'linear' or a callable"
                )
        dep_modules[child] = per_parent
    self.deps = nn.ModuleDict(dep_modules)
    self._dep_spec: dict[str, dict[str, DependencyKind]] = {
        child: dict(parents) for child, parents in deps.items()
    }

conditionals property

conditionals: dict[str, str]

Per-site marginal family, as resolved at construction.

dependencies property

dependencies: dict[str, dict[str, DependencyKind]]

The declared per-edge dependence structure.

rsample

rsample(x: Tensor) -> dict[str, Tensor]

Ancestral reparameterized draw.

Source code in src/quivers/inference/guides/auto_structured.py

def rsample(self, x: torch.Tensor) -> dict[str, torch.Tensor]:
    """Ancestral reparameterized draw."""
    batch = x.shape[0]
    # First sample the shared MVN block (if any) in flat form.
    if self._mvn_names:
        mvn_flat: dict[str, torch.Tensor] = self._sample_mvn_block()
    else:
        mvn_flat = {}
    # Then walk sites in declaration order, sampling each per its
    # marginal family, offsetting the location by every incoming
    # dependency's contribution.
    result: dict[str, torch.Tensor] = {}
    z_flat_by_site: dict[str, torch.Tensor] = {}
    for site in self._registry.sites.values():
        name = site.name
        kind = self._conditionals[name]
        loc = self._loc(name)
        # Dependency shift, if any.
        shift_flat = torch.zeros(_site_flat_length(site), device=loc.device)
        if name in self.deps:
            for parent, dep_mod in self.deps[name].items():
                parent_flat = z_flat_by_site[parent]
                shift_flat = shift_flat + dep_mod(parent_flat)
        shift = shift_flat.reshape(site.unconstrained_shape)
        shifted_loc = loc + shift
        if kind == "delta":
            z_flat = shifted_loc.reshape(-1)
        elif kind == "normal":
            scale = self._normal_scale(name)
            z_shaped = D.Normal(shifted_loc, scale).rsample()
            z_flat = z_shaped.reshape(-1)
        elif kind == "mvn":
            z_flat = mvn_flat[name] + shift_flat
        else:
            raise RuntimeError(
                f"AutoStructured.rsample: unhandled conditional kind {kind!r}"
            )
        z_flat_by_site[name] = z_flat
        v = _push_through_bijector(site, z_flat, batch)
        result[name] = v
    return result

log_prob

log_prob(x: Tensor, sites: dict[str, Tensor]) -> Tensor

Pushforward log-density at the supplied constrained sites.

Source code in src/quivers/inference/guides/auto_structured.py

def log_prob(
    self,
    x: torch.Tensor,
    sites: dict[str, torch.Tensor],
) -> torch.Tensor:
    """Pushforward log-density at the supplied constrained sites."""
    batch = x.shape[0]
    z_flat_by_site: dict[str, torch.Tensor] = {}
    total_scalar = torch.zeros((), device=x.device)
    # First invert every bijector and accumulate the Jacobian terms.
    for site in self._registry.sites.values():
        name = site.name
        if name not in sites:
            raise KeyError(f"AutoStructured.log_prob: missing site {name!r}")
        z_shaped, log_det = self._v_to_unconstrained(site, sites[name])
        z_flat_by_site[name] = z_shaped.reshape(-1)
        total_scalar = total_scalar + log_det.reshape(-1).sum()
    # Next add each site's marginal log-density.
    for site in self._registry.sites.values():
        name = site.name
        kind = self._conditionals[name]
        if kind == "delta":
            continue
        loc = self._loc(name)
        shift_flat = torch.zeros(_site_flat_length(site), device=x.device)
        if name in self.deps:
            for parent, dep_mod in self.deps[name].items():
                shift_flat = shift_flat + dep_mod(z_flat_by_site[parent])
        shift = shift_flat.reshape(site.unconstrained_shape)
        shifted_loc = loc + shift
        if kind == "normal":
            scale = self._normal_scale(name)
            z_shaped = z_flat_by_site[name].reshape(site.unconstrained_shape)
            log_q = D.Normal(shifted_loc, scale).log_prob(z_shaped)
            total_scalar = total_scalar + log_q.reshape(-1).sum()
        # "mvn" sites are scored jointly below.
    # Finally, the joint MVN block, if any.
    if self._mvn_names:
        joint_loc_parts: list[torch.Tensor] = []
        joint_z_parts: list[torch.Tensor] = []
        for n in self._mvn_names:
            site = self._registry.sites[n]
            loc = self._loc(n)
            shift_flat = torch.zeros(_site_flat_length(site), device=x.device)
            if n in self.deps:
                for parent, dep_mod in self.deps[n].items():
                    shift_flat = shift_flat + dep_mod(z_flat_by_site[parent])
            shifted_loc_flat = loc.reshape(-1) + shift_flat
            joint_loc_parts.append(shifted_loc_flat)
            joint_z_parts.append(z_flat_by_site[n])
        joint_loc = torch.cat(joint_loc_parts, dim=0)
        joint_z = torch.cat(joint_z_parts, dim=0)
        dist = D.MultivariateNormal(joint_loc, scale_tril=self._mvn_scale_tril())
        total_scalar = total_scalar + dist.log_prob(joint_z)
    return total_scalar.expand(batch)

Guide

Bases: Module, ABC

Abstract variational guide.

Subclasses MUST implement rsample and log_prob and expose latent_names. They MAY override registry if they construct their registry lazily, but the default implementation expects self._registry to be set in __init__.

registry property

registry: LatentRegistry

The LatentRegistry this guide was built against.

latent_names abstractmethod property

latent_names: list[str]

Names of latent variables this guide covers.

build_registry classmethod

build_registry(model: MonadicProgram, observed_names: set[str] | frozenset[str]) -> LatentRegistry

Convenience wrapper around LatentRegistry.from_model so guide constructors can do self._registry = self.build_registry(model, obs) without an extra import.

Source code in src/quivers/inference/guides/base.py

@classmethod
def build_registry(
    cls,
    model: MonadicProgram,
    observed_names: set[str] | frozenset[str],
) -> LatentRegistry:
    """Convenience wrapper around
    `LatentRegistry.from_model` so guide constructors
    can do ``self._registry = self.build_registry(model, obs)``
    without an extra import."""
    return LatentRegistry.from_model(model, observed_names)

rsample abstractmethod

rsample(x: Tensor) -> dict[str, Tensor]

Reparameterized sample from :math:q_\phi(z \mid x).

PARAMETER	DESCRIPTION
x	Program input. Shape (batch, ...). Used only for its batch dim and device; the variational parameters are stored on the guide itself. TYPE: Tensor

RETURNS	DESCRIPTION
dict[str, Tensor]	Per-site constrained samples shaped to match the model's trace-side convention.

Source code in src/quivers/inference/guides/base.py

@abstractmethod
def rsample(self, x: torch.Tensor) -> dict[str, torch.Tensor]:
    """Reparameterized sample from :math:`q_\\phi(z \\mid x)`.

    Parameters
    ----------
    x : torch.Tensor
        Program input. Shape ``(batch, ...)``. Used only for
        its batch dim and device; the variational parameters
        are stored on the guide itself.

    Returns
    -------
    dict[str, torch.Tensor]
        Per-site constrained samples shaped to match the
        model's trace-side convention.
    """
    ...

log_prob abstractmethod

log_prob(x: Tensor, sites: dict[str, Tensor]) -> Tensor

Log density of sites under :math:q_\phi(z \mid x), with the change-of-variables Jacobian correction baked in.

RETURNS	DESCRIPTION
Tensor	Shape (batch,).

Source code in src/quivers/inference/guides/base.py

@abstractmethod
def log_prob(
    self,
    x: torch.Tensor,
    sites: dict[str, torch.Tensor],
) -> torch.Tensor:
    """Log density of ``sites`` under :math:`q_\\phi(z \\mid x)`,
    with the change-of-variables Jacobian correction baked in.

    Returns
    -------
    torch.Tensor
        Shape ``(batch,)``.
    """
    ...

AutoDeltaGuide

AutoDeltaGuide(model: MonadicProgram, observed_names: set[str], init_value: float = 0.0)

Bases: Guide

Dirac-delta MAP guide with per-site constrained bijector.

PARAMETER	DESCRIPTION
model	Generative model. TYPE: MonadicProgram
observed_names	Variable names treated as observations. TYPE: set[str]
init_value	Initial unconstrained-space coordinate for every latent. Default 0.0; for the standard bijectors this maps to a sensible interior point of each support (the median of a HalfNormal, the centre of the unit interval, the uniform Dirichlet, etc.). Small Gaussian noise is added so two coordinate values don't collide. TYPE: float DEFAULT: 0.0

Source code in src/quivers/inference/guides/delta.py

def __init__(
    self,
    model: MonadicProgram,
    observed_names: set[str],
    init_value: float = 0.0,
) -> None:
    super().__init__()
    self._registry = self.build_registry(model, observed_names)

    for site in self._registry.sites.values():
        self.register_parameter(
            f"unconstrained_{site.name}",
            nn.Parameter(
                torch.full(site.unconstrained_shape, init_value)
                + torch.randn(site.unconstrained_shape) * 0.01
            ),
        )

rsample

rsample(x: Tensor) -> dict[str, Tensor]

Return the learned point estimates in the prior's support.

Source code in src/quivers/inference/guides/delta.py

def rsample(self, x: torch.Tensor) -> dict[str, torch.Tensor]:
    """Return the learned point estimates in the prior's
    support."""
    batch = x.shape[0]
    return {
        site.name: self._push_through_bijector(site, batch)
        for site in self._registry.sites.values()
    }

log_prob

log_prob(x: Tensor, sites: dict[str, Tensor]) -> Tensor

Delta log-density: zero everywhere (the delta term and its Jacobian cancel in the ELBO under the standard score-function trick).

Source code in src/quivers/inference/guides/delta.py

def log_prob(
    self,
    x: torch.Tensor,
    sites: dict[str, torch.Tensor],
) -> torch.Tensor:
    """Delta log-density: zero everywhere (the delta term and
    its Jacobian cancel in the ELBO under the standard
    score-function trick)."""
    return torch.zeros(x.shape[0], device=x.device)

AutoIAFGuide

AutoIAFGuide(model: MonadicProgram, observed_names: set[str], num_flows: int = 4, hidden_dim: int | None = None, num_hidden_layers: int = 2)

Bases: AutoNormalizingFlow

Inverse-autoregressive-flow guide.

Default normalizing-flow guide for variational inference (Kingma-Salimans-Jozefowicz et al. 2016). Stack of InverseAutoregressiveTransform layers, each separated by a reverse permutation so successive layers have different autoregressive orderings.

Sampling is parallel (one MLP forward per layer); density evaluation is sequential (one coordinate at a time per layer), so this guide should be used with objectives that sample more than they score the same flow (ELBO, IWAE).

PARAMETER	DESCRIPTION
model	Generative model. TYPE: MonadicProgram
observed_names	Variable names treated as observations. TYPE: set[str]
num_flows	Number of IAF blocks in the stack. Default 4. TYPE: int DEFAULT: 4
hidden_dim	Hidden width of every MADE inside the stack. Default 2 * D where D is the latent dimension. TYPE: int DEFAULT: None
num_hidden_layers	Number of hidden layers in each MADE. Default 2. TYPE: int DEFAULT: 2

Source code in src/quivers/inference/guides/flow.py

def __init__(
    self,
    model: MonadicProgram,
    observed_names: set[str],
    num_flows: int = 4,
    hidden_dim: int | None = None,
    num_hidden_layers: int = 2,
) -> None:
    if num_flows < 1:
        raise ValueError(f"AutoIAFGuide: num_flows must be >= 1, got {num_flows}")
    registry = Guide.build_registry(model, observed_names)
    D_total = registry.total_unconstrained_dim
    if hidden_dim is None:
        hidden_dim = max(8, 2 * D_total)
    if D_total < 2:
        raise ValueError(
            f"AutoIAFGuide: model must have >= 2 unconstrained "
            f"latent dimensions for an IAF (got {D_total}); use "
            f"AutoNormalGuide for 1-D models"
        )
    layers: list[TransformModule] = []
    for i in range(num_flows):
        made = MADE(
            dim=D_total,
            n_per_dim=2,
            hidden=hidden_dim,
            n_hidden_layers=num_hidden_layers,
        )
        layers.append(InverseAutoregressiveTransform(made))
        if i < num_flows - 1:
            layers.append(_ReversePermutation(D_total))
    super().__init__(model, observed_names, layers)

AutoNeuralSplineGuide

AutoNeuralSplineGuide(model: MonadicProgram, observed_names: set[str], num_flows: int = 4, num_bins: int = 8, tail_bound: float = 3.0, hidden_dim: int | None = None, num_hidden_layers: int = 2)

Bases: AutoNormalizingFlow

Neural-spline-flow guide (Durkan-Bekasov-Murray-Papamakarios 2019).

Stack of monotone rational-quadratic spline coupling layers (NeuralSplineCouplingTransform) with alternating half-masks. Sharper than IAF for posteriors with bounded support or sharp modes; comparable runtime.

PARAMETER	DESCRIPTION
model	Generative model. TYPE: MonadicProgram
observed_names	Variable names treated as observations. TYPE: set[str]
num_flows	Number of coupling layers. Default 4. TYPE: int DEFAULT: 4
num_bins	Number of spline bins per coordinate. Default 8. TYPE: int DEFAULT: 8
tail_bound	Inputs outside [-tail_bound, tail_bound] pass through as identity. Default 3.0. TYPE: float DEFAULT: 3.0
hidden_dim	Hidden width of the coupling MLPs. Default max(64, 2*D). TYPE: int DEFAULT: None
num_hidden_layers	Hidden layers in each coupling MLP. Default 2. TYPE: int DEFAULT: 2

Source code in src/quivers/inference/guides/flow.py

def __init__(
    self,
    model: MonadicProgram,
    observed_names: set[str],
    num_flows: int = 4,
    num_bins: int = 8,
    tail_bound: float = 3.0,
    hidden_dim: int | None = None,
    num_hidden_layers: int = 2,
) -> None:
    if num_flows < 1:
        raise ValueError(
            f"AutoNeuralSplineGuide: num_flows must be >= 1, got {num_flows}"
        )
    registry = Guide.build_registry(model, observed_names)
    D_total = registry.total_unconstrained_dim
    if D_total < 2:
        raise ValueError(
            f"AutoNeuralSplineGuide: model must have >= 2 "
            f"unconstrained latent dimensions for a spline "
            f"coupling flow (got {D_total})"
        )
    if hidden_dim is None:
        hidden_dim = max(64, 2 * D_total)
    layers: list[TransformModule] = []
    for i in range(num_flows):
        mask = alternating_mask(D_total, even=(i % 2 == 0))
        num_unmasked = int(mask.sum().item())
        num_masked = D_total - num_unmasked
        # Net produces (3 * num_bins - 1) parameters per masked coord.
        out_dim = num_masked * (3 * num_bins - 1)
        net = make_coupling_mlp(
            n_in=num_unmasked,
            n_out=out_dim,
            hidden=hidden_dim,
            n_hidden_layers=num_hidden_layers,
        )
        layers.append(
            NeuralSplineCouplingTransform(
                dim=D_total,
                net=net,
                mask=mask,
                num_bins=num_bins,
                tail_bound=tail_bound,
            )
        )
    super().__init__(model, observed_names, layers)

AutoNormalizingFlow

AutoNormalizingFlow(model: MonadicProgram, observed_names: set[str], transforms: list[TransformModule])

Bases: Guide

Normalising-flow variational guide over the flat latent vector.

PARAMETER	DESCRIPTION
model	Generative model to build a guide for. TYPE: MonadicProgram
observed_names	Variable names treated as observations. TYPE: set[str]
transforms	Flow stack applied to the standard-Normal base. Each TransformModule must accept a (..., D)-shaped tensor where D is the registry's total unconstrained dimension. TYPE: list[TransformModule]

Source code in src/quivers/inference/guides/flow.py

def __init__(
    self,
    model: MonadicProgram,
    observed_names: set[str],
    transforms: list[TransformModule],
) -> None:
    super().__init__()
    self._registry = self.build_registry(model, observed_names)
    D_total = self._registry.total_unconstrained_dim
    if D_total == 0:
        raise ValueError(
            f"{type(self).__name__}: registry has zero total "
            f"unconstrained dimension; model has no continuous "
            f"latents to guide"
        )
    if not transforms:
        raise ValueError(
            f"{type(self).__name__}: transforms list must be non-empty"
        )
    self._D = D_total
    self.flow = nn.ModuleList(transforms)
    self.register_buffer("base_loc", torch.zeros(D_total))
    self.register_buffer("base_scale", torch.ones(D_total))

rsample

rsample(x: Tensor) -> dict[str, Tensor]

One flow draw, unflattened and bijected to constrained space.

Source code in src/quivers/inference/guides/flow.py

def rsample(self, x: torch.Tensor) -> dict[str, torch.Tensor]:
    """One flow draw, unflattened and bijected to constrained space."""
    batch = x.shape[0]
    z0 = self._base_dist().rsample()
    z_K, _ = self._forward(z0)
    per_site = self._registry.unflatten_unconstrained(z_K)
    result: dict[str, torch.Tensor] = {}
    for site in self._registry.sites.values():
        z_site = per_site[site.name]
        if not site.is_plate:
            z_site = z_site.unsqueeze(0).expand(batch, *site.unconstrained_shape)
        v = site.bijector(z_site)
        if site.constrained_dim == 1 and v.dim() >= 1 and v.shape[-1] == 1:
            v = v.squeeze(-1)
        result[site.name] = v
    return result

log_prob

log_prob(x: Tensor, sites: dict[str, Tensor]) -> Tensor

Log-density at the supplied constrained sites.

Source code in src/quivers/inference/guides/flow.py

def log_prob(
    self,
    x: torch.Tensor,
    sites: dict[str, torch.Tensor],
) -> torch.Tensor:
    """Log-density at the supplied constrained sites."""
    batch = x.shape[0]
    unconstrained_per_site: dict[str, torch.Tensor] = {}
    bijector_log_det = torch.zeros((), device=x.device)
    for site in self._registry.sites.values():
        if site.name not in sites:
            raise KeyError(
                f"{type(self).__name__}.log_prob: missing site {site.name!r}"
            )
        v = sites[site.name]
        if site.constrained_dim == 1 and v.dim() == (1 if site.is_plate else 1):
            v_e = v.unsqueeze(-1)
        else:
            v_e = v
        if not site.is_plate and v_e.dim() == len(site.unconstrained_shape) + 1:
            v_e = v_e[0]
        z_site = site.bijector.inv(v_e)
        unconstrained_per_site[site.name] = z_site
        bijector_log_det = bijector_log_det + (
            site.bijector.inv.log_abs_det_jacobian(v_e, z_site).sum()
        )

    z_K = self._registry.flatten_unconstrained(unconstrained_per_site)
    z_0, flow_log_det = self._inverse(z_K)
    log_p_base = self._base_dist().log_prob(z_0)
    # log q(z_K) = log p_0(z_0) - log|det dT/dz_0|
    #            = log p_0(z_0) - flow_log_det
    log_q_z = log_p_base - flow_log_det
    return (log_q_z + bijector_log_det).expand(batch)

AutoLaplaceApproximation

AutoLaplaceApproximation(model: MonadicProgram, observed_names: set[str], init_value: float = 0.0)

Bases: Guide

Laplace-approximation guide.

PARAMETER	DESCRIPTION
model	Generative model. TYPE: MonadicProgram
observed_names	Variable names treated as observations. TYPE: set[str]
init_value	Initial unconstrained-space MAP estimate. Default 0.0. TYPE: float DEFAULT: 0.0

Source code in src/quivers/inference/guides/laplace.py

def __init__(
    self,
    model: MonadicProgram,
    observed_names: set[str],
    init_value: float = 0.0,
) -> None:
    super().__init__()
    self._registry = self.build_registry(model, observed_names)
    D_total = self._registry.total_unconstrained_dim
    if D_total == 0:
        raise ValueError(
            f"{type(self).__name__}: registry has zero total "
            f"unconstrained dimension; model has no continuous "
            f"latents to guide"
        )
    self._D = D_total
    # MAP estimate in flat unconstrained space.
    init = torch.full((D_total,), float(init_value))
    init = init + 0.01 * torch.randn(D_total)
    self.map_z = nn.Parameter(init)
    # Hessian-phase parameters; initialized to identity scale_tril
    # but only used after fit_hessian() is called.
    self.register_buffer("_hessian_fitted", torch.zeros((), dtype=torch.bool))
    self.register_buffer(
        "_scale_tril",
        torch.eye(D_total) * 1e-3,
        persistent=True,
    )

hessian_fitted property

hessian_fitted: bool

Whether fit_hessian has been called.

fit_hessian

fit_hessian(model: MonadicProgram, x: Tensor, observations: dict[str, Tensor], *, jitter: float = 0.0001) -> None

Compute and cache the Hessian-derived Cholesky factor.

Solves the eigenproblem of the negative-log-joint Hessian at the current MAP, projects negative eigenvalues to jitter (so the resulting Gaussian is always positive-definite), and stores the matching lower-triangular Cholesky factor of the inverse Hessian as the posterior scale_tril.

Call this after MAP optimisation has converged. Subsequent rsample / log_prob calls sample from :math:\mathcal{N}(z^\star, H^{-1}).

Source code in src/quivers/inference/guides/laplace.py

def fit_hessian(
    self,
    model: MonadicProgram,
    x: torch.Tensor,
    observations: dict[str, torch.Tensor],
    *,
    jitter: float = 1e-4,
) -> None:
    """Compute and cache the Hessian-derived Cholesky factor.

    Solves the eigenproblem of the negative-log-joint Hessian at
    the current MAP, projects negative eigenvalues to ``jitter``
    (so the resulting Gaussian is always positive-definite), and
    stores the matching lower-triangular Cholesky factor of the
    inverse Hessian as the posterior scale_tril.

    Call this after MAP optimisation has converged. Subsequent
    `rsample` / `log_prob` calls sample from
    :math:`\\mathcal{N}(z^\\star, H^{-1})`.
    """

    def neg_log_joint(z_flat: torch.Tensor) -> torch.Tensor:
        per_site_unconstrained = self._registry.unflatten_unconstrained(z_flat)
        constrained: dict[str, torch.Tensor] = {}
        log_det_sum = torch.zeros((), device=z_flat.device)
        for site in self._registry.sites.values():
            z_site = per_site_unconstrained[site.name]
            if not site.is_plate:
                z_site = z_site.unsqueeze(0)
            v = site.bijector(z_site)
            log_det_sum = log_det_sum + (
                site.bijector.log_abs_det_jacobian(z_site, v).sum()
            )
            if site.constrained_dim == 1 and v.dim() >= 1 and v.shape[-1] == 1:
                v = v.squeeze(-1)
            constrained[site.name] = v
        log_p = model.log_joint(x, {**constrained, **observations})
        return -(log_p.sum() + log_det_sum)

    H = torch.autograd.functional.hessian(neg_log_joint, self.map_z.detach())
    H = 0.5 * (H + H.t())
    eigvals, eigvecs = torch.linalg.eigh(H)
    eigvals_clamped = eigvals.clamp(min=jitter)
    # Σ = (V Λ V^T)^{-1} = V Λ^{-1} V^T; scale_tril is its Cholesky.
    inv_eigvals = 1.0 / eigvals_clamped
    sigma = (eigvecs * inv_eigvals.unsqueeze(0)) @ eigvecs.t()
    sigma = 0.5 * (sigma + sigma.t())
    sigma = sigma + jitter * torch.eye(self._D, device=sigma.device)
    L = torch.linalg.cholesky(sigma)
    self._scale_tril.copy_(L)
    self._hessian_fitted.fill_(True)

rsample

rsample(x: Tensor) -> dict[str, Tensor]

Sample from the Laplace posterior, unflatten, and biject.

Source code in src/quivers/inference/guides/laplace.py

def rsample(self, x: torch.Tensor) -> dict[str, torch.Tensor]:
    """Sample from the Laplace posterior, unflatten, and biject."""
    batch = x.shape[0]
    z_flat = self._sample_unconstrained()
    per_site = self._registry.unflatten_unconstrained(z_flat)

    result: dict[str, torch.Tensor] = {}
    for site in self._registry.sites.values():
        z_site = per_site[site.name]
        if not site.is_plate:
            z_site = z_site.unsqueeze(0).expand(batch, *site.unconstrained_shape)
        v = site.bijector(z_site)
        if site.constrained_dim == 1 and v.dim() >= 1 and v.shape[-1] == 1:
            v = v.squeeze(-1)
        result[site.name] = v
    return result

log_prob

log_prob(x: Tensor, sites: dict[str, Tensor]) -> Tensor

Log-density at the supplied constrained sites.

Returns zero before fit_hessian (MAP-phase delta convention); after fit_hessian returns the Gaussian log-density plus the per-site bijector Jacobian correction.

Source code in src/quivers/inference/guides/laplace.py

def log_prob(
    self,
    x: torch.Tensor,
    sites: dict[str, torch.Tensor],
) -> torch.Tensor:
    """Log-density at the supplied constrained sites.

    Returns zero before `fit_hessian` (MAP-phase delta
    convention); after `fit_hessian` returns the Gaussian
    log-density plus the per-site bijector Jacobian correction.
    """
    batch = x.shape[0]
    if not bool(self._hessian_fitted):
        return torch.zeros(batch, device=x.device)

    unconstrained_per_site: dict[str, torch.Tensor] = {}
    bijector_log_det = torch.zeros((), device=x.device)
    for site in self._registry.sites.values():
        if site.name not in sites:
            raise KeyError(
                f"{type(self).__name__}.log_prob: missing site {site.name!r}"
            )
        v = sites[site.name]
        if site.constrained_dim == 1 and v.dim() == (1 if site.is_plate else 1):
            v_e = v.unsqueeze(-1)
        else:
            v_e = v
        if not site.is_plate and v_e.dim() == len(site.unconstrained_shape) + 1:
            v_e = v_e[0]
        z_site = site.bijector.inv(v_e)
        unconstrained_per_site[site.name] = z_site
        bijector_log_det = bijector_log_det + (
            site.bijector.inv.log_abs_det_jacobian(v_e, z_site).sum()
        )
    z_flat = self._registry.flatten_unconstrained(unconstrained_per_site)
    gauss = D.MultivariateNormal(self.map_z, scale_tril=self._scale_tril)
    log_q_z = gauss.log_prob(z_flat)
    return (log_q_z + bijector_log_det).expand(batch)

AutoMixtureGuide

AutoMixtureGuide(components: list[Guide], init_temperature: float = 1.0)

Bases: Guide

Finite mixture variational guide.

PARAMETER	DESCRIPTION
components	Component guides. All components must share the same LatentRegistry (i.e. be built against the same model + observed-name set). TYPE: list[Guide]
init_temperature	Initial Gumbel-Softmax temperature. Default 1.0; anneal toward zero for sharper component selection. TYPE: float DEFAULT: 1.0

Source code in src/quivers/inference/guides/mixture.py

def __init__(
    self,
    components: list[Guide],
    init_temperature: float = 1.0,
) -> None:
    super().__init__()
    if len(components) < 2:
        raise ValueError(
            f"AutoMixtureGuide: need at least 2 components, got {len(components)}"
        )
    reference = components[0]
    ref_names = tuple(reference.registry.names)
    for i, comp in enumerate(components[1:], 1):
        comp_names = tuple(comp.registry.names)
        if comp_names != ref_names:
            raise ValueError(
                f"AutoMixtureGuide: component {i} has different "
                f"latent names {comp_names!r} than component 0 "
                f"{ref_names!r}"
            )
    if init_temperature <= 0.0:
        raise ValueError(
            f"AutoMixtureGuide: init_temperature must be positive, "
            f"got {init_temperature}"
        )
    self._registry = reference.registry
    self.components = nn.ModuleList(components)
    self.mixture_logits = nn.Parameter(torch.zeros(len(components)))
    self._temperature: torch.Tensor
    self.register_buffer("_temperature", torch.tensor(float(init_temperature)))

temperature property

temperature: float

Current Gumbel-Softmax temperature.

set_temperature

set_temperature(value: float) -> None

Anneal the Gumbel-Softmax temperature.

Source code in src/quivers/inference/guides/mixture.py

def set_temperature(self, value: float) -> None:
    """Anneal the Gumbel-Softmax temperature."""
    if value <= 0.0:
        raise ValueError(
            f"AutoMixtureGuide.set_temperature: must be positive, got {value}"
        )
    self._temperature.fill_(float(value))

rsample

rsample(x: Tensor) -> dict[str, Tensor]

Reparameterized mixture draw via Gumbel-Softmax.

Each call samples a Gumbel-Softmax weight vector :math:w \in \Delta^{K-1} and returns :math:\sum_k w_k \cdot v^{(k)} per site, where :math:v^{(k)} is component k's constrained-space sample. Because the constrained-space sites' supports are not in general convex (e.g. a Cholesky factor on torch.distributions.constraints.corr_cholesky), the soft mixture can drift outside any single component's support during training; the categorical-pick fallback in hard_rsample returns a single component's sample for use at inference time.

Source code in src/quivers/inference/guides/mixture.py

def rsample(self, x: torch.Tensor) -> dict[str, torch.Tensor]:
    """Reparameterized mixture draw via Gumbel-Softmax.

    Each call samples a Gumbel-Softmax weight vector
    :math:`w \\in \\Delta^{K-1}` and returns
    :math:`\\sum_k w_k \\cdot v^{(k)}` per site, where
    :math:`v^{(k)}` is component ``k``'s constrained-space
    sample. Because the constrained-space sites' supports are
    not in general convex (e.g. a Cholesky factor on
    `torch.distributions.constraints.corr_cholesky`), the
    soft mixture can drift outside any single component's
    support during training; the categorical-pick fallback in
    `hard_rsample` returns a single component's sample
    for use at inference time.
    """
    gumbel_logits = (
        self.mixture_logits
        - torch.empty_like(self.mixture_logits).exponential_().log()
    )
    w = F.softmax(gumbel_logits / self._temperature, dim=-1)
    component_samples = [comp.rsample(x) for comp in self.components]

    result: dict[str, torch.Tensor] = {}
    for site_name in self._registry.names:
        stacked = torch.stack(
            [comp_samples[site_name] for comp_samples in component_samples],
            dim=0,
        )
        # Broadcast w against the stacked shape.
        broadcast_shape = (self.num_components,) + (1,) * (stacked.dim() - 1)
        result[site_name] = (w.reshape(broadcast_shape) * stacked).sum(dim=0)
    return result

hard_rsample

hard_rsample(x: Tensor) -> dict[str, Tensor]

Categorical-pick variant: sample a component index and return that component's draw verbatim. Use at inference time when soft-mixture interpolation would violate a support constraint.

Source code in src/quivers/inference/guides/mixture.py

def hard_rsample(self, x: torch.Tensor) -> dict[str, torch.Tensor]:
    """Categorical-pick variant: sample a component index and
    return that component's draw verbatim. Use at inference
    time when soft-mixture interpolation would violate a
    support constraint."""
    probs = F.softmax(self.mixture_logits, dim=-1)
    k = int(torch.distributions.Categorical(probs=probs).sample().item())
    return self.components[k].rsample(x)

log_prob

log_prob(x: Tensor, sites: dict[str, Tensor]) -> Tensor

Mixture log-density via logsumexp over components.

Source code in src/quivers/inference/guides/mixture.py

def log_prob(
    self,
    x: torch.Tensor,
    sites: dict[str, torch.Tensor],
) -> torch.Tensor:
    """Mixture log-density via logsumexp over components."""
    log_pi = F.log_softmax(self.mixture_logits, dim=-1)
    component_log_probs = torch.stack(
        [comp.log_prob(x, sites) for comp in self.components], dim=0
    )
    return torch.logsumexp(log_pi.unsqueeze(-1) + component_log_probs, dim=0)

AutoLowRankMultivariateNormalGuide

AutoLowRankMultivariateNormalGuide(model: MonadicProgram, observed_names: set[str], rank: int = 5, init_scale: float = 0.1)

Bases: _MVNCommon

Low-rank-plus-diagonal multivariate-Normal guide.

Covariance :math:\Sigma = W W^\top + \mathrm{diag}(\sigma^2) with W of shape :math:(D, r) and :math:\sigma \in \mathbb{R}^{D}_{>0}. Memory :math:O(Dr); sampling and log-density via Woodbury / matrix-determinant lemma in torch.distributions.LowRankMultivariateNormal.

Captures the dominant r posterior correlation directions while remaining tractable for D in the hundreds-to- thousands range, where full-rank is infeasible.

PARAMETER	DESCRIPTION
model	Generative model to build a guide for. TYPE: MonadicProgram
observed_names	Variable names treated as observations. TYPE: set[str]
rank	Number of correlated directions. Default 5. TYPE: int DEFAULT: 5
init_scale	Initial diagonal scale. W is initialized at zero. TYPE: float DEFAULT: 0.1

Source code in src/quivers/inference/guides/multivariate_normal.py

def __init__(
    self,
    model: MonadicProgram,
    observed_names: set[str],
    rank: int = 5,
    init_scale: float = 0.1,
) -> None:
    super().__init__(model, observed_names)
    if rank < 1:
        raise ValueError(
            f"AutoLowRankMultivariateNormalGuide: rank must be >= 1, got {rank}"
        )
    if rank > self._D:
        raise ValueError(
            f"AutoLowRankMultivariateNormalGuide: rank ({rank}) "
            f"cannot exceed total unconstrained dimension "
            f"({self._D})"
        )
    self._rank = rank
    init_diag = torch.full((self._D,), float(init_scale))
    init_diag_raw = torch.log(torch.expm1(init_diag.clamp(min=1e-6)))
    self.cov_diag_raw = nn.Parameter(init_diag_raw)
    self.cov_factor = nn.Parameter(torch.zeros(self._D, rank))

AutoMultivariateNormalGuide

AutoMultivariateNormalGuide(model: MonadicProgram, observed_names: set[str], init_scale: float = 0.1)

Bases: _MVNCommon

Full-rank multivariate-Normal variational guide.

Parameterises a joint Gaussian over the registry's flat unconstrained vector with a learnable lower-triangular Cholesky factor. Captures every pairwise posterior correlation across every latent site, the right choice when posterior couplings are strong (hierarchical regression with crossed random effects, parameter pairs with multiplicative interaction).

PARAMETER	DESCRIPTION
model	Generative model to build a guide for. TYPE: MonadicProgram
observed_names	Variable names treated as observations. TYPE: set[str]
init_scale	Initial diagonal of the Cholesky factor. Default 0.1; the off-diagonal entries start at 0. TYPE: float DEFAULT: 0.1

Source code in src/quivers/inference/guides/multivariate_normal.py

def __init__(
    self,
    model: MonadicProgram,
    observed_names: set[str],
    init_scale: float = 0.1,
) -> None:
    super().__init__(model, observed_names)
    init_diag = torch.full((self._D,), float(init_scale))
    init_diag_raw = torch.log(torch.expm1(init_diag.clamp(min=1e-6)))
    self.scale_diag_raw = nn.Parameter(init_diag_raw)
    self.scale_offdiag = nn.Parameter(torch.zeros(self._D, self._D))

AutoNormalGuide

AutoNormalGuide(model: MonadicProgram, observed_names: set[str], init_scale: float = 0.1)

Bases: Guide

Mean-field Normal guide with per-site constrained-support bijector.

PARAMETER	DESCRIPTION
model	Generative model to build a guide for. TYPE: MonadicProgram
observed_names	Variable names treated as observations (skipped in the guide; their values flow through the conditioning data dict at trace time). TYPE: set[str]
init_scale	Initial scale (in unconstrained space) of every latent. Default 0.1; small enough to keep the guide near its prior at the start of optimisation. TYPE: float DEFAULT: 0.1

Source code in src/quivers/inference/guides/normal.py

def __init__(
    self,
    model: MonadicProgram,
    observed_names: set[str],
    init_scale: float = 0.1,
) -> None:
    super().__init__()
    self._registry = self.build_registry(model, observed_names)
    init_log_scale = float(torch.tensor(init_scale).log().item())

    for site in self._registry.sites.values():
        self.register_parameter(
            f"loc_{site.name}",
            nn.Parameter(torch.zeros(site.unconstrained_shape)),
        )
        self.register_parameter(
            f"log_scale_{site.name}",
            nn.Parameter(torch.full(site.unconstrained_shape, init_log_scale)),
        )

rsample

rsample(x: Tensor) -> dict[str, Tensor]

Reparameterized mean-field Normal-then-bijector sample.

Source code in src/quivers/inference/guides/normal.py

def rsample(self, x: torch.Tensor) -> dict[str, torch.Tensor]:
    """Reparameterized mean-field Normal-then-bijector sample."""
    batch = x.shape[0]
    result: dict[str, torch.Tensor] = {}
    for site in self._registry.sites.values():
        _, v = self._sample_site(site, batch)
        result[site.name] = v
    return result

log_prob

log_prob(x: Tensor, sites: dict[str, Tensor]) -> Tensor

Pushforward log-density at constrained values sites.

Uses the change-of-variables identity:

log q(v) = log Normal(z; loc, scale) + log|det J_{T^{-1}}(v)|

where z = bijector.inv(v). The plate / scalar shape dispatch matches rsample's convention.

Source code in src/quivers/inference/guides/normal.py

def log_prob(
    self,
    x: torch.Tensor,
    sites: dict[str, torch.Tensor],
) -> torch.Tensor:
    """Pushforward log-density at constrained values ``sites``.

    Uses the change-of-variables identity:

        log q(v) = log Normal(z; loc, scale) + log|det J_{T^{-1}}(v)|

    where ``z = bijector.inv(v)``. The plate / scalar shape
    dispatch matches `rsample`'s convention.
    """
    batch = x.shape[0]
    total = torch.zeros(batch, device=x.device)
    for site in self._registry.sites.values():
        if site.name not in sites:
            continue
        v = sites[site.name]
        if site.constrained_dim == 1 and v.dim() == (1 if site.is_plate else 1):
            v = v.unsqueeze(-1)
        z = site.bijector.inv(v)
        loc = self._loc(site.name)
        scale = self._scale(site.name)
        if site.is_plate:
            # Plate latent: single shared sample, scalar density
            # broadcast against the batch accumulator.
            log_q_z = D.Normal(loc, scale).log_prob(z)
            log_abs_det = site.bijector.inv.log_abs_det_jacobian(v, z)
            contribution = log_q_z.reshape(-1).sum() + log_abs_det.reshape(-1).sum()
            total = total + contribution
        else:
            loc_b = loc.unsqueeze(0).expand(batch, *site.unconstrained_shape)
            scale_b = scale.unsqueeze(0).expand(batch, *site.unconstrained_shape)
            log_q_z = D.Normal(loc_b, scale_b).log_prob(z)
            log_abs_det = site.bijector.inv.log_abs_det_jacobian(v, z)
            while log_q_z.dim() > 1:
                log_q_z = log_q_z.sum(dim=-1)
            while log_abs_det.dim() > 1:
                log_abs_det = log_abs_det.sum(dim=-1)
            total = total + log_q_z + log_abs_det
    return total

Latent registry

Every guide and MCMC kernel reads the model's latent sites through a LatentRegistry.

registry

Latent-site registry: the canonical introspection surface that every variational guide and MCMC kernel consumes.

A LatentRegistry is built once from a quivers.continuous.programs.MonadicProgram and a set of observed site names. It walks the model's steps linear IR exactly once, classifies each latent into "plate" vs "scalar", records the prior support and per-site dimensionality, builds the matching torch.distributions.transforms.Transform via torch.distributions.constraint_registry.biject_to, and caches a fused flatten / unflatten pipeline so a single unconstrained vector can be turned into a fully constrained site dict (and back) in one pass.

The registry is the only place in the inference layer that knows about PlateDraw, the layout of MonadicProgram.steps, or the per-site constraint-to-bijector mapping. Every downstream component (Guide, MCMCKernel, the flow-based / structured / mixture guides) receives a fully populated registry and operates against its flat-vector and per-site dict APIs.

Categorical interpretation

Each latent site corresponds to a draw in the model's Kleisli arrow Γ → 𝒢(τ). The site's prior is a parameterized family F : Θ → 𝒢(B) with a support constraint C ⊆ B. The bijector T : ℝ^d → C is the biject_to lift of C; its log-determinant gives the change-of-variables term in any pushforward density.

For a plate v : A → B the latent lives in C^|A|; the registry handles the leading plate axis (size |A|) separately from the trailing event axis (the support dimension), and the unconstrained flat vector concatenates plate × event across every site in declaration order.

LatentSite dataclass

LatentSite(name: str, morphism: ContinuousMorphism, support: Constraint, bijector: Transform, constrained_dim: int, unconstrained_dim: int, is_plate: bool, plate_index_size: int, spec_index: int, flat_offset: int, flat_length: int)

Frozen record describing a single latent site.

ATTRIBUTE	DESCRIPTION
name	Variable name (the bound name in the program body). TYPE: str
morphism	The site's runtime morphism (the prior). Held by reference so guides / kernels can call morphism.log_prob directly when they need the per-site prior contribution. TYPE: ContinuousMorphism
support	Support of the prior. Drives the choice of bijector. TYPE: Constraint
bijector	biject_to(support), cached at registry construction. TYPE: Transform
constrained_dim	Event dimension on the constrained side (for a scalar site, 1; for a Dirichlet on a 5-simplex, 5). TYPE: int
unconstrained_dim	Event dimension on the unconstrained side (for a Dirichlet on a 5-simplex, 4; otherwise equals constrained_dim). TYPE: int
is_plate	Whether this site is a plate draw (PlateDraw). TYPE: bool
plate_index_size	|A| for a plate site, 0 otherwise. TYPE: int
spec_index	Position in model.steps. TYPE: int
flat_offset	Index of this site's first element in the registry's flat unconstrained vector. TYPE: int
flat_length	Number of elements this site occupies in the flat vector (= plate_index_size * unconstrained_dim for plate sites, unconstrained_dim for scalar sites). TYPE: int

constrained_shape property

constrained_shape: tuple[int, ...]

The trace-side shape of this site's value (after bijector forward / squeezing scalar event dims).

unconstrained_shape property

unconstrained_shape: tuple[int, ...]

The variational-side shape of this site's value in unconstrained space (before the bijector).

LatentRegistry

LatentRegistry(sites: dict[str, LatentSite], observed_names: frozenset[str], model: MonadicProgram)

The complete introspection result for one (model, observed-set) pair.

Constructed once via from_model; immutable thereafter. Provides:

• Iteration over latent sites in declaration order (sites, plate_sites, scalar_sites).
• Total unconstrained dimensionality (total_unconstrained_dim), the dim of the flat vector HMC and full-rank Gaussian guides operate on.
• Round-tripping between a flat unconstrained vector and a dict of constrained per-site tensors (unflatten_unconstrained, flatten_constrained).
• Bijector forward / inverse with Jacobian accumulation (to_constrained, to_unconstrained).

All operations are vectorized, with no per-site Python loops on the hot path beyond the construction-time setup.

Source code in src/quivers/inference/registry.py

def __init__(
    self,
    sites: dict[str, LatentSite],
    observed_names: frozenset[str],
    model: MonadicProgram,
) -> None:
    self._sites: dict[str, LatentSite] = sites
    self._observed_names: frozenset[str] = observed_names
    self._model: MonadicProgram = model
    self._total_unconstrained_dim: int = sum(s.flat_length for s in sites.values())

sites property

sites: dict[str, LatentSite]

Every latent site, keyed by name, in declaration order.

plate_sites property

plate_sites: dict[str, LatentSite]

The subset of sites that are PlateDraw-shaped.

scalar_sites property

scalar_sites: dict[str, LatentSite]

The subset of sites that are not plates.

names property

names: tuple[str, ...]

Latent site names in declaration order.

observed_names property

observed_names: frozenset[str]

The observed-name set used to build this registry.

model property

model: MonadicProgram

The underlying model.

total_unconstrained_dim property

total_unconstrained_dim: int

Sum of every site's flat-vector length.

This is the dim of the unconstrained vector that full-rank Gaussian guides and HMC operate on.

from_model classmethod

from_model(model: MonadicProgram, observed_names: set[str] | frozenset[str]) -> 'LatentRegistry'

Walk model.steps and assemble the registry.

PARAMETER	DESCRIPTION
model	Compiled probabilistic program. TYPE: MonadicProgram
observed_names	Variable names treated as observations. Sample sites with names in this set are skipped (their values come from the conditioning data dict at trace time, not from a guide). TYPE: set[str] | frozenset[str]

RETURNS	DESCRIPTION
LatentRegistry	Fully populated registry with each latent's morphism, support, bijector, dims, and flat-vector offsets.

Source code in src/quivers/inference/registry.py

@classmethod
def from_model(
    cls,
    model: MonadicProgram,
    observed_names: set[str] | frozenset[str],
) -> "LatentRegistry":
    """Walk ``model.steps`` and assemble the registry.

    Parameters
    ----------
    model : MonadicProgram
        Compiled probabilistic program.
    observed_names : set[str] | frozenset[str]
        Variable names treated as observations. Sample sites
        with names in this set are skipped (their values come
        from the conditioning data dict at trace time, not
        from a guide).

    Returns
    -------
    LatentRegistry
        Fully populated registry with each latent's morphism,
        support, bijector, dims, and flat-vector offsets.
    """
    observed: frozenset[str] = frozenset(observed_names)
    sites: dict[str, LatentSite] = {}
    flat_cursor = 0

    for spec_idx, spec in enumerate(model.steps):
        if not isinstance(spec, Draw | Observe):
            continue
        for var in spec.names:
            if var in observed:
                continue
            morph = model.step_module(spec)
            if morph is None:
                raise RuntimeError(
                    f"LatentRegistry: morphism module for site "
                    f"{var!r} is unregistered on the model"
                )
            morph_cm = cast(ContinuousMorphism, morph)
            is_plate = isinstance(morph_cm, PlateDraw)
            if is_plate:
                plate = cast(PlateDraw, morph_cm)
                inner = plate.family
                inner_cod = inner.codomain
                if isinstance(inner_cod, ContinuousSpace):
                    constrained_dim = max(1, int(inner_cod.dim))
                else:
                    constrained_dim = 1
                support = inner.support
                plate_index_size = plate.index_size
            else:
                cod = morph_cm.codomain
                if isinstance(cod, ContinuousSpace):
                    total_dim = int(cod.dim)
                    constrained_dim = max(1, total_dim // max(1, len(spec.names)))
                else:
                    constrained_dim = 1
                support = morph_cm.support
                plate_index_size = 0

            unconstrained_dim = _unconstrained_event_dim(support, constrained_dim)
            bijector = biject_to(support)

            flat_length = (
                plate_index_size * unconstrained_dim
                if is_plate
                else unconstrained_dim
            )

            sites[var] = LatentSite(
                name=var,
                morphism=morph_cm,
                support=support,
                bijector=bijector,
                constrained_dim=constrained_dim,
                unconstrained_dim=unconstrained_dim,
                is_plate=is_plate,
                plate_index_size=plate_index_size,
                spec_index=spec_idx,
                flat_offset=flat_cursor,
                flat_length=flat_length,
            )
            flat_cursor += flat_length

    return cls(sites, observed, model)

unflatten_unconstrained

unflatten_unconstrained(vec: Tensor) -> dict[str, Tensor]

Reshape a flat unconstrained vector into per-site unconstrained tensors.

PARAMETER	DESCRIPTION
vec	Shape (..., total_unconstrained_dim). Leading axes (batch, particles, chains) are preserved on every site's output. TYPE: Tensor

RETURNS	DESCRIPTION
dict[str, Tensor]	{name: site_tensor} where each site_tensor has shape vec.shape[:-1] + site.unconstrained_shape.

Source code in src/quivers/inference/registry.py

def unflatten_unconstrained(self, vec: torch.Tensor) -> dict[str, torch.Tensor]:
    """Reshape a flat unconstrained vector into per-site
    unconstrained tensors.

    Parameters
    ----------
    vec : torch.Tensor
        Shape ``(..., total_unconstrained_dim)``. Leading
        axes (batch, particles, chains) are preserved on every
        site's output.

    Returns
    -------
    dict[str, torch.Tensor]
        ``{name: site_tensor}`` where each ``site_tensor`` has
        shape ``vec.shape[:-1] + site.unconstrained_shape``.
    """
    if vec.shape[-1] != self._total_unconstrained_dim:
        raise ValueError(
            f"LatentRegistry.unflatten_unconstrained: expected "
            f"trailing dim {self._total_unconstrained_dim}, got "
            f"{vec.shape[-1]}"
        )
    out: dict[str, torch.Tensor] = {}
    for site in self._sites.values():
        chunk = vec[..., site.flat_offset : site.flat_offset + site.flat_length]
        out[site.name] = chunk.reshape(*vec.shape[:-1], *site.unconstrained_shape)
    return out

flatten_unconstrained

flatten_unconstrained(sites: dict[str, Tensor]) -> Tensor

Concatenate per-site unconstrained tensors into a single flat vector.

Inverse of unflatten_unconstrained (up to leading batch axes). The trailing batch axes of every site tensor must agree.

Source code in src/quivers/inference/registry.py

def flatten_unconstrained(self, sites: dict[str, torch.Tensor]) -> torch.Tensor:
    """Concatenate per-site unconstrained tensors into a single
    flat vector.

    Inverse of `unflatten_unconstrained` (up to leading
    batch axes). The trailing batch axes of every site tensor
    must agree.
    """
    parts: list[torch.Tensor] = []
    leading_shape: tuple[int, ...] | None = None
    for site in self._sites.values():
        if site.name not in sites:
            raise KeyError(
                f"LatentRegistry.flatten_unconstrained: missing site {site.name!r}"
            )
        t = sites[site.name]
        trail = len(site.unconstrained_shape)
        site_leading = t.shape[: t.dim() - trail] if trail else t.shape
        if leading_shape is None:
            leading_shape = tuple(site_leading)
        elif tuple(site_leading) != leading_shape:
            raise ValueError(
                f"LatentRegistry.flatten_unconstrained: site "
                f"{site.name!r} has leading shape {tuple(site_leading)} "
                f"which disagrees with previously-seen "
                f"{leading_shape}"
            )
        parts.append(t.reshape(*site_leading, site.flat_length))
    if not parts:
        device = next(iter(sites.values())).device if sites else torch.device("cpu")
        return torch.zeros((0,), device=device)
    return torch.cat(parts, dim=-1)

to_constrained

to_constrained(unconstrained: dict[str, Tensor]) -> tuple[dict[str, Tensor], dict[str, Tensor]]

Push every site's unconstrained value through its bijector.

RETURNS	DESCRIPTION
constrained	Per-site constrained values shaped according to LatentSite.constrained_shape. The trailing length-1 event axis of a scalar site is squeezed off to match the trace-side shape convention. TYPE: dict[str, Tensor]
log_abs_dets	Per-site log |det dT/dz| returned by the bijector, at the bijector's natural shape (event axes preserved). Callers aggregate as the objective requires: for ELBO we sum over event axes; for IWAE we keep particles as a leading axis; etc. TYPE: dict[str, Tensor]

Source code in src/quivers/inference/registry.py

def to_constrained(
    self, unconstrained: dict[str, torch.Tensor]
) -> tuple[dict[str, torch.Tensor], dict[str, torch.Tensor]]:
    """Push every site's unconstrained value through its
    bijector.

    Returns
    -------
    constrained : dict[str, torch.Tensor]
        Per-site constrained values shaped according to
        `LatentSite.constrained_shape`. The trailing
        length-1 event axis of a scalar site is squeezed off
        to match the trace-side shape convention.
    log_abs_dets : dict[str, torch.Tensor]
        Per-site ``log |det dT/dz|`` returned by the bijector,
        at the bijector's natural shape (event axes preserved).
        Callers aggregate as the objective requires: for ELBO
        we sum over event axes; for IWAE we keep particles as
        a leading axis; etc.
    """
    constrained: dict[str, torch.Tensor] = {}
    log_abs_dets: dict[str, torch.Tensor] = {}
    for site in self._sites.values():
        if site.name not in unconstrained:
            raise KeyError(
                f"LatentRegistry.to_constrained: missing site {site.name!r}"
            )
        z = unconstrained[site.name]
        v = site.bijector(z)
        log_abs_dets[site.name] = site.bijector.log_abs_det_jacobian(z, v)
        if site.constrained_dim == 1 and v.dim() >= 1 and v.shape[-1] == 1:
            v = v.squeeze(-1)
        constrained[site.name] = v
    return constrained, log_abs_dets

to_unconstrained

to_unconstrained(constrained: dict[str, Tensor]) -> tuple[dict[str, Tensor], dict[str, Tensor]]

Inverse of to_constrained.

Source code in src/quivers/inference/registry.py

def to_unconstrained(
    self, constrained: dict[str, torch.Tensor]
) -> tuple[dict[str, torch.Tensor], dict[str, torch.Tensor]]:
    """Inverse of `to_constrained`."""
    unconstrained: dict[str, torch.Tensor] = {}
    log_abs_dets: dict[str, torch.Tensor] = {}
    for site in self._sites.values():
        if site.name not in constrained:
            raise KeyError(
                f"LatentRegistry.to_unconstrained: missing site {site.name!r}"
            )
        v = constrained[site.name]
        if site.constrained_dim == 1 and v.dim() == (1 if site.is_plate else 0):
            v = v.unsqueeze(-1)
        z = site.bijector.inv(v)
        log_abs_dets[site.name] = site.bijector.inv.log_abs_det_jacobian(v, z)
        unconstrained[site.name] = z
    return unconstrained, log_abs_dets

aggregate_log_abs_det staticmethod

aggregate_log_abs_det(log_abs_dets: dict[str, Tensor], leading_shape: tuple[int, ...] = ()) -> Tensor

Sum every per-site log_abs_det over its event axes and produce a single tensor broadcast to leading_shape.

This is the canonical aggregation an ELBO needs. Objectives that want particle-aware or per-site granularity can skip this helper and aggregate themselves.

Source code in src/quivers/inference/registry.py

@staticmethod
def aggregate_log_abs_det(
    log_abs_dets: dict[str, torch.Tensor],
    leading_shape: tuple[int, ...] = (),
) -> torch.Tensor:
    """Sum every per-site ``log_abs_det`` over its event axes
    and produce a single tensor broadcast to ``leading_shape``.

    This is the canonical aggregation an ELBO needs. Objectives
    that want particle-aware or per-site granularity can skip
    this helper and aggregate themselves.
    """
    if not log_abs_dets:
        return torch.zeros(leading_shape) if leading_shape else torch.zeros(())
    total: torch.Tensor | None = None
    n_leading = len(leading_shape)
    for site_lad in log_abs_dets.values():
        # Sum away every axis beyond the leading shape; what
        # remains is broadcast against the leading shape.
        while site_lad.dim() > n_leading:
            site_lad = site_lad.sum(dim=-1)
        total = site_lad if total is None else total + site_lad
    assert total is not None
    return total

zero_unconstrained

zero_unconstrained(*, leading_shape: tuple[int, ...] = (), device: device | None = None, dtype: dtype | None = None) -> dict[str, Tensor]

Return a dict of zero-initialized unconstrained site tensors with the supplied leading batch shape.

Useful for HMC chain initialisation and AutoLaplace's MAP starting point.

Source code in src/quivers/inference/registry.py

def zero_unconstrained(
    self,
    *,
    leading_shape: tuple[int, ...] = (),
    device: torch.device | None = None,
    dtype: torch.dtype | None = None,
) -> dict[str, torch.Tensor]:
    """Return a dict of zero-initialized unconstrained site
    tensors with the supplied leading batch shape.

    Useful for HMC chain initialisation and AutoLaplace's MAP
    starting point.
    """
    out: dict[str, torch.Tensor] = {}
    for site in self._sites.values():
        shape = leading_shape + site.unconstrained_shape
        out[site.name] = torch.zeros(shape, device=device, dtype=dtype)
    return out

randn_unconstrained

randn_unconstrained(*, leading_shape: tuple[int, ...] = (), device: device | None = None, dtype: dtype | None = None, scale: float = 1.0) -> dict[str, Tensor]

Standard-normal-initialized unconstrained site tensors.

Source code in src/quivers/inference/registry.py

def randn_unconstrained(
    self,
    *,
    leading_shape: tuple[int, ...] = (),
    device: torch.device | None = None,
    dtype: torch.dtype | None = None,
    scale: float = 1.0,
) -> dict[str, torch.Tensor]:
    """Standard-normal-initialized unconstrained site tensors."""
    out: dict[str, torch.Tensor] = {}
    for site in self._sites.values():
        shape = leading_shape + site.unconstrained_shape
        out[site.name] = scale * torch.randn(shape, device=device, dtype=dtype)
    return out

Flow transforms

The learnable bijections that AutoNormalizingFlow and its subclasses stack. A custom flow subclasses TransformModule.

transforms

Bijector and normalizing-flow primitives.

This module collects the differentiable bijections quivers's variational guides chain through. Every primitive is a subclass of torch.distributions.transforms.Transform, so they compose under torch.distributions.transforms.ComposeTransform, plug into torch.distributions.transforms.biject_to-style pipelines, and interop cleanly with torch.distributions.transformed_distribution.TransformedDistribution.

Coverage

• AffineCouplingTransform: RealNVP coupling layer (Dinh, Sohl-Dickstein, Bengio 2017, doi:10.48550/arXiv.1605.08803 <https://doi.org/10.48550/arXiv.1605.08803>_). Splits the input into two halves, scales / shifts one half conditional on the other.
• MaskedAutoregressiveTransform: MAF (Papamakarios, Pavlakou, Murray 2017, doi:10.48550/arXiv.1705.07057 <https://doi.org/10.48550/arXiv.1705.07057>_). Forward pass is parallel (single masked MLP call); inverse is sequential (per-coordinate). The mass-density form used by variational inference.
• InverseAutoregressiveTransform: IAF (Kingma, Salimans, Jozefowicz et al. 2016, doi:10.48550/arXiv.1606.04934 <https://doi.org/10.48550/arXiv.1606.04934>_). The dual of MAF: sampling is parallel, density evaluation is sequential. Preferred when you sample more often than score (which is true for variational guides).
• NeuralSplineCouplingTransform: monotone rational-quadratic spline coupling (Durkan, Bekasov, Murray, Papamakarios 2019, doi:10.48550/arXiv.1906.04032 <https://doi.org/10.48550/arXiv.1906.04032>_). Strictly more expressive than affine coupling at the same parameter budget.
• LULinearTransform: LU-decomposed linear permutation (Kingma, Dhariwal 2018 Glow, doi:10.48550/arXiv.1807.03039 <https://doi.org/10.48550/arXiv.1807.03039>_). Cheap learnable invertible linear layer; the standard inter- coupling-layer mixer.
• BatchNormTransform: running-statistics-based BN as a normalizing-flow transform (Dinh, Sohl-Dickstein, Bengio 2017, same paper as RealNVP).

Every primitive is fully implemented (forward, inverse, log_abs_det_jacobian). No primitive raises NotImplementedError under any input its type signature accepts. The masked-network factories are designed so callers can pass a callable / module rather than configuring topology through string flags.

TransformModule

TransformModule(cache_size: int = 0)

Bases: Transform, Module

Cooperative base for transforms that also need to register torch.nn.Parameter and torch.nn.Module.register_buffer attributes.

torch.distributions.transforms.Transform is not a torch.nn.Module, so subclassing both is required for any flow primitive that holds learnable parameters or buffers (every primitive in this module). The cooperative base calls both __init__\ s in MRO order, nn.Module first so attribute storage is set up before Transform populates the cache-size machinery.

Source code in src/quivers/inference/transforms.py

def __init__(self, cache_size: int = 0) -> None:
    nn.Module.__init__(self)
    Transform.__init__(self, cache_size=cache_size)

AffineCouplingTransform

AffineCouplingTransform(dim: int, net: Module, mask: Tensor)

Bases: TransformModule

RealNVP-style affine coupling transform.

Splits the input vector into two parts: an "identity" part x_a and a "transformed" part x_b. The transformed part is shifted and scaled by parameters output from a neural network that sees only the identity part. The Jacobian is triangular, so its log-determinant is the sum of the log-scales on x_b.

PARAMETER	DESCRIPTION
dim	Total event dimension. TYPE: int
net	Module mapping x_a of shape (..., dim_a) to (log_scale_b, shift_b) of shape (..., 2 * dim_b). The factory make_coupling_mlp builds a sensible default. TYPE: Module
mask	Boolean mask of shape (dim,): True entries form the identity part, False entries form the transformed part. Permuting the mask between consecutive coupling layers is the canonical RealNVP recipe. TYPE: Tensor

Source code in src/quivers/inference/transforms.py

def __init__(self, dim: int, net: nn.Module, mask: torch.Tensor) -> None:
    TransformModule.__init__(self, cache_size=1)
    if mask.dtype != torch.bool:
        raise TypeError(
            f"AffineCouplingTransform: mask must be bool, got {mask.dtype}"
        )
    if mask.shape != (dim,):
        raise ValueError(
            f"AffineCouplingTransform: mask shape {tuple(mask.shape)} "
            f"does not match dim={dim}"
        )
    n_a = int(mask.sum().item())
    n_b = dim - n_a
    if n_a == 0 or n_b == 0:
        raise ValueError(
            f"AffineCouplingTransform: mask must split the input into "
            f"two non-empty parts; got {n_a} identity / {n_b} transformed"
        )
    self._dim = dim
    self._n_a = n_a
    self._n_b = n_b
    self.net = net
    self.register_buffer("_mask", mask, persistent=True)

MADE

MADE(dim: int, n_per_dim: int, *, hidden: int = 64, n_hidden_layers: int = 2, ordering: Tensor | None = None)

Bases: Module

Masked Autoencoder for Distribution Estimation (Germain, Gregor, Mnih, Larochelle 2015, doi:10.48550/arXiv.1502.03509 <https://doi.org/10.48550/arXiv.1502.03509>_).

Outputs n_per_dim parameters per input coordinate. The masking guarantees that output parameters for coordinate j depend only on input coordinates < j in the supplied ordering, which is the autoregressive property MAF / IAF exploit.

Source code in src/quivers/inference/transforms.py

def __init__(
    self,
    dim: int,
    n_per_dim: int,
    *,
    hidden: int = 64,
    n_hidden_layers: int = 2,
    ordering: torch.Tensor | None = None,
) -> None:
    super().__init__()
    if ordering is None:
        ordering = torch.arange(dim)
    if ordering.shape != (dim,) or ordering.dtype != torch.long:
        raise TypeError("MADE: ordering must be a long tensor of shape (dim,)")
    self._dim = dim
    self._n_per_dim = n_per_dim
    self.register_buffer("_ordering", ordering, persistent=True)
    # Assign integer "degrees" to each hidden unit so that the
    # masking enforces autoregressivity.
    hidden_degrees = [
        torch.randint(1, dim, (hidden,)) for _ in range(n_hidden_layers)
    ]
    # Input degrees: ordering itself; output degrees: ordering
    # repeated n_per_dim times (one block per output parameter).
    in_degrees = ordering
    layers: list[nn.Module] = []
    prev_degrees = in_degrees
    prev_size = dim
    for h_degrees in hidden_degrees:
        mask = (h_degrees.unsqueeze(1) >= prev_degrees.unsqueeze(0)).to(
            torch.float32
        )
        layers.append(_MaskedLinear(prev_size, hidden, mask))
        layers.append(nn.ReLU())
        prev_degrees = h_degrees
        prev_size = hidden
    # Final layer: strict-less-than for the autoregressive
    # property; one output block per parameter.
    out_degrees = ordering.repeat(n_per_dim)
    final_mask = (out_degrees.unsqueeze(1) > prev_degrees.unsqueeze(0)).to(
        torch.float32
    )
    final = _MaskedLinear(prev_size, dim * n_per_dim, final_mask)
    nn.init.zeros_(final.weight)
    nn.init.zeros_(final.bias)
    layers.append(final)
    self.net = nn.Sequential(*layers)

MaskedAutoregressiveTransform

MaskedAutoregressiveTransform(made: MADE)

Bases: TransformModule

Masked Autoregressive Flow (MAF) layer.

Forward (density) is parallel, one MADE call. Inverse (sampling) is sequential in the autoregressive order.

PARAMETER	DESCRIPTION
made	Masked network producing per-coordinate (shift, log_scale). TYPE: MADE

Source code in src/quivers/inference/transforms.py

def __init__(self, made: MADE) -> None:
    super().__init__(cache_size=1)
    if made.n_per_dim != 2:
        raise ValueError(
            f"MaskedAutoregressiveTransform: MADE must have "
            f"n_per_dim == 2 (shift, log_scale), got {made.n_per_dim}"
        )
    self.made = made

InverseAutoregressiveTransform

InverseAutoregressiveTransform(made: MADE)

Bases: TransformModule

Inverse Autoregressive Flow (IAF) layer.

Dual of MaskedAutoregressiveTransform: sampling (forward) is parallel, density (inverse) is sequential. The preferred direction for variational guides because we sample more often than we score.

PARAMETER	DESCRIPTION
made	Masked network producing per-coordinate (shift, log_scale). TYPE: MADE

Source code in src/quivers/inference/transforms.py

def __init__(self, made: MADE) -> None:
    super().__init__(cache_size=1)
    if made.n_per_dim != 2:
        raise ValueError(
            f"InverseAutoregressiveTransform: MADE must have "
            f"n_per_dim == 2 (shift, log_scale), got {made.n_per_dim}"
        )
    self.made = made

NeuralSplineCouplingTransform

NeuralSplineCouplingTransform(dim: int, net: Module, mask: Tensor, *, num_bins: int = 8, tail_bound: float = 3.0)

Bases: TransformModule

Monotone rational-quadratic spline coupling (Durkan-Bekasov-Murray-Papamakarios 2019).

Same coupling structure as AffineCouplingTransform, but the per-coordinate transform on x_b is a monotone rational-quadratic spline rather than an affine map. The network outputs (unnormalized_widths, unnormalized_heights, unnormalized_derivatives) per transformed coordinate; num_bins knots per coordinate.

PARAMETER	DESCRIPTION
dim	Event dimension. TYPE: int
net	Maps x_a of shape (..., n_a) to spline parameters of shape (..., n_b * (3 * num_bins - 1)). make_coupling_mlp produces a sensible default with n_out = n_b * (3 * num_bins - 1). TYPE: Module
mask	Boolean mask of shape (dim,). TYPE: Tensor
num_bins	Number of spline knots. TYPE: int DEFAULT: 8
tail_bound	Domain of the spline; outside [-tail_bound, tail_bound] the transform is the identity. TYPE: float DEFAULT: 3.0

Source code in src/quivers/inference/transforms.py

def __init__(
    self,
    dim: int,
    net: nn.Module,
    mask: torch.Tensor,
    *,
    num_bins: int = 8,
    tail_bound: float = 3.0,
) -> None:
    super().__init__(cache_size=1)
    if mask.dtype != torch.bool or mask.shape != (dim,):
        raise ValueError(
            "NeuralSplineCouplingTransform: mask must be bool of shape "
            f"(dim,) = ({dim},); got {tuple(mask.shape)} {mask.dtype}"
        )
    n_a = int(mask.sum().item())
    n_b = dim - n_a
    if n_a == 0 or n_b == 0:
        raise ValueError(
            "NeuralSplineCouplingTransform: mask must split the input "
            f"into two non-empty parts; got {n_a} / {n_b}"
        )
    self._dim = dim
    self._n_a = n_a
    self._n_b = n_b
    self._num_bins = num_bins
    self._tail_bound = tail_bound
    self.net = net
    self.register_buffer("_mask", mask, persistent=True)

LULinearTransform

LULinearTransform(dim: int, perm: Tensor | None = None)

Bases: TransformModule

Invertible linear layer parameterized as :math:A = P L U (Kingma & Dhariwal 2018 Glow).

The unit-lower-triangular L and unit-upper-triangular U factors are stored along with their respective off-diagonal weights; the diagonal of U is stored as a log-scale so the layer remains invertible and the log-determinant is the sum of the log-scales. P is a fixed permutation chosen at construction.

Source code in src/quivers/inference/transforms.py

def __init__(self, dim: int, perm: torch.Tensor | None = None) -> None:
    super().__init__(cache_size=1)
    self._dim = dim
    # Lower-triangular off-diagonal (strict)
    self._lower = nn.Parameter(torch.zeros(dim, dim))
    # Upper-triangular off-diagonal (strict)
    self._upper = nn.Parameter(torch.zeros(dim, dim))
    # Log-magnitude of U's diagonal; sign baked into the
    # signed-zero init below so the layer starts as identity.
    self._log_diag = nn.Parameter(torch.zeros(dim))
    if perm is None:
        perm = torch.arange(dim)
    if perm.shape != (dim,):
        raise ValueError(
            f"LULinearTransform: perm shape {tuple(perm.shape)} != ({dim},)"
        )
    self.register_buffer("_perm", perm.long(), persistent=True)
    # The fixed permutation matrix's log-det is 0 (it's a
    # signed permutation; we take +1 here).
    self.register_buffer(
        "_lower_mask",
        torch.tril(torch.ones(dim, dim), diagonal=-1),
        persistent=True,
    )
    self.register_buffer(
        "_upper_mask",
        torch.triu(torch.ones(dim, dim), diagonal=1),
        persistent=True,
    )
    self.register_buffer("_eye", torch.eye(dim), persistent=True)

BatchNormTransform

BatchNormTransform(dim: int, *, momentum: float = 0.1, eps: float = 1e-05)

Bases: TransformModule

Batch-normalisation as a normalizing-flow layer (Dinh, Sohl-Dickstein, Bengio 2017).

In training mode uses the batch statistics; in eval mode uses the cached running statistics. This module is an torch.nn.Module to register the running buffers, and a torch.distributions.transforms.Transform so it composes with the rest of the flow stack.

Source code in src/quivers/inference/transforms.py

def __init__(
    self,
    dim: int,
    *,
    momentum: float = 0.1,
    eps: float = 1e-5,
) -> None:
    super().__init__(cache_size=1)
    self._dim = dim
    self._momentum = momentum
    self._eps = eps
    self._gamma = nn.Parameter(torch.zeros(dim))  # log-scale
    self._beta = nn.Parameter(torch.zeros(dim))
    self.register_buffer("_running_mean", torch.zeros(dim), persistent=True)
    self.register_buffer("_running_var", torch.ones(dim), persistent=True)
    self._training = True

make_coupling_mlp

make_coupling_mlp(n_in: int, n_out: int, *, hidden: int = 64, n_hidden_layers: int = 2, activation: type[Module] = ReLU) -> Module

Build an MLP for use as the net argument of AffineCouplingTransform or NeuralSplineCouplingTransform.

The output layer is initialized to zero so the transform starts as the identity, which is a standard normalizing-flow initialisation (see Glow, Real NVP).

Source code in src/quivers/inference/transforms.py

def make_coupling_mlp(
    n_in: int,
    n_out: int,
    *,
    hidden: int = 64,
    n_hidden_layers: int = 2,
    activation: type[nn.Module] = nn.ReLU,
) -> nn.Module:
    """Build an MLP for use as the ``net`` argument of
    `AffineCouplingTransform` or
    `NeuralSplineCouplingTransform`.

    The output layer is initialized to zero so the transform
    starts as the identity, which is a standard normalizing-flow
    initialisation (see Glow, Real NVP).
    """
    layers: list[nn.Module] = [nn.Linear(n_in, hidden), activation()]
    for _ in range(n_hidden_layers - 1):
        layers.append(nn.Linear(hidden, hidden))
        layers.append(activation())
    final = nn.Linear(hidden, n_out)
    nn.init.zeros_(final.weight)
    nn.init.zeros_(final.bias)
    layers.append(final)
    return nn.Sequential(*layers)

alternating_mask

alternating_mask(dim: int, *, even: bool = True) -> Tensor

Boolean mask alternating True / False. Even-indexed entries are True when even=True.

Source code in src/quivers/inference/transforms.py

def alternating_mask(dim: int, *, even: bool = True) -> torch.Tensor:
    """Boolean mask alternating ``True`` / ``False``. Even-indexed
    entries are ``True`` when ``even=True``."""
    indices = torch.arange(dim)
    return (indices % 2 == 0) if even else (indices % 2 == 1)

half_mask

half_mask(dim: int, *, first_half_true: bool = True) -> Tensor

Boolean mask with the first / second half True.

Source code in src/quivers/inference/transforms.py

def half_mask(dim: int, *, first_half_true: bool = True) -> torch.Tensor:
    """Boolean mask with the first / second half ``True``."""
    mask = torch.zeros(dim, dtype=torch.bool)
    half = dim // 2
    if first_half_true:
        mask[:half] = True
    else:
        mask[half:] = True
    return mask

Annealed importance sampling

dais

Differentiable Annealed Importance Sampling (DAIS) guide.

DAIS sits between variational inference and MCMC: it runs a short chain of HMC-like leapfrog steps along an annealing path between a tractable base distribution :math:q_0 (typically AutoNormalGuide) and the target posterior. The trajectory's per-step parameters (step size, base mean / scale, and the inverse-temperature schedule) are all variational parameters trained by SVI. Concretely, every leapfrog operation is reparameterized through the momentum and position so gradients flow end-to-end through the trajectory.

DAIS gives an unbiased lower-bound estimator of the model evidence (Geffner-Domke 2021, doi:10.48550/arXiv.2102.07501 <https://doi.org/10.48550/arXiv.2102.07501>; Zhang et al. 2021, doi:10.48550/arXiv.2107.10211 <https://doi.org/10.48550/arXiv.2107.10211>) that strictly dominates the base guide's ELBO for num_steps >= 1. Combined with multimodal base guides (AutoMixtureGuide) it recovers multimodal posteriors that plain VI misses.

The implementation tracks the auxiliary "extended-target" form: a fresh momentum is sampled at the start of each transition and discarded at the end (no Metropolis correction in the differentiable variant; the integrator is deterministic and the trajectory length is fixed). The variational lower bound absorbs the change-of-variables for the leapfrog flow and the momentum prior / posterior contributions.

AutoDAIS

AutoDAIS(base: Guide, model: MonadicProgram, observations: dict[str, Tensor], x: Tensor | None = None, num_steps: int = 8, leapfrog_steps: int = 1, init_step_size: float = 0.1, init_temperature: float = 0.1)

Bases: Guide

Annealed-importance-sampling guide wrapping a base guide.

PARAMETER	DESCRIPTION
base	Base variational guide (the bridge's start). Must be built against the same model + observed-name set the DAIS guide is for. AutoNormalGuide or AutoMultivariateNormalGuide are the canonical choices. TYPE: Guide
model	Generative model. Required because DAIS needs the target log-density at each annealing step. TYPE: MonadicProgram
observations	Observed-site values that pin the model's likelihood. The DAIS bridge anneals between base and the conditional posterior under these observations; if the observations change the guide must be rebuilt. TYPE: dict[str, Tensor]
x	Program input fed to the model's log-joint. Default torch.zeros(1, 1); only its device / leading batch shape matter. TYPE: Tensor DEFAULT: None
num_steps	Number of HMC-style leapfrog transitions along the path. 0 recovers the base guide exactly. Default 8. TYPE: int DEFAULT: 8
leapfrog_steps	Inner leapfrog substeps per transition. Default 1; higher values trade compute for tighter bounds. TYPE: int DEFAULT: 1
init_step_size	Initial leapfrog step size. Adapted as a learnable scalar. TYPE: float DEFAULT: 0.1
init_temperature	Initial first-bridge inverse temperature. The schedule anneals from this value at step 1 to 1.0 at step num_steps. Stored as num_steps learnable parameters so the schedule is data-adaptive. TYPE: float DEFAULT: 0.1

Source code in src/quivers/inference/dais.py

def __init__(
    self,
    base: Guide,
    model: MonadicProgram,
    observations: dict[str, torch.Tensor],
    x: torch.Tensor | None = None,
    num_steps: int = 8,
    leapfrog_steps: int = 1,
    init_step_size: float = 0.1,
    init_temperature: float = 0.1,
) -> None:
    super().__init__()
    if num_steps < 1:
        raise ValueError(f"AutoDAIS: num_steps must be >= 1, got {num_steps}")
    if leapfrog_steps < 1:
        raise ValueError(
            f"AutoDAIS: leapfrog_steps must be >= 1, got {leapfrog_steps}"
        )
    if init_step_size <= 0:
        raise ValueError(
            f"AutoDAIS: init_step_size must be > 0, got {init_step_size}"
        )
    if not 0.0 < init_temperature < 1.0:
        raise ValueError(
            f"AutoDAIS: init_temperature must be in (0, 1), got {init_temperature}"
        )
    self.base = base
    self._registry = base.registry
    self._model = model
    self._observations = observations
    self._x = x if x is not None else torch.zeros(1, 1)
    self._num_steps = num_steps
    self._leapfrog_steps = leapfrog_steps
    self._D = self._registry.total_unconstrained_dim
    if self._D == 0:
        raise ValueError(
            "AutoDAIS: base registry has zero total "
            "unconstrained dimension; no continuous latents "
            "to anneal over"
        )
    # Learnable trajectory params.
    log_step = torch.log(torch.tensor(float(init_step_size)))
    self.log_step_size = nn.Parameter(log_step.expand(num_steps).clone())
    # Inverse temperatures parameterized through a logit so they
    # stay in (0, 1) and are monotone-friendly. The final beta is
    # 1.0 (full target); the earlier ones are learned.
    # Initial schedule: linear in inverse temperature.
    init_betas = torch.linspace(init_temperature, 1.0, num_steps + 1)[
        :-1
    ]  # length num_steps; last one is fixed at 1.0.
    # Store as raw logits and apply a cumulative-softplus to
    # keep them monotone increasing.
    increments = torch.zeros(num_steps)
    increments[0] = init_betas[0]
    for i in range(1, num_steps):
        increments[i] = init_betas[i] - init_betas[i - 1]
    # Store log(increments) so cumsum of softplus(raw) is positive
    # and monotone.
    raw_increments = torch.log(torch.expm1(increments.clamp(min=1e-4)))
    self.beta_increments_raw = nn.Parameter(raw_increments)

log_prob

log_prob(x: Tensor, sites: dict[str, Tensor]) -> Tensor

Effective log-density of the trajectory's endpoint.

DAIS does not admit a closed-form density at the endpoint because the trajectory is stochastic in the momentum. The differentiable lower-bound estimator instead substitutes :math:\log \pi(z) - \mathrm{log\_weight} for :math:\log q(z) in the ELBO, so when this guide is used with an ELBO objective the per-step ELBO is the DAIS lower bound. We return that quantity here.

Source code in src/quivers/inference/dais.py

def log_prob(
    self,
    x: torch.Tensor,
    sites: dict[str, torch.Tensor],
) -> torch.Tensor:
    """Effective log-density of the trajectory's endpoint.

    DAIS does not admit a closed-form density at the endpoint
    because the trajectory is stochastic in the momentum. The
    differentiable lower-bound estimator instead substitutes
    :math:`\\log \\pi(z) - \\mathrm{log\\_weight}` for
    :math:`\\log q(z)` in the ELBO, so when this guide is used
    with an `ELBO` objective the per-step ELBO is the
    DAIS lower bound. We return that quantity here.
    """
    del sites  # The DAIS log-density is path-dependent, not
    # site-conditional; the standard ELBO call path
    # passes sites = rsample's output, and we ignore
    # them in favour of the cached trajectory weight.
    z_final, log_weight = self._trajectory(x)
    log_pi = self._target_log_density(z_final)
    # log q_DAIS(z) ≈ log π(z) - log_weight, so that
    # ELBO = E[log π(z) - log q_DAIS(z)] = E[log_weight].
    return (log_pi - log_weight).expand(x.shape[0])