Indexing | how to index the gray regions

PyTorch Live
1

How to index the gray regions? The indexing has to be generic, i.e., I can add D1, D2, D3 and it should still work.

image

2

Assuming you want blocks of the size mxm in matrix of the size nxn :

for i in range round(n/m):
   elements[:,:,i] = matrix[m*(i-1):i*m ,m*(i-1):i*m]

Another way and maybe easier is using torch.block_diag:

A = torch.tensor([[1, 1], [1, 1]])
B = torch.block_diag(A, A, A)

This will result with :

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

And maybe use B == 1