torch.diag
torch.diag(input, diagonal=0, *, out=None) -> Tensor


If "input" is a vector (1-D tensor), then returns a 2-D square
     tensor with the elements of "input" as the diagonal.


If "input" is a matrix (2-D tensor), then returns a 1-D tensor
     with the diagonal elements of "input".


The argument "diagonal" controls which diagonal to consider:


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


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


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


Parameters:
      * input (Tensor) -- the input tensor.
  * **diagonal** (*int**, **optional*) -- the diagonal to consider

Keyword Arguments:
      out (Tensor, optional) -- the output tensor.
See also:
 "torch.diagonal()" always returns the diagonal of its input.

 "torch.diagflat()" always constructs a tensor with diagonal
 elements specified by the input.

Examples: