Find maximum length of consecutive zeros in each row

Paulo_Nascimento: <p>My goal is to find the maximum length of consecutive zeros in each row. So for instance, if I have a tensor like</p>
<pre><code class="lang-auto">input = torch.tensor([
    [0, 1, 0, 0, 0, 1],
    [0, 0, 1, 0, 1, 0],
    [1, 0, 0, 0, 0, 0]
])
</code></pre>
<p>I expect to get the result</p>
<pre><code class="lang-auto">tensor([3, 2, 5])
</code></pre>
<p>I’ve done this using numpy (listed below assuming “input” is a numpy binary matrix) and it tends to be very efficient, but I cannot seem to find a way using tensors that is at least as efficient. I’ve tried to follow a similar logic using torch operations but the performance is always worse.<br>
Thanks!</p>
<pre><code class="lang-auto"># Pad the matrix with ones
padded_matrix = np.pad(input, ((0, 0), (1, 1)), constant_values=1)

# Compute differences
diffs = np.diff(padded_matrix, axis=1)

# Identify start and end of zero runs
start_indices = np.where(diffs == -1)
end_indices = np.where(diffs == 1)

# Compute lengths of zero runs
run_lengths = end_indices[1] - start_indices[1]

# Create a result array initialized with zeros
max_zeros = np.zeros(binaryGrid.shape[0], dtype=int)

# Use np.maximum.at to find the maximum run length for each row
np.maximum.at(max_zeros, start_indices[0], run_lengths)
</code></pre>

Eduardo_Lawson: <p>When you talk about efficiency what do you mean? Time? Accuracy?<br>
I tried this with torch and it worked:</p>
<pre><code class="lang-auto">import torch

def max_consecutive_zeros(input_tensor):
    # Pad the tensor with ones on both sides
    padded = torch.nn.functional.pad(input_tensor, (1, 1), value=1)
    
    # Compute differences along the columns
    diffs = torch.diff(padded, dim=1)
    
    # Find where zero sequences start (-1) and end (1)
    start_mask = (diffs == -1)
    end_mask = (diffs == 1)
    
    # Get row and column indices of starts and ends
    start_rows, start_cols = torch.where(start_mask)
    end_rows, end_cols = torch.where(end_mask)
    
    # Compute run lengths
    run_lengths = end_cols - start_cols
    
    # Initialize max_zeros and update manually
    max_zeros = torch.zeros(input_tensor.shape[0], dtype=torch.long)
    for row, length in zip(start_rows, run_lengths):
        max_zeros[row] = max(max_zeros[row], length)
    
    return max_zeros

input_tensor = torch.tensor([
    [0, 1, 0, 0, 0, 1],
    [0, 0, 1, 0, 1, 0],
    [1, 0, 0, 0, 0, 0]
])

result = max_consecutive_zeros(input_tensor)
print(result)  # Should output: tensor([3, 2, 5])
</code></pre>

KFrank: <p>Hi Paulo!</p>
<aside class="quote no-group" data-username="Paulo_Nascimento" data-post="1" data-topic="213871" 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/paulo_nascimento/48/72888_2.png" class="avatar"> Paulo_Nascimento:</div>
<blockquote>
<p>My goal is to find the maximum length of consecutive zeros in each row. So for instance, if I have a tensor like</p>
<pre><code class="lang-auto">input = torch.tensor([
    [0, 1, 0, 0, 0, 1],
    [0, 0, 1, 0, 1, 0],
    [1, 0, 0, 0, 0, 0]
])
</code></pre>
<p>I expect to get the result</p>
<pre><code class="lang-auto">tensor([3, 2, 5])
</code></pre>
</blockquote>
</aside>
<p>If you use <code>cumsum()</code> to “count” the lengths of the sums, you can avoid using<br>
<code>torch.where()</code> and the need for a loop to construct the result tensor.</p>
<p>As a general rule, you will get the best performance out of pytorch if you use<br>
pytorch tensor operations and avoid explicit loops.</p>
<p>Here is a loop-free, <code>cumsum()</code>-based version:</p>
<pre data-code-wrap="python"><code class="lang-python">import torch
print (torch.__version__)

t = torch.tensor ([                                               # example input data
    [0, 1, 0, 0, 0, 1],
    [0, 0, 1, 0, 1, 0],
    [1, 0, 0, 0, 0, 0]
])

u = torch.nn.functional.pad ((t == 0).long(), (0, 1))             # put ones where the zeros are and right pad

v = u.cumsum (dim = 1)                                            # use cumsum to count the run lengths

w = v[:, :-1] * (u.diff (dim = 1) == -1).long()                   # zero out all but the summed run lengths

x = torch.nn.functional.pad (w.cummax (dim = 1).values, (1, 0))   # copy run lengths to the right and left pad

y = x.diff (dim = 1)                                              # take diff to get actual (unsummed) run lengths

result = y.max (dim = 1).values                                   # get max run length per row

print ('result:', result)                                         # should be [3, 2, 5]
</code></pre>
<p>And here is the result:</p>
<pre><code class="lang-plaintext">2.5.1
result: tensor([3, 2, 5])
</code></pre>
<p>It is a little clunky, but it’s the best I could come up with – there might be a cleaner<br>
pytorch approach.  Also, pytorch does do some stuff that numpy doesn’t do, so the<br>
fastest (non-gpu) pytorch version won’t necessarily be as fast as the fastest numpy<br>
version.</p>
<p>Best.</p>
<p>K. Frank</p>