How can I set maximum n elements in one row in pytorch tensor to 1 and оther to zero?

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]])

You could use list comprehension. Easy and fast.

First pick out the threshold for each array (based on n).

def tensor_threshold(tensor, n):
    thresholds = [sorted(arr)[-n] for arr in tensor]
    return [[0 if x < th else 1 for x in arr]
            for arr, th in zip(tensor, thresholds)]


T = [[0.8, 0.1, 0.9, 0.2], [0.7, 0.1, 0.4, 0.6]]

tensor_threshold(T, n=2)
>>> [[1, 0, 1, 0], [1, 0, 0, 1]]

You can use torch.topk to get indices and then set it to 1 to new tensor.

t = torch.tensor([[0.8, 0.1, 0.9, 0.2],[0.7, 0.1, 0.4, 0.6]])
tb = torch.zeros(t.shape)                      # create new tensor for 1,0

# set `1` to topk indices
tb[(torch.arange(len(t)).unsqueeze(1), torch.topk(t,2).indices)] = 1
tb
tensor([[1., 0., 1., 0.],
        [1., 0., 0., 1.]])

or n=3

tb = torch.zeros(t.shape)
tb[(torch.arange(len(t)).unsqueeze(1), torch.topk(t,3).indices)] = 1
tb
tensor([[1., 0., 1., 1.],
        [1., 0., 1., 1.]])