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

Question

假设我们有两个 numpy 数组：Awith shape (n,p,q)，Bwith shape (n,q,r)。

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

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

但是是否存在可用于直接计算它的 numpy 函数？

标签： pythonarraysnumpy

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)