burn_mamba/mamba2/
mod.rs

1//! # Mamba-2
2//!
3//! Structured State Space Duality (SSD).  Mamba-2 recasts the selective SSM
4//! recurrence as a chunkwise algorithm built from batched GEMMs, making it
5//! tensor-core friendly for training while remaining exactly equivalent to the
6//! recurrent form used for decoding.  See [`mamba2`] for the full SSD math.
7//!
8//! - [`mamba2`] — the SSD block.  The residual layer stack, the full language
9//!   model, and the bidirectional wrappers are the family-generic types in
10//!   [`crate::generic`] (`MambaLatentNet` / `MambaVocabNet` / `MambaBidiLayers`).
11//! - [`cache`] — the conv-window + SSM-state carried between calls.
12//! - [`ssd`] — the pluggable chunkwise SSD algorithms (Minimal / Serial /
13//!   SerialRecalculated) and the backend extension trait.
14
15pub mod cache;
16pub mod mamba2;
17pub mod ssd;
18
19/// Public re-exports for Mamba-2.
20pub mod prelude {
21    use super::*;
22    pub use cache::{Mamba2Cache, Mamba2CacheConfig, Mamba2Caches, Mamba2CachesConfig};
23    pub use mamba2::{Mamba2, Mamba2Config};
24    #[cfg(feature = "autodiff")]
25    pub use ssd::Mamba2AutodiffBackendExt;
26    pub use ssd::Mamba2BackendExt;
27    pub use ssd::{Mamba2SsdInput, Mamba2SsdPath};
28}
burn_mamba/mamba2/mod.rs

burn_mamba/mamba2/
mod.rs
