Indexing tensors of varying dimensions with topk

username: kleingeo
cooked: <p>I am currently using <code>torch.topk</code> to determine the indices of the of a 2D tensor <code>scores</code> which is of size [Batch, N]. I can get the topk values (6000) from <code>scores</code> with <code>torch.gather</code> (or simply from the <code>torch.topk</code> directly).</p>
<pre><code class="lang-python">idx = torch.topk(scores, 6000, dim=1, sorted=True)
scores = torch.gather(scores, dim=1, index=idx) # Output is of size [B, 6000]
</code></pre>
<p>My issues comes when I am trying to use the same indices on a 3D tensor <code>bbox</code> which is of size [Batch, N, 4]. How could I use the same indices to get something like below without having to resort to for loops</p>
<pre><code class="lang-python">bbox = torch.gather(bbox, dim=1, index=idx) # Outside of size [B, 6000, 4]
</code></pre>

username: ptrblck
cooked: <p>Your code should also work for this shape:</p>
<pre><code class="lang-python">scores = torch.randn(10, 6000)
idx = torch.topk(scores, 6000, dim=1, sorted=True)
out = torch.gather(scores, dim=1, index=idx.indices)
print((torch.sort(scores, dim=1, descending=True).values == out).all())
&gt; tensor(true)

scores = torch.randn(10, 6000, 4)
idx = torch.topk(scores, 6000, dim=1, sorted=True)
out = torch.gather(scores, dim=1, index=idx.indices)
print((torch.sort(scores, dim=1, descending=True).values == out).all())
&gt; tensor(true)
</code></pre>
<p>Do you see an error?</p>

username: kleingeo
cooked: <aside class="quote no-group" data-username="ptrblck" data-post="2" data-topic="72338">
<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/ptrblck/48/1823_2.png" class="avatar"> ptrblck:</div>
<blockquote>
<p>for this</p>
</blockquote>
</aside>
<p>I understand that, but wouldn’t that require a reshape or modification of scores to work obtain indices for both the 2D and 3D tensors, rather than a solution that works for both?</p>
<p>One not-entirely-elegant solution I did was</p>
<pre data-code-wrap="python"><code class="lang-python">idx_stacked= torch.stack([idx] * bbox.share[-1], dim=idx.dim()
bbox = torch.gather(bbox, dim=1, index=idx_stacked)
</code></pre>
<p>Unsure if there a cleaner way to do it.</p>

username: Chuong98
cooked: <p>Hello, can you please help with this example</p>
<pre><code class="lang-auto">x = torch.rand(10,5)
y = torch.rand(10,4,5)
top_x,top_inds = x.topk(3,dim=-1)
top_y= y.gather(dim=-1,index=top_inds)
</code></pre>
<p>and it does not work.</p>
<blockquote>
<p>RuntimeError: Index tensor must have the same number of dimensions as input tensor</p>
</blockquote>
<p>Currently, I have to work around this by a for-loop:</p>
<pre><code class="lang-auto">s_inds = torch.arange(len(x)).long()
top_y = [y[s_inds,:,top_inds[:,i]] for i in range(3)]
top_y = torch.stack(top_y,dim=2)
</code></pre>

username: Chuong98
cooked: <p>My slightly better solution is:</p>
<pre><code class="lang-auto">n_samples, topk = len(x), 3
s_inds = torch.arange(len(x)).long()
top_y = y[s_inds.repeat_interleave(topk),:,top_inds.view(-1)] #(n_samples*topk,C)
top_y = top_y.view(n_samples,topk,-1).permute(0,2,1) #(n_samples,C,topk)
</code></pre>