python - 将高阶矩阵与numpy相乘

我创建了这个玩具问题，反射(reflect)了我更大的问题:

import numpy as np

ind = np.ones((3,2,4)) # shape=(3L, 2L, 4L)
dist = np.array([[0.1,0.3],[1,2],[0,1]]) # shape=(3L, 2L)

ans = np.array([np.dot(dist[i],ind[i]) for i in xrange(dist.shape[0])]) # shape=(3L, 4L)
print ans

""" prints:
   [[ 0.4  0.4  0.4  0.4]
    [ 3.   3.   3.   3. ]
    [ 1.   1.   1.   1. ]]
"""

我想尽快执行此操作，因此使用numpy函数计算ans应该是最好的方法，因为此操作很繁琐且矩阵很大。

我看到了this post，但是形状不同，因此我无法理解应该为该问题使用哪种axes。但是，我确定tensordot应该有答案。有什么建议么？

编辑:我接受了@ajcr's answer，但也请阅读我自己的答案，它可能会对其他人有所帮助...

最佳答案

您可以使用np.einsum进行操作，因为它可以非常小心地控制要乘和求和的轴:

>>> np.einsum('ijk,ij->ik', ind, dist)
array([[ 0.4,  0.4,  0.4,  0.4],
       [ 3. ,  3. ,  3. ,  3. ],
       [ 1. ,  1. ,  1. ,  1. ]])

该函数将ind的第一个轴上的条目与dist的第一个轴上的条目(下标'i')相乘。每个数组第二个轴的同上(下标'j')。代替返回3D数组，我们告诉einsum通过从输出下标中省略沿'j'轴的值来求和，从而返回2D数组。
np.tensordot更难以应用于此问题。它会自动对轴的乘积求和。但是，我们需要两组产品，但只能将其中一组相加。

编写np.tensordot(ind, dist, axes=[1, 1])(如您所链接的答案)会为您计算正确的值，但会返回形状为(3, 4, 3)的3D数组。如果您负担得起较大阵列的内存成本，则可以使用:

np.tensordot(ind, dist, axes=[1, 1])[0].T

这样可以为您提供正确的结果，但是由于tensordot首先会创建一个比所需数组大得多的数组，因此einsum似乎是一个更好的选择。