Composite lets

Normalizing program ASTs before lowering.

Internal module

This module is an implementation detail of the compiler and its tooling. Its __all__ is empty, so none of the names below are part of the public API, and they may change without notice.

composite_lets

Normalize program ASTs before IR lowering.

The pass expands sampled composite bindings into atomic sample and let steps. It preserves marginalize blocks for the lowering stage while expanding composite references inside their scopes.

expand_composite_lets

expand_composite_lets(module: Module, *, target: str | None = None) -> Module

Rewrite program_decl bodies so composite-let sample steps become equivalent chains of atomic sample steps.

A composite let is a DefineDecl whose .expr is an ExprCompose (the prior >> likelihood form). Each SampleStep / ObserveStep whose morphism slot names such a let is rewritten into a sequence of fresh SampleSteps, one per element of the composition chain, with the trailing step keeping the original step's bound variable name.

The returned Module shares vertex identity with the input for every non-rewritten statement; only the program_decls with composite-let references are rebuilt.

The rewrite is target-independent: every backend consumes the same expanded AST, and the choice of whether to enumerate or to sample a marginalize latent is each renderer's, taken against the lowered IRMarginalize. target is accepted so callers can name the backend they are compiling for without the pass branching on it.

Source code in src/quivers/dsl/composite_lets.py

def expand_composite_lets(
    module: Module,
    *,
    target: str | None = None,
) -> Module:
    """Rewrite `program_decl` bodies so composite-let sample steps
    become equivalent chains of atomic sample steps.

    A composite let is a `DefineDecl` whose `.expr` is an `ExprCompose`
    (the `prior >> likelihood` form). Each `SampleStep` /
    `ObserveStep` whose `morphism` slot names such a let is rewritten
    into a sequence of fresh `SampleStep`s, one per element of the
    composition chain, with the trailing step keeping the original
    step's bound variable name.

    The returned `Module` shares vertex identity with the input for
    every non-rewritten statement; only the `program_decl`s with
    composite-let references are rebuilt.

    The rewrite is target-independent: every backend consumes the same
    expanded AST, and the choice of whether to enumerate or to sample a
    marginalize latent is each renderer's, taken against the lowered
    [`IRMarginalize`][quivers.transpile.ir.IRMarginalize]. `target` is
    accepted so callers can name the backend they are compiling for
    without the pass branching on it.
    """
    morphism_table: dict[str, MorphismDecl] = {
        name: s
        for s in module.statements
        if isinstance(s, MorphismDecl)
        for name in s.names
    }
    let_table: dict[str, Expr] = {
        s.name: s.expr for s in module.statements if isinstance(s, DefineDecl)
    }
    program_table: dict[str, ProgramDecl] = {
        s.name: s for s in module.statements if isinstance(s, ProgramDecl)
    }

    del target
    declared = set(morphism_table)
    declared_programs = set(program_table)
    new_statements: list = []
    for stmt in module.statements:
        if isinstance(stmt, ProgramDecl):
            program = _Program(
                morphism_table,
                let_table,
                program_table,
                _domain_names(stmt),
                _domain_factors(stmt.domain),
            )
            new_draws = _expand_draws(stmt.draws, program)
            # A scan's step program is declared before the program that
            # scans with it, so a consumer reading in order sees it.
            for name, declaration in program.programs.items():
                if name not in declared_programs:
                    declared_programs.add(name)
                    new_statements.append(declaration)
            if new_draws is stmt.draws:
                new_statements.append(stmt)
            else:
                new_statements.append(stmt.with_(draws=tuple(new_draws)))
        else:
            new_statements.append(stmt)
    # A replica of a `[replicate=k]` morphism and a copy a `stack`
    # makes are morphisms of their own, with their own parameters;
    # the expanded module declares them so every consumer sees them.
    synthesized = [
        decl for name, decl in morphism_table.items() if name not in declared
    ]
    if synthesized:
        last = max(
            index
            for index, stmt in enumerate(new_statements)
            if isinstance(stmt, MorphismDecl)
        )
        new_statements[last + 1 : last + 1] = synthesized
    return module.with_(statements=tuple(new_statements))