Pytorch: Set Block-Diagonal Matrix Efficiently?
I have a Tensor A of size [N x 3 x 3], and a Matrix B of size [N*3 x N*3]
I want to copy the contents of A -&gt; B, so that the diagonal elements are filled up basically, and I want to do this efficiently:
It should kind of fill up B to look something filled like this:

So each [i,3,3] fills into each [3x3] part in B diagonally down the line.
How do I do this? As efficiently as possible as this is for a real time application. I could write a CUDA kernel to do this, but I would prefer to do it with some special Pytorch function

Use torch.block_diag():
# Setup
A = torch.ones(3,3,3, dtype=int)

# Unpack blocks and apply
B = torch.block_diag(*A)

&gt;&gt;&gt; B
tensor([[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, 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],
        [0, 0, 0, 0, 0, 0, 1, 1, 1],
        [0, 0, 0, 0, 0, 0, 1, 1, 1]])

