CODEWITHSUNDEEP
pythonpytorchtensordot-product
Jack Rolph
Jack Rolph

Reputation: 595

PyTorch: Row-wise Dot Product

Suppose I have two tensors:

a = torch.randn(10, 1000, 1, 4)
b = torch.randn(10, 1000, 6, 4)

Where the third index is the index of a vector.

I want to take the dot product between each vector in b with respect to the vector in a.

To illustrate, this is what I mean:

dots = torch.Tensor(10, 1000, 6, 1)
for b in range(10):
     for c in range(1000):
           for v in range(6):
            dots[b,c,v] = torch.dot(b[b,c,v], a[b,c,0])

How would I achieve this using torch functions?

Upvotes: 8

Views: 10980

Answers (1)

Seyed Sajad Ashrafi
Seyed Sajad Ashrafi

Reputation: 171

a = torch.randn(10, 1000, 1, 4)
b = torch.randn(10, 1000, 6, 4)

c = torch.sum(a * b, dim=-1)

print(c.shape)

torch.Size([10, 1000, 6])

c = c.unsqueeze(-1)
print(c.shape)

torch.Size([10, 1000, 6, 1])

Upvotes: 13

Related Questions