burn_core/module/param/
group.rs

1use crate::module::{Module, ModuleVisitor, Param};
2
3use alloc::string::String;
4use alloc::vec;
5use alloc::vec::Vec;
6#[cfg(feature = "std")]
7use regex::Regex;
8
9use burn_std::id::ParamId;
10use burn_std::sync::Arc;
11use burn_tensor::{Bool, Int, Tensor};
12
13/// Errors tied to [ParamGroup]'s.
14#[derive(Debug)]
15pub enum ParamGroupError {
16    /// Use of an invalid Regex pattern.
17    InvalidPatternError(String),
18}
19
20#[derive(Default)]
21struct ParamIdCollector {
22    ids: Vec<ParamId>,
23}
24
25impl ParamIdCollector {
26    pub fn ids(&self) -> Vec<ParamId> {
27        self.ids.clone()
28    }
29}
30
31impl ModuleVisitor for ParamIdCollector {
32    fn visit_float<const D: usize>(&mut self, param: &Param<Tensor<D>>) {
33        self.ids.push(param.id);
34    }
35
36    fn visit_int<const D: usize>(&mut self, param: &Param<Tensor<D, Int>>) {
37        self.ids.push(param.id);
38    }
39
40    fn visit_bool<const D: usize>(&mut self, param: &Param<Tensor<D, Bool>>) {
41        self.ids.push(param.id);
42    }
43
44    /// A flag is collected like any other parameter value, so `ids_from_module` over a subtree
45    /// includes its module-owned control state.
46    fn visit_flag(&mut self, flag: &crate::module::Param<crate::module::Flag>) {
47        self.ids.push(flag.id);
48    }
49}
50
51/// A way to represent a group of parameter for a Burn module.
52#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)]
53pub struct ParamGroup {
54    matcher: ParamGroupMatcher,
55    excludes: Option<ParamGroupMatcher>,
56}
57
58impl ParamGroup {
59    /// Evaluates whether a given parameter ID and its module path match this group.
60    pub fn matches(&self, id: &ParamId, path: Option<&str>) -> bool {
61        let matched = self.matcher.matches(id, path);
62
63        let excluded = if let Some(exclude_matcher) = &self.excludes {
64            exclude_matcher.matches(id, path)
65        } else {
66            false
67        };
68
69        matched && !excluded
70    }
71
72    /// Returns a parameter group with all of the module's parameters.
73    pub fn ids_from_module<M: Module>(module: M) -> Self {
74        let mut collector = ParamIdCollector::default();
75        module.visit(&mut collector);
76        Self {
77            matcher: ParamGroupMatcher::from_ids(collector.ids()),
78            excludes: None,
79        }
80    }
81
82    /// Matches parameters by exact text path (e.g., "model.backbone.linear.weight")
83    pub fn from_path(path: impl Into<String>) -> Self {
84        ParamGroup::from_paths(vec![path])
85    }
86
87    /// Matches parameters by exact text paths (e.g., "model.backbone.linear.weight", etc.)
88    pub fn from_paths(paths: Vec<impl Into<String>>) -> Self {
89        Self {
90            matcher: ParamGroupMatcher::Path(Arc::new(PathMatcher::Exact(
91                paths.into_iter().map(|p| p.into()).collect(),
92            ))),
93            excludes: None,
94        }
95    }
96
97    /// Matches parameters that include the predicate in their paths (e.g., "backbone")
98    pub fn from_predicate(path: impl Into<String>) -> Self {
99        ParamGroup::from_predicates(vec![path])
100    }
101
102    /// Matches parameters that include all the predicates in their path (AND logic).
103    /// (e.g., parameter path contains "backbone" and "linear")
104    pub fn from_predicates(paths: Vec<impl Into<String>>) -> Self {
105        Self {
106            matcher: ParamGroupMatcher::Path(Arc::new(PathMatcher::Include(
107                paths.into_iter().map(|p| p.into()).collect(),
108            ))),
109            excludes: None,
110        }
111    }
112
113    /// Matches parameters that include any of the predicates in their path (OR logic).
114    /// (e.g., parameter path contains "backbone" or "linear")
115    pub fn from_any_predicates(paths: Vec<impl Into<String>>) -> Self {
116        let mut matchers: Vec<ParamGroupMatcher> = paths
117            .into_iter()
118            .map(|p| ParamGroupMatcher::Path(Arc::new(PathMatcher::Include(vec![p.into()]))))
119            .collect();
120        let mut main_matcher = if let Some(value) = matchers.pop() {
121            value
122        } else {
123            return Self {
124                matcher: ParamGroupMatcher::Path(Arc::new(PathMatcher::Include(vec![]))),
125                excludes: None,
126            };
127        };
128
129        matchers
130            .iter()
131            .for_each(|m| main_matcher = main_matcher.clone().fuse(m));
132
133        Self {
134            matcher: main_matcher,
135            excludes: None,
136        }
137    }
138
139    #[cfg(feature = "std")]
140    /// Matches parameters by regex pattern (e.g., "^model\.layer\.\d+$")
141    ///
142    /// # Errors
143    /// Returns a [ParamGroupError::InvalidPatternError] if the string cannot be compiled into a valid regex.
144    pub fn from_regex<S: AsRef<str>>(pattern: S) -> Result<Self, ParamGroupError> {
145        ParamGroup::from_regexes(vec![pattern])
146    }
147
148    #[cfg(feature = "std")]
149    /// Matches parameters for all the regex patterns (AND logic).
150    /// (e.g., "^encoder\.layer\.\d+", and "bias$" )
151    ///
152    /// # Errors
153    /// Returns a [ParamGroupError::InvalidPatternError] if the strings cannot be compiled into a valid regex.
154    pub fn from_regexes<S: AsRef<str>>(patterns: Vec<S>) -> Result<Self, ParamGroupError> {
155        let mut new_patterns = vec![];
156        for pattern in patterns {
157            match Regex::new(pattern.as_ref()) {
158                Ok(re) => new_patterns.push(re),
159                Err(e) => {
160                    return Err(ParamGroupError::InvalidPatternError(format!(
161                        "Invalid regex pattern: {e}"
162                    )));
163                }
164            }
165        }
166        Ok(Self {
167            matcher: ParamGroupMatcher::Path(Arc::new(PathMatcher::Regex(new_patterns))),
168            excludes: None,
169        })
170    }
171
172    #[cfg(feature = "std")]
173    /// Matches parameters for any the regex patterns (OR logic).
174    /// (e.g., "^encoder\.layer\.\d+$", or "^decoder\.layer\.\d+$" )
175    ///
176    /// # Errors
177    /// Returns a [ParamGroupError::InvalidPatternError] if the strings cannot be compiled into a valid regex.
178    pub fn from_any_regexes<S: AsRef<str>>(patterns: Vec<S>) -> Result<Self, ParamGroupError> {
179        let mut matchers = vec![];
180        for pattern in patterns {
181            match Regex::new(pattern.as_ref()) {
182                Ok(re) => {
183                    matchers.push(ParamGroupMatcher::Path(Arc::new(PathMatcher::Regex(vec![
184                        re,
185                    ]))))
186                }
187                Err(e) => {
188                    return Err(ParamGroupError::InvalidPatternError(format!(
189                        "Invalid regex pattern: {e}"
190                    )));
191                }
192            }
193        }
194
195        let mut main_matcher = if let Some(value) = matchers.pop() {
196            value
197        } else {
198            return Ok(Self {
199                matcher: ParamGroupMatcher::Path(Arc::new(PathMatcher::Include(vec![]))),
200                excludes: None,
201            });
202        };
203
204        matchers
205            .iter()
206            .for_each(|m| main_matcher = main_matcher.clone().fuse(m));
207
208        Ok(Self {
209            matcher: main_matcher,
210            excludes: None,
211        })
212    }
213
214    /// Matches any parameter.
215    pub fn all() -> Self {
216        Self {
217            matcher: ParamGroupMatcher::All,
218            excludes: None,
219        }
220    }
221
222    /// Matches a specific slice of predefined parameter IDs.
223    pub fn from_ids(ids: Vec<ParamId>) -> Self {
224        Self {
225            matcher: ParamGroupMatcher::Explicit(Arc::new(ids)),
226            excludes: None,
227        }
228    }
229
230    /// Fuse two parameter group.
231    pub fn fuse(self, other: &Self) -> Self {
232        Self {
233            matcher: self.matcher.fuse(&other.matcher),
234            excludes: None,
235        }
236    }
237
238    /// Exclude the given group from the current group
239    pub fn exclude(mut self, group: Self) -> Self {
240        self.excludes = match &self.excludes {
241            Some(excluded) => Some(excluded.clone().fuse(&group.matcher)),
242            None => Some(group.matcher.clone()),
243        };
244        self
245    }
246}
247
248mod arc_serde {
249    use serde::{Deserialize, Deserializer, Serialize, Serializer};
250
251    use super::*;
252
253    pub fn serialize<S, T>(val: &Arc<T>, serializer: S) -> Result<S::Ok, S::Error>
254    where
255        S: Serializer,
256        T: Serialize,
257    {
258        val.as_ref().serialize(serializer)
259    }
260
261    pub fn deserialize<'de, D, T>(deserializer: D) -> Result<Arc<T>, D::Error>
262    where
263        D: Deserializer<'de>,
264        T: Deserialize<'de>,
265    {
266        let v = T::deserialize(deserializer)?;
267        Ok(Arc::new(v))
268    }
269}
270
271#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)]
272enum ParamGroupMatcher {
273    All,
274    #[serde(with = "arc_serde")]
275    Explicit(Arc<Vec<ParamId>>),
276    #[serde(with = "arc_serde")]
277    Path(Arc<PathMatcher>),
278    #[serde(with = "arc_serde")]
279    Combined(Arc<Vec<Self>>),
280}
281
282impl ParamGroupMatcher {
283    pub fn from_ids(ids: Vec<ParamId>) -> Self {
284        Self::Explicit(Arc::new(ids))
285    }
286
287    pub(crate) fn matches(&self, id: &ParamId, path: Option<&str>) -> bool {
288        match self {
289            Self::All => true,
290            Self::Explicit(ids) => ids.contains(id),
291            Self::Path(matcher) => path.is_some_and(|p| matcher.matches(p)),
292            Self::Combined(matchers) => matchers.iter().any(|m| m.matches(id, path)),
293        }
294    }
295
296    fn push_combined(self, other: Self) -> Self {
297        match self {
298            ParamGroupMatcher::Combined(param_group_matchers) => match other.clone() {
299                ParamGroupMatcher::All => Self::All,
300                ParamGroupMatcher::Explicit(_) => {
301                    let mut matchers = (*param_group_matchers).clone();
302                    matchers.push(other);
303                    Self::Combined(Arc::new(matchers))
304                }
305                ParamGroupMatcher::Path(_) => {
306                    let mut matchers = (*param_group_matchers).clone();
307                    matchers.push(other);
308                    Self::Combined(Arc::new(matchers))
309                }
310                ParamGroupMatcher::Combined(other_matchers) => {
311                    let mut matchers = (*param_group_matchers).clone();
312                    matchers.append(&mut (*other_matchers).clone());
313                    Self::Combined(Arc::new(matchers))
314                }
315            },
316            _ => panic!(
317                "`push_combined` should only be called on a ParamGroupMatcher::Combined variant."
318            ),
319        }
320    }
321
322    pub(crate) fn fuse(self, other: &Self) -> Self {
323        match (self.clone(), other.clone()) {
324            (ParamGroupMatcher::All, _) => Self::All,
325            (_, ParamGroupMatcher::All) => Self::All,
326            (ParamGroupMatcher::Explicit(_), ParamGroupMatcher::Combined(_)) => {
327                other.clone().push_combined(self)
328            }
329            (ParamGroupMatcher::Path(_), ParamGroupMatcher::Combined(_)) => {
330                other.clone().push_combined(self)
331            }
332            (ParamGroupMatcher::Combined(_), ParamGroupMatcher::Explicit(_)) => {
333                self.push_combined(other.clone())
334            }
335            (ParamGroupMatcher::Combined(_), ParamGroupMatcher::Path(_)) => {
336                self.push_combined(other.clone())
337            }
338            (ParamGroupMatcher::Combined(_), ParamGroupMatcher::Combined(_)) => {
339                self.push_combined(other.clone())
340            }
341            _ => ParamGroupMatcher::Combined(Arc::new(vec![self, other.clone()])),
342        }
343    }
344}
345
346#[cfg(feature = "std")]
347mod regex_serde {
348    use serde::{Deserialize, Deserializer, Serialize, Serializer};
349
350    use super::*;
351
352    pub fn serialize<S>(regexes: &[Regex], serializer: S) -> Result<S::Ok, S::Error>
353    where
354        S: Serializer,
355    {
356        let v: Vec<&str> = regexes.iter().map(|r| r.as_str()).collect();
357        v.serialize(serializer)
358    }
359
360    pub fn deserialize<'de, D>(deserializer: D) -> Result<Vec<Regex>, D::Error>
361    where
362        D: Deserializer<'de>,
363    {
364        let strings: Vec<String> = Vec::deserialize(deserializer)?;
365        strings
366            .into_iter()
367            .map(|s| Regex::new(&s).map_err(serde::de::Error::custom))
368            .collect()
369    }
370}
371
372#[derive(Clone, serde::Serialize, serde::Deserialize)]
373enum PathMatcher {
374    Exact(Vec<String>),
375    #[cfg(feature = "std")]
376    #[serde(with = "regex_serde")]
377    Regex(Vec<Regex>),
378    Include(Vec<String>),
379}
380
381impl PathMatcher {
382    pub(crate) fn matches(&self, path: &str) -> bool {
383        match self {
384            PathMatcher::Exact(paths) => paths.iter().any(|p| p == path),
385            #[cfg(feature = "std")]
386            PathMatcher::Regex(patterns) => patterns.iter().all(|r| r.is_match(path)),
387            PathMatcher::Include(includes) => includes.iter().all(|inc| path.contains(inc)),
388        }
389    }
390}
391
392impl alloc::fmt::Debug for PathMatcher {
393    fn fmt(&self, f: &mut alloc::fmt::Formatter<'_>) -> alloc::fmt::Result {
394        match self {
395            Self::Exact(arg0) => f.debug_tuple("Exact").field(arg0).finish(),
396            #[cfg(feature = "std")]
397            Self::Regex(arg0) => f.debug_tuple("Regex").field(arg0).finish(),
398            Self::Include(arg0) => f.debug_tuple("Include").field(arg0).finish(),
399        }
400    }
401}
402
403#[cfg(test)]
404mod tests {
405    use super::*;
406    use crate::test_utils::SimpleLinear;
407
408    #[test]
409    fn all_matches_any_parameter() {
410        let group = ParamGroup::all();
411        let id = ParamId::new();
412
413        assert!(group.matches(&id, None));
414    }
415
416    #[test]
417    fn explicit_matches_only_selected_ids() {
418        let id = ParamId::new();
419        let other_id = ParamId::new();
420        let group = ParamGroup::from_ids(vec![id]);
421
422        assert!(group.matches(&id, None));
423        assert!(!group.matches(&other_id, None));
424    }
425
426    #[test]
427    fn path_matcher_requires_path_and_matches_exactly() {
428        let group = ParamGroup::from_path("model.backbone.weight");
429        let id = ParamId::new();
430
431        assert!(group.matches(&id, Some("model.backbone.weight")));
432        assert!(!group.matches(&id, Some("model.backbone.bias")));
433        assert!(!group.matches(&id, None));
434    }
435
436    #[test]
437    fn predicate_matcher_matches_substrings() {
438        let group = ParamGroup::from_predicate("backbone");
439        let id = ParamId::new();
440
441        assert!(group.matches(&id, Some("model.backbone.weight")));
442        assert!(!group.matches(&id, Some("model.other.weight")));
443    }
444
445    #[cfg(feature = "std")]
446    #[test]
447    fn regex_matcher_matches_pattern() {
448        let group = ParamGroup::from_regex(r"^model\.layer\.[0-9]+\.weight$").unwrap();
449        let id = ParamId::new();
450
451        assert!(group.matches(&id, Some("model.layer.3.weight")));
452        assert!(!group.matches(&id, Some("model.layer.weight")));
453    }
454
455    #[test]
456    fn ids_from_module_collects_all_param_ids() {
457        let device = crate::test_device();
458        let module = SimpleLinear::new(4, 8, &device);
459        let weight_id = module.weight.id;
460        let bias_id = module.bias.as_ref().unwrap().id;
461        let group = ParamGroup::ids_from_module(module);
462
463        assert!(group.matches(&weight_id, Some("weight")));
464        assert!(group.matches(&bias_id, Some("bias")));
465    }
466
467    #[test]
468    fn fuse_combines_multiple_groups() {
469        let id = ParamId::new();
470        let group1 = ParamGroup::from_ids(vec![id]);
471        let group2 = ParamGroup::from_path("model.layer.weight");
472        let fused = group1.fuse(&group2);
473
474        assert!(fused.matches(&id, Some("model.other.bias")));
475        assert!(fused.matches(&ParamId::new(), Some("model.layer.weight")));
476        assert!(!fused.matches(&ParamId::new(), Some("model.layer.bias")));
477    }
478
479    #[test]
480    fn exclude_removes_matching_ids_from_a_group() {
481        let id = ParamId::new();
482        let excluded_id = ParamId::new();
483        let exclude_group = ParamGroup::from_ids(vec![excluded_id]);
484        let group = ParamGroup::from_ids(vec![id, excluded_id]).exclude(exclude_group);
485
486        assert!(group.matches(&id, None));
487        assert!(!group.matches(&excluded_id, None));
488    }
489
490    #[test]
491    fn exclude_filters_path_matches() {
492        let id = ParamId::new();
493        let group = ParamGroup::from_predicate("backbone")
494            .exclude(ParamGroup::from_path("model.backbone.bias"));
495
496        assert!(group.matches(&id, Some("model.backbone.weight")));
497        assert!(!group.matches(&id, Some("model.backbone.bias")));
498    }
499
500    #[test]
501    fn exclude_ignores_excludes_on_the_provided_group() {
502        let id = ParamId::new();
503        let excluded = ParamGroup::from_path("model.backbone.bias")
504            .exclude(ParamGroup::from_path("model.backbone.weight"));
505        let group = ParamGroup::from_path("model.backbone.weight").exclude(excluded);
506
507        assert!(group.matches(&id, Some("model.backbone.weight")));
508        assert!(!group.matches(&id, Some("model.backbone.bias")));
509    }
510
511    #[test]
512    fn from_any_predicates_matches_any_predicate() {
513        let group = ParamGroup::from_any_predicates(vec!["backbone", "encoder"]);
514        let id = ParamId::new();
515
516        assert!(group.matches(&id, Some("model.backbone.weight")));
517        assert!(group.matches(&id, Some("model.encoder.weight")));
518        assert!(!group.matches(&id, Some("model.decoder.weight")));
519    }
520
521    #[cfg(feature = "std")]
522    #[test]
523    fn from_any_regexes_matches_any_pattern() {
524        let group =
525            ParamGroup::from_any_regexes(vec![r"^model\.layer\.[0-9]+\.weight$", r"^model\.bias$"])
526                .unwrap();
527        let id = ParamId::new();
528
529        assert!(group.matches(&id, Some("model.layer.3.weight")));
530        assert!(group.matches(&id, Some("model.bias")));
531        assert!(!group.matches(&id, Some("model.layer.weight")));
532        assert!(!group.matches(&id, Some("model.other.weight")));
533    }
534
535    #[cfg(feature = "std")]
536    #[test]
537    fn from_regexes_with_invalid_pattern() {
538        let result = ParamGroup::from_regex(r"[invalid(");
539        assert!(result.is_err());
540
541        let result = ParamGroup::from_regexes(vec![r"[invalid("]);
542        assert!(result.is_err());
543
544        let result = ParamGroup::from_any_regexes(vec![r"[invalid("]);
545        assert!(result.is_err());
546    }
547}
burn_core/module/param/group.rs

burn_core/module/param/
group.rs
