`numpy.einsum` 是如何工作的？

Question

根据爱因斯坦求和来编写求和的正确方法对我来说是一个难题，所以我想在我的代码中尝试一下。我在少数情况下取得了成功，但主要是通过反复试验。

现在有一个我想不通的案例。首先，一个基本问题。对于两个矩阵A 和B，它们分别是Nx1 和1xN，AB 是NxN，但BA 是1x1。当我想用np.einsum 计算NxN 的情况时，我可以这样做：

import numpy as np

a = np.asarray([[1,2]])
b = np.asarray([[2,3]])
print np.einsum('ij,ji->ij', a, b)

最后的数组是2x2。不过

a = np.asarray([[1,2]])
b = np.asarray([[2,3]])
print np.einsum('ij,ij->ij', a, b)

返回一个1x2 数组。我不太明白为什么这不能给出正确的结果。 例如，对于上述情况，numpy 的指南说箭头可用于强制求和或阻止它发生。但这对我来说似乎很模糊。在上述情况下，我不明白 numpy 如何根据索引的顺序（显然会发生变化）决定输出数组的最终大小。

形式上我知道以下内容：当箭头右侧没有任何内容时，可以将数学求和写为 $\sum\limits_{i=0}^{N}\sum\limits_{j=0 }^{M} A_{ij}B_{ij}$ 对于np.einsum('ij,ij',A,B)，但是当有箭头时，我不知道如何用正式的数学表达式来解释它。

Answer 1

In [22]: a
Out[22]: array([[1, 2]])
In [23]: b
Out[23]: array([[2, 3]])
In [24]: np.einsum('ij,ij->ij',a,b)
Out[24]: array([[2, 6]])
In [29]: a*b
Out[29]: array([[2, 6]])

在这里，包括输出在内的所有部分中索引的重复被解释为逐个元素的乘法。什么都没有总结。 a[i,j]*b[i,j] = c[i,j] 代表所有 i,j。

In [25]: np.einsum('ij,ji->ij',a,b)
Out[25]: 
array([[2, 4],
       [3, 6]])
In [28]: np.dot(a.T,b).T
Out[28]: 
array([[2, 4],
       [3, 6]])
In [38]: np.outer(a,b)
Out[38]: 
array([[2, 3],
       [4, 6]])

再次没有求和，因为相同的索引出现在左侧和右侧。 a[i,j]*b[j,i] = c[i,j]，换句话说：

[[1*2, 2*2],
 [1*3, 2*3]]

实际上是一种外部产品。看看a 如何针对b.T 进行广播可能会有所帮助：

In [69]: np.broadcast_arrays(a,b.T)
Out[69]: 
[array([[1, 2],
        [1, 2]]), 
 array([[2, 2],
        [3, 3]])]

在语句的左侧，重复的索引指示乘以哪些维度。左右两边的匹配决定了它们是否相加。

np.einsum('ij,ji->j',a,b) # array([ 5, 10]) sum on i only
np.einsum('ij,ji->i',a,b) # array([ 5, 10]) sum on j only
np.einsum('ij,ji',a,b) # 15 sum on i and j

不久前，我开发了一个与einsum 等效的纯 Python，主要关注它如何解析字符串。目标是创建一个nditer，用它来计算产品的总和。但这不是一个简单的脚本，即使在 Python 中也是如此：

https://github.com/hpaulj/numpy-einsum/blob/master/einsum_py.py

---

显示这些求和规则的更简单的序列：

In [53]: c=np.array([[1,2],[3,4]])
In [55]: np.einsum('ij',c)
Out[55]: 
array([[1, 2],
       [3, 4]])
In [56]: np.einsum('ij->i',c)
Out[56]: array([3, 7])
In [57]: np.einsum('ij->j',c)
Out[57]: array([4, 6])
In [58]: np.einsum('ij->',c)
Out[58]: 10

---

使用没有1 维度的数组可以消除广播的复杂性：

In [71]: b2=np.arange(1,7).reshape(2,3)
In [72]: np.einsum('ij,ji',a2,b2)
...
ValueError: operands could not be broadcast together with remapped shapes [original->remapped]: (2,3)->(2,3) (2,3)->(3,2) 

或者我应该说，它暴露了尝试的广播。

省略号为 einsum 解释增加了一定程度的复杂性。当我解决了... 的使用中的一个错误时，我开发了上面提到的 github 代码。但我并没有花太多精力来完善文档。

Ellipsis broadcasting in numpy.einsum

---

当您想要一个可以处理各种大小数组的表达式时，省略号最有用。如果您的数组始终是 2D 的，则它不会做任何额外的事情。

例如，考虑dot 的泛化，将A 的最后一个维度与B 的倒数第二个维度相乘。使用省略号，我们可以编写一个表达式来处理 2d、3D 和更大数组的混合：

np.einsum('...ij,...jk',np.ones((2,3)),np.ones((3,4)))  # (2,4)
np.einsum('...ij,...jk',np.ones((5,2,3)),np.ones((3,4)))  # (5,2,4)
np.einsum('...ij,...jk',np.ones((5,2,3)),np.ones((5,3,4))) # (5,2,4)
np.einsum('...ij,...jk',np.ones((5,2,3)),np.ones((7,5,3,4)))  # (7,5,2,4)
np.einsum('...ij,...jk->...ik',np.ones((5,2,3)),np.ones((7,5,3,4)) # (7, 5, 2, 4)

最后一个表达式使用默认右侧索引...ik、省略号加上非求和索引。

你原来的例子可以写成

np.einsum('...j,j...->...j',a,b)

有效地填充i（或更多维度）以匹配数组的维度。

如果 a 或 b 是 1d 也可以：

np.einsum('...j,j...->...j',a,b[0,:])

---

np.dot 泛化到更大维度的方式不同

dot(a, b)[i,j,k,m] = sum(a[i,j,:] * b[k,:,m])

在einsum中表示为：

np.einsum('ijo,kom->ijkm',np.ones((2,3,4)),np.ones((3,4,2)))

可以概括为

np.einsum('...o,kom->...km',np.ones((4,)),np.ones((3,4,2)))

或

np.einsum('ijo,...om->ij...m',np.ones((2,3,4)),np.ones((3,4,2)))

但我认为我无法在 einsum 中完全复制它。也就是说，我不能告诉它为A 填写索引，然后为B 填写不同的索引。