torch.diag_embed
torch.diag_embed(input, offset=0, dim1=- 2, dim2=- 1) -> Tensor
Creates a tensor whose diagonals of certain 2D planes (specified by
   "dim1" and "dim2") are filled by "input". To facilitate creating
   batched diagonal matrices, the 2D planes formed by the last two
   dimensions of the returned tensor are chosen by default.
The argument "offset" controls which diagonal to consider:


If "offset" = 0, it is the main diagonal.


If "offset" > 0, it is above the main diagonal.


If "offset" < 0, it is below the main diagonal.


The size of the new matrix will be calculated to make the specified
   diagonal of the size of the last input dimension. Note that for
   "offset" other than 0, the order of "dim1" and "dim2" matters.
   Exchanging them is equivalent to changing the sign of "offset".
Applying "torch.diagonal()" to the output of this function with the
   same arguments yields a matrix identical to input. However,