将 3D 矩阵与 2D 矩阵相乘，返回与 2D 矩阵具有相同维度的矩阵

Question

在 Numpy 中，假设我有一个尺寸为 ixjxk 的矩阵 A，和一个尺寸为 kx i 的矩阵 B。 如何在不使用循环的情况下取回尺寸为 kxi 的产品？
例如：

 Let         A = [[[1 2],
                   [3 4]],
                  [[5 6],
                   [7 8]]]

             B = [[a b],
                  [c d]]

我想得到：

             C = [[a+2c  5b+6d],
                  [3a+4c 7b+8d]]

我目前的解决方案是

np.diagonal(np.dot(A, B), axis1=0, axis2=2)

但是，问题是我正在处理一个大数据集（ A和B有很大的尺寸），所以np.dot(A, B)会导致MemoryError 。 因此，我想找出一种更好的方法来解决这个问题，而无需计算点积。

我已经研究过像einsum和tensordot这样的函数，但我没有找到我需要的东西（或者我可能错过了一些东西）。 如果有人可以帮助我，我将不胜感激。 谢谢！

Answer 1

您可以通过分析输入和 output 维度来做到这一点。

A.shape -> i, j, k
B.shape -> k, i

你的例子不是很好，但如果你仔细看看你要求 output 是什么， B的形状必须与A的第一个和最后一个尺寸相匹配。

总和减少发生在A的最后一个轴和B的第一个轴上：

np.einsum('ijk,ki->ji', A, B)

einsum可能令人生畏，但进行此类分析可以为您节省大量因反复试验而导致的挫败感。

Answer 2

看着你的 output， C[:,0]是

A[0] @ B[0].T

同样， C[:,1]是A[1] @ B[1].T 。

考虑到这一点， np.einsum公式为：

C = np.einsum('ijk,ik->ji', A,B)