could someone tell me the best way to get the indices of the first occurrences in a tensor along a specific dimension? By occurrence I mean the indices where the tensor is e.g. larger than a specific value:
For example, consider
my_tensor = torch.rand(10, 100)
my_value = 0.5
ind = (my_tensor < my_value) # Boolean mask whether the values are < 0.5
How do I obtain the first occurrence where ind == 1 along the dimension -1 (in a (10, 1) tensor) ?
For example, ind.nonzero() returns all non-zero values and does not keep the dimensions…
In my shallow view, the .nonzero operation cannot returns the keepdim=True version, because it cannot guarantee each dimension has the same number of elements that fit the nonzero condition, so it performs reduction to 1D and return positions.
But you could get the first occurrence by iteration like this:
As this is still getting some views, here is another approach.
my_tensor = torch.rand(10, 100)
my_value = 0.5
# step by step
bool_mask = (my_tensor < my_value) # Boolean mask with your condition (here < 0.5 but can be anything)```
int_mask = bool_mask.int() # argmax is not working on bool
idx = torch.argmax(int_mask, dim=-1) # argmax returns the first maximum value along a dimension.
# one-liner
idx = (my_tensor < my_value).int().argmax(dim=-1)