2

I have a 2d matrix A[1000*90] and B[90*90*1000]

I would like to calculate C[1000*90]


For i in range(1000)
  C[i,:]=np.matmul(A[i,:],B[:,:,i]

I understand if I use a vectorized formula it's going to be faster, seems like einsum might be the function I'm looking for, but I am having trouble cyphering the syntax of einsum. Is it np.einsum(ij,jki->ik,A,B)?

Julien
  • 13,986
  • 5
  • 29
  • 53
Walterwang201112
  • 41
  • 5
  • 1
    You are right, einsum is what you want. Not expert enough with it myself to tell you by eye if your attempt is correct (looks decent to me). But you could easily convince yourself with a few small examples... – Julien Oct 07 '21 at 03:52
  • 1
    `matmul` handles 'batches' of multiplications. Didn't you see that in the docs? You may need to add dimensions and transpose so `i` is the first of 3 dimensions. – hpaulj Oct 07 '21 at 05:05
  • @Walterwang201112 IMO your `einsum` looks alright. Why didn't you "just try it"? Is there a problem with it? – André Oct 07 '21 at 07:26
  • your einsum is correct, but @hpaulj way is more efficient as it dispatches directly to linked BLAS library (if any). –  Oct 07 '21 at 08:59

1 Answers1

2

Your einsum is correct. But there is a better way as pointed out by hpaulj.

Using Matmul:

import numpy as np
A =np.random.rand(1000,90)
B = np.random.rand(90,90,1000)
C = A[:,np.newaxis,:]@B.transpose(2,0,1) ## Matrix multiplication
C = C = C.reshape(-1,C.shape[2])
np.array_equal(C,np.einsum('ij,jki->ik',A,B)) # check if both give same result
  • Thanks, I have to say the Matmul is really difficult to wrap my head around. I have no idea why you added a new axis to A, transposed B, and reshaped C the way you did, especially with 3d matrix matmul. In terms for performance, both method are about the same, with matmul slightly faster for smaller matrix, and einsum faster for larger matrix. – Walterwang201112 Oct 07 '21 at 13:54
  • 1
    Einsum is faster because you did not link a BLAS library to Numpy. You can check it out in the following link https://stackoverflow.com/questions/21671040/link-atlas-mkl-to-an-installed-numpy It definitely gives you good boost when dealing with larger matrices. you need to perform a multiplication of 1000 matrices from each array. To do that you have to match the dimensions of A and B, so that matmul does the for loop for you. The dimension of A is 2 and and B is 3. so to make it same, I added a new dimension so that you have 1000 (1,90) matrices in A and 1000 (90,90) matrices in B –  Oct 07 '21 at 15:44
  • 1
    When using matmul, the matrix multiplication is performed over last two dimensions, to I have to transpose B so that 1000 becomes my first dimension. Now, A[:,newaxis,:] has (1000,1,90) and B.transpose(2,0,1) has (1000,90,90) dimensions. Now you can perform the matmul directly so that it does the for loop over 1000 automatically for you. after matmul, the dimension is (1000,1,90). I reshaped it to make it (1000,90) –  Oct 07 '21 at 15:47