Function diag 

pub fn diag<Args, Inp>(param: Args) -> <Args as DiagAPI<Inp>>::Out
where
    Args: DiagAPI<Inp>,

Extract a diagonal or construct a diagonal tensor.

• If input is a 2-D tensor, return a copy of its diagonal (with offset).
• If input is a 1-D tensor, construct a 2-D tensor with the input as its diagonal.

See also

• numpy.diag