Parameter Sources

quivers.continuous.param_source provides the map from a conditional family's input to its distribution parameters. A morphism declared ~ Family over a continuous domain is a Kleisli arrow whose parameters are produced by a ParamSource, so the source is where a kernel's dependence on its input is computed, and where any nonlinearity in that dependence lives.

The concrete sources cover the standard architectures: LinearSource is a single nn.Linear; MLPSource is a multi-layer perceptron with configurable widths and activation; AttentionSource is a self-attention head; LookupSource and EmbeddingSource handle discrete domains; IdentitySource, FunctionSource, and ComposeSource cover pass-through, a fixed callable, and composition of two sources.

make_param_source is the factory the families call, and param_source_from_option parses the DSL's [param_source=<kind>] morphism option. The default for a continuous domain is LinearSource, so a kernel is linear unless it asks for something else; a SetObject domain always uses LookupSource regardless of the requested kind. The Bayesian Neural Network example selects the MLP source explicitly and relies on it for its nonlinearity.

The hidden widths come either from the option's arguments or from hidden_dim, one width per hidden layer:

morphism f : X -> Y [param_source=mlp] ~ Normal                      # (64, 64)
morphism f : X -> Y [param_source=mlp(64, 32)] ~ Normal              # (64, 32)
morphism f : X -> Y [param_source=mlp, hidden_dim=[64, 32]] ~ Normal # (64, 32)
morphism f : X -> Y [param_source=mlp, hidden_dim=64] ~ Normal       # (64,)

A width given to a source with no hidden layers to apply it to is an error rather than a silent no-op, and so is param_source on a family whose parameters do not come from a source at all (Horseshoe, GaussianProcess, Independent, Transformed).

param_source

Pluggable parameter sources for conditional distribution families.

A ParamSource produces a (batch, param_dim) tensor from a per-row input. Every ConditionalX family in quivers.continuous.families uses one to convert its input into the flattened parameter vector its underlying distribution needs.

The primitives:

• LinearSource: the default; one nn.Linear, no nonlinearity. Matches the single-linear layer the transpile backends emit exactly, so a kernel morphism on this source is numerically equivalent to its transpiled counterpart.
• MLPSource: a multi-layer perceptron with user-configurable hidden widths and activation. Selected by [param_source=mlp]; a kernel is linear unless it asks for this.
• LookupSource: a learnable per-entry embedding table, the discrete-domain standard.
• EmbeddingSource: an embedding table piped through a downstream ParamSource (embedding + MLP head, the standard categorical-input pattern).
• AttentionSource: single-head self-attention over the input dimension, a useful primitive for set-valued inputs.
• IdentitySource: pass the input through unchanged (parameters supplied as data).
• FunctionSource: wraps an arbitrary Callable, letting a user drop in any nn.Module without subclassing.
• ComposeSource: categorical composition of two param sources, useful for building sequential architectures out of primitives.

The DSL surface accepts a [param_source=...] option on morphism declarations, with the hidden widths given either as the option's arguments or through hidden_dim:

morphism trans : State -> State [param_source=linear] ~ Normal
morphism trans : State -> State [param_source=mlp] ~ Normal
morphism trans : State -> State [param_source=mlp(64, 64)] ~ Normal
morphism trans : State -> State [param_source=mlp, hidden_dim=[64, 32]] ~ Normal

One width per hidden layer, so the sequence says how many as well as how wide; a bare hidden_dim=64 is the one-layer case. The call form is parsed by param_source_from_option.

ParamSource

Bases: Module, ABC

A learnable map from per-row input to a flat parameter vector.

Every conditional distribution family holds a ParamSource instance and defers its parameter computation to it. Subclasses override forward and declare param_dim.

LinearSource

LinearSource(domain_dim: int, param_dim: int, bias: bool = True)

Bases: ParamSource

Single nn.Linear(domain_dim, param_dim); no nonlinearity.

This is the parameter source that matches the transpile backends' emit: the DSL morphism f : A -> B [role=kernel] ~ Family lowers to a single linear layer Family(loc = W x, ...) on every backend. Configuring a runtime kernel with LinearSource(...) makes the runtime numerically equivalent to its transpiled counterpart.

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

def __init__(self, domain_dim: int, param_dim: int, bias: bool = True) -> None:
    super().__init__()
    self._domain_dim = int(domain_dim)
    self._param_dim = int(param_dim)
    self.linear = nn.Linear(domain_dim, param_dim, bias=bias)

MLPSource

MLPSource(domain_dim: int, param_dim: int, hidden_dims: Sequence[int] = (64, 64), activation: type[Module] = Tanh, bias: bool = True)

Bases: ParamSource

Multi-layer perceptron with configurable hidden widths and activation. The default hidden_dims=(64, 64) yields a two-hidden-layer, tanh-activated network; a user who wants a wider, deeper, or differently-activated network drops in the same constructor.

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

def __init__(
    self,
    domain_dim: int,
    param_dim: int,
    hidden_dims: Sequence[int] = (64, 64),
    activation: type[nn.Module] = nn.Tanh,
    bias: bool = True,
) -> None:
    super().__init__()
    self._domain_dim = int(domain_dim)
    self._param_dim = int(param_dim)
    layers: list[nn.Module] = []
    widths = [domain_dim, *hidden_dims, param_dim]
    for i in range(len(widths) - 1):
        layers.append(nn.Linear(widths[i], widths[i + 1], bias=bias))
        if i < len(widths) - 2:
            layers.append(activation())
    self.net = nn.Sequential(*layers)

LookupSource

LookupSource(n_entries: int, param_dim: int)

Bases: ParamSource

Per-entry learnable parameter table indexed by an integer input. The categorical / discrete-input parameter source: the input is a LongTensor of category indices, the output is the per-index parameter vector.

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

def __init__(self, n_entries: int, param_dim: int) -> None:
    super().__init__()
    self._n_entries = int(n_entries)
    self._param_dim = int(param_dim)
    self.table = nn.Parameter(torch.randn(n_entries, param_dim) * 0.01)

EmbeddingSource

EmbeddingSource(n_entries: int, embed_dim: int, head: ParamSource)

Bases: ParamSource

Embedding table followed by a downstream ParamSource. The canonical embedding + MLP head pattern: categorical inputs pass through an nn.Embedding and the resulting dense vector feeds a user-supplied head (linear, MLP, or attention).

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

def __init__(
    self,
    n_entries: int,
    embed_dim: int,
    head: ParamSource,
) -> None:
    super().__init__()
    if head._domain_dim != embed_dim:
        raise ValueError(
            "EmbeddingSource: head's domain_dim must equal embed_dim; "
            f"got head.domain_dim={head._domain_dim}, embed_dim={embed_dim}"
        )
    self.embed = nn.Embedding(n_entries, embed_dim)
    self.head = head
    self._embed_dim = int(embed_dim)

AttentionSource

AttentionSource(domain_dim: int, param_dim: int, num_heads: int = 4, bias: bool = True)

Bases: ParamSource

Single-head self-attention over the input feature dimension, followed by an aggregation and a linear head. Useful when the input is a set of features whose ordering carries no information and the model should be permutation-equivariant.

For a (batch, seq_len, domain_dim) input, computes scaled-dot-product attention across the sequence dimension, aggregates via mean-pool, and projects to param_dim via a linear head. When the input is (batch, domain_dim), treats it as a length-one sequence.

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

def __init__(
    self,
    domain_dim: int,
    param_dim: int,
    num_heads: int = 4,
    bias: bool = True,
) -> None:
    super().__init__()
    self._domain_dim = int(domain_dim)
    self._param_dim = int(param_dim)
    self.attn = nn.MultiheadAttention(
        embed_dim=domain_dim,
        num_heads=num_heads,
        batch_first=True,
        bias=bias,
    )
    self.head = nn.Linear(domain_dim, param_dim, bias=bias)

IdentitySource

IdentitySource(param_dim: int)

Bases: ParamSource

Passes the input through unchanged. Useful when parameters come directly from the host (e.g. a design matrix or an already- computed feature vector) and no further transformation is needed. param_dim equals the input's last-axis size and must be declared at construction.

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

def __init__(self, param_dim: int) -> None:
    super().__init__()
    self._param_dim = int(param_dim)

FunctionSource

FunctionSource(fn: Callable[[Tensor], Tensor], param_dim: int)

Bases: ParamSource

Wraps an arbitrary nn.Module or callable as a ParamSource. The user supplies the module and declares its output dim; the machinery around conditional families sees a uniform interface.

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

def __init__(self, fn: Callable[[Tensor], Tensor], param_dim: int) -> None:
    super().__init__()
    self._param_dim = int(param_dim)
    if isinstance(fn, nn.Module):
        self._module: nn.Module | None = fn
        self._fn: Callable[[Tensor], Tensor] | None = None
    else:
        self._module = None
        self._fn = fn

ComposeSource

ComposeSource(outer: ParamSource, inner: ParamSource)

Bases: ParamSource

Categorical composition of two param sources: outer(inner(x)). inner.param_dim must equal outer._domain_dim. Enables building sequential architectures from primitives (e.g. Compose(MLPSource(...), LinearSource(...))).

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

def __init__(self, outer: ParamSource, inner: ParamSource) -> None:
    super().__init__()
    if inner.param_dim != outer._domain_dim:
        raise ValueError(
            "ComposeSource: inner.param_dim must equal outer.domain_dim; "
            f"got inner.param_dim={inner.param_dim}, "
            f"outer.domain_dim={outer._domain_dim}"
        )
    self.outer = outer
    self.inner = inner

make_param_source

make_param_source(domain: AnySpace, param_dim: int, kind: str = _DEFAULT_SOURCE_KIND, **kwargs) -> ParamSource

Factory that dispatches the [param_source=...] DSL option to the concrete class.

The default is LinearSource, so a morphism declared f : X -> Y ~ Normal maps its input to the family's parameters the way its arrow reads: linearly. A model that wants a nonlinearity between its input and its parameters asks for one, and the fact then appears in the source rather than in a default.

Recognised kinds: * "lookup" — always used when the domain is a SetObject, regardless of the requested kind. * "linear" — the default; one nn.Linear. * "mlp" — hidden_dims=(64, 64) unless overridden by hidden_dims or hidden_dim kwargs. * "identity" — pass through. * "attention" — self-attention head.

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

def make_param_source(
    domain: AnySpace,
    param_dim: int,
    kind: str = _DEFAULT_SOURCE_KIND,
    **kwargs,
) -> ParamSource:
    """Factory that dispatches the `[param_source=...]` DSL option
    to the concrete class.

    The default is `LinearSource`, so a morphism declared
    ``f : X -> Y ~ Normal`` maps its input to the family's parameters
    the way its arrow reads: linearly. A model that wants a
    nonlinearity between its input and its parameters asks for one,
    and the fact then appears in the source rather than in a default.

    Recognised kinds:
    * ``"lookup"`` — always used when the domain is a `SetObject`,
      regardless of the requested kind.
    * ``"linear"`` — the default; one `nn.Linear`.
    * ``"mlp"`` — `hidden_dims=(64, 64)` unless overridden by
      ``hidden_dims`` or ``hidden_dim`` kwargs.
    * ``"identity"`` — pass through.
    * ``"attention"`` — self-attention head.
    """
    if isinstance(domain, SetObject):
        return LookupSource(domain.size, param_dim)
    if not isinstance(domain, ContinuousSpace):
        raise TypeError(
            f"make_param_source: unsupported domain type {type(domain).__name__}"
        )
    dim = domain.dim
    if kind == "linear":
        return LinearSource(dim, param_dim, **kwargs)
    if kind == "identity":
        return IdentitySource(param_dim)
    if kind == "attention":
        if "heads" in kwargs:
            kwargs["num_heads"] = kwargs.pop("heads")
        return AttentionSource(dim, param_dim, **kwargs)
    if kind == "mlp":
        hidden_dims = kwargs.pop("hidden_dims", None)
        if hidden_dims is None:
            hidden_dim = kwargs.pop("hidden_dim", 64)
            hidden_dims = (int(hidden_dim), int(hidden_dim))
        return MLPSource(dim, param_dim, hidden_dims=hidden_dims, **kwargs)
    raise ValueError(f"make_param_source: unknown kind {kind!r}")

param_source_from_option

param_source_from_option(domain: AnySpace, param_dim: int, option_value: str | None) -> ParamSource

Parse a [param_source=...] option string into a ParamSource. Accepts:

• "mlp", "linear", "identity", "attention"
• "mlp(64, 64)" — parenthesised hidden widths
• "mlp(32)" — single hidden width
• "attention(heads=4)" — keyword-only options

Unrecognised syntax raises ValueError so parse errors surface at compile time rather than as silent identity fallthrough.

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

def param_source_from_option(
    domain: AnySpace,
    param_dim: int,
    option_value: str | None,
) -> ParamSource:
    """Parse a `[param_source=...]` option string into a
    `ParamSource`. Accepts:

    * ``"mlp"``, ``"linear"``, ``"identity"``, ``"attention"``
    * ``"mlp(64, 64)"`` — parenthesised hidden widths
    * ``"mlp(32)"`` — single hidden width
    * ``"attention(heads=4)"`` — keyword-only options

    Unrecognised syntax raises `ValueError` so parse errors surface
    at compile time rather than as silent identity fallthrough.
    """
    kind, kwargs = _split_source_option(option_value)
    return make_param_source(domain, param_dim, kind=kind, **kwargs)