rstsr-core 0.7.2

An n-Dimension Rust Tensor Toolkit
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107

#![cfg(not(feature = "col_major"))]

use rand::rngs::StdRng;
use rand::{Rng, SeedableRng};

#[cfg(test)]
mod test {
    use std::time::Instant;

    use super::*;

    #[test]
    fn test_tensor_sum_leading() {
        let n = 512;
        let time = Instant::now();
        let t_rstsr = {
            use rstsr_core::prelude::*;
            let mut rng = StdRng::seed_from_u64(42);

            let vec_b: Vec<f64> = (0..4 * n * n).map(|_| rng.gen()).collect::<_>();
            let b_full = rt::asarray((vec_b, [4, n, n], &DeviceCpuSerial::default()));
            b_full.sum_axes(0)
        };
        println!("{:} usec", time.elapsed().as_micros());
        let t_rstsr = t_rstsr.reshape(-1).to_vec();

        let time = Instant::now();
        let t_rayon = {
            use rstsr_core::prelude::*;
            let mut rng = StdRng::seed_from_u64(42);

            let vec_b: Vec<f64> = (0..4 * n * n).map(|_| rng.gen()).collect::<_>();
            let b_full = rt::asarray((vec_b, [4, n, n]));
            b_full.sum_axes(0)
        };
        println!("{:} usec", time.elapsed().as_micros());
        let t_rayon = t_rayon.reshape(-1).to_vec();

        // let time = Instant::now();
        let t_ndarray = {
            use ndarray::prelude::*;

            let mut rng = StdRng::seed_from_u64(42);
            let vec_b: Vec<f64> = (0..4 * n * n).map(|_| rng.gen()).collect::<_>();
            let b_full = Array::from_shape_vec((4, n, n), vec_b).unwrap();
            b_full.sum_axis(Axis(0))
        };
        println!("{:} usec", time.elapsed().as_micros());
        let t_ndarray = t_ndarray.into_raw_vec();

        let diff =
            t_rstsr.iter().zip(t_ndarray.iter()).map(|(a, b)| (a - b).abs()).max_by(|a, b| a.total_cmp(b)).unwrap();
        assert!(diff < 1e-6);

        let diff =
            t_rayon.iter().zip(t_ndarray.iter()).map(|(a, b)| (a - b).abs()).max_by(|a, b| a.total_cmp(b)).unwrap();
        assert!(diff < 1e-6);
    }

    #[test]
    fn test_tensor_sum_last() {
        let n = 512;
        let time = Instant::now();
        let t_rstsr = {
            use rstsr_core::prelude::*;
            let mut rng = StdRng::seed_from_u64(42);

            let vec_b: Vec<f64> = (0..4 * n * n).map(|_| rng.gen()).collect::<_>();
            let b_full = rt::asarray((vec_b, [4, n, n], &DeviceCpuSerial::default()));
            b_full.sum_axes([-1, -2])
        };
        println!("{:} usec", time.elapsed().as_micros());
        let t_rstsr = t_rstsr.reshape(-1).to_vec();

        let time = Instant::now();
        let t_rayon = {
            use rstsr_core::prelude::*;
            let mut rng = StdRng::seed_from_u64(42);

            let vec_b: Vec<f64> = (0..4 * n * n).map(|_| rng.gen()).collect::<_>();
            let b_full = rt::asarray((vec_b, [4, n, n]));
            b_full.sum_axes([-1, -2])
        };
        println!("{:} usec", time.elapsed().as_micros());
        let t_rayon = t_rayon.reshape(-1).to_vec();

        let time = Instant::now();
        let t_ndarray = {
            use ndarray::prelude::*;

            let mut rng = StdRng::seed_from_u64(42);
            let vec_b: Vec<f64> = (0..4 * n * n).map(|_| rng.gen()).collect::<_>();
            let b_full = Array::from_shape_vec((4, n, n), vec_b).unwrap();
            b_full.sum_axis(Axis(2)).sum_axis(Axis(1))
        };
        println!("{:} usec", time.elapsed().as_micros());
        let t_ndarray = t_ndarray.into_raw_vec();

        let diff =
            t_rstsr.iter().zip(t_ndarray.iter()).map(|(a, b)| (a - b).abs()).max_by(|a, b| a.total_cmp(b)).unwrap();
        assert!(diff < 1e-6);

        let diff =
            t_rayon.iter().zip(t_ndarray.iter()).map(|(a, b)| (a - b).abs()).max_by(|a, b| a.total_cmp(b)).unwrap();
        assert!(diff < 1e-6);
    }
}