如何沿某个特定轴相乘两个 3d numpy 数组？

Question

假设我们有两个 numpy 数组：A，形状为 (n,p,q)，B，形状为 (n,q,r)。

如何将它们相乘以获得形状为(n,p,r) 的数组C？我的意思是保留轴 0 并将它们乘以轴 1 和 2。

我知道它可以通过以下方式计算： C = np.stack([np.dot(a[i], b[i]) for i in range(A.shape[0])])

但是有没有可以直接计算的numpy函数呢？

Answer 1

我觉得你可以np.einsum:

# sample data
n,p,q,r = 2,3,4,5
A = np.zeros((n,p,q))
B = np.zeros((n,p,r))

out = np.einsum('npq,nqr->npr',A,B)

out.shape
# (2, 3, 5)