torch.diagonal
torch.diagonal(input, offset=0, dim1=0, dim2=1) -> Tensor
Returns a partial view of "input" with the its diagonal elements
   with respect to "dim1" and "dim2" appended as a dimension at the
   end of the shape.
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.


Applying "torch.diag_embed()" to the output of this function with
   the same arguments yields a diagonal matrix with the diagonal
   entries of the input. However, "torch.diag_embed()" has different
   default dimensions, so those need to be explicitly specified.
Parameters:
      * input (Tensor) -- the input tensor. Must be at least
        2-dimensional.
  * **offset** (*int**, **optional*) -- which diagonal to
    consider. Default: 0 (main diagonal).

  * **dim1** (*int**, **optional*) -- first dimension with respect
