Einops

Enables reshaping and transposing tensors across NumPy, PyTorch, TensorFlow, JAX, and other frameworks using a unified Einstein-inspired notation (e.g., 'batch height width channels -> batch (height width) channels'). The implementation uses a two-stage compilation pipeline: ParsedExpression extracts axis names and composite axes from pattern strings, then TransformRecipe generates optimized backend-specific transformation instructions. Dual-level LRU caching (256 recipe entries, 1024 shape entries) eliminates recompilation overhead for repeated operations.

Performs reductions (sum, mean, max, min) along specified dimensions using named axes in Einstein notation (e.g., 'batch height width channels -> batch channels' reduces over height and width). The pattern parser identifies which axes to reduce, and the backend layer translates this into framework-specific reduction operations. Runtime validation ensures all named axes in the pattern match the input tensor's dimensions, preventing silent reduction errors that occur with positional indexing.

Implements support for the Array API standard, enabling einops to work with any framework that implements the Array API specification (NumPy 2.0+, PyTorch, TensorFlow, JAX, etc.). This provides a path toward true framework independence by relying on standardized array operations rather than framework-specific APIs. The implementation detects Array API compliance and uses standard operations when available, falling back to framework-specific implementations when necessary.

Includes an extensive test infrastructure that validates einops operations across all supported frameworks (NumPy, PyTorch, TensorFlow, JAX, MLX) with systematic shape testing, edge case coverage, and numerical correctness verification. The test suite uses parameterized tests to cover combinations of frameworks, tensor shapes, and operation types, ensuring consistent behavior across backends. CI/CD pipelines run tests on multiple Python versions and framework versions to catch compatibility issues early.

Replicates tensor data along new or existing dimensions using Einstein notation (e.g., 'batch height width -> batch height width repeat_count' repeats along a new axis). The pattern parser identifies which axes are new (appear in output but not input) and generates backend-specific repeat/broadcast instructions. This avoids manual broadcasting and explicit repeat calls, providing a declarative alternative to framework-specific APIs like torch.repeat or tf.tile.

Parses Einstein notation patterns to extract axis names, composite axes (e.g., '(height width)'), and ellipsis operators, then validates that the pattern matches the input tensor's shape at runtime. The ParsedExpression class decomposes patterns into semantic components, and the validation layer checks that all named axes have consistent dimensions across input and output. This prevents silent shape mismatches and provides clear error messages when patterns are invalid.

Compiles patterns into optimized TransformRecipe objects that encode the exact transformation steps, then caches recipes using a 256-entry LRU cache to avoid recompilation on repeated operations. The caching layer operates at two levels: recipe caching (pattern → transformation instructions) and shape caching (1024 entries) for frequently seen tensor shapes. This architecture eliminates parsing and compilation overhead for operations that use the same pattern multiple times, critical for performance in training loops.

Automatically detects the input tensor's framework (NumPy, PyTorch, TensorFlow, JAX, MLX, etc.) and dispatches operations to the appropriate backend implementation without user configuration. The backend abstraction layer wraps framework-specific operations (reshape, transpose, reduce, etc.) with a unified interface, enabling identical einops code to execute on any supported framework. This design eliminates the need for framework-specific imports or conditional logic in user code.

Provides Einstein summation notation for complex tensor contractions (e.g., 'ij,jk->ik' for matrix multiplication) with automatic translation to framework-specific einsum implementations. The pattern parser converts Einstein notation into the appropriate backend einsum call, supporting implicit and explicit summation semantics. This unifies einsum syntax across frameworks that have different einsum APIs (PyTorch vs TensorFlow vs NumPy).

Provides pack() and unpack() operations to efficiently handle variable-length sequences by concatenating them into a single tensor with a companion packed_shapes array, then reconstructing the original sequences. This avoids padding and masking overhead for operations on ragged/variable-length data. The pack operation flattens multiple tensors along a specified axis and records shape information, while unpack reverses this process using the recorded shapes.

Provides drop-in neural network layer implementations (einops.layers.torch.Rearrange, einops.layers.tensorflow.Rearrange, etc.) that apply einops operations within model architectures. These layers wrap rearrange, reduce, and repeat operations as nn.Module subclasses, enabling einops patterns to be used directly in model definitions without explicit function calls. The EinMix layer provides learnable linear transformations along specified axes, combining einops rearrangement with weight matrices.

Enables einops operations to be compiled and optimized by PyTorch's torch.compile() mechanism, allowing patterns to be fused with surrounding operations for improved performance. The implementation ensures that einops operations are compatible with PyTorch's symbolic tracing and graph compilation, enabling end-to-end model compilation without einops being a compilation bottleneck. This is particularly valuable for production inference where compilation overhead is amortized across many forward passes.