Function extract_diagonal 

pub fn extract_diagonal(
    input: &Tensor,
    axis_a: usize,
    axis_b: usize,
) -> Result<Tensor>