Numpy multiplication of matrix and 3 dimensional array

Question

I am trying to multiply ak by k matrix, let's say,

W=np.array([[W_11,...,W_1k],...,[W_k1,...W_kk]])

(where the W_ij are numbers) and a (k,m,m) multidimensional array, let's say,

A=np.array([A_1,...,A_k])

where A_i are m by m matrices.

If

A_i=[a_i]

where the a_i are numbers then the numpy.dot

C=np.dot(W,A) just yields the normal matrix vector product, ie C has shape (k,1) and one has that

C[i]=np.array([W_i1 a_1+W_i2 a_2+...W_ik*a_k])

What I would like to know is what is the best way to multiply W and A where now A is not necessarily a vector, ie A_i are m by m matrices in such a way that it mimics the product as if A_i=[a_i], ie I would like C=np.dot(W,A) to have shape (k,m,m) and C[i] should be the m by m matrix

W_i1 A_1+...W_ik A_k

Of course I can do this with a loop but I'm looking for an efficient solution.

Answer 1

From what I understand from your question, you can use numpy.einsum :

C = np.einsum('ij,jkl->ikl', W, A)

Should do the job.