Overview

Effect handlers for probabilistic programs, executed as lexical kernel handlers.

effects

Algebraic effect handlers for probabilistic programs.

Effect handlers in the Pyro poutine / NumPyro handlers shape, grounded in the algebraic-effects calculus of Plotkin and Pretnar 2009 and its probabilistic application in Scibior et al. 2018 and Nguyen et al. 2023, and executed as lexical handlers of the Quivers Indexed Effect Core: a MonadicProgram runs on the reference machine as the kernel computation quivers.effects.program_module encodes it to, and every handler here is a handler of that computation's random, score, or param instance.

The EffectHandler ABC and the thread-local handler stack live in quivers.effects.base. Concrete handlers (TraceHandler, ClampHandler, DoHandler, MaskHandler, ScaleHandler, BlockHandler, ReplayHandler, LiftHandler, CollapseHandler) live in per-file modules and are re-exported here along with their short-name factories. The reparam subpackage collects the reparameterisation strategies.

The clamp handler is the effect-stack analogue of Pyro's condition; the name avoids collision with the top-level quivers.inference.conditioning.condition factory, which returns a Conditioned model wrapper. Both compose observations onto a model; pick clamp when writing a handler stack, condition when writing a top-level Conditioned object.

The handler-aware interpreter run_program lives in quivers.effects.interpreter; the thin quivers.inference.trace.trace wrapper stacks a TraceHandler on and returns the recorded trace.

CONJUGATE_SOLVERS module-attribute

CONJUGATE_SOLVERS: Mapping[tuple[str, str], Solver] = {('Normal', 'Normal'): normal_normal, ('Beta', 'Bernoulli'): beta_bernoulli, ('Gamma', 'Poisson'): gamma_poisson, ('Dirichlet', 'Categorical'): dirichlet_categorical}

The closed-form solvers, keyed by (prior family, likelihood family).

SiteKey

SiteKey = tuple[str, tuple[tuple[str, str | int], ...], tuple[int, ...]]

The dynamic address of one occurrence of a site: its static identity, its dynamic address frames, and its resumption path.

Solver

Solver = Callable[[Mapping[str, Tensor], Tensor, TorchSampleable, Tensor], Tensor]

A closed-form marginal: given the parent's parameters, the parent's drawn value, the child's sampleable at that value, and the child's value, the log marginal density of the child's value.

SiteAddress

SiteAddress = tuple[str, tuple[tuple[str, str | int], ...], tuple[int, ...]]

The kernel's dynamic address of one occurrence of a site.

Contribution dataclass

Contribution(key: SiteKey, role: str, weight: Tensor, path: tuple[str | int, ...] = ())

One score contribution that reached the run's accumulator.

PARAMETER	DESCRIPTION
key	The address of the site the contribution was derived from. TYPE: SiteKey
role	What the contribution was for, from the emitting handler's provenance: a scored draw, a conditioned observation, a replay. TYPE: str
weight	The contribution as it arrived, after every transformer. TYPE: Tensor
path	The structural path of the request's origin; a score step of the program has the path of its step, a derived contribution the path it was generated at. TYPE: tuple[str | int, ...] DEFAULT: ()

RunContext dataclass

RunContext(kernel: ProgramKernel, validator: RuntimeValidator, contributions: list[Contribution] = list(), interventions: set[SiteKey] = set(), annotations: dict[str, dict[str, object]] = dict(), parameters: dict[str, Tensor] = dict(), observations: dict[str, Tensor] = dict(), supplied: dict[str, Tensor] = dict())

What one run shares with the handlers installed on it.

PARAMETER	DESCRIPTION
kernel	The program's kernel encoding. TYPE: ProgramKernel
validator	The host value validator every handler's clauses check with. TYPE: RuntimeValidator
contributions	The score contributions that reached the accumulator, in order. TYPE: list[Contribution] DEFAULT: list()
interventions	The sites an intervention fixed, by address. TYPE: set[SiteKey] DEFAULT: set()
annotations	Per-site annotations handlers attach, keyed by site label. TYPE: dict[str, dict[str, object]] DEFAULT: dict()
parameters	The learned parameters of the program's morphisms by qualified name, as the parameter store answers them. TYPE: dict[str, Tensor] DEFAULT: dict()
observations	The run's observations at declared sites, by site label. TYPE: dict[str, Tensor] DEFAULT: dict()
supplied	Values standing in for let bindings, by name, which their steps read rather than compute. TYPE: dict[str, Tensor] DEFAULT: dict()

result_type property

result_type: TypeExpr

The type of the program's result.

RETURNS	DESCRIPTION
TypeExpr	The entry computation's result type, which every transparent handler answers with.

key_of

key_of(request: RuntimeRequest) -> SiteKey

The address of the site a request belongs to.

PARAMETER	DESCRIPTION
request	A sample request, or a score request derived from one. TYPE: RuntimeRequest

RETURNS	DESCRIPTION
SiteKey	The request's own address for a sample request; for a derived request, the address of the site it was derived from.

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

def key_of(self, request: RuntimeRequest) -> SiteKey:
    """The address of the site a request belongs to.

    Parameters
    ----------
    request : RuntimeRequest
        A sample request, or a score request derived from one.

    Returns
    -------
    SiteKey
        The request's own address for a sample request; for a derived
        request, the address of the site it was derived from.
    """
    static, dynamic, path = request.address
    parents = request.core.origin.parents
    if parents:
        # A derived request's declared dynamic path is the address of
        # the site it was derived from.
        static = str(parents[0])
        dynamic = tuple(
            (frame.scope, frame.key) for frame in request.core.origin.dynamic_path
        )
    return (static, dynamic, path)

Installation dataclass

Installation(instance: NamedEffectInstance, runtime: RuntimeHandler)

One lexical handler a stack entry installs on a run.

PARAMETER	DESCRIPTION
instance	The instance the handler handles: the program's random, score, or param instance. TYPE: NamedEffectInstance
runtime	The executable handler, whose definition joins the run's module. TYPE: RuntimeHandler

named

named(name: str) -> Installation

The installation with its handler declared under a name.

A module's handler names and identities are unique, while the prelude factories name every handler of a kind alike and derive its identity from the stack entry that made it; a run renames each installation it makes and derives its identity from the name, so two runs installing the same kinds of handler in the same order build the same module.

PARAMETER	DESCRIPTION
name	The declaration's name, unique within the run. TYPE: str

RETURNS	DESCRIPTION
Installation	The installation with the renamed declaration, the rename carried into any per-installation factory the handler has.

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

def named(self, name: str) -> Installation:
    """The installation with its handler declared under a name.

    A module's handler names and identities are unique, while the
    prelude factories name every handler of a kind alike and derive
    its identity from the stack entry that made it; a run renames
    each installation it makes and derives its identity from the
    name, so two runs installing the same kinds of handler in the
    same order build the same module.

    Parameters
    ----------
    name : str
        The declaration's name, unique within the run.

    Returns
    -------
    Installation
        The installation with the renamed declaration, the rename
        carried into any per-installation factory the handler has.
    """
    definition = replace(
        self.runtime.definition,
        id=HandlerId.derive("run", name),
        name=name,
    )
    runtime = replace(self.runtime, definition=definition)
    factory = self.runtime.context_factory
    if factory is not None:

        def context_factory() -> RuntimeHandler:
            """Build a fresh installation under the renamed declaration.

            Returns
            -------
            RuntimeHandler
                The handler the original factory builds, renamed.
            """
            return replace(factory(), definition=definition)

        runtime.context_factory = context_factory
    return Installation(self.instance, runtime)

EffectHandler

Bases: ABC

A handler of a program's canonical instances, stacked by with.

Subclasses implement :meth:install, which builds the lexical handlers the entry contributes to one run, and may implement :meth:finish, which runs after the program returns with the run's context. The handler activates by being used as a context manager:

with clamp({"z": z_val}):
    samples = predictive.rsample(x)

Nested with blocks stack handlers outer-first; the innermost handler of an instance is installed closest to the program.

__enter__

__enter__() -> Self

Push the handler on the thread's stack.

RETURNS	DESCRIPTION
Self	The handler, so with handler as h binds it.

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

def __enter__(self) -> Self:
    """Push the handler on the thread's stack.

    Returns
    -------
    Self
        The handler, so ``with handler as h`` binds it.
    """
    _handler_stack().append(self)
    return self

__exit__

__exit__(exc_type: type[BaseException] | None, exc: BaseException | None, tb: object) -> None

Pop the handler off the thread's stack.

PARAMETER	DESCRIPTION
exc_type	The exception type leaving the block, if any. TYPE: type[BaseException] | None
exc	The exception, if any. TYPE: BaseException | None
tb	Its traceback, if any. TYPE: object

RAISES	DESCRIPTION
RuntimeError	If the handler is not the top of the stack, which means the blocks were exited out of order.

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

def __exit__(
    self,
    exc_type: type[BaseException] | None,
    exc: BaseException | None,
    tb: object,
) -> None:
    """Pop the handler off the thread's stack.

    Parameters
    ----------
    exc_type : type[BaseException] | None
        The exception type leaving the block, if any.
    exc : BaseException | None
        The exception, if any.
    tb : object
        Its traceback, if any.

    Raises
    ------
    RuntimeError
        If the handler is not the top of the stack, which means the
        blocks were exited out of order.
    """
    del exc_type, exc, tb
    stack = _handler_stack()
    if not stack or stack[-1] is not self:
        raise RuntimeError(
            "EffectHandler.__exit__: handler stack corruption; "
            "handlers must be entered / exited in strict LIFO order."
        )
    stack.pop()

install abstractmethod

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

Build the lexical handlers this entry installs on a run.

PARAMETER	DESCRIPTION
run	The run being prepared. TYPE: RunContext

RETURNS	DESCRIPTION
tuple[Installation, ...]	The handlers, each naming the instance it handles.

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

@abstractmethod
def install(self, run: RunContext) -> tuple[Installation, ...]:
    """Build the lexical handlers this entry installs on a run.

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

    Returns
    -------
    tuple[Installation, ...]
        The handlers, each naming the instance it handles.
    """

finish

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

Observe the completed run.

PARAMETER	DESCRIPTION
run	The run, with every contribution the accumulator received. TYPE: RunContext
output	The program's output. TYPE: object

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

def finish(self, run: RunContext, output: object) -> None:
    """Observe the completed run.

    Parameters
    ----------
    run : RunContext
        The run, with every contribution the accumulator received.
    output : object
        The program's output.
    """
    del run, output

BlockHandler

BlockHandler(hide: list[str] | None = None, expose: list[str] | None = None)

Bases: EffectHandler

Hide sites from outer handlers.

PARAMETER	DESCRIPTION
hide	Site names to hide; every site when both arguments are absent. TYPE: list[str] or None DEFAULT: None
expose	Site names to leave visible, hiding every other. TYPE: list[str] or None DEFAULT: None

RAISES	DESCRIPTION
ValueError	If both hide and expose are given.

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

def __init__(
    self, hide: list[str] | None = None, expose: list[str] | None = None
) -> None:
    if hide is not None and expose is not None:
        raise ValueError(
            "BlockHandler: pass at most one of `hide` and `expose`, not both."
        )
    self.hide = list(hide) if hide is not None else None
    self.expose = list(expose) if expose is not None else None

install

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

Install a blocking handler of the random instance.

PARAMETER	DESCRIPTION
run	The run being prepared. TYPE: RunContext

RETURNS	DESCRIPTION
tuple[Installation, ...]	The blocking handler.

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

def install(self, run: RunContext) -> tuple[Installation, ...]:
    """Install a blocking handler of the ``random`` instance.

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

    Returns
    -------
    tuple[Installation, ...]
        The blocking handler.
    """
    kernel = run.kernel
    return (
        Installation(
            kernel.random,
            block_handler(
                self._should_block,
                score_instance=kernel.score.entry.instance,
                result_validator=run.validator,
                answer_type=run.result_type,
                key=f"block-{id(self):x}",
            ),
        ),
    )

CheckedProgram

CheckedProgram(domain: AnySpace, codomain: AnySpace, module: QiecModule, name: str)

Bases: ContinuousMorphism

A program run as its elaborated computation on the reference machine.

PARAMETER	DESCRIPTION
domain	The program's declared domain. TYPE: AnySpace
codomain	The program's declared codomain. TYPE: AnySpace
module	The checked module holding the program's entry point. TYPE: QiecModule
name	The program's name. TYPE: str

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

def __init__(
    self, domain: AnySpace, codomain: AnySpace, module: QiecModule, name: str
) -> None:
    super().__init__(domain, codomain)
    self._module = module
    self._name = name
    self._entry: ProgramEntry = program_entry(module, name)

module property

module: QiecModule

The checked module the program runs in.

name property

name: str

The program's name.

entry property

entry: ProgramEntry

The program's elaborated entry point.

observed_names

observed_names() -> set[str]

The names the program observes.

RETURNS	DESCRIPTION
set[str]	The labels of its observe sites.

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

def observed_names(self) -> set[str]:
    """The names the program observes.

    Returns
    -------
    set[str]
        The labels of its ``observe`` sites.
    """
    return {site.name for site in self._entry.sites if site.kind == "observe"}

rsample

rsample(x: Tensor, sample_shape: Size = Size(), observations: Mapping[str, Tensor] | None = None) -> Tensor

Run the program forward, once per input row and sample.

PARAMETER	DESCRIPTION
x	Program input. Shape (batch, ...). TYPE: Tensor
sample_shape	Additional leading sample dimensions; each element is one independent run. TYPE: Size DEFAULT: Size()
observations	The program's data by parameter name, and values for the sample sites the runs condition on; every other site is drawn. TYPE: Mapping[str, Tensor] or None DEFAULT: None

RETURNS	DESCRIPTION
Tensor	The returned values, shape (*sample_shape, batch, ...).

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

def rsample(  # type: ignore[override]
    self,
    x: torch.Tensor,
    sample_shape: torch.Size = torch.Size(),
    observations: Mapping[str, torch.Tensor] | None = None,
) -> torch.Tensor:
    """Run the program forward, once per input row and sample.

    Parameters
    ----------
    x : torch.Tensor
        Program input. Shape ``(batch, ...)``.
    sample_shape : torch.Size
        Additional leading sample dimensions; each element is one
        independent run.
    observations : Mapping[str, torch.Tensor] or None
        The program's data by parameter name, and values for the
        sample sites the runs condition on; every other site is
        drawn.

    Returns
    -------
    torch.Tensor
        The returned values, shape ``(*sample_shape, batch, ...)``.
    """
    given = dict(observations or {})
    batch = int(x.shape[0]) if x.dim() >= 1 else 1
    runs = math.prod(sample_shape)
    rows: list[torch.Tensor] = []
    for _ in range(max(runs, 1)):
        values: list[torch.Tensor] = []
        for row in range(batch):
            data, sites = self._split(x, row, given)
            run = sample_program(self._module, self._name, data=data, sites=sites)
            values.append(_tensor(run.value))
        rows.append(torch.stack(values))
    stacked = torch.stack(rows)
    if not sample_shape:
        return stacked[0]
    return stacked.reshape((*sample_shape, *stacked.shape[1:]))

has_conditional_density

has_conditional_density() -> bool

No: a program's density at its output marginalizes its draws.

RETURNS	DESCRIPTION
bool	Always False; log_joint scores the draws themselves.

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

def has_conditional_density(self) -> bool:
    """No: a program's density at its output marginalizes its draws.

    Returns
    -------
    bool
        Always ``False``;
        [`log_joint`][quivers.effects.checked_program.CheckedProgram.log_joint]
        scores the draws themselves.
    """
    return False

log_prob

log_prob(x: Tensor, y: Tensor) -> Tensor

Log-probability is not defined for a program's output alone.

PARAMETER	DESCRIPTION
x	Program input. TYPE: Tensor
y	Program output. TYPE: Tensor

RAISES	DESCRIPTION
NotImplementedError	Always: the density at the output marginalizes every draw.

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

def log_prob(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
    """Log-probability is not defined for a program's output alone.

    Parameters
    ----------
    x : torch.Tensor
        Program input.
    y : torch.Tensor
        Program output.

    Raises
    ------
    NotImplementedError
        Always: the density at the output marginalizes every draw.
    """
    del x, y
    raise NotImplementedError(
        "log_prob is not supported for programs; computing p(y | x) "
        "requires marginalizing over all intermediate draws. Use "
        "log_joint with a value for every site."
    )

log_joint

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

The joint log-density at given values for every site.

PARAMETER	DESCRIPTION
x	Program input. Shape (batch, ...). TYPE: Tensor
intermediates	The program's data by parameter name and a value for every sample site, by label. TYPE: Mapping[str, Tensor]

RETURNS	DESCRIPTION
Tensor	The log joint of each input row's run. Shape (batch,). The machine scores numbers, so the result carries no gradient path to the given tensors.

RAISES	DESCRIPTION
ExecutionFailure	If a site the program reaches is given no value, or a given site is never reached.

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

def log_joint(
    self, x: torch.Tensor, intermediates: Mapping[str, torch.Tensor]
) -> torch.Tensor:
    """The joint log-density at given values for every site.

    Parameters
    ----------
    x : torch.Tensor
        Program input. Shape ``(batch, ...)``.
    intermediates : Mapping[str, torch.Tensor]
        The program's data by parameter name and a value for every
        sample site, by label.

    Returns
    -------
    torch.Tensor
        The log joint of each input row's run. Shape ``(batch,)``.
        The machine scores numbers, so the result carries no
        gradient path to the given tensors.

    Raises
    ------
    ExecutionFailure
        If a site the program reaches is given no value, or a given
        site is never reached.
    """
    batch = int(x.shape[0]) if x.dim() >= 1 else 1
    totals: list[torch.Tensor] = []
    for row in range(batch):
        data, sites = self._split(x, row, intermediates)
        run = run_program(self._module, self._name, data=data, sites=sites)
        totals.append(
            torch.as_tensor(run.log_joint, dtype=torch.get_default_dtype())
        )
    return torch.stack(totals)

ClampHandler

ClampHandler(data: dict[str, Tensor])

Bases: EffectHandler

Condition named sample sites on given values.

PARAMETER	DESCRIPTION
data	Site name to the value the site is clamped to. Other sites pass through untouched. TYPE: dict[str, Tensor]

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

def __init__(self, data: dict[str, torch.Tensor]) -> None:
    self.data = dict(data)

install

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

Install a conditioning handler of the random instance.

PARAMETER	DESCRIPTION
run	The run being prepared. TYPE: RunContext

RETURNS	DESCRIPTION
tuple[Installation, ...]	The conditioning handler, which forwards sites it has no value for and ignores values at sites the program never samples.

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

def install(self, run: RunContext) -> tuple[Installation, ...]:
    """Install a conditioning handler of the ``random`` instance.

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

    Returns
    -------
    tuple[Installation, ...]
        The conditioning handler, which forwards sites it has no
        value for and ignores values at sites the program never
        samples.
    """
    kernel = run.kernel
    return (
        Installation(
            kernel.random,
            condition_handler(
                self.data,
                score_instance=kernel.score.entry.instance,
                result_validator=run.validator,
                missing=MissingValuePolicy.FORWARD,
                extra=ExtraValuePolicy.IGNORE,
                answer_type=run.result_type,
                key=f"clamp-{id(self):x}",
            ),
        ),
    )

CollapseHandler

CollapseHandler(pairs: dict[str, str] | None = None)

Bases: EffectHandler

Integrate conjugate parents out of the joint.

PARAMETER	DESCRIPTION
pairs	Parent site name to child site name. None collapses nothing and only carries the solver registry for introspection. TYPE: dict[str, str] or None DEFAULT: None

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

def __init__(self, pairs: dict[str, str] | None = None) -> None:
    self.pairs = dict(pairs) if pairs is not None else {}
    self.registry = CONJUGATE_SOLVERS

install

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

Install the collapsing handler of the random instance.

PARAMETER	DESCRIPTION
run	The run being prepared, whose observations supply each child's value. TYPE: RunContext

RETURNS	DESCRIPTION
tuple[Installation, ...]	The handler, which draws parents without scoring them and scores children under the pair's marginal.

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

def install(self, run: RunContext) -> tuple[Installation, ...]:
    """Install the collapsing handler of the ``random`` instance.

    Parameters
    ----------
    run : RunContext
        The run being prepared, whose observations supply each
        child's value.

    Returns
    -------
    tuple[Installation, ...]
        The handler, which draws parents without scoring them and
        scores children under the pair's marginal.
    """
    kernel = run.kernel
    children = {child: parent for parent, child in self.pairs.items()}
    parents: dict[str, tuple[TorchSampleable, torch.Tensor]] = {}
    score_instance = kernel.score.entry.instance
    definition = _handler_def(
        "Random.collapse",
        RANDOM,
        ((RANDOM_SAMPLE, ResumptionGrade.LINEAR),),
        key=f"collapse-{id(self):x}",
        input_type=run.result_type,
        output_type=run.result_type,
        introduced=EffectRow((RowEntry(score_instance, SCORE),)),
        total=False,
    )

    def sample(
        request: RuntimeRequest, resume: Resumption, context: ClauseContext
    ) -> object:
        """Answer a parent by an unscored draw and a child by its marginal.

        Parameters
        ----------
        request : RuntimeRequest
            The request being answered and its arguments.
        resume : Resumption
            The continuation.
        context : ClauseContext
            Runtime services available to the clause.

        Returns
        -------
        object
            What the resumed computation produced, or a `Forward`
            for a site the pairs do not name.

        Raises
        ------
        InvalidHandlerError
            If a child is reached before its parent, or the run has
            no observation for it.
        """
        site, sampleable = _site_and_sampleable(request, definition.name)
        assert isinstance(site, str)
        assert isinstance(sampleable, TorchSampleable)
        if site in self.pairs:
            value = sampleable.rsample()
            parents[site] = (sampleable, value)
            run.annotations.setdefault(site, {})["collapse"] = self.pairs[site]
            run.interventions.add(run.key_of(request))
            return resume(value)
        if site in children:
            parent_name = children[site]
            if parent_name not in parents:
                raise InvalidHandlerError(
                    f"collapse reached child {site!r} before its parent "
                    f"{parent_name!r}"
                )
            if site not in run.observations:
                raise InvalidHandlerError(
                    f"collapse needs an observation for child {site!r}"
                )
            parent_sampleable, parent_value = parents[parent_name]
            solver = conjugate_solver(
                family_of(parent_sampleable), family_of(sampleable)
            )
            value = run.observations[site]
            weight = solver(
                parent_sampleable.parameters(), parent_value, sampleable, value
            )
            _emit_score(
                context,
                request,
                score_instance,
                weight,
                run.validator,
                role="condition-score",
            )
            return resume(value)
        return Forward()

    runtime = RuntimeHandler(
        definition, {RANDOM_SAMPLE: RuntimeClause(sample, run.validator)}
    )
    return (Installation(kernel.random, runtime),)

DoHandler

DoHandler(data: dict[str, Tensor])

Bases: EffectHandler

Intervene on named sites.

PARAMETER	DESCRIPTION
data	Site name to the value the site is fixed to. TYPE: dict[str, Tensor]

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

def __init__(self, data: dict[str, torch.Tensor]) -> None:
    self.data = dict(data)

install

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

Install an intervening handler of the random instance.

PARAMETER	DESCRIPTION
run	The run being prepared; each intervened site's address is recorded on it, so a trace marks the site deterministic. TYPE: RunContext

RETURNS	DESCRIPTION
tuple[Installation, ...]	The intervening handler.

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

def install(self, run: RunContext) -> tuple[Installation, ...]:
    """Install an intervening handler of the ``random`` instance.

    Parameters
    ----------
    run : RunContext
        The run being prepared; each intervened site's address is
        recorded on it, so a trace marks the site deterministic.

    Returns
    -------
    tuple[Installation, ...]
        The intervening handler.
    """

    def mark(request: RuntimeRequest, value: object) -> None:
        """Record an intervened site on the run.

        Parameters
        ----------
        request : RuntimeRequest
            The site's request.
        value : object
            The value fixed, unused.
        """
        del value
        run.interventions.add(run.key_of(request))

    return (
        Installation(
            run.kernel.random,
            intervene_handler(
                self.data,
                result_validator=run.validator,
                on_intervene=mark,
                answer_type=run.result_type,
                key=f"do-{id(self):x}",
            ),
        ),
    )

LiftHandler

LiftHandler(prior_scale: float = 1.0)

Bases: EffectHandler

Draw every reachable parameter from a Normal prior.

PARAMETER	DESCRIPTION
prior_scale	Standard deviation of the Normal prior around each parameter's current value. TYPE: float DEFAULT: 1.0

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

def __init__(self, prior_scale: float = 1.0) -> None:
    self.prior_scale = float(prior_scale)
    self.sampled_params: dict[str, torch.Tensor] = {}

install

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

Install a prior-drawing handler of the param instance.

PARAMETER	DESCRIPTION
run	The run being prepared, whose parameter store holds each parameter's current value. TYPE: RunContext

RETURNS	DESCRIPTION
tuple[Installation, ...]	The handler, which answers each parameter's first read with a fresh draw and later reads with the same draw.

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

def install(self, run: RunContext) -> tuple[Installation, ...]:
    """Install a prior-drawing handler of the ``param`` instance.

    Parameters
    ----------
    run : RunContext
        The run being prepared, whose parameter store holds each
        parameter's current value.

    Returns
    -------
    tuple[Installation, ...]
        The handler, which answers each parameter's first read with
        a fresh draw and later reads with the same draw.
    """
    scale = self.prior_scale
    sampled = self.sampled_params

    def draw(name: str, request: RuntimeRequest) -> object:
        """Answer a parameter read with its prior draw.

        Parameters
        ----------
        name : str
            The parameter's qualified name.
        request : RuntimeRequest
            The read, unused.

        Returns
        -------
        object
            The draw made for the name on this run.
        """
        del request
        if name not in sampled:
            current = run.parameters[name]
            sampled[name] = current.detach() + scale * torch.randn_like(current)
        return sampled[name]

    return (
        Installation(
            run.kernel.param,
            param_handler(
                draw,
                result_validator=run.validator,
                total=False,
                answer_type=run.result_type,
                key=f"lift-{id(self):x}",
            ),
        ),
    )

reset

reset() -> None

Clear the draws so the next run makes fresh ones.

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

def reset(self) -> None:
    """Clear the draws so the next run makes fresh ones."""
    self.sampled_params.clear()

MaskHandler

MaskHandler(mask_tensor: Tensor)

Bases: EffectHandler

Gate every density contribution elementwise.

PARAMETER	DESCRIPTION
mask_tensor	The mask, broadcast against each contribution. TYPE: Tensor

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

def __init__(self, mask_tensor: torch.Tensor) -> None:
    self.mask = mask_tensor

install

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

Install a masking transformer of the score instance.

PARAMETER	DESCRIPTION
run	The run being prepared. TYPE: RunContext

RETURNS	DESCRIPTION
tuple[Installation, ...]	The transformer.

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

def install(self, run: RunContext) -> tuple[Installation, ...]:
    """Install a masking transformer of the ``score`` instance.

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

    Returns
    -------
    tuple[Installation, ...]
        The transformer.
    """
    mask = self.mask

    def masked(weight: object) -> object:
        """Mask one contribution.

        Parameters
        ----------
        weight : object
            The contribution.

        Returns
        -------
        object
            The contribution times the mask.
        """
        return torch.as_tensor(weight) * mask.to(torch.as_tensor(weight).dtype)

    return (
        Installation(
            run.kernel.score,
            reweight_handler(
                masked,
                score_instance=run.kernel.score.entry.instance,
                answer_type=run.result_type,
                key=f"mask-{id(self):x}",
            ),
        ),
    )

HostStep dataclass

HostStep(instance: EffectInstanceId, effect: EffectRef, kind: str, spec: Step | None, environment: tuple[str, ...], answer_type: TypeExpr, site: str | None = None)

One host part of the program, as the runtime provider serves it.

PARAMETER	DESCRIPTION
instance	The Compute instance the step performs on. TYPE: EffectInstanceId
effect	The applied Compute[X, A] interface of the instance. TYPE: EffectRef
kind	"input" for the program's parameter split, "draw" for a morphism applied to its input, "apply" for a categorical morphism's tensor contracted against its input, "bind" for a destructuring split, "let" for a deterministic binding, "score" for a scored binding. TYPE: str
spec	The program's step record, or None for the input split. TYPE: Step | None
environment	The names of the locals the step reads, in the order the environment tuple carries them. TYPE: tuple[str, ...]
answer_type	The type the host part returns. TYPE: TypeExpr
site	The site label of a draw step, else None. TYPE: str | None DEFAULT: None

ProgramKernel dataclass

ProgramKernel(module: QiecModule, program: MonadicProgram, steps: tuple[HostStep, ...], sites: tuple[str, ...], site_types: dict[str, TypeExpr], data: tuple[str, ...], random: NamedEffectInstance, score: NamedEffectInstance, param: NamedEffectInstance, result_type: TypeExpr, locals: dict[str, Local] = dict())

A program's computation with the record the runtime needs to serve it.

PARAMETER	DESCRIPTION
module	The kernel module: the interfaces, the instances, and the entry computation, without handlers; a run adds its own. TYPE: QiecModule
program	The program encoded. TYPE: MonadicProgram
steps	The host steps, in body order. TYPE: tuple[HostStep, ...]
sites	The site label of every draw, in order. TYPE: tuple[str, ...]
site_types	The value type of each site, by label. TYPE: dict[str, TypeExpr]
data	The names of the host data inputs after the program's input, in parameter order. TYPE: tuple[str, ...]
random	The canonical Random instance. TYPE: NamedEffectInstance
score	The canonical Score instance. TYPE: NamedEffectInstance
param	The canonical Param instance. TYPE: NamedEffectInstance
result_type	The entry computation's result type. TYPE: TypeExpr

ReplayHandler

ReplayHandler(trace: Trace)

Bases: EffectHandler

Replay the values of a recorded trace.

PARAMETER	DESCRIPTION
trace	The trace whose site values are reinstalled. TYPE: Trace

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

def __init__(self, trace: Trace) -> None:
    self._trace = trace

install

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

Install a scoring replay handler of the random instance.

PARAMETER	DESCRIPTION
run	The run being prepared. TYPE: RunContext

RETURNS	DESCRIPTION
tuple[Installation, ...]	The replay handler, which forwards sites the trace lacks and ignores recorded sites the program never samples.

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

def install(self, run: RunContext) -> tuple[Installation, ...]:
    """Install a scoring replay handler of the ``random`` instance.

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

    Returns
    -------
    tuple[Installation, ...]
        The replay handler, which forwards sites the trace lacks and
        ignores recorded sites the program never samples.
    """
    kernel = run.kernel
    values = {
        name: site.value
        for name, site in self._trace.sites.items()
        if not site.is_deterministic or site.morphism is not None
    }
    return (
        Installation(
            kernel.random,
            replay_handler(
                values,
                result_validator=run.validator,
                policy=ReplayPolicy.CLAMP_AND_SCORE,
                score_instance=kernel.score.entry.instance,
                extra=ExtraValuePolicy.IGNORE,
                answer_type=run.result_type,
                key=f"replay-{id(self):x}",
            ),
        ),
    )

ScaleHandler

ScaleHandler(factor: float)

Bases: EffectHandler

Multiply every density contribution by a scalar.

PARAMETER	DESCRIPTION
factor	The multiplier. TYPE: float

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

def __init__(self, factor: float) -> None:
    self.factor = float(factor)

install

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

Install a scaling transformer of the score instance.

PARAMETER	DESCRIPTION
run	The run being prepared. TYPE: RunContext

RETURNS	DESCRIPTION
tuple[Installation, ...]	The transformer.

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

def install(self, run: RunContext) -> tuple[Installation, ...]:
    """Install a scaling transformer of the ``score`` instance.

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

    Returns
    -------
    tuple[Installation, ...]
        The transformer.
    """
    factor = self.factor

    def scaled(weight: object) -> object:
        """Scale one contribution.

        Parameters
        ----------
        weight : object
            The contribution.

        Returns
        -------
        object
            The contribution times the factor.
        """
        return torch.as_tensor(weight) * factor

    return (
        Installation(
            run.kernel.score,
            reweight_handler(
                scaled,
                score_instance=run.kernel.score.entry.instance,
                answer_type=run.result_type,
                key=f"scale-{id(self):x}",
            ),
        ),
    )

TorchSampleable dataclass

TorchSampleable(morphism: ContinuousMorphism, input: Tensor, declared: ContinuousMorphism | None = None)

A morphism applied to an input, offered to the kernel as a sampleable.

PARAMETER	DESCRIPTION
morphism	The morphism the site draws from, under whatever parameters the run's Param handlers answered. TYPE: ContinuousMorphism
input	The input tensor the morphism is conditioned on, carrying the batch axis. TYPE: Tensor
declared	The morphism as the program declares it, which a trace records and a strategy reads the family of; the drawing morphism itself when no parameter store rebinds it. TYPE: ContinuousMorphism | None DEFAULT: None

__post_init__

__post_init__() -> None

Default the declared morphism to the drawing one.

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

def __post_init__(self) -> None:
    """Default the declared morphism to the drawing one."""
    if self.declared is None:
        object.__setattr__(self, "declared", self.morphism)

rsample

rsample() -> Tensor

Draw a reparameterised sample.

RETURNS	DESCRIPTION
Tensor	One draw per batch row.

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

def rsample(self) -> torch.Tensor:
    """Draw a reparameterised sample.

    Returns
    -------
    torch.Tensor
        One draw per batch row.
    """
    return self.morphism.rsample(self.input)

sample

sample() -> Tensor

Draw a sample.

RETURNS	DESCRIPTION
Tensor	One draw per batch row.

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

def sample(self) -> torch.Tensor:
    """Draw a sample.

    Returns
    -------
    torch.Tensor
        One draw per batch row.
    """
    return self.morphism.rsample(self.input)

log_prob

log_prob(value: Tensor) -> Tensor

Score a value under the morphism at the input.

PARAMETER	DESCRIPTION
value	The value, one per batch row. TYPE: Tensor

RETURNS	DESCRIPTION
Tensor	The log density per batch row.

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

def log_prob(self, value: torch.Tensor) -> torch.Tensor:
    """Score a value under the morphism at the input.

    Parameters
    ----------
    value : torch.Tensor
        The value, one per batch row.

    Returns
    -------
    torch.Tensor
        The log density per batch row.
    """
    return self.morphism.log_prob(self.input, value)

parameters

parameters() -> dict[str, Tensor]

The distribution parameters the morphism computes at the input.

RETURNS	DESCRIPTION
dict[str, Tensor]	The named parameters of the torch distribution the morphism builds, or loc and scale for a morphism that exposes only a location-scale pair.

RAISES	DESCRIPTION
TypeError	If the morphism exposes neither a distribution builder nor a location-scale pair.

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

def parameters(self) -> dict[str, torch.Tensor]:
    """The distribution parameters the morphism computes at the input.

    Returns
    -------
    dict[str, torch.Tensor]
        The named parameters of the torch distribution the morphism
        builds, or ``loc`` and ``scale`` for a morphism that exposes
        only a location-scale pair.

    Raises
    ------
    TypeError
        If the morphism exposes neither a distribution builder nor a
        location-scale pair.
    """
    get_dist = getattr(self.morphism, "_get_dist", None)
    if callable(get_dist):
        distribution = get_dist(self.input)
        return {
            name: getattr(distribution, name)
            for name in distribution.arg_constraints
            if hasattr(distribution, name)
        }
    get_params = getattr(self.morphism, "_get_params", None)
    if callable(get_params):
        loc, scale = get_params(self.input)
        return {"loc": loc, "scale": scale}
    raise TypeError(
        f"{type(self.morphism).__name__} exposes no distribution parameters"
    )

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

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.

ConjugateReparam

ConjugateReparam(parent: str, parent_family: str, child_family: str)

Bases: Reparam

Score a child site under the conjugate marginal of its parent.

Applied to the child site of a conjugate pair, the strategy answers the site with the run's observation for it, or a draw from the child's own sampleable when there is none, and scores it under the closed-form marginal that integrates the parent out, read from the parent's parameters as the orchestrator saw them. The parent's own density is left in the joint; collapse the parent with quivers.effects.collapse when the marginal joint is wanted.

PARAMETER	DESCRIPTION
parent	The parent site's name. TYPE: str
parent_family	The parent's family. TYPE: str
child_family	The child's family. TYPE: str

RAISES	DESCRIPTION
KeyError	If the pair has no registered marginal.

Source code in src/quivers/effects/reparam/conjugate.py

def __init__(self, parent: str, parent_family: str, child_family: str) -> None:
    self.parent = parent
    self.parent_family = parent_family
    self.child_family = child_family
    self.solver = conjugate_solver(parent_family, child_family)

apply

apply(site: SiteRequest) -> tuple[Tensor, Tensor]

Rewrite the child site.

PARAMETER	DESCRIPTION
site	The child site. TYPE: SiteRequest

RETURNS	DESCRIPTION
tuple[Tensor, Tensor]	The value and its marginal log density.

RAISES	DESCRIPTION
KeyError	If the orchestrator has not seen the parent site.
ValueError	If the parent or the child is not of the declared family.

Source code in src/quivers/effects/reparam/conjugate.py

def apply(self, site: SiteRequest) -> tuple[torch.Tensor, torch.Tensor]:
    """Rewrite the child site.

    Parameters
    ----------
    site : SiteRequest
        The child site.

    Returns
    -------
    tuple[torch.Tensor, torch.Tensor]
        The value and its marginal log density.

    Raises
    ------
    KeyError
        If the orchestrator has not seen the parent site.
    ValueError
        If the parent or the child is not of the declared family.
    """
    try:
        parent = site.seen[self.parent]
    except KeyError as error:
        raise KeyError(
            f"ConjugateReparam: parent site {self.parent!r} was not seen "
            f"before child {site.name!r}"
        ) from error
    if family_of(parent) != self.parent_family:
        raise ValueError(
            f"ConjugateReparam: parent {self.parent!r} draws from "
            f"{family_of(parent)!r}, not {self.parent_family!r}"
        )
    if family_of(site.sampleable) != self.child_family:
        raise ValueError(
            f"ConjugateReparam: child {site.name!r} draws from "
            f"{family_of(site.sampleable)!r}, not {self.child_family!r}"
        )
    if site.given is not None:
        value = site.given
    elif site.observation is not None:
        value = site.observation
    else:
        value = site.sampleable.rsample()
    parent_value = site.sampleable.input
    return value, self.solver(
        parent.parameters(), parent_value, site.sampleable, value
    )

LocScaleReparam

LocScaleReparam(centered: float = 0.0)

Bases: Reparam

Non-centred rewrite for location-scale sample sites.

The site's morphism must expose a _get_params(x) method returning (loc, scale) tensors of the same shape, the convention every ConditionalNormal-shaped family in quivers.continuous.families follows.

PARAMETER	DESCRIPTION
centered	Interpolation between fully centred (1.0) and fully non-centred (0.0) parameterisations. Values in between produce a partial reparam, matching the centered parameter of Pyro's LocScaleReparam. TYPE: float DEFAULT: 0.0

RAISES	DESCRIPTION
ValueError	If centered is outside [0, 1].

Source code in src/quivers/effects/reparam/loc_scale.py

def __init__(self, centered: float = 0.0) -> None:
    if not 0.0 <= centered <= 1.0:
        raise ValueError(
            f"LocScaleReparam: centered must be in [0, 1], got {centered}."
        )
    self.centered = float(centered)

apply

apply(site: SiteRequest) -> tuple[Tensor, Tensor]

Rewrite one site.

PARAMETER	DESCRIPTION
site	The site. TYPE: SiteRequest

RETURNS	DESCRIPTION
tuple[Tensor, Tensor]	The value and its log density under Normal(loc, scale), summed over the trailing feature axis.

RAISES	DESCRIPTION
TypeError	If the site's morphism exposes no _get_params.

Source code in src/quivers/effects/reparam/loc_scale.py

def apply(self, site: SiteRequest) -> tuple[torch.Tensor, torch.Tensor]:
    """Rewrite one site.

    Parameters
    ----------
    site : SiteRequest
        The site.

    Returns
    -------
    tuple[torch.Tensor, torch.Tensor]
        The value and its log density under ``Normal(loc, scale)``,
        summed over the trailing feature axis.

    Raises
    ------
    TypeError
        If the site's morphism exposes no ``_get_params``.
    """
    morphism = site.sampleable.morphism
    get_params = getattr(morphism, "_get_params", None)
    if get_params is None:
        raise TypeError(
            f"LocScaleReparam: site '{site.name}' morphism "
            f"{type(morphism).__name__} does not expose `_get_params(x)`; "
            f"LocScaleReparam requires a Normal-family morphism."
        )
    loc, scale = get_params(site.sampleable.input)
    # Partial centring interpolates the effective scale used to
    # push the base sample forward: ``centered=1`` recovers the
    # original sample-then-score path, ``centered=0`` the fully
    # non-centred rewrite.
    effective = scale.pow(self.centered)
    base_scale = scale / effective
    if site.given is None:
        eps = torch.randn_like(loc)
        y_raw = effective * eps
        y = loc + base_scale * y_raw
    else:
        y = site.given
    residual = (y - loc) / scale
    log_p = -0.5 * residual.pow(2) - scale.log() - 0.5 * math.log(2 * math.pi)
    return y, log_p.sum(dim=-1)

NeuTraReparam

NeuTraReparam(autoguide: Guide)

Bases: Reparam

Reparameterise a site through a trained normalizing-flow guide.

PARAMETER	DESCRIPTION
autoguide	A trained AutoIAFGuide, or any Guide whose sample returns per-site values, whose draw at apply time has the site's support. TYPE: Guide

Source code in src/quivers/effects/reparam/neutra.py

def __init__(self, autoguide: Guide) -> None:
    self.autoguide = autoguide

apply

apply(site: SiteRequest) -> tuple[Tensor, Tensor]

Rewrite one site.

PARAMETER	DESCRIPTION
site	The site. TYPE: SiteRequest

RETURNS	DESCRIPTION
tuple[Tensor, Tensor]	The guide's draw for the site, or the site's own draw when the guide does not cover it, and its log density under the site's sampleable; scoring uses the original sampleable throughout, so a partially covered NeuTra still yields a well-defined joint.

RAISES	DESCRIPTION
KeyError	If the guide covers the site but its draw does not carry it.

Source code in src/quivers/effects/reparam/neutra.py

def apply(self, site: SiteRequest) -> tuple[torch.Tensor, torch.Tensor]:
    """Rewrite one site.

    Parameters
    ----------
    site : SiteRequest
        The site.

    Returns
    -------
    tuple[torch.Tensor, torch.Tensor]
        The guide's draw for the site, or the site's own draw when
        the guide does not cover it, and its log density under the
        site's sampleable; scoring uses the original sampleable
        throughout, so a partially covered NeuTra still yields a
        well-defined joint.

    Raises
    ------
    KeyError
        If the guide covers the site but its draw does not carry it.
    """
    registry = getattr(self.autoguide, "registry", None)
    names = set(registry.names()) if registry is not None else set()
    if site.given is not None:
        y = site.given
    elif site.name not in names:
        y = site.sampleable.rsample()
    else:
        samples = self.autoguide.sample(site.sampleable.input)
        if site.name not in samples:
            raise KeyError(
                f"NeuTraReparam: autoguide sample did not contain site "
                f"'{site.name}'."
            )
        y = samples[site.name]
    return y, site.sampleable.log_prob(y)

Reparam

Bases: ABC

Base class for site-level reparameterisation strategies.

A strategy is a policy, not a handler. The ReparamOrchestrator handler dispatches each covered site to the matching strategy and installs the strategy's value and density on the site. Subclasses override :meth:apply.

apply abstractmethod

apply(site: SiteRequest) -> tuple[Tensor, Tensor]

Rewrite one site.

PARAMETER	DESCRIPTION
site	The site, with a value fixed when one is given. TYPE: SiteRequest

RETURNS	DESCRIPTION
tuple[Tensor, Tensor]	The value the site takes, in the site's original space so downstream sites see the intended value, and the density a downstream sampler should score the site with.

Source code in src/quivers/effects/reparam/base.py

@abstractmethod
def apply(self, site: SiteRequest) -> tuple[torch.Tensor, torch.Tensor]:
    """Rewrite one site.

    Parameters
    ----------
    site : SiteRequest
        The site, with a value fixed when one is given.

    Returns
    -------
    tuple[torch.Tensor, torch.Tensor]
        The value the site takes, in the site's original space so
        downstream sites see the intended value, and the density a
        downstream sampler should score the site with.
    """

ReparamOrchestrator

ReparamOrchestrator(strategies: dict[str, Reparam])

Bases: EffectHandler

Dispatch sample sites to per-site reparam strategies.

PARAMETER	DESCRIPTION
strategies	Site name to strategy. TYPE: dict[str, Reparam]

Source code in src/quivers/effects/reparam/base.py

def __init__(self, strategies: dict[str, Reparam]) -> None:
    self.strategies = dict(strategies)

install

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

Install the dispatching handler of the random instance.

PARAMETER	DESCRIPTION
run	The run being prepared. TYPE: RunContext

RETURNS	DESCRIPTION
tuple[Installation, ...]	The handler, which answers covered sites and forwards the rest.

Source code in src/quivers/effects/reparam/base.py

def install(self, run: RunContext) -> tuple[Installation, ...]:
    """Install the dispatching handler of the ``random`` instance.

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

    Returns
    -------
    tuple[Installation, ...]
        The handler, which answers covered sites and forwards the
        rest.
    """
    kernel = run.kernel
    score_instance = kernel.score.entry.instance
    seen: dict[str, TorchSampleable] = {}
    definition = _handler_def(
        "Random.reparam",
        RANDOM,
        ((RANDOM_SAMPLE, ResumptionGrade.LINEAR),),
        key=f"reparam-{id(self):x}",
        input_type=run.result_type,
        output_type=run.result_type,
        introduced=EffectRow((RowEntry(score_instance, SCORE),)),
        total=False,
    )

    def sample(
        request: RuntimeRequest, resume: Resumption, context: ClauseContext
    ) -> object:
        """Answer a covered site with its strategy's value and density.

        Parameters
        ----------
        request : RuntimeRequest
            The request being answered and its arguments.
        resume : Resumption
            The continuation.
        context : ClauseContext
            Runtime services available to the clause.

        Returns
        -------
        object
            What the resumed computation produced, or a `Forward`
            for a site no strategy covers.
        """
        site, sampleable = _site_and_sampleable(request, definition.name)
        assert isinstance(site, str)
        assert isinstance(sampleable, TorchSampleable)
        strategy = self.strategies.get(site)
        if strategy is None:
            seen[site] = sampleable
            return Forward()
        value, log_prob = strategy.apply(
            SiteRequest(
                site,
                sampleable,
                seen=dict(seen),
                observation=run.observations.get(site),
            )
        )
        seen[site] = sampleable
        _emit_score(
            context,
            request,
            score_instance,
            log_prob,
            run.validator,
            role="draw-score",
        )
        return resume(value)

    runtime = RuntimeHandler(
        definition, {RANDOM_SAMPLE: RuntimeClause(sample, run.validator)}
    )
    return (Installation(kernel.random, runtime),)

SiteRequest dataclass

SiteRequest(name: str, sampleable: TorchSampleable, given: Tensor | None = None, seen: Mapping[str, TorchSampleable] = None, observation: Tensor | None = None)

What a strategy is given for one sample site.

PARAMETER	DESCRIPTION
name	The site's name. TYPE: str
sampleable	The site's distribution at its input. TYPE: TorchSampleable
given	A value the site is to take, when an enclosing strategy or a caller fixed one; None when the strategy draws. TYPE: Tensor | None DEFAULT: None
seen	The sampleables of the sites the orchestrator saw before this one, by name, so a strategy may read a parent's parameters. TYPE: Mapping[str, TorchSampleable] DEFAULT: None
observation	The run's observation for the site, if any. TYPE: Tensor | None DEFAULT: None

with_value

with_value(value: Tensor) -> SiteRequest

The same request with a value fixed.

PARAMETER	DESCRIPTION
value	The value the site is to take. TYPE: Tensor

RETURNS	DESCRIPTION
SiteRequest	The request with given set.

Source code in src/quivers/effects/reparam/base.py

def with_value(self, value: torch.Tensor) -> SiteRequest:
    """The same request with a value fixed.

    Parameters
    ----------
    value : torch.Tensor
        The value the site is to take.

    Returns
    -------
    SiteRequest
        The request with ``given`` set.
    """
    return replace(self, given=value)

TransformReparam

TransformReparam(bijector: Bijector)

Bases: Reparam

Reparameterise a sample site through a fixed bijector.

The reparameterised base sample lives in the bijector's domain and the site's value in its codomain. The joint density remains that of the original distribution: the strategy scores the value under the site's own sampleable, and the bijector reshapes the sampling geometry a downstream sampler works in.

PARAMETER	DESCRIPTION
bijector	The measurable bijection b the site's geometry is reshaped through. TYPE: Bijector

Source code in src/quivers/effects/reparam/transform.py

def __init__(self, bijector: Bijector) -> None:
    self.bijector = bijector

apply

apply(site: SiteRequest) -> tuple[Tensor, Tensor]

Rewrite one site.

PARAMETER	DESCRIPTION
site	The site. TYPE: SiteRequest

RETURNS	DESCRIPTION
tuple[Tensor, Tensor]	The value, drawn from the site's sampleable unless given, and its log density under the sampleable.

Source code in src/quivers/effects/reparam/transform.py

def apply(self, site: SiteRequest) -> tuple[torch.Tensor, torch.Tensor]:
    """Rewrite one site.

    Parameters
    ----------
    site : SiteRequest
        The site.

    Returns
    -------
    tuple[torch.Tensor, torch.Tensor]
        The value, drawn from the site's sampleable unless given, and
        its log density under the sampleable.
    """
    y = site.given if site.given is not None else site.sampleable.rsample()
    return y, site.sampleable.log_prob(y)

family_of

family_of(sampleable: TorchSampleable) -> str

The registry family a sampleable's morphism draws from.

PARAMETER	DESCRIPTION
sampleable	The sampleable. TYPE: TorchSampleable

RETURNS	DESCRIPTION
str	The family's registry name.

RAISES	DESCRIPTION
TypeError	If the morphism's class is registered for no family.

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

def family_of(sampleable: TorchSampleable) -> str:
    """The registry family a sampleable's morphism draws from.

    Parameters
    ----------
    sampleable : TorchSampleable
        The sampleable.

    Returns
    -------
    str
        The family's registry name.

    Raises
    ------
    TypeError
        If the morphism's class is registered for no family.
    """
    inner = sampleable.declared
    assert inner is not None
    for name, meta in FAMILY_META.items():
        if meta.quivers_class is not None and isinstance(inner, meta.quivers_class):
            return name
    get_dist = getattr(inner, "_get_dist", None)
    if callable(get_dist):
        distribution = get_dist(sampleable.input)
        for name, meta in FAMILY_META.items():
            if type(distribution) is meta.distribution_class:
                return name
    raise TypeError(
        f"{type(inner).__name__} draws from no registered family; a conjugate "
        "marginal needs the parent's and the child's families"
    )

normal_normal

normal_normal(parent: Mapping[str, Tensor], parent_value: Tensor, child: TorchSampleable, value: Tensor) -> Tensor

The marginal of y ~ Normal(z, tau) under z ~ Normal(mu, sigma).

PARAMETER	DESCRIPTION
parent	The parent's loc and scale. TYPE: Mapping[str, Tensor]
parent_value	The parent's drawn value. TYPE: Tensor
child	The child at the parent's drawn value, whose scale is tau. TYPE: TorchSampleable
value	The child's value. TYPE: Tensor

RETURNS	DESCRIPTION
Tensor	log Normal(y; mu, sqrt(sigma^2 + tau^2)) summed over the event.

RAISES	DESCRIPTION
ValueError	If the child's location is not the parent's value.

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

def normal_normal(
    parent: Mapping[str, torch.Tensor],
    parent_value: torch.Tensor,
    child: TorchSampleable,
    value: torch.Tensor,
) -> torch.Tensor:
    """The marginal of ``y ~ Normal(z, tau)`` under ``z ~ Normal(mu, sigma)``.

    Parameters
    ----------
    parent : Mapping[str, torch.Tensor]
        The parent's ``loc`` and ``scale``.
    parent_value : torch.Tensor
        The parent's drawn value.
    child : TorchSampleable
        The child at the parent's drawn value, whose ``scale`` is ``tau``.
    value : torch.Tensor
        The child's value.

    Returns
    -------
    torch.Tensor
        ``log Normal(y; mu, sqrt(sigma^2 + tau^2))`` summed over the event.

    Raises
    ------
    ValueError
        If the child's location is not the parent's value.
    """
    _identity_link(child, "loc", parent_value)
    tau = child.parameters()["scale"].reshape(parent["scale"].shape)
    scale = torch.sqrt(parent["scale"] ** 2 + tau**2)
    residual = (value - parent["loc"]) / scale
    log_density = -0.5 * residual**2 - scale.log() - 0.5 * math.log(2 * math.pi)
    return log_density.sum(dim=-1) if log_density.dim() > 1 else log_density

beta_bernoulli

beta_bernoulli(parent: Mapping[str, Tensor], parent_value: Tensor, child: TorchSampleable, value: Tensor) -> Tensor

The marginal of y ~ Bernoulli(z) under z ~ Beta(a, b).

PARAMETER	DESCRIPTION
parent	The parent's concentration1 and concentration0. TYPE: Mapping[str, Tensor]
parent_value	The parent's drawn value. TYPE: Tensor
child	The child at the parent's drawn value. TYPE: TorchSampleable
value	The child's value, zero or one. TYPE: Tensor

RETURNS	DESCRIPTION
Tensor	log(a / (a + b)) for one and log(b / (a + b)) for zero.

RAISES	DESCRIPTION
ValueError	If the child's probability is not the parent's value.

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

def beta_bernoulli(
    parent: Mapping[str, torch.Tensor],
    parent_value: torch.Tensor,
    child: TorchSampleable,
    value: torch.Tensor,
) -> torch.Tensor:
    """The marginal of ``y ~ Bernoulli(z)`` under ``z ~ Beta(a, b)``.

    Parameters
    ----------
    parent : Mapping[str, torch.Tensor]
        The parent's ``concentration1`` and ``concentration0``.
    parent_value : torch.Tensor
        The parent's drawn value.
    child : TorchSampleable
        The child at the parent's drawn value.
    value : torch.Tensor
        The child's value, zero or one.

    Returns
    -------
    torch.Tensor
        ``log(a / (a + b))`` for one and ``log(b / (a + b))`` for zero.

    Raises
    ------
    ValueError
        If the child's probability is not the parent's value.
    """
    _identity_link(child, "probs", parent_value)
    a = parent["concentration1"]
    b = parent["concentration0"]
    probability = a / (a + b)
    log_density = torch.where(
        value.to(probability.dtype) > 0.5, probability.log(), (1 - probability).log()
    )
    return log_density.sum(dim=-1) if log_density.dim() > 1 else log_density

gamma_poisson

gamma_poisson(parent: Mapping[str, Tensor], parent_value: Tensor, child: TorchSampleable, value: Tensor) -> Tensor

The marginal of y ~ Poisson(z) under z ~ Gamma(alpha, beta).

PARAMETER	DESCRIPTION
parent	The parent's concentration and rate. TYPE: Mapping[str, Tensor]
parent_value	The parent's drawn value. TYPE: Tensor
child	The child at the parent's drawn value. TYPE: TorchSampleable
value	The child's count. TYPE: Tensor

RETURNS	DESCRIPTION
Tensor	The negative binomial log mass with alpha successes and success probability beta / (beta + 1).

RAISES	DESCRIPTION
ValueError	If the child's rate is not the parent's value.

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

def gamma_poisson(
    parent: Mapping[str, torch.Tensor],
    parent_value: torch.Tensor,
    child: TorchSampleable,
    value: torch.Tensor,
) -> torch.Tensor:
    """The marginal of ``y ~ Poisson(z)`` under ``z ~ Gamma(alpha, beta)``.

    Parameters
    ----------
    parent : Mapping[str, torch.Tensor]
        The parent's ``concentration`` and ``rate``.
    parent_value : torch.Tensor
        The parent's drawn value.
    child : TorchSampleable
        The child at the parent's drawn value.
    value : torch.Tensor
        The child's count.

    Returns
    -------
    torch.Tensor
        The negative binomial log mass with ``alpha`` successes and
        success probability ``beta / (beta + 1)``.

    Raises
    ------
    ValueError
        If the child's rate is not the parent's value.
    """
    _identity_link(child, "rate", parent_value)
    alpha = parent["concentration"]
    beta = parent["rate"]
    count = value.to(alpha.dtype)
    log_density = (
        torch.lgamma(alpha + count)
        - torch.lgamma(alpha)
        - torch.lgamma(count + 1)
        + alpha * (beta.log() - torch.log1p(beta))
        - count * torch.log1p(beta)
    )
    return log_density.sum(dim=-1) if log_density.dim() > 1 else log_density

dirichlet_categorical

dirichlet_categorical(parent: Mapping[str, Tensor], parent_value: Tensor, child: TorchSampleable, value: Tensor) -> Tensor

The marginal of y ~ Categorical(z) under z ~ Dirichlet(alpha).

PARAMETER	DESCRIPTION
parent	The parent's concentration. TYPE: Mapping[str, Tensor]
parent_value	The parent's drawn value. TYPE: Tensor
child	The child at the parent's drawn value. TYPE: TorchSampleable
value	The child's category. TYPE: Tensor

RETURNS	DESCRIPTION
Tensor	log(alpha_y / sum(alpha)).

RAISES	DESCRIPTION
ValueError	If the child's probabilities are not the parent's value.

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

def dirichlet_categorical(
    parent: Mapping[str, torch.Tensor],
    parent_value: torch.Tensor,
    child: TorchSampleable,
    value: torch.Tensor,
) -> torch.Tensor:
    """The marginal of ``y ~ Categorical(z)`` under ``z ~ Dirichlet(alpha)``.

    Parameters
    ----------
    parent : Mapping[str, torch.Tensor]
        The parent's ``concentration``.
    parent_value : torch.Tensor
        The parent's drawn value.
    child : TorchSampleable
        The child at the parent's drawn value.
    value : torch.Tensor
        The child's category.

    Returns
    -------
    torch.Tensor
        ``log(alpha_y / sum(alpha))``.

    Raises
    ------
    ValueError
        If the child's probabilities are not the parent's value.
    """
    _identity_link(child, "probs", parent_value)
    alpha = parent["concentration"]
    probabilities = alpha / alpha.sum(dim=-1, keepdim=True)
    index = value.long().unsqueeze(-1)
    return probabilities.gather(-1, index).squeeze(-1).log()

conjugate_solver

conjugate_solver(parent_family: str, child_family: str) -> Solver

The solver of a conjugate pair.

PARAMETER	DESCRIPTION
parent_family	The parent's family. TYPE: str
child_family	The child's family. TYPE: str

RETURNS	DESCRIPTION
Solver	The closed-form marginal.

RAISES	DESCRIPTION
KeyError	If the pair has no registered marginal.

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

def conjugate_solver(parent_family: str, child_family: str) -> Solver:
    """The solver of a conjugate pair.

    Parameters
    ----------
    parent_family : str
        The parent's family.
    child_family : str
        The child's family.

    Returns
    -------
    Solver
        The closed-form marginal.

    Raises
    ------
    KeyError
        If the pair has no registered marginal.
    """
    try:
        return CONJUGATE_SOLVERS[(parent_family, child_family)]
    except KeyError as error:
        registered = ", ".join(f"{a}-{b}" for a, b in CONJUGATE_SOLVERS)
        raise KeyError(
            f"no conjugate marginal is registered for ({parent_family}, "
            f"{child_family}); the registered pairs are {registered}"
        ) from error

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)}

program_kernel

program_kernel(program: MonadicProgram, data: Sequence[str] = ()) -> ProgramKernel

Encode a program as a kernel module.

PARAMETER	DESCRIPTION
program	The program. TYPE: MonadicProgram
data	The names of the host data inputs a run supplies beside the program's input, in parameter order. TYPE: Sequence[str] DEFAULT: ()

RETURNS	DESCRIPTION
ProgramKernel	The module and the record the runtime serves it by.

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

def program_kernel(program: MonadicProgram, data: Sequence[str] = ()) -> ProgramKernel:
    """Encode a program as a kernel module.

    Parameters
    ----------
    program : MonadicProgram
        The program.
    data : Sequence[str]
        The names of the host data inputs a run supplies beside the
        program's input, in parameter order.

    Returns
    -------
    ProgramKernel
        The module and the record the runtime serves it by.
    """
    return _Builder(program, data).build()