Graphical einops: bridging tensor networks and computation graphs
This work provides a formal framework for tensor programming, which could benefit researchers and practitioners in deep learning by simplifying proofs and optimizing implementations.
This paper introduces a formal graphical calculus for tensor programming, enabling architecture diagrams to serve as proof-enabling tools rather than mere representations. It allows for short diagrammatic derivations of standard equivariance proofs and can convert attention masks into efficient pre-processing operations for sparse attention blocks.
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.