einops

einops é uma biblioteca Python de código aberto que fornece operações de tensor flexíveis e poderosas por meio de uma notação inspirada em Einstein, projetada para tornar o código de manipulação de tensores mais legível e confiável. O nome significa Notação Inspirada em Einstein para operações, vagamente inspirada em numpy.einsum. A API principal é minimalista, porém poderosa, construída em torno de três funções principais: rearrange (reorganiza elementos de acordo com um padrão), reduce (combina reorganização com redução, como média ou máximo) e repeat (copia ao longo de um novo eixo). Adições posteriores incluem pack e unpack para combinar reversivelmente múltiplos tensores de diferentes dimensionalidades, e einsum com suporte a nomes de eixos com múltiplas letras. O einops também inclui uma camada EinMix, uma camada linear genérica adequada para MLP Mixers e arquiteturas semelhantes, além de wrappers de camadas específicos para cada framework (Rearrange, Reduce, etc.) para PyTorch, TensorFlow, Flax e Paddle. A notação enfatiza informações semânticas: em vez de escrever x.view(x.shape[0], -1), os usuários escrevem rearrange(x, 'b c h w -> b (c h w)'), que documenta tanto as formas de entrada quanto as de saída. Eixos nomeados como 'time c h w -> time (c h w)' indicam que os dados são uma sequência de vídeo em vez de um lote independente. Os usuários também podem fixar dimensões específicas (por exemplo, c=256, h=19, w=19) para que a operação valide as formas em tempo de execução em vez de depender de comentários. Os resultados são estritamente determinados, portanto operações ambíguas como depth-to-space podem ser escritas sem ambiguidade. A mesma sintaxe funciona uniformemente em pooling 1D, 2D e 3D, e o comportamento é consistente entre frameworks, mesmo onde as APIs nativas diferem (por exemplo, a semântica de flatten varia entre NumPy/PyTorch e TensorFlow/Keras). Einops suporta NumPy, PyTorch, TensorFlow, JAX, CuPy, Flax, Paddle, OneFlow, tinygrad e PyTensor, e funciona com qualquer framework que implemente o padrão Python array API, incluindo NumPy >= 2.0, MLX, pydata/sparse, cubed e nndonnx. Ele suporta torch.compile e torch.jit.script para camadas do PyTorch. A biblioteca requer Python 3.10 ou posterior e é instalada via pip install einops. O projeto foi apresentado oralmente no ICLR 2022 e é usado por mais de 10.000 projetos no GitHub.