Method embedding

cake-core/src/backends/mod.rs:513–528

Embedding lookup: select rows from weight matrix by token IDs.

(&self, ids: &Tensor, weight: &Tensor)

    /// Embedding lookup: select rows from weight matrix by token IDs.
    fn embedding(&self, ids: &Tensor, weight: &Tensor) -> Result<Tensor> {
        let hidden_size = weight.dim(1)?;
        let dims = ids.dims();
        // Fast path: 1D input — no reshape needed
        if dims.len() == 1 {
            let selected = weight.index_select(ids, 0)?;
            return Ok(selected);
        }
        // General path: flatten, select, reshape
        let elem_count: usize = dims.iter().product();
        let flat_ids = ids.reshape(elem_count)?;
        let selected = weight.index_select(&flat_ids, 0)?;
        let mut out_shape = dims.to_vec();
        out_shape.push(hidden_size);
        selected.reshape(out_shape.as_slice())
    }

    /// Create a causal attention mask. Returns a U8 tensor of shape `(seq_len, kv_len)`
    /// where 1 = masked (future position), 0 = attend.