Program

Top-level probabilistic program definitions and execution.

program

Program: compile a morphism expression into a trainable nn.Module.

Program

Program(morphism: Morphism | ContinuousMorphism | Module | None = None)

Bases: Module

Wraps a morphism expression as a trainable nn.Module.

Traverses the morphism DAG, collects all learnable parameters and observed buffers, and provides a forward() that materializes the composed tensor.

Supports both discrete Morphism instances (which produce a membership tensor via forward()) and ContinuousMorphism instances (which expose rsample / log_prob).

PARAMETER	DESCRIPTION
morphism	The root morphism expression (possibly a composition tree). TYPE: Morphism or ContinuousMorphism DEFAULT: None

Examples:

>>> from quivers import FinSet, morphism, Program
>>> X = FinSet(name="X", cardinality=3)
>>> Y = FinSet(name="Y", cardinality=4)
>>> f = morphism(X, Y)
>>> prog = Program(f)
>>> out = prog()  # shape (3, 4)

Source code in src/quivers/program.py

def __init__(
    self,
    morphism: Morphism | ContinuousMorphism | nn.Module | None = None,
) -> None:
    super().__init__()
    self._morphism = morphism
    self._is_continuous = isinstance(morphism, ContinuousMorphism)
    self._is_callable_module = (
        isinstance(morphism, nn.Module)
        and not isinstance(morphism, ContinuousMorphism)
        and not isinstance(morphism, Morphism)
    )

    # registering the module tree makes all parameters visible
    # to optimizer via self.parameters()
    if morphism is None:
        # A morphism-less Program is a container for structural
        # artifacts (signatures, encoders, decoders, losses)
        # in a module that declares no top-level exported morphism.
        self._root = None
    elif self._is_continuous or self._is_callable_module:
        # continuous morphisms and parser modules are already nn.Modules
        self._root = morphism
    else:
        self._root = cast(Morphism, morphism).module()

morphism property

morphism: Morphism | ContinuousMorphism | Module | None

The underlying morphism expression, or None for a morphism-less module (one declaring only signatures / encoders / decoders / losses).

domain property

domain

Domain of the underlying morphism.

codomain property

codomain

Codomain of the underlying morphism.

rsample

rsample(x: Tensor, sample_shape: Size = Size()) -> Tensor

Reparameterized sample (continuous programs only).

PARAMETER	DESCRIPTION
x	Input tensor. TYPE: Tensor
sample_shape	Extra leading sample dimensions. TYPE: Size DEFAULT: Size()

RETURNS	DESCRIPTION
Tensor	Sampled output.

RAISES	DESCRIPTION
TypeError	If the underlying morphism is not continuous.

Source code in src/quivers/program.py

def rsample(
    self,
    x: torch.Tensor,
    sample_shape: torch.Size = torch.Size(),
) -> torch.Tensor:
    """Reparameterized sample (continuous programs only).

    Parameters
    ----------
    x : torch.Tensor
        Input tensor.
    sample_shape : torch.Size
        Extra leading sample dimensions.

    Returns
    -------
    torch.Tensor
        Sampled output.

    Raises
    ------
    TypeError
        If the underlying morphism is not continuous.
    """
    if not isinstance(self._morphism, ContinuousMorphism):
        raise TypeError("rsample is only available for continuous programs")

    return cast(torch.Tensor, self._morphism.rsample(x, sample_shape))

log_prob

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

Log-probability (continuous programs only).

PARAMETER	DESCRIPTION
x	Input tensor. TYPE: Tensor
y	Output tensor. TYPE: Tensor

RETURNS	DESCRIPTION
Tensor	Log-probability of y given x.

RAISES	DESCRIPTION
TypeError	If the underlying morphism is not continuous.

Source code in src/quivers/program.py

def log_prob(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
    """Log-probability (continuous programs only).

    Parameters
    ----------
    x : torch.Tensor
        Input tensor.
    y : torch.Tensor
        Output tensor.

    Returns
    -------
    torch.Tensor
        Log-probability of y given x.

    Raises
    ------
    TypeError
        If the underlying morphism is not continuous.
    """
    if not isinstance(self._morphism, ContinuousMorphism):
        raise TypeError("log_prob is only available for continuous programs")

    return cast(torch.Tensor, self._morphism.log_prob(x, y))

forward

forward(n_steps: int | None = None) -> Tensor

Materialize the composed tensor (discrete programs only).

PARAMETER	DESCRIPTION
n_steps	If provided, sets the step count on all RepeatMorphism instances in the morphism tree before materializing. This allows runtime-variable sequence lengths for models using repeat(f). TYPE: int or None DEFAULT: None

RETURNS	DESCRIPTION
Tensor	Tensor of shape (*domain.shape, *codomain.shape) with values in [0, 1].

RAISES	DESCRIPTION
TypeError	If the underlying morphism is continuous.

Source code in src/quivers/program.py

def forward(self, n_steps: int | None = None) -> torch.Tensor:
    """Materialize the composed tensor (discrete programs only).

    Parameters
    ----------
    n_steps : int or None
        If provided, sets the step count on all
        ``RepeatMorphism`` instances in the morphism tree
        before materializing. This allows runtime-variable
        sequence lengths for models using ``repeat(f)``.

    Returns
    -------
    torch.Tensor
        Tensor of shape ``(*domain.shape, *codomain.shape)`` with
        values in [0, 1].

    Raises
    ------
    TypeError
        If the underlying morphism is continuous.
    """
    if self._morphism is None:
        raise TypeError(
            "forward() is not supported on a Program with no exported "
            "morphism; this module declares structural artifacts only"
        )
    if self._is_continuous:
        raise TypeError(
            "forward() is not supported for continuous programs; "
            "use rsample() or log_prob() instead"
        )

    if n_steps is not None:
        self._set_repeat_steps(n_steps)

    return cast(Morphism, self._morphism).tensor

log_membership

log_membership() -> Tensor

Log of the membership tensor.

RETURNS	DESCRIPTION
Tensor	Log-membership values. Entries near 0 map to large negative values.

Source code in src/quivers/program.py

def log_membership(self) -> torch.Tensor:
    """Log of the membership tensor.

    Returns
    -------
    torch.Tensor
        Log-membership values.  Entries near 0 map to large
        negative values.
    """
    t = cast(Morphism, self._morphism).tensor
    return torch.log(t.clamp(min=1e-7))

nll_loss

nll_loss(domain_indices: Tensor, codomain_indices: Tensor) -> Tensor

Negative log-likelihood loss for observed (domain, codomain) pairs.

For each pair (x, y), computes -log(tensor[x, y]). Suitable when the morphism represents fuzzy membership and observed pairs should have high membership.

PARAMETER	DESCRIPTION
domain_indices	Integer tensor of shape (batch,) or (batch, n_domain_dims) indexing into the domain. TYPE: Tensor
codomain_indices	Integer tensor of shape (batch,) or (batch, n_codomain_dims) indexing into the codomain. TYPE: Tensor

RETURNS	DESCRIPTION
Tensor	Scalar mean negative log-likelihood.

Source code in src/quivers/program.py

def nll_loss(
    self,
    domain_indices: torch.Tensor,
    codomain_indices: torch.Tensor,
) -> torch.Tensor:
    """Negative log-likelihood loss for observed (domain, codomain) pairs.

    For each pair (x, y), computes -log(tensor[x, y]).
    Suitable when the morphism represents fuzzy membership and
    observed pairs should have high membership.

    Parameters
    ----------
    domain_indices : torch.Tensor
        Integer tensor of shape (batch,) or (batch, n_domain_dims)
        indexing into the domain.
    codomain_indices : torch.Tensor
        Integer tensor of shape (batch,) or (batch, n_codomain_dims)
        indexing into the codomain.

    Returns
    -------
    torch.Tensor
        Scalar mean negative log-likelihood.
    """
    t = cast(Morphism, self._morphism).tensor

    # handle both flat and multi-dimensional indexing
    if domain_indices.ndim == 1:
        domain_indices = domain_indices.unsqueeze(-1)

    if codomain_indices.ndim == 1:
        codomain_indices = codomain_indices.unsqueeze(-1)

    # combine indices
    indices = torch.cat([domain_indices, codomain_indices], dim=-1)

    # index into the tensor
    idx_tuple = tuple(indices[:, i] for i in range(indices.shape[1]))
    values = t[idx_tuple]

    return -torch.log(values.clamp(min=1e-7)).mean()

bce_loss

bce_loss(target: Tensor) -> Tensor

Binary cross-entropy between the morphism tensor and a target.

PARAMETER	DESCRIPTION
target	Target tensor of the same shape as the morphism tensor, with values in [0, 1]. TYPE: Tensor

RETURNS	DESCRIPTION
Tensor	Scalar BCE loss.

Source code in src/quivers/program.py

def bce_loss(self, target: torch.Tensor) -> torch.Tensor:
    """Binary cross-entropy between the morphism tensor and a target.

    Parameters
    ----------
    target : torch.Tensor
        Target tensor of the same shape as the morphism tensor,
        with values in [0, 1].

    Returns
    -------
    torch.Tensor
        Scalar BCE loss.
    """
    t = cast(Morphism, self._morphism).tensor
    return F.binary_cross_entropy(t, target)