WG211/M26Shan

Formalizing einsum with named tensor notation

Yafei Yang and Chung-chieh Shan, Indiana University

Einstein notation (einsum) is a widely adopted tensor operation that has proven useful in areas such as machine learning and quantum computing. But it had no formal semantics until recently, and its properties had no mechanized verification. We formalize the denotational semantics of a language of tensors (multi-dimensional arrays) that supports einsum and repeated tensors (expand). With this semantics, we verify equational rewriting rules about einsum and expand that improve the performance of tensor computation. Whereas typical languages use unnamed tensors, whose shape is a sequence of natural numbers, we found it advantageous for mechanization to use named tensors, whose shape is a finite map from names to natural numbers. To our knowledge, this is the first time that the semantics and rewriting rules of einsum are mechanized in a proof assistant. Our solution is agnostic to proof assistants; we have implemented it in Lean and partially in Agda.