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)

Record of a single sample 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). TYPE: ContinuousMorphism or None
value	The sampled or observed value. TYPE: Tensor
log_prob	Log-density of the value under the morphism. Shape (batch,). Zero for let bindings. TYPE: Tensor
is_observed	Whether this site was clamped to an observed value. TYPE: bool DEFAULT: False
is_deterministic	Whether this is a deterministic let binding. TYPE: bool DEFAULT: False

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 sample 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 stochastic sites. Shape (batch,). TYPE: Tensor DEFAULT: None

stochastic_sites property

stochastic_sites: dict[str, SampleSite]

Return only stochastic (non-deterministic) sites.

latent_sites property

latent_sites: dict[str, SampleSite]

Return only latent (non-observed, non-deterministic) sites.

observed_sites property

observed_sites: dict[str, SampleSite]

Return only observed sites.

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