It's All Just Vectorization: einx, a Universal Notation for Tensor Operations

2026-07-30

AI Summary

This work addresses the verbosity and error-proneness of NumPy-style APIs in existing tensor frameworks, as well as the limited generality of alternatives like einsum. The authors propose einx, a universal, declarative notation for tensor operations grounded in vectorization. By unifying diverse tensor computations through loop-like pointwise expressions, einx leverages vectorization as its core abstraction, combining declarative element-wise specifications with dimension lifting techniques. This approach reconstructs the entire tensor API using only a small set of primitive operations. The resulting einx library is compatible with major deep learning frameworks and integrates seamlessly into Python, substantially improving code readability and consistency while reducing cognitive load and shape-related errors.

Abstract

Tensor operations represent a cornerstone of modern scientific computing. However, the Numpy-like notation adopted by predominant tensor frameworks is often difficult to read and write and prone to so-called shape errors, i.a., due to following inconsistent rules across a large, complex collection of operations. Alternatives like einsum and einops have gained popularity, but are inherently restricted to few operations and lack the generality required for a universal model of tensor programming. To derive a better paradigm, we revisit vectorization as a function for transforming tensor operations, and use it to both lift lower-order operations to higher-order operations, and conceptually decompose higher-order operations to lower-order operations and their vectorization. Building on the universal nature of vectorization, we introduce einx, a universal notation for tensor operations. It uses declarative, pointful expressions that are defined by analogy with loop notation and represent the vectorization of tensor operations. The notation reduces the large APIs of existing frameworks to a small set of elementary operations, applies consistent rules across all operations, and enables a clean, readable and writable representation in code. We provide an implementation of einx that is embedded in Python and integrates seamlessly with existing tensor frameworks: https://github.com/fferflo/einx

Problem

• tensor operations
• shape errors
• vectorization
• notation
• universal programming

Innovation

• vectorization
• tensor operations
• einx
• declarative notation
• universal programming model