Multiplying 3d tensor by another 3d tensor, dimension-wise

1

Hello!

My question is the following:

I have 2 3D tensors, and with one dimension in common.

1st tensor dimensions: a x b x c
2nd tensor dimensions: a x d x e

I want to multiply these two tensors along the dimension they have in common (a), and such that their output will be a 4 d tensor (e.g. dimensions b x c x d x e). Pretty much multiplying each vector (with a elements) of tensor 1 by each vector (with a elements) of tensor 2.

I can’t think of an easy built-in way to accomplish this. I can think of iterating over the dimensions and performing torch.dot, but that does not look very good. Does any one know of a better way?

Thank you for reading.

2

Hi Pedro!

torch.einsum is the general-purpose tool for “contracting” various
indices of tensors. Here is an example:

>>> import torch
>>> torch.__version__
'1.7.1'
>>> a = 3
>>> b = 2
>>> c = 5
>>> d = 3
>>> e = 4
>>> t1 = torch.randn ((a, b, c))
>>> t2 = torch.randn ((a, d, e))
>>> tt = torch.einsum ('ijk,ilm -> jklm', t1, t2)
>>> tt.shape
torch.Size([2, 5, 3, 4])

Best.

K. Frank

1 Like
3

Thank you, that’s exaclty what I wanted!