Interpreter

Running a program on the reference machine under a handler stack.

interpreter

Running a MonadicProgram on the reference machine under a handler stack.

run_program encodes the program as a kernel computation, installs the handlers on the thread-local stack as lexical handlers of the program's random and score instances, serves the program's host steps through Compute handlers over torch, and evaluates the whole on the reference machine. The value it returns is the program's output; handlers accumulate state on the side. A caller that wants a Trace stacks a TraceHandler before run_program and reads its trace attribute afterwards; the thin quivers.inference.trace.trace wrapper does exactly this.

The nesting of the installed handlers, outermost first, is fixed by the kernel's rules rather than by the order of the stack alone: the run's score accumulator, then every stack handler of the score instance in stack order, then the default scoring draw and the run's observations, then every stack handler of the random instance in stack order, then the parameter store with every stack handler of the param instance inside it, then the host steps. A density is a Score.add the answering handler emits outward, so a score transformer anywhere on the stack reaches every site, and an inner random handler sees the answer an outer one gives.

host_value

host_value(value: object) -> bool

Whether a value is one a host step may produce or consume.

PARAMETER	DESCRIPTION
value	The value. TYPE: object

RETURNS	DESCRIPTION
bool	True for tensors, numbers, strings, sampleables, None, and tuples of such values. The kernel checks the terms statically; the host convention is that every tensor carries the batch axis in front of the shape its type names.

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

def host_value(value: object) -> bool:
    """Whether a value is one a host step may produce or consume.

    Parameters
    ----------
    value : object
        The value.

    Returns
    -------
    bool
        True for tensors, numbers, strings, sampleables, ``None``, and
        tuples of such values. The kernel checks the terms statically;
        the host convention is that every tensor carries the batch axis
        in front of the shape its type names.
    """
    if isinstance(value, tuple):
        return all(host_value(item) for item in value)
    return value is None or isinstance(
        value, torch.Tensor | int | float | bool | str | TorchSampleable
    )

run_program

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

Execute a program on the reference machine under the handler stack.

PARAMETER	DESCRIPTION
program	Program to execute. TYPE: MonadicProgram
x	Program input. Shape (batch, ...). TYPE: Tensor
observations	Values to condition observed sites on, keyed by site name. Undeclared keys are host data the program's steps read by name. TYPE: Mapping[str, Tensor] or None DEFAULT: None

RETURNS	DESCRIPTION
Tensor or dict[str, Tensor]	The program's return value: the single returned variable, or the returned variables keyed by their labels.

RAISES	DESCRIPTION
ValueError	If an observation names one variable of a destructuring draw.

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

def run_program(
    program: MonadicProgram,
    x: torch.Tensor,
    observations: Mapping[str, torch.Tensor] | None = None,
) -> torch.Tensor | dict[str, torch.Tensor]:
    """Execute a program on the reference machine under the handler stack.

    Parameters
    ----------
    program : MonadicProgram
        Program to execute.
    x : torch.Tensor
        Program input. Shape ``(batch, ...)``.
    observations : Mapping[str, torch.Tensor] or None
        Values to condition observed sites on, keyed by site name.
        Undeclared keys are host data the program's steps read by name.

    Returns
    -------
    torch.Tensor or dict[str, torch.Tensor]
        The program's return value: the single returned variable, or the
        returned variables keyed by their labels.

    Raises
    ------
    ValueError
        If an observation names one variable of a destructuring draw.
    """
    sites, supplied, data = _split_observations(program, observations or {})
    kernel = _cached_kernel(program, tuple(sorted(data)))
    run = RunContext(kernel, host_value, observations=dict(sites), supplied=supplied)
    stack = list(_handler_stack())
    steps = _HostSteps(run, x)
    result_type = kernel.result_type

    def record(request: RuntimeRequest, weight: object) -> None:
        """Record a contribution reaching the accumulator.

        Parameters
        ----------
        request : RuntimeRequest
            The score request.
        weight : object
            The contribution.
        """
        assert isinstance(weight, torch.Tensor | float | int)
        run.contributions.append(
            Contribution(
                run.key_of(request),
                request.core.origin.origin.role,
                torch.as_tensor(weight),
                request.core.origin.origin.structural_path,
            )
        )

    accumulator, _ = score_handler(
        answer_type=result_type, observer=record, key="run-accumulator"
    )
    score_installations = [Installation(kernel.score, accumulator)]
    random_installations: list[Installation] = []
    param_installations: list[Installation] = []
    user_random: list[Installation] = []
    step_observers: list[Installation] = []
    step_instances = {item.entry.instance for item in kernel.module.instances}
    finishers: list[EffectHandler] = []
    for handler in stack:
        for installation in handler.install(run):
            if installation.instance is kernel.score:
                score_installations.append(installation)
            elif installation.instance is kernel.random:
                user_random.append(installation)
            elif installation.instance is kernel.param:
                param_installations.append(installation)
            elif installation.instance.entry.instance in step_instances:
                step_observers.append(installation)
            else:
                raise ValueError(
                    f"{type(handler).__name__} installs on an instance the run "
                    "does not serve"
                )
        finishers.append(handler)
    random_installations.append(
        Installation(
            kernel.random,
            draw_scoring_handler(
                score_instance=kernel.score.entry.instance,
                result_validator=host_value,
                answer_type=result_type,
                key="run-draw",
            ),
        )
    )
    if sites:
        clamped = {
            label: _canonical(
                steps.program.step_module(
                    next(
                        spec
                        for spec in program.steps
                        if isinstance(spec, Draw | Observe)
                        and ",".join(spec.names) == label
                    )
                ),  # type: ignore[arg-type]
                value,
            )
            for label, value in sites.items()
        }
        random_installations.append(
            Installation(
                kernel.random,
                condition_handler(
                    clamped,
                    score_instance=kernel.score.entry.instance,
                    result_validator=host_value,
                    missing=MissingValuePolicy.FORWARD,
                    extra=ExtraValuePolicy.IGNORE,
                    answer_type=result_type,
                    key="run-observations",
                ),
            )
        )
    random_installations.extend(user_random)
    store: dict[str, torch.Tensor] = {}
    for step in kernel.steps:
        if step.kind != "draw" or step.site is None:
            continue
        assert isinstance(step.spec, Draw | Observe)
        morphism = program.step_module(step.spec)
        assert isinstance(morphism, ContinuousMorphism)
        for name, parameter in morphism.named_parameters():
            store[f"{step.site}.{name}"] = parameter
    run.parameters = store
    param_installations.insert(
        0,
        Installation(
            kernel.param,
            param_handler(
                lambda name, _request: store[name],
                result_validator=host_value,
                answer_type=result_type,
                key="run-parameters",
            ),
        ),
    )
    step_installations: list[Installation] = []
    for step in kernel.steps:
        instance = next(
            item
            for item in kernel.module.instances
            if item.entry.instance == step.instance
        )
        introduced = (
            EffectRow((kernel.param.entry,)) if step.kind == "draw" else EffectRow()
        )
        step_installations.append(
            Installation(
                instance,
                compute_handler(
                    steps.function(step),
                    effect=step.effect,
                    result_validator=host_value,
                    introduced=introduced,
                    answer_type=result_type,
                    key=f"run-step-{instance.name}",
                ),
            )
        )
    entry = kernel.module.computations[0]
    arguments = [x, *(data[name] for name in kernel.data)]
    parameters = [Local(local.name, local.type) for local in entry.parameters]
    call = Call(
        entry.id,
        entry.name,
        (),
        tuple(Var(parameter) for parameter in parameters),
        result_type,
        entry.type.effects,
        SourceOrigin(kernel.module.module, ("run", "call"), "call", SOURCE_PROTOCOL),
    )
    # A handler is installed only over an effect the entry performs:
    # a program with no draw performs neither Random nor Param, and
    # one with no draw and no score step performs no Score either.
    has_draws = any(step.kind == "draw" for step in kernel.steps)
    performs_score = entry.type.effects.lookup(kernel.score.entry.instance) is not None
    if not has_draws:
        random_installations = []
        param_installations = []
        if not performs_score:
            score_installations = []
    installations = [
        installation.named(f"{installation.runtime.definition.name}#{position}")
        for position, installation in enumerate(
            [
                *score_installations,
                *random_installations,
                *param_installations,
                *step_installations,
                *step_observers,
            ]
        )
    ]
    body = _handle_chain(installations, call)
    wrapper = NamedComputation(
        ComputationId.derive(kernel.module.module, "computation", RUN),
        RUN,
        (),
        tuple(parameters),
        body,
        ComputationType(EffectRow(), result_type),
        SourceOrigin(kernel.module.module, ("run",), "computation", SOURCE_PROTOCOL),
    )
    module = QiecModule(
        kernel.module.module,
        kernel.module.source_protocol,
        effects=kernel.module.effects,
        instances=kernel.module.instances,
        handlers=tuple(
            installation.runtime.definition for installation in installations
        ),
        computations=(entry, wrapper),
    )
    registry = _checked_registry(program, kernel, module)
    attachments = RuntimeAttachments()
    for installation in installations:
        attachments.bind_handler(installation.runtime)
    evaluator = Evaluator(
        attachments, module=module, type_validator=lambda _type: host_value
    )
    # The run owns its handlers: a program a host step runs in turn (a
    # score closing over another program's joint, a sub-program drawn as
    # a step) is a separate invocation, whose sites this run's handlers
    # neither observe nor condition. The stack is cleared for the
    # evaluation and restored after it.
    active = _handler_stack()
    entered = active[:]
    active.clear()
    try:
        output = evaluator.evaluate_validated(
            body, registry, dict(zip(parameters, arguments, strict=True))
        )
    finally:
        active.extend(entered)
    for handler in finishers:
        handler.finish(run, output)
    if program._return_is_single:
        assert isinstance(output, torch.Tensor)
        return output
    assert isinstance(output, tuple)
    keys = program._return_labels if program._return_labels else program._return_vars
    return {key: value for key, value in zip(keys, output, strict=True)}

sample_program

sample_program(program: MonadicProgram, x: Tensor, sample_shape: Size = Size(), observations: Mapping[str, Tensor] | None = None) -> Tensor | dict[str, Tensor]

Run a program forward under the handler stack, with sample dimensions.

PARAMETER	DESCRIPTION
program	The program. TYPE: MonadicProgram
x	Program input. Shape (batch, ...). TYPE: Tensor
sample_shape	Leading sample dimensions. Each element of the shape is one independent run of the program; the returned values stack the runs ahead of their batch axis. TYPE: Size DEFAULT: Size()
observations	Values to condition observed sites on, keyed by site name, and host data the steps read by name. TYPE: Mapping[str, Tensor] or None DEFAULT: None

RETURNS	DESCRIPTION
Tensor or dict[str, Tensor]	The program's return value, with sample_shape in front.

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

def sample_program(
    program: MonadicProgram,
    x: torch.Tensor,
    sample_shape: torch.Size = torch.Size(),
    observations: Mapping[str, torch.Tensor] | None = None,
) -> torch.Tensor | dict[str, torch.Tensor]:
    """Run a program forward under the handler stack, with sample dimensions.

    Parameters
    ----------
    program : MonadicProgram
        The program.
    x : torch.Tensor
        Program input. Shape ``(batch, ...)``.
    sample_shape : torch.Size
        Leading sample dimensions. Each element of the shape is one
        independent run of the program; the returned values stack the
        runs ahead of their batch axis.
    observations : Mapping[str, torch.Tensor] or None
        Values to condition observed sites on, keyed by site name, and
        host data the steps read by name.

    Returns
    -------
    torch.Tensor or dict[str, torch.Tensor]
        The program's return value, with ``sample_shape`` in front.
    """
    if not sample_shape:
        return run_program(program, x, observations)
    count = 1
    for extent in sample_shape:
        count *= int(extent)
    runs = [run_program(program, x, observations) for _ in range(count)]
    first = runs[0]
    if isinstance(first, torch.Tensor):
        stacked = torch.stack([run for run in runs if isinstance(run, torch.Tensor)])
        return stacked.reshape(*sample_shape, *first.shape)
    assert isinstance(first, dict)
    return {
        key: torch.stack([cast_dict(run)[key] for run in runs]).reshape(
            *sample_shape, *first[key].shape
        )
        for key in first
    }

cast_dict

cast_dict(value: Tensor | dict[str, Tensor]) -> dict[str, Tensor]

Read a program's dictionary return.

PARAMETER	DESCRIPTION
value	A program's return value. TYPE: Tensor | dict[str, Tensor]

RETURNS	DESCRIPTION
dict[str, Tensor]	The value, which is a dictionary.

RAISES	DESCRIPTION
TypeError	If the value is a single tensor.

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

def cast_dict(value: torch.Tensor | dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
    """Read a program's dictionary return.

    Parameters
    ----------
    value : torch.Tensor | dict[str, torch.Tensor]
        A program's return value.

    Returns
    -------
    dict[str, torch.Tensor]
        The value, which is a dictionary.

    Raises
    ------
    TypeError
        If the value is a single tensor.
    """
    if isinstance(value, dict):
        return value
    raise TypeError("the program returns a single tensor, not named values")

log_joint

log_joint(program: MonadicProgram, x: Tensor, intermediates: Mapping[str, Tensor]) -> Tensor

The joint log density of a program at values for all of its draws.

The program runs on the reference machine under the handler stack with every draw conditioned on its given value, and the joint is the sum of the site densities the run accumulated, one term per site, over the input's batch axis.

PARAMETER	DESCRIPTION
program	The program. TYPE: MonadicProgram
x	Program input. Shape (batch, ...). TYPE: Tensor
intermediates	A value for every draw the program makes, keyed by variable name or by return label, together with any host data its steps read by name. TYPE: Mapping[str, Tensor]

RETURNS	DESCRIPTION
Tensor	The joint log density, shaped (batch,).

RAISES	DESCRIPTION
KeyError	If a draw the program makes is given no value.

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

def log_joint(
    program: MonadicProgram,
    x: torch.Tensor,
    intermediates: Mapping[str, torch.Tensor],
) -> torch.Tensor:
    """The joint log density of a program at values for all of its draws.

    The program runs on the reference machine under the handler stack
    with every draw conditioned on its given value, and the joint is
    the sum of the site densities the run accumulated, one term per
    site, over the input's batch axis.

    Parameters
    ----------
    program : MonadicProgram
        The program.
    x : torch.Tensor
        Program input. Shape ``(batch, ...)``.
    intermediates : Mapping[str, torch.Tensor]
        A value for every draw the program makes, keyed by variable
        name or by return label, together with any host data its steps
        read by name.

    Returns
    -------
    torch.Tensor
        The joint log density, shaped ``(batch,)``.

    Raises
    ------
    KeyError
        If a draw the program makes is given no value.
    """
    values = dict(intermediates)
    if program._return_labels:
        for label, name in zip(
            program._return_labels, program._return_vars, strict=True
        ):
            if label in values and name not in values:
                values[name] = values[label]
    with TraceHandler() as handler:
        run_program(program, x, values)
    for name, site in handler.trace.sites.items():
        if (
            site.sampleable is not None
            and not site.is_observed
            and not site.is_deterministic
        ):
            raise KeyError(name)
    # The accumulator takes the widest dtype among the densities: a
    # zero-dimensional float64 density added to a float32 batch of zeros
    # would otherwise be demoted to float32.
    dtype = torch.get_default_dtype()
    for site in handler.trace.sites.values():
        if site.log_prob.is_floating_point():
            dtype = torch.promote_types(dtype, site.log_prob.dtype)
    total = torch.zeros(x.shape[0], device=x.device, dtype=dtype)
    for site in handler.trace.sites.values():
        total = total + site.log_prob
    return total