Shapes

Ready-made signatures, encoders, and decoders for the three principal compressible shapes: sequences, trees, and graphs. Each factory builds on the generic Encoder and Decoder runtimes, so the objects it returns can be inspected and composed like any hand-built one.

seq

Sequence-shaped signatures and canonical encoders / decoders.

Surface algebra::

Seq[A] = Nil | Cons(A, Seq[A])

The principal shape behind autoregressive LMs (the "compress a prefix to a vector, decode next-token from the vector" loop). All encoders and decoders below are realised on top of the generic quivers.structural.Encoder and quivers.structural.Decoder runtimes, with sequence- specific per-op functions.

seq_signature

seq_signature(name: str = 'Seq', dim: int = 64) -> Signature

Return a cons-list signature with a data element sort.

The object sort Seq has the constructors Nil and Cons, the latter over an element of data sort A and a tail.

PARAMETER	DESCRIPTION
name	The signature's name. TYPE: str DEFAULT: 'Seq'
dim	The embedding dimension of every sort. TYPE: int DEFAULT: 64

RETURNS	DESCRIPTION
Signature	The sequence signature.

Source code in src/quivers/structural/shapes/seq.py

def seq_signature(name: str = "Seq", dim: int = 64) -> Signature:
    """Return a cons-list signature with a data element sort.

    The object sort ``Seq`` has the constructors ``Nil`` and ``Cons``,
    the latter over an element of data sort ``A`` and a tail.

    Parameters
    ----------
    name : str
        The signature's name.
    dim : int
        The embedding dimension of every sort.

    Returns
    -------
    Signature
        The sequence signature.
    """
    sorts = {
        "Seq": Sort(name="Seq", kind="object", dim=dim),
        "A": Sort(name="A", kind="data", dim=dim),
    }
    constructors = {
        "Nil": Constructor(name="Nil", domain=(), codomain="Seq"),
        "Cons": Constructor(
            name="Cons",
            domain=("A", "Seq"),
            codomain="Seq",
        ),
    }
    return Signature(
        name=name,
        sorts_t=tuple(sorts.values()),
        constructors_t=tuple(constructors.values()),
    )

rnn_encoder

rnn_encoder(sig: Signature | None = None, dim: int = 64) -> Encoder

A GRU-cell encoder: Cons(head, tail) updates a hidden state from the tail's compressed vector and the head's element embedding.

PARAMETER	DESCRIPTION
sig	A sequence signature with object sort Seq and data sort A; seq_signature at dimension dim when omitted. TYPE: Signature or None DEFAULT: None
dim	The embedding dimension of both sorts. TYPE: int DEFAULT: 64

RETURNS	DESCRIPTION
Encoder	The sequence encoder.

Source code in src/quivers/structural/shapes/seq.py

def rnn_encoder(sig: Signature | None = None, dim: int = 64) -> Encoder:
    """A GRU-cell encoder: ``Cons(head, tail)`` updates a hidden
    state from the tail's compressed vector and the head's element
    embedding.

    Parameters
    ----------
    sig : Signature or None
        A sequence signature with object sort ``Seq`` and data sort
        ``A``; `seq_signature` at dimension ``dim`` when omitted.
    dim : int
        The embedding dimension of both sorts.

    Returns
    -------
    Encoder
        The sequence encoder.
    """
    sig = sig or seq_signature(dim=dim)
    cell = nn.GRUCell(dim, dim)
    nil_mod, nil_const = _learnable_nil(dim)
    table, embedder = _data_embedder(dim)

    def cons_fn(a: torch.Tensor, tail: torch.Tensor) -> torch.Tensor:
        return cell(a.reshape(1, dim), tail.reshape(1, dim)).reshape(dim)

    def nil_fn() -> torch.Tensor:
        return nil_const

    op_fns = {
        "Nil": PerOpFn("Nil", "plain", (), nil_fn),
        "Cons": PerOpFn("Cons", "plain", ("head", "tail"), cons_fn),
    }
    return Encoder(
        name="RNN",
        signature=sig,
        sort_dims={"Seq": dim, "A": dim},
        op_fns=op_fns,
        var_init_fns={},
        data_embedders={"A": embedder},
        modules_owned=[cell, table, nil_mod],
    )

transformer_encoder

transformer_encoder(sig: Signature | None = None, dim: int = 64) -> Encoder

A transformer-style encoder: each Cons step combines the head's element embedding with the tail's running compressed vector through a learned MLP over the concatenation.

PARAMETER	DESCRIPTION
sig	A sequence signature with object sort Seq and data sort A; seq_signature at dimension dim when omitted. TYPE: Signature or None DEFAULT: None
dim	The embedding dimension of both sorts. TYPE: int DEFAULT: 64

RETURNS	DESCRIPTION
Encoder	The sequence encoder.

Source code in src/quivers/structural/shapes/seq.py

def transformer_encoder(sig: Signature | None = None, dim: int = 64) -> Encoder:
    """A transformer-style encoder: each ``Cons`` step combines
    the head's element embedding with the tail's running compressed
    vector through a learned MLP over the concatenation.

    Parameters
    ----------
    sig : Signature or None
        A sequence signature with object sort ``Seq`` and data sort
        ``A``; `seq_signature` at dimension ``dim`` when omitted.
    dim : int
        The embedding dimension of both sorts.

    Returns
    -------
    Encoder
        The sequence encoder.
    """
    sig = sig or seq_signature(dim=dim)
    head_proj = nn.Linear(dim, dim)
    tail_proj = nn.Linear(dim, dim)
    fuse = nn.Sequential(nn.Linear(2 * dim, dim), nn.GELU(), nn.Linear(dim, dim))
    nil_mod, nil_const = _learnable_nil(dim)
    table, embedder = _data_embedder(dim)

    def cons_fn(a: torch.Tensor, tail: torch.Tensor) -> torch.Tensor:
        h = head_proj(a)
        t = tail_proj(tail)
        cat = torch.cat([h, t], dim=-1)
        return fuse(cat)

    def nil_fn() -> torch.Tensor:
        return nil_const

    op_fns = {
        "Nil": PerOpFn("Nil", "plain", (), nil_fn),
        "Cons": PerOpFn("Cons", "plain", ("head", "tail"), cons_fn),
    }
    return Encoder(
        name="Tfm",
        signature=sig,
        sort_dims={"Seq": dim, "A": dim},
        op_fns=op_fns,
        var_init_fns={},
        data_embedders={"A": embedder},
        modules_owned=[head_proj, tail_proj, fuse, table, nil_mod],
    )

bow_encoder

bow_encoder(sig: Signature | None = None, dim: int = 64) -> Encoder

An order-independent sum-of-element-embeddings encoder.

Each Cons(head, tail) step adds the head's element embedding to the running tail embedding; Nil is a learnable zero. Equivalent to a bag-of-tokens fixed-length representation.

PARAMETER	DESCRIPTION
sig	A sequence signature with object sort Seq and data sort A; seq_signature at dimension dim when omitted. TYPE: Signature or None DEFAULT: None
dim	The embedding dimension of both sorts. TYPE: int DEFAULT: 64

RETURNS	DESCRIPTION
Encoder	The sequence encoder.

Source code in src/quivers/structural/shapes/seq.py

def bow_encoder(sig: Signature | None = None, dim: int = 64) -> Encoder:
    """An order-independent sum-of-element-embeddings encoder.

    Each ``Cons(head, tail)`` step adds the head's element embedding
    to the running tail embedding; ``Nil`` is a learnable zero.
    Equivalent to a bag-of-tokens fixed-length representation.

    Parameters
    ----------
    sig : Signature or None
        A sequence signature with object sort ``Seq`` and data sort
        ``A``; `seq_signature` at dimension ``dim`` when omitted.
    dim : int
        The embedding dimension of both sorts.

    Returns
    -------
    Encoder
        The sequence encoder.
    """
    sig = sig or seq_signature(dim=dim)
    nil_mod, nil_const = _learnable_nil(dim)
    table, embedder = _data_embedder(dim)

    def cons_fn(a: torch.Tensor, tail: torch.Tensor) -> torch.Tensor:
        return a + tail

    def nil_fn() -> torch.Tensor:
        return nil_const

    op_fns = {
        "Nil": PerOpFn("Nil", "plain", (), nil_fn),
        "Cons": PerOpFn("Cons", "plain", ("head", "tail"), cons_fn),
    }
    return Encoder(
        name="BoW",
        signature=sig,
        sort_dims={"Seq": dim, "A": dim},
        op_fns=op_fns,
        var_init_fns={},
        data_embedders={"A": embedder},
        modules_owned=[table, nil_mod],
    )

ar_decoder

ar_decoder(sig: Signature | None = None, dim: int = 64, vocab: list[DataLeaf] | None = None, depth: int = 64) -> Decoder

An autoregressive decoder over a sequence signature.

vocab is the closed token set the primitive head selects from. Required; the runtime raises on an empty vocabulary.

PARAMETER	DESCRIPTION
sig	A sequence signature with object sort Seq and data sort A; seq_signature at dimension dim when omitted. TYPE: Signature or None DEFAULT: None
dim	The embedding dimension of both sorts. TYPE: int DEFAULT: 64
vocab	The closed set of element values the decoder emits. TYPE: list[DataLeaf] or None DEFAULT: None
depth	The recursion-depth bound, which caps the sampled length. TYPE: int DEFAULT: 64

RETURNS	DESCRIPTION
Decoder	The sequence decoder.

RAISES	DESCRIPTION
ValueError	If vocab is missing or empty.

Source code in src/quivers/structural/shapes/seq.py

def ar_decoder(
    sig: Signature | None = None,
    dim: int = 64,
    vocab: list[DataLeaf] | None = None,
    depth: int = 64,
) -> Decoder:
    """An autoregressive decoder over a sequence signature.

    ``vocab`` is the closed token set the primitive head selects
    from. Required; the runtime raises on an empty vocabulary.

    Parameters
    ----------
    sig : Signature or None
        A sequence signature with object sort ``Seq`` and data sort
        ``A``; `seq_signature` at dimension ``dim`` when omitted.
    dim : int
        The embedding dimension of both sorts.
    vocab : list[DataLeaf] or None
        The closed set of element values the decoder emits.
    depth : int
        The recursion-depth bound, which caps the sampled length.

    Returns
    -------
    Decoder
        The sequence decoder.

    Raises
    ------
    ValueError
        If ``vocab`` is missing or empty.
    """
    if not vocab:
        raise ValueError("ar_decoder requires a non-empty vocabulary")
    sig = sig or seq_signature(dim=dim)
    # Two structural choices: Nil, Cons. (Seq has no binders and so
    # no BoundVar candidate is ever activated.)
    structure_head = nn.Linear(dim, 2)
    primitive = nn.Linear(dim, len(vocab))
    # Factor for Cons (arity 2 over sort Seq): split into head + tail.
    factor_lin = nn.Linear(dim, 2 * dim)
    # Bilinear binder-select. Seq itself never invokes this, but
    # signatures that compose Seq into a larger algebra (with a
    # binder reaching into Seq scope) do; the scorer is a real
    # learned bilinear key/query pair.
    bs_q = nn.Linear(dim, dim)
    bs_k = nn.Linear(dim, dim)

    def structure_fn(v: torch.Tensor) -> torch.Tensor:
        return structure_head(v.reshape(-1))

    def primitive_fn(v: torch.Tensor) -> torch.Tensor:
        return primitive(v.reshape(-1))

    def factor_2(v: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        out = factor_lin(v.reshape(-1))
        return out[:dim], out[dim:]

    def binder_select_fn(v: torch.Tensor, embeds: list[torch.Tensor]) -> torch.Tensor:
        qv = bs_q(v.reshape(-1))
        keys = torch.stack([bs_k(e.reshape(-1)) for e in embeds], dim=0)
        return keys @ qv

    return Decoder(
        name="AR",
        signature=sig,
        sort_dims={"Seq": dim, "A": dim},
        depth=depth,
        structure_fns={"Seq": structure_fn},
        primitive_fns={"A": primitive_fn},
        factor_fns={"Seq": {2: factor_2}},
        binder_select_fn=binder_select_fn,
        data_vocab={"A": vocab},
        modules_owned=[structure_head, primitive, factor_lin, bs_q, bs_k],
    )

list_to_term

list_to_term(elements: list[DataLeaf]) -> Term

Convert a Python list to a right-nested Cons(..., Nil) term.

PARAMETER	DESCRIPTION
elements	The sequence elements, in order. TYPE: list[DataLeaf]

RETURNS	DESCRIPTION
Term	The cons-list term over seq_signature's constructors.

Source code in src/quivers/structural/shapes/seq.py

def list_to_term(elements: list[DataLeaf]) -> Term:
    """Convert a Python list to a right-nested ``Cons(..., Nil)`` term.

    Parameters
    ----------
    elements : list[DataLeaf]
        The sequence elements, in order.

    Returns
    -------
    Term
        The cons-list term over `seq_signature`'s constructors.
    """
    term = Term(op="Nil", args=())
    for e in reversed(elements):
        term = Term(op="Cons", args=(e, term))
    return term

term_to_list

term_to_list(term: Term) -> list[DataLeaf]

Inverse of list_to_term.

PARAMETER	DESCRIPTION
term	A cons-list term. TYPE: Term

RETURNS	DESCRIPTION
list[DataLeaf]	The heads of the Cons chain, outermost first; the walk stops at the first term that is not a Cons.

Source code in src/quivers/structural/shapes/seq.py

def term_to_list(term: Term) -> list[DataLeaf]:
    """Inverse of `list_to_term`.

    Parameters
    ----------
    term : Term
        A cons-list term.

    Returns
    -------
    list[DataLeaf]
        The heads of the ``Cons`` chain, outermost first; the walk
        stops at the first term that is not a ``Cons``.
    """
    out: list[DataLeaf] = []
    cur = term
    while isinstance(cur, Term) and cur.op == "Cons":
        head, tail = cur.args
        out.append(head)
        cur = tail
    return out

tree

Tree-shaped signatures and canonical encoders / decoders.

Surface algebra::

Tree[L, B] = Leaf(L) | Node(B, Tree, Tree)

tree_signature

tree_signature(name: str = 'Tree', dim: int = 64) -> Signature

Return a binary-tree signature with leaf and branch data sorts.

The object sort Tree has two constructors: Leaf over the data sort L and Node over a branch label of data sort B and two subtrees.

PARAMETER	DESCRIPTION
name	The signature's name. TYPE: str DEFAULT: 'Tree'
dim	The embedding dimension of every sort. TYPE: int DEFAULT: 64

RETURNS	DESCRIPTION
Signature	The tree signature.

Source code in src/quivers/structural/shapes/tree.py

def tree_signature(name: str = "Tree", dim: int = 64) -> Signature:
    """Return a binary-tree signature with leaf and branch data sorts.

    The object sort ``Tree`` has two constructors: ``Leaf`` over the
    data sort ``L`` and ``Node`` over a branch label of data sort
    ``B`` and two subtrees.

    Parameters
    ----------
    name : str
        The signature's name.
    dim : int
        The embedding dimension of every sort.

    Returns
    -------
    Signature
        The tree signature.
    """
    sorts = {
        "Tree": Sort(name="Tree", kind="object", dim=dim),
        "L": Sort(name="L", kind="data", dim=dim),
        "B": Sort(name="B", kind="data", dim=dim),
    }
    constructors = {
        "Leaf": Constructor(name="Leaf", domain=("L",), codomain="Tree"),
        "Node": Constructor(
            name="Node",
            domain=("B", "Tree", "Tree"),
            codomain="Tree",
        ),
    }
    return Signature(
        name=name,
        sorts_t=tuple(sorts.values()),
        constructors_t=tuple(constructors.values()),
    )

tree_lstm_encoder

tree_lstm_encoder(sig: Signature | None = None, dim: int = 64) -> Encoder

A binary-tree LSTM encoder (Tai et al. 2015 child-sum variant).

A Leaf projects its token embedding. A Node gates between an MLP over its label and both children and the mean of the two children's embeddings.

PARAMETER	DESCRIPTION
sig	A tree signature with object sort Tree and data sorts L and B; tree_signature at dimension dim when omitted. TYPE: Signature or None DEFAULT: None
dim	The embedding dimension of every sort. TYPE: int DEFAULT: 64

RETURNS	DESCRIPTION
Encoder	The tree encoder.

Source code in src/quivers/structural/shapes/tree.py

def tree_lstm_encoder(sig: Signature | None = None, dim: int = 64) -> Encoder:
    """A binary-tree LSTM encoder (Tai et al. 2015 child-sum
    variant).

    A ``Leaf`` projects its token embedding. A ``Node`` gates between
    an MLP over its label and both children and the mean of the two
    children's embeddings.

    Parameters
    ----------
    sig : Signature or None
        A tree signature with object sort ``Tree`` and data sorts
        ``L`` and ``B``; `tree_signature` at dimension ``dim`` when
        omitted.
    dim : int
        The embedding dimension of every sort.

    Returns
    -------
    Encoder
        The tree encoder.
    """
    sig = sig or tree_signature(dim=dim)

    leaf_proj = nn.Linear(dim, dim)
    label_proj = nn.Linear(dim, dim)
    gate_mlp = nn.Sequential(
        nn.Linear(3 * dim, dim),
        nn.Tanh(),
        nn.Linear(dim, dim),
    )
    out_mlp = nn.Sequential(
        nn.Linear(3 * dim, dim),
        nn.Tanh(),
        nn.Linear(dim, dim),
    )
    leaf_tbl, leaf_embed = _data_embedder(dim)
    label_tbl, label_embed = _data_embedder(dim)

    def leaf_fn(token: torch.Tensor) -> torch.Tensor:
        return leaf_proj(token)

    def node_fn(
        label: torch.Tensor, left: torch.Tensor, right: torch.Tensor
    ) -> torch.Tensor:
        b = label_proj(label)
        cat = torch.cat([b, left, right], dim=-1)
        gate = torch.sigmoid(gate_mlp(cat))
        h = torch.tanh(out_mlp(cat))
        return gate * h + (1 - gate) * (0.5 * (left + right))

    op_fns = {
        "Leaf": PerOpFn("Leaf", "plain", ("t",), leaf_fn),
        "Node": PerOpFn("Node", "plain", ("b", "l", "r"), node_fn),
    }
    return Encoder(
        name="TreeLSTM",
        signature=sig,
        sort_dims={"Tree": dim, "L": dim, "B": dim},
        op_fns=op_fns,
        var_init_fns={},
        data_embedders={"L": leaf_embed, "B": label_embed},
        modules_owned=[leaf_proj, label_proj, gate_mlp, out_mlp, leaf_tbl, label_tbl],
    )

tree_decoder

tree_decoder(sig: Signature | None = None, dim: int = 64, leaf_vocab: list[DataLeaf] | None = None, label_vocab: list[DataLeaf] | None = None, depth: int = 12) -> Decoder

Top-down structural decoder over a tree signature.

PARAMETER	DESCRIPTION
sig	A tree signature with object sort Tree and data sorts L and B; tree_signature at dimension dim when omitted. TYPE: Signature or None DEFAULT: None
dim	The embedding dimension of every sort. TYPE: int DEFAULT: 64
leaf_vocab	The closed set of leaf values the decoder emits. TYPE: list[DataLeaf] or None DEFAULT: None
label_vocab	The closed set of branch labels the decoder emits. TYPE: list[DataLeaf] or None DEFAULT: None
depth	The recursion-depth bound on sampled trees. TYPE: int DEFAULT: 12

RETURNS	DESCRIPTION
Decoder	The tree decoder.

RAISES	DESCRIPTION
ValueError	If leaf_vocab or label_vocab is missing or empty.

Source code in src/quivers/structural/shapes/tree.py

def tree_decoder(
    sig: Signature | None = None,
    dim: int = 64,
    leaf_vocab: list[DataLeaf] | None = None,
    label_vocab: list[DataLeaf] | None = None,
    depth: int = 12,
) -> Decoder:
    """Top-down structural decoder over a tree signature.

    Parameters
    ----------
    sig : Signature or None
        A tree signature with object sort ``Tree`` and data sorts
        ``L`` and ``B``; `tree_signature` at dimension ``dim`` when
        omitted.
    dim : int
        The embedding dimension of every sort.
    leaf_vocab : list[DataLeaf] or None
        The closed set of leaf values the decoder emits.
    label_vocab : list[DataLeaf] or None
        The closed set of branch labels the decoder emits.
    depth : int
        The recursion-depth bound on sampled trees.

    Returns
    -------
    Decoder
        The tree decoder.

    Raises
    ------
    ValueError
        If ``leaf_vocab`` or ``label_vocab`` is missing or empty.
    """
    if not leaf_vocab or not label_vocab:
        raise ValueError("tree_decoder requires non-empty leaf_vocab and label_vocab")
    sig = sig or tree_signature(dim=dim)
    structure = nn.Linear(dim, 3)  # Leaf / Node / BoundVar (unused)
    leaf_head = nn.Linear(dim, len(leaf_vocab))
    label_head = nn.Linear(dim, len(label_vocab))
    factor_1 = nn.Linear(dim, dim)
    factor_3 = nn.Linear(dim, 3 * dim)
    bs_q = nn.Linear(dim, dim)
    bs_k = nn.Linear(dim, dim)

    def structure_fn(v: torch.Tensor) -> torch.Tensor:
        return structure(v.reshape(-1))

    def f1(v: torch.Tensor) -> tuple[torch.Tensor]:
        return (factor_1(v.reshape(-1)),)

    def f3(v: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        out = factor_3(v.reshape(-1))
        return out[:dim], out[dim : 2 * dim], out[2 * dim :]

    def binder_select_fn(v: torch.Tensor, embeds: list[torch.Tensor]) -> torch.Tensor:
        qv = bs_q(v.reshape(-1))
        keys = torch.stack([bs_k(e.reshape(-1)) for e in embeds], dim=0)
        return keys @ qv

    return Decoder(
        name="TreeAR",
        signature=sig,
        sort_dims={"Tree": dim, "L": dim, "B": dim},
        depth=depth,
        structure_fns={"Tree": structure_fn},
        primitive_fns={
            "L": lambda v, h=leaf_head: h(v.reshape(-1)),
            "B": lambda v, h=label_head: h(v.reshape(-1)),
        },
        factor_fns={"Tree": {1: f1, 3: f3}},
        binder_select_fn=binder_select_fn,
        data_vocab={"L": leaf_vocab, "B": label_vocab},
        modules_owned=[
            structure,
            leaf_head,
            label_head,
            factor_1,
            factor_3,
            bs_q,
            bs_k,
        ],
    )

graph

Graph-shaped signatures and canonical encoder.

A graph signature is declared with vertex kinds and edge kinds; the canonical GNN encoder runs iterations rounds of message passing with per-edge-kind message functions and per-vertex-kind update functions.

graph_signature

graph_signature(name: str, vertex_kinds: dict[str, int], edge_kinds: dict[str, tuple[str, str, bool]]) -> Signature

Construct a graph signature.

PARAMETER	DESCRIPTION
name	The signature's name. TYPE: str
vertex_kinds	Each vertex kind's embedding dimension, keyed by kind name. Every vertex kind is data-kinded. TYPE: dict[str, int]
edge_kinds	Each edge kind's (src_kind, tgt_kind, directed), keyed by kind name. TYPE: dict[str, tuple[str, str, bool]]

RETURNS	DESCRIPTION
Signature	The graph signature.

Source code in src/quivers/structural/shapes/graph.py

def graph_signature(
    name: str,
    vertex_kinds: dict[str, int],
    edge_kinds: dict[str, tuple[str, str, bool]],
) -> Signature:
    """Construct a graph signature.

    Parameters
    ----------
    name : str
        The signature's name.
    vertex_kinds : dict[str, int]
        Each vertex kind's embedding dimension, keyed by kind name.
        Every vertex kind is data-kinded.
    edge_kinds : dict[str, tuple[str, str, bool]]
        Each edge kind's ``(src_kind, tgt_kind, directed)``, keyed by
        kind name.

    Returns
    -------
    Signature
        The graph signature.
    """
    vk = {n: VertexKind(name=n, kind="data", dim=d) for n, d in vertex_kinds.items()}
    ek = {
        n: EdgeKind(name=n, src=s, tgt=t, directed=directed)
        for n, (s, t, directed) in edge_kinds.items()
    }
    return Signature(
        name=name, vertex_kinds_t=tuple(vk.values()), edge_kinds_t=tuple(ek.values())
    )

gnn_encoder

gnn_encoder(sig: Signature, iterations: int = 4, dim: int = 64, readout: str = 'mean') -> Encoder

A GNN encoder: per-edge-kind message MLP, per-vertex-kind GRU update, mean / sum / max readout.

Each vertex kind embeds its payload through a lazily grown table of learned vectors keyed by the payload's string form.

PARAMETER	DESCRIPTION
sig	A graph signature, as built by graph_signature. TYPE: Signature
iterations	The number of message-passing rounds. TYPE: int DEFAULT: 4
dim	The embedding dimension of a vertex kind that declares none. TYPE: int DEFAULT: 64
readout	How the final vertex embeddings are reduced: "mean", "sum", or "max". TYPE: str DEFAULT: 'mean'

RETURNS	DESCRIPTION
Encoder	The graph encoder; call its Encoder.forward_graph.

RAISES	DESCRIPTION
ValueError	Raised by the encoder's readout when readout is not one of the supported reductions.

Source code in src/quivers/structural/shapes/graph.py

def gnn_encoder(
    sig: Signature,
    iterations: int = 4,
    dim: int = 64,
    readout: str = "mean",
) -> Encoder:
    """A GNN encoder: per-edge-kind message MLP, per-vertex-kind
    GRU update, mean / sum / max readout.

    Each vertex kind embeds its payload through a lazily grown table
    of learned vectors keyed by the payload's string form.

    Parameters
    ----------
    sig : Signature
        A graph signature, as built by `graph_signature`.
    iterations : int
        The number of message-passing rounds.
    dim : int
        The embedding dimension of a vertex kind that declares none.
    readout : str
        How the final vertex embeddings are reduced: ``"mean"``,
        ``"sum"``, or ``"max"``.

    Returns
    -------
    Encoder
        The graph encoder; call its `Encoder.forward_graph`.

    Raises
    ------
    ValueError
        Raised by the encoder's readout when ``readout`` is not one of
        the supported reductions.
    """
    init_fns: dict[str, Callable[[DataLeaf], torch.Tensor]] = {}
    message_fns: dict[str, Callable[[torch.Tensor, torch.Tensor], torch.Tensor]] = {}
    update_fns: dict[str, Callable[[torch.Tensor, torch.Tensor], torch.Tensor]] = {}
    modules_owned: list[nn.Module] = []

    for vname, v in sig.vertex_kinds.items():
        vdim = v.dim or dim
        table = nn.ParameterDict()
        modules_owned.append(table)

        def make_init(table=table, vdim=vdim):
            def init(payload):
                key = str(payload).replace(".", "_")
                if key not in table:
                    table[key] = nn.Parameter(torch.randn(vdim) * 0.1)
                return table[key]

            return init

        init_fns[vname] = make_init()

        gru = nn.GRUCell(vdim, vdim)
        modules_owned.append(gru)

        def make_update(gru=gru, vdim=vdim):
            def upd(self_e, msg):
                return gru(msg.reshape(1, vdim), self_e.reshape(1, vdim)).reshape(vdim)

            return upd

        update_fns[vname] = make_update()

    for ename, e in sig.edge_kinds.items():
        src_dim = sig.vertex_kinds[e.src].dim or dim
        tgt_dim = sig.vertex_kinds[e.tgt].dim or dim
        out_dim = tgt_dim
        mlp = nn.Sequential(
            nn.Linear(src_dim + tgt_dim, max(out_dim, 64)),
            nn.Tanh(),
            nn.Linear(max(out_dim, 64), out_dim),
        )
        modules_owned.append(mlp)

        def make_msg(mlp=mlp, src_dim=src_dim, tgt_dim=tgt_dim):
            def msg(s, t):
                cat = torch.cat([s.reshape(-1), t.reshape(-1)], dim=0)
                return mlp(cat)

            return msg

        message_fns[ename] = make_msg()

    def readout_fn(embeds, mode=readout):
        stack = torch.stack(embeds, dim=0)
        if mode == "mean":
            return stack.mean(dim=0)
        if mode == "sum":
            return stack.sum(dim=0)
        if mode == "max":
            return stack.max(dim=0).values
        raise ValueError(f"unknown readout {mode!r}")

    sort_dims = {n: (v.dim or dim) for n, v in sig.vertex_kinds.items()}

    return Encoder(
        name="GNN",
        signature=sig,
        sort_dims=sort_dims,
        op_fns={},
        var_init_fns={},
        data_embedders={},
        modules_owned=modules_owned,
        iterations=iterations,
        init_fns=init_fns,
        message_fns=message_fns,
        update_fns=update_fns,
        readout=readout_fn,
    )