https://user-images.githubusercontent.com/6318811/177030658-66f0eb5d-e136-44d8-99c9-86ae298ead5b.mp4

einops

灵活且强大的张量操作，让代码更易读、更可靠。
支持 NumPy、PyTorch、JAX、MLX 等多种框架。

最新更新：

• einops 演示环境 可以直接在浏览器中运行其中的 2 个示例笔记本
• 0.8.2：新增 MLX 后端
• 0.8.0：新增 Tinygrad 后端，并进行了一些小修复
• 0.7.0：无需繁琐配置即可使用 torch.compile，支持 Array API 标准，以及其他改进
• 10,000🎉：GitHub 报告称已有超过 1 万个项目在使用 einops
• einops 0.6.1：新增 Paddle 后端
• einops 0.6 引入了 打包与解包
• einops 0.5：einsum 现已成为 einops 的一部分
• Einops 论文 被 ICLR 2022 接受为口头报告（确实值得一读）。演讲录像已公开

推文

如果你还需要理由来花时间学习 einsum 和 einops... Tim Rocktäschel

使用 PyTorch 和 einops 编写更好的代码 👌 Andrej Karpathy

einops 正在慢慢渗透到我代码的每一个角落。如果你经常需要处理高维张量，它可能会改变你的工作方式。 Nasim Rahaman

更多评价

目录

• 安装
• 文档
• 教程
• API 微参考
• 为什么使用 einops
• 支持的框架
• 引用
• 仓库 和 讨论区

安装

非常简单：

pip install einops

(uv pip install einops 也同样适用)

教程

教程是了解 einops 实际应用的最佳方式

• 第一部分：einops 基础
• 第二部分：einops 用于深度学习
• 第三部分：打包与解包
• 第四部分：用 einops 改进 PyTorch 代码

Kapil Sachdeva 录制了一个简短的 einops 入门视频。

API

einops 拥有极简但功能强大的 API。

提供了三项核心操作（einops 教程展示了这些操作可以覆盖堆叠、重塑、转置、挤压/扩展、重复、平铺、拼接、视图以及多种约简操作）：

from einops import rearrange, reduce, repeat
# 按照指定模式重新排列元素
output_tensor = rearrange(input_tensor, 't b c -> b c t')
# 结合重新排列和约简操作
output_tensor = reduce(input_tensor, 'b c (h h2) (w w2) -> b h w c', 'mean', h2=2, w2=2)
# 沿新轴复制
output_tensor = repeat(input_tensor, 'h w -> h w c', c=3)

后来又增加了 pack 和 unpack 函数（比传统的堆叠/拆分/拼接更优）：

from einops import pack, unpack
# pack 和 unpack 可以将多个张量可逆地“打包”成一个张量。
# 被打包的张量可以具有不同的维度：
packed, ps = pack([class_token_bc, image_tokens_bhwc, text_tokens_btc], 'b * c')
class_emb_bc, image_emb_bhwc, text_emb_btc = unpack(transformer(packed), ps, 'b * c')

最后，einops 还提供了支持多字母变量名的 einsum：

from einops import einsum, pack, unpack
# einsum 就像……einsum，是一种通用且灵活的点积运算，
# 不过 1) 轴可以使用多字母名称 2) 模式放在最后 3) 支持多种框架
C = einsum(A, B, 'b t1 head c, b t2 head c -> b head t1 t2')

EinMix

EinMix 是一种通用的线性层，非常适合 MLP Mixer 等架构。

层

Einops 提供了层（einops 为每个框架保留单独的版本），这些层反映了相应的函数：

from einops.layers.torch      import Rearrange, Reduce
from einops.layers.tensorflow import Rearrange, Reduce
from einops.layers.flax       import Rearrange, Reduce
from einops.layers.paddle     import Rearrange, Reduce

from torch.nn import Sequential, Conv2d, MaxPool2d, Linear, ReLU
from einops.layers.torch import Rearrange

model = Sequential(
    ...,
    Conv2d(6, 16, kernel_size=5),
    MaxPool2d(kernel_size=2),
    # 无需编写 forward 方法即可完成展平
    Rearrange('b c h w -> b (c h w)'),
    Linear(16*5*5, 120),
    ReLU(),
    Linear(120, 10),
)

名称

einops 代表受爱因斯坦启发的操作记号法， 尽管“爱因斯坦操作”听起来更吸引人且更容易记住。

这种记号法大致受到爱因斯坦求和约定的启发（特别是 numpy.einsum 操作）。

为什么使用 einops 记号法？！

语义信息（明确期望）

y = x.view(x.shape[0], -1)
y = rearrange(x, 'b c h w -> b (c h w)')

虽然这两行代码在 某些 情境下执行相同的操作， 但第二行提供了关于输入和输出的信息。换句话说，einops 关注接口：输入和输出是什么，而不是 如何计算输出。

下一个操作看起来类似：

y = rearrange(x, 'time c h w -> time (c h w)')

但它给读者一个提示： 我们处理的不是独立的一批图像， 而是一个序列（视频）。

语义信息使代码更易于阅读和维护。

方便的检查

再看同一个例子：

y = x.view(x.shape[0], -1) # x: (batch, 256, 19, 19)
y = rearrange(x, 'b c h w -> b (c h w)')

第二行会检查输入是否具有四个维度， 你还可以指定具体的维度。 这与仅仅写形状注释不同，因为注释并不能防止错误， 而且如果没有经过代码审查，很容易过时。

y = x.view(x.shape[0], -1) # x: (batch, 256, 19, 19)
y = rearrange(x, 'b c h w -> b (c h w)', c=256, h=19, w=19)

结果严格确定

下面我们至少有两种方式来定义深度转空间操作：

# 深度转空间
rearrange(x, 'b c (h h2) (w w2) -> b (c h2 w2) h w', h2=2, w2=2)
rearrange(x, 'b c (h h2) (w w2) -> b (h2 w2 c) h w', h2=2, w2=2)

实际上还有至少四种其他方式。那么框架到底使用哪一种呢？

这些细节通常会被忽略，因为通常不会产生影响， 但在某些情况下（例如在下一阶段使用分组卷积时）可能会有很大差异， 而你希望在代码中明确指定这一点。

统一性

reduce(x, 'b c (x dx) -> b c x', 'max', dx=2)
reduce(x, 'b c (x dx) (y dy) -> b c x y', 'max', dx=2, dy=3)
reduce(x, 'b c (x dx) (y dy) (z dz) -> b c x y z', 'max', dx=2, dy=3, dz=4)

这些示例表明，我们不需要为 1D/2D/3D 池化分别定义不同的操作， 它们都可以用统一的方式定义。

空间转深度和深度转空间在许多框架中都有定义，那宽度转高度呢？如下所示：

rearrange(x, 'b c h (w w2) -> b c (h w2) w', w2=2)

框架无关的行为

即使是简单的函数，在不同框架中也有不同的定义：

y = x.flatten() # 或 flatten(x)

假设 x 的形状是 (3, 4, 5)，那么 y 的形状将是……

• NumPy、PyTorch、CuPy、Chainer、JAX：(60,)
• Keras、TensorFlow Layers、Gluon：(3, 20)

einops 在所有框架中都以相同的方式工作。

框架术语的独立性

例如，“tile”和“repeat”常常让人困惑。要沿宽度复制图像：

np.tile(image, (1, 2))    # 在 NumPy 中
image.repeat(1, 2)        # PyTorch 的 repeat 类似于 NumPy 的 tile

使用 einops，你就不需要去弄清楚是哪个轴被重复了：

repeat(image, 'h w -> h (tile w)', tile=2)  # 在 NumPy 中
repeat(image, 'h w -> h (tile w)', tile=2)  # 在 PyTorch 中
repeat(image, 'h w -> h (tile w)', tile=2)  # 在 TensorFlow 中
repeat(image, 'h w -> h (tile w)', tile=2)  # 在 JAX 中
repeat(image, 'h w -> h (tile w)', tile=2)  # 在 CuPy 中
... (等等)

用户评价提供了用户对同一问题的看法。

支持的框架

Einops 支持以下框架：

• NumPy
• PyTorch
• TensorFlow
• JAX
• CuPy
• Flax（社区版）
• Paddle（社区版）
• OneFlow（社区版）
• TinyGrad（社区版）
• PyTensor（社区版）

from einops import rearrange  
=> from einops.array_api import rearrange

但实际上更好：einops 可以与任何支持 Python 数组 API 标准 的框架一起使用， 例如：

• NumPy ≥ 2.0
• MLX # 是的，einops 可以与苹果的框架一起使用
• PyData/Sparse ≥ 0.15 # 并且支持稀疏张量
• Cubed # 还支持分布式张量
• Quantco/ndonnx
• JAX
• CuPy
• Dask 也可以通过 array-api-compat 来支持

开发

提供了 Devcontainer 环境，可以在本地、服务器上或 GitHub Codespaces 中使用。 要在 VS Code 中使用 Devcontainer，只需克隆仓库并点击“在 Devcontainer 中重新打开”。

从 einops 0.8.1 开始，einops 将测试作为包的一部分分发。

# pip install einops pytest
python -m einops.tests.run_tests numpy pytorch jax --pip-install

numpy pytorch jax 只是一个 示例，可以提供任何可测试框架的子集。 每个框架都会与 NumPy 进行对比测试，因此这是测试的前提条件。

指定 --pip-install 会在当前虚拟环境中安装依赖项， 如果本地已安装依赖，则应省略该选项。

构建和测试文档：

hatch run docs:serve  # 在 http://localhost:8000/ 上提供服务

引用 einops

请使用以下 BibTeX 条目：

@inproceedings{
    rogozhnikov2022einops,
    title={Einops: 基于爱因斯坦求和约定的清晰可靠的张量操作},
    author={Alex Rogozhnikov},
    booktitle={国际表示学习会议},
    year={2022},
    url={https://openreview.net/forum?id=oapKSVM2bcj}
}

支持的 Python 版本

einops 支持 Python 3.10 及更高版本。

Einops 快速上手指南

Einops 是一个灵活且强大的张量操作库，旨在编写可读性高、可靠性强的代码。它支持 NumPy、PyTorch、JAX、TensorFlow、MLX 等多种主流深度学习框架，通过统一的符号表示法简化张量的重塑、转置、重复和归约等操作。

环境准备

• 操作系统：Linux, macOS, Windows
• Python 版本：支持 Python 3.7+（推荐最新稳定版）
• 前置依赖：无需特殊系统依赖，只需安装目标深度学习框架（如 torch, tensorflow, jax 等）之一即可。

安装步骤

使用 pip 进行安装（推荐使用国内镜像源加速）：

pip install einops -i https://pypi.tuna.tsinghua.edu.cn/simple

或者使用 uv 工具安装：

uv pip install einops

基本使用

Einops 的核心 API 非常简洁，主要包含 rearrange（重排）、reduce（归约）和 repeat（重复）三个函数。

1. 导入库

from einops import rearrange, reduce, repeat

2. 张量重排 (Rearrange)

替代繁琐的 view, transpose, permute 等操作，直观定义输入输出维度。

# 将形状为 (t, b, c) 的张量转换为 (b, c, t)
# 模式字符串 't b c -> b c t' 清晰表达了维度变换逻辑
output_tensor = rearrange(input_tensor, 't b c -> b c t')

3. 组合重排与归约 (Reduce)

在一个操作中完成维度调整和数据聚合（如均值池化）。

# 对高度和宽度进行 2x2 的平均池化
# h2=2, w2=2 指定了池化核的大小
output_tensor = reduce(input_tensor, 'b c (h h2) (w w2) -> b h w c', 'mean', h2=2, w2=2)

4. 维度复制 (Repeat)

沿新轴复制数据，替代各框架中不一致的 tile 或 repeat 用法。

# 将灰度图 (h, w) 复制通道维变为 RGB 格式 (h, w, 3)
output_tensor = repeat(input_tensor, 'h w -> h w c', c=3)

5. 在 PyTorch 模型中使用

Einops 提供了可直接嵌入神经网络模型的 Layer 封装。

from torch.nn import Sequential, Conv2d, Linear, ReLU, MaxPool2d
from einops.layers.torch import Rearrange

model = Sequential(
    Conv2d(6, 16, kernel_size=5),
    MaxPool2d(kernel_size=2),
    # 无需手动编写 forward 函数进行 flatten，直接描述维度变化
    Rearrange('b c h w -> b (c h w)'),
    Linear(16*5*5, 120),
    ReLU(),
    Linear(120, 10),
)