arXiv · 2605.31485
Graphical einops: bridging tensor networks and computation graphs
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.
Explore related subjects
Keep this discovery
Explore connections, maps & timelines
Vincent Wang-Maścianica, Nikhil Khatri. 2026-05-29. Graphical einops: bridging tensor networks and computation graphs. https://arxiv.org/abs/2605.31485
Cite the original work for its findings. Save a collection to share your selection of sources.