Make (off-)diagonal matrix from a given tensor of shape `(batch_size, seq_len)`

1

I have a tensor of shape (batch_size, seq_len). From this, I want to make a tensor of shape (bs, seq_len+1, seq_len+1) where each tensor corresponds to a single batch is an off-diagonal matrix with entries from given tensor.

More precisely, from a given tensor
x = [[1, 2, 3], [4, 5, 6]],
I want to make a function that returns
y = [[[0, 1, 0, 0], [0, 0, 2, 0], [0, 0, 0, 3], [0, 0, 0, 0]],
[[0, 4, 0, 0], [0, 0, 5, 0], [0, 0, 0, 6], [0, 0, 0, 0]]].

I can’t directly apply torch.diag since it gives diagonals of given tensor when the tensor has dimension > 1. I think it is possible to do with torch.diagflat, but it is highly memory-inefficient since it constructs a tensor of shape (batch_size * seq_len - 1, batch_size * seq_len - 1).

EDIT: I just found that, for my purpose, it is enough to construct following function: transform
x = [[1, 2, 3], [4, 5, 6]]
to
y = [[[1, 0, 0], [0, 2, 0], [0, 0, 3]],
[[4, 0, 0], [0, 5, 0], [0, 0, 6]]]
i.e. transform each row vectors into diagonal matrices, not off-diagonals.

2

This code should work:

x = torch.tensor([[1, 2, 3], [4, 5, 6]]).float()
idx = torch.arange(3).view(1, 1, 3).repeat(2, 1, 1)
res = torch.zeros(2, 3, 3).scatter_(1, idx , x.unsqueeze(1))