burn_mamba/mamba2/ssd/serial_recalculated/mod.rs
1//! Serial SSD with a custom, memory-efficient backward.
2//!
3//! The forward + [`Mamba2BackendExt`] impl live in `serial_recalculated`; the
4//! registered autodiff [`backward`] node and the recompute-based gradient math
5//! in [`combined_backward`] together mirror the official `ssd_combined.py`,
6//! saving ~⅓ of the training memory versus storing every intermediate.
7
8/// The registered custom `Backward` node (autodiff op).
9#[cfg(feature = "autodiff")]
10pub mod backward;
11/// Recompute-based gradient math (the memory-efficient backward).
12pub mod combined_backward;
13mod serial_recalculated;
14
15pub use serial_recalculated::Mamba2BackendExt;
16
17#[cfg(feature = "autodiff")]
18pub use serial_recalculated::Mamba2AutodiffBackendExt;