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