einsum是做什么的？

假设我们有两个多维数组，A和B。现在假设我们要…

• 乘 A用B一种特殊的方式来创造新的产品阵列; 然后也许
• 沿特定轴求和该新数组；然后也许
• 以特定顺序转置新数组的轴。

有一个很好的机会，einsum可以帮助我们做到这一点更快，内存更是有效的NumPy的功能组合，喜欢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这两个轴应该相乘。换句话说，就像A数组B一样，我们将array 与array 的每一列相乘A[:, np.newaxis] * B。

• 注意，j在我们期望的输出中，它没有显示为标签；我们刚刚使用过i（我们想以一维数组结尾）。通过省略标签，我们告诉einsum来总结沿着这条轴线。换句话说，我们就像对行进行求和.sum(axis=1)。

基本上，这是你需要了解的所有信息einsum。玩一会会有所帮助；如果我们将两个标签都留在输出中，则会'i,ij->ij'返回2D产品数组（与相同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)。这是一张图片，显示了从函数获得的A和B和输出数组的标签：

你会看到j重复的标签-这意味着我们会将的行A与的列相乘B。此外，j输出中不包含标签-我们对这些产品求和。标签i和k被保留用于输出，因此我们得到一个2D数组。

这一结果与其中标签阵列比较可能是更加明显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)或np.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)