I have two 3d matrices X and Y both of which have the shape (5, 1825, 77). I'd like to do five 2d matrix multiplications, i.e., X[i, :, :]@Y[i, :, :].T without using a for loop. Is there a way to do this in numpy?
This is interesting for those (like me) who try to avoid for loops at any cost:
shape = 5, 1825, 77
X = np.random.random(shape)
Y = np.random.random(shape)
p_for = np.empty((shape[0], shape[1], shape[1]))
for i in range(shape[0]):
p_for[i] = X[i] @ Y[i].T
p_matmul = X @ np.moveaxis(Y, -1, 1)
assert np.allclose(p_for, p_matmul)
p_einsum = np.einsum("ijk,ilk->ijl", X, Y)
assert np.allclose(p_for, p_einsum)
The tree methods produce the same result, but, as @JérômeRichard points out:
%%timeit
prod = np.empty((shape[0], shape[1], shape[1]))
for i in range(5):
prod[i] = X[i, :, :] @ Y[i, :, :].T
50.4 ms ± 7.18 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
%timeit X @ np.moveaxis(Y, -1, 1)
115 ms ± 1.6 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
%timeit np.einsum("ijk,ilk->ijl", X, Y)
544 ms ± 3.22 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)