对于许多数据科学家和开发者而言，

einsum
  • 1.

通常被视为NumPy文档中那个晦涩难懂的高级函数——功能强大但难以理解。不过一旦掌握其基本原理，

einsum
  • 1.

将成为Python科学计算生态系统中处理多维数组运算最为优雅高效的工具之一。它不仅语法简洁，表达力强，而且在众多应用场景中通常比常规方法更具计算效率。

本文将全面介绍

einsum()
  • 1.

函数——其数学基础、实现原理以及实际应用场景。我们将深入剖析其符号系统，通过实用示例展示其功能，探讨性能优化策略，并提供一个完整的参考速查表辅助实际应用。

无论您是深度学习研究人员、数值计算工程师，还是追求代码简洁高效的Python开发者，本文都将帮助您充分理解并有效应用

einsum()
  • 1.

函数。

1、爱因斯坦求和约定基础

我们从一个经典案例入手：矩阵乘法。在线性代数中，两个矩阵A和B相乘的标准定义是计算A中每一行与B中每一列的点积。图形化表示如下：

高效处理多维数组：einsum()函数从入门到精通_numpy

图1：标准矩阵乘法示意图。

注意观察，对于结果矩阵C中的每个元素，我们取A的第i行和B的第k列对应位置的元素相乘，然后对索引j求和。

这种元素级乘法和求和模式在张量运算中极为常见，尤其在阿尔伯特·爱因斯坦的广义相对论研究中。爱因斯坦为简化复杂张量表达式，提出了一种简洁表示法，省略显式的求和符号。其核心原则是：当一个索引在表达式中出现两次，则默认对该索引进行求和。

因此，对于矩阵乘法，我们可以表示为：

高效处理多维数组：einsum()函数从入门到精通_深度学习_02

图2：使用爱因斯坦表示法的矩阵乘法。

在NumPy中(后文将讨论其他框架如PyTorch、TensorFlow等)，

np.einsum()
  • 1.

函数允许直接应用这种表示法进行数组操作。下面是一个基本示例，展示了如何使用此函数进行矩阵乘法：

importnumpyasnp  
   
 A=np.random.randn(4,3)  
 B=np.random.randn(3,3)  
   
 C=np.einsum("ij,jk->ik", A, B)
  • 1.
  • 2.
  • 3.
  • 4.
  • 5.
  • 6.

该操作等效于

np.dot(A, B)
  • 1.

或

A @ B
  • 1.

，但

einsum()
  • 1.

的真正价值在于它能够将这一原理推广到更复杂的多维张量运算，这正是本文将要深入探讨的内容。

关于索引的技术说明：j是求和索引(也称为哑索引)。它可以是任何字母，不会改变表达式的实际含义，故称为"哑"索引。它在多个张量中出现，但不出现在结果中。而索引i和k称为自由索引，它们不参与求和运算，在每个张量中仅出现一次，并保留在输出结果中。

接下来，我们将详细分析

einsum()
  • 1.

的语法结构。

2、理解基本的 einsum() 语法

einsum()函数的语法初见时可能显得复杂，但一旦理解其逻辑，您会发现它是一个极为强大的数组操作工具。以下分析主要基于NumPy实现，但这些概念同样适用于PyTorch和TensorFlow等其他科学计算库，只有少量实现细节的差异。

本文将重点关注显式语法，因为从工程实践角度看，显式语法更不易出错且更具可读性。所谓显式语法，是指在einsum字符串中明确指定输出张量的索引结构。

虽然隐式语法更为简洁，但它缺乏灵活性，且容易因索引字母的顺序问题导致错误。这主要是个人偏好问题，但建议初学者先掌握显式语法。如需讨论隐式语法的细节，可在评论区交流。

以下是我们在矩阵乘法示例中使用的einsum字符串结构：

高效处理多维数组：einsum()函数从入门到精通_pytorch_03

图3：einsum()基本语法结构。

这张图清晰展示了einsum字符串的基本结构。现在让我们专注于理解这一字符串的核心含义。

在einsum中，输入和输出张量的每个维度都由一个索引字母表示。例如，3维张量需要3个字母，4维张量需要4个字母，依此类推。这些索引实质上代表着循环变量。在NumPy实现中，可以使用任何小写或大写字母作为索引，且在显式语法中，字母的顺序并不影响计算结果。

einsum()
  • 1.

执行的具体操作取决于索引在字符串中的出现位置和方式。主要有四种情况：

高效处理多维数组：einsum()函数从入门到精通_人工智能_04

图4：显式

einsum()
  • 1.

语法规则。

在张量计算中，“收缩”(contraction)这一术语经常与"求和"一起使用，但两者有细微差别。求和专指对特定索引进行规约操作，例如对索引j求和。而收缩是一个更广泛的概念，它包含求和操作，但同时也表示张量维度的减少。例如，矩阵乘法可视为在一个维度上的收缩。因此，每个求和操作都构成一种收缩，但并非所有收缩都是简单的求和。

为了建立直观理解，我们来分析几个基本示例：

外积

外积是线性代数中的基本运算，它通过将第一个向量的每个元素与第二个向量的每个元素相乘，生成一个二维矩阵。

高效处理多维数组：einsum()函数从入门到精通_pytorch_05

图5：

einsum()
  • 1.

表示的外积运算。

观察einsum字符串

"i,k->ik"
  • 1.

，我们发现输入中的所有索引都保留在输出中，因此没有执行求和操作。这表明第一个向量中的每个元素都与第二个向量中的每个元素相乘，形成结果矩阵。

批处理外积

批处理外积是外积的扩展应用，适用于同时计算多对向量的外积情况。

高效处理多维数组：einsum()函数从入门到精通_深度学习_06

图6：

<