Double indexing/slicing using tensors as indices

username: Minh_Nguyen
cooked: <p>Hello, I’m having a hard time understanding double indexing using tensors as indices; as an example:</p>
<pre><code class="lang-auto">import torch

points = torch.randint(0, 4, (4, 5, 2))
batch_indices = torch.randint(0, 4, (4, 4, 3))
idx = torch.randint(0, 4, (4, 4, 3))

new_points = points[batch_indices, idx, :]

print('points shape: {}'.format(points.shape))
print('batch_indices shape: {}'.format(batch_indices.shape))
print('idx shape: {}'.format(idx.shape))
print('new_points shape: {}'.format(new_points.shape))
</code></pre>
<p>The shapes printing results are as follows:</p>
<pre><code class="lang-auto">points shape: torch.Size([4, 5, 2])
batch_indices shape: torch.Size([4, 4, 3])
idx shape: torch.Size([4, 4, 3])
new_points shape: torch.Size([4, 4, 3, 2])
</code></pre>
<p>Can someone help to provide an equivalent set of simpler operations?</p>

username: KFrank
cooked: <p>Hi Minh!</p>
<aside class="quote no-group" data-username="Minh_Nguyen" data-post="1" data-topic="160540" data-full="true">
<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/minh_nguyen/48/32466_2.png" class="avatar"> Minh_Nguyen:</div>
<blockquote>
<p>I’m having a hard time understanding double indexing using tensors as indices;</p>
</blockquote>
</aside>
<p>Let’s start by looking at an example of so-called <em>advanced indexing</em><br>
where we use two index tensors to index into a two-dimensional<br>
tensor.  This is similar to your example, except we simplify it by leaving<br>
out the part where you use slicing (the colon, “<code>:</code>”) on the third dimension.</p>
<p>The two index tensors have the same shape, but their shape is not<br>
related to the shape of the tensor into which you are indexing.  Instead,<br>
the shape of the index tensors becomes the shape of the resulting tensor.</p>
<p>This is the first part of the example given below.</p>
<p>Then when you put the colon in for your third dimension, you are saying<br>
that you want all of the third dimension, and advanced indexing simply<br>
applies your two index tensors to each (two-dimensional) slice along the<br>
third dimension.</p>
<p>This is done explicitly in the second part of the example.</p>
<p>Finally, you can (not that you should) build three index tensors so that<br>
you can use “plain” advanced indexing to index directly into your<br>
three-dimensional tensor without slicing.</p>
<p>This is illustrated in the third part of the example.</p>
<p>Consider:</p>
<pre><code class="lang-plaintext">&gt;&gt;&gt; import torch
&gt;&gt;&gt; print (torch.__version__)
1.12.0
&gt;&gt;&gt;
&gt;&gt;&gt; _ = torch.manual_seed (2022)
&gt;&gt;&gt;
&gt;&gt;&gt; points = torch.randint(0, 4, (4, 5, 2))
&gt;&gt;&gt; batch_indices = torch.randint(0, 4, (4, 4, 3))
&gt;&gt;&gt; idx = torch.randint(0, 4, (4, 4, 3))
&gt;&gt;&gt;
&gt;&gt;&gt; points2d = points[..., 0]   # make a 2d tensor
&gt;&gt;&gt; points2d.shape
torch.Size([4, 5])
&gt;&gt;&gt;
&gt;&gt;&gt; # use two index tensors to index into 2d tensor
&gt;&gt;&gt; new_points2d = points2d[batch_indices, idx]   # the index tensors are 3d so new_points2d is 3d
&gt;&gt;&gt; new_points2d.shape
torch.Size([4, 4, 3])
&gt;&gt;&gt;
&gt;&gt;&gt; # use two index tensors plus slicing to index into 3d tensor
&gt;&gt;&gt; new_points = points[batch_indices, idx, :]
&gt;&gt;&gt; new_points.shape
torch.Size([4, 4, 3, 2])
&gt;&gt;&gt;
&gt;&gt;&gt; # build new_points by indexing into slices of points
&gt;&gt;&gt; s0 = points[..., 0][batch_indices, idx]
&gt;&gt;&gt; s1 = points[..., 1][batch_indices, idx]
&gt;&gt;&gt; new_pointsB = torch.stack ((s0, s1), -1)
&gt;&gt;&gt; torch.equal (new_points, new_pointsB)
True
&gt;&gt;&gt;
&gt;&gt;&gt; # use three index tensors to index into 3d tensor
&gt;&gt;&gt; batch_indicesC = batch_indices.unsqueeze (-1).expand (-1, -1, -1, 2)
&gt;&gt;&gt; idxC = idx.unsqueeze (-1).expand (-1, -1, -1, 2)
&gt;&gt;&gt; slice_idxC = torch.stack ((torch.zeros (4, 4, 3).long(), torch.ones (4, 4, 3).long()), -1)
&gt;&gt;&gt; batch_indicesC.shape
torch.Size([4, 4, 3, 2])
&gt;&gt;&gt; idxC.shape
torch.Size([4, 4, 3, 2])
&gt;&gt;&gt; slice_idxC.shape
torch.Size([4, 4, 3, 2])
&gt;&gt;&gt; new_pointsC = points[batch_indicesC, idxC, slice_idxC]
&gt;&gt;&gt; torch.equal (new_points, new_pointsC)
True
</code></pre>
<p>Best.</p>
<p>K. Frank</p>

username: Minh_Nguyen
cooked: <p>Hello KFrank,</p>
<p>Thank you very much for your efforts in making such a detailed answer! I will spend time reading and understanding all of them.</p>