How to randomly set a variable number of elements in each row of a tensor in PyTorch

1

I want to create a zero-one tensor of dimension (n, n). The ones should be placed randomly, with a cap on the number of ones in each row. Let us say I have a list of length n that has the value of cap for each of the n rows.

One way very close to what I want is to create a random tensor and mask it using torch.topk. However, in torch.topk, k is an integer and remains same for the dimension specified. I need different value of k for every row, so that I can sample variable number of top k elements and turn them into 1s.

2

Hope could help you.

import torch 

n = 5 
top_k = torch.randint(n, (n,))
print(top_k)

# Step 1, generate a variable number of elements in each row
top_k = torch.nn.functional.one_hot(top_k, num_classes=n)
top_k = 1 - torch.cumsum(top_k, dim=1)  # 
print(top_k)

# Step 2, shuffle each row 
indices = torch.argsort(torch.randn((n, n)), dim=1)
result = top_k[torch.arange(top_k.shape[0])[..., None], indices]
print(result)