没有for循环的子阵列的点积

Question

当我们有：

array 1: A, shape (49998,3,3) 
array 2: B, shape (3, 49998)

我希望将它们的子阵列相乘得到

array 3: C, shape(3,49998)

我使用发电机：

def genC(A,B):
    for a,b in itertools.izip(A,B.T):
        c=np.dot(a,b)
        yield c.T[0]

C=np.array([c for c in genC()]).T

那么我怎样才能进行A，B的数组乘法，而不用for循环来得到数组C？

我试图使用np.tensordot，但我无法得到它

注意：

这只是一个基本的例子，对于原始数据中的一些测试cos，我有

4 * 3 * 37阵列A（500 000,3,3）B（3,500 000）

要做，对于循环sems对我来说不是pythonic方式xD

Answer 1

如果我的代码正确，你想用3矢量执行3x3矩阵的49998个点积，对吧？使用np.einsum非常容易：

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