Block Diagonal in a 5D tensor

I have a vector A of [16,6,3,27,27] dimension and I want to create a block diagonal of [16,6, 81,81] efficiently using A matrix.

torch. block_diag cannot be used because my vector has more than 2 dimensions.
Do you know how can this be implemented?
Thank you in advance.

Hi etc!

If you create the appropriate index tensors, you can do this with tensor
indexing:

>>> import torch
>>> print (torch.__version__)
1.10.2
>>>
>>> _ = torch.manual_seed (2022)
>>>
>>> A = torch.randn (16, 6, 3, 27, 27)
>>>
>>> B = torch.zeros (16, 6, 81, 81)
>>>
>>> i0 = torch.arange (16).expand (6, 3, 27, 27, 16).permute (4, 0, 1, 2, 3)
>>> i1 = torch.arange (6).expand (16, 3, 27, 27, 6).permute (0, 4, 1, 2, 3)
>>> i2 = torch.arange (81).expand (16, 6, 27, 81).transpose (2, 3).view (16, 6, 3, 27, 27)
>>> i3 = i2.transpose (3, 4)
>>>
>>> B[i0, i1, i2, i3] = A
>>>
>>> torch.equal (A[0, 0, 0], B[0, 0, 0:27, 0:27])
True
>>> torch.equal (A[7, 2, 1], B[7, 2, 27:54, 27:54])
True
>>> torch.equal (A[11, 4, 2], B[11, 4, 54:81, 54:81])
True
>>>
>>> torch.all (B[:, :, 0:27, 27:81] == 0)
tensor(True)

Best.

K. Frank

2 Likes