Select specific columns of each row in a torch Tensor

1

There’s probably a simple way to do this, but owing to my noobness I do not know who to do this in PyTorch.

Basically, let’s say I have a torch tensor like so:

m = Variable(torch.randn(4,2))

Furthermore, I have a bunch of indices, given by inds, where inds is:

inds
Variable containing
 1
 1
 0
 0
 [torch.LongTensor of size 4]

What I would like to do, is use the elements of inds, to index into the columns of m1. So I would like to get back a 4x1 vector, composed of m[0,1], m[1,1], m[2,0], m[3,0]

How can I do that?

Thanks!

10 Likes
2

There is a gather function for that.

m = torch.randn(4,2)
ids = torch.Tensor([1,1,0,0]).long()
print(m.gather(1, ids.view(-1,1)))
28 Likes
3

Thank you that’s exactly what I needed. Why cant we use normal slicing as in numpy?

4

I think we haven’t yet had the time to implement advanced indexing as the ones present in numpy. But @apaszke can better comment on that, might not be high in the priority list.

1 Like
5

How would you like to index it? We support indexing with LongTensors, and it works in the same way as numpy, but it’s not equivalent to gather, but rather to index_select.

6

@apaszke I think this is the first example in the Advanced Indexing from numpy

7

Not really, this would require passing in a 2D list.

9

https://docs.scipy.org/doc/numpy-1.13.0/user/basics.indexing.html#indexing-multi-dimensional-arrays

10

A follow up question: is the columns selection in Pytorch differentiable? Say the column index is obtained though a network, and the loss is computed based on the selected columns from the matrix. Will the back propagation go through?

7 Likes
11

Yes indexing is fully differentiable in PyTorch!

You can see this if you print out a tensor that you just performed indexing (selecting) on: tensor(..., grad_fn=<SelectBackward>)

Obviously this only back-propagates into the indexed elements with the other elements receiving zero gradient. This is necessary with some techniques such as Policy Gradients however.

Also thanks again to @fmassa for the clean answer!

2 Likes
12

How would you do the same when assigning to an existing tensor?

Example:

tensor.gather(-1, indices) = new_values_tensor

Basically, need an index assign version of gather.

EDIT:
It’s scatter.
Still testing to figure out it’s usability in autograd graphs. I’m basically bootlegging a sparse representation with an occasional dense step in the middle.

4 Likes
13

Oh man, you’re a life saver. Been searching this for over an hour. :slight_smile:

14

Hi, Can the selected indexes of a tensor itself, but not the selected elements of indexes differentiable in PyTorch? So that I can select different indexes of a tensor by back propagation.