Module utils 

Virtual-layer/LR scheduling, class tokens, and custom-backward plumbing.

Shared utilities

Building blocks reused across the Mamba families: custom activations and norms (often fp16-stable variants Burn lacks), loss functions, the segsum / gqa tensor helpers, the custom-backward plumbing (backend_macros / combined_grad / primitive), runtime sanity guards, LR schedulers, and the per-dtype numerical constants below.

Re-exports

pub use class::ClassLatent;

pub use class::ClassToken;

pub use schedule::BidiSchedule;

pub use schedule::Schedule;

pub use scheduler::ConstantLr;

pub use scheduler::CosineAnnealingLr;

pub use scheduler::Lr;

Modules

backend_macros

Macros emitting per-backend BackendExt impls + autodiff marker traits. Macros to cut down the per-backend boilerplate for *SsdBackendExt traits.

class

Learnable [CLS]-style class tokens/latents spliced into the sequence.

combined_grad

Flatten/unflatten (y, final_state) into one tracked tensor for the custom backward. Helpers for the “two-output, one autodiff node” pattern used by both crate::mamba2::ssd::Mamba2BackendExt::ssd_serial_recalculated and crate::mamba3::double_ssd::ssd::Mamba3DoubleSsdBackendExt::double_ssd_serial_recalculated/crate::mamba3::single_ssd::ssd::Mamba3SingleSsdBackendExt::single_ssd_serial_recalculated.

fprim 🔒

Rank-tagged FloatTensor primitive wrapper mirroring the Tensor method API, used by the custom-backward gradient math.

schedule

Virtual-layer → real-weight index scheduling shared by all families.

scheduler

Learning-rate schedulers (cosine-annealing + warmup, constant). Learning rate schedulers for controlling the optimization process.

Functions

div_eps

A small dtype-specific epsilon for safe division (x / (y + eps)), returned as f32.