Having some trouble to vectorize these loops ([N, T, K] and [K, C] to [N, T, C])

Hello !

I’m currently having some trouble with a quite simple task, so here I am.

I have a tensor A of size [N, T, K], and a tensor B of size [K, C].
I want to compute a tensor Y of size [N, T, C] according to the following algorithm:

for i in range(N):
    for j in range(T):
        for l in range(C):
            Y[i, j, l] = 0
            for k in range(K):
                Y[i, j, l] += A[i, j, k] * B[k, l]

I’ve tried many strange thing (always ending with a “RuntimeError: the size of … must match the size of …”), but cannot find an efficient way to vectorize this.
I feel like it’s a perfect job for einsum, but didn’t find the correct notation.

Thanks for reading !


torch.matmul(A, B) will do the job :slight_smile:

Oh goood, that was so easy, thank you !

By the way, for future reader : einsum notation is torch.einsum('ijk,kl->ijl', (A, B))