Pytorch get "reduced" tensor by indices

Question

I have a tensor a = torch.arange(6).reshape(2,3), and another tensor b=(torch.rand(a.size())> 0.5).int().nonzero().

I want to create a new tensor that contains only values from a of the indices that are indicated by b.

For example:

a = torch.arange(6).reshape(2,3) # tensor([[0, 1, 2],
                                 #         [3, 4, 5]])
b = (torch.rand(a.size())> 0.5).int().nonzero() # tensor([[0, 1],
                                                #         [0, 2],
                                                #         [1, 0],
                                                #         [1, 1]])

The desired output is:

tensor([1,2,3,4])

I know that I can iterate over the values of b and access those values in a as indices but I wanted to know if there is a better Pytorch way to to this (using tensor operations only).

** The shape of the output tensor doesn't really matter, I just need to have a tensor with only the values indicated by b.

Answer 1

If I understand you correctly, you can do:

a[b[:,0], b[:,1]]

This will produce a 1D tensor with the values at the indices specified by b. Note that the output might not be the same from run to run of the program since the indices are selected nondeterministically.

---

If you don't know the number of dimension in advance, you'll need to use map() to generate the desired slices:

a[tuple(map(lambda x: b[:,x], range(a.dim())))]