Representing Tensors as Compositions for Efficient Computation on Sparse and Structured Data

Valérian Rey, Praneth Paruchuri

We propose a new mathematical formalism in which all kinds of tensors and of operations on those tensors are represented as simple function compositions. Tensors are functions from indices to values. An element-wise operation (e.g. an activation function) is a composition with a function applied to values; re-indexing (e.g. reshape, transposition) is a composition with a function applied to indices; and a reduction is a composition with both an index function, mapping each index to several indices, and a reduction operation (e.g. product, sum, max). Combined with tensor stacking, these compositions can represent most interesting kinds of tensors (e.g. dense, jagged, all kinds of sparsity patterns, matrix factorizations), as well as lazy versions of almost every operation available in modern tensor computation frameworks, including einsums between tensors of any kind. Identities between compositions can then be used to rewrite such expressions into equivalent ones, without evaluating them. For instance, the product of two sparse matrices can be rewritten so that only products between non-zero values are computed, turning a sparse problem into a dense one. Using this algebra as the…