python - 没有 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

那么我如何在没有for循环的情况下在A，B内部进行数组乘法以获得数组C？

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

笔记：

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

4*3*37 数组 A(500 000,3,3) B(3,500 000)

要做，对于我来说，for loop sems不是pythonic方式xD

score 4 · Accepted Answer

如果我的代码正确，你想用 3 向量执行 3x3 矩阵的 49998 点积，对吗？这很容易做到np.einsum：

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

1 回答 1