How to randomly set a fixed number of elements in each row of a tensor

username: sisaman
cooked: <p>I was wondering if there is any more efficient alternative for the below code, without using the “for” loop in the 4th line?</p>
<pre><code class="lang-auto">import torch
n, d = 37700, 7842
k = 4
sample = torch.cat([torch.randperm(d)[:k] for _ in range(n)]).view(n, k)
mask = torch.zeros(n, d, dtype=torch.bool)
mask.scatter_(dim=1, index=sample, value=True)
</code></pre>
<p>Basically, what I am trying to do is to create a <code>n x d</code> mask tensor, such that in each row exactly <code>k</code> random elements are True.</p>

username: sisaman
cooked: <p>Here’s a potential solution:<br>
</p><aside class="onebox stackexchange">
  <header class="source">
      <a href="https://stackoverflow.com/questions/64162672/how-to-randomly-set-a-fixed-number-of-elements-in-each-row-of-a-tensor-in-pytorc/64170973#64170973" target="_blank" rel="noopener nofollow ugc">stackoverflow.com</a>
  </header>
  <article class="onebox-body">
      <a href="https://stackoverflow.com/users/4757715/gil-pinsky" target="_blank" rel="noopener nofollow ugc">
    <img alt="Gil Pinsky" src="https://graph.facebook.com/10205551145809669/picture?type=large" class="thumbnail onebox-avatar" width="60" height="60">
  </a>
<h4>
  <a href="https://stackoverflow.com/questions/64162672/how-to-randomly-set-a-fixed-number-of-elements-in-each-row-of-a-tensor-in-pytorc/64170973#64170973" target="_blank" rel="noopener nofollow ugc">How to randomly set a fixed number of elements in each row of a tensor in PyTorch</a>
</h4>

<div class="tags">
  <strong>pytorch</strong>
</div>

<div class="date">
  
  answered by
  <a href="https://stackoverflow.com/users/4757715/gil-pinsky" target="_blank" rel="noopener nofollow ugc">
    Gil Pinsky
  </a>
  on <a href="https://stackoverflow.com/questions/64162672/how-to-randomly-set-a-fixed-number-of-elements-in-each-row-of-a-tensor-in-pytorc/64170973#64170973" target="_blank" rel="noopener nofollow ugc">11:45AM - 02 Oct 20 UTC</a>
</div>

  </article>
  <div class="onebox-metadata">
    
    
  </div>
  <div style="clear: both"></div>
</aside>

<p>More efficient approaches are welcome.</p>

username: tom
cooked: <p>The obvious thing would be to use <code>torch.multinomial(torch.ones(n, d), 4)</code>, which would take out ~1/3 of the time for me but is somewhat slow.<br>
I can half the time (relative to your code, on my machine, on CPU etc.) by using rand + topk.</p>
<pre><code class="lang-auto">    sample = torch.rand(n, d).topk(4, dim=1).indices
    mask = torch.zeros(n, d, dtype=torch.bool)
    mask.scatter_(dim=1, index=sample, value=True)
</code></pre>
<p>I imagine for k much smaller than d you could be blazingly fast writing your own GPU kernel that just loops until it has found four new things because the 37700 is just one giant opportunity for parallelization on the GPU.</p>
<p>Best regards</p>
<p>Thomas</p>

username: tom
cooked: <p>Had I known you went to stackoverflow to get the same answer, I would have not went through the trouble of benchmarking things.<br>
I spent 15 minutes on your problem trying to help you only to find that you would have been OK without.</p>

username: sisaman
cooked: <p>Thank you for your answer. Actually, your answer works better for me than the one in StackOverflow both in terms of running time and similarity to my own code, with only one line changed.</p>
<p>But I think I have the right to seek my answer from different sources, right? It’s not against the community guidelines, I suppose. Also, someone else could have posted an answer here (instead of StackOverflow) a few minutes before you, so I don’t think it’s my fault. Anyway, I really appreciate your time, and your answer not only helped me but could also help other people in the future.</p>
<p>Best regards,<br>
Sina</p>

username: tom
cooked: <p>Oh, you have every right to ask where ever you want how often you want and so.<br>
I would think that is is common courtesy and sound use of a resource to not ask the same question in multiple places causing many people to invest time helping you when they could have helped the next person instead with much of the same outcome.</p>
<p>I feel I should have done something else with my time in retrospect, but yeah, maybe then I should not be answering questions.</p>
<p>Best regards and good luck with your project</p>
<p>Thomas</p>

username: sisaman
cooked: <p>Again, I am so grateful for your time and effort and sorry for any inconvenience I might have caused.</p>
<p>Best,<br>
Sina</p>