How to get a matrix top-n in its second axis?

username: acgtyrant
cooked: <p>For example, I use numpy:</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/2X/8/8b9ff6b93a40d959242b89d04e8befd7a46ec75d.png" data-download-href="https://discuss.pytorch.org/uploads/default/8b9ff6b93a40d959242b89d04e8befd7a46ec75d" title="2017-12-05-170222_544x733_scrot"><img src="https://discuss.pytorch.org/uploads/default/optimized/2X/8/8b9ff6b93a40d959242b89d04e8befd7a46ec75d_2_371x500.png" alt="2017-12-05-170222_544x733_scrot" data-base62-sha1="jVb5wN6lznDyaZaVeBA1JKuh3ZP" width="371" height="500" srcset="https://discuss.pytorch.org/uploads/default/optimized/2X/8/8b9ff6b93a40d959242b89d04e8befd7a46ec75d_2_371x500.png, https://discuss.pytorch.org/uploads/default/original/2X/8/8b9ff6b93a40d959242b89d04e8befd7a46ec75d.png 1.5x, https://discuss.pytorch.org/uploads/default/original/2X/8/8b9ff6b93a40d959242b89d04e8befd7a46ec75d.png 2x" data-dominant-color="252B2D"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">2017-12-05-170222_544x733_scrot</span><span class="informations">544×733 41.1 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>Then I would like to do it in PyTorch, but it failed, I do not know how to solve it.</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/2X/a/ac05e46d799dd65b0bbb27496e82b0e1e8ce3f54.png" data-download-href="https://discuss.pytorch.org/uploads/default/ac05e46d799dd65b0bbb27496e82b0e1e8ce3f54" title="2017-12-05-170351_1042x765_scrot"><img src="https://discuss.pytorch.org/uploads/default/optimized/2X/a/ac05e46d799dd65b0bbb27496e82b0e1e8ce3f54_2_681x500.png" alt="2017-12-05-170351_1042x765_scrot" data-base62-sha1="oxMKrxZjvrpJDb07fZSS2cETY2w" width="681" height="500" srcset="https://discuss.pytorch.org/uploads/default/optimized/2X/a/ac05e46d799dd65b0bbb27496e82b0e1e8ce3f54_2_681x500.png, https://discuss.pytorch.org/uploads/default/optimized/2X/a/ac05e46d799dd65b0bbb27496e82b0e1e8ce3f54_2_1021x750.png 1.5x, https://discuss.pytorch.org/uploads/default/original/2X/a/ac05e46d799dd65b0bbb27496e82b0e1e8ce3f54.png 2x" data-dominant-color="232526"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">2017-12-05-170351_1042x765_scrot</span><span class="informations">1042×765 47 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>

username: tom
cooked: <p>You can use <a href="http://pytorch.org/docs/master/torch.html#torch.topk" rel="nofollow noopener"><code>torch.topk</code></a>.<br>
It returns a tuple with the maximum entries (like sort) and their indices (like argsort).</p>
<p>Best regards</p>
<p>Thomas</p>

username: mu2
cooked: <p>Hey <a class="mention" href="/u/tom">@tom</a>, torch.topk is significantly slower than torch.sort</p>
<p>I’m running it on a tensor of dimension 300k and taking k=200 out.</p>
<p>sample code -</p>
<pre><code class="lang-auto">import torch
x = torch.rand(300000, device='cuda')

%timeit torch.topk(x, 200) # 1000 loops, best of 5: 2.96 ms per loop
%timeit torch.sort(x, descending=True)[0][:200] # 1000 loops, best of 5: 812 µs per loop
</code></pre>
<p>Can you please help me here?</p>
<p>Regards<br>
Nikhil</p>

username: tom
cooked: <p>So what happens is that <code>torch.sort</code> (for large tensors if memory serves me well) goes to the Thrust library’s sort which is likely better optimized than PyTorch’s own kernels.<br>
Note that there are memory implications of using sort: The tensor you get from it is a view of the sorted 300k element tensor. Depending on what you do (in particular when involving autograd), this can be a disadvantage.</p>
<p>Best regards</p>
<p>Thomas</p>