Skip to main content

burn_mamba/mamba3/double_ssd/ssd/
mod.rs

1//! Standard MIMO-first SSD kernels reused by the double-SSD γ and β passes.
2//!
3//! Three exact reformulations (Minimal / Serial / SerialRecalculated) that
4//! agree on values and gradients; selected via [`Mamba3SsdPath`] and dispatched
5//! by [`Mamba3DoubleSsdInput::run`].
6//!
7//! [`Mamba3SsdPath`]: crate::mamba3::ssd_path::Mamba3SsdPath
8
9/// Matmul/`segsum` MIMO-first SSD with plain autodiff backward.
10pub mod minimal;
11/// Serial-over-chunks SSD with plain autodiff backward.
12pub mod serial;
13/// Serial-over-chunks SSD with a custom recompute backward.
14pub mod serial_recalculated;
15/// The [`Mamba3DoubleSsdInput`] bundle and its `run()` dispatch.
16pub mod ssd_path;
17
18#[cfg(feature = "autodiff")]
19pub use serial_recalculated::Mamba3DoubleSsdAutodiffBackendExt;
20pub use serial_recalculated::Mamba3DoubleSsdBackendExt;
21pub use ssd_path::Mamba3DoubleSsdInput;