# Graphical einops: bridging tensor networks and computation graphs

arXiv: https://arxiv.org/abs/2605.31485 · HTML: https://mlsanity.com/paper/2605.31485/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 17 of its 34 references that are in the map.
- First posted on arXiv 29 May 2026.

## Builds on (most cited first, 17 of 17)

- [Efficient Memory Management for Large Language Model Serving with PagedAttention](https://mlsanity.com/paper/2309.06180/) (2023)
- [Self-Consistency Improves Chain of Thought Reasoning in Language Models](https://mlsanity.com/paper/2203.11171/) (2022)
- [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://mlsanity.com/paper/2205.14135/) (2022)
- [Tree of Thoughts: Deliberate Problem Solving with Large Language Models](https://mlsanity.com/paper/2305.10601/) (2023)
- [FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning](https://mlsanity.com/paper/2307.08691/) (2023)
- [Efficient Streaming Language Models with Attention Sinks](https://mlsanity.com/paper/2309.17453/) (2023)
- [Big Bird: Transformers for Longer Sequences](https://mlsanity.com/paper/2007.14062/) (2020)
- [Fast Inference from Transformers via Speculative Decoding](https://mlsanity.com/paper/2211.17192/) (2022)
- [Generating Long Sequences with Sparse Transformers](https://mlsanity.com/paper/1904.10509/) (2019)
- [SGLang: Efficient Execution of Structured Language Model Programs](https://mlsanity.com/paper/2312.07104/) (2023)
- [Transformer Dissection: A Unified Understanding of Transformer's Attention via the Lens of Kernel](https://mlsanity.com/paper/1908.11775/) (2019)
- [Transformer Language Models without Positional Encodings Still Learn Positional Information](https://mlsanity.com/paper/2203.16634/) (2022)
- [Hydragen: High-Throughput LLM Inference with Shared Prefixes](https://mlsanity.com/paper/2402.05099/) (2024)
- [FlashMask: Efficient and Rich Mask Extension of FlashAttention](https://mlsanity.com/paper/2410.01359/) (2024)
- [An introduction to graphical tensor notation for mechanistic interpretability](https://mlsanity.com/paper/2402.01790/) (2024)
- [Neural Circuit Diagrams: Robust Diagrams for the Communication, Implementation, and Analysis of Deep Learning Architectures](https://mlsanity.com/paper/2402.05424/) (2024)
- [Tape Diagrams for Monoidal Monads](https://mlsanity.com/paper/2503.22819/) (2025)

## Abstract (arXiv)

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.