Hey! I implemented the following algorithm using python for loops, which are, afaik, not very efficient:
sorted, indices = torch.sort(torch.abs(torch.reshape(x, (-1,))), descending=True)
running_sum = -r # r is some positive number
for idx, element in enumerate(sorted):
    running_sum += element
    if element <= (running_sum) / (idx + 1):
         running_sum -= element
         idx -= 1
         break

Is there any way to do this faster using Pytorch operations? The crucial part is finding the first entry of a large sorted tensor such that said entry is less-or-equal than the mean of all previous entries.
Looking forward for your ideas! Thanks a lot!