Continuous Programs

Probabilistic programs in continuous domains.

programs

Monadic programs: sequenced probabilistic programs as ContinuousMorphisms.

A MonadicProgram defines a ContinuousMorphism via monadic sequencing of draw steps. Each step samples from a named morphism, optionally conditioned on previously drawn variables, and binds the result. The program returns one or more of the bound variables as its output.

This corresponds to the Kleisli composition pattern used in probabilistic programming languages like PDS (Grove & White), where a sequence of let' x ~ D in ... bindings threads probabilistic state through a generative model.

Features

• Single and tuple returns: return x or return (x, y, z)
• Named input parameters for product-domain sub-programs
• Multi-argument draw steps: draw z ~ f(x, y)
• Destructuring draws from tuple-returning sub-programs: draw (a, b) ~ sub_prog(x)

Example

Given morphisms f : A -> B and g : B -> C, the monadic program::

program p : A -> C
    draw x ~ f
    draw y ~ g(x)
    return y

is equivalent to the composition f >> g, but the program form allows fan-out (using the input in multiple draws) and non-linear variable dependency graphs.

PDS-style nested programs::

program cg_update(y, z) : Belief * Belief -> Truth * Truth
    draw c ~ bern_c(y)
    draw d ~ bern_d(z)
    return (c, d)

program factivityPrior : Entity -> Truth * Truth * Truth
    draw x ~ prior_x
    draw y ~ prior_y
    draw z ~ prior_z
    draw b ~ bern_b(x)
    draw (c, d) ~ cg_update(y, z)
    return (b, c, d)

ProgramEvaluator

Bases: Protocol

What runs a program: the reference machine under the handler stack.

The evaluator lives in quivers.effects, which builds on programs; it installs itself here when imported, so a program's own methods reach it without the program package depending on the effects package at import.

sample_program

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

Run the program forward.

PARAMETER	DESCRIPTION
program	The program. TYPE: MonadicProgram
x	Program input. TYPE: Tensor
sample_shape	Leading sample dimensions. TYPE: Size
observations	Values to clamp observed variables to, and host data. TYPE: Mapping[str, Tensor] or None

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

Source code in src/quivers/continuous/programs.py

def sample_program(
    self,
    program: MonadicProgram,
    x: torch.Tensor,
    sample_shape: torch.Size,
    observations: Mapping[str, torch.Tensor] | None,
) -> torch.Tensor | dict[str, torch.Tensor]:
    """Run the program forward.

    Parameters
    ----------
    program : MonadicProgram
        The program.
    x : torch.Tensor
        Program input.
    sample_shape : torch.Size
        Leading sample dimensions.
    observations : Mapping[str, torch.Tensor] or None
        Values to clamp observed variables to, and host data.

    Returns
    -------
    torch.Tensor or dict[str, torch.Tensor]
        The program's return value.
    """
    ...

log_joint

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

Score the program at values for all of its draws.

PARAMETER	DESCRIPTION
program	The program. TYPE: MonadicProgram
x	Program input. TYPE: Tensor
intermediates	A value for every draw, and host data. TYPE: Mapping[str, Tensor]

RETURNS	DESCRIPTION
Tensor	The joint log density.

Source code in src/quivers/continuous/programs.py

def log_joint(
    self,
    program: MonadicProgram,
    x: torch.Tensor,
    intermediates: Mapping[str, torch.Tensor],
) -> torch.Tensor:
    """Score the program at values for all of its draws.

    Parameters
    ----------
    program : MonadicProgram
        The program.
    x : torch.Tensor
        Program input.
    intermediates : Mapping[str, torch.Tensor]
        A value for every draw, and host data.

    Returns
    -------
    torch.Tensor
        The joint log density.
    """
    ...

MonadicProgram

MonadicProgram(domain: AnySpace, codomain: AnySpace, steps: Sequence[Step], return_vars: tuple[str, ...], params: tuple[str, ...] | None = None, return_labels: tuple[str, ...] | None = None, effect_set: frozenset[str] | None = None)

Bases: ContinuousMorphism

A probabilistic program defined by monadic sequencing of draw steps.

Each draw step samples from a ContinuousMorphism and binds the result to one or more named variables. Later steps can reference earlier bindings as their input. The program's output is the value(s) of the designated return variable(s).

PARAMETER	DESCRIPTION
domain	The program's input space. TYPE: SetObject or ContinuousSpace
codomain	The program's output space. TYPE: SetObject or ContinuousSpace
steps	The program body, in order: Draw, Observe, Let, and Score records. TYPE: Sequence[Step]
return_vars	Name(s) of the bound variable(s) whose value(s) are the program output. TYPE: tuple[str, ...]
params	Named input parameters for product-domain programs. When set, the program input is split along the feature dimension and each component is pre-bound in the env. TYPE: tuple[str, ...] or None DEFAULT: None
return_labels	Optional labels for tuple return fields. When set, the output dict uses these labels as keys instead of the variable names. Length must match return_vars. TYPE: tuple[str, ...] or None DEFAULT: None
effect_set	The declared effect row (Sample, Score, Marginal, Pure), or None when unannotated. TYPE: frozenset[str] or None DEFAULT: None

RAISES	DESCRIPTION
ValueError	If a step binds no name, or two draws or observes bind the same name.

Source code in src/quivers/continuous/programs.py

def __init__(
    self,
    domain: AnySpace,
    codomain: AnySpace,
    steps: Sequence[Step],
    return_vars: tuple[str, ...],
    params: tuple[str, ...] | None = None,
    return_labels: tuple[str, ...] | None = None,
    effect_set: frozenset[str] | None = None,
) -> None:
    super().__init__(domain, codomain)
    self._return_vars = return_vars
    self._return_is_single = len(return_vars) == 1
    self._params = params
    self._return_labels = return_labels
    # The effect-row annotation. None when unannotated;
    # otherwise carries the declared capability set
    # (Sample / Score / Marginal / Pure) for introspection by
    # downstream inference / dispatch code.
    self.effect_set: frozenset[str] | None = effect_set
    # compute input component dimensions for param splitting
    if params is not None and len(params) > 1:
        self._param_dims = self._compute_component_dims(domain)
        self._param_is_continuous = self._compute_component_continuous(domain)

    else:
        self._param_dims = None
        self._param_is_continuous = None

    # Each draw's morphism registers as a named submodule so its
    # parameters are visible to optimizers.
    for step in steps:
        if not _bound_names(step):
            raise ValueError(f"{step.kind} step binds no name")
        if isinstance(step, Draw | Observe):
            key = step_key(step)
            if key in self._modules:
                raise ValueError(
                    f"two {step.kind} steps bind {step.names[0]!r}; each "
                    f"draw or observe binds a name no other one does"
                )
            # A backend-agnostic `Morphism` registers through its
            # parameter container, which keeps the morphism under
            # ``_morphism`` for `extract_morphism`.
            self.add_module(key, as_torch_module(step.morphism))
    self._steps: tuple[Step, ...] = tuple(steps)

steps property

steps: tuple[Step, ...]

The program body, in order.

RETURNS	DESCRIPTION
tuple[Step, ...]	The records the program was built from.

observed_names property

observed_names: set[str]

Return the set of variable names marked as observed in the DSL.

step_module

step_module(step: Draw | Observe) -> Module

The registered module a draw or observe step applies.

For a ContinuousMorphism this is the morphism itself; for a backend-agnostic Morphism it is the parameter container as_torch_module wraps it in.

PARAMETER	DESCRIPTION
step	One of the program's steps. TYPE: Draw | Observe

RETURNS	DESCRIPTION
Module	The submodule the program registered for the step.

RAISES	DESCRIPTION
KeyError	If the program registered no module for the step.

Source code in src/quivers/continuous/programs.py

def step_module(self, step: Draw | Observe) -> nn.Module:
    """The registered module a draw or observe step applies.

    For a [`ContinuousMorphism`][quivers.continuous.ContinuousMorphism]
    this is the morphism itself; for a backend-agnostic
    [`Morphism`][quivers.core.Morphism] it is the parameter container
    [`as_torch_module`][quivers.core.as_torch_module] wraps it in.

    Parameters
    ----------
    step : Draw | Observe
        One of the program's steps.

    Returns
    -------
    nn.Module
        The submodule the program registered for the step.

    Raises
    ------
    KeyError
        If the program registered no module for the step.
    """
    module = self._modules[step_key(step)]
    if module is None:
        raise KeyError(step_key(step))
    return module

rsample

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

Run the program forward, returning the designated output(s).

The program runs on the reference machine under the active handler stack, so a clamp, do, or trace in scope applies to the draws the run makes.

PARAMETER	DESCRIPTION
x	Program input. TYPE: Tensor
sample_shape	Additional leading sample dimensions; each element is one independent run of the program. TYPE: Size DEFAULT: Size()
observations	Values to clamp observed variables to, keyed by variable name, together with any host data the steps read by name. TYPE: dict[str, Tensor] or None DEFAULT: None

RETURNS	DESCRIPTION
Tensor or dict[str, Tensor]	The value of the return variable(s). Returns a tensor for single-variable returns, or a dict keyed by variable name for tuple returns.

Source code in src/quivers/continuous/programs.py

def rsample(  # type: ignore[override]
    self,
    x: torch.Tensor,
    sample_shape: torch.Size = torch.Size(),
    observations: dict[str, torch.Tensor] | None = None,
) -> torch.Tensor | dict[str, torch.Tensor]:
    """Run the program forward, returning the designated output(s).

    The program runs on the reference machine under the active
    handler stack, so a `clamp`, `do`, or `trace` in scope applies
    to the draws the run makes.

    Parameters
    ----------
    x : torch.Tensor
        Program input.
    sample_shape : torch.Size
        Additional leading sample dimensions; each element is one
        independent run of the program.
    observations : dict[str, torch.Tensor] or None
        Values to clamp observed variables to, keyed by variable
        name, together with any host data the steps read by name.

    Returns
    -------
    torch.Tensor or dict[str, torch.Tensor]
        The value of the return variable(s). Returns a tensor
        for single-variable returns, or a dict keyed by variable
        name for tuple returns.
    """
    return _evaluator().sample_program(self, x, sample_shape, observations)

has_conditional_density

has_conditional_density() -> bool

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

log_prob raises for exactly this reason, and a caller choosing between constructions reads the answer here rather than by catching that raise. log_joint is the exact route: it scores the draws themselves, given values for all of them.

Source code in src/quivers/continuous/programs.py

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

    [`log_prob`][quivers.continuous.programs.MonadicProgram.log_prob]
    raises for exactly this reason, and a caller choosing between
    constructions reads the answer here rather than by catching
    that raise.
    [`log_joint`][quivers.continuous.programs.MonadicProgram.log_joint]
    is the exact route: it scores the draws themselves, given
    values for all of them.
    """
    return False

log_prob

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

Log-probability is not supported for monadic programs.

Computing log p(y | x) for a monadic program requires marginalizing over all intermediate variables, which is intractable in general. Use rsample for forward sampling and condition via score function estimators or variational methods.

RAISES	DESCRIPTION
NotImplementedError	Always.

Source code in src/quivers/continuous/programs.py

def log_prob(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
    """Log-probability is not supported for monadic programs.

    Computing log p(y | x) for a monadic program requires
    marginalizing over all intermediate variables, which is
    intractable in general. Use ``rsample`` for forward sampling
    and condition via score function estimators or variational
    methods.

    Raises
    ------
    NotImplementedError
        Always.
    """
    raise NotImplementedError(
        "log_prob is not supported for monadic programs; "
        "computing p(y | x) requires marginalizing over all "
        "intermediate draws, which is intractable in general. "
        "use rsample() for forward sampling."
    )

log_joint

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

Joint log-density given all intermediate values.

When all intermediate variables are observed (e.g. during inference with HMC/NUTS), computes the joint log-density:

log p(x_1, ..., x_n | input) = sum_i log p(x_i | pa(x_i))

where pa(x_i) is the parent variable of step i (either the program input or a previously drawn variable). The program runs on the reference machine with every draw conditioned on its given value, under the active handler stack, so the joint is the one trace reports.

PARAMETER	DESCRIPTION
x	Program input. TYPE: Tensor
intermediates	Values for ALL bound variables (keyed by variable name or by return label if labels are set), together with any host data the steps read by name. TYPE: dict[str, Tensor]

RETURNS	DESCRIPTION
Tensor	Joint log-density. Shape (batch,).

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

Source code in src/quivers/continuous/programs.py

def log_joint(
    self,
    x: torch.Tensor,
    intermediates: dict[str, torch.Tensor],
) -> torch.Tensor:
    """Joint log-density given all intermediate values.

    When all intermediate variables are observed (e.g. during
    inference with HMC/NUTS), computes the joint log-density:

        log p(x_1, ..., x_n | input) = sum_i log p(x_i | pa(x_i))

    where pa(x_i) is the parent variable of step i (either the
    program input or a previously drawn variable). The program
    runs on the reference machine with every draw conditioned on
    its given value, under the active handler stack, so the joint
    is the one [`trace`][quivers.inference.trace.trace] reports.

    Parameters
    ----------
    x : torch.Tensor
        Program input.
    intermediates : dict[str, torch.Tensor]
        Values for ALL bound variables (keyed by variable name
        or by return label if labels are set), together with any
        host data the steps read by name.

    Returns
    -------
    torch.Tensor
        Joint log-density. Shape (batch,).

    Raises
    ------
    KeyError
        If a draw the program makes is given no value.
    """
    return _evaluator().log_joint(self, x, intermediates)

step_key

step_key(step: Draw | Observe) -> str

The submodule name a program registers a draw's morphism under.

PARAMETER	DESCRIPTION
step	A draw or observe step. TYPE: Draw | Observe

RETURNS	DESCRIPTION
str	_step_ followed by the step's first bound name.

Source code in src/quivers/continuous/programs.py

def step_key(step: Draw | Observe) -> str:
    """The submodule name a program registers a draw's morphism under.

    Parameters
    ----------
    step : Draw | Observe
        A draw or observe step.

    Returns
    -------
    str
        ``_step_`` followed by the step's first bound name.
    """
    return f"_step_{step.names[0]}"

install_evaluator

install_evaluator(evaluator: ProgramEvaluator) -> None

Install the evaluator every program's methods run through.

PARAMETER	DESCRIPTION
evaluator	The evaluator. TYPE: ProgramEvaluator

Source code in src/quivers/continuous/programs.py

def install_evaluator(evaluator: ProgramEvaluator) -> None:
    """Install the evaluator every program's methods run through.

    Parameters
    ----------
    evaluator : ProgramEvaluator
        The evaluator.
    """
    global _EVALUATOR
    _EVALUATOR = evaluator