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
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135

#![recursion_limit = "256"]

//! Shared Dragon core math and state contracts for the focused language-training repo.

pub mod constants;
pub mod experimental;
pub mod kernel;
pub mod model;
pub mod objective;
pub mod positional;

pub use model::LanguageModuleLrScaleTarget;

pub mod api {
    pub mod config {
        pub use crate::model::LanguageModuleLrScaleTarget;
        pub use crate::{
            AttentionResidualConfig, BlockAttentionResidualConfig,
            BlockAttentionResidualSummaryMode, ClockedSlowMemoryConfig, DragonConfig,
            DragonFiringTargetConfig, DragonFiringTargetKind, DragonInitializationConfig,
            DragonInitializationKind, DragonNeuronGainConfig, DragonNeuronGainKind,
            DragonNormConfig, DragonNormKind, DragonReservoirInitializationConfig,
            DragonResidualScalingConfig, DragonResidualScalingKind, DragonTopologyPriorConfig,
            DragonTopologyPriorKind, FusedAttentionExecutor, FusedKernelConfig,
            FusedProjectionExecutor, LanguageHeadConfig, LatentFanoutScheduleConfig,
            MambaSequenceConfig, ManifoldHyperConnectionCoefficientPolicy,
            ManifoldHyperConnectionsConfig, ResidualConnectorKind, SequenceKernelConfig,
            SequenceMemorySystem, SequenceTrainingExecutor, SummaryMemoryConfig,
            YNeuronRecurrenceConfig,
        };
        pub use burn_dragon_kernel::api::projection::LowrankGradInputExecutor;
    }

    pub mod state {
        pub use crate::ModelState;
    }

    pub mod recurrent {
        pub use crate::{
            DragonModel, DragonNorm, HaltHead, LanguageMhcLayerDiagnostics, LanguagePipelineState,
            LogitsProjectionProfileSnapshot, LowRankResidualMemoryProfileSnapshot,
            LowRankResidualMemoryStageSnapshot, LowRankResidualOutput,
            LowRankResidualProfileSnapshot, logits_projection_profile_reset,
            logits_projection_profile_snapshot, lowrank_residual_memory_profile_reset,
            lowrank_residual_memory_profile_snapshot, lowrank_residual_profile_reset,
            lowrank_residual_profile_snapshot, lowrank_residual_step, lowrank_residual_step_next,
        };
        #[cfg(any(feature = "probe", test))]
        pub use crate::{
            HeadTensorComparisonDiagnostics, HeadTensorGeometryDiagnostics,
            LanguageDragonInitLayerDiagnostics, LanguageLayerStateDeltaDiagnostics,
            LanguageLayerStateSummaryDiagnostics, LanguageLowRankLayerComparisonDiagnostics,
            LanguageLowRankLayerGeometryDiagnostics, TensorComparisonDiagnostics,
            TensorDistributionDiagnostics, TensorStateDeltaDiagnostics,
            TensorStateSummaryDiagnostics, compare_model_states, summarize_model_state,
        };
    }

    pub mod mhc {
        pub use crate::{
            ManifoldHyperConnectionCoefficients, ManifoldHyperConnectionStreamCoefficients,
            ManifoldHyperConnectionStreamOutput, ManifoldHyperConnectionWidthOutput,
            ManifoldHyperConnections, mhc_merge, mhc_merge_with_coefficients, mhc_passthrough,
            mhc_passthrough_with_coefficients, mhc_split, mhc_split_with_coefficients,
        };
    }

    pub mod experimental {
        pub mod connectors {
            pub use crate::{
                AttentionResidual, AttentionResidualConfig, BlockAttentionResidual,
                BlockAttentionResidualConfig, BlockAttentionResidualSummaryMode,
                ResidualConnectorKind,
            };
        }

        pub mod sequence {
            pub use crate::{
                MambaSequenceConfig, SequenceKernelConfig, SequenceMemorySystem,
                SequenceTrainingExecutor,
            };
        }
    }

    pub mod expert {
        pub use crate::constants;
        pub use crate::kernel;
        pub use crate::model;
        pub use crate::objective;
        pub use crate::positional;
    }
}

pub use burn_dragon_kernel::api::projection::LowrankGradInputExecutor;
pub use kernel::{BlockPattern1d, BlockPattern2d, BlockSparseConfig};
#[cfg(any(feature = "probe", test))]
pub use model::LanguageDragonInitLayerDiagnostics;
#[cfg(any(feature = "viz", feature = "probe"))]
pub use model::LayerVizState;
pub use model::{
    AttentionResidual, AttentionResidualConfig, BlockAttentionResidual,
    BlockAttentionResidualConfig, BlockAttentionResidualSummaryMode, ClockedSlowMemoryConfig,
    DragonActivationThresholds, DragonConfig, DragonFiringTargetConfig, DragonFiringTargetKind,
    DragonInitializationConfig, DragonInitializationKind, DragonInitializer, DragonModel,
    DragonNeuronGainConfig, DragonNeuronGainKind, DragonNorm, DragonNormConfig, DragonNormKind,
    DragonProjectionRole, DragonReservoirInitializationConfig, DragonResidualScalingConfig,
    DragonResidualScalingKind, DragonTopologyPriorConfig, DragonTopologyPriorKind,
    FusedAttentionExecutor, FusedKernelConfig, FusedProjectionExecutor, HaltHead,
    LanguageHeadConfig, LanguageMhcLayerDiagnostics, LanguagePipelineState,
    LatentFanoutScheduleConfig, LogitsProjectionProfileSnapshot,
    LowRankResidualMemoryProfileSnapshot, LowRankResidualMemoryStageSnapshot,
    LowRankResidualOutput, LowRankResidualProfileSnapshot, MambaSequenceConfig,
    ManifoldHyperConnectionCoefficientPolicy, ManifoldHyperConnectionCoefficients,
    ManifoldHyperConnectionStreamCoefficients, ManifoldHyperConnectionStreamOutput,
    ManifoldHyperConnectionWidthOutput, ManifoldHyperConnections, ManifoldHyperConnectionsConfig,
    MicroTransformerBlock, ModelState, ResidualConnectorKind, SequenceKernelConfig,
    SequenceMemorySystem, SequenceTrainingExecutor, SharedLowrankActivationBatchStats,
    SharedLowrankContinualBackpropRuntime, SharedLowrankFeatureMetrics, SharedLowrankParamIds,
    SummaryMemoryConfig, YNeuronRecurrenceConfig, logits_projection_profile_reset,
    logits_projection_profile_snapshot, lowrank_residual_memory_profile_reset,
    lowrank_residual_memory_profile_snapshot, lowrank_residual_profile_reset,
    lowrank_residual_profile_snapshot, lowrank_residual_step, lowrank_residual_step_next,
    mhc_merge, mhc_merge_with_coefficients, mhc_passthrough, mhc_passthrough_with_coefficients,
    mhc_split, mhc_split_with_coefficients, near_critical_embedding_initializer,
    near_critical_embedding_std, near_critical_projection_std, near_critical_residual_output_std,
};
#[cfg(any(feature = "probe", test))]
pub use model::{
    HeadTensorComparisonDiagnostics, HeadTensorGeometryDiagnostics,
    LanguageLayerStateDeltaDiagnostics, LanguageLayerStateSummaryDiagnostics,
    LanguageLowRankLayerComparisonDiagnostics, LanguageLowRankLayerGeometryDiagnostics,
    TensorComparisonDiagnostics, TensorDistributionDiagnostics, TensorStateDeltaDiagnostics,
    TensorStateSummaryDiagnostics, compare_model_states, summarize_model_state,
};
pub use positional::RotaryEmbedding;