How to do dot product along a given dimension

1

There are 2 tensors: q with dimension(64, 100, 500) and key with dimension(64, 500). I want to do dot product of key and q along the dimension of 500. Keras has a function dot() where we can give specific axes values. How can I achieve this in pytorch?

2

i think this is what you are looking for: bmm

3

thanks! but I tried bmm and what It says that the 2 tensors should be of the same exact dimension

4

x = Tensor(B,C,D)
y = Tensor(B,D)

If you want simplicity, you can just element-wise multiply them, then sum along that dimension:
(x * y.reshape(B,1,D)).sum(dim=2)

If you want to save memory, you can either use einsum():
torch.einsum('bcd,bd->bc', x, y)

or use bmm():
torch.bmm(x, y.rehsape(B,D,1)).reshape(B,C)

2 Likes