Einops is a tensor manipulation library that provides a framework-agnostic interface for reshaping, Einstein summation, and multi-dimensional array operations. It serves as an abstraction layer that works across NumPy, PyTorch, TensorFlow, and JAX, allowing for tensor transformations without changing the API.
The library distinguishes itself through a declarative notation system that uses readable string patterns to describe tensor rearrangements and reductions. This approach includes an extended Einstein summation interface that supports multi-letter axis names and a named dimension mapping system to eliminate manual integer indexing.
The toolkit covers a broad range of capabilities including tensor packing, unpacking, and element repetition. It provides utilities for dimension reduction, axis rearrangement, and input shape validation. Additionally, it offers pluggable layer wrappers that integrate these transformations and linear mixing operations directly into sequential neural network model definitions.
The library integrates tensor manipulation functions into compiled execution graphs to maintain runtime performance.
