Graphical einops: bridging tensor networks and computation graphs

In deep learning, architecture diagrams are everywhere. But they are almost always just pictures. When we design a transformer, we draw boxes and arrows to represent data flow. The actual logic—the complex dance of reshaping, broadcasting, and contracting tensor axes—is still handled by engineers. We manipulate indices manually or write prose to explain why a transpose does not break the math.

Currently, we rely on libraries like einops for readable syntax. However, there is no formal way to prove that two different sequences of operations are actually doing the same thing. We operate on intuition. This is risky until a subtle index error causes a silent divergence in a training run. This paper introduces a formal graphical calculus. It turns these diagrams into actual proofs. By representing tensor axes as nested "tubes," the authors bridge the gap between undirected tensor networks and directed computation graphs. This provides a mathematical foundation for the "tricks" used to speed up inference.

The Problem

The status quo in tensor programming is split between two incompatible worlds. On one side is the undirected tradition of tensor networks. These treat wires as indices and focus on the topology (the connectivity) of contractions. This is great for physics. However, it struggles with the operational reality of deep learning. It fails to easily handle batching, sequence axes, and the temporal flow of a computation graph.

On the other side is the directed tradition of computation graphs. Here, operations are boxes and data flows left-to-right. This matches how we build models. Yet, it lacks a rigorous way to handle the "structural" layer of code. This includes the reshapes, repeats, and pads that define memory layout.

Proving identities like "permuting queries permutes outputs" requires manual index manipulation. We lack a unified language. We need one that handles both topological connectivity and procedural logic. Existing tools like tensorgrad focus on symbolic simplification. They are not a complete proof calculus for the structural rearrangements that dominate modern LLM implementations.

How It Works

The authors propose a "graded-monad tube calculus" to unify these views. The central architectural choice is the use of tubes. Instead of drawing dozens of individual wires for a high-dimensional tensor, they wrap bundles of parallel wires into a single colored tube. The "grade" of the tube (its integer size) represents the dimension of that axis.

The system works through several key mechanisms:

1. Nested Tubing: To represent a multi-axis tensor, you nest tubes. An $m \times n$ tensor is an $m$-grade tube wrapped around an $n$-grade tube. This captures the hierarchical structure of the data.
2. Dual Representation: The calculus maintains a "double view." The upper boundary of every tube recovers the undirected tensor-network view. Meanwhile, the directed interior retains the operational reading of the computation graph.
3. The Rewrite Engine: The core of the proof system is a single rule called grade-naturality, or "sliding spectacles." If you have a pointwise operation (like an activation function) inside a tube, you can "slide" that operation along the tube boundary. Mathematically, the operation commutes with the reindexing of the axes.
4. Padding as Partial Functions: To handle "ragged" data (tensors with varying sequence lengths), the authors treat padding as a way to extend mappings to partial functions (functions that may not be defined for all inputs). As shown in, padding allows us to represent irregular data structures within a consistent, dense framework.

Using these rules, complex identities can be reduced to short, visual derivations.

Numbers

The authors use the calculus to derive and validate major LLM optimization patterns. They report performance gains when applying these theoretical corollaries to real-world workloads. These were validated on an Apple M4 Max.

The most striking results come from the Mask-Augment Duality theorem. This proves that a masked attention mechanism is mathematically equivalent to performing unmasked attention on augmented data. This leads to several measurable wins:

• Prefix Caching: This technique reuses Key/Value states for shared prefixes. For queries sharing a common prefix, the paper predicts speedups for shared-context decoding. On a Hermes 4 70B model, they measured a 4.08x speedup for the DAG case. This is a significant improvement over naive implementations.
• Sequence Packing: This involves compacting disjoint "fibres" (individual documents) into dense slabs. The authors predict a speedup proportional to the data density. In synthetic tests with a density $\rho \approx 0.099$, they measured a 9.25x speedup. This demonstrates how much efficiency can be gained by packing multiple short sequences together.
• Sliding-Window Attention: This restricts each query to a local neighborhood. The calculus predicts massive efficiency gains for models like Mistral or Gemma. For a GPT-OSS 120B model with a window size of 128 and a sequence length of 8192, they measured a 32x speedup per sliding layer. This tracks closely with theoretical predictions in .

What's Missing

The theoretical framework is robust, but there are gaps. First, the calculus is currently limited to the structural fragment of tensor programming. It handles reshapes, repeats, and reductions. However, it does not yet formally cover non-linear functions or complex base algebras. You cannot use this to prove the correctness of an entire end-to-end training loop yet.

Second, the empirical validation is limited to Apple Silicon. The results on the M4 Max are impressive. However, we do not yet have a breakdown of how these "compiled displays" perform on NVIDIA CUDA kernels. The paper suggests lowering these abstractions to FlexAttention on CUDA as a follow-up.

Finally, the validation for sliding-window attention relied on synthetic data. Real-world memory bandwidth bottlenecks on a GPU might affect the theoretical speedups seen in synthetic benchmarks.

Should You Prototype This

Not yet.

This is a foundational paper, not a library release. The "calculus" is currently a mathematical tool for proving identities. It is not a compiler you can drop into a PyTorch project today. However, the implications for the next generation of tensor compilers are massive. If you work on low-level kernel optimization, the "Mask-Augment Duality" is a highly actionable insight. It provides a formal way to think about implementing sparsity via data augmentation. Keep an eye on this as it moves toward a practical tool for verifying einops rewrites.

Figures from the paper