Finding the maximum elements whose sum is less than a given value in PyTorch
Given any PyTorch 2D tensor, what will be the most efficient way to compute the number of top-K values for each row whose sum is less than a given value?
Input:
tensor([[0.6607, 0.1165, 0.0278, 0.1950],
        [0.0529, 0.4607, 0.2729, 0.2135],
        [0.3267, 0.0902, 0.4578, 0.1253]])

Required Output for the given value 0.8:
tensor([[1], #as 0.6607+0.1950 &gt; 0.8
        [2], #as 0.4607+0.2729+0.2135 &gt; 0.8
        [2]]) #as 0.4578+0.3267+0.1253 &gt; 0.8


You can manage such operation by using a combination of sorting, cumulative sum, and max pooling.
First sort the values by row with torch.Tensor.sort
&gt;&gt;&gt; v = x.sort(dim=1, descending=True).values
tensor([[0.6607, 0.1950, 0.1165, 0.0278],
        [0.4607, 0.2729, 0.2135, 0.0529],
        [0.4578, 0.3267, 0.1253, 0.0902]])

Then construct a mask on the cumulative sorted values that you get from applying torch.cumsum:
&gt;&gt;&gt; torch.cumsum(v, dim=1) &gt; .8
tensor([[False,  True,  True,  True],
        [False, False,  True,  True],
        [False, False,  True,  True]])

Applying a torch.Tensor.max on that mask will return the index of the first occurring True value, i.e. the location of the cumulative element which is above the threshold value 0.8:
&gt;&gt;&gt; mask.max(1, True).indices
tensor([[1],
        [2],
        [2]])

