Wiring

Einsum wiring rules that generalize binary composition to contractions over any number of tensors under a composition rule.

wiring

Operadic n-ary tensor contractions for V-enriched categories.

A CompositionRule defines binary composition: (f >> g)[i, k] = ⋁_j f[i, j] ⊗ g[j, k]. That covers the 2-cells of a 2-category. Useful but narrow — many models combine three or more tensors at a shared reduction (e.g. two argument tensors plus a kernel folded at a single contraction):

out[s, c] = ⋁_{p, q} (
    arg1[s, p] ⊗ arg2[s, q] ⊗ kernel[p, q, c]
)

This isn't expressible as a chain of binary >>: it's an n-input wiring diagram whose centre contracts two axes at once under a fixed tensor_op and join. The general framework is operads / wiring diagrams (Spivak 2014); this module exposes a Python surface for them.

The unifying construct is WiringRule: an n-ary generalization of binary composition that pairs a CompositionRule (supplying tensor_op and join) with an einsum-style wiring spec naming which input axes are contracted vs preserved.

Categorically a WiringRule is a single operation of a colored operad whose colors are the V-Cat object types; the CompositionRule supplies the operad's enriching algebra.

WiringRule

Bases: ABC

An n-ary morphism contraction under a fixed CompositionRule.

Generalizes binary composition >> to arbitrary-arity tensor networks. A WiringRule carries:

• a composition_rule whose tensor_op and join drive the contraction;
• a wiring specification that names which axes of which inputs participate in which contractions, and which axes survive to the output.

Subclasses implement apply to perform the contraction given a list of input tensors.

composition_rule abstractmethod property

composition_rule: CompositionRule

The underlying binary composition rule used as the enriching algebra for this n-ary contraction.

input_arity abstractmethod property

input_arity: int

Number of input tensors this wiring expects.

apply abstractmethod

apply(*tensors: Tensor) -> Tensor

Run the contraction on the given input tensors.

Source code in src/quivers/core/wiring.py

@abstractmethod
def apply(self, *tensors: torch.Tensor) -> torch.Tensor:
    """Run the contraction on the given input tensors."""

EinsumWiring

EinsumWiring(composition_rule: CompositionRule, spec: str)

Bases: WiringRule

A multi-input contraction specified by an einsum-style string.

Spec syntax::

"input_1, input_2, ..., input_n -> output"

where each input_i and output is a sequence of axis letters identifying which dimensions of that tensor align. Axes that appear in some input but not in the output are contracted (reduced via the composition rule's join); axes that appear in the output are preserved (broadcast as needed).

Examples::

EinsumWiring(product_fuzzy, "ij, jk -> ik")
    # binary composition: contract over j, preserve (i, k)
EinsumWiring(product_fuzzy, "sp, sq, pqo -> so")
    # ternary contraction: two argument tensors over
    # (s, *) and a kernel contracted over (p, q),
    # preserving (s, o)

The contraction algorithm:

1. Determine the universe of axis letters from inputs and output.
2. Broadcast each input tensor against the full axis universe (singleton dims for axes the input doesn't carry).
3. Fold all broadcast inputs with tensor_op (left fold).
4. Reduce over the contracted axes using join.
5. Permute to the output's letter order.

Steps 3-4 correspond to the categorical content; steps 1, 2, and 5 are the einsum bookkeeping that places axes in a common shape so the binary tensor_op can fold over the n inputs.

Source code in src/quivers/core/wiring.py

def __init__(self, composition_rule: CompositionRule, spec: str) -> None:
    if not isinstance(composition_rule, CompositionRule):
        raise TypeError(
            f"EinsumWiring: composition_rule must be a CompositionRule, "
            f"got {type(composition_rule).__name__}"
        )
    if "->" not in spec:
        raise ValueError(f"EinsumWiring: spec must contain '->', got {spec!r}")
    lhs, rhs = spec.split("->", maxsplit=1)
    input_specs = tuple(s.strip() for s in lhs.split(","))
    output_spec = rhs.strip()
    # Validate that every input is non-empty and uses distinct
    # letters.  Duplicate letters within one input express a
    # trace / diagonal contraction that this implementation
    # does not realise; refusing them at construction time
    # surfaces the limitation honestly rather than producing
    # wrong shapes at apply time.
    for i, ispec in enumerate(input_specs):
        if not ispec:
            raise ValueError(f"EinsumWiring: input {i} of spec {spec!r} is empty")
        if len(set(ispec)) != len(ispec):
            raise ValueError(
                f"EinsumWiring: input {i} of spec {spec!r} repeats "
                f"an axis letter ({ispec!r}); diagonal/trace "
                f"contractions are not supported by this wiring"
            )
    # Output letters must also be distinct (a repeated output
    # letter would require axis-duplication which einsum does
    # not realise).
    if len(set(output_spec)) != len(output_spec):
        raise ValueError(
            f"EinsumWiring: output {output_spec!r} of spec "
            f"{spec!r} repeats an axis letter"
        )
    # Build the axis universe (preserving first-mention order).
    seen: dict[str, None] = {}
    for ispec in input_specs:
        for ch in ispec:
            seen.setdefault(ch, None)
    for ch in output_spec:
        if ch not in seen:
            raise ValueError(
                f"EinsumWiring: output letter {ch!r} not present "
                f"in any input of spec {spec!r}"
            )
    self._rule = composition_rule
    self._spec = spec
    self._input_specs = input_specs
    self._output_spec = output_spec
    self._axis_universe: tuple[str, ...] = tuple(seen)
    # Indexed mapping for fast access.
    self._axis_index = {ch: i for i, ch in enumerate(self._axis_universe)}
    # Axes that get reduced (in universe but not in output).
    out_set = set(output_spec)
    self._reduce_axes: tuple[int, ...] = tuple(
        i for i, ch in enumerate(self._axis_universe) if ch not in out_set
    )

einsum_wiring

einsum_wiring(composition_rule: CompositionRule, spec: str) -> EinsumWiring

Convenience constructor for EinsumWiring.

PARAMETER	DESCRIPTION
composition_rule	The rule supplying tensor_op and join. TYPE: CompositionRule
spec	An einsum specification such as "ij,jk->ik". TYPE: str

RETURNS	DESCRIPTION
EinsumWiring	The wiring rule for spec.

RAISES	DESCRIPTION
TypeError	If composition_rule is not a CompositionRule.
ValueError	If spec is malformed.

Source code in src/quivers/core/wiring.py

def einsum_wiring(composition_rule: CompositionRule, spec: str) -> EinsumWiring:
    """Convenience constructor for `EinsumWiring`.

    Parameters
    ----------
    composition_rule : CompositionRule
        The rule supplying ``tensor_op`` and ``join``.
    spec : str
        An einsum specification such as ``"ij,jk->ik"``.

    Returns
    -------
    EinsumWiring
        The wiring rule for ``spec``.

    Raises
    ------
    TypeError
        If ``composition_rule`` is not a `CompositionRule`.
    ValueError
        If ``spec`` is malformed.
    """
    return EinsumWiring(composition_rule, spec)

contract

contract(rule: WiringRule, *tensors: Tensor) -> Tensor

Apply a WiringRule to a sequence of input tensors.

Thin wrapper around WiringRule.apply for call-site readability::

result = contract(my_rule, arg1, arg2, kernel)

is an alternative to the method-call form.

PARAMETER	DESCRIPTION
rule	The wiring rule to apply. TYPE: WiringRule
*tensors	The input tensors, in the order the rule expects. TYPE: Tensor DEFAULT: ()

RETURNS	DESCRIPTION
Tensor	The contracted tensor.

Source code in src/quivers/core/wiring.py

def contract(
    rule: WiringRule,
    *tensors: torch.Tensor,
) -> torch.Tensor:
    """Apply a `WiringRule` to a sequence of input tensors.

    Thin wrapper around `WiringRule.apply` for call-site
    readability::

        result = contract(my_rule, arg1, arg2, kernel)

    is an alternative to the method-call form.

    Parameters
    ----------
    rule : WiringRule
        The wiring rule to apply.
    *tensors : torch.Tensor
        The input tensors, in the order the rule expects.

    Returns
    -------
    torch.Tensor
        The contracted tensor.
    """
    return rule.apply(*tensors)