Function embed_diagonal 

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