Find maximum length of consecutive zeros in each row

My goal is to find the maximum length of consecutive zeros in each row. So for instance, if I have a tensor like

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

I expect to get the result

tensor([3, 2, 5])

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.
Thanks!

# 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)

When you talk about efficiency what do you mean? Time? Accuracy?
I tried this with torch and it worked:

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])

Hi Paulo!

If you use cumsum() to “count” the lengths of the sums, you can avoid using
torch.where() and the need for a loop to construct the result tensor.

As a general rule, you will get the best performance out of pytorch if you use
pytorch tensor operations and avoid explicit loops.

Here is a loop-free, cumsum()-based version:

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]

And here is the result:

2.5.1
result: tensor([3, 2, 5])

It is a little clunky, but it’s the best I could come up with – there might be a cleaner
pytorch approach. Also, pytorch does do some stuff that numpy doesn’t do, so the
fastest (non-gpu) pytorch version won’t necessarily be as fast as the fastest numpy
version.

Best.

K. Frank