python：有效地將矩陣堆棧的切片i乘以矩陣的列i

Question

import numpy as np
M,N,R = 2,3,4

# For a 3-dimensional array A:
A = np.reshape(np.r_[0:M*N*R], [M,N,R], order = 'C')

# and 2-dimensional array B:
B = np.reshape(np.r_[0:M*R], [R,M], order = 'C')

我想N*M矩陣，從乘以切片導致i的A通過柱i的B 。 我已經嘗試過np.dot和np.einsum ，但是無法獲得我所需要的。

有人可以幫忙嗎？ 謝謝！

Answer 1

使用np.einsum ，我們將-

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

讓我們使用給定的樣本並使用np.dot與矩陣相乘來驗證結果-

In [35]: A.shape
Out[35]: (2, 3, 4)

In [36]: B.shape
Out[36]: (4, 2)

In [37]: A[0].dot(B[:,0])
Out[37]: array([ 28,  76, 124])

In [38]: A[1].dot(B[:,1])
Out[38]: array([226, 290, 354])

In [39]: np.einsum('ijk,ki->ji',A,B)
Out[39]: 
array([[ 28, 226],
       [ 76, 290],
       [124, 354]])

對於有關何時使用方面einsum了dot-based的工具，如np.dot / np.tensordot ，這里有一個related post 。