3d matrix multiplication into multiple 2d matrix multiplications-CodePudding

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?

CodePudding user response：

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)