Tensormultiplication with einsum

Question

I have a tensor phi = np.random.rand(n, n, 3) and a matrix D = np.random.rand(3, 3). I want to multiply the matrix D along the last axis of phi so that the output has shape (n, n, 3). I have tried this

np.einsum("klj,ij->kli", phi, D)

But I am not confident in this notation at all. Basically I want to do

res = np.zeros_like(phi)           
for i in range(n):
    for j in range(n):
        res[i, j, :] = D.dot(phi[i, j, :])

Answer 1

You are treating phi as an n, n array of vectors, each of which is to be left-multiplied by D. So you want to keep the n, n portion of the shape exactly as-is. The last (only) dimension of the vectors should be multiplied and summed with the last dimension of the matrix (the vectors are implicitly 3x1):

np.einsum('ijk,lk->ijl', phi, D)

OR

np.einsum('ij,klj->kli', D, phi)

It's likely much simpler to use broadcasting with np.matmul (the @ operator):

np.squeeze(D @ phi[..., None])

You can omit the squeeze if you don't mind the extra unit dimension at the end.