burn_core/module/param/
primitive.rs

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