numpy einsum: Elementwise product between 3D matrix and 2D matrix

1
votes

I have two numy matrices :

  1. A with shape (N,M,d)
  2. B with shape (N,d)

So, I am trying to get a Matrix with shape (N,M,d) in such manner that I do element wise product between B and each element of A (which are M elements). I used numpy's einsum as follows :

product = np.einsum('ijk,ik->ijk', A, B)

I GET THE GOOD shape but I suspect that the operation is wrong and that I am not doing element wise product as intended.

Is my use of einsum correct ?

1
pythonnumpynumpy-einsum
You can test this in about 2 seconds with a pair of toy matrices. Voting to close because you didn't look at the data. - Mad Physicist
Also, use A * B[:, None, :] if you're not sure. - Mad Physicist
@MadPhysicist, made the same edit to my answer as well at the same time. - Akshay Sehgal

1 Answers

3
votes

I am trying to get a Matrix with shape (N,M,d) in such manner that I do element wise product between B and each element of A (which are M elements).

The operation you are trying to perform is a broadcasted element wise product over axis = 1 of A (M size) -

C1 = A * B[:,None,:]

You can always do a quick test to check whether what you are expecting, as your result, is actually what you have implemented.

A = np.random.randint(0,5,(2,3,1)) # N,M,d
B = np.random.randint(0,5,(1,1))   # N,d

C2 = np.einsum('ijk,ik->ijk', A, B)
print(C2)

[[[4]
  [0]
  [8]]

 [[4]
  [4]
  [4]]]

To double check whether both operations are equal =

np.allclose(C1, C2)

##True

More details on how np.einsum works can be found here.