将 numpy 3d 数组中的第 i 个二维矩阵与二维数组中的第 i 列相乘

Question

假设我有一个 3d 数组 A 和一个二维数组 B。A 具有维度 (s,m,m) 而 B 具有维度 (m,s)。

我想为二维数组 C 编写代码，其维度为 (m,s)，使得 C[:,i] = A[i,:,:] @ B[:,i]。

有没有一种方法可以优雅地做到这一点，而无需在 numpy 中使用 for 循环？

我想到的一种解决方案是将 B 重塑为具有维度 (m,s,1) 的 3d 阵列，通过 A@B 将 A 和 B 相乘，然后将生成的 3d 阵列重塑为二维阵列。 这听起来有点乏味，想知道是否可以在这里应用 tensordot 或 einsum。

建议表示赞赏。 谢谢！

Answer 1

在这里使用einsum很简单：

A = np.arange(18).reshape(2,3,3)
B = np.arange(6).reshape(3,2)

C = np.einsum("ijk,ki->ji",A,B)

for i in range(2):
   A[i]@B[:,i]==C[:,i]

# array([ True,  True,  True])
# array([ True,  True,  True])