# Select specific columns of each row in a torch Tensor

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

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

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

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

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`.

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

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

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

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

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

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