python - 矩阵将沿第一维的每对二维数组与 einsum 相乘

Question

我有两个大小相同的 3-D 数组 a 和 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中获取第一个数组

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]]])