2

I have two tensors containing batches of matrices of the same batch size (first dimension) but different matrix structure (all other dimensions). For example A of shape (n,d,d) and B (n,e,e). Now I would like to build block diagonals of A and B for all n. So that the output shape (n,(d+e),(d+e)).

Is there an implementation for a problem like this?

I could only find torch.block_diag which is not suited for dimensions higher than 2.

yeich
  • 23
  • 2

1 Answers1

0

Unfortunately there's no vectorized implementation, you'd have to loop through the batch:

A = torch.rand((2, 2, 2))
B = torch.rand((2, 3, 3))
C = torch.zeros((2, 5, 5))
for i in range(2):
    C[i] = torch.block_diag(A[i], B[i])
DMcC
  • 321
  • 2
  • 7