假设我有一个3d数组A和一个2d数组B。A的尺寸为(s,m,m),而B的尺寸为(m,s)。
我想为尺寸为(m,s)的二维数组C编写代码,使得C [:,i] = A [i,:::] @ B [:,i]。
有没有一种方法可以优雅地做到这一点,而无需在numpy中使用for循环?
我想到的一种解决方案是将B重塑为尺寸为(m,s,1)的3d阵列,通过A @ B将A和B相乘,然后将所得3d阵列重塑为2d阵列。这听起来有点乏味,并且想知道是否可以在此处应用tensordot或einsum。
建议表示赞赏。谢谢!
einsum
这里的用法很简单:
A = np.arange(18).reshape(2,3,3)
B = np.arange(6).reshape(3,2)
C = np.einsum("ijk,ki->ji",A,B)
for i in range(2):
A[i]@B[:,i]==C[:,i]
# array([ True, True, True])
# array([ True, True, True])
非常感谢!这看起来非常干净高效。