矩阵乘以每对2-d阵列沿着第一维度与einsum

Question

我具有相同的尺寸a的两个3 d阵列和b

np.random.seed([3,14159]) 
a = np.random.randint(10, size=(4, 3, 2)) 
b = np.random.randint(10, size=(4, 3, 2)) 

---

print(a) 

[[[4 8] 
    [1 1] 
    [9 2]] 

[[8 1] 
    [4 2] 
    [8 2]] 

[[8 4] 
    [9 4] 
    [3 4]] 

[[1 5] 
    [1 2] 
    [6 2]]] 

---

print(b) 

[[[7 7] 
    [1 1] 
    [7 8]] 

[[7 4] 
    [8 0] 
    [0 9]] 

[[3 8] 
    [7 7] 
    [2 6]] 

[[3 1] 
    [9 3] 
    [0 5]]] 

我想从

a[0] 

[[4 8] 
[1 1] 
[9 2]] 

而且从b

b[0] 

[[7 7] 
[1 1] 
[7 8]] 

的第一个，返回此

a[0].T.dot(b[0]) 

[[ 92 101] 
[ 71 73]] 

但我想这样做，在整个第一个维度。我想我可以使用np.einsum

np.einsum('abc,ade->ace', a, b) 

[[[210 224] 
    [165 176]] 

[[300 260] 
    [ 75 65]] 

[[240 420] 
    [144 252]] 

[[ 96 72] 
    [108 81]]] 

这是正确的形状，但不是值。

我希望得到这样的：

np.array([x.T.dot(y).tolist() for x, y in zip(a, b)]) 

[[[ 92 101] 
    [ 71 73]] 

[[ 88 104] 
    [ 23 22]] 

[[ 93 145] 
    [ 48 84]] 

[[ 12 34] 
    [ 33 21]]] 

Answer 1

的矩阵乘法数额的，其中总和被在中轴线的乘积之和，因此指数b应该是两个阵列相同的：（即改变ade到abe）：

In [40]: np.einsum('abc,abe->ace', a, b) 
Out[40]: 
array([[[ 92, 101], 
     [ 71, 73]], 

     [[ 88, 104], 
     [ 23, 22]], 

     [[ 93, 145], 
     [ 48, 84]], 

     [[ 12, 34], 
     [ 33, 21]]]) 

---

当输入阵列具有缺少输出数组中的索引标， 他们被独立地总结。也就是说，

np.einsum('abc,ade->ace', a, b) 

相当于

In [44]: np.einsum('abc,ade->acebd', a, b).sum(axis=-1).sum(axis=-1) 
Out[44]: 
array([[[210, 224], 
     [165, 176]], 

     [[300, 260], 
     [ 75, 65]], 

     [[240, 420], 
     [144, 252]], 

     [[ 96, 72], 
     [108, 81]]]) 

Answer 2

这里有一个与np.matmul因为我们需要的a第二轴推回至年底，因此它会得到sum-reduced对从第二轴b，同时保持他们的第一轴对准 -

np.matmul(a.swapaxes(1,2),b) 

示意性地提出：

启动时：

a : M x N X R1 
b : M x N X R2 

与用于交换轴：

a : M x R1 X [N] 
b : M x [N] X R2 

括号内的轴得到sum-reduced，留给我们：

out : M x R1 X R2 

---

关于Python 3 。x，matmul由@ operator负责 -

a.swapaxes(1,2) @ b