lumen_core/ops/
indexer.rs

1use std::{fmt::Display, ops::Deref};
2use crate::{AutogradMetaT, Dim, Error, IntTensor, NumDType, Result, WithDType, D};
3use crate::Tensor;
4
5impl<T: WithDType> Tensor<T> {
6    pub fn indexes(&self, indexers: &[Indexer]) -> Result<Self> {
7        let mut x = self.clone();
8        let mut current_dim = 0;
9        for indexer in indexers.iter() {
10            x = match indexer {
11                Indexer::Select(n) => x.narrow(current_dim, *n, 1)?.squeeze(current_dim)?,
12                Indexer::SelectD(d) => {
13                    let dim_size = x.dim(current_dim)?;
14                    let n = d.to_real_index(dim_size, "index")?;
15                    x.narrow(current_dim, n, 1)?.squeeze(current_dim)?
16                }
17                Indexer::Slice(range) => {
18                    let out = x.slice(current_dim, range)?;
19                    current_dim += 1;
20                    out
21                }
22                Indexer::Boolean(index) => {
23                    if index.dims1()? != x.dims()[current_dim] {
24                        return Err(Error::BooleanIndexShouldLikeVector(index.shape().clone()))
25                    }
26
27                    // index is  a 1'dim, len == x.dims(current_dim)
28                    // TODO: better
29                    let index = index.to_vec()?
30                        .into_iter()
31                        .enumerate()
32                        .filter(|(_, v)| *v)
33                        .map(|(i, _)| i as u32)
34                        .collect::<Vec<_>>();
35                    let index = Tensor::new(index)?;
36                    let out = x.index_select(index, current_dim)?;
37                    current_dim += 1;
38                    out
39                }
40            };
41        }
42        Ok(x)
43    }
44
45    /// Returns the sub-tensor fixing the index at `i` on the first dimension.
46    ///
47    /// ```rust
48    /// use lumen_core::Tensor;
49    /// let tensor = Tensor::<f32>::new(&[[0f32, 1.], [2., 3.], [4., 5.]]).unwrap();
50    /// let t = tensor.get(0).unwrap();
51    /// assert_eq!(t.to_vec().unwrap(), &[0., 1.]);
52    /// let t = tensor.get(1).unwrap();
53    /// assert_eq!(t.to_vec().unwrap(), &[2., 3.]);
54    /// ```
55    pub fn get(&self, i: usize) -> Result<Self> {
56        let dims = self.dims();
57        if dims.is_empty() {
58            Ok(self.clone())
59        } else {
60            self.narrow(0, i, 1)?.reshape(&dims[1..])
61        }
62    }
63
64    pub fn index_select<D: Dim>(&self, indexes: impl Into<IntTensor>, dim: D) -> Result<Self> {   
65        let indexes: IntTensor = indexes.into();
66        let dim = dim.to_index(self.shape(), "index-select")?;
67        let indexes_len = indexes.shape().dims1()?;
68        let mut dims = self.dims().to_vec();
69        dims[dim] = indexes_len;
70        let meta = T::AutogradMeta::on_index_select_op(self, &indexes, dim);
71        let storage = match indexes {
72            IntTensor::I32(indexes) => self.storage_read()?.index_select(
73                self.layout(),
74                indexes.storage_read()?.deref(),
75                indexes.layout(),
76                dim,
77            )?,
78            IntTensor::U32(indexes) => self.storage_read()?.index_select(
79                self.layout(),
80                indexes.storage_read()?.deref(),
81                indexes.layout(),
82                dim,
83            )?,
84            IntTensor::U8(indexes) => self.storage_read()?.index_select(
85                self.layout(),
86                indexes.storage_read()?.deref(),
87                indexes.layout(),
88                dim,
89            )?,
90        };
91
92        Ok(Self::from_storage(storage, dims, meta))
93    }
94
95    /// Gather values across the target dimension.
96    ///
97    /// # Arguments
98    ///
99    /// * `self` - The input tensor.
100    /// * `indexes` - The indices of elements to gather, this should have same number of dimensions as `self`
101    ///   and indexes.dims()[d] <= self.dims()[d] for all dimensions d != dim
102    /// * `dim` - the target dimension.
103    ///
104    /// The resulting tensor has the same shape as `indexes` and use values from `self` indexed on
105    /// dimension `dim` by the values in `indexes`.
106    pub fn gather<D: Dim>(&self, indexes: impl Into<IntTensor>, dim: D) -> Result<Self> {
107        let indexes = indexes.into();
108        let dim = dim.to_index(self.shape(), "gather")?;
109        let self_dims = self.dims();
110        let indexes_dims = indexes.dims();
111        let mismatch = if indexes_dims.len() != self_dims.len() {
112            true
113        } else {
114            let mut mismatch = false;
115            for (i, (&d1, &d2)) in self_dims.iter().zip(indexes_dims.iter()).enumerate() {
116                if i != dim && d1 < d2 {
117                    mismatch = true;
118                    break;
119                }
120            }
121            mismatch
122        };
123        if mismatch {
124            Err(Error::ShapeMismatchBinaryOp {
125                op: "gather",
126                lhs: self.shape().clone(),
127                rhs: indexes.shape().clone(),
128            })?
129        }
130
131        let storage = match &indexes {
132            IntTensor::I32(idx) => self.storage_read()?.gather(self.layout(), idx.storage_read()?.deref(), idx.layout(), dim)?,
133            IntTensor::U32(idx) => self.storage_read()?.gather(self.layout(), idx.storage_read()?.deref(), idx.layout(), dim)?,
134            IntTensor::U8(idx) => self.storage_read()?.gather(self.layout(), idx.storage_read()?.deref(), idx.layout(), dim)?,
135        };
136
137        let meta = T::AutogradMeta::on_gather_op(self, &indexes, dim);
138        Ok(Self::from_storage(storage, indexes.shape(), meta))
139    }
140}
141
142impl<T: NumDType> Tensor<T> {
143    pub fn index_add<D: Dim>(&self, indexes: impl Into<IntTensor>, source: &Tensor<T>, dim: D) -> Result<Self> {
144        let indexes: IntTensor = indexes.into();
145        let dim = dim.to_index(self.shape(), "index-add")?;
146        
147        let source_dims = source.dims();
148        let self_dims = self.dims();
149        if source_dims.len() != self_dims.len() {
150             return Err(Error::ShapeMismatchBinaryOp { 
151                 op: "index-add", 
152                 lhs: self.shape().clone(), 
153                 rhs: source.shape().clone() 
154             }.into());
155        }
156
157        let indexes_len = indexes.shape().dims1()?;
158        for (i, (&d_self, &d_src)) in self_dims.iter().zip(source_dims.iter()).enumerate() {
159            if i == dim {
160                if d_src != indexes_len {
161                    return Err(Error::ShapeMismatchBinaryOp { op: "index-add (dim mismatch)", lhs: self.shape().clone(), rhs: source.shape().clone() }.into());
162                }
163            } else if d_self != d_src {
164                return Err(Error::ShapeMismatchBinaryOp { op: "index-add", lhs: self.shape().clone(), rhs: source.shape().clone() }.into());
165            }
166        }
167
168        let storage = match &indexes {
169            IntTensor::I32(idx) => self.storage_read()?.index_add(
170                self.layout(),
171                idx.storage_read()?.deref(),
172                idx.layout(),
173                source.storage_read()?.deref(),
174                source.layout(),
175                dim,
176            )?,
177            IntTensor::U32(idx) => self.storage_read()?.index_add(
178                self.layout(),
179                idx.storage_read()?.deref(),
180                idx.layout(),
181                source.storage_read()?.deref(),
182                source.layout(),
183                dim,
184            )?,
185            IntTensor::U8(idx) => self.storage_read()?.index_add(
186                self.layout(),
187                idx.storage_read()?.deref(),
188                idx.layout(),
189                source.storage_read()?.deref(),
190                source.layout(),
191                dim,
192            )?,
193        };
194
195        let meta = T::AutogradMeta::on_index_add_op(self, &indexes, source, dim);
196        Ok(Self::from_storage(storage, self_dims.to_vec(), meta))
197    } 
198
199    pub fn scatter_add<D: Dim>(&self, indexes: impl Into<IntTensor>, source: &Self, dim: D) -> Result<Self> {
200        let indexes = indexes.into();
201        let dim = dim.to_index(self.shape(), "scatter-add")?;
202        self.scatter_checks(&indexes, source, dim)?;
203
204        let storage = match &indexes {
205            IntTensor::I32(idx) => self.storage_read()?.scatter_add(
206                self.layout(),
207                idx.storage_read()?.deref(),
208                idx.layout(),
209                source.storage_read()?.deref(),
210                source.layout(),
211                dim,
212            )?,
213            IntTensor::U32(idx) => self.storage_read()?.scatter_add(
214                self.layout(),
215                idx.storage_read()?.deref(),
216                idx.layout(),
217                source.storage_read()?.deref(),
218                source.layout(),
219                dim,
220            )?,
221            IntTensor::U8(idx) => self.storage_read()?.scatter_add(
222                self.layout(),
223                idx.storage_read()?.deref(),
224                idx.layout(),
225                source.storage_read()?.deref(),
226                source.layout(),
227                dim,
228            )?,
229        };
230
231        let meta = T::AutogradMeta::on_scatter_add_op(self, &indexes, source, dim);
232        Ok(Self::from_storage(storage, self.shape(), meta))
233    }
234
235    fn scatter_checks(&self, indexes: &IntTensor, source: &Self, dim: usize) -> Result<()> {
236        let source_dims = source.dims();
237        let self_dims = self.dims();
238        let mismatch = if source_dims.len() != self_dims.len() {
239            true
240        } else {
241            let mut mismatch = false;
242            for (i, (&d1, &d2)) in self_dims.iter().zip(source_dims.iter()).enumerate() {
243                if i != dim && d1 != d2 {
244                    mismatch = true;
245                    break;
246                }
247            }
248            mismatch
249        };
250        if mismatch {
251            Err(Error::ShapeMismatchBinaryOp {
252                op: "scatter (self, src)",
253                lhs: self.shape().clone(),
254                rhs: source.shape().clone(),
255            })?
256        }
257        if indexes.dims() != source.dims() {
258            Err(Error::ShapeMismatchBinaryOp {
259                op: "scatter (indexes, src)",
260                lhs: indexes.shape().clone(),
261                rhs: source.shape().clone(),
262            })?
263        }
264        Ok(())
265    }
266}
267
268impl<T: WithDType> Tensor<T> {
269    pub fn matrix_get(&self, row: usize, col: usize) -> Result<T> {
270        self.index((row, col))?.to_scalar()
271    }
272
273    pub fn matrix_set(&self, row: usize, col: usize, val: T) -> Result<()> {
274        self.index((row, col))?.set_scalar(val)
275    }
276
277    pub fn vector_get(&self, n: usize) -> Result<T> {
278        self.index(n)?.to_scalar()
279    }
280}
281
282#[derive(Debug, Clone, PartialEq, Eq)]
283pub enum Indexer {
284    Select(usize),
285    SelectD(D),
286    Slice(Slice),
287    Boolean(Tensor<bool>),
288}
289
290#[derive(Debug, Clone, PartialEq, Eq)]
291pub struct Slice {
292    pub start: usize, 
293    pub end: Option<isize>, 
294    pub step: usize
295}
296
297impl Slice {
298    pub fn new(start: usize, end: Option<isize>, step: usize) -> Self {
299        Self { start, end, step }
300    }
301
302    pub fn len(&self) -> usize {
303        self.clone().count()
304    }
305}
306
307impl Iterator for Slice {
308    type Item = usize;
309    fn next(&mut self) -> Option<Self::Item> {
310        match self.end {
311            Some(end) if end < 0 => {
312                let value = self.start;
313                self.start += self.step;
314                Some(value)
315            }
316            Some(end) => {
317                if self.start < end as usize {
318                    let value = self.start;
319                    self.start += self.step;
320                    Some(value)
321                } else {
322                    None
323                }
324            }
325            None => {
326                let value = self.start;
327                self.start += self.step;
328                Some(value)
329            }
330        }
331    }
332}
333
334impl Display for Slice {
335    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
336        let step_part = match self.step {
337            1 => format!(""),
338            _ => format!(":{}", self.step),
339        };
340        match self.end {
341            Some(end) => write!(f, "{}:{}{}", self.start, end, step_part),
342            None => write!(f, "{}:{}", self.start, step_part),
343        }
344    }
345}
346
347impl From<usize> for Indexer {
348    fn from(index: usize) -> Self {
349        Indexer::Select(index)
350    }
351}
352
353impl From<D> for Indexer {
354    fn from(index: D) -> Self {
355        Indexer::SelectD(index)
356    }
357}
358
359impl From<Slice> for Indexer {
360    fn from(value: Slice) -> Self {
361        Indexer::Slice(value)
362    }
363}
364
365impl From<&Tensor<bool>> for Indexer {
366    fn from(value: &Tensor<bool>) -> Self {
367        Indexer::Boolean(value.clone())
368    }
369}
370
371impl From<Tensor<bool>> for Indexer {
372    fn from(value: Tensor<bool>) -> Self {
373        Indexer::Boolean(value)
374    }
375}
376
377impl From<std::ops::Range<usize>> for Indexer {
378    fn from(value: std::ops::Range<usize>) -> Self {
379        let range = Slice::new(value.start, Some(value.end as isize), 1);
380        range.into()
381    }
382}
383
384impl From<std::ops::RangeFrom<usize>> for Indexer {
385    fn from(value: std::ops::RangeFrom<usize>) -> Self {
386        let range = Slice::new(value.start, None, 1);
387        range.into()
388    }
389}
390
391impl From<std::ops::RangeTo<usize>> for Indexer {
392    fn from(value: std::ops::RangeTo<usize>) -> Self {
393        let range = Slice::new(0, Some(value.end as isize), 1);
394        range.into()
395    }
396}
397
398
399impl From<std::ops::RangeFull> for Indexer {
400    fn from(_: std::ops::RangeFull) -> Self {
401        let range = Slice::new(0, None, 1);
402        range.into()
403    }
404}
405
406pub trait IndexOp<T, D: WithDType> {
407    fn index(&self, index: T) -> Result<Tensor<D>>;
408}
409
410impl<I: Into<Indexer>, D: WithDType> IndexOp<I, D> for Tensor<D> {
411    fn index(&self, index: I) -> Result<Tensor<D>> {
412        self.indexes(&[index.into()])
413    }
414}
415
416impl<I: Into<Indexer>, D: WithDType> IndexOp<(I,), D> for Tensor<D> {
417    fn index(&self, (index,): (I,)) -> Result<Tensor<D>> {
418        self.indexes(&[index.into()])
419    }
420}
421
422macro_rules! index_op_tuple {
423    ($($t:ident),+) => {
424        #[allow(non_snake_case)]
425        impl<$($t),*, D: WithDType> IndexOp<($($t,)*), D> for Tensor<D>
426        where
427            $($t: Into<Indexer>,)*
428        {
429            fn index(&self, ($($t,)*): ($($t,)*)) -> Result<Tensor<D>> {
430                self.indexes(&[$($t.into(),)*])
431            }
432        }
433    };
434}
435
436index_op_tuple!(I1, I2);
437index_op_tuple!(I1, I2, I3);
438index_op_tuple!(I1, I2, I3, I4);
439index_op_tuple!(I1, I2, I3, I4, I5);
440
441impl<I: Into<Indexer>, D: WithDType> IndexOp<Vec<I>, D> for Tensor<D> {
442    fn index(&self, index: Vec<I>) -> Result<Tensor<D>> {
443        let indexs = index.into_iter().map(|i| i.into()).collect::<Vec<Indexer>>();
444        self.indexes(&indexs)
445    }
446}
447
448#[macro_export]
449macro_rules! s {
450    // s!(start:end)
451    ($start:tt : $end:expr) => {
452        Slice::new($start as usize, Some($end as isize), 1)
453    };
454    // s!(start:end:step)
455    ($start:tt : $end:tt : $step:expr) => {
456        Slice::new($start as usize, Some($end as isize), $step as usize)
457    };
458    // s!(start:)
459    ($start:tt :) => {
460        Slice::new($start as usize, None, 1)
461    };
462    // s!(start::step)
463    ($start:tt :: $step:expr) => {
464        Slice::new($start as usize, None, $step as usize)
465    };
466    // s!(:$end)
467    (: $end:tt) => {
468        Slice::new(0, Some($end as isize), 1)
469    };
470    // s!(:$end:$step)
471    (: $end:tt : $step:expr) => {
472        Slice::new(0, Some($end as isize), $step as usize)
473    };
474    // s!(::$step)
475    (:: $step:expr) => {
476        Slice::new(0, None, $step as usize)
477    };
478    // s!(:)
479    (:) => {
480        Slice::new(0, None, 1)
481    };
482}
483
484#[cfg(test)]
485#[allow(unused)]
486mod test {
487    use crate::DType;
488    use super::*;
489
490    #[test]
491    fn test_index_select_basic() {
492        // [[ 0,  1,  2,  3],
493        //  [ 4,  5,  6,  7],
494        //  [ 8,  9, 10, 11]]
495        let arr = Tensor::arange(0, 12).unwrap().reshape((3, 4)).unwrap();
496
497        let indices = Tensor::new(&[0, 2]).unwrap();
498        let selected = arr.index_select(indices, 0).unwrap();
499        
500        assert_eq!(selected.shape().dims(), &[2, 4]);
501        let data = selected.to_vec().unwrap();
502        assert_eq!(data, vec![0, 1, 2, 3, 8, 9, 10, 11]);
503
504        let indices_col = Tensor::new(&[1]).unwrap();
505        let selected_col = arr.index_select(indices_col, 1).unwrap();
506
507        assert_eq!(selected_col.shape().dims(), &[3, 1]);
508        let data_col = selected_col.to_vec().unwrap();
509        assert_eq!(data_col, vec![1, 5, 9]);
510    }
511
512    #[test]
513    fn test_boolean_index() {
514        // [[ 0,  1,  2,  3],
515        //  [ 4,  5,  6,  7],
516        //  [ 8,  9, 10, 11]]
517        let arr = Tensor::arange(0, 12).unwrap().reshape((3, 4)).unwrap();
518        let index = Tensor::new([true, false, true]).unwrap();
519
520        let selected = arr.index(index).unwrap();
521        println!("{}", selected);
522    }
523
524    #[test]
525    fn test_index_select_1d() {
526        let scores = Tensor::<f64>::arange(0.0, 100.0).unwrap();
527        let indices = Tensor::new(&[25, 34, 12, 90]).unwrap();
528
529        let selected = scores.index_select(indices, 0).unwrap();
530        println!("{}", selected);
531    }
532
533    #[test]
534    fn test_index_select_duplicates_and_reorder() {
535        let arr = Tensor::arange(0, 5).unwrap(); // [0, 1, 2, 3, 4]
536
537        let indices = Tensor::new(&[4, 0, 0, 1]).unwrap();
538        let selected = arr.index_select(indices, 0).unwrap();
539
540        assert_eq!(selected.shape().dims(), &[4]);
541        let data = selected.to_vec().unwrap();
542        assert_eq!(data, vec![4, 0, 0, 1]);
543    }
544
545    #[test]
546    fn test_index_select_out_of_bounds() {
547        let arr = Tensor::arange(0, 10).unwrap();
548        let indices = Tensor::new(&[0, 10]).unwrap(); 
549        
550        let result = arr.index_select(indices, 0);
551        assert!(result.is_err());
552    }
553
554    #[test]
555    fn test_index_add_basic() {
556        let dst = Tensor::<i32>::zeros((3, 3)).unwrap();
557        let src = Tensor::<i32>::ones((2, 3)).unwrap();
558        let indices = Tensor::new(&[0, 2]).unwrap();
559        
560        let result = dst.index_add(indices, &src, 0).unwrap();
561        
562        // [[1, 1, 1],
563        //  [0, 0, 0],
564        //  [1, 1, 1]]
565        let data = result.to_vec().unwrap();
566        assert_eq!(data, vec![
567            1, 1, 1, 
568            0, 0, 0, 
569            1, 1, 1
570        ]);
571    }
572
573    #[test]
574    fn test_index_add_accumulate() {
575        let dst = Tensor::<i32>::zeros((5,)).unwrap(); // [0, 0, 0, 0, 0]        
576        let src = Tensor::new(&[10, 20, 30]).unwrap();
577        let indices = Tensor::new(&[1, 1, 3]).unwrap();
578        
579        let result = dst.index_add(indices, &src, 0).unwrap();
580        
581        // dst[0] = 0
582        // dst[1] = 0 + 10 + 20 = 30
583        // dst[2] = 0
584        // dst[3] = 0 + 30 = 30
585        // dst[4] = 0
586        let data = result.to_vec().unwrap();
587        assert_eq!(data, vec![0, 30, 0, 30, 0]);
588    }
589
590    #[test]
591    fn test_index_add_dim_mismatch() {
592        let dst = Tensor::<i32>::zeros((3, 3)).unwrap();        
593        let src = Tensor::<i32>::ones((2, 3)).unwrap();
594        let indices = Tensor::new(&[0, 1, 2]).unwrap();
595        
596        let result = dst.index_add(indices, &src, 0);
597        assert!(result.is_err());
598    }
599
600    #[test]
601    fn test_index_add_inner_dim() {
602        // Shape: 2x3
603        // [[0, 0, 0],
604        //  [0, 0, 0]]
605        let dst = Tensor::<i32>::zeros((2, 3)).unwrap();
606        // Source: 2x1
607        let src = Tensor::new(&[
608            [5],
609            [5]
610        ]).unwrap(); // Shape 2x1
611        let indices = Tensor::new(&[1]).unwrap();        
612        let result = dst.index_add(indices, &src, 1).unwrap();
613        
614        // 预期:
615        // [[0, 5, 0],
616        //  [0, 5, 0]]
617        let data = result.to_vec().unwrap();
618        assert_eq!(data, vec![0, 5, 0, 0, 5, 0]);
619    }
620
621
622    #[test]
623    fn test_gather_dim_1() {
624        // Source: 2x2
625        // [[1, 2],
626        //  [3, 4]]
627        let src = Tensor::new(&[
628            [1, 2],
629            [3, 4]
630        ]).unwrap();
631        
632        // 我们想从第 1 维 (列) 取值。
633        // indexes 形状必须与输出一致。
634        // [[0, 0],  -> 取 src[0,0], src[0,0]
635        //  [1, 0]]  -> 取 src[1,1], src[1,0]
636        let indices = Tensor::new(&[
637            [0, 0],
638            [1, 0]
639        ]).unwrap();
640
641        let result = src.gather(&indices, 1).unwrap();
642
643        // 预期结果:
644        // [[1, 1],
645        //  [4, 3]]
646        let data = result.to_vec().unwrap();
647        assert_eq!(data, vec![1, 1, 4, 3]);
648    }
649
650    #[test]
651    fn test_gather_dim_0() {
652        // Source: 3x2
653        // [[10, 20],
654        //  [30, 40],
655        //  [50, 60]]
656        let src = Tensor::new(&[
657            [10, 20],
658            [30, 40],
659            [50, 60]
660        ]).unwrap();
661
662        // 沿维度 0 (行) gather
663        // [[1, 2], -> src[1,0]=30, src[2,1]=60
664        //  [0, 1]] -> src[0,0]=10, src[1,1]=40
665        let indices = Tensor::new(&[
666            [1, 2],
667            [0, 1]
668        ]).unwrap();
669
670        let result = src.gather(&indices, 0).unwrap();
671
672        // 预期结果:
673        // [[30, 60],
674        //  [10, 40]]
675        let data = result.to_vec().unwrap();
676        assert_eq!(data, vec![30, 60, 10, 40]);
677    }
678
679    #[test]
680    fn test_gather_3d() {
681        // Shape: 2x2x2
682        // Block 0: [[0, 1], [2, 3]]
683        // Block 1: [[4, 5], [6, 7]]
684        let src = Tensor::new(&[0, 1, 2, 3, 4, 5, 6, 7]).unwrap().reshape((2, 2, 2)).unwrap();
685
686        // Gather dim 1 (中间维度)
687        // 我们保持 dim 0 和 dim 2 不变，只改变 dim 1 的索引
688        // index shape: 2x1x2 (我们在 dim 1 上做降维提取)
689        let indices = Tensor::<u32>::zeros((2, 1, 2)).unwrap(); // 全是 0
690
691        // 逻辑:
692        // out[0, 0, 0] = src[0, idx[0,0,0], 0] = src[0,0,0] = 0
693        // out[0, 0, 1] = src[0, idx[0,0,1], 1] = src[0,0,1] = 1
694        // out[1, 0, 0] = src[1, idx[1,0,0], 0] = src[1,0,0] = 4
695        // out[1, 0, 1] = src[1, idx[1,0,1], 1] = src[1,0,1] = 5
696        
697        let result = src.gather(&indices, 1).unwrap();
698        
699        assert_eq!(result.dims(), &[2, 1, 2]);
700        let data = result.to_vec().unwrap();
701        assert_eq!(data, vec![0, 1, 4, 5]);
702    }
703
704
705    #[test]
706    fn test_scatter_add_1d_accumulate() {
707        // 类似于直方图统计
708        let dst = Tensor::<i32>::zeros((5,)).unwrap();
709        
710        // Source 和 Indices 形状一致
711        let src = Tensor::new(&[1, 1, 1, 1]).unwrap();
712        let indices = Tensor::new(&[0, 2, 0, 4]).unwrap();
713
714        // scatter_add(dim=0)
715        // dst[0] += 1
716        // dst[2] += 1
717        // dst[0] += 1 (累加)
718        // dst[4] += 1
719        let result = dst.scatter_add(indices, &src, 0).unwrap();
720
721        // 预期: [2, 0, 1, 0, 1]
722        let data = result.to_vec().unwrap();
723        assert_eq!(data, vec![2, 0, 1, 0, 1]);
724    }
725
726    #[test]
727    fn test_scatter_add_2d_dim1() {
728        // Dst: 2x3 zeros
729        let dst = Tensor::<i32>::zeros((2, 3)).unwrap();
730        
731        // Src: 2x2
732        // [[10, 20],
733        //  [30, 40]]
734        let src = Tensor::new(&[
735            [10, 20],
736            [30, 40]
737        ]).unwrap();
738
739        // Indices: 2x2
740        // [[0, 2], -> 将 10 加到 dst[0,0], 20 加到 dst[0,2]
741        //  [1, 1]] -> 将 30 加到 dst[1,1], 40 加到 dst[1,1] (累加)
742        let indices = Tensor::new(&[
743            [0, 2],
744            [1, 1]
745        ]).unwrap();
746
747        let result = dst.scatter_add(indices, &src, 1).unwrap();
748
749        // 预期:
750        // Row 0: [10, 0, 20]
751        // Row 1: [0, 70, 0] (30+40=70)
752        let data = result.to_vec().unwrap();
753        assert_eq!(data, vec![10, 0, 20, 0, 70, 0]);
754    }
755
756    #[test]
757    fn test_scatter_add_3d() {
758        // Dst: 2x2x2 zeros
759        let dst = Tensor::<i32>::zeros((2, 2, 2)).unwrap();
760        
761        // Src: 2x1x2 (dim 1 is smaller)
762        let src = Tensor::ones((2, 1, 2)).unwrap(); // All ones
763        
764        // Indices: 2x1x2
765        // Block 0: [[1, 0]] -> dst[0, 1, 0] += 1, dst[0, 0, 1] += 1
766        // Block 1: [[0, 0]] -> dst[1, 0, 0] += 1, dst[1, 0, 1] += 1
767        let indices = Tensor::new(&[1, 0, 0, 0]).unwrap().reshape((2, 1, 2)).unwrap();
768
769        let result = dst.scatter_add(indices, &src, 1).unwrap();
770
771        // Check specifics:
772        // dst[0, 1, 0] should be 1
773        // dst[0, 0, 1] should be 1
774        // dst[1, 0, 0] should be 1
775        // dst[1, 0, 1] should be 1
776        // Others 0
777        let res_vec = result.to_vec().unwrap();
778        // Flattened index check:
779        // 2x2x2 -> stride [4, 2, 1]
780        // [0,1,0] -> 2 -> val 1
781        // [0,0,1] -> 1 -> val 1
782        // [1,0,0] -> 4 -> val 1
783        // [1,0,1] -> 5 -> val 1
784        
785        // Vec: [0, 1, 1, 0, 1, 1, 0, 0]
786        assert_eq!(res_vec, vec![0, 1, 1, 0, 1, 1, 0, 0]);
787    }
788
789    #[test]
790    fn test_scatter_add_shape_mismatch() {
791        let dst = Tensor::<i32>::zeros((2, 2)).unwrap();
792        let src = Tensor::<i32>::ones((2, 2)).unwrap();
793        // Indices shape mismatch with Src
794        let indices = Tensor::new(&[0]).unwrap(); 
795
796        let result = dst.scatter_add(indices, &src, 0);
797        assert!(result.is_err());
798    }
799
800    #[test]
801    fn test_index_scalar_dim_reduction() {
802        let arr = Tensor::arange(0, 125).unwrap().reshape((5, 5, 5)).unwrap();
803        let sub = arr.index(1).unwrap();
804        assert_eq!(sub.shape().dims(), &[5, 5]);
805
806        let sub = arr.index((2, 3)).unwrap();
807        assert_eq!(sub.shape().dims(), &[5]); 
808    }
809
810    #[test]
811    fn test_index_range_basic() {
812        let arr = Tensor::arange(0, 125).unwrap().reshape((5, 5, 5)).unwrap();
813
814        let sub = arr.index(s!(1:3)).unwrap();
815        assert_eq!(sub.shape().dims(), &[2, 5, 5]);
816
817        let sub = arr.index((s!(1:3), s!(3:4), 1)).unwrap();
818        assert_eq!(sub.shape().dims(), &[2, 1]);
819    }
820
821    #[test]
822    fn test_index_full_and_mixed() {
823        let arr = Tensor::<i32>::zeros((5, 5, 5)).unwrap();
824
825        let sub = arr.index((s!(1:3), .., 1..2)).unwrap();
826        assert_eq!(sub.shape().dims(), &[2, 5, 1]);
827
828        let sub = arr.index((2, .., s!(0:2))).unwrap();
829        assert_eq!(sub.shape().dims(), &[5, 2]);
830
831        let sub = arr.index((s!(0:2), s!(2:5), s!(1:3))).unwrap();
832        assert_eq!(sub.shape().dims(), &[2, 3, 2]);
833    }
834
835    #[test]
836    fn test_index_out_of_bounds() {
837        let arr = Tensor::<i32>::zeros((5, 5, 5)).unwrap();
838        let result = arr.index(10);
839        assert!(result.is_err());
840
841        let result = arr.index(s!(3:10));
842        assert!(result.is_err());
843    }
844    
845    #[test]
846    fn test_index_scalar_and_values() {
847        let arr = Tensor::arange(0, 125).unwrap().reshape((5, 5, 5)).unwrap();
848
849        let sub = arr.index(1).unwrap();
850        let expected = Tensor::arange(25, 50).unwrap().reshape((5, 5)).unwrap();
851        assert!(sub.allclose(&expected, 0.0, 0.0).unwrap());
852    }
853
854    #[test]
855    fn test_index_range_values() {
856        let arr = Tensor::arange(0, 125).unwrap().reshape((5, 5, 5)).unwrap();
857
858        let sub = arr.index(s!(1:3)).unwrap();
859        let expected = Tensor::arange(25, 75).unwrap().reshape((2, 5, 5)).unwrap();
860        assert!(sub.allclose(&expected, 0.0, 0.0).unwrap());
861    }
862
863    #[test]
864    fn test_index_mixed_values() {
865        let arr = Tensor::arange(0, 125).unwrap().reshape((5, 5, 5)).unwrap();
866        let sub = arr.index((2, 3)).unwrap();
867        let expected = Tensor::arange(65, 70).unwrap();
868        assert!(sub.allclose(&expected, 0.0, 0.0).unwrap());
869
870        let sub = arr.index((s!(1:3), s!(3:5), 2)).unwrap();
871        let mut vals = Vec::new();
872        for i in 1..3 {
873            for j in 3..5 {
874                vals.push(i * 25 + j * 5 + 2);
875            }
876        }
877        let expected = Tensor::from_vec(vals, (2, 2)).unwrap();
878        assert!(sub.allclose(&expected, 0.0, 0.0).unwrap());
879    }
880
881    #[test]
882    fn test_index_with_full_dim() {
883        let arr = Tensor::arange(0, 125).unwrap().reshape((5, 5, 5)).unwrap();
884        let sub = arr.index((s!(1:3), .., 1..2)).unwrap();
885
886        let expected = arr.index((s!(1:3), s!(0:5), s!(1:2))).unwrap();
887        assert!(sub.allclose(&expected, 0.0, 0.0).unwrap());
888    }
889
890    #[test]
891    fn test_index_dim() {
892        let arr = Tensor::<f32>::zeros((3, 3)).unwrap();
893        // last col
894        let sub = arr.index((.., D::Minus1)).unwrap();
895        sub.copy_(5.2).unwrap();
896        println!("{}", arr);
897    }
898
899    #[test]
900    fn test_macro() {
901        let t = (0..12usize);
902        let t = (2usize..);
903        assert_eq!(s!(1:10), Slice {start:1, end: Some(10), step:1});
904
905        assert!(
906            s!(1:20).zip((1..20))
907                .all(|(a, b)| a == b)
908        );
909    
910        assert!(
911            s!(1:13:3).zip((1..13).step_by(3))
912                .all(|(a, b)| a == b)
913        );
914    
915        assert!(
916            s!(1:).zip((1..).take(100))
917                .all(|(a, b)| a == b)
918        );
919
920        assert!(
921            s!(1::2).zip((1..).step_by(2).take(100))
922                .all(|(a, b)| a == b)
923        );
924
925        assert!(
926            s!(:20).zip((0..20usize))
927                .all(|(a, b)| a == b)
928        );
929
930        assert!(
931            s!(:20:5).zip((0..20usize).step_by(5))
932                .all(|(a, b)| a == b)
933        );
934
935        assert!(
936            s!(::2).zip((0..).step_by(2).take(100))
937                .all(|(a, b)| a == b)
938        );
939
940        assert!(
941            s!(:).zip((0..).take(100))
942                .all(|(a, b)| a == b)
943        );
944    }
945}
946
lumen_core/ops/indexer.rs

lumen_core/ops/
indexer.rs
