How can I set maximum n elements in one row in pytorch tensor to 1 and оther to zero?
For example, if I have tensor 

0.8, 0.1, 0.9, 0.2
0.7, 0.1, 0.4, 0.6


and n = 2, I want to get 

1, 0, 1, 0
1, 0, 0, 1


Or maybe it's better to do it numpy, but the question is the same.

For performance efficiency, we can use np.argpartition -

def partition_assign(a, n):
    idx = np.argpartition(a,-n,axis=1)[:,-n:]
    out = np.zeros(a.shape, dtype=int)
    np.put_along_axis(out,idx,1,axis=1)
    return out


Or we can use np.argpartition(-a,n,axis=1)[:,:n] at that argpartition step.

Sample runs -

In [56]: a
Out[56]: 
array([[0.8, 0.1, 0.9, 0.2],
       [0.7, 0.1, 0.4, 0.6]])

In [57]: partition_assign(a, n=2)
Out[57]: 
array([[1, 0, 1, 0],
       [1, 0, 0, 1]])

In [58]: partition_assign(a, n=3)
Out[58]: 
array([[1, 0, 1, 1],
       [1, 0, 1, 1]])

