Einops: Clear and Reliable Tensor Manipulations with Einstein-like Notation
Alex Rogozhnikov
Abstract
Tensor computations underlie modern scientific computing and deep learning.A number of tensor frameworks emerged varying in execution model, hardware support, memory management, model definition, etc.However, tensor operations in all frameworks follow the same paradigm.Recent neural network architectures demonstrate demand for higher expressiveness of tensor operations.The current paradigm is not suited to write readable, reliable, or easy-to-modify code for multidimensional tensor manipulations. Moreover, some commonly used operations do not provide sufficient checks and can break a tensor structure.These mistakes are elusive as no tools or tests can detect them.Independently, API discrepancies complicate code transfer between frameworks.We propose einops notation: a uniform and generic way to manipulate tensor structure, that significantly improves code readability and flexibility by focusing on the structure of input and output tensors.We implement einops notation in a Python package that efficiently supports multiple widely used frameworks and provides framework-independent minimalist API for tensor manipulations.
Ask about this paper
Ask your agent about it.
Lune has read the top-tier papers around this one, so every answer names the papers it rests on.
Your agent calls
Lunesearch_papers
Free to start. No credit card required.
Terminal
Install the CLIlune papers get 13fef035-d30f-43b8-9363-e42696b408e0Cited by top-tier papers32
- L4GM: Large 4D Gaussian Reconstruction ModelJiawei Ren, Cheng Xie, Ashkan Mirzaei, Hanxue Liang et al.NeurIPS 2024 · 173 citations
- Permutation Equivariant Neural FunctionalsAllan Zhou, Kaien Yang, Kaylee Burns, Adriano Cardace et al.NeurIPS 2023 · 84 citations
- From Zero to Turbulence: Generative Modeling for 3D Flow SimulationMarten Lienen, David Lüdke, Jan Hansen-Palmus, Stephan GünnemannICLR 2024 · 56 citations
- Progress measures for grokking via mechanistic interpretabilityNeel Nanda, Lawrence Chan, Tom Lieberum, Jess Smith et al.ICLR 2023 · 54 citations
- Vivid-ZOO: Multi-View Video Generation with Diffusion ModelBing Li, Cheng Zheng, Wenxuan Zhu, Jinjie Mai et al.NeurIPS 2024 · 48 citations
Related papers
- It's All Just Vectorization: einx, a Universal Notation for Tensor OperationsFlorian Fervers, Sebastian Bullinger, Christoph Bodensteiner, Michael ArensICLR 2026 · 1 citation
- EinDecomp: Decomposition of Declaratively-Specified Machine Learning and Numerical Computations for Parallel ExecutionDaniel Bourgeois, Zhimin Ding, Dimitrije Jankov, Jiehui Li et al.VLDB 2025 · 4 citations
- A Simple and Efficient Tensor CalculusSören Laue, Matthias Mitterreiter, Joachim GiesenAAAI 2020 · 40 citations
- FreeTensor: a free-form DSL with holistic optimizations for irregular tensor programsShizhi Tang, Jidong Zhai, Haojie Wang, Lin Jiang et al.PLDI 2022 · 16 citations
- EINNET: Optimizing Tensor Programs with Derivation-Based TransformationsLiyan Zheng, Haojie Wang, Jidong Zhai, Muyan Hu et al.OSDI 2023 · 19 citations
