python - Matrix multiplication with extra dimensions in NumPy

Question

Given an N-by-L array A and an N-by-P-by-L array B, I would like to calculate the P-by-L array C where C[j, :] = sum(A[i, :] * B[i, j, :] for i in range(N)).

Mathematically speaking, this is somewhat equivalent of multiplying the matrix A from the left by B, where the elements of matrices are vectors of length L and the multiplication of two vectors is indeed defined as their point-wise multiplication.

A for-based solution to do the above calculation is:

C = np.zeros((P, L))
for i in range(N): 
    for j in range(P):
        C[j, :] += B[i] * A[i, j] 

Is there a 'vectorized' (and hopefully faster) way to do this in NumPy?

Note that in my application L is relatively large, so I would need a memory-efficient solution.

标签： pythonarraysnumpy

Answer 1

Following up from @hpaulj comment, you have different options using numpy. Most intuitive to me in this case are the following:

import numpy as np

np.random.seed(1)
N,L,P = 100,200,300

A = np.random.rand(N, L)
B = np.random.rand(N, P, L)

C1 = np.einsum('nl,npl->pl',A,B)
C2 = np.sum(A[:,None,:]*B, axis=0)

With numpy.einsum being (much) faster:

%timeit np.einsum('nl,npl->pl',A,B)
7.44 ms ± 520 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

%timeit np.sum(A[:,None,:]*B, axis=0)
46.1 ms ± 2.54 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)