How to select particular elements of a 3D tensor based on indices along dim 1 in PyTorch?

Soumyajain29: <p>I have a 3D tensor A of shape (M, N, K) where M is the batch size, N is sequence length, and K is embedding dimension. I also have a list of indices along dimension 1 (i.e., indices for sequence length dimension). I am unable to figure out how to get elements of A based on this indexing.</p>
<blockquote>
<p>For ex- M = 2 , N = 4 , K = 2<br>
A = [[[1,2],[3,4],[5,6],[7,8]] , [[9,10],[11,12],[13,14],[15,16]]]<br>
Index = [0,2]</p>
<p>output = [[1,2] ,[13,14]]</p>
</blockquote>
<p>Note that for each element of the batch, I have a different index for sequence. Basically, I want to extract embeddings of &lt;\s&gt; token. And for each sample in the batch, this &lt;\s&gt; can lie at a different index.</p>

ptrblck: <p>This should work:</p>
<pre><code class="lang-python">A = torch.tensor([[[1,2],[3,4],[5,6],[7,8]] , [[9,10],[11,12],[13,14],[15,16]]])
index = torch.tensor([0,2])

out = A[torch.arange(A.size(0)), index]
print(out)
&gt; tensor([[ 1,  2],
          [13, 14]])
</code></pre>

Soumyajain29: <p><a class="mention" href="/u/ptrblck">@ptrblck</a>, It worked. Thanks a lot.<br>
It would be very helpful if you could explain the working of this. I am not getting how torch.arange is working here.</p>

ptrblck: <p><code>torch.arange(A.size(0))</code> will return <code>tensor([0, 1])</code> in this case, so the indexing will work as:</p>
<pre><code class="lang-python">out[0] = A[0, index[0]]
out[1] = A[1, index[1]]
</code></pre>
<p>and will thus yield the desired results. On the other hand using <code>A[:, index]</code> will not work (you might be more familiar with this indexing operation), as it would use <code>index</code> on “all” values on dim0, so it’ll return:</p>
<pre><code class="lang-python">A[:, index]
&gt; tensor([[[ 1,  2],
           [ 5,  6]],

          [[ 9, 10],
           [13, 14]]])
</code></pre>
<p>As you can see, <code>index</code> (containing <code>[0, 2]</code>) was now applied to all values in <code>dim0</code> as:</p>
<pre><code class="lang-python">A[0, index]
A[1, index]
</code></pre>
<p>thus returning more values.</p>
<p>EDIT: <a href="https://numpy.org/doc/stable/reference/arrays.indexing.html#integer-array-indexing">this numpy doc</a> might be also helpful in case my explanation is confusing.</p>

Soumyajain29: <p>Understood. Thanks a lot.</p>

Soumyajain29: <p><a class="mention" href="/u/ptrblck">@ptrblck</a> , How can I do it if I have indices like-</p>
<blockquote>
<p>index = torch.tensor([[0,2],[1,3] , [0,1]]) ?</p>
</blockquote>
<p>Basically, for each batch element, I have a different list of indices.</p>

ptrblck: <p>Could you post the desired output using this new <code>index</code> and the previously defined <code>A</code> tensor?</p>

Soumyajain29: <blockquote>
<p>A = torch.tensor([[[1,2],[3,4],[5,6],[7,8]] , [[9,10],[11,12],[13,14],[15,16]] , [[21,22],[23,24] ,[25,26],[27,28]]])<br>
index = torch.tensor([[0,2],[1,3] , [0,1]])<br>
output = [[[1,2] , [5,6]] ,[[11,12] , [15,16]] , [[21,22] , [23,24]]]</p>
</blockquote>

ptrblck: <p>This should work:</p>
<pre><code class="lang-python">A = torch.tensor([[[1,2],[3,4],[5,6],[7,8]] , [[9,10],[11,12],[13,14],[15,16]] , [[21,22],[23,24] ,[25,26],[27,28]]])
index = torch.tensor([[0,2], [1,3], [0,1]])
A[torch.arange(A.size(0)).unsqueeze(1), index]
&gt; tensor([[[ 1,  2],
           [ 5,  6]],

          [[11, 12],
           [15, 16]],

          [[21, 22],
           [23, 24]]])
</code></pre>