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?

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

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)`

1 Like