Traces

The trace handler and the trace records.

trace_handler

Trace handler: records every site the program visits.

TraceHandler installs an observing handler of the program's random instance and of its let and score steps at its position on the stack, so it records each site's value as the handlers outside it answered it and sees no site a handler inside it hid. A site's density is read after the run from the contributions that reached the run's accumulator under the site's provenance, after every score transformer on the stack. The thin quivers.inference.trace.trace wrapper stacks this handler on and returns its accumulated Trace.

TraceHandler

TraceHandler()

Bases: EffectHandler

Record every site visited during a program's execution.

Produces a Trace whose sites dict is keyed by variable name, whose output is the program's return value, and whose log_joint is the sum of every site's log_prob. A TraceHandler is single-use: run the program under one instance, read trace, then discard. The handler takes no configuration.

Source code in src/quivers/effects/trace_handler.py

def __init__(self) -> None:
    self.trace: Trace = Trace()
    self._recorder = TraceRecorder()
    self._steps: dict[object, HostStep] = {}

install

install(run: RunContext) -> tuple[Installation, ...]

Install observers of the sites and of the let and score steps.

PARAMETER	DESCRIPTION
run	The run being prepared. TYPE: RunContext

RETURNS	DESCRIPTION
tuple[Installation, ...]	One forwarding observer of the random instance and one of each let and score step's Compute instance.

Source code in src/quivers/effects/trace_handler.py

def install(self, run: RunContext) -> tuple[Installation, ...]:
    """Install observers of the sites and of the let and score steps.

    Parameters
    ----------
    run : RunContext
        The run being prepared.

    Returns
    -------
    tuple[Installation, ...]
        One forwarding observer of the ``random`` instance and one of
        each let and score step's ``Compute`` instance.
    """
    kernel = run.kernel
    installations = [
        Installation(
            kernel.random,
            trace_handler(
                RANDOM,
                (RANDOM_SAMPLE,),
                self._recorder,
                validators={RANDOM_SAMPLE: run.validator},
                answer_type=run.result_type,
                key=f"trace-random-{id(self):x}",
            ),
        )
    ]
    for step in kernel.steps:
        if step.kind not in ("let", "score"):
            continue
        instance = next(
            item
            for item in kernel.module.instances
            if item.entry.instance == step.instance
        )
        self._steps[step.instance] = step
        installations.append(
            Installation(
                instance,
                trace_handler(
                    step.effect,
                    (COMPUTE_APPLY,),
                    self._recorder,
                    validators={COMPUTE_APPLY: run.validator},
                    answer_type=run.result_type,
                    key=f"trace-step-{instance.name}-{id(self):x}",
                ),
            )
        )
    return tuple(installations)

finish

finish(run: RunContext, output: object) -> None

Build the trace from the recorded events and the run's contributions.

PARAMETER	DESCRIPTION
run	The run, with every contribution the accumulator received. TYPE: RunContext
output	The program's output, which the trace does not read; the wrapper that stacks this handler sets trace.output from the value it returns. TYPE: object

Source code in src/quivers/effects/trace_handler.py

def finish(self, run: RunContext, output: object) -> None:
    """Build the trace from the recorded events and the run's contributions.

    Parameters
    ----------
    run : RunContext
        The run, with every contribution the accumulator received.
    output : object
        The program's output, which the trace does not read; the
        wrapper that stacks this handler sets ``trace.output`` from
        the value it returns.
    """
    del output
    for event in self._recorder.events:
        if event.operation == RANDOM_SAMPLE:
            self._record_site(run, event)
        else:
            self._record_step(run, event)

total_log_joint

total_log_joint(batch_size: int, device: device) -> Tensor

Sum every site's log_prob into the joint density.

The accumulator uses the narrowest right-aligned shape shared by the nonzero site contributions. Wider axes are summed before addition so broadcasting cannot duplicate scalar terms. Zero contributions do not determine the target shape but are still added.

PARAMETER	DESCRIPTION
batch_size	Leading extent of the program input. Not consulted: the joint's shape follows from the sites themselves, and a program whose sites reduce their own plate axes has a joint narrower than its input. TYPE: int
device	Device the accumulator seeds on. TYPE: device

RETURNS	DESCRIPTION
Tensor	The joint log-density.

Source code in src/quivers/effects/trace_handler.py

def total_log_joint(self, batch_size: int, device: torch.device) -> torch.Tensor:
    """Sum every site's ``log_prob`` into the joint density.

    The accumulator uses the narrowest right-aligned shape shared
    by the nonzero site contributions. Wider axes are summed before
    addition so broadcasting cannot duplicate scalar terms. Zero
    contributions do not determine the target shape but are still
    added.

    Parameters
    ----------
    batch_size : int
        Leading extent of the program input. Not consulted: the
        joint's shape follows from the sites themselves, and a
        program whose sites reduce their own plate axes has a
        joint narrower than its input.
    device : torch.device
        Device the accumulator seeds on.

    Returns
    -------
    torch.Tensor
        The joint log-density.
    """
    del batch_size
    contributions = [site.log_prob for site in self.trace.sites.values()]
    target = _narrowest_shape(
        [lp for lp in contributions if not bool(torch.all(lp == 0))]
    )
    total = torch.zeros((), device=device)
    for log_prob in contributions:
        total = total + _reduce_to(log_prob, target)
    return total