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 !