Python: vectorization over numpy array

Question

X and Y are both 3d arrays with dimensions (a,b,c). My goal is to do a dot product.

Consider that case where index i and j are scalar, and (X[i,:,j].T).dot(Y[i,:,j]) would be simple and return a scalar.

However, if I try to do vectorization, i and j become 1d arrays, and (X[i,:,j].T).dot(Y[i,:,j]) return a matrix but I am expecting a 1d array as result. How do I get around this problem ?

• Aly Hosny: you can use np.einsum to calculate the product. if you can, add some random values with expected output
• bjornsing: it would be nice if you could give us coherent code

Answer 1

Naive implementation using list comprehension:

a,b,c = X.shape
r1 = [(X[i,:,j].T).dot(Y[i,:,j]) for i in range(a) for j in range(c)]

Implementation using np.einsum:

r2 = np.einsum('ijk,ijk->ik', X,Y).flatten()