einops

einops es una biblioteca de Python de código abierto que proporciona operaciones tensoriales flexibles y potentes mediante una notación inspirada en Einstein, diseñada para hacer que el código de manipulación de tensores sea más legible y confiable. El nombre significa Notación Inspirada en Einstein para operaciones, inspirada libremente en numpy.einsum. La API central es minimalista pero potente, construida alrededor de tres funciones principales: rearrange (reorganiza elementos según un patrón), reduce (combina la reorganización con reducción como media o máximo) y repeat (copia a lo largo de un nuevo eje). Adiciones posteriores incluyen pack y unpack para combinar reversiblemente múltiples tensores de diferente dimensionalidad, y einsum con soporte para nombres de ejes de múltiples letras. Einops también incluye una capa EinMix, una capa lineal genérica adecuada para MLP Mixers y arquitecturas similares, además de envoltorios de capas específicos del framework (Rearrange, Reduce, etc.) para PyTorch, TensorFlow, Flax y Paddle. La notación enfatiza la información semántica: en lugar de escribir x.view(x.shape[0], -1), los usuarios escriben rearrange(x, 'b c h w -> b (c h w)'), que documenta tanto las formas de entrada como de salida. Ejes con nombre como 'time c h w -> time (c h w)' sugieren que los datos son una secuencia de video en lugar de un lote independiente. Los usuarios también pueden fijar dimensiones específicas (por ejemplo, c=256, h=19, w=19) para que la operación valide las formas en tiempo de ejecución en lugar de depender de comentarios. Los resultados están estrictamente determinados, por lo que operaciones ambiguas como depth-to-space pueden escribirse sin ambigüedad. La misma sintaxis funciona uniformemente en pooling 1D, 2D y 3D, y el comportamiento es consistente entre frameworks incluso donde las API nativas difieren (por ejemplo, la semántica de flatten varía entre NumPy/PyTorch y TensorFlow/Keras). einops soporta NumPy, PyTorch, TensorFlow, JAX, CuPy, Flax, Paddle, OneFlow, tinygrad y PyTensor, y funciona con cualquier framework que implemente el estándar de API de arrays de Python, incluyendo NumPy >= 2.0, MLX, pydata/sparse, cubed y nndonnx. Soporta torch.compile y torch.jit.script para capas de PyTorch. La biblioteca requiere Python 3.10 o posterior y se instala mediante pip install einops. El proyecto se presentó como presentación oral en ICLR 2022 y es utilizado por más de 10,000 proyectos en GitHub.