Create a matrix based on vector indices

1

I have a one-dimensional tensor a = torch.tensor([0,2,3,4]).
From this tensor, I want to create another two-dimensional tensor which looks like this:

tensor([[1, 1, 0, 0, 0, 0, 0, 0, 0],
        [1, 1, 0, 0, 0, 0, 0, 0, 0],
        [0, 0, 1, 1, 1, 0, 0, 0, 0],
        [0, 0, 1, 1, 1, 0, 0, 0, 0],
        [0, 0, 1, 1, 1, 0, 0, 0, 0],
        [0, 0, 0, 0, 0, 1, 1, 1, 1],
        [0, 0, 0, 0, 0, 1, 1, 1, 1],
        [0, 0, 0, 0, 0, 1, 1, 1, 1],
        [0, 0, 0, 0, 0, 1, 1, 1, 1]])

The output tensor’s dimensions should be equal to (torch.sum(a), torch.sum(a)) which is (9,9). The tensor a determines which elements of the output should be 1. Since, the first two elements of a are 0 and 2, the values in output[0:0+2, 0:0+2] should be 1. And output[2:2+3, 2:2+3] should be 1 since the next two elements are 2 and 3. And finally output[2+3:2+3+4, 2+3:2+3+4] should be 1.
If there an efficient way of doing this in PyTorch?

2

repeat_interleave does the trick
torch.repeat_interleave(torch.repeat_interleave(torch.eye(3), a[1:], dim=0), a[1:], dim=1)

3

torch.blog_diag should also work:

a = torch.tensor([0,2,3,4])
blocks = []
for a_ in a:
    tmp = torch.ones(a_, a_)
    blocks.append(tmp)

torch.block_diag(*[b for b in blocks])
# tensor([[1., 1., 0., 0., 0., 0., 0., 0., 0.],
#         [1., 1., 0., 0., 0., 0., 0., 0., 0.],
#         [0., 0., 1., 1., 1., 0., 0., 0., 0.],
#         [0., 0., 1., 1., 1., 0., 0., 0., 0.],
#         [0., 0., 1., 1., 1., 0., 0., 0., 0.],
#         [0., 0., 0., 0., 0., 1., 1., 1., 1.],
#         [0., 0., 0., 0., 0., 1., 1., 1., 1.],
#         [0., 0., 0., 0., 0., 1., 1., 1., 1.],
#         [0., 0., 0., 0., 0., 1., 1., 1., 1.]])
1 Like
4

Thanks a lot @ptrblck. Is there any way we can do it without the loop? This function has to go inside a neural network.

5

You have to pass the tensors as arguments somehow to block_diag. If you can create these tensors separately somehow without a loop, it would also work.
Btw. you can also just use torch.block_diag(*blocks). In my previous code I had another condition to filter for the empty tensor, which wasn’t needed.

1 Like