burn_core/module/param/
primitive.rs

1use crate::module::{
2    Content, Module, ModuleDisplay, ModuleDisplayDefault, ModuleMapper, ModuleVisitor,
3};
4
5use alloc::{
6    boxed::Box,
7    format,
8    string::{String, ToString},
9    vec::Vec,
10};
11
12use burn_tensor::Device;
13use core::fmt::Debug;
14
15// Box is transparent to module paths, like Option: owning a child on the heap
16// must not change the child's checkpoint keys or parameter-group paths.
17impl<T: Module> Module for Box<T> {
18    fn num_params(&self) -> usize {
19        (**self).num_params()
20    }
21
22    fn visit<V: ModuleVisitor>(&self, visitor: &mut V) {
23        (**self).visit(visitor)
24    }
25
26    fn map<M: ModuleMapper>(self, mapper: &mut M) -> Self {
27        Box::new((*self).map(mapper))
28    }
29
30    fn collect_devices(&self, devices: Vec<Device>) -> Vec<Device> {
31        (**self).collect_devices(devices)
32    }
33
34    fn to_device(self, device: &Device) -> Self {
35        Box::new((*self).to_device(device))
36    }
37
38    fn fork(self, device: &Device) -> Self {
39        Box::new((*self).fork(device))
40    }
41
42    fn valid(&self) -> Self {
43        Box::new((**self).valid())
44    }
45
46    fn train(self) -> Self {
47        Box::new((*self).train())
48    }
49
50    fn materialize(self) -> Self {
51        Box::new((*self).materialize())
52    }
53}
54
55impl<T: ModuleDisplay + ?Sized> ModuleDisplayDefault for Box<T> {
56    fn content(&self, content: Content) -> Option<Content> {
57        (**self).content(content)
58    }
59
60    fn num_params(&self) -> usize {
61        (**self).num_params()
62    }
63}
64
65impl<T: ModuleDisplay + ?Sized> ModuleDisplay for Box<T> {
66    fn format(&self, settings: crate::module::DisplaySettings) -> String {
67        (**self).format(settings)
68    }
69
70    fn custom_settings(&self) -> Option<crate::module::DisplaySettings> {
71        (**self).custom_settings()
72    }
73
74    fn custom_content(&self, content: Content) -> Option<Content> {
75        (**self).custom_content(content)
76    }
77}
78
79impl<T> Module for Option<T>
80where
81    T: Module + Debug + Send + Clone,
82{
83    fn visit<V: ModuleVisitor>(&self, visitor: &mut V) {
84        if let Some(module) = self {
85            module.visit(visitor)
86        }
87    }
88
89    fn map<M: ModuleMapper>(self, mapper: &mut M) -> Self {
90        self.map(|module| module.map(mapper))
91    }
92
93    fn to_device(self, device: &Device) -> Self {
94        self.map(|module| module.to_device(device))
95    }
96
97    fn fork(self, device: &Device) -> Self {
98        self.map(|module| module.fork(device))
99    }
100
101    fn collect_devices(&self, mut devices: Vec<Device>) -> Vec<Device> {
102        if let Some(module) = self.as_ref() {
103            devices = module.collect_devices(devices);
104        }
105
106        devices
107    }
108
109    fn valid(&self) -> Self {
110        self.as_ref().map(|module| module.valid())
111    }
112
113    fn train(self) -> Self {
114        self.map(Module::train)
115    }
116
117    fn materialize(self) -> Self {
118        self.map(Module::materialize)
119    }
120}
121
122impl<T: ModuleDisplay> ModuleDisplayDefault for Option<T> {
123    fn content(&self, content: Content) -> Option<Content> {
124        match self {
125            Some(module) => content.add_single(module).optional(),
126            None => content.add_single("None").optional(),
127        }
128    }
129}
130
131impl<T: ModuleDisplay> ModuleDisplay for Option<T> {}
132
133impl<T> Module for Vec<T>
134where
135    T: Module + Debug + Send + Clone,
136{
137    fn num_params(&self) -> usize {
138        let mut num_params = 0;
139        for module in self.iter() {
140            num_params += module.num_params();
141        }
142
143        num_params
144    }
145
146    fn visit<V: ModuleVisitor>(&self, visitor: &mut V) {
147        for (i, module) in self.iter().enumerate() {
148            let index_str = alloc::format!("{}", i);
149            visitor.enter_module(&index_str, "Vec");
150            module.visit(visitor);
151            visitor.exit_module(&index_str, "Vec");
152        }
153    }
154
155    fn map<M: ModuleMapper>(self, mapper: &mut M) -> Self {
156        self.into_iter()
157            .enumerate()
158            .map(|(i, module)| {
159                let index_str = alloc::format!("{}", i);
160                mapper.enter_module(&index_str, "Vec");
161                let mapped = module.map(mapper);
162                mapper.exit_module(&index_str, "Vec");
163                mapped
164            })
165            .collect()
166    }
167
168    fn to_device(self, device: &Device) -> Self {
169        self.into_iter()
170            .map(|module| module.to_device(device))
171            .collect()
172    }
173
174    fn fork(self, device: &Device) -> Self {
175        self.into_iter().map(|module| module.fork(device)).collect()
176    }
177
178    fn collect_devices(&self, mut devices: Vec<Device>) -> Vec<Device> {
179        for module in self.iter() {
180            devices = module.collect_devices(devices);
181        }
182
183        devices
184    }
185
186    fn valid(&self) -> Self {
187        self.iter().map(|module| module.valid()).collect()
188    }
189
190    fn train(self) -> Self {
191        self.into_iter().map(Module::train).collect()
192    }
193
194    fn materialize(self) -> Self {
195        self.into_iter().map(Module::materialize).collect()
196    }
197}
198
199impl<T: ModuleDisplay> ModuleDisplayDefault for Vec<T> {
200    fn content(&self, content: Content) -> Option<Content> {
201        self.iter()
202            .enumerate()
203            .fold(content, |acc, (i, module)| {
204                let index = format!("{i}");
205                acc.add(&index, module)
206            })
207            .set_top_level_type(format!("Vec<0..{}>", self.len()).as_str())
208            .optional()
209    }
210}
211
212impl<T: ModuleDisplay> ModuleDisplay for Vec<T> {}
213
214impl<const N: usize, T> Module for [T; N]
215where
216    T: Module + Debug + Send + Clone,
217{
218    fn collect_devices(&self, mut devices: Vec<Device>) -> Vec<Device> {
219        for module in self.iter() {
220            devices = module.collect_devices(devices);
221        }
222
223        devices
224    }
225
226    fn num_params(&self) -> usize {
227        let mut num_params = 0;
228        for module in self.iter() {
229            num_params += module.num_params();
230        }
231
232        num_params
233    }
234
235    fn visit<V: ModuleVisitor>(&self, visitor: &mut V) {
236        for (i, module) in self.iter().enumerate() {
237            let index_str = alloc::format!("{}", i);
238            visitor.enter_module(&index_str, "Array");
239            module.visit(visitor);
240            visitor.exit_module(&index_str, "Array");
241        }
242    }
243
244    fn map<M: ModuleMapper>(self, mapper: &mut M) -> Self {
245        let mut result = Vec::with_capacity(N);
246        for (i, module) in IntoIterator::into_iter(self).enumerate() {
247            let index_str = alloc::format!("{}", i);
248            mapper.enter_module(&index_str, "Array");
249            let mapped = module.map(mapper);
250            mapper.exit_module(&index_str, "Array");
251            result.push(mapped);
252        }
253        result
254            .try_into()
255            .unwrap_or_else(|v: Vec<T>| panic!("Expected array of length {}, got {}", N, v.len()))
256    }
257
258    fn to_device(self, device: &Device) -> Self {
259        self.map(|module| module.to_device(device))
260    }
261
262    fn fork(self, device: &Device) -> Self {
263        self.map(|module| module.fork(device))
264    }
265
266    fn valid(&self) -> Self {
267        self.clone().map(|module| module.valid())
268    }
269
270    fn train(self) -> Self {
271        self.map(Module::train)
272    }
273
274    fn materialize(self) -> Self {
275        self.map(Module::materialize)
276    }
277}
278
279impl<const N: usize, T: ModuleDisplay> ModuleDisplayDefault for [T; N] {
280    fn content(&self, content: Content) -> Option<Content> {
281        self.iter()
282            .enumerate()
283            .fold(content, |acc, (i, module)| {
284                let index = format!("{i}");
285                acc.add(&index, module)
286            })
287            .set_top_level_type(format!("[0..{}]", self.len()).as_str())
288            .optional()
289    }
290}
291
292impl<const N: usize, T: ModuleDisplay> ModuleDisplay for [T; N] {}
293
294/// A macro for generating implementations for tuple modules of different sizes.
295/// For example: `impl_module_tuple!([L0, L1][0, 1])`.
296/// Would generate an implementation for a tuple of size 2.
297/// For this macro to work properly, please adhere to the convention:
298/// `impl_module_tuple!([L0, L1, ..., Ln][0, 1, ..., n])`.
299macro_rules! impl_module_tuple {
300    // `$l` represents the generic modules.
301    // `$i` represents the indices of the modules in the tuple.
302    ([$($l:ident),*][$($i:tt),*]) => {
303        impl<$($l,)*> Module for ($($l,)*)
304        where
305            $($l: Module + Debug + Send + Clone,)*
306        {
307            fn collect_devices(&self, mut devices: Vec<Device>) -> Vec<Device> {
308                $(devices = self.$i.collect_devices(devices);)*
309                devices
310            }
311
312            fn fork(self, device: &Device) -> Self {
313                ($(self.$i.fork(device),)*)
314            }
315
316            fn to_device(self, device: &Device) -> Self {
317                ($(self.$i.to_device(device),)*)
318            }
319
320            fn visit<V: ModuleVisitor>(&self, visitor: &mut V) {
321                $(
322                    let index_str = $i.to_string();
323                    visitor.enter_module(&index_str, "Tuple");
324                    self.$i.visit(visitor);
325                    visitor.exit_module(&index_str, "Tuple");
326                )*
327            }
328
329            fn map<M: ModuleMapper>(self, mapper: &mut M) -> Self {
330                ($(
331                    {
332                        let index_str = $i.to_string();
333                        mapper.enter_module(&index_str, "Tuple");
334                        let mapped = self.$i.map(mapper);
335                        mapper.exit_module(&index_str, "Tuple");
336                        mapped
337                    }
338                ,)*)
339            }
340
341            fn valid(&self) -> Self {
342                ($(self.$i.valid(),)*)
343            }
344
345            fn train(self) -> Self {
346                ($(self.$i.train(),)*)
347            }
348
349            fn materialize(self) -> Self {
350                ($(self.$i.materialize(),)*)
351            }
352        }
353
354        impl<$($l,)*> ModuleDisplayDefault for ($($l,)*)
355        where
356            $($l: ModuleDisplay,)*
357        {
358            fn content(&self, content: Content) -> Option<Content> {
359                let content = content
360                    $(.add(&format!("{}", $i), &self.$i))*
361                    .set_top_level_type(format!("({})", stringify!($($l),*)).as_str());
362                content.optional()
363            }
364        }
365
366        impl<$($l,)*> ModuleDisplay for ($($l,)*) where $($l: ModuleDisplay,)* {}
367
368    };
369}
370
371impl_module_tuple!([L0, L1][0, 1]);
372impl_module_tuple!([L0, L1, L2][0, 1, 2]);
373impl_module_tuple!([L0, L1, L2, L3][0, 1, 2, 3]);
374impl_module_tuple!([L0, L1, L2, L3, L4][0, 1, 2, 3, 4]);
375impl_module_tuple!([L0, L1, L2, L3, L4, L5][0, 1, 2, 3, 4, 5]);
376impl_module_tuple!([L0, L1, L2, L3, L4, L5, L6][0, 1, 2, 3, 4, 5, 6]);
377impl_module_tuple!([L0, L1, L2, L3, L4, L5, L6, L7][0, 1, 2, 3, 4, 5, 6, 7]);
378impl_module_tuple!([L0, L1, L2, L3, L4, L5, L6, L7, L8][0, 1, 2, 3, 4, 5, 6, 7, 8]);
379impl_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
