我有两个多维NumPy数组A和B,其中A.shape =(K,d,N)和B.shape =(K,N,d).我想在轴0(K)上执行逐元素操作,该操作是在轴1和2(d,N和N,d)上的矩阵乘法.因此结果应该是具有C.shape =(K,d)的多维数组C,因此C [k] = np.dot(A [k],B [k]).一个天真的实现看起来像这样：

C = np.vstack([np.dot(A[k],B[k])[np.newaxis,:,:] for k in xrange(K)])

但这种实施很慢.稍微快一点的方法看起来像这样：

C = np.dot(A,B)[:,:]

它在多维数组上使用np.dot的默认行为,给我一个形状为(K,K,d)的数组.但是,这种方法计算所需答案K次(沿轴2的每个条目都相同).渐渐地,它将比第一种方法慢,但开销要小得多.我也知道以下方法：

from numpy.core.umath_tests import matrix_multiply
C = matrix_multiply(A,B)

但我不能保证这个功能可用.因此,我的问题是,NumPy是否提供了一种有效执行此操作的标准方法？一般来说,适用于多维数组的答案是完美的,但仅针对这种情况的答案也是很好的.

编辑：正如@Juh_所指出的,第二种方法是错误的.正确的版本是：

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

但是增加的开销使得它比第一种方法慢,即使对于小矩阵也是如此.对于小型和大型矩阵,最后一种方法是在我的所有时序测试中获胜.我现在正在考虑使用这个,如果没有更好的解决方案,即使这意味着将numpy.core.umath_tests库(用C语言编写)复制到我的项目中.