CODEWITHSUNDEEP
pythonnumpy
user39628
user39628

Reputation: 11

How to compute the scalar product of matrices as fast as possible

I want to compute (as fast as possible)

Av = v[0] * A[0, :, :] + ... + v[M-1] * A[M-1, :, :]

where v and A are np.ndarrays with shape (M,) and (M, N, N). Here's a minimal example what I have so far:

import numpy as np

N = 1000
M = 100

A = np.random.randint(-10, 10, size=(M, N, N))
v = np.random.randint(-10, 10, size=(M,))

%timeit np.sum(A * v[:, None, None], axis=0)

works as expected and gives

591 ms ± 7.42 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

However, I was wondering whether there's a faster way to calculate it?

Upvotes: 0

Views: 172

Answers (1)

joni
joni

Reputation: 7157

You can use np.einsum, see this question and the awesome answers for more details. Timing both on my machine yields:

In [57]: %timeit np.sum(A * v[:, None, None], axis=0)
441 ms ± 7.5 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

In [58]: %timeit np.einsum('ijk,i', A, v)
74.6 ms ± 725 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

Note that np.einsum('ijk,i', A, v) is the same as np.einsum('ijk,i -> jk', A, v).

Upvotes: 2

Related Questions