Blame for src/ssd.rs

//! Slot Mamba-3 (1-de-6 do tronco 5:1) — scan recorrente diagonal seletivo
//! (SSD) com adjunto EXATO via CustomOp2. Compatível com o Teo D.1(i):
//! fluxo estritamente esquerda→direita (a camada é causal por construção);
//! a bidirecionalidade intra-nível vem das camadas ADBA (H_D1 — r-04 §3).
//!
//! Recorrência: h_t = a_t ⊙ h_{t−1} + (1 − a_t) ⊙ x_t;  y_t = h_t
//! com a_t = σ(W_a·x_t + b_a) ∈ (0,1)^d (portas por canal).
//!
//! Adjunto (exato): dh_t = dy_t + a_{t+1} ⊙ dh_{t+1};
//! dx_t = (1 − a_t) ⊙ dh_t;  da_t = ⟨dh_t, h_{t−1} − x_t⟩ (escalar por canal).

use candle_core::{CpuStorage, CustomOp2, Device, Layout, Result, Shape, Tensor, Var};
use crate::camadas::{Linear, RmsNorm, Swiglu};

pub struct ScanDiagonal;

impl CustomOp2 for ScanDiagonal {
    fn name(&self) -> &'static str { "ssd_scan_diagonal" }

    fn cpu_fwd(&self, s1: &CpuStorage, l1: &Layout, s2: &CpuStorage, l2: &Layout) -> Result<(CpuStorage, Shape)> {
        if !(l1.is_contiguous() && l2.is_contiguous()) {
            return Err(candle_core::Error::Msg(
                "ssd_scan: entradas devem ser contíguas (chame .contiguous())".into(),
            ));
        }
        let x = s1.as_slice::<f32>()?; // (n, d)
        let a = s2.as_slice::<f32>()?; // (n, d) em (0,1)
        let (n, d) = (l1.shape().dims()[0], l1.shape().dims()[1]);
        let mut y = vec![0f32; n * d];
        for j in 0..d {
            let mut h = 0f32;
            for i in 0..n {
                let ai = a[i * d + j];
                h = ai * h + (1.0 - ai) * x[i * d + j];
                y[i * d + j] = h;
            }
        }
        Ok((CpuStorage::F32(y), (n, d).into()))
    }

    fn bwd(&self, arg1: &Tensor, arg2: &Tensor, _res: &Tensor, grad_res: &Tensor) -> Result<(Option<Tensor>, Option<Tensor>)> {
        // recomputa h (forward sem gradiente) e percorre reverso
        let (x, a) = (arg1, arg2);
        let h = Tensor::apply_op2(x, a, ScanDiagonal)?.detach();
        let x_v = x.flatten_all()?.to_vec1::<f32>()?;
        let a_v = a.flatten_all()?.to_vec1::<f32>()?;
        let h_v = h.flatten_all()?.to_vec1::<f32>()?;
        let dy_v = grad_res.flatten_all()?.to_vec1::<f32>()?;
        let (n, d) = (x.dim(0)?, x.dim(1)?);
        let mut dx = vec![0f32; n * d];
        let mut da = vec![0f32; n * d];
        let mut dh = vec![0f32; d];
        for i in (0..n).rev() {
            for j in 0..d {
                let idx = i * d + j;
                let g = dy_v[idx] + dh[j];
                dx[idx] = (1.0 - a_v[idx]) * g;
                let h_ant = if i > 0 { h_v[(i - 1) * d + j] } else { 0.0 };
                da[idx] = g * (h_ant - x_v[idx]);
                dh[j] = a_v[idx] * g;
            }
        }
        let dev = x.device();
        let gx = Tensor::from_vec(dx, (n, d), dev)?;
        let ga = Tensor::from_vec(da, (n, d), dev)?;
        Ok((Some(gx), Some(ga)))
    }
}

/// Bloco SSD (a camada recorrente do tronco): portas por canal + scan + projeção.
pub struct BlocoSsd {
    pub norm1: RmsNorm,
    pub porta_a: Linear, // d → d (logitas do decaimento)
    pub proj_x: Linear,  // d → d
    pub proj_o: Linear,  // d → d
    pub norm2: RmsNorm,
    pub ffn: Swiglu,
}

impl BlocoSsd {
    pub fn novo(d: usize, d_ff: usize, dev: &Device) -> Result<Self> {
        let desvio = 1.0 / ((d * 2) as f64).sqrt();
        Ok(BlocoSsd {
            norm1: RmsNorm::novo(d, 1e-5, dev)?,
            porta_a: Linear::novo(d, d, desvio, dev)?,
            proj_x: Linear::sem_vies(d, d, desvio, dev)?,
            proj_o: Linear::sem_vies(d, d, desvio, dev)?,
            norm2: RmsNorm::novo(d, 1e-5, dev)?,
            ffn: Swiglu::novo(d, d_ff, dev)?,
        })
    }

    /// x (b, n, d) → (b, n, d). O scan é aplicado por elemento do lote.
    /// `cache_estado`: estado h anterior (b, d) para streaming incremental.
    pub fn fwd(&self, x: &Tensor, cache_estado: &mut Option<Tensor>) -> Result<Tensor> {
        let b = x.dim(0)?;
        let n = x.dim(1)?;
        let d = x.dim(2)?;
        let xb = self.norm1.fwd(x)?;
        let a = candle_nn::ops::sigmoid(&self.porta_a.fwd(&xb)?)?;
        let xp = self.proj_x.fwd(&xb)?;
        // achata (b·n, d), aplica o scan por lote e recompõe
        let xf = xp.reshape((b * n, d))?.contiguous()?;
        let af = a.reshape((b * n, d))?.contiguous()?;
        // pré-estado: se houver cache, subtrai efeito — simplificação: scan
        // independente por segmento com estado inicial zero (documentado; o
        // estado cruzado entre segmentos é ponto de extensão do streaming).
        let _ = cache_estado;
        let mut saidas = Vec::with_capacity(b);
        for i in 0..b {
            let xi = xf.narrow(0, i * n, n)?;
            let ai = af.narrow(0, i * n, n)?;
            let yi = Tensor::apply_op2(&xi, &ai, ScanDiagonal)?;
            saidas.push(yi);
        }
        let y = Tensor::cat(&saidas, 0)?.reshape((b, n, d))?;
        let y = self.proj_o.fwd(&y)?;
        let x2 = (x + &y)?;
        let f = self.ffn.fwd(&self.norm2.fwd(&x2)?)?;
        (&x2 + &f)
    }

    pub fn vars(&self) -> Vec<Var> {
        [
            self.norm1.vars(), self.porta_a.vars(), self.proj_x.vars(),
            self.proj_o.vars(), self.norm2.vars(), self.ffn.vars(),
        ].concat()
    }
}

#[cfg(test)]
mod testes {
    use super::*;

    #[test]
    fn adjunto_do_scan_vs_finitas_diferencas() -> Result<()> {
        let dev = Device::Cpu;
        let n = 9;
        let d = 5;
        let xv = Var::from_tensor(&Tensor::randn(0f32, 1f32, (n, d), &dev)?)?;
        let av = Var::from_tensor(&candle_nn::ops::sigmoid(&Tensor::randn(0f32, 1f32, (n, d), &dev)?)?)?;
        let pesos = Tensor::randn(0f32, 1f32, (n, d), &dev)?;
        let y = Tensor::apply_op2(xv.as_tensor(), &av.as_tensor(), ScanDiagonal)?;
        let perda = (&y * &pesos)?.sum_all()?;
        let gs = perda.backward()?;
        let gx = gs.get(&xv).ok_or_else(|| candle_core::Error::Msg("sem grad x".into()))?;
        let ga = gs.get(&av).ok_or_else(|| candle_core::Error::Msg("sem grad a".into()))?;
        let gxv = gx.flatten_all()?.to_vec1::<f32>()?;
        let gav = ga.flatten_all()?.to_vec1::<f32>()?;
        // FD
        let eps = 1e-3f32;
        let f = |x: &Tensor, a: &Tensor| -> Result<f64> {
            let y = Tensor::apply_op2(x, a, ScanDiagonal)?;
            Ok((&y * &pesos)?.sum_all()?.to_scalar::<f32>()? as f64)
        };
        let x0 = xv.as_tensor();
        let a0 = av.as_tensor();
        let x_plano = x0.flatten_all()?.to_vec1::<f32>()?;
        let a_plano = a0.flatten_all()?.to_vec1::<f32>()?;
        let mut erro = 0f64;
        for idx in [3usize, 11, 17, 24, 30] {
            let mut xp = x_plano.clone();
            let mut xm = x_plano.clone();
            xp[idx] += eps;
            xm[idx] -= eps;
            let tp = Tensor::from_vec(xp, (n, d), &dev)?;
            let tm = Tensor::from_vec(xm, (n, d), &dev)?;
            let fd = (f(&tp, a0)? - f(&tm, a0)?) / (2.0 * eps as f64);
            erro = erro.max((fd - gxv[idx] as f64).abs());
            // porta a
            let mut ap = a_plano.clone();
            let mut am = a_plano.clone();
            ap[idx] += eps;
            am[idx] -= eps;
            let tap = Tensor::from_vec(ap, (n, d), &dev)?;
            let tam = Tensor::from_vec(am, (n, d), &dev)?;
            let fda = (f(x0, &tap)? - f(x0, &tam)?) / (2.0 * eps as f64);
            erro = erro.max((fda - gav[idx] as f64).abs());
        }
        assert!(erro < 5e-3, "erro adjunto SSD = {erro}");
        Ok(())
    }
}