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 !