How to randomly set a variable number of elements in each row of a tensor in PyTorch
I want to create a zero-one matrix 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. How can I do this in PyTorch?
My question is similar to this previous question. The only change I am looking for is, there should be n values of k, corresponding to n rows.

As explained by @Marcel in the comments above, you can first set the first m values to value k then index by permuted indices in order to get a shuffle tensor:
&gt;&gt;&gt; n = 10; m = 3; k = 1
&gt;&gt;&gt; x = torch.zeros(n, n)

&gt;&gt;&gt; x[:, :m] = k
tensor([[1., 1., 1., 0., 0., 0., 0., 0., 0., 0.],
        [1., 1., 1., 0., 0., 0., 0., 0., 0., 0.],
        [1., 1., 1., 0., 0., 0., 0., 0., 0., 0.],
        [1., 1., 1., 0., 0., 0., 0., 0., 0., 0.],
        [1., 1., 1., 0., 0., 0., 0., 0., 0., 0.],
        [1., 1., 1., 0., 0., 0., 0., 0., 0., 0.],
        [1., 1., 1., 0., 0., 0., 0., 0., 0., 0.],
        [1., 1., 1., 0., 0., 0., 0., 0., 0., 0.],
        [1., 1., 1., 0., 0., 0., 0., 0., 0., 0.],
        [1., 1., 1., 0., 0., 0., 0., 0., 0., 0.]])

Use torch.randperm to get row-wise column permutations:
&gt;&gt;&gt; perm = torch.stack([torch.randperm(10) for _ in range(len(x))])
tensor([[8, 0, 3, 2, 1, 6, 9, 4, 5, 7],
        [5, 7, 1, 4, 8, 0, 6, 9, 2, 3],
        [2, 1, 9, 7, 0, 8, 6, 3, 5, 4],
        [1, 3, 5, 8, 7, 6, 9, 4, 2, 0],
        [7, 6, 0, 5, 2, 9, 1, 8, 4, 3],
        [5, 0, 6, 8, 1, 9, 2, 4, 3, 7],
        [4, 0, 6, 5, 8, 1, 3, 7, 2, 9],
        [5, 3, 4, 9, 0, 1, 7, 6, 8, 2],
        [5, 7, 9, 3, 2, 6, 8, 0, 4, 1],
        [2, 7, 4, 6, 3, 0, 9, 8, 5, 1]])

Then use torch.gather to index the tensor x with perm:
&gt;&gt;&gt; x.gather(dim=0, index=perm)
tensor([[0., 1., 0., 1., 1., 0., 0., 0., 0., 0.],
        [0., 0., 1., 0., 0., 1., 0., 0., 1., 0.],
        [1., 1., 0., 0., 1., 0., 0., 0., 0., 0.],
        [1., 0., 0., 0., 0., 0., 0., 0., 1., 1.],
        [0., 0., 1., 0., 1., 0., 1., 0., 0., 0.],
        [0., 1., 0., 0., 1., 0., 1., 0., 0., 0.],
        [0., 1., 0., 0., 0., 1., 0., 0., 1., 0.],
        [0., 0., 0., 0., 1., 1., 0., 0., 0., 1.],
        [0., 0., 0., 0., 1., 0., 0., 1., 0., 1.],
        [1., 0., 0., 0., 0., 1., 0., 0., 0., 1.]])

Alternatively you can use torch.scatter straight way with the value keyword argument:
&gt;&gt;&gt; torch.zeros(n, n).scatter(dim=0, index=perm, value=1)
tensor([[0., 1., 0., 1., 1., 0., 0., 0., 0., 0.],
        [0., 0., 1., 0., 0., 1., 0., 0., 1., 0.],
        [1., 1., 0., 0., 1., 0., 0., 0., 0., 0.],
        [1., 0., 0., 0., 0., 0., 0., 0., 1., 1.],
        [0., 0., 1., 0., 1., 0., 1., 0., 0., 0.],
        [0., 1., 0., 0., 1., 0., 1., 0., 0., 0.],
        [0., 1., 0., 0., 0., 1., 0., 0., 1., 0.],
        [0., 0., 0., 0., 1., 1., 0., 0., 0., 1.],
        [0., 0., 0., 0., 1., 0., 0., 1., 0., 1.],
        [1., 0., 0., 0., 0., 1., 0., 0., 0., 1.]])


If m is a tensor itself, you can find a workaround using a combination of torch.arange and torch.where:
First encode the positions:
&gt;&gt;&gt; d = torch.arange(n)[None].repeat(n,1)
&gt;&gt;&gt; x = torch.where(d+m&gt;n, 0, 1)
tensor([[1, 1, 1, 1, 1, 1, 0, 0, 0, 0],
        [1, 1, 1, 0, 0, 0, 0, 0, 0, 0],
        [1, 1, 1, 1, 1, 1, 1, 0, 0, 0],
        [1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
        [1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
        [1, 1, 1, 1, 1, 1, 1, 1, 0, 0],
        [1, 1, 1, 1, 1, 1, 1, 1, 0, 0],
        [1, 1, 0, 0, 0, 0, 0, 0, 0, 0],
        [1, 1, 1, 1, 1, 1, 1, 1, 1, 0],
        [1, 1, 1, 1, 1, 1, 1, 1, 1, 1]])

Construct the permutation as before:
&gt;&gt;&gt; perm = torch.stack([torch.randperm(10) for _ in range(n)])
tensor([[2, 5, 7, 0, 4, 1, 3, 6, 8, 9],
        [7, 4, 9, 5, 6, 0, 3, 1, 2, 8],
        [5, 1, 4, 9, 0, 3, 2, 6, 7, 8],
        [9, 6, 0, 2, 3, 1, 7, 5, 4, 8],
        [3, 5, 4, 6, 0, 7, 9, 8, 2, 1],
        [5, 7, 8, 6, 9, 2, 0, 4, 3, 1],
        [8, 3, 9, 0, 6, 2, 5, 7, 4, 1],
        [2, 9, 4, 3, 7, 8, 1, 0, 6, 5],
        [5, 4, 8, 3, 2, 9, 7, 1, 6, 0],
        [8, 7, 3, 6, 5, 4, 2, 0, 9, 1]])

Then scatter on x:
&gt;&gt;&gt; x.scatter(dim=0, index=perm, value=1)
tensor([[1, 1, 1, 1, 1, 1, 1, 1, 0, 1],
        [1, 1, 1, 0, 0, 1, 1, 1, 0, 1],
        [1, 1, 1, 1, 1, 1, 1, 0, 1, 0],
        [1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
        [1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
        [1, 1, 1, 1, 1, 1, 1, 1, 0, 1],
        [1, 1, 1, 1, 1, 1, 1, 1, 1, 0],
        [1, 1, 1, 0, 1, 1, 1, 1, 1, 0],
        [1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
        [1, 1, 1, 1, 1, 1, 1, 1, 1, 1]])

