CODEWITHSUNDEEP
pythonpytorchtorchdiagonal
yeich
yeich

Reputation: 23

Pytorch Batchwise block diagonal

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.

Upvotes: 2

Views: 604

Answers (1)

DMcC
DMcC

Reputation: 567

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])

Upvotes: 0

Related Questions