Graphical einops:
bridging tensor networks and computation graphs
May 29, 2026
Abstract
Architecture diagrams are ubiquitous in deep learning, but they are usually only representational: the tensor-program identities they suggest are still proved by prose and tensor-axis manipulation. We introduce a formal graphical calculus for the
structural fragment of tensor programming underlying einops, making such diagrams proof-enabling. Our calculus represents tensor axes as nested graded tubes around a base type. The tube boundary recovers the undirected tensor-network view of
axes, while the directed interior retains the operational reading of computation graphs. The key rewrite is grade-naturality: sliding spectacles over tubes. Standard equivariance proofs become short diagrammatic derivations. We additionally demonstrate how
our rewrite system may be applied to convert attention masks into pre-processing operations, recovering efficient implementations of sparse attention blocks.
1 Introduction↩︎
Tensor-program identities — “transpose then reshape equals reshape then transpose”, “permuting an attention layer’s queries permutes its outputs”, “masked attention equals augmented unmasked attention” — are routine in deep learning but proved largely
by hand: index manipulation, prose, or appeal to operational intuition. The de facto readable syntax for the structural layer (axis reshapes, broadcasts, contractions, packing) is einops: a procedural surface that disambiguates index
choreography but supplies no proof system. We give that layer a graphical proof calculus, where einops-level diagrams are typed proof terms, the rewrite engine is a single naturality square, and structural equality is decidable.
Figure 1: 
.
The reason such a calculus has not appeared earlier is mathematical, not stylistic. Tensor data has an undirected, topological syntax, whereas computation has a directed left-to-right flow. The two views correspond to graphical traditions that disagree on what parallel composition should mean: \(\otimes\) or \(\oplus\).
The undirected tradition (Penrose tensor diagrams [1], named-tensor notation [2], and ZX-calculus and its quantum-process descendants [3]) treats wires as tensor indices and reads diagrams up to free deformation. Pavlovic, Coecke, and Vicary [4] show that the underlying primitive on every wire type is a special dagger Frobenius algebra (SDFA), and that an SDFA structure on a Hilbert space is equivalent to a choice of orthonormal basis. Among other consequences,
parallel composition of diagrams is forced to be the Kronecker product of vector spaces \(\otimes\), and the rewrites available bind to that interpretation. Operational structure (concatenation, sequence axes, batching)
sits awkwardly here: named-tensor papers explicitly mark pack/unpack-like operations as out of scope, and nonlinear functions are difficult to address formally. Recent DL-flavoured offerings adapt undirected tensor-network diagrams for ML audiences.
Taylor’s expository diagrams [5] and Ahle’s tensorgrad [6] are notable examples; the latter performs the simplifications of the Tensor Cookbook symbolically and lowers diagrams to PyTorch with full higher-order-derivative support. Both treat
the diagram chiefly as a frontend for symbolic manipulation and compiler lowering rather than as a proof calculus. At an adjacent abstraction layer, [7] give a formal tensor calculus for evaluating and differentiating Einstein-notation contractions: a calculus about derivatives of contractions rather than about the index-and-axis structure of the surrounding tensor
program.
The directed tradition (computation graphs, architecture flowcharts) imposes a temporal reading order: operations are boxes, data are wires connecting past outputs to sequentially-composed future inputs. When all such directed data admit copying and deleting maps, Fox’s theorem [8] forces the meaning of parallel composition to be the categorical product: when we deal with functions between Euclidean spaces in the setting of Deep Learning, parallel composition must mean direct sum. Recent DL-side graphical work in this tradition formalises the layout with explicit SIMD-style boxes for parallel operations along tensor axes: neural circuit diagrams [9] and the “Anatomy of Attention” framework [10] are the closest precedents to our calculus and inspired the wires-and-tubes layer used here.
einops [11] is the practitioner-side bridge between the two views. It is a procedural language for tensor
manipulation (rearrange, repeat, reduce, pack, unpack) that sits in the directed tradition, so by Fox its data carries copy and delete. Via the display-map correspondence, each
structural move determines a finite read map; reductions additionally require a chosen algebra on values. We formalise this bridge as a graphical calculus by wrapping wires in tubes that carry dimension data; nested tubes give multi-axis tensors,
partial-function reads on a dense surface handle pack/unpack, and the upper boundary of every tube recovers the undirected tensor-network reading of its axes (2.2). The paper’s theorem bundle is: (A) structural
einops programs embed in a \(\mathsf{FinSet}^{\mathrm{op}}\)-graded tube calculus; (B) equality in the structural fragment reduces to equality of composite functions between finite sets; (C) base morphisms
acting pointwise along a grade satisfy a grade-naturality rewrite, sliding spectacles; and (D) every attention mask induces a query-indexed display of visible K/V fibres, so masked attention is fibrewise unmasked attention over that display.
Causal masking, padding masks, and sliding-window masks are special cases.
5pt
\(^a\)Special-Frobenius copy on the SDFA fragment, distinct from cartesian copy. \(^b\)Topological deformation plus Frobenius/special-Frobenius rewrites. \(^c\)Marked out-of-scope by the named-tensor authors. \(^d\)Symbolic simplification within the tensor-cookbook fragment, not a global decision procedure. \(^e\)Procedural DSL, not a graphical formalism. \(^f\)Decidable on the structural fragment (1); reductions and softmax interact with declared base algebras and lie outside this guarantee.
1 situates our work among neighbouring approaches. Unlike prior diagram systems for machine learning, which primarily serve as notation, exposition, or compiler frontends, this calculus provides a formal equality theory for the einops-level structural fragment, together with a normalisation procedure.
2 Background: building the calculus↩︎
The directed flowchart tradition reads a diagram temporally: operations are boxes, data are wires, and a diagram drawn left-to-right is read in evaluation order. Informal flowcharts in DL papers often differ up to a Poincaré dual (treating data as boxes
and processes as wires) or mix conventions altogether. The string-diagram convention is the formally sound version: Joyal and Street [12]
prove that the desired topological manipulations (sliding boxes along wires, bending wires, redrawing without changing meaning) agree exactly with the algebraic semantics of symmetric monoidal categories. The interchange law (right) is the headline
example: the four parallel-and-sequential arrangements of two boxes \(f, g, h, k\) all compute the same morphism. This section walks pedagogically from string-diagram primitives to a calculus that captures
einops. The path: a wire denotes \(\mathbb{R}\); copy, delete, swap generate every (partial) function on indices via the display-map correspondence; tubes wrap parallel wires; nesting handles multi-axis
tensors; reshape, rearrange, cat, split are index bijections; pack/unpack reduces to padding; repeat, reduce, dot product, and einsum
follow. The summary table is 2.
2.1 From maps to diagrams↩︎
We start with a single wire type. A black wire denotes the field \(\mathbb{R}\); \(n\) wires in parallel denote \(\mathbb{R}^n\). Three structural maps recur in tensor work: swap (\(\mathbb{R}^2\to\mathbb{R}^2\), \((a,b)\mapsto(b,a)\), transpose two indices), copy (\(\mathbb{R}\to\mathbb{R}^2\), \(a\mapsto(a,a)\), broadcast), and delete (\(\mathbb{R}\to\mathbf{1}\), \(a\mapsto\star\), discard).
The diagrammatic forms of the comonoid axioms say that copy is coassociative (order-independent), cocommutative (branch-order-independent), and counital (unital with respect to delete). All functions are naturally copyable and deletable. These axioms are the entire algebraic content of working with copy and delete, and they are exactly the topological identities that you would intuitively expect of a fork-and-discard structure. The choice to interpret parallel composition as direct sum is what gives us copy and delete in the first place: Fox’s theorem [8] shows that a symmetric monoidal category in which the tensor product is the categorical product is exactly one in which every object carries a natural comonoid (copy, delete) compatible with the structural maps.
Display-map correspondence. Tag each input wire with an index, follow the diagram, and read off where each output index lands, going backwards. Swap alone gives bijections; adding delete gives injections; adding copy gives every total function between finite sets. Allowing a wire to be padded (some output positions have no source and are filled by an external value) extends the correspondence to every partial function. We return to padding once ragged data is on the table.
2.2 Tubes wrap parallel wires↩︎
Drawing \(\mathbb{R}^{37}\) as 37 parallel wires would be cumbersome. We compress the bundle into a single coloured tube of grade \(37\) carrying one wire; the tube means direct-sum my contents 37 times. Two helpful conventions: a tube of grade 1 is invisible (it pops in and out of existence freely), and a tube of grade 0 does not appear at all (neither the tube nor its contents). \[\scalebox{0.9}{\tikzfig{matmuls/pita}}\quad\scalebox{0.9}{\tikzfig{matmuls/pitamul}}\quad\scalebox{0.9}{\tikzfig{matmuls/SIMD}}\] (left) \(k\) parallel wires compress into one tube of grade \(k\). (middle) A matrix-shaped bundle of operations from \(n\) wires to \(m\) wires becomes a single tube-typed box. (right) A base operation applied independently to every wire in a bundle becomes one box inside the tube: this is the SIMD pattern.
Highlighting the upper boundary. We will keep the upper boundary of every tube visible. The reason is that highlighting just the upper boundary recovers the undirected tensor-network drawing of the same axes: data moves along wires
inside the tube (operational reading), while the boundary records only the indexed type with no temporal ordering (tensor-network reading). einops chooses a directional presentation that coheres with the operational view, but the undirected
tensor-network reading remains visible in the same picture as the tube boundary. The two traditions of 1 therefore live alongside each other in a single drawing, sharing geometry but not algebraic structure.
Nested tubes for multi-axis tensors. Iterate the wrapping. A tensor \(\mathbf{X}\) of shape \([m]\times[n]\times[p]\) becomes an \(\textcolor{orange}{m}\)-tube around an \(\textcolor{blue}{n}\)-tube around a \(\textcolor{red}{p}\)-tube; each axis is one layer of nesting, with the outermost tube as the topmost iterator. Tube colours read inside-out as the axis hierarchy.
2.3 Reshape, reorder, cat, split↩︎
We are now in einops territory. Four operations are common in practice: reshape and reorder, cat (concatenation), and split (slicing). All four share the property that they are bijections
between indices. By the display-map correspondence, every such bijection is realisable as a wire-swapping braid of black \(\mathbb{R}\) wires: nothing happens to the data, only the metadata commitment about how the wires
are bundled changes.
Reorder swaps two tensor indices: a generalised transposition. In our calculus this is a braid that swaps the order of two tubes via their boundaries. The data wires inside are unaffected; the boundary tells the reader which axis is
which.
Reshape can change the number of nesting layers but is otherwise an identity process. The cardinality isomorphism \([mn]\cong[m]\times[n]\) rewires data between a flat \(mn\)-tube and a nested \([m]\)-outer, \([n]\)-inner pair: a tubing can split its integer grade among child tubes, or gather child grades into a parent. Reshape
acts on a single axis at a time, with the other dimensions matching: the standard axis constraint of einops.rearrange.
cat and split are concatenation and slicing along an axis. Categorically, they are ordered-sum operations on the grade, not tensor contractions on the data. Given segment metadata, cat embeds two consecutive
index blocks into one longer ordered block, and split recovers the two summands. Diagrammatically these are pants (cat) and copants (split). The pants/copants pair is a strict inverse only relative to the retained segment lengths.
2.4 Repeat, reduce, dot product, and matrix multiplication↩︎
We can now move to operations that touch the data, building up to einsum in steps.
repeat is a sock. The simplest content-bearing move is einops.repeat: copy a (possibly tubed) wiring a number of times. Diagrammatically this is a sock that wraps the input in an outer tube of grade equal to
the repeat count. On the index side, repeat corresponds to delete: the index map for a repeat operation forgets the new outer grade, since every output position reads the same input cell.
reduce is an anti-sock. einops.reduce is the converse of repeat: it sheds an outermost tubing layer via an algebra on the base, most commonly \(\mathrm{sum}\), \(\max\), or \(\mathrm{mean}\). The diagrammatic dual of the broadcasting sock is a sealing cap on the outermost tube, which the algebra computes the value of.
Dot product. The dot product of two vectors is built from gadgets we already have: zip them together with pants, SIMD elementwise multiply, then reduce with sum. This is the canonical example of a tube “closing”: two tubes of grade \(n\) enter, one wire of grade \(1\) exits.
Tensor contraction. A general tensor contraction \(\sum_n A_{m,n} B_{k,n}\) is built from the dot-product gadget by tubing-up the non-contracted axes. We use repeat-socks to broadcast each operand to the
joint axis layout, pants to pair the two contracted tubes, then a SIMDed dot product. This is einsum pictorially.
Corollary 1 (Matrix multiplication). Matrix multiplication \(AB^\top\) at shape \(m\times k\) is the special case of tensor contraction with two two-axis operands sharing one axis. The repeat-socks collapse to a single broadcast on each side, and the calculus simplifies the contraction to the standard form \(C_{i,k} = \sum_n A_{i,n}\, B_{k,n}\).
A visual simplification of 1 is provable inside the calculus. Since broadcast (sock) corresponds to index-deletion and cat (pants) to
index-copy, counitality simplifies the diagram, exposing which index is contracted. With no nontrivial nesting or direct-summation interaction, we lose nothing by depicting only the topmost boundary and treating data wires as separating guardrails. \[\tikzfig{ML/towards-matrix-contraction/simplerform}\]
2.5 (un)/pack: ragged data via padding↩︎
For ragged tensors, einops provides pack and unpack. We have two presentations: we can (left) decorate tubes with packing metadata (the bracket-label refinement of 9), or we can (right) cast every ragged tensor as a dense tensor with padding elements and recover the original data by partial reads. The latter is closer to practice and keeps tube grades as positive
integers, so we adopt it here. \[\tikzfig{raggedoptions}\]
Padding extends display maps to partial functions. A wire that has been padded gives a partial function on indices: each output cell either reads from one input cell, or, if it falls in a padded position, requests an external filler. The filler is not determined by the partial function itself. It belongs to a padding interpretation for the relevant base object or to the operation that consumes the padded wire.
(un)pad with an arbitrary \(\star\) requires bookkeeping for how the filler propagates through subsequent operations. In practice we choose \(-\infty\), \(0\), or \(1\), since they have neutral or absorbing behaviour under \(+\), \(\times\), and \(\sigma\).
Addition absorbs \(0\) paddings. Because \(0\) is the unit of the addition semiring, \(0\)-padding is visually absorbed by summation algebras.
Multiplication absorbs \(1\) paddings. Because \(1\) is the unit of the multiplication semiring, \(1\)-padding is visually absorbed by products.
Softmax transforms -\(\infty\) padding into \(0\) padding. Softmax (whether normalising or not) is an \(k \in \mathbb{N}\)-indexed family of operations \(\sigma_k: \mathbb{R}^k \rightarrow \mathbb{R}^k\) with the property that a -\(\infty\) logit contributes \(0\) output, since \(e^{-\infty} = 0\). Hence -\(\infty\) paddings can slide through softmax, turning into \(0\).
Proposition 1 (pack/unpack is a special case of pad/unpad). For ragged components \(\{X_i\}\) in any dense bounding shape \(B\), \[\texttt{pad}\colon \bigsqcup_i X_i \hookrightarrow B,\qquad \texttt{unpad}\colon B \rightharpoonup \bigsqcup_i X_i,\qquad \texttt{unpad}\circ \texttt{pad} \;=\; \mathrm{id}.\] Ordinary
pack/unpack is the case \(B = [\sum_i |X_i|]\) with contiguous concatenation; non-contiguous variants are captured by the metadata refinement of 9.
Sketch. By construction: pad places each valid cell at its embedded coordinate in \(B\); unpad is the partial inverse, undefined on padded coordinates. ◻
3 Method: a graded-monad calculus and sliding spectacles↩︎
The formalism behind 2 is a graded semantics over the category \(\mathsf{FinSet}^{\mathrm{op}}\) of finite sets and (op-of-)functions: full categorical foundations are in 8; the body needs only the consequence below.
Proposition 2 (Grade-naturality, “sliding spectacles” (P3)). For every base morphism \(h\colon A\to B\) and every grade morphism \(r\colon Y\to X\) in \(\mathsf{FinSet}^{\mathrm{op}}\), the reindexing \(r^*\) commutes with \(T(h)\): \[T_Y(h) \circ r^* \;=\; r^* \circ T_X(h)\colon\; A^X\to B^Y.\] For partial read maps with padding interpretation \(\bot\), additionally \(h(\bot_A)=\bot_B\).
A grade morphism \(r\) acts at both boundaries of a tube, so it is drawn as a spectacles bracketing tubing, potentially changing the tubing before-and-after. Naturality says spectacles may slide freely along tubes unless blocked along the tube boundary by a tube-to-tube base morphism: this single rewrite is a proof-engine. When we depict only the top boundary of tubes, gradings are just string diagrams for partial functions on finite sets.
4 Equivariance, in one line per claim↩︎
The statements in this section are deliberately familiar. They calibrate the calculus against facts whose ordinary-index proofs are well-understood: their value here is that each becomes a one-line application of grade-naturality (2) within the rewrite system of 2. The main stress test of the calculus is the mask-augment theorem of 5, where the same rewrite principle controls a less obvious structural transformation.
4.1 Transformer permutation equivariance↩︎
Many data transformations have tubed wires as inputs and outputs, and usually they represent boundaries that indexing gadgets on tube boundaries do not interact with. Interesting mathematics occurs when there is interaction. Let \(f: \mathbb{R}^k \rightarrow \mathbb{R}^k\) and let \(\pi: k \rightarrow k\) denote a bijection; i.e. a permutation of indices, or a reorder.
Definition 3 (Permutation equivariance). A process \(f\colon \textcolor{Purple}{A}^{\textcolor{blue}{T}}\to\textcolor{ForestGreen}{B}^{\textcolor{blue}{T}}\) is \(\pi\)-equivariant* if it is equal whether \(\pi\)-reorder occurs before or after, as drawn opposite. Visually, \(\pi\) “hops over". If this
holds for all \(\pi\), then \(f\) is *permutation-equivariant**
Definition 4 (Permutation invariance). A process \(f\colon \textcolor{ForestGreen}{A}^{\textcolor{blue}{T}}\to \textcolor{Purple}{B}\) is \(\pi\)-invariant* if
\(\pi\)-reorder on the inputs has no effect, as drawn opposite. Visually, \(\pi\) is “absorbed" reading one way, and”freely generated" the other. If this holds for all \(\pi\), then \(f\) is permutation invariant.*
Remark 5. For \(f: \mathbb{R}^j \rightarrow \mathbb{R}^k\) where \(j\) and \(k\) both support representations of a (finite) group \(G\), the definitions above are adaptable to \(G\)-equivariance and \(G\)-invariance with the same visual consequences.
Proposition 6 (Attention is permutation-equivariant in queries).
**Proof.* \(\sigma\) denotes softmax throughout. Blue is for embedding dimension and orange is for sequence length. A permutation on the queries’ sequence axis slides through to the output. \[\scalebox{0.75}{\tikzfig{ML/equi0}}\;{=}\; \scalebox{0.75}{\tikzfig{ML/equi1}}\;{=}\; \scalebox{0.75}{\tikzfig{ML/equi2}}.\] ◻*
Proposition 7 (Attention is invariant under simultaneous key/value permutation).
Proof.* A simultaneous permutation on the shared K,V sequence axis is absorbed by the contraction. Reading bottom-up; (a): \(\sigma\) is perm.-equi.; (b): perms. have inverses; (c): pants copy/merge indexing; (d)**: \(\Sigma\) is perm.-inv. \[\scalebox{0.5}{\tikzfig{ML/inv0vert}}\;\raisebox{0.75cm}{\overset{}{=}}\; \scalebox{0.5}{\tikzfig{ML/inv1vert}}\;\raisebox{0.75cm}{\overset{(a)}{=}}\;\scalebox{0.5}{\tikzfig{ML/inv2vert}}\;\raisebox{0.75cm}{\overset{(b)}{=}}\; \scalebox{0.5}{\tikzfig{ML/inv3vert}}\;\raisebox{0.75cm}{\overset{(c)}{=}}\;\scalebox{0.5}{\tikzfig{ML/inv4vert}}\;\raisebox{0.75cm}{\overset{(b)}{=}}\; \scalebox{0.5}{\tikzfig{ML/inv5vert}}\;\raisebox{0.75cm}{\overset{(d)}{=}}\;\scalebox{0.5}{\tikzfig{ML/inv6vert}}\] ◻*
5 Mask-augment duality↩︎
The folklore explanation of causal masking and positional encoding is that they break the permutation equivariance of attention. [13] demonstrated empirically that causal masking is sufficient for this purpose, confirming the theoretical observation of [14]. A naïve method to achieve the same goal for an unmasked attention is to augment the data going through the Key and Value streams: take the original sequence \(\langle a, b, c, d, ...\rangle\), and create new “chunked" tokens corresponding to prefixes \(\langle [a], [a,b], [a,b,c], [a,b,c,d], ... \rangle\); this is evidently not permutation equivariant, and intuitively there is some shape-level relationship to causal masking. This construction is in fact equal to causal masking, as a corollary of the following general theorem:
Theorem 1 (Mask-Augment Duality). For mask \(\mu\colon k \to \mathcal{P}(k)\) encoding the visible positions per query, let \(E_\mu = \{(i,j)\in k\times k : j \in \mu(i)\}\) with projections \(p,s\colon E_\mu \to k\), \(p(i,j) = i\), \(s(i,j) = j\). Masked attention equals fibrewise unmasked attention over \(p^{-1}(i)\) using \(K\circ s\) and \(V\circ s\): \[\mathrm{Attn}_\mu(Q,K,V)_i \;=\; \sum_{j\in\mu(i)} \frac{\exp(\langle Q_i, K_j\rangle / \sqrt{d})}{\sum_{\ell\in\mu(i)} \exp(\langle Q_i, K_\ell\rangle / \sqrt{d})}\, V_j .\] Equivalently, the K/V streams are augmented along the tube-function \(\mathcal{A}\colon k \twoheadleftarrow \bigsqcup_{i \leq k}\mu(i)\) with \(\mathcal{A}^{-1}(i) = \mu(i)\). Rows with \(\mu(i)=\varnothing\) are fixed to a default, e.g.the zero vector.
Proof. Full lemma statements and a more detailed re-display of the chain are in 10; an einops-code transcription accompanies it.
| \(\rightarrow\) | \(\rightarrow\) | \(\rightarrow\) | \(\rightarrow\) |
| \(\rightarrow\) | \(\rightarrow\) | \(\rightarrow\) | \(\rightarrow\) | \(\rightarrow\) |
◻
Three threads of attention-time optimisations are special cases of 1. Attribution, speedup formulae, and on-device validation are in 12; 14 (Apple M4 Max) reports the headline numbers.
Corollary 2 (Hoisting common fibres). A common sub-fibre factors out of the per-query computation. The strict-left-prefix case is prompt-cache prefix sharing [15], [16]; the DAG case is shared-context decoding [17]–[19].
Corollary 3 (Compacting disjoint fibres). Pairwise-disjoint fibres compact into dense slabs; one SDPA per slab is exact. This is sequence packing [20], [21] and the variable-length path of FlashAttention, xFormers, and FlexAttention [22]–[24].
Corollary 4 (Bounded-component scheduling). If each row sees at most \(c\) intervals of total width \(w\), the mask compiles to a row-state vocabulary with \(\mathcal{O}(c\,T\,w/B^2)\) tile visits versus \(\mathcal{O}(T^2/B^2)\) dense. The \(c{=}1\) case is sliding-window attention [25]–[28]; \(c{\geq}2\) covers local+global mixtures. Kernel-side compilation: FlashMask and FlexAttention [24], [29].
6 Discussion and conclusion↩︎
We have presented a formal diagrammatic proof calculus for the structural layer of tensor programming. This is the layer manipulated by reshape/reorder/repeat/reduce/pack/unpack/pad patterns in einops-style code, ubiquitous in implementations of machine learning architectures. Prior ML diagram systems are useful as notation, communication devices, or compiler frontends, but they do not provide a formal graphical proof calculus for structural equality.
Our calculus formally bridges undirected tensor network notation, and directed computation graphs. Undirected tensor-network calculi are the right syntax for Kronecker-product structure, contraction topology, and Frobenius-basis reasoning. The structural layer of deep-learning programs is different. Batches, sequence axes, broadcasts, repeats, reshapes, packing, and padding are cartesian phenomena: they are about copying, deleting, and reindexing positions. Tubes make this cartesian structure visible while preserving the operational left-to-right reading of computation graphs. Their upper boundaries retain the familiar tensor-network view of axes, but the interior records the directed computation.
The main rewrite in our diagrams is grade-naturality. Pointwise maps do not care how the surrounding finite index set has been rearranged; hence the “spectacles” slide. This simple fact accounts for a surprisingly large class of tensor-program equalities. The examples in the paper are intentionally familiar at first: transformer permutation equivariance and key/value permutation invariance become one-line diagrammatic derivations. The mask-augment duality then shows that the same principle also explains a less obvious equivalence: a mask is a query-indexed display of visible key/value fibres, and masked attention is equal to unmasked attention computed fibrewise over that display.
In the appendix, we present several corollaries of this single rewrite. These recover known architecture implementation patterns from the literature. These serve as validation targets, demonstrating the breadth of results captured in one succinct derivation.
The accompanying code transcription in torch with einops proves the same theorem using only explicit reshapes, packs, pads, unpads, broadcasts, masks, softmax-support changes, and reductions. It is verbose and difficult to audit locally. The diagrammatic proof compresses this derivation into typed local rewrites, each licensed by grade-naturality or an explicitly declared algebra/padding law. Our diagrams are thus more than a presentation layer, and constitute an ergonomic tool to compress useful proofs.
Several directions for further development and application of our contribution follow naturally. A small normaliser for the structural fragment (WIP) would turn the calculus into a practical checker for einops-style rewrites. Derivative tubes would extend the language to backpropagation while preserving the separation between structural reindexing and base arithmetic. Multi-head attention, multi-query attention, tiled attention, and FlashAttention-style IO rewrites are natural next tests: each combines finite display structure with a small number of declared algebraic or systems-level invariants. More speculatively, the same wires/tubes split may be useful in mechanistic interpretability, where residual-stream read/write structure is often discussed diagrammatically but rarely given a formal axis semantics.
7 Cheat sheet↩︎
8 The graded monad: \(\mathsf{FinSet}^{\mathrm{op}}\) over \(\mathsf{Euc}\)↩︎
The formal algebraic semantics of tubes cohering with string diagrams is a symmetric monoidal graded monad; these proofs assume familiarity with symmetric monoidal categories.
Definition 8 (Symmetric monoidal graded monad). A graded monad consists of a grading category* \(\mathcal{C}\) over a base category \(\mathcal{D}\), along with a functor \(\mathbb{G}\) from the grading into the endofunctor category of \(\mathcal{D}\), \(\mathbb{G}: \mathcal{C} \rightarrow [\mathcal{D},\mathcal{D}]\). Endofunctor categories always carry monoidal structure \(([\mathcal{D},\mathcal{D}],\circ,\mathsf{id}_{\mathcal{D}})\), where only some pairs of endofunctors \(\mathbf{J},\mathbf{K}: \mathcal{D} \rightarrow \mathcal{D}\) may braid symmetric monoidally (in the presence of natural isomorphisms \(\mathbf{J} \circ \mathbf{K} \Leftrightarrow \mathbf{K} \circ \mathbf{J}\)). When the grading category \(\mathcal{C}\) and the functor \(\mathbb{G}\) are symmetric monoidal we have a Symmetric monoidal graded monad.*
Throughout the body, we have used a \(\mathsf{FinSet}^{\mathrm{op}}\) graded-monad over \(\mathsf{Euc}\); we will handle the padded fragment \(\mathsf{FinPart}^{\mathrm{op}}\) by a reduction to a \(\mathsf{FinSet}^{\mathrm{op}}\)-grading. \(\mathsf{FinSet}^{\mathrm{op}}\) is the opposite of the category of finite sets and partial functions, with cartesian product as symmetric monoidal structure. \(\mathsf{Euc}\) is the category of Euclidean spaces and arbitrary functions, with direct sum as symmetric monoidal structure. The connection between the algebraic categorical content and the diagrammatic notation is already established by Melliès’ Functor Boxes [30], where tubes are an evident extension to the symmetric monoidal case [31]. Hence all we have to elaborate and verify is that we indeed have a symmetric monoidal graded monad.
We proceed in two steps. First, we construct and verify the \((\mathsf{FinSet},\times,1)^{\mathrm{op}}\) grading over \(([\mathsf{Euc},\mathsf{Euc}],\circ,\mathrm{id}_{\mathsf{Euc}})\). Second, we demonstrate that padding bookkeeping of \(\mathsf{FinPart}^{\mathrm{op}}\) is notational: a \(\mathsf{FinSet}^{\mathrm{op}}\) grading suffices.
8.1 \(\mathsf{FinSet}^{\mathrm{op}}\) graded monad over \(\mathsf{Euc}\)↩︎
We will proceed by describing the desired functor \(\mathbb{G}\), and then we verify its required properties.
8.1.0.1 The object part of \(\mathbb{G}\)
The functor \(\mathbb{G}: (\mathsf{FinSet},\times,1)^{\mathrm{op}} \rightarrow ([\mathsf{Euc},\mathsf{Euc}],\circ,\mathrm{id}_{\mathsf{Euc}})\) sends \(\mathsf{FinSet}\ni X\) to the endofunctor \((-)^{\oplus |X|}: \mathsf{Euc}\rightarrow \mathsf{Euc}\). Equivalently, in the cartesian-base form most useful for the body’s calculus, write \(T_X(A) := A^X\) for the indexed-product object. The unit and product cases are then \[T_1(A) \cong A, \qquad T_{X\times Y}(A) \cong T_X(T_Y(A)),\] which is the strong graded structure: the cartesian product on grades corresponds to the composition on endofunctors. Respecting our outside-in reading convention for tubing, the functor sends \(X \times Y\) to the composite endofunctor \((-)^{\oplus |X|} \circ (-)^{\oplus |Y|}\). On morphisms, a finite read map \(r\colon Y\to X\) in \(\mathsf{FinSet}\) acts contravariantly as the reindexing \[r^*\colon A^X \to A^Y, \qquad (r^* a)_y \;=\; a_{r(y)},\] recovering the display-map view of the body: the read map points opposite to data flow, sending each output index to the input index it reads from.
8.1.0.2 The morphism part of \(\mathbb{G}\)
In this section we must formally justify the display-map argument that relates copy-delete to functions between finite sets. We proceed by the following steps; first, we establish a helper lemma concerning \(\mathsf{Euc}\) that renders its objects as positive integers; second, we establish a normal-form factorisation of functions between finite sets (viewable as integers) in terms of bijections, surjections, and injections; third, we establish a one-to-one correspondence between the normal form factorisation and composites of copies, deletes, and braidings in \(\mathsf{Euc}\); fourth, we establish that the correspondence holds for tensor products, which are equivalently sequentially composed tupling endomorphisms, or cartesian products of finite sets in the grading category.
First, the intermediary lemma concerning \(\mathsf{Euc}\), justifying our convention throughout of only using a single base wire \(\mathbb{R}\).
Lemma 1 (One wire type suffices for \(\mathsf{Euc}\)). Working in the ambient setting of (not necessarily smooth) functions between Euclidean spaces, only one wire type \(\mathbb{R}\) is needed.
**Proof.* The direct sum is the categorical product of Euclidean spaces, which Fox’s theorem licenses as the parallel composition in the presence of natural copy-delete maps. Since every Euclidean space is a direct sum of \(\mathbb{R}\), we are done. ◻*
Lemma 1 allows us to deal with just a single copy map \(\delta: \mathbb{R} \rightarrow \mathbb{R} \oplus \mathbb{R}\) and a single delete map \(\epsilon: \mathbb{R} \rightarrow \{\star\}\); the 0-dimensional space is isomorphic to the singleton set.
Second, the normal form. Where \(n\) is a positive integer, let \([n]\) denote the set \(\{0,1,...,(n-1)\}\); without loss of generality for the arguments to follow, we may consider finite sets of this form 3
Proposition 9. Every function \(f: [n] \rightarrow [m]\) admits a canonical factorisation of two bijections, a surjection, and an injection \([n] \leftrightarrow [n] \twoheadrightarrow [|\mathsf{Im}(f)|] \leftrightarrow [|\mathsf{Im}(f)|] \hookrightarrow [m]\), where the surjective and injective legs are isotone (order-respecting).
**Proof.* The surjection and injection arise for free from the usual epi-mono factorisation available in \(\mathsf{FinSet}\). The bijections rearrange \([n]\) and \([|\mathsf{Im}(f)|]\) such that the fibres of the surjection and injection respectively are contiguous, which render the surjection and injection isotone whilst preserving composite equality with the original \(f\). ◻*
Third, the correspondence to copy-deletes.
Proposition 10. Up to isomorphism, every morphism \(g: \mathbb{R}^m \rightarrow \mathbb{R}^{n}\) that is formed of a composite of \(\delta\), \(\epsilon\), and braiding corresponds to a unique function \(f: [n] \rightarrow [m]\).
**Proof.* By counitality, every \(g\) is isomorphic to \(g'\) where there are no \(\epsilon\) sequentially composed with any leg of a \(\delta\). By braid naturality and expanding/contracting identity maps we may organise all the \(\epsilon\) to occur sequentially before all \(\delta\): \(g'\) is isomorphic to a braiding, followed by a tensor of the form \(\mathrm{id}_{\mathbb{R}^x} \oplus (\epsilon)^y\), followed by a second braiding, followed by a tensored composite of \(\delta\) (by coassociativity, the order of \(\delta\) is immaterial). Up to an isomorphism in \(\mathsf{FinSet}\) recasting coproducts (disjoint unions) as \([n]\), the \(\delta\)-block and \(\epsilon\)-block agree with coproduct factorisations of the desired surjection and injection respectively in \(\mathsf{FinSet}^\mathrm{op}\) obtained via Proposition 9. ◻*
Now we may state (\(\times\)-free) action of \(\mathbb{G}\) on morphisms.
Definition 11 (\(\mathbb{G}\) on \(\times\)-free morphisms). \(\mathbb{G}\) sends finite functions \(f: Y \rightarrow X\) in \(\mathsf{FinSet}^\mathrm{op}\) to the function \(\mathbb{G}(f): \mathbb{R}^{|X|} \rightarrow \mathbb{R}^{|Y|}\) in \(\mathsf{Euc}\) composed of \(\delta\), \(\epsilon\) and braiding, given by Propositions 9 and 10, which guarantee well-definedness and uniqueness up to isomorphism.
Finally, we address the behaviour of \(\mathbb{G}\) on morphisms with respect to \(\times\). For this it suffices to note that pairs of integers can be ordered lexicographically, which establish isomorphisms through which Propositions 9 and 10 carry through.
8.2 Completing the verification↩︎
Proposition 12. \(\mathbb{G}\) is (strong) symmetric monoidal.
**Proof.* We have two tasks; first we must establish the symmetric monoidality of \(\mathbb{G}\); second we must verify naturality in the image \(\mathbb{G}\). For the first task, the
only point of interest is the specification that the braids \(\gamma_{A,B}: B \times A \leftarrow A \times B\) of \(\mathsf{FinSet}^\mathrm{op}\) are sent to the (reorder)
natural isomorphisms \(((-)^{|B|})^{|A|} \Leftrightarrow ((-)^{|A|})^{|B|}\) in \([\mathsf{Euc},\mathsf{Euc}]\); from this the monoidality and Yang-Baxter equations evidently follow. For the
second task, it suffices to observe that braids are definitionally natural transformations, and that by Fox’s theorem and corollary 1 all morphisms in \(\mathsf{Euc}\) are cohomomorphisms with respect to the \(\delta\)-\(\epsilon\) comonoid. ◻*
Corollary 5. \(\mathbb{G}\) is a semirig functor, additionally sending the coproduct of \(\mathsf{FinSet}^\mathrm{op}\) to \((- \;\oplus =)\).
**Proof.* The same proof strategy above carries over to the monoidal structures \((\mathsf{FinSet}^\mathrm{op},+,0)\) and \(([\mathsf{Euc},\mathsf{Euc}],\oplus,\mathsf{Euc}\rightarrow \mathbb{R}^0)\). The verification of the interaction (e.g. distributivity) of the monoidal products is routine. ◻*
8.3 Handling (un)pad↩︎
Concretely in \(\mathsf{Euc}\), the following is just the observation that we can explicitly introduce paddings \(\star\) via a direct sum of copies of \(\star\), at which point we are dealing with tensor-axes manipulation that falls within the \(\mathsf{FinSet}^\mathrm{op}\) grading.
8.3.0.1 Explicit form.
Fix a padding interpretation \(\bot_A \in A\) for each base object. A partial read map \(r\colon Y\rightharpoonup X\) acts on \(A^X\) as \[(r^*_{\bot_A} a)_y \;=\; \begin{cases} a_{r(y)} & y \in \mathrm{dom}(r),\\ \bot_A & y \notin \mathrm{dom}(r). \end{cases}\] This is exactly the dense action \(r^*\) of the previous subsection, extended to undefined positions by reading from the chosen filler. Grade-naturality (2) then has a single compatibility condition: a base morphism \(h\colon A\to B\) slides past a partial read \(r\) iff \[h(\bot_A) \;=\; \bot_B,\] i.e.\(h\) preserves the chosen padding element. This is the body’s “compatible padding” condition spelt out at the level of the indexed-product action.
8.3.0.2 Reduction to \(\mathsf{FinSet}^{\mathrm{op}}\).
Because we parameterise padding with a single element at a time, we can exploit two observations. First, because partial functions arise as the Kleisli category of the Maybe monad \((- + \bot)\), \(r\colon A \rightharpoonup B\) are viewable as functions \(\hat{r}\colon A \to B + \{\bot\}\), where elements in \(A\) outside the defined domain of \(r\) are mapped to a distinguished additional formal element \(\bot\); the action of \(\hat{r}\) on the preimage of \(\bot\) corresponds to copying the padding element that may be introduced as an explicit direct-summed morphism in \(\mathsf{Euc}\). Second, by Corollary 5, we can handle coproducts in the grading category \(\mathsf{FinSet}^\mathrm{op}\) algebraically, though we have opted in the body not to depict them graphically to keep tubing visually simple. To eliminate partial functions altogether, it suffices to note that Kleisli composition can be made explicit in the grading \(\mathsf{FinSet}^\mathrm{op}\), holding aside a formal element \(\bot\); although in practice, we do not come across instances in which padded data is subject to arbitrary transformations; usually they are absorbed or bookkept for removal.
9 Metadata displays↩︎
This appendix is optional reading. The dense-padding proof of 1 (5) does not depend on it: the body uses only the \(\mathsf{FinSet}^{\mathrm{op}}\) grading of 8 and the dense-with-partial-reads presentation of 1. Metadata displays supply the alternative grading category used when one wants diagrams to remember the logical block decomposition of a packed buffer (rather than compiling it away into a dense partial read). Readers focused on the body theorems may skip this appendix; readers interested in non-contiguous packing, structured ragged batches, or block-aware rewriting will find the construction here.
The dense calculus grades by finite sets of tensor positions. For ragged and packed tensors this is insufficient: two packed buffers may have the same underlying finite set of cells while carrying different logical decompositions into blocks. We therefore replace a bare finite set by a finite display \[p:X\to B,\] where \(X\) is the finite set of dense cells and \(B\) is the finite set of logical blocks. The fibre \(p^{-1}(b)\) is the block of cells belonging to \(b\).
The purpose of the construction is threefold. First, we define a category \(\mathsf{Disp}(\mathsf{Met})\) whose morphisms preserve this block metadata. Second, we check that \(\mathsf{Disp}(\mathsf{Met})\) is symmetric monoidal, so it can be used as a grading category for tubed diagrams. Third, we relate it back to ordinary \(\mathsf{FinSet}^\mathrm{op}\) grading by an adjunction: forgetting the display recovers the dense set of cells, while the free embedding equips a finite set with trivial display structure.
9.1 The metadata category \(\mathsf{Met}\)↩︎
Let \([n]=\{0,\ldots,n-1\}\), regarded as a finite total order. We define \(\mathsf{Met}\) to have as objects the least class of finite posets containing all \([n]\) and closed under \[P,Q\longmapsto P\times Q, \qquad P,Q\longmapsto P+Q, \qquad P,Q\longmapsto P+_{\mathrm{ord}}Q.\] Here \(P\times Q\) is the cartesian product (lexicographic) order, \(P+Q\) is disjoint union with no new order relations between the summands, and \(P+_{\mathrm{ord}}Q\) is ordered disjoint union: every element of \(P\) is placed below every element of \(Q\).
Morphisms in \(\mathsf{Met}\) are monotone maps. The intended readings are: \[\begin{array}{cc} \boldsymbol{metadata construction} & \boldsymbol{tensor reading}\\ \hline [n] & \text{ordinary dense axis of length }n\\ P\times Q & \text{rectangular product of axes}\\ P+Q & \text{unordered collection of regions}\\ P+_{\mathrm{ord}}Q & \text{contiguous concatenation of regions} \end{array}\]
The subcollection generated only by finite total orders and cartesian products is the dense rectangular fragment. The additional coproduct-like operations are used only to remember how a packed dense surface decomposes into logical pieces.
9.2 Displays↩︎
A metadata display is a monotone map \[p:X\to B\] in \(\mathsf{Met}\). We call \(X\) the dense surface and \(B\) the block poset. The fibre \(p^{-1}(b)\) is the logical piece of the packed object lying over \(b\).
Example 1 (Dense axis). The ordinary dense axis of length \(n\), with no nontrivial packing metadata, may be represented either as the discrete display \[\mathrm{id}_{[n]}:[n]\to[n],\] where every cell is its own block, or as the one-block display \[!:[n]\to[1],\] where all cells form a single dense block. These two embeddings have different universal properties below.
Example 2 (Packed ragged batch). Let \[X=[2]+_{\mathrm{ord}}([3]\times[4])+_{\mathrm{ord}}[5].\] The map \[p:X\to[3]\] sending each ordered summand to its corresponding element of \([3]\) represents a packed ragged batch with three logical fibres of shapes \(2\), \(3\times4\), and \(5\). The underlying dense buffer has cardinality \[|X|=2+12+5=19,\] but the display remembers the three-block decomposition.
9.3 The category \(\mathsf{Disp}(\mathsf{Met})\)↩︎
The category \(\mathsf{Disp}(\mathsf{Met})\) has displays \(p:X\to B\) as objects. A morphism \[(a,b):(p:X\to B)\longrightarrow(q:Y\to D)\] is a commutative square in \(\mathsf{Met}\): \[\begin{tikzcd} X \arrow[r,"a"] \arrow[d,"p"'] & Y \arrow[d,"q"] \\ B \arrow[r,"b"'] & D . \end{tikzcd}\] Equivalently, \[q\circ a=b\circ p.\] The map \(a\) sends dense cells to dense cells, while \(b\) sends source blocks to target blocks. The commutativity condition says that cell movement respects the recorded packing metadata. Where \(b\) is a bijection, the logical blocks are conserved up to reordering. More generally, we allow the metadata of logical blocks to be rewritten.
Composition is by pasting commutative squares. Identities are identity squares. Hence \(\mathsf{Disp}(\mathsf{Met})\) is the arrow category of \(\mathsf{Met}\), restricted to the display objects under consideration.
Remark 13 (Relation to partial reads). The definition above records packing metadata, not padding values. Partial reads and padding fillers are additional structure used by the operational semantics. They may be added by replacing the dense-cell map \(a\) with a partial map, or by passing to a partial-map enrichment. The present appendix only constructs the metadata grading category and its relation to ordinary finite-set grading.
9.4 Symmetric monoidal structure↩︎
The monoidal product of displays is given componentwise: \[(p:X\to B)\boxtimes(q:Y\to D) := (p\times q:X\times Y\to B\times D).\] On morphisms, \[(a,b)\boxtimes(c,d):=(a\times c,\; b\times d).\] The monoidal unit is \[\mathbb{1}:=\mathrm{id}_{[1]}:[1]\to[1].\]
Proposition 14. \(\mathsf{Disp}(\mathsf{Met})\), equipped with \(\boxtimes\), is symmetric monoidal.
Proof. Since \(\mathsf{Met}\) has finite products, the cartesian product supplies associators, unitors, and symmetries: \[(X\times Y)\times Z \cong X\times(Y\times Z), \qquad X\times[1]\cong X, \qquad X\times Y\cong Y\times X.\] Applying these isomorphisms simultaneously on dense surfaces and block posets gives the corresponding structure maps in \(\mathsf{Disp}(\mathsf{Met})\). For example, the symmetry \[(p:X\to B)\boxtimes(q:Y\to D) \longrightarrow (q:Y\to D)\boxtimes(p:X\to B)\] is the square \[\begin{tikzcd} X\times Y \arrow[r,"\sigma_{X,Y}"] \arrow[d,"p\times q"'] & Y\times X \arrow[d,"q\times p"] \\ B\times D \arrow[r,"\sigma_{B,D}"'] & D\times B . \end{tikzcd}\] The square commutes by naturality of the cartesian symmetry. The associator and unitors are identical componentwise arguments. The pentagon, triangle, and hexagon coherence conditions hold because they hold in \(\mathsf{Met}\) for cartesian product, and \(\mathsf{Disp}(\mathsf{Met})\) inherits them levelwise. ◻
Thus \(\mathsf{Disp}(\mathsf{Met})\) is eligible to grade the tube calculus: tensoring grades corresponds to taking product structure both on the dense cells and on their displayed metadata.
9.5 Forgetting and freely adding display structure↩︎
We may characterise the relationship between the plain \(\mathsf{FinSet}^\mathrm{op}\) grading and the metadata variant by a pair of adjunctions.
Let \[U:\mathsf{Disp}(\mathsf{Met})\to\mathsf{FinSet}\] be the forgetful functor sending a display \(p:X\to B\) to the underlying finite set \(|X|\) of dense cells. On a morphism \[(a,b):(p:X\to B)\to(q:Y\to D),\] \(U\) sends \((a,b)\) to the underlying finite function \[|a|:|X|\to|Y|.\]
There are two canonical ways to regard a finite set as a display.
First, the discrete display functor \[\Delta:\mathsf{FinSet}\to\mathsf{Disp}(\mathsf{Met})\] sends a finite set \(S\) to \[\Delta S := \mathrm{id}_S:S\to S.\] This equips \(S\) with the finest possible trivial metadata: every cell is its own block.
Second, the codiscrete or one-block display functor \[\nabla:\mathsf{FinSet}\to\mathsf{Disp}(\mathsf{Met})\] sends \(S\) to \[\nabla S := !:S\to[1].\] This equips \(S\) with the coarsest possible trivial metadata: all cells lie in a single block.
Proposition 15. There are adjunctions \[\Delta \dashv U \dashv \nabla .\]
Proof. For the left adjunction, let \(S\in\mathsf{FinSet}\) and let \(p:X\to B\) be a display. A morphism \[\Delta S=(\mathrm{id}_S:S\to S)\longrightarrow(p:X\to B)\] is a commutative square \[\begin{tikzcd} S \arrow[r,"a"] \arrow[d,"\mathrm{id}_S"'] & X \arrow[d,"p"] \\ S \arrow[r,"b"'] & B . \end{tikzcd}\] The commutativity condition is \(p\circ a=b\). Hence \(b\) is uniquely determined by \(a\). Therefore \[\mathsf{Disp}(\mathsf{Met})(\Delta S,p)\cong \mathsf{FinSet}(S,U p),\] naturally in \(S\) and \(p\). Thus \(\Delta\dashv U\).
For the right adjunction, a morphism \[(p:X\to B)\longrightarrow \nabla S=(!:S\to[1])\] is a commutative square \[\begin{tikzcd} X \arrow[r,"a"] \arrow[d,"p"'] & S \arrow[d,"!"] \\ B \arrow[r,"!"'] & [1] . \end{tikzcd}\] The bottom map is unique, and the square commutes automatically. Hence such morphisms are exactly finite functions \(a:X\to S\), giving a natural bijection \[\mathsf{Disp}(\mathsf{Met})(p,\nabla S)\cong \mathsf{FinSet}(U p,S).\] Thus \(U\dashv\nabla\). ◻
10 Mask-augment duality: helper lemmas and detailed proof↩︎
This appendix records the four helper lemmas that license individual transitions in the eleven-frame chain proving 1, and reproduces the chain itself with each step’s licence stated alongside it. The lemma proofs are routine pointwise verifications and are deferred for space.
Lemma 2 (Augment \(\Leftrightarrow\) pad/unpad). There exists a pad/unpad pair of functions (varying together in the padding parameter) such that the augmented K/V tube
\(\mathcal{A}_M(\cdot)\) equals a dense pad followed by the corresponding unpad.
10.0.0.1 The chain in detail.
The chain reproduces 5, displayed full-size; each equality is licensed by the lemma or move discussed inline.
\[\tikzfig{maskprep/augment0} = \tikzfig{maskprep/augment1}\] The first step follows definitionally from the padding. \[= \tikzfig{maskprep/augment2} = \tikzfig{maskprep/augment3}\] These two steps follow from diagram-manipulation. Copy-maps nested inside copants allow copying of indexing-morphisms, and since tubings are string-diagrams in \(\mathsf{FinSet}^\mathrm{op}\) (which has copy and delete), we may use counitality for the second equation. \[= \tikzfig{maskprep/augment4}\] This step introducing a \(0\)-padding comes from the following lemma, from the middle frame to the right.
Lemma 3. \[\tikzfig{lemmas/FR2}\]
**Proof.* The first equality results from straightforward application of Lemma 5. The second equality may be proven by observing that the initial socks in both diagram create \(m\) and \(n\) copies of the same input data, and the unpad in the second diagram selects a subset of size \(m\), from copies, guaranteeing that the resulting matrix contains the same entries. ◻*
\[= \tikzfig{maskprep/augment5}\] This equation comes from the following interaction between padding and products.
Lemma 4. \[\tikzfig{lemmas/pants1}\]
**Proof.* The pants match padded indexes in both input tensors, feeding these in to \(f\), producing \(f(u, v)\). ◻*
\[= \tikzfig{maskprep/augment6}\]
This equation (introducing a new pad-unpad pair) arises by the following definitional property of pad-unpad.
Lemma 5. \[\tikzfig{lemmas/FR}\]
**Proof.* See the definition of the pad and unpad, and their interaction defined in 1. ◻*
\[= \tikzfig{maskprep/augment7}\]
This step relies on padding-transformation properties of softmax:
Lemma 6. \[\tikzfig{lemmas/FRsmax}\]
**Proof.* This lemma falls out of the fact that \(e^{-\infty} = 0\), and padding an array with \(-\infty\) preserve the softmax output of the existing array, sending the padded elements to 0. \[\begin{align} \sigma(\vec{x} || -\infty \cdots -\infty) = \sigma(\vec{x}) || 0 \cdots 0. \end{align}\] ◻*
\[= \tikzfig{maskprep/augment8}\]
This step relies on addition being invariant to 0-padding.
Lemma 7. \[\tikzfig{lemmas/maskinvadd}\]
**Proof.* The sum of an array is invariant under padding with \(0\)s. \[\begin{align} \sum \vec{x}|| 0 \cdots 0 = \sum \vec{x} \end{align}\] ◻*
\[= \tikzfig{maskprep/augment9}\]
This step relies on an interaction law between paddings and general binary operations across tensors.
Lemma 8. \[\tikzfig{lemmas/RFmult}\]
**Proof.* Using 4, it is possible to push the pad on the output backwards through the product, producing a \(\texttt{pad}_0\) on one
input, and an arbitrary pad on the other. The pad value here maybe left unspecified, since \(\forall u. 0 \cdot u = 0\). ◻*
\[= \tikzfig{maskprep/augment10}.\]
The final step is another interaction law of softmax with padding.
Lemma 9. \[\tikzfig{lemmas/FRsmax2}\]
**Proof.* The equivalence is a consequence of the following known property of the softmax function. \[\begin{align} \sigma(\vec{x} || -\infty \cdots -\infty) = \sigma(\vec{x}) || 0 \cdots 0. \end{align}\] Inputs to softmax may be padded with \(-\infty\) leaving the output unchanged, and sending the padding to 0. This is a standard trick used to implement masked self-attention. ◻*
11 Mask-augment duality: Comparison with code↩︎
Below we repeat the eleven frames of the derivation, each paired with the corresponding forward function, implemented in torch + einops. The code transcription proves the same identity without diagrams. Its length is the point:
the graphical proof is a typed compression of a structural tensor-program derivation, not an informal illustration of it.
| Diagram | torch + einops |
Figure 3: 
.
Figure 4: 
.
Figure 5: 
.
Figure 6: 
.
Figure 7: 
.
Figure 8: 
.
Figure 9: 
.
Figure 10: 
.
Figure 11: 
.
Figure 12: 
.
Figure 13: 
.
12 Three systems optimisations as corollaries↩︎
The three corollaries below recover known systems optimisations as instances of the fibrewise decomposition in 1. Prefix and DAG hoisting, packed-document compaction, and bounded-component scheduling come from disparate threads of the systems literature and were engineered independently; the duality places them on a common footing and yields a quantitative speedup prediction for each. We attribute the corollaries to their known instances and validate the predictions on a single Apple M4 Max (128 GB unified memory). The following corollaries are not priority claims over the cited systems optimisations. They are validation targets: independently discovered efficient implementations should be recovered by any structural theory that correctly explains attention masks.
12.1 Corollary 1: hoisting common fibres↩︎
12.1.0.1 Known instances.
For \(n\) queries sharing a contiguous left-prefix, prompt-prefix caching computes the prefix’s K/V once and reuses it: PagedAttention with prefix sharing in vLLM [15], RadixAttention in SGLang [16], and llama.cpp’s prompt cache. For queries sharing an interior segment (tree-of-thoughts [18], self-consistency sampling [32], shared-context RAG), Hydragen [17] runs varying suffixes against the shared K/V; speculative-decoding tree caches do the same [19].
12.1.0.2 Validation.
Hermes 4 70B (MLX 4-bit [33]) via mlx-lm [34]. Strict-prefix sweep: shared prompt of length \(P\), \(n\) candidate suffixes of length \(s\), predicted speedup \(n(P+s)^2 / (P^2 + 2nPs + ns^2)\) over the grid \(P{\in}\{1024,2048\}\), \(n{\in}\{2,4,8\}\); \(\arg\max\)-of-best-of-\(n\) parity throughout, \(6.82\times\) measured at \(P{=}1024,\,s{=}64,\,n{=}8\) against \(4.45\times\) predicted (14, left). DAG bench (root, \(m\) intermediates, \(k\) leaves; \(R{=}256,\,I{=}32,\,L{=}16,\,m{=}4,\,k{=}4\)): \(1.20\times\) above strict-prefix, \(4.08\times\) above naive.
12.2 Corollary 2: compacting disjoint fibres into dense slabs↩︎
12.2.0.1 Known instances.
Sequence packing concatenates short documents into a fixed-length training sequence with a block-diagonal mask [20], [21]. The same pattern is the canonical use of FlashAttention’s variable-length API [22], [35], xFormers’ BlockDiagonalMask [23], and FlexAttention’s block_diag mask_mod [24].
12.2.0.2 Validation.
Synthetic Q/K/V at \(T\in\{4096,8192,16384\}\), document counts \(n_d\in\{4,8,16\}\), average visible density \(\rho = \sum_i \ell_i^2 / (\sum_i
\ell_i)^2\). Predicted speedup \(\approx 1/\rho\); best cell \(9.25\times\) measured against \(10.06\times\) predicted at \(\rho\approx 0.099\), parity \(\max|\Delta|\le 4\times 10^{-3}\) at fp16 (14, middle). A whole-model patch on Hermes 4 70B’s mlx_lm.models.llama
attention (exp02_compaction/hermes_mlx_patch.py) reproduces the pattern with bit-exact parity on a single full-causal segment.
12.3 Corollary 3: bounded-component masks compile to row-state machines↩︎
12.3.0.1 Known instances.
Sliding-window attention restricts each query to a bounded neighbourhood: Sparse Transformer [36], Longformer [25], BigBird [26], Mistral 7B [27], alternating local-global layers in Gemma 3/4 [33], [37], GPT-OSS 120B [33], Streaming-LLM [28]. Kernel-side, FlashAttention-2 supports
custom masks [22]; FlexAttention compiles score_mod/mask_mod into a BlockMask [24]; FlashMask records each row’s visible columns as an interval list [29].
12.3.0.2 Validation.
Two paths share one compile_program: measured on synthetic SWA at varying \(T\) and window \(w\); predicted on four GGUF frontier models (Gemma 4 31B/26B-A4B, Llama 3.3 70B,
GPT-OSS 120B) by reading sliding-window metadata from each GGUF header. Measured tracks \(T/(c\,w)\) over \(T/w\in[4,128]\) at \(73\)–\(95\%\) of theoretical maximum; the gap is host-side gather/scatter. At \(T{=}8192\): \(32\times\) per sliding layer for GPT-OSS 120B (window 128, 18/36 sliding);
\(7.1\times\) for Gemma 4 31B (window 1024, 50/60 sliding). Whole-model effective speedup scales by sliding-layer fraction (14, right).
Figure 14: Three corollaries of the mask-augment duality. Left: Corollary 1, strict-prefix hoisting on Hermes 4 70B. Middle: Corollary 2, packed-document compaction on synthetic Q/K/V at fp16 (mean \(\pm\) stdev across three \(T\) values per \(n_d\)). Right: Corollary 3, bounded-component scheduling: predicted (open markers, dashed) on four GGUF frontier models from architectural metadata, measured (filled, solid) on synthetic SWA at three window sizes..
12.4 Methodology and limitations↩︎
Corollary 1’s strict-prefix sweep is theorem-prediction validation, not a competitive benchmark; production stacks [15], [16] already implement strict-prefix caching. The genuine gap is the DAG case.
Corollary 3’s measured anchor is synthetic SWA. Local Mistral-7B-v0.1 and Longformer weights were absent or partially incompatible (Longformer’s split global-token projections break a unified-path patch); the predicted-side reads published architectural metadata.
Hermes 4 70B’s parity drift is 4-bit-bounded (\(\sim 0.5\)–\(1.0\) nats per cell on log-prob sums of 64–256 tokens). The \(\arg\max\)-of-best-of-\(n\) parity claim holds throughout; per-token claims would tighten to \(\sim 10^{-2}\) nats at fp16.
No comparison against custom block-sparse kernels on non-Apple silicon. Lowering compiled displays to FlexAttention or FlashMask on CUDA is the natural follow-up; the harness emits the abstract
Segment/RowStateobjects such a lowering would consume.
References↩︎
tensorgrad: Machine learning with symbolic tensors.” 2024.Laboratory for Human-Centered AI, Department of Philosophy, University of Oxford.
vincent.wang-mascianica@philosophy.ox.ac.uk↩︎Machine Learning Research Group, Department of Engineering Science, University of Oxford
nikhil@robots.ox.ac.uk↩︎We cannot in general only consider finite sets of this form (the skeleton of \(\mathsf{FinSet}^\mathrm{op}\)) as we must also distinguish the data of disjoint unions and cartesian products, which are bookkept by isomorphisms in the ambient \(\mathsf{FinSet}\). It is folklore that the free PROP generated by a cocommutative comonoid is precisely this skeleton \((\mathsf{FinSet},+,0)^\mathrm{op}\); this folklore essentially settles the morphism part of the functor, but we provide a proof anyway.↩︎