（注意：这个答案是基于我不久前写的一篇简短的博客文章。
）einsum

做什么einsum？

假设我们有两个多维数组，A并且B. 现在让我们假设我们想要…

• __ 以特定方式 相乘A以创建新的产品数组；B然后也许
• __沿特定轴对这个新数组 求和； 然后也许
• __以特定顺序 转置新数组的轴。

与 NumPy 函数（如 ,
）的组合相比，这很有可能einsum会帮助我们更快、更高效地执行此操作multiply，sum并且transpose将允许这样做。

如何einsum工作？

这是一个简单（但并非完全无关紧要）的示例。取以下两个数组：

A = np.array([0, 1, 2])

B = np.array([[ 0,  1,  2,  3],
              [ 4,  5,  6,  7],
              [ 8,  9, 10, 11]])

我们将乘以元素，A然后B沿新数组的行求和。在“正常”的 NumPy 中，我们会这样写：

>>> (A[:, np.newaxis] * B).sum(axis=1)
array([ 0, 22, 76])

所以在这里，索引操作A将两个数组的第一个轴对齐，以便可以广播乘法。然后对产品数组的行求和以返回答案。

现在，如果我们想einsum改用，我们可以这样写：

>>> np.einsum('i,ij->i', A, B)
array([ 0, 22, 76])

签名 字符串是这里的'i,ij->i'关键，需要稍微解释一下。你可以把它分成两半。在左侧（ 的左侧->），我们标记了两个输入数组。在
的右侧->，我们已经标记了我们想要结束的数组。

以下是接下来会发生的事情：

• A有一个轴；我们已经给它贴上了标签i。并且B有两个轴；我们将轴 0 标记为i，轴 1 标记为j。

• 通过在两个输入数组中 重复 标签i，我们告诉einsum这两个轴应该 相乘 。换句话说，我们将 array 与 arrayA的每一列相乘B，就像这样A[:, np.newaxis] * B做一样。

• 请注意，j它不会在我们想要的输出中显示为标签；我们刚刚使用过i（我们希望得到一个一维数组）。通过 省略 标签，我们告诉einsum沿该轴 求和。 换句话说，我们正在对产品的行求和，就像这样.sum(axis=1)做一样。

这基本上就是您使用einsum. 它有助于玩一点；如果我们将两个标签都留在输出中，'i,ij->ij'，我们将返回一个二维产品数组（与
相同A[:, np.newaxis] * B）。如果我们说没有输出标签，'i,ij->我们会返回一个数字（与做相同(A[:, np.newaxis] * B).sum()）。

然而，伟大的事情einsum是它不会首先构建一个临时的产品阵列。它只是对产品进行汇总。这可以大大节省内存使用。

一个稍微大一点的例子

为了解释点积，这里有两个新数组：

A = array([[1, 1, 1],
           [2, 2, 2],
           [5, 5, 5]])

B = array([[0, 1, 0],
           [1, 1, 0],
           [1, 1, 1]])

我们将使用 计算点积np.einsum('ij,jk->ik', A, B)。这是一张图片，显示了我们从函数中获得的Aand和输出数组的标签：B

您可以看到标签j重复了 - 这意味着我们将 的行A与 的列相乘B。此外，标签j不包含在输出中 -
我们正在对这些产品求和。标签i和k被保留用于输出，所以我们得到一个二维数组。

将此结果与标签未求和的数组进行比较可能会j更清楚 。 下面，在左侧，您可以看到写入结果的 3D 数组np.einsum('ij,jk->ijk', A, B)（即我们保留了 label j）：

求和轴j给出了预期的点积，如右图所示。

一些练习

为了获得更多的感觉einsum，使用下标符号来实现熟悉的 NumPy 数组操作会很有用。任何涉及乘法和求和轴组合的东西都可以使用 einsum.

设 A 和 B 是两个长度相同的一维数组。例如，A = np.arange(10)和B = np.arange(5, 15)。

• 的总和A可以写成：

  np.einsum('i->', A)
  

• 逐元素乘法 , A * B, 可以写成：

  np.einsum('i,i->i', A, B)
  

• 内积或点积np.inner(A, B)ornp.dot(A, B)可以写成：

  np.einsum('i,i->', A, B) # or just use 'i,i'
  

• 外积 ,np.outer(A, B)可以写成：

  np.einsum('i,j->ij', A, B)
  

对于 2D 数组C和D，前提是轴的长度兼容（长度相同或其中之一的长度为 1），以下是一些示例：

• C(主对角线之和)的迹线np.trace(C)可以写成：

  np.einsum('ii', C)
  

• , 的逐元素乘法C和转置D可以C * D.T写成：

  np.einsum('ij,ji->ij', C, D)
  

• 将 的每个元素乘以C数组D（形成一个 4D 数组），C[:, :, None, None] * D可以写成：

  np.einsum('ij,kl->ijkl', C, D)