Trace

Program trace data structures and trace-based inference.

trace

Execution trace for monadic programs.

A Trace records every sample site visited during program execution, capturing the morphism, sampled or observed value, and log-density at each site. This is the foundation for all inference algorithms: SVI uses traces to compute the ELBO, and conditioning operates by clamping trace sites to observed data.

The trace(program, x, observations) free function is a thin wrapper on top of the effect-handler machinery: it pushes a quivers.effects.trace_handler.TraceHandler onto the active handler stack, invokes quivers.effects.interpreter.run_program, and returns the accumulated Trace. Any other handlers active in the enclosing scope compose with TraceHandler in the standard outer-to-inner order.

The SampleSite and Trace data types are defined in quivers.effects.trace_types and re-exported here.

SampleSite dataclass

SampleSite(name: str, morphism: ContinuousMorphism | None, value: Tensor, log_prob: Tensor, is_observed: bool = False, is_deterministic: bool = False, sampleable: TorchSampleable | None = None, address: SiteAddress | None = None, metadata: dict[str, object] = dict())

Record of a single site in a program trace.

Holds a torch.Tensor per site; not a value type.

PARAMETER	DESCRIPTION
name	Variable name bound at this site. TYPE: str
morphism	The distribution morphism (None for let bindings and scores). TYPE: ContinuousMorphism or None
value	The sampled, observed, or computed value. TYPE: Tensor
log_prob	The density the site contributed to the run's score, after every score transformer on the stack. Zero for let bindings and for intervened sites. TYPE: Tensor
is_observed	Whether this site was conditioned on an observed value. TYPE: bool DEFAULT: False
is_deterministic	Whether this is a let binding, a score step, or an intervened site. TYPE: bool DEFAULT: False
sampleable	The sampleable the site's request carried, with the morphism and the input it was conditioned on; None for let bindings and scores. TYPE: TorchSampleable or None DEFAULT: None
address	The kernel's dynamic address of the site's request, which is one identity for the run's trace and for every handler that answered it. TYPE: SiteAddress or None DEFAULT: None
metadata	Annotations handlers attached to the site. TYPE: dict[str, object] DEFAULT: dict()

Trace dataclass

Trace(sites: dict[str, SampleSite] = dict(), output: Tensor | dict[str, Tensor] | None = None, log_joint: Tensor | None = None)

Complete execution trace of a monadic program.

Mutable accumulator: sites grows as the program executes; not a value type.

PARAMETER	DESCRIPTION
sites	All sites keyed by variable name. TYPE: dict[str, SampleSite] DEFAULT: dict()
output	The program's return value. TYPE: Tensor or dict[str, Tensor] DEFAULT: None
log_joint	Sum of log-densities across all sites. Shape (batch,). TYPE: Tensor DEFAULT: None

stochastic_sites property

stochastic_sites: dict[str, SampleSite]

The sites that are not deterministic.

RETURNS	DESCRIPTION
dict[str, SampleSite]	Every sampled or observed site.

latent_sites property

latent_sites: dict[str, SampleSite]

The sites that are neither observed nor deterministic.

RETURNS	DESCRIPTION
dict[str, SampleSite]	Every sampled latent site.

observed_sites property

observed_sites: dict[str, SampleSite]

The observed sites.

RETURNS	DESCRIPTION
dict[str, SampleSite]	Every site conditioned on an observation.

trace

trace(program: MonadicProgram, x: Tensor, observations: dict[str, Tensor] | None = None) -> Trace

Execute a program and record all sample sites.

A thin wrapper around the handler-aware interpreter: pushes a quivers.effects.trace_handler.TraceHandler onto the active handler stack, delegates to quivers.effects.interpreter.run_program, and returns the accumulated Trace (with output and log_joint filled in). Any other handlers already on the stack compose with trace in the standard outer-to-inner order.

PARAMETER	DESCRIPTION
program	The program to trace. TYPE: MonadicProgram
x	Program input. Shape (batch, ...). TYPE: Tensor
observations	Values to clamp observed variables to. Keys are variable names, values are tensors of the appropriate shape. TYPE: dict[str, Tensor] or None DEFAULT: None

RETURNS	DESCRIPTION
Trace	Complete execution trace with all sites, output, and log-joint.

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

def trace(
    program: MonadicProgram,
    x: torch.Tensor,
    observations: dict[str, torch.Tensor] | None = None,
) -> Trace:
    """Execute a program and record all sample sites.

    A thin wrapper around the handler-aware interpreter: pushes a
    `quivers.effects.trace_handler.TraceHandler` onto the active
    handler stack, delegates to
    `quivers.effects.interpreter.run_program`, and returns the
    accumulated `Trace` (with ``output`` and ``log_joint`` filled
    in). Any other handlers already on the stack compose with
    trace in the standard outer-to-inner order.

    Parameters
    ----------
    program : MonadicProgram
        The program to trace.
    x : torch.Tensor
        Program input. Shape ``(batch, ...)``.
    observations : dict[str, torch.Tensor] or None
        Values to clamp observed variables to. Keys are variable
        names, values are tensors of the appropriate shape.

    Returns
    -------
    Trace
        Complete execution trace with all sites, output, and
        log-joint.
    """
    with TraceHandler() as handler:
        output = run_program(program, x, observations)
    handler.trace.output = output
    handler.trace.log_joint = handler.total_log_joint(x.shape[0], x.device)
    return handler.trace

Package-level exports

quivers.inference also re-exports the variational guides of quivers.inference.guides and the kernels and runner of quivers.inference.mcmc, so that from quivers.inference import AutoNormalGuide, NUTSKernel works. The entries below record those import paths with their signatures; the Guides and MCMC pages document each class in full.

AutoNormal module-attribute

AutoNormal = AutoNormalGuide

AutoMultivariateNormal module-attribute

AutoMultivariateNormal = AutoMultivariateNormalGuide

AutoLowRankMVN module-attribute

AutoLowRankMVN = AutoLowRankMultivariateNormalGuide

AutoDelta module-attribute

AutoDelta = AutoDeltaGuide

AutoLaplace module-attribute

AutoLaplace = AutoLaplaceApproximation

AutoIAFNormal module-attribute

AutoIAFNormal = AutoIAFGuide

Guide

AutoNormalGuide

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

AutoMultivariateNormalGuide

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

AutoLowRankMultivariateNormalGuide

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

AutoDeltaGuide

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

AutoLaplaceApproximation

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

AutoNormalizingFlow

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

AutoIAFGuide

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

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)

AutoMixtureGuide

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

AutoGuideList

AutoGuideList(parts: dict[str, Guide])

AutoStructured

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

MCMCKernel

HMCKernel

HMCKernel(step_size: float = 0.1, num_steps: int = 10, mass_matrix: MassMatrixKind = 'identity', target_accept: float = 0.65, divergence_threshold: float = 1000.0, adapt_step_size: bool = True, adapt_mass_matrix: bool = True)

NUTSKernel

NUTSKernel(step_size: float = 0.1, max_tree_depth: int = 10, mass_matrix: MassMatrixKind = 'diagonal', target_accept: float = 0.8, divergence_threshold: float = 1000.0, adapt_step_size: bool = True, adapt_mass_matrix: bool = True)

MCMC

MCMC(kernel: MCMCKernel, num_warmup: int, num_samples: int, num_chains: int = 4, init_strategy: InitStrategy = 'prior', *, seed: int | None = None, generator: Generator | None = None)

MCMCResult dataclass

MCMCResult(samples: dict[str, Tensor], log_densities: Tensor, acceptance_rates: Tensor, divergence_counts: Tensor, r_hat: dict[str, Tensor], ess: dict[str, Tensor], num_warmup: int, num_samples: int)