luma_tensor/ops/
indexer.rs

1use crate::{Bool, DTypeKind, Device, Dim, Float, Int, Layout, Shape, Tensor, TensorMeta, dtype::Storage};
2
3pub trait IndexingDTypeKind<D: Device>: DTypeKind<D> + Sized {
4    fn index_select_dispatch(
5        x: &Self::Storage,
6        x_l: &Layout,
7        idx: &D::IntStorage,
8        idx_l: &Layout,
9        dim: usize,
10    ) -> crate::Result<(Self::Storage, Shape)>;
11
12    fn gather_dispatch(
13        x: &Self::Storage,
14        x_l: &Layout,
15        idx: &D::IntStorage,
16        idx_l: &Layout,
17        dim: usize,
18    ) -> crate::Result<(Self::Storage, Shape)>;
19
20    fn index_add_dispatch(
21        init: &Self::Storage,
22        init_l: &Layout,
23        idx: &D::IntStorage,
24        idx_l: &Layout,
25        src: &Self::Storage,
26        src_l: &Layout,
27        dim: usize,
28    ) -> crate::Result<Self::Storage>;
29
30    fn scatter_add_dispatch(
31        init: &Self::Storage,
32        init_l: &Layout,
33        idx: &D::IntStorage,
34        idx_l: &Layout,
35        src: &Self::Storage,
36        src_l: &Layout,
37        dim: usize,
38    ) -> crate::Result<Self::Storage>;
39}
40
41impl<D: Device> IndexingDTypeKind<D> for Float {
42    fn index_select_dispatch(
43        x: &Self::Storage,
44        x_l: &Layout,
45        idx: &D::IntStorage,
46        idx_l: &Layout,
47        dim: usize,
48    ) -> crate::Result<(Self::Storage, Shape)> {
49        D::f_index_select(x, x_l, idx, idx_l, dim)
50    }
51
52    fn gather_dispatch(
53        x: &Self::Storage,
54        x_l: &Layout,
55        idx: &D::IntStorage,
56        idx_l: &Layout,
57        dim: usize,
58    ) -> crate::Result<(Self::Storage, Shape)> {
59        D::f_gather(x, x_l, idx, idx_l, dim)
60    }
61
62    fn index_add_dispatch(
63        init: &Self::Storage,
64        init_l: &Layout,
65        idx: &D::IntStorage,
66        idx_l: &Layout,
67        src: &Self::Storage,
68        src_l: &Layout,
69        dim: usize,
70    ) -> crate::Result<Self::Storage> {
71        D::f_index_add(init, init_l, idx, idx_l, src, src_l, dim)
72    }
73
74    fn scatter_add_dispatch(
75        init: &Self::Storage,
76        init_l: &Layout,
77        idx: &D::IntStorage,
78        idx_l: &Layout,
79        src: &Self::Storage,
80        src_l: &Layout,
81        dim: usize,
82    ) -> crate::Result<Self::Storage> {
83        D::f_scatter_add(init, init_l, idx, idx_l, src, src_l, dim)
84    }
85}
86
87impl<D: Device> IndexingDTypeKind<D> for Int {
88    fn index_select_dispatch(
89        x: &Self::Storage,
90        x_l: &Layout,
91        idx: &D::IntStorage,
92        idx_l: &Layout,
93        dim: usize,
94    ) -> crate::Result<(Self::Storage, Shape)> {
95        D::i_index_select(x, x_l, idx, idx_l, dim)
96    }
97
98    fn gather_dispatch(
99        x: &Self::Storage,
100        x_l: &Layout,
101        idx: &D::IntStorage,
102        idx_l: &Layout,
103        dim: usize,
104    ) -> crate::Result<(Self::Storage, Shape)> {
105        D::i_gather(x, x_l, idx, idx_l, dim)
106    }
107
108    fn index_add_dispatch(
109        init: &Self::Storage,
110        init_l: &Layout,
111        idx: &D::IntStorage,
112        idx_l: &Layout,
113        src: &Self::Storage,
114        src_l: &Layout,
115        dim: usize,
116    ) -> crate::Result<Self::Storage> {
117        D::i_index_add(init, init_l, idx, idx_l, src, src_l, dim)
118    }
119
120    fn scatter_add_dispatch(
121        init: &Self::Storage,
122        init_l: &Layout,
123        idx: &D::IntStorage,
124        idx_l: &Layout,
125        src: &Self::Storage,
126        src_l: &Layout,
127        dim: usize,
128    ) -> crate::Result<Self::Storage> {
129        D::i_scatter_add(init, init_l, idx, idx_l, src, src_l, dim)
130    }
131}
132
133impl<D: Device, K: IndexingDTypeKind<D>> Tensor<D, K> {
134    /// Select slices along `dim` at the given 1-D `indices`.
135    pub fn index_select<Dm: Dim>(&self, indices: &Tensor<D, Int>, dim: Dm) -> crate::Result<Self> {
136        let dim = dim.to_index(self.shape(), "index_select")?;
137        let (storage, shape) =
138            K::index_select_dispatch(&*self.storage_read()?, self.layout(), &*indices.storage_read()?, indices.layout(), dim)?;
139        let meta = K::Meta::on_index_select(self, indices, dim);
140        assert_eq!(self.dtype(), storage.dtype());
141        Ok(Self::from_storage(storage, shape, meta))
142    }
143
144    /// Gather along `dim` using an index tensor of the same rank.
145    pub fn gather<Dm: Dim>(&self, indices: &Tensor<D, Int>, dim: Dm) -> crate::Result<Self> {
146        let dim = dim.to_index(self.shape(), "gather")?;
147        let (storage, shape) = K::gather_dispatch(&*self.storage_read()?, self.layout(), &*indices.storage_read()?, indices.layout(), dim)?;
148        let meta = K::Meta::on_gather(self, indices, dim);
149        assert_eq!(self.dtype(), storage.dtype());
150        Ok(Self::from_storage(storage, shape, meta))
151    }
152
153    /// `out = self; out[.., idx[i], ..] += src[.., i, ..]`.
154    pub fn index_add<Dm: Dim>(&self, indices: &Tensor<D, Int>, src: &Tensor<D, K>, dim: Dm) -> crate::Result<Self> {
155        let dim = dim.to_index(self.shape(), "index_add")?;
156        let storage = K::index_add_dispatch(
157            &*self.storage_read()?,
158            self.layout(),
159            &*indices.storage_read()?,
160            indices.layout(),
161            &*src.storage_read()?,
162            src.layout(),
163            dim,
164        )?;
165        let meta = K::Meta::on_index_add(self, indices, src, dim);
166        assert_eq!(self.dtype(), storage.dtype());
167        Ok(Self::from_storage(storage, self.shape().clone(), meta))
168    }
169
170    /// `out = self; out[.., idx[i,j,k], k] += src[i,j,k]`.
171    pub fn scatter_add<Dm: Dim>(&self, indices: &Tensor<D, Int>, src: &Tensor<D, K>, dim: Dm) -> crate::Result<Self> {
172        let dim = dim.to_index(self.shape(), "scatter_add")?;
173        let storage = K::scatter_add_dispatch(
174            &*self.storage_read()?,
175            self.layout(),
176            &*indices.storage_read()?,
177            indices.layout(),
178            &*src.storage_read()?,
179            src.layout(),
180            dim,
181        )?;
182        let meta = K::Meta::on_scatter_add(self, indices, src, dim);
183        assert_eq!(self.dtype(), storage.dtype());
184        Ok(Self::from_storage(storage, self.shape().clone(), meta))
185    }
186}
187
188// ---- Slice -------------------------------------------------------------------
189
190/// A slice range with `start`, optional `end` (negative = from end, `None` = to end), and `step`.
191#[derive(Debug, Clone, PartialEq, Eq)]
192pub struct Slice {
193    pub start: usize,
194    pub end: Option<isize>,
195    pub step: usize,
196}
197
198impl Slice {
199    pub fn new(start: usize, end: Option<isize>, step: usize) -> Self {
200        Self { start, end, step }
201    }
202
203    /// Resolve `end` against a concrete dimension size.
204    pub fn resolve(&self, dim_size: usize) -> (usize, usize, usize) {
205        let end_abs = match self.end {
206            None => dim_size,
207            Some(e) if e < 0 => {
208                let abs = (-e) as usize;
209                if abs > dim_size { 0 } else { dim_size - abs }
210            }
211            Some(e) => {
212                let e = e as usize;
213                if e > dim_size { dim_size } else { e }
214            }
215        };
216        (self.start, end_abs, self.step)
217    }
218
219    pub fn len(&self) -> usize {
220        self.clone().count()
221    }
222}
223
224impl Iterator for Slice {
225    type Item = usize;
226    fn next(&mut self) -> Option<Self::Item> {
227        match self.end {
228            Some(end) if end < 0 => {
229                let value = self.start;
230                self.start += self.step;
231                Some(value)
232            }
233            Some(end) => {
234                if self.start < end as usize {
235                    let value = self.start;
236                    self.start += self.step;
237                    Some(value)
238                } else {
239                    None
240                }
241            }
242            None => {
243                let value = self.start;
244                self.start += self.step;
245                Some(value)
246            }
247        }
248    }
249}
250
251impl std::fmt::Display for Slice {
252    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
253        let step_part = if self.step == 1 { String::new() } else { format!(":{}", self.step) };
254        match self.end {
255            Some(end) => write!(f, "{}:{}{}", self.start, end, step_part),
256            None => write!(f, "{}:{}", self.start, step_part),
257        }
258    }
259}
260
261// ---- s! macro ----------------------------------------------------------------
262
263/// Create a [`Slice`] with python-like syntax.
264///
265/// ```ignore
266/// s!(1..5)    // Slice { start: 1, end: Some(5), step: 1 }
267/// s!(1:)      // Slice { start: 1, end: None, step: 1 }
268/// s!(1::2)    // Slice { start: 1, end: None, step: 2 }
269/// s!(..5)     // Slice { start: 0, end: Some(5), step: 1 }
270/// s!(:)       // Slice { start: 0, end: None, step: 1 }
271/// s!(::3)     // Slice { start: 0, end: None, step: 3 }
272/// ```
273#[macro_export]
274macro_rules! s {
275    ($start:tt : $end:expr) => {
276        $crate::ops::Slice::new($start as usize, Some($end as isize), 1)
277    };
278    ($start:tt : $end:tt : $step:expr) => {
279        $crate::ops::Slice::new($start as usize, Some($end as isize), $step as usize)
280    };
281    ($start:tt :) => {
282        $crate::ops::Slice::new($start as usize, None, 1)
283    };
284    ($start:tt :: $step:expr) => {
285        $crate::ops::Slice::new($start as usize, None, $step as usize)
286    };
287    (: $end:tt) => {
288        $crate::ops::Slice::new(0, Some($end as isize), 1)
289    };
290    (: $end:tt : $step:expr) => {
291        $crate::ops::Slice::new(0, Some($end as isize), $step as usize)
292    };
293    (:: $step:expr) => {
294        $crate::ops::Slice::new(0, None, $step as usize)
295    };
296    (:) => {
297        $crate::ops::Slice::new(0, None, 1)
298    };
299}
300
301// ---- Indexer -----------------------------------------------------------------
302
303/// One element of a fancy-indexing operation.
304#[derive(Clone)]
305pub enum Indexer<D: Device> {
306    /// Select a single index (removes the dimension).
307    Select(usize),
308    /// Select via a signed dimension index (removes the dimension).
309    SelectD(crate::D),
310    /// Slice a range (keeps the dimension).
311    Slice(Slice),
312    /// Boolean mask filtering (keeps the dimension).
313    Boolean(Tensor<D, Bool>),
314}
315
316// From impls: single values + ranges → Indexer
317
318impl<D: Device> From<usize> for Indexer<D> {
319    fn from(index: usize) -> Self {
320        Indexer::Select(index)
321    }
322}
323
324impl<D: Device> From<crate::D> for Indexer<D> {
325    fn from(index: crate::D) -> Self {
326        Indexer::SelectD(index)
327    }
328}
329
330impl<D: Device> From<Slice> for Indexer<D> {
331    fn from(value: Slice) -> Self {
332        Indexer::Slice(value)
333    }
334}
335
336impl<D: Device> std::fmt::Debug for Indexer<D> {
337    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
338        match self {
339            Self::Select(n) => f.debug_tuple("Select").field(n).finish(),
340            Self::SelectD(d) => f.debug_tuple("SelectD").field(d).finish(),
341            Self::Slice(s) => f.debug_tuple("Slice").field(s).finish(),
342            Self::Boolean(_) => f.debug_tuple("Boolean").field(&"Tensor<..>").finish(),
343        }
344    }
345}
346
347impl<D: Device> From<Tensor<D, Bool>> for Indexer<D> {
348    fn from(value: Tensor<D, Bool>) -> Self {
349        Indexer::Boolean(value)
350    }
351}
352
353impl<D: Device> From<&Tensor<D, Bool>> for Indexer<D> {
354    fn from(value: &Tensor<D, Bool>) -> Self {
355        Indexer::Boolean(value.clone())
356    }
357}
358
359impl<D: Device> From<std::ops::Range<usize>> for Indexer<D> {
360    fn from(value: std::ops::Range<usize>) -> Self {
361        Indexer::Slice(Slice::new(value.start, Some(value.end as isize), 1))
362    }
363}
364
365impl<D: Device> From<std::ops::RangeFrom<usize>> for Indexer<D> {
366    fn from(value: std::ops::RangeFrom<usize>) -> Self {
367        Indexer::Slice(Slice::new(value.start, None, 1))
368    }
369}
370
371impl<D: Device> From<std::ops::RangeTo<usize>> for Indexer<D> {
372    fn from(value: std::ops::RangeTo<usize>) -> Self {
373        Indexer::Slice(Slice::new(0, Some(value.end as isize), 1))
374    }
375}
376
377impl<D: Device> From<std::ops::RangeFull> for Indexer<D> {
378    fn from(_: std::ops::RangeFull) -> Self {
379        Indexer::Slice(Slice::new(0, None, 1))
380    }
381}
382
383// ---- indexes() ---------------------------------------------------------------
384
385impl<D: Device, K: IndexingDTypeKind<D> + crate::ops::shape::ShapeDTypeKind<D>> Tensor<D, K> {
386    /// Apply a sequence of [`Indexer`]s, one per dimension, in order.
387    ///
388    /// - [`Indexer::Select`] / [`Indexer::SelectD`]: narrow + squeeze (removes dim).
389    /// - [`Indexer::Slice`]: narrow (step=1) or slice (step>1) — keeps dim.
390    /// - [`Indexer::Boolean`]: boolean mask → `index_select` — keeps dim.
391    pub fn indexes(&self, indexers: &[Indexer<D>]) -> crate::Result<Self> {
392        let mut x = self.clone();
393        let mut current_dim = 0;
394        for idx in indexers {
395            x = match idx {
396                Indexer::Select(n) => x.narrow(current_dim, *n, 1)?.squeeze(current_dim)?,
397                Indexer::SelectD(d) => {
398                    let dim_size = x.dim(current_dim)?;
399                    let n = d.to_real_index(dim_size, "index")?;
400                    x.narrow(current_dim, n, 1)?.squeeze(current_dim)?
401                }
402                Indexer::Slice(s) => {
403                    let dim_size = x.dim(current_dim)?;
404                    let (start, end, step) = s.resolve(dim_size);
405                    let out = if step == 1 { x.narrow(current_dim, start, end - start)? } else { x.slice(current_dim, start, end, step)? };
406                    current_dim += 1;
407                    out
408                }
409                Indexer::Boolean(mask) => {
410                    let indices: Vec<i64> = mask.to_vec()?.into_iter().enumerate().filter(|(_, v)| *v).map(|(i, _)| i as i64).collect();
411                    let idx_tensor = Tensor::<D, Int>::from_slice(&indices, indices.len(), ())?;
412                    let out = x.index_select(&idx_tensor, current_dim)?;
413                    current_dim += 1;
414                    out
415                }
416            };
417        }
418        Ok(x)
419    }
420}
421
422// ---- IndexOp trait (for .i() syntax) -----------------------------------------
423
424/// Trait for fancy indexing via the `.i()` method.
425pub trait IndexOp<T, D: Device, K: DTypeKind<D>> {
426    fn i(&self, index: T) -> crate::Result<Tensor<D, K>>;
427}
428
429// Single indexer → .i(0) or .i(s!(1..3))
430impl<I, D, K> IndexOp<I, D, K> for Tensor<D, K>
431where
432    I: Into<Indexer<D>>,
433    D: Device,
434    K: IndexingDTypeKind<D> + crate::ops::shape::ShapeDTypeKind<D>,
435{
436    fn i(&self, index: I) -> crate::Result<Tensor<D, K>> {
437        self.indexes(&[index.into()])
438    }
439}
440
441// Tuple of indexers → .i((0, s!(1..)))
442macro_rules! index_op_tuple {
443    ($($t:ident),+) => {
444        #[allow(non_snake_case)]
445        impl<$($t),*, D, K> IndexOp<($($t,)*), D, K> for Tensor<D, K>
446        where
447            $($t: Into<Indexer<D>>,)*
448            D: Device,
449            K: IndexingDTypeKind<D> + crate::ops::shape::ShapeDTypeKind<D>,
450        {
451            fn i(&self, ($($t,)*): ($($t,)*)) -> crate::Result<Tensor<D, K>> {
452                self.indexes(&[$($t.into(),)*])
453            }
454        }
455    };
456}
457
458index_op_tuple!(I1);
459index_op_tuple!(I1, I2);
460index_op_tuple!(I1, I2, I3);
461index_op_tuple!(I1, I2, I3, I4);
462index_op_tuple!(I1, I2, I3, I4, I5);
463
464// Vec<Indexer> → .i(vec![...])
465impl<I, D, K> IndexOp<Vec<I>, D, K> for Tensor<D, K>
466where
467    I: Into<Indexer<D>>,
468    D: Device,
469    K: IndexingDTypeKind<D> + crate::ops::shape::ShapeDTypeKind<D>,
470{
471    fn i(&self, index: Vec<I>) -> crate::Result<Tensor<D, K>> {
472        let idxs: Vec<Indexer<D>> = index.into_iter().map(|i| i.into()).collect();
473        self.indexes(&idxs)
474    }
475}
476
477// ---- get() convenience -------------------------------------------------------
478
479impl<D: Device, K: crate::ops::shape::ShapeDTypeKind<D>> Tensor<D, K> {
480    /// Returns the sub-tensor fixing the index at `i` on the first dimension.
481    ///
482    /// ```ignore
483    /// let t = Tensor::<Cpu>::new(&[[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]]).unwrap();
484    /// let row1 = t.get(1).unwrap(); // [3.0, 4.0]
485    /// ```
486    pub fn get(&self, i: usize) -> crate::Result<Self> {
487        let dims = self.dims();
488        if dims.is_empty() { Ok(self.clone()) } else { self.narrow(0, i, 1)?.reshape(&dims[1..]) }
489    }
490}
luma_tensor/ops/indexer.rs

luma_tensor/ops/
indexer.rs
