Can this for loop be avoided using Pytorch operations?

mzimmer: <p>Hey! I implemented the following algorithm using python for loops, which are, afaik, not very efficient:</p>
<pre><code class="lang-auto">sorted, indices = torch.sort(torch.abs(torch.reshape(x, (-1,))), descending=True)
running_sum = -r # r is some positive number
for idx, element in enumerate(sorted):
    running_sum += element
    if element &lt;= (running_sum) / (idx + 1):
         running_sum -= element
         idx -= 1
         break
</code></pre>
<p>Is there any way to do this faster using Pytorch operations? The crucial part is finding the first entry of a large sorted tensor such that said entry is less-or-equal than the mean of all previous entries.</p>
<p>Looking forward for your ideas! Thanks a lot!</p>

KFrank: <p>Hi Max!</p>
<aside class="quote no-group" data-username="mzimmer" data-post="1" data-topic="95452" 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/mzimmer/48/31510_2.png" class="avatar"> mzimmer:</div>
<blockquote>
<p>Is there any way to do this faster using Pytorch operations? The crucial part is finding the first entry of a large sorted tensor such that said entry is less-or-equal than the mean of all previous entries.</p>
</blockquote>
</aside>
<p>Yes, you can do this with pytorch tensor operations.  Use <code>cumsum()</code><br>
divided by something like <code>arange()</code> to get the running mean.  Then<br>
test that against your sorted tensor to find when your condition first<br>
holds.</p>
<p>Here is a pytorch version 0.3.0 script:</p>
<pre data-code-wrap="python"><code class="lang-python">import torch
torch.__version__

torch.manual_seed (2020)

x = torch.randn ((3, 2))
r = 2.0

# pytorch tensor operations
sorted, indices = torch.sort(torch.abs( x.view (-1,) ), descending=True)               # sort data
running_mean = (torch.cumsum (sorted, 0) - r) / torch.arange (1, sorted.numel() + 1)   # get cumulative mean
is_le = sorted &lt;= running_mean          # test condition
idx = is_le.numel() - is_le.sum() - 1   # get index
print ('idx =', idx)
print ('sorted[idx + 1] =', sorted[idx + 1])
print ('running_mean[idx + 1] =', running_mean[idx + 1])


# loop
# sorted, indices = torch.sort(torch.abs(torch.reshape(x, (-1,))), descending=True)
sorted, indices = torch.sort(torch.abs( x.view (-1,) ), descending=True)
running_sum = -r # r is some positive number
for idx, element in enumerate(sorted):
    running_sum += element
    if element &lt;= (running_sum) / (idx + 1):
         running_sum -= element
         idx -= 1
         break

print ('x =', x)
print ('sorted =', sorted)
print ('running_sum =', running_sum)
print ('idx =', idx)
print ('sorted[idx + 1] =', sorted[idx + 1])
print ('(running_sum + sorted[idx + 1]) / (idx + 2) =', (running_sum + sorted[idx + 1]) / (idx + 2))
</code></pre>
<p>And here is its output</p>
<pre><code class="lang-plaintext">&gt;&gt;&gt; import torch
&gt;&gt;&gt; torch.__version__
'0.3.0b0+591e73e'
&gt;&gt;&gt;
&gt;&gt;&gt; torch.manual_seed (2020)
&lt;torch._C.Generator object at 0x0000020EC7916630&gt;
&gt;&gt;&gt;
&gt;&gt;&gt; x = torch.randn ((3, 2))
&gt;&gt;&gt; r = 2.0
&gt;&gt;&gt;
&gt;&gt;&gt; # pytorch tensor operations
... sorted, indices = torch.sort(torch.abs( x.view (-1,) ), descending=True)               # sort data
&gt;&gt;&gt; running_mean = (torch.cumsum (sorted, 0) - r) / torch.arange (1, sorted.numel() + 1)   # get cumulative mean
&gt;&gt;&gt; is_le = sorted &lt;= running_mean          # test condition
&gt;&gt;&gt; idx = is_le.numel() - is_le.sum() - 1   # get index
&gt;&gt;&gt; print ('idx =', idx)
idx = 3
&gt;&gt;&gt; print ('sorted[idx + 1] =', sorted[idx + 1])
sorted[idx + 1] = 0.40786537528038025
&gt;&gt;&gt; print ('running_mean[idx + 1] =', running_mean[idx + 1])
running_mean[idx + 1] = 0.6055201292037964
&gt;&gt;&gt;
&gt;&gt;&gt;
&gt;&gt;&gt; # loop
... # sorted, indices = torch.sort(torch.abs(torch.reshape(x, (-1,))), descending=True)
... sorted, indices = torch.sort(torch.abs( x.view (-1,) ), descending=True)
&gt;&gt;&gt; running_sum = -r # r is some positive number
&gt;&gt;&gt; for idx, element in enumerate(sorted):
...     running_sum += element
...     if element &lt;= (running_sum) / (idx + 1):
...          running_sum -= element
...          idx -= 1
...          break
...
&gt;&gt;&gt; print ('x =', x)
x =
 1.2372 -0.9604
 1.5415 -0.4079
 0.8806  0.0529
[torch.FloatTensor of size 3x2]

&gt;&gt;&gt; print ('sorted =', sorted)
sorted =
 1.5415
 1.2372
 0.9604
 0.8806
 0.4079
 0.0529
[torch.FloatTensor of size 6]

&gt;&gt;&gt; print ('running_sum =', running_sum)
running_sum = 2.619735360145569
&gt;&gt;&gt; print ('idx =', idx)
idx = 3
&gt;&gt;&gt; print ('sorted[idx + 1] =', sorted[idx + 1])
sorted[idx + 1] = 0.40786537528038025
&gt;&gt;&gt; print ('(running_sum + sorted[idx + 1]) / (idx + 2) =', (running_sum + sorted[idx + 1]) / (idx + 2))
(running_sum + sorted[idx + 1]) / (idx + 2) = 0.6055201470851899
</code></pre>
<p>Best.</p>
<p>K. Frank</p>

mzimmer: <p>Thanks for the fast and intuitive answer! One question remains: I was kind of surprised by your way to find <code>idx</code>. Is that really the most efficient way of finding the index of the largest False element of <code>is_le</code>? <code>numel()</code> is probably constant in the size of <code>is_le</code>, but <code>sum()</code> should be linear?</p>

KFrank: <p>Hi Max!</p>
<aside class="quote no-group" data-username="mzimmer" data-post="3" data-topic="95452" 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/mzimmer/48/31510_2.png" class="avatar"> mzimmer:</div>
<blockquote>
<p>Is that really the most efficient way of finding the index of the largest False element of <code>is_le</code>? <code>numel()</code> is probably constant in the size of <code>is_le</code>, but <code>sum()</code> should be linear?</p>
</blockquote>
</aside>
<p>First, yes, <code>sum()</code> is linear.  (But, as an aside, finding the first <code>True</code><br>
index is also linear, but see below).</p>
<p>The key point is that <code>sum()</code> is implemented as a pytorch tensor<br>
operation, so you get that potentially important benefit.</p>
<p>Note, many libraries offer a <code>find_first</code> type of function for searching<br>
strings or arrays.  I’m not aware that pytorch has such (at least not<br>
in my old version).  One could consider using pytorch’s <code>argmax()</code><br>
function, but in this case several elements of <code>is_le</code> take on the<br>
maximum value, and I don’t think that pytorch guarantees that<br>
<code>argmax()</code> will return the index of the <em>first</em> maximum element.</p>
<p>Now (other than the initial sort) this entire algorithm is linear.  But<br>
this algorithm looks at the <em>entire</em> <code>sorted</code> vector – whether or not<br>
it needs to.  You only have to calculate your running mean until your<br>
condition is satisfied.  So if your condition is equally likely to be<br>
satisfied anywhere in the vector, then, on average, you would only<br>
have to look at half of <code>sorted</code> and compute only half of the running<br>
mean.  Cheaper, but still linear.</p>
<p>If your condition were strongly-enough biased to occur towards the<br>
beginning of your vector, the algorithm could even become order 1<br>
(“O(1)”), for the average case.  But, given that pytorch doesn’t have<br>
anything built into it that does this for you, you’re almost certainly<br>
better of using the linear – but optimized – pytorch tensor operations,<br>
than writing your own non-tensor loop (even in the hypothetical<br>
average-case O(1) example).</p>
<p>The same comments apply to finding the first <code>True</code> index in <code>is_le</code>.<br>
This, in general, is linear, but could be cheaper in operations if you<br>
run a loop that stops “early” when it finds that index.  But the tensor<br>
operation <code>sum()</code> is likely to be cheaper in time (but not operations),<br>
even in the hypothetical average-case O(1) situation.</p>
<p>Best.</p>
<p>K. Frank</p>