Selecting from a 2d tensor with rows of column indexes

sam_l: <p>I have a 2d input tensor and a 2d index tensor with the same number of rows where each value represents the column index I want to select.</p>
<p>I want as output a 2d tensor with the same number of rows as the input tensor and cols as the index tensor.</p>
<p>I can do this by creating row and column index tensors, selecting from the input tensor, and then reshaping, but is there a more efficient way to do this?</p>
<p>Example:</p>
<pre><code class="lang-auto">In [2]: a = torch.arange(12).view(3,4);a
Out[2]:
tensor([[ 0,  1,  2,  3],
        [ 4,  5,  6,  7],
        [ 8,  9, 10, 11]])

In [3]: idx = torch.tensor([[1,3],[0,1],[2,3]]);idx #all elements are &lt;4, the number of columns of a
   ...:
Out[3]:
tensor([[1, 3], #want element 1,3 of row 0
        [0, 1], # want element 0, 1 of row 1
        [2, 3]]) #want element 2, 3 of row 2

In [4]: row_idx = torch.arange(3).repeat_interleave(2) 
   ...: col_idx = idx.view(-1)
   ...: (row_idx, col_idx)
Out[4]: (tensor([0, 0, 1, 1, 2, 2]), tensor([1, 3, 0, 1, 2, 3]))

In [5]: a[row_idx, col_idx].view(3, -1)
Out[5]:
tensor([[ 1,  3],
        [ 4,  5],
        [10, 11]])
</code></pre>

ptrblck: <p>Direct indexing with <code>torch.arange</code> should work:</p>
<pre><code class="lang-python">a[torch.arange(a.size(0)).unsqueeze(1), idx]
# tensor([[ 1,  3],
#         [ 4,  5],
#         [10, 11]])
</code></pre>

sam_l: <p>Ah thanks!</p>
<p>To see if I understand correctly, what’s going on is:</p>
<ol>
<li>
<code>torch.arange(a.size(0)).unsqueeze(1)</code> returns a 3x1 tensor of <code>[[0,1,2]]</code>
</li>
<li>that gets broadcast to the shape of idx, (3x2) by ‘stretching’ (numpy’s term) in the 2nd dimension. call this <code>r_idx</code>
</li>
<li>The i,jth element of the output is the element of <code>a</code> in the <code>r_idx[i,j]</code>th row and <code>idx[i,j]</code>th col. That is, <code>out[i,j] = a[r_idx[i], idx[j]]</code>
</li>
</ol>
<p>Is my understanding correct?</p>
<p>Where can I find this in the pytorch docs? I have been combing over them, particularly <a href="https://pytorch.org/docs/stable/torch.html#indexing-slicing-joining" rel="noopener nofollow ugc">this section on indexing and slicing</a> but haven’t found anything on how <code>Tensor.__getitem__</code> works.</p>

ptrblck: <p>I think your explanation is correct.<br>
I’m unsure if this behavior is explicitly documented and think it’s cloned from numpy’s advanced indexing logic. <a href="https://docs.scipy.org/doc/numpy-1.10.4/reference/arrays.indexing.html#advanced-indexing">These docs</a> explain the numpy indexing in more detail.</p>

sam_l: <p>Great, thank you!</p>
<p>I see in the <a href="https://pytorch.org/docs/stable/tensor_view.html" rel="noopener nofollow ugc">Tensor view docs</a>, there’s a note saying</p>
<blockquote>
<p>When accessing the contents of a tensor via indexing, PyTorch follows Numpy behaviors that basic indexing returns views, while advanced indexing returns a copy. Assignment via either basic or advanced indexing is in-place. See more examples in <a href="https://docs.scipy.org/doc/numpy/reference/arrays.indexing.html" rel="noopener nofollow ugc">Numpy indexing documentation</a>.</p>
</blockquote>
<p>I will open an issue on github to request that this is made explicit in the indexing/slicing docs.</p>