Search code examples
pythonnumpymatrix-multiplication

Element-wise multiply of two numpy array of different shapes


I have two numpy arrays F and C of dimensions NxM and MxB. How can I get a matrix with elements Result[n,m,b] = F[n,m]*C[m,b] with dimensions NxMxB in an efficient manner?


Solution

  • While the solution by @Nick Odell comments works and is fast, it is perhaps not very readable. Another option is to use numpy.einsum, which allows you to explicitely state which dimensions have to be multiplied and which ones must be summed. In your case, you could do

    import numpy as np
    
    N = 20
    M = 40
    B = 10
    
    F = np.random.rand(N,M)
    C = np.random.rand(M,B)
    
    Result = np.einsum('nm,mb->nmb', F, C) 
    
    print(Result.shape) # (20,40,10)
    

    Here 'nm,mb->nmb' means:

    • nm: first matrix has dimensions with subscripts n and m respectively,
    • mb: second matrix has dimenisons with subscripts m and b respectively,
    • -->nmb: multiply n x m x b, do not sum over any subscript.