Function embedding 

pub fn embedding(
    ids: &Tensor,
    wte: &Tensor,
    wpe: Option<&Tensor>,
    pos: usize,
) -> Result<Tensor>

Fused token + positional embedding: out[t] = wte[ids[t]] + wpe[t + pos].