Creating a Block-Diagonal Matrix

1

Hey,

I am wondering what the fastest way to create a block-diagonal matrix would be using PyTorch?

As an example, given

a = torch.FloatTensor([[1,2],[3,4]])
b = torch.FloatTensor([[5,6],[7,8]])

I would like to create
c = torch.FloatTensor([[1,2,0,0],[3,4,0,0],[0,0,5,6],[0,0,7,8]])

Any help much appreciated.
Many thanks,
Max

2 Likes
2

Hey,
I figured out a solution. It’s not a perfect one but it works.
Create your blocks:

b_1 = torch.FloatTensor([[1,2],[3,4]])
b_2 = torch.FloatTensor([[5,6],[7,8]])

By concatenating the b_i’s with zeros you to get rows:
r_1 = torch.cat((b_1, torch.zeros((2, 2))), dim = 1)
r_2= torch.cat((torch.zeros((2, 2)), b_2), dim = 1)

And then you concatenate the rows to get c
c = torch.cat((r_1, r_2), dim = 0)

I hope one Max was able to help the other Max :wink:

1 Like
3

I made a vectorized implementation on Gist: block_diag.py at https://gist.github.com/yulkang/2e4fc3061b45403f455d7f4c316ab168

For an up-to-date version, check block_diag() in numpytorch.py in my pylabyk library: https://github.com/yulkang/pylabyk/blob/master/numpytorch.py

1 Like
4

PyTorch has an implementation for this already. Checkout this:
https://pytorch.org/docs/stable/generated/torch.block_diag.html

An example from the docs:

>>> import torch
>>> A = torch.tensor([[0, 1], [1, 0]])
>>> B = torch.tensor([[3, 4, 5], [6, 7, 8]])
>>> C = torch.tensor(7)
>>> D = torch.tensor([1, 2, 3])
>>> E = torch.tensor([[4], [5], [6]])
>>> torch.block_diag(A, B, C, D, E)
tensor([[0, 1, 0, 0, 0, 0, 0, 0, 0, 0],
        [1, 0, 0, 0, 0, 0, 0, 0, 0, 0],
        [0, 0, 3, 4, 5, 0, 0, 0, 0, 0],
        [0, 0, 6, 7, 8, 0, 0, 0, 0, 0],
        [0, 0, 0, 0, 0, 7, 0, 0, 0, 0],
        [0, 0, 0, 0, 0, 0, 1, 2, 3, 0],
        [0, 0, 0, 0, 0, 0, 0, 0, 0, 4],
        [0, 0, 0, 0, 0, 0, 0, 0, 0, 5],
        [0, 0, 0, 0, 0, 0, 0, 0, 0, 6]])