kopitiam_tensor/tensor/
gather.rs

1//! Embedding-style row gather.
2
3use kopitiam_core::{DType, Error, Result, Shape};
4
5use crate::storage::Storage;
6
7use super::Tensor;
8
9impl Tensor {
10    /// Gathers rows of a `[vocab, hidden]` embedding table by token id.
11    ///
12    /// `indices` (dtype [`DType::I32`], any shape) holds token ids in
13    /// `[0, vocab)`. The result has shape `indices.shape() + [hidden]` —
14    /// every index becomes one full row of `self`.
15    ///
16    /// # Why this is not `torch.gather`
17    ///
18    /// PyTorch's `gather(dim, index)` requires `index` to have the *same
19    /// rank* as the source and picks one scalar per output position along
20    /// `dim`. That is a general primitive this crate deliberately does not
21    /// implement: nothing in a forward pass needs it. What every
22    /// transformer *does* need is embedding lookup — pick whole rows out
23    /// of a `[vocab, hidden]` table by a batch of token ids — which is a
24    /// different (simpler, more common) shape of operation with its own
25    /// name here rather than a special case of a more general `gather`
26    /// that this crate does not otherwise use.
27    pub fn gather_rows(&self, indices: &Tensor) -> Result<Tensor> {
28        self.require_dtype(DType::F32)?;
29        indices.require_dtype(DType::I32)?;
30        if self.rank() != 2 {
31            return Err(Error::ShapeMismatch { expected: self.shape.clone(), actual: self.shape.clone() });
32        }
33        let vocab = self.shape.dims()[0];
34        let hidden = self.shape.dims()[1];
35
36        let Storage::F32(data) = self.storage.as_ref() else { unreachable!() };
37        let Storage::I32(idx_data) = indices.storage.as_ref() else { unreachable!() };
38
39        let mut out = Vec::with_capacity(indices.elem_count() * hidden);
40        for offset in indices.logical_offsets() {
41            let id = idx_data[offset];
42            if id < 0 || id as usize >= vocab {
43                return Err(Error::IndexOutOfBounds {
44                    dim: 0,
45                    index: id.max(0) as usize,
46                    len: vocab,
47                });
48            }
49            let row_start = self.offset + id as usize * self.strides[0];
50            for c in 0..hidden {
51                out.push(data[row_start + c * self.strides[1]]);
52            }
53        }
54
55        let mut out_dims = indices.shape.dims().to_vec();
56        out_dims.push(hidden);
57        Tensor::from_f32(out, Shape::new(out_dims))
58    }
59}
60
61#[cfg(test)]
62mod tests {
63    use super::*;
64
65    #[test]
66    fn gather_rows_picks_the_requested_embedding_rows() {
67        // vocab=4, hidden=2: rows are [0,1], [2,3], [4,5], [6,7].
68        let table = Tensor::from_f32((0..8).map(|v| v as f32).collect(), [4, 2]).unwrap();
69        let ids = Tensor::from_i32(vec![2, 0, 3], [3]).unwrap();
70        let out = table.gather_rows(&ids).unwrap();
71        assert_eq!(out.shape().dims(), &[3, 2]);
72        assert_eq!(out.to_vec_f32().unwrap(), vec![4.0, 5.0, 0.0, 1.0, 6.0, 7.0]);
73    }
74
75    #[test]
76    fn gather_rows_preserves_the_indices_shape_for_batched_lookup() {
77        // A [batch=2, seq=2] block of token ids gathers into [2, 2, hidden].
78        let table = Tensor::from_f32((0..6).map(|v| v as f32).collect(), [3, 2]).unwrap();
79        let ids = Tensor::from_i32(vec![0, 1, 2, 0], [2, 2]).unwrap();
80        let out = table.gather_rows(&ids).unwrap();
81        assert_eq!(out.shape().dims(), &[2, 2, 2]);
82        assert_eq!(out.to_vec_f32().unwrap(), vec![0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 0.0, 1.0]);
83    }
84
85    #[test]
86    fn gather_rows_works_through_a_transposed_view_of_the_table() {
87        // table stored as [hidden, vocab], transposed to [vocab, hidden]
88        // before gathering — exercises the strides[]/offset-aware path.
89        let table = Tensor::from_f32(vec![0.0, 2.0, 4.0, 6.0, 1.0, 3.0, 5.0, 7.0], [2, 4]).unwrap();
90        let transposed = table.transpose(0, 1).unwrap(); // [4, 2]: rows [0,1],[2,3],[4,5],[6,7]
91        let ids = Tensor::from_i32(vec![1], [1]).unwrap();
92        let out = transposed.gather_rows(&ids).unwrap();
93        assert_eq!(out.to_vec_f32().unwrap(), vec![2.0, 3.0]);
94    }
95
96    #[test]
97    fn gather_rows_rejects_an_out_of_range_id() {
98        let table = Tensor::from_f32(vec![1.0, 2.0], [1, 2]).unwrap();
99        let ids = Tensor::from_i32(vec![5], [1]).unwrap();
100        assert!(matches!(table.gather_rows(&ids), Err(Error::IndexOutOfBounds { .. })));
101    }
102
103    #[test]
104    fn gather_rows_requires_i32_indices() {
105        let table = Tensor::from_f32(vec![1.0, 2.0], [1, 2]).unwrap();
106        let bad_ids = Tensor::from_f32(vec![0.0], [1]).unwrap();
107        assert!(matches!(table.gather_rows(&bad_ids), Err(Error::DTypeMismatch { .. })));
108    }
109}
kopitiam_tensor/tensor/gather.rs

kopitiam_tensor/tensor/
gather.rs
