Selecting entries of a pytorch tensor with another tensor

user_6209598 0
Viewed 134

I have a tensor a with with float entries and torch.Size([64,2]) and I also have a tensor b with torch.Size([64]). The entries of b are only 0 or 1.

I would like to get a new tensor c with torch.Size([64]) such that c[i] == a[i,b[i]] for every index i. How can I do that?

My attempt I tried with torch.gather but without success. The following code gives me RuntimeError: Index tensor must have the same number of dimensions as input tensor

import torch
a = torch.zeros([64,2])
b = torch.ones(64).long()
torch.gather(input=a, dim=1,index=b)

Any help will be highly appreciated!

tensor pytorch indexing python
2 Answers

You can perform this straight with an indexing of a on both dimensions:

  • On dimension=0: a "sequential" indexing using torch.arange.

  • On dimension=1: indexing using b.

All in all, this gives:

>>> a[torch.arange(len(a)), b]

Alternatively you can use torch.gather, the operation you are looking for is:

# c[i] == a[i,b[i]]

The provided gather operation when applied on dim=1 provides something like:

# c[i,j] == a[i,b[i,j]]

As you can see, we need to account for the difference in shapes between between a and b. To do so, you can unsqueeze a singleton dimension on b (annotated with the letter j above) such than #b=(64, 1), for instance with b.unsqueeze(-1) or b[...,None]:

>>> a.gather(dim=1, index=b[...,None]).flatten()
user_63313690

Not sure I understand your question, but I think ou can loop over your tensor

a = torch.zeros([64,2])
b = torch.ones(64).long()
c = torch.empty([64])
for i, _ in enumerate(a):
  c[i] = a[i,b[i]]
c
user_49943520
Related