burn_core/module/param/
primitive.rs

1use crate::module::{
2    Content, Module, ModuleDisplay, ModuleDisplayDefault, ModuleMapper, ModuleVisitor,
3};
4
5use alloc::{format, string::ToString, vec::Vec};
6
7use burn_tensor::Device;
8use core::fmt::Debug;
9
10impl<T> Module for Option<T>
11where
12    T: Module + Debug + Send + Clone,
13{
14    fn visit<V: ModuleVisitor>(&self, visitor: &mut V) {
15        if let Some(module) = self {
16            module.visit(visitor)
17        }
18    }
19
20    fn map<M: ModuleMapper>(self, mapper: &mut M) -> Self {
21        self.map(|module| module.map(mapper))
22    }
23
24    fn to_device(self, device: &Device) -> Self {
25        self.map(|module| module.to_device(device))
26    }
27
28    fn fork(self, device: &Device) -> Self {
29        self.map(|module| module.fork(device))
30    }
31
32    fn collect_devices(&self, mut devices: Vec<Device>) -> Vec<Device> {
33        if let Some(module) = self.as_ref() {
34            devices = module.collect_devices(devices);
35        }
36
37        devices
38    }
39
40    fn valid(&self) -> Self {
41        self.as_ref().map(|module| module.valid())
42    }
43
44    fn train(self) -> Self {
45        self.map(Module::train)
46    }
47}
48
49impl<T: ModuleDisplay> ModuleDisplayDefault for Option<T> {
50    fn content(&self, content: Content) -> Option<Content> {
51        match self {
52            Some(module) => content.add_single(module).optional(),
53            None => content.add_single("None").optional(),
54        }
55    }
56}
57
58impl<T: ModuleDisplay> ModuleDisplay for Option<T> {}
59
60impl<T> Module for Vec<T>
61where
62    T: Module + Debug + Send + Clone,
63{
64    fn num_params(&self) -> usize {
65        let mut num_params = 0;
66        for module in self.iter() {
67            num_params += module.num_params();
68        }
69
70        num_params
71    }
72
73    fn visit<V: ModuleVisitor>(&self, visitor: &mut V) {
74        for (i, module) in self.iter().enumerate() {
75            let index_str = alloc::format!("{}", i);
76            visitor.enter_module(&index_str, "Vec");
77            module.visit(visitor);
78            visitor.exit_module(&index_str, "Vec");
79        }
80    }
81
82    fn map<M: ModuleMapper>(self, mapper: &mut M) -> Self {
83        self.into_iter()
84            .enumerate()
85            .map(|(i, module)| {
86                let index_str = alloc::format!("{}", i);
87                mapper.enter_module(&index_str, "Vec");
88                let mapped = module.map(mapper);
89                mapper.exit_module(&index_str, "Vec");
90                mapped
91            })
92            .collect()
93    }
94
95    fn to_device(self, device: &Device) -> Self {
96        self.into_iter()
97            .map(|module| module.to_device(device))
98            .collect()
99    }
100
101    fn fork(self, device: &Device) -> Self {
102        self.into_iter().map(|module| module.fork(device)).collect()
103    }
104
105    fn collect_devices(&self, mut devices: Vec<Device>) -> Vec<Device> {
106        for module in self.iter() {
107            devices = module.collect_devices(devices);
108        }
109
110        devices
111    }
112
113    fn valid(&self) -> Self {
114        self.iter().map(|module| module.valid()).collect()
115    }
116
117    fn train(self) -> Self {
118        self.into_iter().map(Module::train).collect()
119    }
120}
121
122impl<T: ModuleDisplay> ModuleDisplayDefault for Vec<T> {
123    fn content(&self, content: Content) -> Option<Content> {
124        self.iter()
125            .enumerate()
126            .fold(content, |acc, (i, module)| {
127                let index = format!("{i}");
128                acc.add(&index, module)
129            })
130            .set_top_level_type(format!("Vec<0..{}>", self.len()).as_str())
131            .optional()
132    }
133}
134
135impl<T: ModuleDisplay> ModuleDisplay for Vec<T> {}
136
137impl<const N: usize, T> Module for [T; N]
138where
139    T: Module + Debug + Send + Clone,
140{
141    fn collect_devices(&self, mut devices: Vec<Device>) -> Vec<Device> {
142        for module in self.iter() {
143            devices = module.collect_devices(devices);
144        }
145
146        devices
147    }
148
149    fn num_params(&self) -> usize {
150        let mut num_params = 0;
151        for module in self.iter() {
152            num_params += module.num_params();
153        }
154
155        num_params
156    }
157
158    fn visit<V: ModuleVisitor>(&self, visitor: &mut V) {
159        for (i, module) in self.iter().enumerate() {
160            let index_str = alloc::format!("{}", i);
161            visitor.enter_module(&index_str, "Array");
162            module.visit(visitor);
163            visitor.exit_module(&index_str, "Array");
164        }
165    }
166
167    fn map<M: ModuleMapper>(self, mapper: &mut M) -> Self {
168        let mut result = Vec::with_capacity(N);
169        for (i, module) in IntoIterator::into_iter(self).enumerate() {
170            let index_str = alloc::format!("{}", i);
171            mapper.enter_module(&index_str, "Array");
172            let mapped = module.map(mapper);
173            mapper.exit_module(&index_str, "Array");
174            result.push(mapped);
175        }
176        result
177            .try_into()
178            .unwrap_or_else(|v: Vec<T>| panic!("Expected array of length {}, got {}", N, v.len()))
179    }
180
181    fn to_device(self, device: &Device) -> Self {
182        self.map(|module| module.to_device(device))
183    }
184
185    fn fork(self, device: &Device) -> Self {
186        self.map(|module| module.fork(device))
187    }
188
189    fn valid(&self) -> Self {
190        self.clone().map(|module| module.valid())
191    }
192
193    fn train(self) -> Self {
194        self.map(Module::train)
195    }
196}
197
198impl<const N: usize, T: ModuleDisplay> ModuleDisplayDefault for [T; N] {
199    fn content(&self, content: Content) -> Option<Content> {
200        self.iter()
201            .enumerate()
202            .fold(content, |acc, (i, module)| {
203                let index = format!("{i}");
204                acc.add(&index, module)
205            })
206            .set_top_level_type(format!("[0..{}]", self.len()).as_str())
207            .optional()
208    }
209}
210
211impl<const N: usize, T: ModuleDisplay> ModuleDisplay for [T; N] {}
212
213/// A macro for generating implementations for tuple modules of different sizes.
214/// For example: `impl_module_tuple!([L0, L1][0, 1])`.
215/// Would generate an implementation for a tuple of size 2.
216/// For this macro to work properly, please adhere to the convention:
217/// `impl_module_tuple!([L0, L1, ..., Ln][0, 1, ..., n])`.
218macro_rules! impl_module_tuple {
219    // `$l` represents the generic modules.
220    // `$i` represents the indices of the modules in the tuple.
221    ([$($l:ident),*][$($i:tt),*]) => {
222        impl<$($l,)*> Module for ($($l,)*)
223        where
224            $($l: Module + Debug + Send + Clone,)*
225        {
226            fn collect_devices(&self, mut devices: Vec<Device>) -> Vec<Device> {
227                $(devices = self.$i.collect_devices(devices);)*
228                devices
229            }
230
231            fn fork(self, device: &Device) -> Self {
232                ($(self.$i.fork(device),)*)
233            }
234
235            fn to_device(self, device: &Device) -> Self {
236                ($(self.$i.to_device(device),)*)
237            }
238
239            fn visit<V: ModuleVisitor>(&self, visitor: &mut V) {
240                $(
241                    let index_str = $i.to_string();
242                    visitor.enter_module(&index_str, "Tuple");
243                    self.$i.visit(visitor);
244                    visitor.exit_module(&index_str, "Tuple");
245                )*
246            }
247
248            fn map<M: ModuleMapper>(self, mapper: &mut M) -> Self {
249                ($(
250                    {
251                        let index_str = $i.to_string();
252                        mapper.enter_module(&index_str, "Tuple");
253                        let mapped = self.$i.map(mapper);
254                        mapper.exit_module(&index_str, "Tuple");
255                        mapped
256                    }
257                ,)*)
258            }
259
260            fn valid(&self) -> Self {
261                ($(self.$i.valid(),)*)
262            }
263
264            fn train(self) -> Self {
265                ($(self.$i.train(),)*)
266            }
267        }
268
269        impl<$($l,)*> ModuleDisplayDefault for ($($l,)*)
270        where
271            $($l: ModuleDisplay,)*
272        {
273            fn content(&self, content: Content) -> Option<Content> {
274                let content = content
275                    $(.add(&format!("{}", $i), &self.$i))*
276                    .set_top_level_type(format!("({})", stringify!($($l),*)).as_str());
277                content.optional()
278            }
279        }
280
281        impl<$($l,)*> ModuleDisplay for ($($l,)*) where $($l: ModuleDisplay,)* {}
282
283    };
284}
285
286impl_module_tuple!([L0, L1][0, 1]);
287impl_module_tuple!([L0, L1, L2][0, 1, 2]);
288impl_module_tuple!([L0, L1, L2, L3][0, 1, 2, 3]);
289impl_module_tuple!([L0, L1, L2, L3, L4][0, 1, 2, 3, 4]);
290impl_module_tuple!([L0, L1, L2, L3, L4, L5][0, 1, 2, 3, 4, 5]);
291impl_module_tuple!([L0, L1, L2, L3, L4, L5, L6][0, 1, 2, 3, 4, 5, 6]);
292impl_module_tuple!([L0, L1, L2, L3, L4, L5, L6, L7][0, 1, 2, 3, 4, 5, 6, 7]);
293impl_module_tuple!([L0, L1, L2, L3, L4, L5, L6, L7, L8][0, 1, 2, 3, 4, 5, 6, 7, 8]);
294impl_module_tuple!([L0, L1, L2, L3, L4, L5, L6, L7, L8, L9][0, 1, 2, 3, 4, 5, 6, 7, 8, 9]);
burn_core/module/param/primitive.rs

burn_core/module/param/
primitive.rs
