了解NumPy的einsum

怎么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)

问题来源于stack overflow