rustytorch/crates/training/rtx-transformers/tests/mamba_step_equivalence.rs

//! D309 equivalence pin — the single-step `MambaRecurrence` is numerically
//! identical to the full-sequence `MambaBlock::forward`.
//!
//! This is the precondition the SMT plan flags as the highest remaining risk:
//! before a recurrent updater can be behaviorally cloned, its exposed
//! `(state, x_t) → (state', y_t)` step MUST reproduce the forward exactly.
//! We run a window through `forward`, then through `step` one timestep at a
//! time from a zero state, and require the outputs to match to the f32 floor.
//!
//! The block uses the `active_block` trick (wide `Δ`) so the scan genuinely
//! contributes — otherwise the near-identity default init would make the test
//! pass trivially without exercising the recurrence.

use std::collections::HashMap;

use rtx_tensor::{Device, Tensor};
use rtx_transformers::layers::mamba::{MambaBlock, MambaConfig};
use rtx_transformers::layers::mamba_step::MambaRecurrence;

const D_MODEL: usize = 8;
const D_STATE: usize = 16;
const D_CONV: usize = 4;
const L: usize = 7;

fn max_abs_diff(a: &[f32], b: &[f32]) -> f32 {
    a.iter()
        .zip(b)
        .map(|(x, y)| (x - y).abs())
        .fold(0.0, f32::max)
}

/// A seeded block with `dt_bias` zeroed (`Δ = softplus(0) ≈ 0.69`) so the scan
/// genuinely drives the output — mirrors `real_selective_scan.rs::active_block`.
fn active_block(dev: &Device, seed: u64) -> MambaBlock {
    let cfg = MambaConfig::new(D_MODEL, D_STATE, D_CONV);
    let base = MambaBlock::new_seeded(cfg.clone(), dev, seed).expect("base");
    let d = cfg.get_d_inner();
    let mut map = HashMap::new();
    for (name, t) in base.persistence_tensors() {
        if name == "dt_bias" {
            map.insert(
                name.to_string(),
                Tensor::from_vec(vec![0.0f32; d], &[d], dev).expect("dtb"),
            );
        } else {
            map.insert(name.to_string(), t.clone());
        }
    }
    MambaBlock::from_persistence_tensors(cfg, dev, map).expect("rebuild")
}

/// Deterministic input window `[1, L, D_MODEL]`.
fn input(dev: &Device) -> (Tensor, Vec<f32>) {
    let mut s = 0x57E9u64;
    let data: Vec<f32> = (0..L * D_MODEL)
        .map(|_| {
            s = s
                .wrapping_mul(6_364_136_223_846_793_005)
                .wrapping_add(1_442_695_040_888_963_407);
            ((s >> 40) as f32 / (1u64 << 24) as f32) * 2.0 - 1.0
        })
        .collect();
    let t = Tensor::from_vec(data.clone(), &[1, L, D_MODEL], dev).expect("input");
    (t, data)
}

#[test]
fn step_matches_full_forward() {
    let dev = Device::cpu();
    let block = active_block(&dev, 0xD309);
    let (x, flat) = input(&dev);

    // Full-sequence reference.
    let full = block
        .forward(&x)
        .expect("forward")
        .output
        .to_vec()
        .expect("vec");
    assert_eq!(full.len(), L * D_MODEL);

    // Stepwise: feed one token at a time from a zero state.
    let rec = MambaRecurrence::from_block(&block).expect("recurrence");
    assert_eq!(rec.d_model(), D_MODEL);
    let mut state = rec.init_state();
    let mut stepwise = Vec::with_capacity(L * D_MODEL);
    for li in 0..L {
        let x_t = &flat[li * D_MODEL..(li + 1) * D_MODEL];
        stepwise.extend(rec.step(x_t, &mut state));
    }

    let diff = max_abs_diff(&full, &stepwise);
    eprintln!("D309 step-vs-forward max_abs_diff = {diff:.2e}");
    assert!(
        diff < 1e-4,
        "stepwise recurrence must match the full forward (max_abs_diff {diff:.2e})"
    );

    // The state must be non-trivial (the scan actually ran).
    assert_eq!(rec.state_size(), state.hidden().len());
    assert!(
        state.hidden().iter().any(|&v| v.abs() > 1e-6),
        "SSM state should be populated after stepping"
    );
}

#[test]
fn fresh_state_is_zero_and_deterministic() {
    let dev = Device::cpu();
    let block = active_block(&dev, 7);
    let rec = MambaRecurrence::from_block(&block).expect("recurrence");
    let s = rec.init_state();
    assert!(s.hidden().iter().all(|&v| v == 0.0), "init state h is zero");

    // Same input from a fresh state → identical output (stateless determinism).
    let (_, flat) = input(&dev);
    let run = |rec: &MambaRecurrence| {
        let mut st = rec.init_state();
        let mut out = Vec::new();
        for li in 0..L {
            out.extend(rec.step(&flat[li * D_MODEL..(li + 1) * D_MODEL], &mut st));
        }
        out
    };
    assert_eq!(run(&rec), run(&rec), "stepping must be deterministic");
}