Pytorch/Python-如何使用索引张量获取结果张量

Question

Given a 1-D tensor:

A = torch.tensor([1, 2, 3, 4])

suppose we have some "index tensor"

ind1 = torch.tensor([3, 0, 1])
ind2 = torch.tensor([[3, 0], [1, 2]])

as we run A[ind1] & A[ind2]

we get tensor([4, 1, 2]) & tensor([[4, 1],[2, 3]]) which is the same shape of index tensor and its value are mapped from tensor A.

Now if I have ND tensor, I want to get a result tensor using an index map, how can I do that?

Answer 1

Let's take a minimal multi-dimensional example with A being a 2-D tensor defined as:

>>> A = torch.tensor([[1, 2, 3, 4],
                      [5, 6, 7, 8]])

From your description of the problem:

the same shape of index tensor and its value are mapped from tensor A .

Then the indexer tensor would require to have the same shape as the indexed tensor A , since this one is no longer flat. Otherwise, what would the result of A (shaped (2, 4) ) indexed by ind1 (shape (3,) ) be?

If you are indexing on a single dimension then you can utilize torch.gather :

>>> A.gather(1, ind2)
tensor([[4, 1],
        [6, 7]])