Is there a better way to perform row-corresponding indexing?

u7122029: <p>So let’s say I have array <code>A</code> and the indexing array <code>idxs</code> as follows</p>
<pre data-code-wrap="py"><code class="lang-py">A = torch.Tensor([[8, 3, 5],
                  [7, 6, 1],
                  [2, 4, 9]])
idxs = torch.Tensor([[0, 2],
                     [1, 2],
                     [2, 0]]).int()
</code></pre>
<p>I want to find a fast method of indexing where the first row of <code>idxs</code> indexes the first row of <code>A</code>, the second row of <code>idxs</code> indexes the second row, and so on. In this way I’d have the following output</p>
<pre data-code-wrap="py"><code class="lang-py">[[8, 5],
 [6, 1],
 [9, 2]]
</code></pre>
<p>I have a solution as shown below</p>
<pre data-code-wrap="py"><code class="lang-py">A[:, idxs][torch.eye(len(A))]
</code></pre>
<p>Which works quite well, but I am worried that if <code>A</code> has tens of thousands of rows there will be memory issues due to the <code>torch.eye</code> function.</p>
<ol>
<li>Is pytorch’s storage of tensors sufficiently efficient that there should be no issues at all? Or,</li>
<li>Should I be worried about this and change my solution to something better?</li>
<li>Regardless of both questions, is there a better indexing method than what I have shown?</li>
</ol>
<p>Thanks in advance!</p>

ptrblck: <aside class="quote no-group" data-username="u7122029" data-post="1" data-topic="195797">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/u7122029/48/64593_2.png" class="avatar"> u7122029:</div>
<blockquote>
<p>Which works quite well,</p>
</blockquote>
</aside>
<p>Your code doesn’t work for me and fails with:</p>
<pre data-code-wrap="python"><code class="lang-python">A = torch.Tensor([[8, 3, 5],
                  [7, 6, 1],
                  [2, 4, 9]])
idxs = torch.Tensor([[0, 2],
                     [1, 2],
                     [2, 0]]).int()
A[:, idxs][torch.eye(len(A))]
# IndexError: tensors used as indices must be long, int, byte or bool tensors
</code></pre>
<p>Fixing this by calling <code>torch.eye(...).long()</code> gives:</p>
<pre data-code-wrap="python"><code class="lang-python">A[:, idxs][torch.eye(len(A)).long()]
# tensor([[[[7., 1.],
#           [6., 1.],
#           [1., 7.]],

#          [[8., 5.],
#           [3., 5.],
#           [5., 8.]],

#          [[8., 5.],
#           [3., 5.],
#           [5., 8.]]],


#         [[[8., 5.],
#           [3., 5.],
#           [5., 8.]],

#          [[7., 1.],
#           [6., 1.],
#           [1., 7.]],

#          [[8., 5.],
#           [3., 5.],
#           [5., 8.]]],


#         [[[8., 5.],
#           [3., 5.],
#           [5., 8.]],

#          [[8., 5.],
#           [3., 5.],
#           [5., 8.]],

#          [[7., 1.],
#           [6., 1.],
#           [1., 7.]]]])
</code></pre>
<p>so also not your expected output, I’m thus unsure how your approach worked.</p>
<p>In any case, <code>torch.gather</code> should work:</p>
<pre data-code-wrap="python"><code class="lang-python">A = torch.tensor([[8, 3, 5],
                  [7, 6, 1],
                  [2, 4, 9]])
idxs = torch.tensor([[0, 2],
                     [1, 2],
                     [2, 0]])
torch.gather(A, 1, idxs)
# tensor([[8, 5],
#         [6, 1],
#         [9, 2]])
</code></pre>