@wouter’s solution above was what I was looking for.
I would like to point out two things this though.
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.
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.
One way to improve this is to find the places at which the sum of (x>0) along axis equals zero. This indicates that all elements along the axis were negative. We can then set those to some different value, e.g. -1 or `float(‘nan’)’. We can use the already computed cumulative sum to do this.
def first_nonnegative(tensor, axis=0, value_for_all_nonnegative=-1):
nonnegative = (tensor > 0)
cumsum = nonnegative.cumsum(axis=axis)
all_negative = cumsum[-1] == 0 # Any dimensions where all are negative
nonnegative_idx = ((cumsum == 1) & nonnegative).max(0).indices
nonnegative_idx[all_negative] = value_for_all_nonnegative
return nonnegative_idx