burn_dragon_core 0.21.0

burn dragon core model and utilities
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115

use burn::module::{AutodiffModule, Devices, Module, ModuleMapper, ModuleVisitor};
use burn::tensor::backend::{AutodiffBackend, Backend};
use serde::{Deserialize, Serialize};

const MHC_EPS: f32 = 1e-6;

#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum ManifoldHyperConnectionCoefficientPolicy {
    #[default]
    StaticSinkhorn,
    DynamicPositive,
}

impl ManifoldHyperConnectionCoefficientPolicy {
    pub const fn as_str(self) -> &'static str {
        match self {
            Self::StaticSinkhorn => "static_sinkhorn",
            Self::DynamicPositive => "dynamic_positive",
        }
    }

    pub const fn uses_dynamic_stream_controller(self) -> bool {
        matches!(self, Self::DynamicPositive)
    }
}

impl<B: Backend> Module<B> for ManifoldHyperConnectionCoefficientPolicy {
    type Record = ();

    fn collect_devices(&self, devices: Devices<B>) -> Devices<B> {
        devices
    }

    fn fork(self, _device: &B::Device) -> Self {
        self
    }

    fn to_device(self, _device: &B::Device) -> Self {
        self
    }

    fn visit<Visitor: ModuleVisitor<B>>(&self, _visitor: &mut Visitor) {}

    fn map<Mapper: ModuleMapper<B>>(self, _mapper: &mut Mapper) -> Self {
        self
    }

    fn load_record(self, _record: Self::Record) -> Self {
        self
    }

    fn into_record(self) -> Self::Record {}
}

impl<B: AutodiffBackend> AutodiffModule<B> for ManifoldHyperConnectionCoefficientPolicy {
    type InnerModule = ManifoldHyperConnectionCoefficientPolicy;

    fn valid(&self) -> Self::InnerModule {
        *self
    }

    fn from_inner(module: Self::InnerModule) -> Self {
        module
    }
}

/// Configuration for manifold-constrained hyper-connections (mHC).
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct ManifoldHyperConnectionsConfig {
    pub enabled: bool,
    pub num_streams: usize,
    pub num_views: usize,
    pub last_layers: Option<usize>,
    #[serde(default)]
    pub coefficient_policy: ManifoldHyperConnectionCoefficientPolicy,
    pub mhc_iters: usize,
    pub mhc_tau: f32,
    pub add_branch_out_to_residual: bool,
    pub dropout: f64,
}

impl Default for ManifoldHyperConnectionsConfig {
    fn default() -> Self {
        Self {
            enabled: false,
            num_streams: 1,
            num_views: 1,
            last_layers: None,
            coefficient_policy: ManifoldHyperConnectionCoefficientPolicy::StaticSinkhorn,
            mhc_iters: 10,
            mhc_tau: 0.05,
            add_branch_out_to_residual: true,
            dropout: 0.0,
        }
    }
}

impl ManifoldHyperConnectionsConfig {
    pub fn resolved_num_streams(&self) -> usize {
        self.num_streams.max(1)
    }

    pub fn resolved_num_views(&self) -> usize {
        self.num_views.max(1)
    }

    pub fn resolved_tau(&self) -> f32 {
        self.mhc_tau.max(MHC_EPS)
    }

    pub fn resolved_iters(&self) -> usize {
        self.mhc_iters.max(1)
    }
}