einops

einops是一个开源Python库，通过受爱因斯坦启发的记号提供灵活而强大的张量操作，旨在使张量处理代码更具可读性和可靠性。其名称代表Einstein-Inspired Notation for operations（受爱因斯坦启发的操作记号），大致受numpy.einsum启发。 核心API简约而强大，围绕三个主要函数构建：rearrange（按模式重新排列元素）、reduce（将重排与归约（如均值或最大值）结合）和repeat（沿新轴复制）。后续增加了pack和unpack，用于可逆地组合多个不同维度的张量，以及支持多字母轴名称的einsum。Einops还附带EinMix层，这是一个通用线性层，非常适合MLP Mixers及类似架构，此外还为PyTorch、TensorFlow、Flax和Paddle提供了特定框架的层包装器（Rearrange、Reduce等）。 该记号强调语义信息：无需编写x.view(x.shape[0], -1)，用户可写rearrange(x, 'b c h w -> b (c h w)')，同时记录输入和输出形状。命名轴如'time c h w -> time (c h w)'表明数据是视频序列而非独立批次。用户还可以固定特定维（如c=256, h=19, w=19），使操作在运行时验证形状，而不依赖注释。结果严格确定，因此像depth-to-space这样有歧义的操作也可无歧义地书写。同一语法统一适用于1D、2D和3D池化，并且在原生API不同的框架之间行为也一致（例如，NumPy/PyTorch与TensorFlow/Keras之间的flatten语义不同）。 Einops支持NumPy、PyTorch、TensorFlow、JAX、CuPy、Flax、Paddle、OneFlow、tinygrad和PyTensor，并能与任何实现Python数组API标准的框架配合使用，包括NumPy >= 2.0、MLX、pydata/sparse、cubed和nndonnx。它支持对PyTorch层使用torch.compile和torch.jit.script。该库要求Python 3.10或更高版本，可通过pip install einops安装。该项目在ICLR 2022上以口头报告形式展示，并被超过10,000个GitHub项目使用。