First nonzero index

dizcza: <p>I have a batch of N rows each of M values that are sorted along dim=1. For each row, I want to find the first nonzero element index from M sorted values. I’d like to do it efficiently without the for-loop.</p>
<pre><code class="lang-auto">x = torch.randn(5, 7)
x[x&lt;0] = 0
x = x.sort(dim=1)
first_nonzero = f(x)
</code></pre>

kelam_goutam: <p>Hi, I coded a solution for your problem. You can check it in the link below.<br>
<aside class="onebox whitelistedgeneric">
  <header class="source">
      <img src="https://assets-cdn.github.com/favicon.ico" class="site-icon" width="16" height="16">
      <a href="https://github.com/Goutam-Kelam/PytorchForum/tree/master/FirstNonZeroIndex" target="_blank" rel="nofollow noopener">GitHub</a>
  </header>
  <article class="onebox-body">
    <img src="https://avatars3.githubusercontent.com/u/20452424?s=400&amp;v=4" class="thumbnail onebox-avatar" width="60" height="60">

<h3><a href="https://github.com/Goutam-Kelam/PytorchForum/tree/master/FirstNonZeroIndex" target="_blank" rel="nofollow noopener">Goutam-Kelam/PytorchForum</a></h3>

<p>Contains my trials for answering the questions asked in the forum - Goutam-Kelam/PytorchForum</p>


  </article>
  <div class="onebox-metadata">
    
    
  </div>
  <div style="clear: both"></div>
</aside>
</p>
<p>Hope this works out for you.</p>

InnovArul: <p>Here is another try, especially if you want to only use torch APIs.</p>
<pre><code class="lang-auto">import torch

def f(x):
    # non zero values mask
    non_zero_mask = x != 0

    # operations on the mask to find first nonzero values in the rows
    mask_max_values, mask_max_indices = torch.max(non_zero_mask, dim=1)

    # if the max-mask is zero, there is no nonzero value in the row
    mask_max_indices[mask_max_values == 0] = -1
    return mask_max_indices

x = torch.randn(4, 5)
x[x&lt;0] = 0
x, sort_indices = x.sort(dim=1)
print('x', x)
first_nonzero = f(x)

print(first_nonzero)
</code></pre>

adrianjav: <p>It is easier if you count the number of zero elements in that dimension</p>
<pre><code class="lang-python">x = torch.randn(5, 7)
x[x&lt;0] = 0
x = x.sort(dim=1)[0] # You forgot that sort returns a pair
first_nonzero = (x == 0).sum(dim=1)
</code></pre>
<p>Even easier, you can skip the <code>x[x&lt;0] = 0</code> line and count the non-positive elements:</p>
<pre><code class="lang-python">x = torch.randn(5, 7)
x = x.sort(dim=1)[0]
first_positive = (x &lt;= 0).sum(dim=1)
</code></pre>

InnovArul: <p>In the case that there are no positive numbers in a row, the answer might be misleading. Just a minor caveat to be handled I guess.</p>

adrianjav: <p>True, in that case <code>first_nonzero[i] == x.size(1)</code>. I don’t think is a caveat tho, since you’ll have to mark those cases somehow.</p>

dizcza: <p>Thanks for replies.<br>
I’m interested in pytorch-only implementation to keep my code eco-friendly. Sorry for not mentioning this beforehand.<br>
<a class="mention" href="/u/innovarul">@InnovArul</a>, your solution implies that <code>torch.max</code> always returns the first occurrence. I don’t find the docs mention that. Though it works with current version 0.4.1, I think it’s better to avoid accounting for internal implementation of pytorch functional. For that reason, <a class="mention" href="/u/adrianjav">@adrianjav</a> answer looks more compatible as long as ByteTensor supports sum operation.</p>

InnovArul: <p>yes. I like <a class="mention" href="/u/adrianjav">@adrianjav</a>’s answer as well. It looks more fail safe to me. All the best!</p>

wouter: <p>I was looking for a solution that does not assume the input is already sorted. Since this assumption was not mentioned in the topic title, it may be useful to have a solution that works in that case as well. I think the following should work in the general case. The idea is that an element is the <em>first</em> nonzero element if it is nonzero and the cumulative sum of a nonzero indicator is 1.</p>
<pre><code class="lang-auto">import torch

def first_nonzero(x, axis=0):
    nonz = (x &gt; 0)
    return ((nonz.cumsum(axis) == 1) &amp; nonz).max(axis)

x = (torch.rand(10, 5) * 10 - 6).int().clamp(0, 10)
print (x)

# Function returns if there are any nonzero's and the index of the first nonzero (0 if no nonzero)
any_nonz, idx_first_nonz = first_nonzero(x, axis=1)
print (any_nonz, idx_first_nonz)

# If you want -1 for rows with no nonzero's
idx_first_nonz[any_nonz == 0] = -1
print (idx_first_nonz)
</code></pre>

JakobHavtorn: <p><a class="mention" href="/u/wouter">@wouter</a>’s solution above was what I was looking for.</p>
<p>I would like to point out two things this though.</p>
<p>First, it actually finds the first nonnegative element, not the first nonzero element. I.e. it finds the first element that is zero or larger.</p>
<p>Second, it has a failure case when all elements are negative along the axis. In this case it returns 0, which indicates that there is a nonnegative element at index zero, which is wrong.</p>
<p>One way to improve this is to find the places at which the sum of <code>(x&gt;0)</code> along <code>axis</code> equals zero. This indicates that all elements along the axis were negative. We can then set those to some different value, e.g. <code>-1</code> or `float(‘nan’)’. We can use the already computed cumulative sum to do this.</p>
<pre><code class="lang-auto">def first_nonnegative(tensor, axis=0, value_for_all_nonnegative=-1):
    nonnegative = (tensor &gt; 0)
    cumsum = nonnegative.cumsum(axis=axis)
    all_negative = cumsum[-1] == 0  # Any dimensions where all are negative
    nonnegative_idx = ((cumsum == 1) &amp; nonnegative).max(0).indices
    nonnegative_idx[all_negative] = value_for_all_nonnegative
    return nonnegative_idx

</code></pre>

wouter: <p>Glad that it was of help to you! You are right it finds the first non-negative value (I guess I implicitly assumed only non-negative values), but you can easily change <code>nonz = (x &gt; 0)</code> for <code>nonz = (x != 0)</code>.</p>
<p>The failure case you describe I already accounted for by returning <code>any_nonz</code> as well, which you can use to set rows without any nonzero to <code>-1</code> by <code>idx_first_nonz[any_nonz == 0] = -1</code> as done in my example.</p>

JakobHavtorn: <p>Haha! Don’t know how I missed that in your example, sorry for ranting on with that.</p>
<p>And yea the other thing is pretty easy to change <img src="https://discuss.pytorch.org/images/emoji/apple/ok_hand.png?v=9" title=":ok_hand:" class="emoji" alt=":ok_hand:"></p>