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 00 , the order of dim1
and di