The Syntax and Semantics of einsum
This work addresses a foundational problem for researchers and developers in machine learning and related fields by formalizing a widely used but theoretically underdeveloped notation, though it is incremental in building upon existing practical implementations.
The paper tackles the lack of a solid theoretical basis and unification for the einsum notation used in tensor expressions across frameworks like NumPy, PyTorch, and TensorFlow, by providing a formal definition and proving equivalence rules to enable formal reasoning and optimization.
In 2011, einsum was introduced to NumPy as a practical and convenient notation for tensor expressions in machine learning, quantum circuit simulation, and other fields. It has since been implemented in additional Python frameworks such as PyTorch and TensorFlow, as well as in other programming languages such as Julia. Despite its practical success, the einsum notation still lacks a solid theoretical basis, and is not unified across the different frameworks, limiting opportunities for formal reasoning and systematic optimization. In this work, we discuss the terminology of tensor expressions and provide a formal definition of the einsum language. Based on this definition, we formalize and prove important equivalence rules for tensor expressions and highlight their relevance in practical applications.