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)