Tensor product problem

1

Hi,

I’m new to pytorch, and I have a problem here:

Suppose I have a tensor a with shape (seq_len, batch_size, feature_size), and another tensor b with shape (batch_size, feature_size). I want to product tensor a and b, so that I can obtain a tensor c with shape (seq_len, batch_size). The entry (i,j) to tensor c is the inner product of a[i, j, :] and b[j, :].

Is there any convenient way to calculate this?

Thanks

2

These approaches should work:

# Setup
seq, batch, feature = 2, 3, 4
a = torch.randn(seq, batch, feature)
b = torch.randn(batch, feature)

# Einsum
b_ = b.permute(1, 0).unsqueeze(0).expand(2, -1, -1)
out1 = torch.einsum('sbf, sfb->sb', a, b_)

# Mul + sum
out2 = (a * b).sum(-1)

# Manual
out3 = torch.zeros(seq, batch)
for i in range(seq):
    for j in range(batch):
        a_ = a[i, j]
        b_ = b[j]
        out3[i, j] = torch.dot(a_, b_)

# Check
print(torch.allclose(out1, out2) and torch.allclose(out1, out3))
> True

Depending on the workload one might be faster.

3

Seems like torch.einsum can do a lot of this kind of tensor operation.
Many thanks!