ndrs 0.3.0

A tensor library with GPU support
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
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180

//! 切片宏实现

#[macro_export]
macro_rules! s {
    // 基规则：解析完成
    (@parse [$($elems:expr),*] ) => { $crate::view::SliceInfo::new(vec![$($elems),*]) };

    // ==================== 单独常数（整数） ====================
    (@parse [$($elems:expr),*] $num:literal , $($rest:tt)*) => {
        s!(@parse [$($elems,)* $crate::view::SliceArg::Index($num)] $($rest)*)
    };
    (@parse [$($elems:expr),*] $num:literal ) => {
        s!(@parse [$($elems,)* $crate::view::SliceArg::Index($num)] )
    };

    // ==================== a..=b（包含右边界） ====================
    (@parse [$($elems:expr),*] $s:literal ..= $e:literal , $($rest:tt)*) => {
        s!(@parse [$($elems,)* $crate::view::SliceArg::RangeInclusive($s, $e)] $($rest)*)
    };
    (@parse [$($elems:expr),*] $s:literal ..= $e:literal ) => {
        s!(@parse [$($elems,)* $crate::view::SliceArg::RangeInclusive($s, $e)] )
    };
    (@parse [$($elems:expr),*] - $s:literal ..= $e:literal , $($rest:tt)*) => {
        s!(@parse [$($elems,)* $crate::view::SliceArg::RangeInclusive(-$s, $e)] $($rest)*)
    };
    (@parse [$($elems:expr),*] - $s:literal ..= $e:literal ) => {
        s!(@parse [$($elems,)* $crate::view::SliceArg::RangeInclusive(-$s, $e)] )
    };
    (@parse [$($elems:expr),*] $s:literal ..= - $e:literal , $($rest:tt)*) => {
        s!(@parse [$($elems,)* $crate::view::SliceArg::RangeInclusive($s, -$e)] $($rest)*)
    };
    (@parse [$($elems:expr),*] $s:literal ..= - $e:literal ) => {
        s!(@parse [$($elems,)* $crate::view::SliceArg::RangeInclusive($s, -$e)] )
    };
    (@parse [$($elems:expr),*] - $s:literal ..= - $e:literal , $($rest:tt)*) => {
        s!(@parse [$($elems,)* $crate::view::SliceArg::RangeInclusive(-$s, -$e)] $($rest)*)
    };
    (@parse [$($elems:expr),*] - $s:literal ..= - $e:literal ) => {
        s!(@parse [$($elems,)* $crate::view::SliceArg::RangeInclusive(-$s, -$e)] )
    };

    // ==================== a..b（普通范围，步长1） ====================
    (@parse [$($elems:expr),*] $s:literal .. $e:literal , $($rest:tt)*) => {
        s!(@parse [$($elems,)* $crate::view::SliceArg::Range($s, $e, 1)] $($rest)*)
    };
    (@parse [$($elems:expr),*] $s:literal .. $e:literal ) => {
        s!(@parse [$($elems,)* $crate::view::SliceArg::Range($s, $e, 1)] )
    };
    (@parse [$($elems:expr),*] - $s:literal .. $e:literal , $($rest:tt)*) => {
        s!(@parse [$($elems,)* $crate::view::SliceArg::Range(-$s, $e, 1)] $($rest)*)
    };
    (@parse [$($elems:expr),*] - $s:literal .. $e:literal ) => {
        s!(@parse [$($elems,)* $crate::view::SliceArg::Range(-$s, $e, 1)] )
    };
    (@parse [$($elems:expr),*] $s:literal .. - $e:literal , $($rest:tt)*) => {
        s!(@parse [$($elems,)* $crate::view::SliceArg::Range($s, -$e, 1)] $($rest)*)
    };
    (@parse [$($elems:expr),*] $s:literal .. - $e:literal ) => {
        s!(@parse [$($elems,)* $crate::view::SliceArg::Range($s, -$e, 1)] )
    };
    (@parse [$($elems:expr),*] - $s:literal .. - $e:literal , $($rest:tt)*) => {
        s!(@parse [$($elems,)* $crate::view::SliceArg::Range(-$s, -$e, 1)] $($rest)*)
    };
    (@parse [$($elems:expr),*] - $s:literal .. - $e:literal ) => {
        s!(@parse [$($elems,)* $crate::view::SliceArg::Range(-$s, -$e, 1)] )
    };

    // ==================== a..b;s（带步长） ====================
    (@parse [$($elems:expr),*] $s:literal .. $e:literal ; $step:literal , $($rest:tt)*) => {
        s!(@parse [$($elems,)* $crate::view::SliceArg::Range($s, $e, $step)] $($rest)*)
    };
    (@parse [$($elems:expr),*] $s:literal .. $e:literal ; $step:literal ) => {
        s!(@parse [$($elems,)* $crate::view::SliceArg::Range($s, $e, $step)] )
    };
    (@parse [$($elems:expr),*] - $s:literal .. $e:literal ; $step:literal , $($rest:tt)*) => {
        s!(@parse [$($elems,)* $crate::view::SliceArg::Range(-$s, $e, $step)] $($rest)*)
    };
    (@parse [$($elems:expr),*] - $s:literal .. $e:literal ; $step:literal ) => {
        s!(@parse [$($elems,)* $crate::view::SliceArg::Range(-$s, $e, $step)] )
    };
    (@parse [$($elems:expr),*] $s:literal .. - $e:literal ; $step:literal , $($rest:tt)*) => {
        s!(@parse [$($elems,)* $crate::view::SliceArg::Range($s, -$e, $step)] $($rest)*)
    };
    (@parse [$($elems:expr),*] $s:literal .. - $e:literal ; $step:literal ) => {
        s!(@parse [$($elems,)* $crate::view::SliceArg::Range($s, -$e, $step)] )
    };
    (@parse [$($elems:expr),*] - $s:literal .. - $e:literal ; $step:literal , $($rest:tt)*) => {
        s!(@parse [$($elems,)* $crate::view::SliceArg::Range(-$s, -$e, $step)] $($rest)*)
    };
    (@parse [$($elems:expr),*] - $s:literal .. - $e:literal ; $step:literal ) => {
        s!(@parse [$($elems,)* $crate::view::SliceArg::Range(-$s, -$e, $step)] )
    };
    // step 为负
    (@parse [$($elems:expr),*] $s:literal .. $e:literal ; - $step:literal , $($rest:tt)*) => {
        s!(@parse [$($elems,)* $crate::view::SliceArg::Range($s, $e, -$step)] $($rest)*)
    };
    (@parse [$($elems:expr),*] $s:literal .. $e:literal ; - $step:literal ) => {
        s!(@parse [$($elems,)* $crate::view::SliceArg::Range($s, $e, -$step)] )
    };
    (@parse [$($elems:expr),*] - $s:literal .. $e:literal ; - $step:literal , $($rest:tt)*) => {
        s!(@parse [$($elems,)* $crate::view::SliceArg::Range(-$s, $e, -$step)] $($rest)*)
    };
    (@parse [$($elems:expr),*] - $s:literal .. $e:literal ; - $step:literal ) => {
        s!(@parse [$($elems,)* $crate::view::SliceArg::Range(-$s, $e, -$step)] )
    };
    (@parse [$($elems:expr),*] $s:literal .. - $e:literal ; - $step:literal , $($rest:tt)*) => {
        s!(@parse [$($elems,)* $crate::view::SliceArg::Range($s, -$e, -$step)] $($rest)*)
    };
    (@parse [$($elems:expr),*] $s:literal .. - $e:literal ; - $step:literal ) => {
        s!(@parse [$($elems,)* $crate::view::SliceArg::Range($s, -$e, -$step)] )
    };
    (@parse [$($elems:expr),*] - $s:literal .. - $e:literal ; - $step:literal , $($rest:tt)*) => {
        s!(@parse [$($elems,)* $crate::view::SliceArg::Range(-$s, -$e, -$step)] $($rest)*)
    };
    (@parse [$($elems:expr),*] - $s:literal .. - $e:literal ; - $step:literal ) => {
        s!(@parse [$($elems,)* $crate::view::SliceArg::Range(-$s, -$e, -$step)] )
    };

    // ==================== a..（只有起点） ====================
    (@parse [$($elems:expr),*] $s:literal .. , $($rest:tt)*) => {
        s!(@parse [$($elems,)* $crate::view::SliceArg::From($s)] $($rest)*)
    };
    (@parse [$($elems:expr),*] $s:literal .. ) => {
        s!(@parse [$($elems,)* $crate::view::SliceArg::From($s)] )
    };
    (@parse [$($elems:expr),*] - $s:literal .. , $($rest:tt)*) => {
        s!(@parse [$($elems,)* $crate::view::SliceArg::From(-$s)] $($rest)*)
    };
    (@parse [$($elems:expr),*] - $s:literal .. ) => {
        s!(@parse [$($elems,)* $crate::view::SliceArg::From(-$s)] )
    };

    // ==================== ..（全范围） ====================
    (@parse [$($elems:expr),*] .. , $($rest:tt)*) => {
        s!(@parse [$($elems,)* $crate::view::SliceArg::All] $($rest)*)
    };
    (@parse [$($elems:expr),*] .. ) => {
        s!(@parse [$($elems,)* $crate::view::SliceArg::All] )
    };

    // 入口
    ($($rest:tt)*) => {
        s!(@parse [] $($rest)*)
    };
}

#[macro_export]
macro_rules! tensor {
    ($($tt:tt)*) => {{
        let s = stringify!($($tt)*);
        $crate::Tensor::from_string_literal(s)
            .expect("Failed to create tensor from string literal")
    }};
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::dtype::{DTYPE_FLOAT32, DTYPE_INT32};

    #[test]
    fn test_tensor_macro() {
        let t = tensor!([[1, -2], [3, 4]]);
        assert_eq!(t.shape(), &[2, 2]);
        assert_eq!(t.dtype(), DTYPE_INT32);
        assert_eq!(t.to_vec::<i32>().unwrap(), vec![1, -2, 3, 4]);

        let t = tensor!([[1.0, 2.0], [3.0, 4.0]]);
        assert_eq!(t.dtype(), DTYPE_FLOAT32);
        assert_eq!(t.to_vec::<f32>().unwrap(), vec![1.0, 2.0, 3.0, 4.0]);

        let t = tensor!([[1, 2.0], [3, -4]]);
        assert_eq!(t.dtype(), DTYPE_FLOAT32);
        assert_eq!(t.to_vec::<f32>().unwrap(), vec![1.0, 2.0, 3.0, -4.0]);

        let t = tensor!([]);
        assert_eq!(t.shape(), &[0]);
        assert_eq!(t.size(), 0);
    }
}