python - 如何沿第一维在两个 3D 张量之间执行矩阵乘法？

Question

我想计算两个 3D 张量沿第一维的点积。我尝试了以下 einsum 表示法:

import numpy as np

a = np.random.randn(30).reshape(3, 5, 2)
b = np.random.randn(30).reshape(3, 2, 5)

# Expecting shape: (3, 5, 5)
np.einsum("ijk,ikj->ijj", a, b)

遗憾的是它返回了这个错误:

ValueError: einstein sum subscripts string includes output subscript 'j' multiple times

在用 np.tensordot 求和失败后，我选择了爱因斯坦求和。非常欢迎提出想法和跟进问题!

Answer 1

您的两个尺寸 5 和 5 不对应于相同的轴。因此，您需要使用两个不同的下标来指定它们。例如，您可以这样做:

>>> res = np.einsum('ijk,ilm->ijm', a, b)

>>> res.shape
(3, 5, 5)

---

请注意，您还需要更改大小为 2 和 2 的轴的下标。这是因为您正在计算批量外积(即我们同时在两个轴上迭代)，而不是点积(即我们在两个轴上同时迭代).

• 外积:

  >>> np.einsum('ijk,ilm->ijm', a, b)
  

• 下标 k 上的点积，即 a 和 axis=1 的 axis=2 b 的 >:

  >>> np.einsum('ijk,ikm->ijm', a, b)
  

  相当于a@b。