einops

收录理由

张量代码里到处是琐碎的形状整理，reshape、transpose、squeeze 常常要连成一串，einops 正是很多做模型的人用来简化这类操作的库。它允许你用爱因斯坦记法写出一个可读的表达式，这个表达式在 NumPy、PyTorch、JAX、TensorFlow 和 MLX 上都能直接运行，不用改动。对于先在一种框架里做原型、再到另一种框架部署的团队，这种可移植性很有价值。einops 还提供了类似 einsum 的归约操作，以及用于通道和特征混合的 EinMix 层，构建模型时涉及的线性代数运算，很大一部分都能用它覆盖。新手能从清晰的语法和现成的例子里快速上手，而凡是被静默的形状不匹配坑过的人，都会喜欢它早早指出错误的报错信息。