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.