How to perform Matrix multiplication along dimension

1

I tried to figure it out myself but I just don’t get the summation right. I have an (n ,d^2,m) Tensor A and an (n ,d , m ) tensor B. I want the following summation result C (with size (n ,d , m) )

C[i,j,k] = A[i,:,k]B[i,:,k][j]

where the last product is meant as a matrix multiplication (the second dimension has to split into 2 of size d). How do I achieve this?

2

Nevermind got it. The following does the job:

A=A.reshape(n,d,d,m)
C=torch.einsum('bijk,bjk->bik', A,B)