I created this toy problem that reflects my much bigger problem:
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. ]]
"""
I want to do it as fast as possible, so using numpy's functions to calculate ans
should be the best approach, since this operation is heavy and my matrices are quite big.
I saw this post, but the shapes are different and I cannot understand which axes
I should use for this problem. However, I'm certain that tensordot should have the answer. Any suggestions?
EDIT: I accepted @ajcr's answer, but please read my own answer as well, it may help others...
You could use
np.einsum
to do the operation since it allows for very careful control over which axes are multiplied and which are summed:The function multiplies the entries in the first axis of
ind
with the entries in the first axis ofdist
(subscript'i'
). Ditto for the second axis of each array (subscript'j'
). Instead of returning a 3D array, we tell einsum to sum the values along axis'j'
by omitting it from the output subscripts, thereby returning a 2D array.np.tensordot
is more difficult to apply to this problem. It automatically sums the products of axes. However, we want two sets of products but to sum only one of them.Writing
np.tensordot(ind, dist, axes=[1, 1])
(as in the answer you linked to) computes the correct values for you, but returns a 3D array with shape(3, 4, 3)
. If you can afford the memory cost of a larger array, you could use:This gives you the correct result, but because
tensordot
creates a much larger-than-necessary array first,einsum
seems to be a better option.