假设您希望沿最后两个轴对每个元素执行矩阵乘法,我们可以使用
np.einsum
-
np.einsum('ijkl,jmkl->imkl',a,a)
样品运行验证-
In [43]: np.random.seed(0)
In [44]: a = np.random.rand(3,3,4,5)
In [45]: a[:,:,0,0].dot(a[:,:,0,0])
Out[45]:
array([[0.71750146, 1.17057872, 1.11135764],
[0.62938365, 0.86437796, 0.74541383],
[1.04636618, 1.62011127, 1.35483565]])
In [46]: np.einsum('ijkl,jmkl->imkl',a,a)[:,:,0,0]
Out[46]:
array([[0.71750146, 1.17057872, 1.11135764],
[0.62938365, 0.86437796, 0.74541383],
[1.04636618, 1.62011127, 1.35483565]])