Selection/Filter with indices using pytorch

Question

I have a 3-dimensional numpy array, for example:

 x = [[[0.3, 0.2, 0.5],
       [0.1, 0.2, 0.7],
       [0.2, 0.2, 0.6]]]

The indices array is also 3-dimensional, like:

indices = [[[0],
            [1],
            [2]]]

I expect the output is:

 output= [[[0.3],
           [0.2],
           [0.6]]]

I tried the torch.index_select and torch.gather function, but I couldn't find a right way to deal with the dimension. Thanks for any help!

Umang Gupta · Accepted Answer · 2019-01-03 20:48:28Z

1

How about using x.gather(dim=2, indices)? That works for me.

answered Jan 3, 2019 at 20:48

Umang Gupta

16.6k7 gold badges51 silver badges69 bronze badges

Sign up to request clarification or add additional context in comments.

2 Comments

yes, it works!!! Thank you! I must I lost my mind after trying too much!

@beepretty No worries! Happens with most of us :)

beepretty · Accepted Answer · 2019-01-03 20:26:33Z

0

I found the answer. Please let me know if there is any better solution.

torch.cat([torch.index_select(a.view(1, -1), 1, i.view(1, -1)[0]) 
                                         for a, i in zip(x, indices)])

answered Jan 3, 2019 at 20:26

beepretty

1,1274 gold badges15 silver badges21 bronze badges