Tensor contraction

How can I do tensor contractions with pytorch? The most similar operation that I found is torch.matmul, but actually it is a batched matrix product rather than a contraction. I cannot choose the indicies over which the contraction is done with torch.matmul. Eigen::Tensor has the operation. I also found the operation in numpy whereas it is slightly restricted than the one in Eigen::Tensor. I’m looking for a corresponding one in pytorch.

1 Like

There was a einsum PR but that got abandoned (https://github.com/pytorch/pytorch/pull/2008). However, stay tuned, there will be something released soon that can do this.

Thank you, SimonW. If possible, would you provide a link to the related commits or discussions?

Unfortunately no. For now you can use a series of .transpose(_, __) and .view(*) calls and use matmul as a workaround :slight_smile:

FYI, this will be helpful https://research.fb.com/announcing-tensor-comprehensions/. Pytorch integration is coming soon.