lumen_core/ops/
reduce.rs

1use std::cmp::Ordering;
2use crate::{AutogradMetaT, Dim, NumDType, Result, Storage, Tensor, WithDType};
3use paste::paste;
4
5/*
6
7        let (storage, shape) = self.compute_reduce_op(reduce_dim, op, keepdim)?;
8        Ok(Tensor::<R>::from_storage(storage, shape, meta))
9*/
10macro_rules! reduce_impl {
11    ($fn_name:ident, $reduce:ident, $op:ident) => {
12        paste! {
13            pub fn $fn_name<D: Dim>(&self, axis: D) -> Result<Self> {
14                let (storage, dims) = self.compute_reduce_op(axis, $reduce, false)?;
15                let meta = T::AutogradMeta::on_reduce_op(self, &dims, crate::ReduceOp::$op);
16                Ok(Self::from_storage(storage, dims, meta))
17            }
18        
19            pub fn [< $fn_name _keepdim >]<D: Dim>(&self, axis: D) -> Result<Self> {
20                let (storage, dims) = self.compute_reduce_op(axis, $reduce, true)?;
21                let meta = T::AutogradMeta::on_reduce_op(self, &dims, crate::ReduceOp::$op);
22                Ok(Self::from_storage(storage, dims, meta))
23            }
24
25            pub fn [< $fn_name _all >](&self) -> Result<Self> {
26                self.flatten_all()?.$fn_name(0)
27            }
28        }
29    };
30}
31
32impl<T: NumDType> Tensor<T> {
33    reduce_impl!(sum, ReduceSum, Sum);
34    reduce_impl!(min, ReduceMin, Min);
35    reduce_impl!(max, ReduceMax, Max);
36    reduce_impl!(mean, ReduceMean, Mean);
37
38    pub fn var_keepdim<D: Dim>(&self, axis: D) -> Result<Self> {
39        let mean = self.mean_keepdim(axis)?; // (..., 1, ...)
40        let delta = self.broadcast_sub(&mean)?; // (..., n, ...)
41        let delta_pow = &delta * &delta; // (..., n, ...)
42
43        delta_pow.mean_keepdim(axis)
44    }
45
46    pub fn var_unbiased_keepdim<D: Dim>(&self, axis: D) -> Result<Self> {
47        let n = T::from_usize(self.dim(axis)?);
48        let biased_var = self.var_keepdim(axis)?;
49        
50        // cor scale: N / (N - 1)
51        let correction = n / (n - T::one());
52        Ok(correction * biased_var)
53    }
54    
55    pub fn var<D: Dim>(&self, axis: D) -> Result<Self> {
56        let v = self.var_keepdim(axis)?;
57        let v = v.squeeze(axis)?;
58        Ok(v)
59    }
60
61    pub fn var_unbiased<D: Dim>(&self, axis: D) -> Result<Self> {
62        let v = self.var_unbiased_keepdim(axis)?;
63        let v = v.squeeze(axis)?;
64        Ok(v)
65    }
66
67    pub fn var_all(&self) -> Result<Self> {
68        self.flatten_all()?.var(0)
69    }
70
71    pub fn var_unbiased_all(&self) -> Result<Self> {
72        self.flatten_all()?.var_unbiased(0)
73    }
74
75    pub fn argmin_keepdim<D: Dim>(&self, axis: D) -> Result<Tensor<u32>> {
76        let (storage, dims) = self.compute_reduce_op(axis, ReduceArgMin, true)?;
77        Ok(Tensor::from_storage(storage, dims, Default::default()))
78    }
79
80    pub fn argmin<D: Dim>(&self, axis: D) -> Result<Tensor<u32>> {
81        let (storage, dims) = self.compute_reduce_op(axis, ReduceArgMin, false)?;
82        Ok(Tensor::from_storage(storage, dims, Default::default()))
83    }
84
85    pub fn argmax_keepdim<D: Dim>(&self, axis: D) -> Result<Tensor<u32>> {
86        let (storage, dims) = self.compute_reduce_op(axis, ReduceArgMax, true)?;
87        Ok(Tensor::from_storage(storage, dims, Default::default()))
88    }
89
90    pub fn argmax<D: Dim>(&self, axis: D) -> Result<Tensor<u32>> {
91        let (storage, dims) = self.compute_reduce_op(axis, ReduceArgMax, false)?;
92        Ok(Tensor::from_storage(storage, dims, Default::default()))
93    }
94}
95
96impl Tensor<bool> {
97    pub fn all_all(&self) -> crate::Result<bool> {
98        self.iter().map(|mut i| i.all(|a| a))
99    }
100
101    pub fn any_all(&self) -> crate::Result<bool> {
102        self.iter().map(|mut i| i.any(|a| a))
103    } 
104
105    pub fn all_keepdim<D: Dim>(&self, axis: D) -> Result<Tensor<bool>> {
106        self.reduce_op(axis, ReduceAll, true, Default::default())
107    }
108
109    pub fn any_keepdim<D: Dim>(&self, axis: D) -> Result<Tensor<bool>> {
110        self.reduce_op(axis, ReduceAny, true, Default::default())
111    }
112
113    pub fn all<D: Dim>(&self, axis: D) -> Result<Tensor<bool>> {
114        self.reduce_op(axis, ReduceAll, false, Default::default())
115    }
116
117    pub fn any<D: Dim>(&self, axis: D) -> Result<Tensor<bool>> {
118        self.reduce_op(axis, ReduceAny, false, Default::default())
119    }
120}
121
122impl<T: WithDType> Tensor<T> {
123    pub(crate) fn reduce_op<'a, R, Op, D>(&'a self, reduce_dim: D, op: Op, keepdim: bool, meta: R::AutogradMeta) -> Result<Tensor<R>> 
124    where 
125        R: WithDType,
126        D: Dim,
127        Op: ReduceOp<T, Output = R>,
128    {
129        let (storage, shape) = self.compute_reduce_op(reduce_dim, op, keepdim)?;
130        Ok(Tensor::<R>::from_storage(storage, shape, meta))
131    }
132
133    pub(crate) fn compute_reduce_op<'a, R, Op, D>(&'a self, reduce_dim: D, _: Op, keepdim: bool) -> Result<(Storage<R>, Vec<usize>)> 
134    where 
135        R: WithDType,
136        D: Dim,
137        Op: ReduceOp<T, Output = R>,
138    {
139        let reduce_dim = reduce_dim.to_index(self.shape(), Op::NAME)?;
140        assert!(reduce_dim < self.layout().dims().len());
141        let reduce_dim_stride = self.layout().stride()[reduce_dim];
142        let reduce_dim_size = self.layout().dims()[reduce_dim];
143   
144        let dst_storage = self.storage_read()?;
145        let dst_data = dst_storage.data();
146
147        let dst = if self.is_contiguous() && reduce_dim_stride == 1 {
148            let dst_data = &dst_data[self.layout().start_offset..];
149            (0..self.layout().element_count() / reduce_dim_size).into_iter()
150                .map(|i| {
151                    let chunk = &dst_data[i*reduce_dim_size..i*reduce_dim_size + reduce_dim_size];
152                    let v = Op::op(chunk.iter().copied());
153                    v
154                }).collect()
155        } else {
156            let dst_len = self.layout().element_count() / reduce_dim_size;
157            let mut dst: Vec<R> = Vec::with_capacity(dst_len);
158            unsafe { dst.set_len(dst_len) };
159            // let dst_to_set = dst.spare_capacity_mut();
160
161            let layout = self.layout().narrow(reduce_dim, 0, 1)?;
162            if reduce_dim_stride == 1 {
163                for (dst_index, src_index) in layout.storage_indices().enumerate() {
164                    let end_index = src_index + reduce_dim_size;
165                    let chunk = &dst_data[src_index..end_index];
166                    let v = Op::op(chunk.iter().copied());
167                    dst[dst_index] = v;
168                }
169            } else {
170                for (dst_index, src_index) in layout.storage_indices().enumerate() {
171                    let arr: DimArray<'_, T> = DimArray {
172                        src: &dst_data[src_index..],
173                        size: reduce_dim_size,
174                        stride: reduce_dim_stride
175                    };
176                    let iter: DimArrayIter<'_, T> = arr.into_iter();
177                    let v = Op::op(iter);
178                    dst[dst_index] = v;
179                }
180            }
181
182            dst
183        };
184    
185        let storage = Storage::new(dst);
186        let mut shape = self.dims().to_vec();
187
188        if keepdim {
189            shape[reduce_dim] = 1;
190        } else {
191            shape.remove(reduce_dim);
192        }
193
194        Ok((storage, shape))
195    }
196}
197
198impl<T: NumDType> Tensor<T> {
199    pub fn topk<D: Dim>(&self, k: usize, dim: D) -> crate::Result<(Tensor<T>, Tensor<u32>)> {
200        let reduce_dim = dim.to_index(self.shape(), "topk")?;
201        assert!(reduce_dim < self.layout().dims().len());
202        let reduce_dim_stride = self.layout().stride()[reduce_dim];
203        let reduce_dim_size = self.layout().dims()[reduce_dim];
204        if k > reduce_dim_size {
205            return Err(crate::Error::IndexOutOfRange { max_size: reduce_dim_size, index: k, op: "topk"});
206        }
207
208        // 输出目标：值和数组的总大小
209        let dst_len = self.layout().element_count() / reduce_dim_size * k;
210        // 预先开辟足够的空间
211        let mut indices: Vec<u32> = Vec::with_capacity(dst_len);
212        let mut values: Vec<T> = Vec::with_capacity(dst_len);
213        let mut indexed: Vec<(u32, T)> = Vec::with_capacity(reduce_dim_size);
214
215        // 先去掉 dim 维度，遍历其他所有的可能的排列
216        let storage = self.storage_read()?;
217        let storage_slice = storage.data();
218        let layout = self.layout().narrow(reduce_dim, 0, 1)?;
219        for src_index in layout.storage_indices() {
220            // 取出这个维度的一组数据
221            let arr: DimArray<'_, T> = DimArray {
222                src: &storage_slice[src_index..],
223                size: reduce_dim_size,
224                stride: reduce_dim_stride
225            };
226            // 索引 + 值进行排序
227            indexed.clear();
228            indexed.extend(arr.into_iter().enumerate().map(|(i, v)| (i as u32, v)));
229
230            assert!(indexed.len() >= k);
231            if k < indexed.len() {
232                // 先在 O(N) 时间内找出前 k 大的划分
233                indexed.select_nth_unstable_by(k - 1, |a, b| {
234                    b.1.partial_cmp(&a.1).unwrap_or(Ordering::Equal)
235                });
236                // 然后只对前 k 个进行排序 O(k log k)
237                indexed[..k].sort_unstable_by(|a, b| {
238                    b.1.partial_cmp(&a.1).unwrap_or(Ordering::Equal)
239                });
240            } else {
241                indexed.sort_unstable_by(|a, b| {
242                    b.1.partial_cmp(&a.1).unwrap_or(Ordering::Equal)
243                });
244            }
245
246            let topk_slice = &indexed[..k];
247
248            values.extend(topk_slice.iter().map(|(_, v)| *v));
249            indices.extend(topk_slice.iter().map(|(i, _)| *i));
250        }
251
252        assert_eq!(indices.len(), dst_len);
253        assert_eq!(values.len(), dst_len);
254
255        let mut physical_shape = self.dims().to_vec();
256        physical_shape.remove(reduce_dim);
257        physical_shape.push(k);
258        
259        let temp_values = Tensor::from_vec(values, physical_shape.clone())?;
260        let temp_indices = Tensor::from_vec(indices, physical_shape.clone())?;
261
262        // 如果 reduce_dim 本来就在最后，不需要转置
263        if reduce_dim == self.dims().len() - 1 {
264            return Ok((temp_values, temp_indices));
265        }
266
267        // 把最后一个维度移动回 reduce_dim 的位置
268        let ndims = physical_shape.len();
269        let mut permute_axes: Vec<usize> = (0..ndims - 1).collect(); // [0, 1, 2, ..., N-2]
270        permute_axes.insert(reduce_dim, ndims - 1);
271
272        let final_values = temp_values.permute(permute_axes.as_ref())?.contiguous()?;
273        let final_indices = temp_indices.permute(permute_axes)?.contiguous()?;
274
275        Ok((final_values, final_indices))
276    }
277}
278
279pub trait ReduceOp<D: WithDType> {
280    type Output: WithDType;
281    const NAME: &'static str;
282    fn op<I>(iter: I) -> Self::Output
283    where
284        I: Iterator<Item = D> + ExactSizeIterator + Clone;
285}
286
287pub struct ReduceAll;
288impl ReduceOp<bool> for ReduceAll {
289    type Output = bool;
290    const NAME: &'static str = "all";
291    fn op<I>(mut iter: I) -> Self::Output 
292    where I: Iterator<Item = bool> + ExactSizeIterator + Clone {
293        iter.all(|b| b)
294    }
295}
296
297pub struct ReduceAny;
298impl ReduceOp<bool> for ReduceAny {
299    type Output = bool;
300    const NAME: &'static str = "any";
301    fn op<I>(mut iter: I) -> Self::Output 
302    where I: Iterator<Item = bool> + ExactSizeIterator + Clone {
303        iter.any(|b| b)
304    }
305}
306
307pub struct ReduceSum;
308impl<D: NumDType> ReduceOp<D> for ReduceSum {
309    const NAME: &'static str = "sum";
310    type Output = D;
311    fn op<I>(iter: I) -> Self::Output 
312    where I: Iterator<Item = D> + ExactSizeIterator + Clone {
313        iter.sum::<D>()
314    }
315} 
316
317pub struct ReduceMean;
318impl<D: NumDType> ReduceOp<D> for ReduceMean {
319    type Output = D;
320    const NAME: &'static str = "mean";
321    fn op<I>(iter: I) -> Self::Output 
322    where I: Iterator<Item = D> + ExactSizeIterator + Clone {
323        let len = iter.len();
324        iter.sum::<D>() / D::from_usize(len)
325    }
326} 
327
328pub struct ReduceVar;
329impl<D: NumDType> ReduceOp<D> for ReduceVar {
330    type Output = D;
331    const NAME: &'static str = "var";
332    fn op<I>(iter: I) -> Self::Output 
333    where I: Iterator<Item = D> + ExactSizeIterator + Clone {
334        let len = iter.len();
335        if len == 0 { return D::zero(); }
336
337        // 利用 Clone 给 Mean 使用，本迭代器保留用于第二次遍历
338        let mean = ReduceMean::op(iter.clone());
339        
340        let mut sum_sq_diff = D::zero();
341        for v in iter {
342            let diff = v - mean;
343            sum_sq_diff += diff * diff;
344        }
345
346        sum_sq_diff / D::from_usize(len)
347    }
348} 
349
350pub struct ReduceProduct;
351impl<D: NumDType> ReduceOp<D> for ReduceProduct {
352    type Output = D;
353    const NAME: &'static str = "product";
354    fn op<I>(iter: I) -> Self::Output 
355    where I: Iterator<Item = D> + ExactSizeIterator + Clone {
356        iter.product::<D>()
357    }
358} 
359
360pub struct ReduceMin;
361impl<D: NumDType> ReduceOp<D> for ReduceMin {
362    type Output = D;
363    const NAME: &'static str = "min";
364    fn op<I>(iter: I) -> Self::Output 
365    where I: Iterator<Item = D> + ExactSizeIterator + Clone {
366        iter.reduce(|a, b| D::minimum(a, b)).unwrap()
367    }
368} 
369
370pub struct ReduceMax;
371impl<D: NumDType> ReduceOp<D> for ReduceMax {
372    type Output = D;
373    const NAME: &'static str = "max";
374    fn op<I>(iter: I) -> Self::Output 
375    where I: Iterator<Item = D> + ExactSizeIterator + Clone {
376        iter.reduce(|a, b| D::maximum(a, b)).unwrap()
377    }
378} 
379
380pub struct ReduceArgMin;
381impl<D: NumDType> ReduceOp<D> for ReduceArgMin {
382    type Output = u32;
383    const NAME: &'static str = "arg_min";
384    fn op<I>(iter: I) -> Self::Output 
385    where I: Iterator<Item = D> + ExactSizeIterator + Clone {
386        iter.enumerate()
387            .reduce(|(ia, a), (ib, b)| {
388                if a.partial_cmp(&b) == Some(std::cmp::Ordering::Less) {
389                    (ia, a)
390                } else {
391                    (ib, b)
392                }
393            }).unwrap().0 as u32
394    }
395} 
396
397pub struct ReduceArgMax;
398impl<D: NumDType> ReduceOp<D> for ReduceArgMax {
399    type Output = u32;
400    const NAME: &'static str = "arg_max";
401    fn op<I>(iter: I) -> Self::Output 
402    where I: Iterator<Item = D> + ExactSizeIterator + Clone {
403        iter.enumerate()
404            .reduce(|(ia, a), (ib, b)| {
405                if a.partial_cmp(&b) == Some(std::cmp::Ordering::Greater) {
406                    (ia, a)
407                } else {
408                    (ib, b)
409                }
410            }).unwrap().0 as u32
411    }
412}
413
414#[derive(Clone)]
415pub struct DimArray<'a, T> {
416    pub(crate) src: &'a [T],
417    pub(crate) size: usize,
418    pub(crate) stride: usize
419}
420
421impl<'a, T: WithDType> DimArray<'a, T> {
422    pub fn get(&self, index: usize) -> T {
423        self.src[index * self.stride]
424    }
425
426    #[allow(unused)]
427    pub fn to_vec(&self) -> Vec<T> {
428        let mut v = vec![];
429        for i in 0..self.size {
430            v.push(self.get(i));
431        }
432        v
433    }
434}
435
436impl<'a, T: WithDType> IntoIterator for DimArray<'a, T> {
437    type IntoIter = DimArrayIter<'a, T>;
438    type Item = T;
439    fn into_iter(self) -> Self::IntoIter {
440        DimArrayIter::<'a, T> {
441            array: self,
442            index: 0,
443        }
444    }
445}
446
447#[derive(Clone)]
448pub struct DimArrayIter<'a, T> {
449    array: DimArray<'a, T>,
450    index: usize,
451}
452
453impl<'a, T: WithDType> Iterator for DimArrayIter<'a, T> {
454    type Item = T;
455    fn next(&mut self) -> Option<T> {
456        if self.index >= self.array.size {
457            None
458        } else {
459            let index = self.index;
460            self.index += 1;
461            Some(self.array.get(index))
462        }
463    }
464}
465
466impl<'a, T: WithDType> ExactSizeIterator for DimArrayIter<'a, T> {
467    fn len(&self) -> usize {
468        self.array.size
469    }
470}
471
472#[cfg(test)]
473mod tests {
474    use super::*;
475
476    #[test]
477    fn test_sum_matrix_axis0() -> crate::Result<()> {
478        // [[1, 2, 3],
479        //  [3, 4, 5]]
480        // sum_all(axis=0) -> [4, 6, 8]
481        let arr = Tensor::new(&[[1, 2, 3], [3, 4, 5]]).unwrap();
482        let s = arr.sum(0).unwrap();
483        let expected = Tensor::new(&[4, 6, 8]).unwrap();
484        assert!(s.allclose(&expected, 1e-5, 1e-8).unwrap());
485
486        let arr = arr.transpose_last()?.copy()?.transpose_last()?;
487        let s = arr.sum(0).unwrap();
488        assert!(s.allclose(&expected, 1e-5, 1e-8).unwrap());
489        
490        Ok(())
491    }
492
493    #[test]
494    fn test_sum_matrix_axis1() -> crate::Result<()> {
495        // [[1, 2, 3],
496        //  [3, 4, 5]]
497        // sum_all(axis=1) -> [6, 12]
498        let arr = Tensor::new(&[[1, 2, 3], [3, 4, 5]]).unwrap();
499        let s = arr.sum(1).unwrap();
500        let expected = Tensor::new(&[6, 12]).unwrap();
501        assert!(s.allclose(&expected, 1e-5, 1e-8).unwrap());
502
503        let arr = arr.transpose_last()?.copy()?.transpose_last()?;
504        let s = arr.sum(1).unwrap();
505        assert!(s.allclose(&expected, 1e-5, 1e-8).unwrap());
506
507        Ok(())
508    }
509
510    #[test]
511    fn test_sum_ones_axis() {
512        // ones( (2,3), dtype=I32 )
513        // [[1,1,1],
514        //  [1,1,1]]
515        let arr = Tensor::ones((2, 3)).unwrap();
516        let s0 = arr.sum(0).unwrap(); // -> [2,2,2]
517        let s1 = arr.sum(1).unwrap(); // -> [3,3]
518
519        let expected0 = Tensor::new(&[2, 2, 2]).unwrap();
520        let expected1 = Tensor::new(&[3, 3]).unwrap();
521
522        assert!(s0.allclose(&expected0, 1e-5, 1e-8).unwrap());
523        assert!(s1.allclose(&expected1, 1e-5, 1e-8).unwrap());
524    }
525
526    #[test]
527    fn test_min_matrix_axis0() {
528        // [[1, 2, 3],
529        //  [3, 1, 0]]
530        // min_all(axis=0) -> [1, 1, 0]
531        let arr = Tensor::new(&[[1, 2, 3], [3, 1, 0]]).unwrap();
532        let m = arr.min(0).unwrap();
533        let expected = Tensor::new(&[1, 1, 0]).unwrap();
534        assert!(m.allclose(&expected, 1e-5, 1e-8).unwrap());
535    }
536
537    #[test]
538    fn test_max_matrix_axis1() {
539        // [[1, 2, 3],
540        //  [3, 1, 0]]
541        // max_all(axis=1) -> [3, 3]
542        let arr = Tensor::new(&[[1, 2, 3], [3, 1, 0]]).unwrap();
543        let m = arr.max(1).unwrap();
544        let expected = Tensor::new(&[3, 3]).unwrap();
545        assert!(m.allclose(&expected, 1e-5, 1e-8).unwrap());
546    }
547
548    #[test]
549    fn test_aragmin_matrix_axis0() {
550        // [[1, 2, 3],
551        //  [3, 1, 0]]
552        // min_all(axis=0) -> [1, 1, 0]
553        let arr = Tensor::new(&[[1, 2, 3], [3, 1, 0]]).unwrap();
554        let m = arr.argmin(0).unwrap();
555        let expected = Tensor::new(&[0, 1, 1]).unwrap();
556        assert!(m.allclose(&expected, 1e-5, 1e-8).unwrap());
557    }
558
559    #[test]
560    fn test_sum_all() {
561        // [[1, 2], [3, 4]] -> 1+2+3+4 = 10
562        let arr = Tensor::new(&[[1, 2], [3, 4]]).unwrap();
563        let s = arr.sum_all().unwrap();        
564        let expected = Tensor::new(10).unwrap(); 
565        assert!(s.allclose(&expected, 1e-5, 1e-8).unwrap());
566    }
567
568    #[test]
569    fn test_mean_all() {
570        // [[1.0, 2.0], [3.0, 4.0]] -> Sum=10.0, Count=4 -> Mean=2.5
571        let arr = Tensor::new(&[[1.0, 2.0], [3.0, 4.0]]).unwrap();
572        let m = arr.mean_all().unwrap();
573        let expected = Tensor::new(2.5).unwrap();
574        assert!(m.allclose(&expected, 1e-5, 1e-8).unwrap());
575    }
576
577    #[test]
578    fn test_min_max_all() {
579        // [[10, 2, 5], [8, 1, 9]]
580        // Global Min: 1
581        // Global Max: 10
582        let arr = Tensor::new(&[[10, 2, 5], [8, 1, 9]]).unwrap();
583        
584        let min_val = arr.min_all().unwrap();
585        let max_val = arr.max_all().unwrap();
586        
587        let expected_min = Tensor::new(1).unwrap();
588        let expected_max = Tensor::new(10).unwrap();
589        
590        assert!(min_val.allclose(&expected_min, 1e-5, 1e-8).unwrap());
591        assert!(max_val.allclose(&expected_max, 1e-5, 1e-8).unwrap());
592    }
593    
594    #[test]
595    fn test_argmax_matrix_axis1() {
596        // [[1, 2, 3],
597        //  [3, 1, 0]]
598        // max_all(axis=1) -> [3, 3]
599        let arr = Tensor::new(&[[1, 2, 3], [3, 1, 0]]).unwrap();
600        let m = arr.argmax(1).unwrap();
601        let expected = Tensor::new(&[2, 0]).unwrap();
602        assert!(m.allclose(&expected, 1e-5, 1e-8).unwrap());
603    }
604
605    #[test]
606    fn test_reductions_with_negatives() {
607        // [[-2.0, 0.0, 2.0]]
608        // sum_all = 0.0
609        // mean_all = 0.0
610        // var(axis=1) -> 4.0 (unbiased: (4+0+4)/2)
611        
612        let arr = Tensor::new(&[[-2.0, 0.0, 2.0]]).unwrap();
613        
614        assert!(arr.sum_all().unwrap().allclose(&Tensor::new(0.0).unwrap(), 1e-5, 1e-8).unwrap());
615        assert!(arr.mean_all().unwrap().allclose(&Tensor::new(0.0).unwrap(), 1e-5, 1e-8).unwrap());
616        
617        let expected_var = Tensor::new(2.66666666666666666).unwrap();
618        assert!(arr.var_all().unwrap().allclose(&expected_var, 1e-5, 1e-8).unwrap());
619    }
620
621    #[test]
622    fn test_topk_1d() {
623        // [10, 50, 20, 40, 30]
624        // topk(k=3, axis=0) -> 找最大的 3 个
625        // 期望 Values:  [50, 40, 30]
626        // 期望 Indices: [ 1,  3,  4]
627        let arr = Tensor::new(&[10, 50, 20, 40, 30]).unwrap();
628        let (val, idx) = arr.topk(3, 0).unwrap();
629
630        let expected_val = Tensor::new(&[50, 40, 30]).unwrap();
631        let expected_idx = Tensor::new(&[1u32, 3, 4]).unwrap(); // 注意索引是 u32
632
633        assert!(val.allclose(&expected_val, 1e-5, 1e-8).unwrap());
634        assert!(idx.allclose(&expected_idx, 1e-5, 1e-8).unwrap());
635    }
636
637    #[test]
638    fn test_topk_matrix_axis1() {
639        // 测试按行求 TopK (axis=1，也是连续内存)
640        // [[1, 5, 2],
641        //  [8, 4, 6]]
642        // topk(k=2, axis=1)
643        // 第一行 top2 -> [5, 2], 索引 -> [1, 2]
644        // 第二行 top2 -> [8, 6], 索引 -> [0, 2]
645        let arr = Tensor::new(&[[1, 5, 2], [8, 4, 6]]).unwrap();
646        let (val, idx) = arr.topk(2, 1).unwrap();
647
648        let expected_val = Tensor::new(&[[5, 2], [8, 6]]).unwrap();
649        let expected_idx = Tensor::new(&[[1u32, 2], [0, 2]]).unwrap();
650
651        assert!(val.allclose(&expected_val, 1e-5, 1e-8).unwrap());
652        assert!(idx.allclose(&expected_idx, 1e-5, 1e-8).unwrap());
653    }
654
655    #[test]
656    fn test_topk_matrix_axis0() {
657        // 【关键测试】测试按列求 TopK (axis=0，跨步幅提取)
658        // 用来验证内存 Layout 是否修复正确！
659        // [[1, 5, 2],
660        //  [8, 4, 6],
661        //  [3, 7, 9]]
662        // topk(k=2, axis=0) -> 找每列最大的 2 个
663        // 第 0 列: [1, 8, 3] -> top2 是 8(idx:1), 3(idx:2)
664        // 第 1 列: [5, 4, 7] -> top2 是 7(idx:2), 5(idx:0)
665        // 第 2 列: [2, 6, 9] -> top2 是 9(idx:2), 6(idx:1)
666        // 所以重组后的预期矩阵应该长这样：
667        // Values:  [[8, 7, 9], 
668        //           [3, 5, 6]]
669        // Indices: [[1, 2, 2], 
670        //           [2, 0, 1]]
671        let arr = Tensor::new(&[
672            [1, 5, 2], 
673            [8, 4, 6], 
674            [3, 7, 9]
675        ]).unwrap();
676        let (val, idx) = arr.topk(2, 0).unwrap();
677
678        let expected_val = Tensor::new(&[[8, 7, 9], [3, 5, 6]]).unwrap();
679        let expected_idx = Tensor::new(&[[1u32, 2, 2], [2, 0, 1]]).unwrap();
680
681        assert!(val.allclose(&expected_val, 1e-5, 1e-8).unwrap());
682        assert!(idx.allclose(&expected_idx, 1e-5, 1e-8).unwrap());
683    }
684
685    #[test]
686    fn test_topk_full_sort() {
687        // 测试当 k 等于该维度大小时，相当于进行沿轴的全量降序排序
688        // [[3, 1, 2]] -> topk(3, axis=1)
689        let arr = Tensor::new(&[[3, 1, 2]]).unwrap();
690        let (val, idx) = arr.topk(3, 1).unwrap();
691
692        let expected_val = Tensor::new(&[[3, 2, 1]]).unwrap();
693        let expected_idx = Tensor::new(&[[0u32, 2, 1]]).unwrap();
694
695        assert!(val.allclose(&expected_val, 1e-5, 1e-8).unwrap());
696        assert!(idx.allclose(&expected_idx, 1e-5, 1e-8).unwrap());
697    }
698
699    #[test]
700    fn test_large_ops() {
701        let a = Tensor::randn(0.0, 1.0, (1000, 4096)).unwrap();
702        let now = std::time::Instant::now();
703        let s = a.sum_keepdim(1).unwrap();
704        println!("{:?}", std::time::Instant::now() - now);
705        println!("{:?}", s.dims());
706    }
707}
lumen_core/ops/reduce.rs

lumen_core/ops/
reduce.rs
