mod attention;
mod attention_residual;
mod bdh;
mod bdh_support;
mod config;
mod halt;
mod init;
mod low_bit;
mod low_bit_runtime;
mod mhc;
mod micro_transformer;
mod norm;
mod residual_stream;
mod sequence;
mod state;
mod structured_mode;
mod structured_routing;
mod structured_state;
mod structured_step;
pub use attention_residual::{
AttentionResidual, AttentionResidualConfig, BlockAttentionResidual,
BlockAttentionResidualConfig, BlockAttentionResidualSummaryMode, ResidualConnectorKind,
};
pub use bdh::{
BDH, BdhBitNetDeployScaffold, LanguageModuleLrScaleTarget, SharedLowrankActivationBatchStats,
SharedLowrankContinualBackpropRuntime, SharedLowrankFeatureMetrics, SharedLowrankParamIds,
};
#[cfg(any(feature = "probe", test))]
pub use bdh::{
HeadTensorComparisonDiagnostics, HeadTensorGeometryDiagnostics,
LanguageLayerStateDeltaDiagnostics, LanguageLayerStateSummaryDiagnostics,
LanguageLowRankLayerComparisonDiagnostics, LanguageLowRankLayerGeometryDiagnostics,
TensorComparisonDiagnostics, TensorDistributionDiagnostics, TensorStateDeltaDiagnostics,
TensorStateSummaryDiagnostics, compare_model_states, summarize_model_state,
};
#[cfg(any(feature = "probe", test))]
pub use bdh_support::LanguageBdhInitLayerDiagnostics;
pub use bdh_support::{
LanguageMhcLayerDiagnostics, LanguagePipelineState, LogitsProjectionProfileSnapshot,
logits_projection_profile_reset, logits_projection_profile_snapshot,
};
pub use burn_dragon_kernel::api::projection::LowrankGradInputExecutor;
pub use config::{
BDHConfig, ClockedSlowMemoryConfig, FusedAttentionExecutor, FusedKernelConfig,
FusedProjectionExecutor, LanguageHeadConfig, LatentFanoutScheduleConfig, SummaryMemoryConfig,
YNeuronRecurrenceConfig,
};
pub use halt::HaltHead;
pub use init::{
BdhActivationThresholds, BdhFiringTargetConfig, BdhFiringTargetKind, BdhInitializationConfig,
BdhInitializationKind, BdhInitializer, BdhNeuronGainConfig, BdhNeuronGainKind,
BdhProjectionRole, BdhResidualScalingConfig, BdhResidualScalingKind, BdhTopologyPriorConfig,
BdhTopologyPriorKind, near_critical_embedding_initializer, near_critical_embedding_std,
near_critical_projection_std, near_critical_residual_output_std,
};
pub use low_bit::{
BitNetLowBitProtocol, LowBitActivationFormat, LowBitActivationGrouping, LowBitInferenceMode,
LowBitQuantizationConfig, LowBitRhoConfig, LowBitSavedActivationConfig,
LowBitSavedActivationMode, LowBitTargetModule, LowBitTrainingMode, LowBitWeightFormat,
LowBitWeightGrouping, RhoCompressionConfig, RhoCompressionInterval, RhoPrecisionConfig,
};
pub use low_bit_runtime::{
LowBitKernelCapabilities, LowBitKernelFallbackReason, LowBitKernelPlan,
LowBitKernelRuntimeKind, LowBitMemoryBucketEstimate, LowBitMemoryEstimateInput,
LowBitNativeProjectionProfileSnapshot, LowBitProjectionPlan, LowBitSavedActivationInventory,
LowBitSavedActivationRecomputePolicy, LowBitSavedActivationTensorInventoryEntry,
LowBitTrainingProjectionMemoryProfileSnapshot, LowBitTrainingProjectionMemoryStageSnapshot,
LowBitTrainingQuantizeProfileSnapshot, PackedLowBitProjectionArtifacts, PackedRhoBlockEncoding,
PackedRhoBlockState, PackedRhoInt8DeviceState, PackedSavedActivationBuffer,
PackedSavedActivationState, RhoCompressionQualityGate, RhoCompressionStatsSnapshot,
build_low_bit_saved_activation_inventory, estimate_low_bit_memory_buckets,
fake_quantize_activation_ste, fake_quantize_weight_ste, fraction_nonzero,
low_bit_kernel_capabilities, low_bit_kernel_capabilities_for_backend_name,
low_bit_native_decoder_tail_profile_snapshot, low_bit_native_lowrank_profile_snapshot,
low_bit_native_projection_profile_reset, low_bit_training_lowrank_memory_profile_snapshot,
low_bit_training_quantize_profile_snapshot, pack_saved_activation_state,
resolve_low_bit_kernel_plan, resolve_low_bit_kernel_plan_for_backend_name,
rho_compression_profile_reset, rho_compression_profile_snapshot,
rho_compression_snapshot_passes_gate, unpack_saved_activation_state,
};
pub use mhc::{
ManifoldHyperConnectionCoefficientPolicy, ManifoldHyperConnectionCoefficients,
ManifoldHyperConnectionStreamCoefficients, ManifoldHyperConnectionStreamOutput,
ManifoldHyperConnectionWidthOutput, ManifoldHyperConnections, ManifoldHyperConnectionsConfig,
mhc_merge, mhc_merge_with_coefficients, mhc_passthrough, mhc_passthrough_with_coefficients,
mhc_split, mhc_split_with_coefficients,
};
pub use micro_transformer::MicroTransformerBlock;
pub use norm::{DragonNorm, DragonNormConfig, DragonNormKind};
pub use residual_stream::{
LowBitSavedActivationCache, LowRankResidualMemoryProfileSnapshot,
LowRankResidualMemoryStageSnapshot, LowRankResidualOutput, LowRankResidualProfileSnapshot,
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,
};
pub use sequence::{
MambaSequenceConfig, SequenceKernelConfig, SequenceMemorySystem, SequenceTrainingExecutor,
};
#[cfg(any(feature = "viz", feature = "probe"))]
pub use state::LayerVizState;
pub use state::{LayerState, ModelState};
pub use structured_mode::StructuredStepMode;
pub use structured_routing::{
StructuredBankRole, StructuredRouteOperation, StructuredRoutePattern, StructuredRouteSpec,
StructuredRoutingSpec,
};
pub use structured_state::{BankedRhoState, StructuredGridState, StructuredTopologyState};
pub use structured_step::{
StructuredDenseUpdateOutput, structured_dense_update_tokens, structured_predict_decay,
target_major_apply_decay, target_major_decay_add, target_major_identity_read,
target_major_identity_write, target_major_outer_product,
};