Struct MambaBackbone 

pub struct MambaBackbone { /* private fields */ }

Complete Mamba backbone: input_proj -> N layers -> norm_f.

Owns all weights. Provides both single-step recurrent inference and access to raw weights for training integration.

use mamba_rs::module::MambaBackbone;
use mamba_rs::MambaConfig;

let cfg = MambaConfig::default();
let backbone = MambaBackbone::init(cfg, 128, 42);

let mut state = backbone.alloc_state();
let mut scratch = backbone.alloc_scratch();
let mut output = vec![0.0f32; backbone.config().d_model];

let input = vec![0.1f32; 128];
backbone.forward_step(&input, &mut output, &mut state, &mut scratch);

Implementations

impl MambaBackbone

pub fn init(cfg: MambaConfig, input_dim: usize, seed: u64) -> Self

Create a backbone with Mamba-specific weight initialization.

Uses Kaiming uniform for projections, log-space init for A, inverse-softplus init for dt_proj bias (Gu & Dao, Section 3.5).

Examples found in repository

examples/inference.rs (line 8)

fn main() {
    let cfg = MambaConfig::default();
    let input_dim = cfg.d_model;

    // Initialize backbone with paper-default weights
    let backbone = MambaBackbone::init(cfg, input_dim, 42);
    println!(
        "Mamba: {} layers, d_model={}, d_inner={}, {} params",
        backbone.n_layers(),
        backbone.config().d_model,
        backbone.config().d_inner(),
        backbone.param_count(),
    );

    // Allocate recurrent state + scratch (once, reuse across steps)
    let mut state = backbone.alloc_state();
    let mut scratch = backbone.alloc_scratch();
    let mut output = vec![0.0f32; backbone.config().d_model];

    // Run 10 inference steps
    for step in 0..10 {
        let input = vec![0.1 * step as f32; input_dim];
        backbone.forward_step(&input, &mut output, &mut state, &mut scratch);

        let norm: f32 = output.iter().map(|x| x * x).sum::<f32>().sqrt();
        println!("step {step}: output L2 norm = {norm:.6}");
    }

    // Reset state for new sequence
    state.reset();
    println!("state reset");
}

examples/train_and_infer.rs (line 32)

fn main() {
    let cfg = MambaConfig::default(); // d_model=128, 3 layers, 366K params
    let input_dim = cfg.d_model;
    let seq_len = 16;
    let lr = 1e-3_f32;
    let steps = 50;

    println!("mamba-rs: train -> save -> load -> inference");
    println!("=============================================");
    println!(
        "Config: d_model={}, layers={}, d_inner={}, params={}",
        cfg.d_model,
        cfg.n_layers,
        cfg.d_inner(),
        MambaBackbone::init(cfg, input_dim, 42).param_count()
    );
    println!();

    // =====================================================================
    // Step 1: Initialize weights
    // =====================================================================
    let inf_weights = MambaWeights::init(&cfg, input_dim, 42);
    let mut tw = train_weights_from_inference(&inf_weights);

    let dims = MambaDims::from_config(&cfg, seq_len, input_dim);
    let di = cfg.d_inner();
    let ds = cfg.d_state;
    let dc = cfg.d_conv;
    let dm = cfg.d_model;

    // Pre-allocate scratch (reused every step)
    let mut acts = MambaBackboneFlat::zeros(dims);
    let mut fwd_scratch = PhaseScratch::zeros(&dims);
    let mut bwd_scratch = BackwardPhaseScratch::zeros(&dims);

    // Synthetic training data: random input, target = zeros (regression)
    let input: Vec<f32> = (0..seq_len * input_dim)
        .map(|i| ((i * 7 + 13) % 100) as f32 * 0.01 - 0.5)
        .collect();

    // =====================================================================
    // Step 2: Training loop (SGD, MSE loss)
    // =====================================================================
    println!("Training ({steps} steps, lr={lr}, seq_len={seq_len}):");

    for step in 0..steps {
        // Precompute a_neg = -exp(a_log)
        let mut a_neg = vec![0.0f32; cfg.n_layers * di * ds];
        for (l, lw) in tw.layers.iter().enumerate() {
            for i in 0..di * ds {
                a_neg[l * di * ds + i] = -lw.a_log[i].exp();
            }
        }

        // Forward pass
        let mut conv = vec![0.0f32; cfg.n_layers * di * dc];
        let mut ssm = vec![0.0f32; cfg.n_layers * di * ds];
        let mut state = MambaRecurrentState {
            conv: &mut conv,
            ssm: &mut ssm,
            a_neg: &a_neg,
        };
        let mut temporal = vec![0.0f32; seq_len * dm];
        forward_mamba_backbone_batched(
            &mut temporal,
            &mut acts,
            &tw,
            &input,
            &mut state,
            &mut fwd_scratch,
            &dims,
        );

        // MSE loss
        let loss: f32 = temporal.iter().map(|v| v * v).sum::<f32>() / (seq_len * dm) as f32;

        // Backward pass
        let scale = 2.0 / (seq_len * dm) as f32;
        let mut d_temporal: Vec<f32> = temporal.iter().map(|v| v * scale).collect();
        let mut grads = TrainMambaWeights::zeros_from_dims(&dims);
        backward_mamba_backbone_batched(
            &mut d_temporal,
            &mut grads,
            &acts,
            &tw,
            &a_neg,
            &mut bwd_scratch,
            &dims,
        );

        // SGD weight update
        sgd_update_all(&mut tw, &grads, lr);

        if step % 10 == 0 || step == steps - 1 {
            println!("  step {step:3}: loss = {loss:.6}");
        }
    }

    // =====================================================================
    // Step 3: Save trained weights
    // =====================================================================
    let save_path = std::env::temp_dir().join("mamba_rs_trained.safetensors");
    let trained_inf_weights = inference_weights_from_train(&tw, &inf_weights);
    mamba_rs::serialize::save(&save_path, &trained_inf_weights, &cfg, input_dim)
        .expect("save failed");
    println!("\nSaved to: {}", save_path.display());

    // =====================================================================
    // Step 4: Load and run inference
    // =====================================================================
    let (loaded_w, loaded_cfg, loaded_input_dim) =
        mamba_rs::serialize::load(Path::new(&save_path)).expect("load failed");
    let bb = MambaBackbone::from_weights(loaded_cfg, loaded_w).expect("from_weights failed");
    assert_eq!(bb.input_dim(), loaded_input_dim);
    println!(
        "Loaded: d_model={}, layers={}, params={}",
        loaded_cfg.d_model,
        loaded_cfg.n_layers,
        bb.param_count()
    );

    // Run inference step-by-step
    let mut inf_state = bb.alloc_state();
    let mut scratch = bb.alloc_scratch();
    let mut output = vec![0.0f32; dm];

    println!("\nInference (10 steps):");
    for step in 0..10 {
        let inp: Vec<f32> = (0..input_dim)
            .map(|i| ((step * input_dim + i) * 7 + 13) as f32 % 100.0 * 0.01 - 0.5)
            .collect();
        bb.forward_step(&inp, &mut output, &mut inf_state, &mut scratch);

        let out_norm: f32 = output.iter().map(|v| v * v).sum::<f32>().sqrt();
        println!("  step {step}: output L2 norm = {out_norm:.6}");
    }

    // Cleanup
    std::fs::remove_file(&save_path).ok();
    println!("\nDone.");
}

pub fn from_weights( cfg: MambaConfig, weights: MambaWeights, ) -> Result<Self, String>

Create a backbone from pre-loaded weights.

Validates dimensions against config. Returns Err on mismatch.

Examples found in repository

examples/train_and_infer.rs (line 130)

fn main() {
    let cfg = MambaConfig::default(); // d_model=128, 3 layers, 366K params
    let input_dim = cfg.d_model;
    let seq_len = 16;
    let lr = 1e-3_f32;
    let steps = 50;

    println!("mamba-rs: train -> save -> load -> inference");
    println!("=============================================");
    println!(
        "Config: d_model={}, layers={}, d_inner={}, params={}",
        cfg.d_model,
        cfg.n_layers,
        cfg.d_inner(),
        MambaBackbone::init(cfg, input_dim, 42).param_count()
    );
    println!();

    // =====================================================================
    // Step 1: Initialize weights
    // =====================================================================
    let inf_weights = MambaWeights::init(&cfg, input_dim, 42);
    let mut tw = train_weights_from_inference(&inf_weights);

    let dims = MambaDims::from_config(&cfg, seq_len, input_dim);
    let di = cfg.d_inner();
    let ds = cfg.d_state;
    let dc = cfg.d_conv;
    let dm = cfg.d_model;

    // Pre-allocate scratch (reused every step)
    let mut acts = MambaBackboneFlat::zeros(dims);
    let mut fwd_scratch = PhaseScratch::zeros(&dims);
    let mut bwd_scratch = BackwardPhaseScratch::zeros(&dims);

    // Synthetic training data: random input, target = zeros (regression)
    let input: Vec<f32> = (0..seq_len * input_dim)
        .map(|i| ((i * 7 + 13) % 100) as f32 * 0.01 - 0.5)
        .collect();

    // =====================================================================
    // Step 2: Training loop (SGD, MSE loss)
    // =====================================================================
    println!("Training ({steps} steps, lr={lr}, seq_len={seq_len}):");

    for step in 0..steps {
        // Precompute a_neg = -exp(a_log)
        let mut a_neg = vec![0.0f32; cfg.n_layers * di * ds];
        for (l, lw) in tw.layers.iter().enumerate() {
            for i in 0..di * ds {
                a_neg[l * di * ds + i] = -lw.a_log[i].exp();
            }
        }

        // Forward pass
        let mut conv = vec![0.0f32; cfg.n_layers * di * dc];
        let mut ssm = vec![0.0f32; cfg.n_layers * di * ds];
        let mut state = MambaRecurrentState {
            conv: &mut conv,
            ssm: &mut ssm,
            a_neg: &a_neg,
        };
        let mut temporal = vec![0.0f32; seq_len * dm];
        forward_mamba_backbone_batched(
            &mut temporal,
            &mut acts,
            &tw,
            &input,
            &mut state,
            &mut fwd_scratch,
            &dims,
        );

        // MSE loss
        let loss: f32 = temporal.iter().map(|v| v * v).sum::<f32>() / (seq_len * dm) as f32;

        // Backward pass
        let scale = 2.0 / (seq_len * dm) as f32;
        let mut d_temporal: Vec<f32> = temporal.iter().map(|v| v * scale).collect();
        let mut grads = TrainMambaWeights::zeros_from_dims(&dims);
        backward_mamba_backbone_batched(
            &mut d_temporal,
            &mut grads,
            &acts,
            &tw,
            &a_neg,
            &mut bwd_scratch,
            &dims,
        );

        // SGD weight update
        sgd_update_all(&mut tw, &grads, lr);

        if step % 10 == 0 || step == steps - 1 {
            println!("  step {step:3}: loss = {loss:.6}");
        }
    }

    // =====================================================================
    // Step 3: Save trained weights
    // =====================================================================
    let save_path = std::env::temp_dir().join("mamba_rs_trained.safetensors");
    let trained_inf_weights = inference_weights_from_train(&tw, &inf_weights);
    mamba_rs::serialize::save(&save_path, &trained_inf_weights, &cfg, input_dim)
        .expect("save failed");
    println!("\nSaved to: {}", save_path.display());

    // =====================================================================
    // Step 4: Load and run inference
    // =====================================================================
    let (loaded_w, loaded_cfg, loaded_input_dim) =
        mamba_rs::serialize::load(Path::new(&save_path)).expect("load failed");
    let bb = MambaBackbone::from_weights(loaded_cfg, loaded_w).expect("from_weights failed");
    assert_eq!(bb.input_dim(), loaded_input_dim);
    println!(
        "Loaded: d_model={}, layers={}, params={}",
        loaded_cfg.d_model,
        loaded_cfg.n_layers,
        bb.param_count()
    );

    // Run inference step-by-step
    let mut inf_state = bb.alloc_state();
    let mut scratch = bb.alloc_scratch();
    let mut output = vec![0.0f32; dm];

    println!("\nInference (10 steps):");
    for step in 0..10 {
        let inp: Vec<f32> = (0..input_dim)
            .map(|i| ((step * input_dim + i) * 7 + 13) as f32 % 100.0 * 0.01 - 0.5)
            .collect();
        bb.forward_step(&inp, &mut output, &mut inf_state, &mut scratch);

        let out_norm: f32 = output.iter().map(|v| v * v).sum::<f32>().sqrt();
        println!("  step {step}: output L2 norm = {out_norm:.6}");
    }

    // Cleanup
    std::fs::remove_file(&save_path).ok();
    println!("\nDone.");
}

pub fn into_weights(self) -> MambaWeights

Extract owned weights (consuming self).

pub fn weights(&self) -> &MambaWeights

Read-only weight access.

pub fn weights_mut(&mut self) -> &mut MambaWeights

Mutable weight access (for optimizer updates).

pub fn layer(&self, index: usize) -> &MambaLayerWeights

Read-only access to a specific layer’s weights.

pub fn layer_mut(&mut self, index: usize) -> &mut MambaLayerWeights

Mutable access to a specific layer’s weights.

pub fn n_layers(&self) -> usize

Number of layers.

Examples found in repository

examples/inference.rs (line 11)

fn main() {
    let cfg = MambaConfig::default();
    let input_dim = cfg.d_model;

    // Initialize backbone with paper-default weights
    let backbone = MambaBackbone::init(cfg, input_dim, 42);
    println!(
        "Mamba: {} layers, d_model={}, d_inner={}, {} params",
        backbone.n_layers(),
        backbone.config().d_model,
        backbone.config().d_inner(),
        backbone.param_count(),
    );

    // Allocate recurrent state + scratch (once, reuse across steps)
    let mut state = backbone.alloc_state();
    let mut scratch = backbone.alloc_scratch();
    let mut output = vec![0.0f32; backbone.config().d_model];

    // Run 10 inference steps
    for step in 0..10 {
        let input = vec![0.1 * step as f32; input_dim];
        backbone.forward_step(&input, &mut output, &mut state, &mut scratch);

        let norm: f32 = output.iter().map(|x| x * x).sum::<f32>().sqrt();
        println!("step {step}: output L2 norm = {norm:.6}");
    }

    // Reset state for new sequence
    state.reset();
    println!("state reset");
}

pub fn param_count(&self) -> usize

Total parameter count.

Examples found in repository

examples/inference.rs (line 14)

fn main() {
    let cfg = MambaConfig::default();
    let input_dim = cfg.d_model;

    // Initialize backbone with paper-default weights
    let backbone = MambaBackbone::init(cfg, input_dim, 42);
    println!(
        "Mamba: {} layers, d_model={}, d_inner={}, {} params",
        backbone.n_layers(),
        backbone.config().d_model,
        backbone.config().d_inner(),
        backbone.param_count(),
    );

    // Allocate recurrent state + scratch (once, reuse across steps)
    let mut state = backbone.alloc_state();
    let mut scratch = backbone.alloc_scratch();
    let mut output = vec![0.0f32; backbone.config().d_model];

    // Run 10 inference steps
    for step in 0..10 {
        let input = vec![0.1 * step as f32; input_dim];
        backbone.forward_step(&input, &mut output, &mut state, &mut scratch);

        let norm: f32 = output.iter().map(|x| x * x).sum::<f32>().sqrt();
        println!("step {step}: output L2 norm = {norm:.6}");
    }

    // Reset state for new sequence
    state.reset();
    println!("state reset");
}

examples/train_and_infer.rs (line 32)

fn main() {
    let cfg = MambaConfig::default(); // d_model=128, 3 layers, 366K params
    let input_dim = cfg.d_model;
    let seq_len = 16;
    let lr = 1e-3_f32;
    let steps = 50;

    println!("mamba-rs: train -> save -> load -> inference");
    println!("=============================================");
    println!(
        "Config: d_model={}, layers={}, d_inner={}, params={}",
        cfg.d_model,
        cfg.n_layers,
        cfg.d_inner(),
        MambaBackbone::init(cfg, input_dim, 42).param_count()
    );
    println!();

    // =====================================================================
    // Step 1: Initialize weights
    // =====================================================================
    let inf_weights = MambaWeights::init(&cfg, input_dim, 42);
    let mut tw = train_weights_from_inference(&inf_weights);

    let dims = MambaDims::from_config(&cfg, seq_len, input_dim);
    let di = cfg.d_inner();
    let ds = cfg.d_state;
    let dc = cfg.d_conv;
    let dm = cfg.d_model;

    // Pre-allocate scratch (reused every step)
    let mut acts = MambaBackboneFlat::zeros(dims);
    let mut fwd_scratch = PhaseScratch::zeros(&dims);
    let mut bwd_scratch = BackwardPhaseScratch::zeros(&dims);

    // Synthetic training data: random input, target = zeros (regression)
    let input: Vec<f32> = (0..seq_len * input_dim)
        .map(|i| ((i * 7 + 13) % 100) as f32 * 0.01 - 0.5)
        .collect();

    // =====================================================================
    // Step 2: Training loop (SGD, MSE loss)
    // =====================================================================
    println!("Training ({steps} steps, lr={lr}, seq_len={seq_len}):");

    for step in 0..steps {
        // Precompute a_neg = -exp(a_log)
        let mut a_neg = vec![0.0f32; cfg.n_layers * di * ds];
        for (l, lw) in tw.layers.iter().enumerate() {
            for i in 0..di * ds {
                a_neg[l * di * ds + i] = -lw.a_log[i].exp();
            }
        }

        // Forward pass
        let mut conv = vec![0.0f32; cfg.n_layers * di * dc];
        let mut ssm = vec![0.0f32; cfg.n_layers * di * ds];
        let mut state = MambaRecurrentState {
            conv: &mut conv,
            ssm: &mut ssm,
            a_neg: &a_neg,
        };
        let mut temporal = vec![0.0f32; seq_len * dm];
        forward_mamba_backbone_batched(
            &mut temporal,
            &mut acts,
            &tw,
            &input,
            &mut state,
            &mut fwd_scratch,
            &dims,
        );

        // MSE loss
        let loss: f32 = temporal.iter().map(|v| v * v).sum::<f32>() / (seq_len * dm) as f32;

        // Backward pass
        let scale = 2.0 / (seq_len * dm) as f32;
        let mut d_temporal: Vec<f32> = temporal.iter().map(|v| v * scale).collect();
        let mut grads = TrainMambaWeights::zeros_from_dims(&dims);
        backward_mamba_backbone_batched(
            &mut d_temporal,
            &mut grads,
            &acts,
            &tw,
            &a_neg,
            &mut bwd_scratch,
            &dims,
        );

        // SGD weight update
        sgd_update_all(&mut tw, &grads, lr);

        if step % 10 == 0 || step == steps - 1 {
            println!("  step {step:3}: loss = {loss:.6}");
        }
    }

    // =====================================================================
    // Step 3: Save trained weights
    // =====================================================================
    let save_path = std::env::temp_dir().join("mamba_rs_trained.safetensors");
    let trained_inf_weights = inference_weights_from_train(&tw, &inf_weights);
    mamba_rs::serialize::save(&save_path, &trained_inf_weights, &cfg, input_dim)
        .expect("save failed");
    println!("\nSaved to: {}", save_path.display());

    // =====================================================================
    // Step 4: Load and run inference
    // =====================================================================
    let (loaded_w, loaded_cfg, loaded_input_dim) =
        mamba_rs::serialize::load(Path::new(&save_path)).expect("load failed");
    let bb = MambaBackbone::from_weights(loaded_cfg, loaded_w).expect("from_weights failed");
    assert_eq!(bb.input_dim(), loaded_input_dim);
    println!(
        "Loaded: d_model={}, layers={}, params={}",
        loaded_cfg.d_model,
        loaded_cfg.n_layers,
        bb.param_count()
    );

    // Run inference step-by-step
    let mut inf_state = bb.alloc_state();
    let mut scratch = bb.alloc_scratch();
    let mut output = vec![0.0f32; dm];

    println!("\nInference (10 steps):");
    for step in 0..10 {
        let inp: Vec<f32> = (0..input_dim)
            .map(|i| ((step * input_dim + i) * 7 + 13) as f32 % 100.0 * 0.01 - 0.5)
            .collect();
        bb.forward_step(&inp, &mut output, &mut inf_state, &mut scratch);

        let out_norm: f32 = output.iter().map(|v| v * v).sum::<f32>().sqrt();
        println!("  step {step}: output L2 norm = {out_norm:.6}");
    }

    // Cleanup
    std::fs::remove_file(&save_path).ok();
    println!("\nDone.");
}

pub fn config(&self) -> &MambaConfig

The config this backbone was built with.

Examples found in repository

examples/inference.rs (line 12)

fn main() {
    let cfg = MambaConfig::default();
    let input_dim = cfg.d_model;

    // Initialize backbone with paper-default weights
    let backbone = MambaBackbone::init(cfg, input_dim, 42);
    println!(
        "Mamba: {} layers, d_model={}, d_inner={}, {} params",
        backbone.n_layers(),
        backbone.config().d_model,
        backbone.config().d_inner(),
        backbone.param_count(),
    );

    // Allocate recurrent state + scratch (once, reuse across steps)
    let mut state = backbone.alloc_state();
    let mut scratch = backbone.alloc_scratch();
    let mut output = vec![0.0f32; backbone.config().d_model];

    // Run 10 inference steps
    for step in 0..10 {
        let input = vec![0.1 * step as f32; input_dim];
        backbone.forward_step(&input, &mut output, &mut state, &mut scratch);

        let norm: f32 = output.iter().map(|x| x * x).sum::<f32>().sqrt();
        println!("step {step}: output L2 norm = {norm:.6}");
    }

    // Reset state for new sequence
    state.reset();
    println!("state reset");
}

pub fn input_dim(&self) -> usize

External input dimension.

Examples found in repository

examples/train_and_infer.rs (line 131)

fn main() {
    let cfg = MambaConfig::default(); // d_model=128, 3 layers, 366K params
    let input_dim = cfg.d_model;
    let seq_len = 16;
    let lr = 1e-3_f32;
    let steps = 50;

    println!("mamba-rs: train -> save -> load -> inference");
    println!("=============================================");
    println!(
        "Config: d_model={}, layers={}, d_inner={}, params={}",
        cfg.d_model,
        cfg.n_layers,
        cfg.d_inner(),
        MambaBackbone::init(cfg, input_dim, 42).param_count()
    );
    println!();

    // =====================================================================
    // Step 1: Initialize weights
    // =====================================================================
    let inf_weights = MambaWeights::init(&cfg, input_dim, 42);
    let mut tw = train_weights_from_inference(&inf_weights);

    let dims = MambaDims::from_config(&cfg, seq_len, input_dim);
    let di = cfg.d_inner();
    let ds = cfg.d_state;
    let dc = cfg.d_conv;
    let dm = cfg.d_model;

    // Pre-allocate scratch (reused every step)
    let mut acts = MambaBackboneFlat::zeros(dims);
    let mut fwd_scratch = PhaseScratch::zeros(&dims);
    let mut bwd_scratch = BackwardPhaseScratch::zeros(&dims);

    // Synthetic training data: random input, target = zeros (regression)
    let input: Vec<f32> = (0..seq_len * input_dim)
        .map(|i| ((i * 7 + 13) % 100) as f32 * 0.01 - 0.5)
        .collect();

    // =====================================================================
    // Step 2: Training loop (SGD, MSE loss)
    // =====================================================================
    println!("Training ({steps} steps, lr={lr}, seq_len={seq_len}):");

    for step in 0..steps {
        // Precompute a_neg = -exp(a_log)
        let mut a_neg = vec![0.0f32; cfg.n_layers * di * ds];
        for (l, lw) in tw.layers.iter().enumerate() {
            for i in 0..di * ds {
                a_neg[l * di * ds + i] = -lw.a_log[i].exp();
            }
        }

        // Forward pass
        let mut conv = vec![0.0f32; cfg.n_layers * di * dc];
        let mut ssm = vec![0.0f32; cfg.n_layers * di * ds];
        let mut state = MambaRecurrentState {
            conv: &mut conv,
            ssm: &mut ssm,
            a_neg: &a_neg,
        };
        let mut temporal = vec![0.0f32; seq_len * dm];
        forward_mamba_backbone_batched(
            &mut temporal,
            &mut acts,
            &tw,
            &input,
            &mut state,
            &mut fwd_scratch,
            &dims,
        );

        // MSE loss
        let loss: f32 = temporal.iter().map(|v| v * v).sum::<f32>() / (seq_len * dm) as f32;

        // Backward pass
        let scale = 2.0 / (seq_len * dm) as f32;
        let mut d_temporal: Vec<f32> = temporal.iter().map(|v| v * scale).collect();
        let mut grads = TrainMambaWeights::zeros_from_dims(&dims);
        backward_mamba_backbone_batched(
            &mut d_temporal,
            &mut grads,
            &acts,
            &tw,
            &a_neg,
            &mut bwd_scratch,
            &dims,
        );

        // SGD weight update
        sgd_update_all(&mut tw, &grads, lr);

        if step % 10 == 0 || step == steps - 1 {
            println!("  step {step:3}: loss = {loss:.6}");
        }
    }

    // =====================================================================
    // Step 3: Save trained weights
    // =====================================================================
    let save_path = std::env::temp_dir().join("mamba_rs_trained.safetensors");
    let trained_inf_weights = inference_weights_from_train(&tw, &inf_weights);
    mamba_rs::serialize::save(&save_path, &trained_inf_weights, &cfg, input_dim)
        .expect("save failed");
    println!("\nSaved to: {}", save_path.display());

    // =====================================================================
    // Step 4: Load and run inference
    // =====================================================================
    let (loaded_w, loaded_cfg, loaded_input_dim) =
        mamba_rs::serialize::load(Path::new(&save_path)).expect("load failed");
    let bb = MambaBackbone::from_weights(loaded_cfg, loaded_w).expect("from_weights failed");
    assert_eq!(bb.input_dim(), loaded_input_dim);
    println!(
        "Loaded: d_model={}, layers={}, params={}",
        loaded_cfg.d_model,
        loaded_cfg.n_layers,
        bb.param_count()
    );

    // Run inference step-by-step
    let mut inf_state = bb.alloc_state();
    let mut scratch = bb.alloc_scratch();
    let mut output = vec![0.0f32; dm];

    println!("\nInference (10 steps):");
    for step in 0..10 {
        let inp: Vec<f32> = (0..input_dim)
            .map(|i| ((step * input_dim + i) * 7 + 13) as f32 % 100.0 * 0.01 - 0.5)
            .collect();
        bb.forward_step(&inp, &mut output, &mut inf_state, &mut scratch);

        let out_norm: f32 = output.iter().map(|v| v * v).sum::<f32>().sqrt();
        println!("  step {step}: output L2 norm = {out_norm:.6}");
    }

    // Cleanup
    std::fs::remove_file(&save_path).ok();
    println!("\nDone.");
}

pub fn forward_step( &self, input: &[f32], output: &mut [f32], state: &mut MambaState, scratch: &mut MambaStepScratch, )

Single-step recurrent forward through the full backbone.

input_proj(input) -> N x layer_step -> norm_f -> output

Zero allocations per call. Delegates to mamba_step.

Examples found in repository

examples/inference.rs (line 25)

fn main() {
    let cfg = MambaConfig::default();
    let input_dim = cfg.d_model;

    // Initialize backbone with paper-default weights
    let backbone = MambaBackbone::init(cfg, input_dim, 42);
    println!(
        "Mamba: {} layers, d_model={}, d_inner={}, {} params",
        backbone.n_layers(),
        backbone.config().d_model,
        backbone.config().d_inner(),
        backbone.param_count(),
    );

    // Allocate recurrent state + scratch (once, reuse across steps)
    let mut state = backbone.alloc_state();
    let mut scratch = backbone.alloc_scratch();
    let mut output = vec![0.0f32; backbone.config().d_model];

    // Run 10 inference steps
    for step in 0..10 {
        let input = vec![0.1 * step as f32; input_dim];
        backbone.forward_step(&input, &mut output, &mut state, &mut scratch);

        let norm: f32 = output.iter().map(|x| x * x).sum::<f32>().sqrt();
        println!("step {step}: output L2 norm = {norm:.6}");
    }

    // Reset state for new sequence
    state.reset();
    println!("state reset");
}

examples/train_and_infer.rs (line 149)

fn main() {
    let cfg = MambaConfig::default(); // d_model=128, 3 layers, 366K params
    let input_dim = cfg.d_model;
    let seq_len = 16;
    let lr = 1e-3_f32;
    let steps = 50;

    println!("mamba-rs: train -> save -> load -> inference");
    println!("=============================================");
    println!(
        "Config: d_model={}, layers={}, d_inner={}, params={}",
        cfg.d_model,
        cfg.n_layers,
        cfg.d_inner(),
        MambaBackbone::init(cfg, input_dim, 42).param_count()
    );
    println!();

    // =====================================================================
    // Step 1: Initialize weights
    // =====================================================================
    let inf_weights = MambaWeights::init(&cfg, input_dim, 42);
    let mut tw = train_weights_from_inference(&inf_weights);

    let dims = MambaDims::from_config(&cfg, seq_len, input_dim);
    let di = cfg.d_inner();
    let ds = cfg.d_state;
    let dc = cfg.d_conv;
    let dm = cfg.d_model;

    // Pre-allocate scratch (reused every step)
    let mut acts = MambaBackboneFlat::zeros(dims);
    let mut fwd_scratch = PhaseScratch::zeros(&dims);
    let mut bwd_scratch = BackwardPhaseScratch::zeros(&dims);

    // Synthetic training data: random input, target = zeros (regression)
    let input: Vec<f32> = (0..seq_len * input_dim)
        .map(|i| ((i * 7 + 13) % 100) as f32 * 0.01 - 0.5)
        .collect();

    // =====================================================================
    // Step 2: Training loop (SGD, MSE loss)
    // =====================================================================
    println!("Training ({steps} steps, lr={lr}, seq_len={seq_len}):");

    for step in 0..steps {
        // Precompute a_neg = -exp(a_log)
        let mut a_neg = vec![0.0f32; cfg.n_layers * di * ds];
        for (l, lw) in tw.layers.iter().enumerate() {
            for i in 0..di * ds {
                a_neg[l * di * ds + i] = -lw.a_log[i].exp();
            }
        }

        // Forward pass
        let mut conv = vec![0.0f32; cfg.n_layers * di * dc];
        let mut ssm = vec![0.0f32; cfg.n_layers * di * ds];
        let mut state = MambaRecurrentState {
            conv: &mut conv,
            ssm: &mut ssm,
            a_neg: &a_neg,
        };
        let mut temporal = vec![0.0f32; seq_len * dm];
        forward_mamba_backbone_batched(
            &mut temporal,
            &mut acts,
            &tw,
            &input,
            &mut state,
            &mut fwd_scratch,
            &dims,
        );

        // MSE loss
        let loss: f32 = temporal.iter().map(|v| v * v).sum::<f32>() / (seq_len * dm) as f32;

        // Backward pass
        let scale = 2.0 / (seq_len * dm) as f32;
        let mut d_temporal: Vec<f32> = temporal.iter().map(|v| v * scale).collect();
        let mut grads = TrainMambaWeights::zeros_from_dims(&dims);
        backward_mamba_backbone_batched(
            &mut d_temporal,
            &mut grads,
            &acts,
            &tw,
            &a_neg,
            &mut bwd_scratch,
            &dims,
        );

        // SGD weight update
        sgd_update_all(&mut tw, &grads, lr);

        if step % 10 == 0 || step == steps - 1 {
            println!("  step {step:3}: loss = {loss:.6}");
        }
    }

    // =====================================================================
    // Step 3: Save trained weights
    // =====================================================================
    let save_path = std::env::temp_dir().join("mamba_rs_trained.safetensors");
    let trained_inf_weights = inference_weights_from_train(&tw, &inf_weights);
    mamba_rs::serialize::save(&save_path, &trained_inf_weights, &cfg, input_dim)
        .expect("save failed");
    println!("\nSaved to: {}", save_path.display());

    // =====================================================================
    // Step 4: Load and run inference
    // =====================================================================
    let (loaded_w, loaded_cfg, loaded_input_dim) =
        mamba_rs::serialize::load(Path::new(&save_path)).expect("load failed");
    let bb = MambaBackbone::from_weights(loaded_cfg, loaded_w).expect("from_weights failed");
    assert_eq!(bb.input_dim(), loaded_input_dim);
    println!(
        "Loaded: d_model={}, layers={}, params={}",
        loaded_cfg.d_model,
        loaded_cfg.n_layers,
        bb.param_count()
    );

    // Run inference step-by-step
    let mut inf_state = bb.alloc_state();
    let mut scratch = bb.alloc_scratch();
    let mut output = vec![0.0f32; dm];

    println!("\nInference (10 steps):");
    for step in 0..10 {
        let inp: Vec<f32> = (0..input_dim)
            .map(|i| ((step * input_dim + i) * 7 + 13) as f32 % 100.0 * 0.01 - 0.5)
            .collect();
        bb.forward_step(&inp, &mut output, &mut inf_state, &mut scratch);

        let out_norm: f32 = output.iter().map(|v| v * v).sum::<f32>().sqrt();
        println!("  step {step}: output L2 norm = {out_norm:.6}");
    }

    // Cleanup
    std::fs::remove_file(&save_path).ok();
    println!("\nDone.");
}

pub fn forward_sequence( &self, inputs: &[f32], outputs: &mut [f32], state: &mut MambaState, scratch: &mut MambaStepScratch, seq_len: usize, )

Run T inference steps sequentially, collecting all outputs.

inputs: [T * input_dim] — T sequential inputs. outputs: [T * d_model] — T sequential outputs (written in-place). State carries across all T steps (warm-up, offline eval, etc.).

pub fn forward_step_batch( &self, inputs: &[f32], outputs: &mut [f32], states: &mut [MambaState], scratches: &mut [MambaStepScratch], )

Batched single-step forward through the backbone.

Processes B independent samples with the same weights. inputs: [B * input_dim], outputs: [B * d_model].

pub fn alloc_state(&self) -> MambaState

Allocate zeroed recurrent state matching this backbone.

Examples found in repository

examples/inference.rs (line 18)

fn main() {
    let cfg = MambaConfig::default();
    let input_dim = cfg.d_model;

    // Initialize backbone with paper-default weights
    let backbone = MambaBackbone::init(cfg, input_dim, 42);
    println!(
        "Mamba: {} layers, d_model={}, d_inner={}, {} params",
        backbone.n_layers(),
        backbone.config().d_model,
        backbone.config().d_inner(),
        backbone.param_count(),
    );

    // Allocate recurrent state + scratch (once, reuse across steps)
    let mut state = backbone.alloc_state();
    let mut scratch = backbone.alloc_scratch();
    let mut output = vec![0.0f32; backbone.config().d_model];

    // Run 10 inference steps
    for step in 0..10 {
        let input = vec![0.1 * step as f32; input_dim];
        backbone.forward_step(&input, &mut output, &mut state, &mut scratch);

        let norm: f32 = output.iter().map(|x| x * x).sum::<f32>().sqrt();
        println!("step {step}: output L2 norm = {norm:.6}");
    }

    // Reset state for new sequence
    state.reset();
    println!("state reset");
}

examples/train_and_infer.rs (line 140)

fn main() {
    let cfg = MambaConfig::default(); // d_model=128, 3 layers, 366K params
    let input_dim = cfg.d_model;
    let seq_len = 16;
    let lr = 1e-3_f32;
    let steps = 50;

    println!("mamba-rs: train -> save -> load -> inference");
    println!("=============================================");
    println!(
        "Config: d_model={}, layers={}, d_inner={}, params={}",
        cfg.d_model,
        cfg.n_layers,
        cfg.d_inner(),
        MambaBackbone::init(cfg, input_dim, 42).param_count()
    );
    println!();

    // =====================================================================
    // Step 1: Initialize weights
    // =====================================================================
    let inf_weights = MambaWeights::init(&cfg, input_dim, 42);
    let mut tw = train_weights_from_inference(&inf_weights);

    let dims = MambaDims::from_config(&cfg, seq_len, input_dim);
    let di = cfg.d_inner();
    let ds = cfg.d_state;
    let dc = cfg.d_conv;
    let dm = cfg.d_model;

    // Pre-allocate scratch (reused every step)
    let mut acts = MambaBackboneFlat::zeros(dims);
    let mut fwd_scratch = PhaseScratch::zeros(&dims);
    let mut bwd_scratch = BackwardPhaseScratch::zeros(&dims);

    // Synthetic training data: random input, target = zeros (regression)
    let input: Vec<f32> = (0..seq_len * input_dim)
        .map(|i| ((i * 7 + 13) % 100) as f32 * 0.01 - 0.5)
        .collect();

    // =====================================================================
    // Step 2: Training loop (SGD, MSE loss)
    // =====================================================================
    println!("Training ({steps} steps, lr={lr}, seq_len={seq_len}):");

    for step in 0..steps {
        // Precompute a_neg = -exp(a_log)
        let mut a_neg = vec![0.0f32; cfg.n_layers * di * ds];
        for (l, lw) in tw.layers.iter().enumerate() {
            for i in 0..di * ds {
                a_neg[l * di * ds + i] = -lw.a_log[i].exp();
            }
        }

        // Forward pass
        let mut conv = vec![0.0f32; cfg.n_layers * di * dc];
        let mut ssm = vec![0.0f32; cfg.n_layers * di * ds];
        let mut state = MambaRecurrentState {
            conv: &mut conv,
            ssm: &mut ssm,
            a_neg: &a_neg,
        };
        let mut temporal = vec![0.0f32; seq_len * dm];
        forward_mamba_backbone_batched(
            &mut temporal,
            &mut acts,
            &tw,
            &input,
            &mut state,
            &mut fwd_scratch,
            &dims,
        );

        // MSE loss
        let loss: f32 = temporal.iter().map(|v| v * v).sum::<f32>() / (seq_len * dm) as f32;

        // Backward pass
        let scale = 2.0 / (seq_len * dm) as f32;
        let mut d_temporal: Vec<f32> = temporal.iter().map(|v| v * scale).collect();
        let mut grads = TrainMambaWeights::zeros_from_dims(&dims);
        backward_mamba_backbone_batched(
            &mut d_temporal,
            &mut grads,
            &acts,
            &tw,
            &a_neg,
            &mut bwd_scratch,
            &dims,
        );

        // SGD weight update
        sgd_update_all(&mut tw, &grads, lr);

        if step % 10 == 0 || step == steps - 1 {
            println!("  step {step:3}: loss = {loss:.6}");
        }
    }

    // =====================================================================
    // Step 3: Save trained weights
    // =====================================================================
    let save_path = std::env::temp_dir().join("mamba_rs_trained.safetensors");
    let trained_inf_weights = inference_weights_from_train(&tw, &inf_weights);
    mamba_rs::serialize::save(&save_path, &trained_inf_weights, &cfg, input_dim)
        .expect("save failed");
    println!("\nSaved to: {}", save_path.display());

    // =====================================================================
    // Step 4: Load and run inference
    // =====================================================================
    let (loaded_w, loaded_cfg, loaded_input_dim) =
        mamba_rs::serialize::load(Path::new(&save_path)).expect("load failed");
    let bb = MambaBackbone::from_weights(loaded_cfg, loaded_w).expect("from_weights failed");
    assert_eq!(bb.input_dim(), loaded_input_dim);
    println!(
        "Loaded: d_model={}, layers={}, params={}",
        loaded_cfg.d_model,
        loaded_cfg.n_layers,
        bb.param_count()
    );

    // Run inference step-by-step
    let mut inf_state = bb.alloc_state();
    let mut scratch = bb.alloc_scratch();
    let mut output = vec![0.0f32; dm];

    println!("\nInference (10 steps):");
    for step in 0..10 {
        let inp: Vec<f32> = (0..input_dim)
            .map(|i| ((step * input_dim + i) * 7 + 13) as f32 % 100.0 * 0.01 - 0.5)
            .collect();
        bb.forward_step(&inp, &mut output, &mut inf_state, &mut scratch);

        let out_norm: f32 = output.iter().map(|v| v * v).sum::<f32>().sqrt();
        println!("  step {step}: output L2 norm = {out_norm:.6}");
    }

    // Cleanup
    std::fs::remove_file(&save_path).ok();
    println!("\nDone.");
}

pub fn alloc_scratch(&self) -> MambaStepScratch

Allocate inference scratch buffers matching this backbone.

Examples found in repository

examples/inference.rs (line 19)

fn main() {
    let cfg = MambaConfig::default();
    let input_dim = cfg.d_model;

    // Initialize backbone with paper-default weights
    let backbone = MambaBackbone::init(cfg, input_dim, 42);
    println!(
        "Mamba: {} layers, d_model={}, d_inner={}, {} params",
        backbone.n_layers(),
        backbone.config().d_model,
        backbone.config().d_inner(),
        backbone.param_count(),
    );

    // Allocate recurrent state + scratch (once, reuse across steps)
    let mut state = backbone.alloc_state();
    let mut scratch = backbone.alloc_scratch();
    let mut output = vec![0.0f32; backbone.config().d_model];

    // Run 10 inference steps
    for step in 0..10 {
        let input = vec![0.1 * step as f32; input_dim];
        backbone.forward_step(&input, &mut output, &mut state, &mut scratch);

        let norm: f32 = output.iter().map(|x| x * x).sum::<f32>().sqrt();
        println!("step {step}: output L2 norm = {norm:.6}");
    }

    // Reset state for new sequence
    state.reset();
    println!("state reset");
}

examples/train_and_infer.rs (line 141)

fn main() {
    let cfg = MambaConfig::default(); // d_model=128, 3 layers, 366K params
    let input_dim = cfg.d_model;
    let seq_len = 16;
    let lr = 1e-3_f32;
    let steps = 50;

    println!("mamba-rs: train -> save -> load -> inference");
    println!("=============================================");
    println!(
        "Config: d_model={}, layers={}, d_inner={}, params={}",
        cfg.d_model,
        cfg.n_layers,
        cfg.d_inner(),
        MambaBackbone::init(cfg, input_dim, 42).param_count()
    );
    println!();

    // =====================================================================
    // Step 1: Initialize weights
    // =====================================================================
    let inf_weights = MambaWeights::init(&cfg, input_dim, 42);
    let mut tw = train_weights_from_inference(&inf_weights);

    let dims = MambaDims::from_config(&cfg, seq_len, input_dim);
    let di = cfg.d_inner();
    let ds = cfg.d_state;
    let dc = cfg.d_conv;
    let dm = cfg.d_model;

    // Pre-allocate scratch (reused every step)
    let mut acts = MambaBackboneFlat::zeros(dims);
    let mut fwd_scratch = PhaseScratch::zeros(&dims);
    let mut bwd_scratch = BackwardPhaseScratch::zeros(&dims);

    // Synthetic training data: random input, target = zeros (regression)
    let input: Vec<f32> = (0..seq_len * input_dim)
        .map(|i| ((i * 7 + 13) % 100) as f32 * 0.01 - 0.5)
        .collect();

    // =====================================================================
    // Step 2: Training loop (SGD, MSE loss)
    // =====================================================================
    println!("Training ({steps} steps, lr={lr}, seq_len={seq_len}):");

    for step in 0..steps {
        // Precompute a_neg = -exp(a_log)
        let mut a_neg = vec![0.0f32; cfg.n_layers * di * ds];
        for (l, lw) in tw.layers.iter().enumerate() {
            for i in 0..di * ds {
                a_neg[l * di * ds + i] = -lw.a_log[i].exp();
            }
        }

        // Forward pass
        let mut conv = vec![0.0f32; cfg.n_layers * di * dc];
        let mut ssm = vec![0.0f32; cfg.n_layers * di * ds];
        let mut state = MambaRecurrentState {
            conv: &mut conv,
            ssm: &mut ssm,
            a_neg: &a_neg,
        };
        let mut temporal = vec![0.0f32; seq_len * dm];
        forward_mamba_backbone_batched(
            &mut temporal,
            &mut acts,
            &tw,
            &input,
            &mut state,
            &mut fwd_scratch,
            &dims,
        );

        // MSE loss
        let loss: f32 = temporal.iter().map(|v| v * v).sum::<f32>() / (seq_len * dm) as f32;

        // Backward pass
        let scale = 2.0 / (seq_len * dm) as f32;
        let mut d_temporal: Vec<f32> = temporal.iter().map(|v| v * scale).collect();
        let mut grads = TrainMambaWeights::zeros_from_dims(&dims);
        backward_mamba_backbone_batched(
            &mut d_temporal,
            &mut grads,
            &acts,
            &tw,
            &a_neg,
            &mut bwd_scratch,
            &dims,
        );

        // SGD weight update
        sgd_update_all(&mut tw, &grads, lr);

        if step % 10 == 0 || step == steps - 1 {
            println!("  step {step:3}: loss = {loss:.6}");
        }
    }

    // =====================================================================
    // Step 3: Save trained weights
    // =====================================================================
    let save_path = std::env::temp_dir().join("mamba_rs_trained.safetensors");
    let trained_inf_weights = inference_weights_from_train(&tw, &inf_weights);
    mamba_rs::serialize::save(&save_path, &trained_inf_weights, &cfg, input_dim)
        .expect("save failed");
    println!("\nSaved to: {}", save_path.display());

    // =====================================================================
    // Step 4: Load and run inference
    // =====================================================================
    let (loaded_w, loaded_cfg, loaded_input_dim) =
        mamba_rs::serialize::load(Path::new(&save_path)).expect("load failed");
    let bb = MambaBackbone::from_weights(loaded_cfg, loaded_w).expect("from_weights failed");
    assert_eq!(bb.input_dim(), loaded_input_dim);
    println!(
        "Loaded: d_model={}, layers={}, params={}",
        loaded_cfg.d_model,
        loaded_cfg.n_layers,
        bb.param_count()
    );

    // Run inference step-by-step
    let mut inf_state = bb.alloc_state();
    let mut scratch = bb.alloc_scratch();
    let mut output = vec![0.0f32; dm];

    println!("\nInference (10 steps):");
    for step in 0..10 {
        let inp: Vec<f32> = (0..input_dim)
            .map(|i| ((step * input_dim + i) * 7 + 13) as f32 % 100.0 * 0.01 - 0.5)
            .collect();
        bb.forward_step(&inp, &mut output, &mut inf_state, &mut scratch);

        let out_norm: f32 = output.iter().map(|v| v * v).sum::<f32>().sqrt();
        println!("  step {step}: output L2 norm = {out_norm:.6}");
    }

    // Cleanup
    std::fs::remove_file(&save_path).ok();
    println!("\nDone.");
}