So, I am trying to modify some of my code so that it supports the batch dimension. One of the lines uses torch.diagflat, and I was wondering what would be the batched version of it?
I see there are the torch.diag and torch.diagonal functions, but it’s not clear if they replicate torch.diagflat?
import torch
x = torch.randn(2, 3) # batch of 2
print(torch.diagflat(x).shape)
# size is torch.Size([6, 6]) instead of torch.Size([2, 3, 3])