$ show arogozhnikov/einops
Einops is a Python library that provides a uniform, readable, and reliable API for tensor manipulation across multiple deep learning frameworks (NumPy, PyTorch, TensorFlow, JAX, and others). It uses Einstein-like notation (e.g., 'b c h w -> b h w c') to express axis reordering, composition, decomposition, reduction, and repetition in a single, explicit string. The library is built around a core transformation engine that parses axis patterns into a cached 'TransformRecipe', which is then applied via a backend abstraction layer. This architecture allows einops to support 10+ frameworks with minimal code duplication, while also providing framework-specific layer wrappers (Rearrange, Reduce, EinMix) for seamless integration into model definitions. The library is designed for production and research use, with a focus on reducing errors from manual view/permute/transpose operations.
The LLM identified 4 files worth highlighting for reuse:
Defines an abstract backend interface (AbstractBackend) and concrete implementations for NumPy, PyTorch, TensorFlow, JAX, CuPy, OneFlow, Paddle, TinyGrad, PyTensor, and MLX. Each backend implements a small set of primitive operations (from_numpy, to_numpy, reduce, transpose, reshape, tile, concat, einsum) and is lazily loaded to avoid import overhead. This pattern enables einops to work across 10+ frameworks with a single code path, and can be extracted as a standalone adapter library for any tensor-manipulation tool.
Contains the core transformation logic: TransformRecipe encapsulates a parsed axis pattern into a sequence of elementary operations (permutations, reductions, reshapes, added axes). The recipe is built by _prepare_transformation_recipe and cached via lru_cache for performance. Functions like rearrange, reduce, repeat, and einsum delegate to this engine. The pattern is self-contained and could be extracted to build a generic tensor transformation DSL.
Implements ParsedExpression and AnonymousAxis to parse Einstein-like axis notation strings into an intermediate representation. Handles named axes, anonymous axes (using numbers), ellipsis, and parenthesized compositions. This parser is independent of any tensor framework and could be reused in any project that needs to interpret axis transformation patterns.
Provides mixin classes (RearangeMixin, ReduceMixin, _EinmixMixin) that implement the logic for framework-specific layer wrappers (e.g., einops.layers.torch.Rearrange). Each mixin handles pickling, weight initialization, and forward pass logic, while the concrete layer classes only need to implement parameter creation and the actual tensor operation. This pattern cleanly separates framework-agnostic logic from framework-specific boilerplate, making it easy to add support for new frameworks.