优雅的Numpy张量产品

Question

我需要在numpy (或py手电筒)中取两个张量的乘积：

我有过

A = np.arange(1024).reshape(8,1,128)
B = np.arange(9216).reshape(8, 128, 9)

想要得到C，用点积加在A (axis=2)的最后一个昏暗和B的中间暗(axis=1)上。这应该有维度8x9。目前，我正在做：

C = np.zeros([8, 9])
for i in range(8):
    C[i,:] = np.matmul(A[i,:,:], B[i,:,:])

如何优雅地做这件事？

我试过：

np.tensordot(weights, features, axes=(2,1)).

但是它返回8x1x8x9。

Answer 1

一种方法是使用numpy.einsum。

C = np.einsum('ijk,ikl->il', A, B)

或者你可以使用广播矩阵乘法。

C = (A @ B).squeeze(axis=1)
# equivalent: C = np.matmul(A, B).squeeze(axis=1)