Masked argmax in Pytorch?

yg-smile: <p>May I ask how to do masked argmax in Pytorch?</p>
<p>For example I have a tensor <code>t</code> and a mask <code>m</code>:<br>
<code>t = torch.tensor([20, 10, 50, 40])</code><br>
<code>m = numpy.array([True, True, False, True])</code><br>
The masked argmax is 3 (corresponds to 40 in <code>t</code>)</p>
<p>Thank you very much.</p>

ooodragon: <p>hmm</p>
<p>can’t you just multiply them to make</p>
<p>[20, 10, 0, 40]</p>
<p>and then do argmax?</p>

yg-smile: <p>Thanks. I need to convert <code>m</code> to tensor then.</p>

KFrank: <p>Hello Dragon!</p>
<aside class="quote no-group" data-username="ooodragon" data-post="2" data-topic="105341" 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/ooodragon/48/24206_2.png" class="avatar"> ooodragon:</div>
<blockquote>
<p>can’t you just multiply them to make</p>
<p>[20, 10, 0, 40]</p>
<p>and then do argmax?</p>
</blockquote>
</aside>
<p>This won’t work if tensor <code>t</code> is negative (or, more precisely, if its<br>
largest unmasked element is negative).</p>
<p>I would do this:</p>
<pre data-code-wrap="python"><code class="lang-python">large = torch.finfo (t.dtype).max   # assumes t is a kind of float
# assume msk has zeros where elements t should be masked out
# and ones where they should be kept
(t - large * (1 - msk) - large * (1 - msk)).argmax()
</code></pre>
<p>Best.</p>
<p>K. Frank</p>

Naman-ntc: <p>If you do not want to convert m to a tensor you can use</p>
<pre><code class="lang-auto">def fn(): 
    t = torch.randn(10000) 
    m = np.random.rand(10000) &lt; 0.5 
    return (t==t[m].max()).type(torch.FloatTensor).argmax()  

timeit.timeit(fn, number = 10000)                                                                                                                                                    
   3.285314051026944

</code></pre>
<p>while if you convert m to tensor (using K. Franks code : )</p>
<pre><code class="lang-auto">In [43]: def gn(): 
    ...:     t = torch.randn(10000) 
    ...:     m = torch.from_numpy(np.random.rand(10000) &lt; 0.5) 
    ...:     large = torch.finfo(t.dtype).max 
    ...:     return (t - large * (~m) - large * (~m)).argmax() 

timeit.timeit(gn, number = 10000)                                                                                                                                                    
    2.864162279991433

</code></pre>
<p>Note - benchmarking done on cpu</p>

yg-smile: <p>Thank you.<br>
Do we need two <code>- large * (1 - msk)</code>'s or one is good enough?</p>

yg-smile: <p>Thank you for the solution. Converting <code>m</code> to tensor is faster in my case as well.</p>

KFrank: <p>Hi Smile!</p>
<aside class="quote no-group" data-username="yg-smile" data-post="6" data-topic="105341" 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/yg-smile/48/32036_2.png" class="avatar"> yg-smile:</div>
<blockquote>
<p>Do we need two <code>- large * (1 - msk)</code>'s or one is good enough?</p>
</blockquote>
</aside>
<p>The second <code>- large * (1 - msk)</code> is protection against an edge case.<br>
If one of the masked values in your tensor were equal to (or close<br>
enough to) <code>large</code>, then a single <code>- large * (1 - msk)</code> would reduce<br>
it only to zero, so if all of the unmasked values in your tensor were<br>
negative, you would incorrectly get <code>0.0</code> as the maximum value.</p>
<p>Best.</p>
<p>K. Frank</p>