wai/wai-rs/src/jpegai_pipeline.rs

@@ -12,6 +12,94 @@
 //! **Sigma selection.** [`sigma_index`] closes the scale path: it turns the
 //! hyper-scale decoder's index `i_s` ([`crate::jpegai_synthesis::decode_hyper_scale`])
 //! into the per-position distribution index the mANS `sgm` decoder consumes.
+//!
+//! **Latent reconstruction.** [`reconstruct_yhat`] is the SGMM autoregressive spine:
+//! it turns the dequantized residual back into the synthesis-input latent `y_hat` by
+//! running the four MCM context stages over a 2×2 checkerboard, each stage predicting
+//! a mean from the hyper params plus the already-reconstructed stages.
+
+use crate::jpegai_synthesis::{mcm_stage0, mcm_stage_n, Weights};
+
+/// The 2×2 sub-lattice each MCM stage owns, in reference stage order:
+/// stage 0 → (0,0), stage 1 → (1,1), stage 2 → (0,1), stage 3 → (1,0).
+/// (The reference's `down_shuffle` returns `(k00, k11, k01, k10)`; `up_shuffle` is
+/// its exact inverse. The diagonal pair is decoded first — a checkerboard.)
+const SUBPOS: [(usize, usize); 4] = [(0, 0), (1, 1), (0, 1), (1, 0)];
+
+/// Split `[c, 2h, 2w]` into the four `[c, h, w]` sub-lattices (reference `down_shuffle`).
+fn down_shuffle(y: &[i64], c: usize, h2: usize, w2: usize) -> [Vec<i64>; 4] {
+    let (h, w) = (h2 / 2, w2 / 2);
+    let mut parts = std::array::from_fn(|_| vec![0i64; c * h * w]);
+    for (pi, &(ky, kx)) in SUBPOS.iter().enumerate() {
+        for ch in 0..c {
+            for i in 0..h {
+                for j in 0..w {
+                    parts[pi][(ch * h + i) * w + j] = y[(ch * h2 + 2 * i + ky) * w2 + 2 * j + kx];
+                }
+            }
+        }
+    }
+    parts
+}
+
+/// Interleave the four `[c, h, w]` sub-lattices back into `[c, 2h, 2w]` (`up_shuffle`).
+fn up_shuffle(parts: &[Vec<i64>; 4], c: usize, h: usize, w: usize) -> Vec<i64> {
+    let (h2, w2) = (h * 2, w * 2);
+    let mut out = vec![0i64; c * h2 * w2];
+    for (pi, &(ky, kx)) in SUBPOS.iter().enumerate() {
+        for ch in 0..c {
+            for i in 0..h {
+                for j in 0..w {
+                    out[(ch * h2 + 2 * i + ky) * w2 + 2 * j + kx] = parts[pi][(ch * h + i) * w + j];
+                }
+            }
+        }
+    }
+    out
+}
+
+/// **SGMM autoregressive reconstruction** — rebuild the synthesis-input latent `y_hat`
+/// from the dequantized residual and the hyper params.
+///
+/// `resi` is `[c, 2h, 2w]` (Q16); `psi` is `[4c, h, w]` (Q16, the hyper decoder's
+/// output); `mcm` are the four MCM stages' own weights. Returns `y_hat` as `[c, 2h, 2w]`.
+///
+/// ```text
+/// p        = chunk(psi, 4)                       // per-stage hyper params, ascending
+/// r        = down_shuffle(resi)                  // 2×2 checkerboard sub-lattices
+/// for i in 0..4:
+///     mu_i = MCM_i(p_i, cat(y_0 .. y_{i-1}))     // ctx grows: 0 / c / 2c / 3c channels
+///     y_i  = r_i + mu_i
+/// y_hat    = up_shuffle(y_0 .. y_3)
+/// ```
+///
+/// Stage 0 has no context (the reference seeds `spatial_params` with zeros and its
+/// stage-0 model ignores them). The context concatenation is **ascending** (`y_0`
+/// first) — the reference does `cat(y_hat_list)`, despite the ONNX input being
+/// labelled `y[1-0]`.
+pub fn reconstruct_yhat(resi: &[i64], psi: &[i64], c: usize, h2: usize, w2: usize, mcm: [&Weights; 4]) -> Vec<i64> {
+    let (h, w) = (h2 / 2, w2 / 2);
+    let plane = c * h * w;
+    assert_eq!(resi.len(), c * h2 * w2, "resi must be [c, 2h, 2w]");
+    assert_eq!(psi.len(), 4 * plane, "psi must be [4c, h, w]");
+
+    let r = down_shuffle(resi, c, h2, w2);
+    let mut ctx: Vec<i64> = Vec::with_capacity(3 * plane); // accumulated y_hat, ascending
+    let mut parts: [Vec<i64>; 4] = std::array::from_fn(|_| Vec::new());
+
+    for i in 0..4 {
+        let p = &psi[i * plane..(i + 1) * plane];
+        let mu = if i == 0 {
+            mcm_stage0(p, h, w, mcm[0])
+        } else {
+            mcm_stage_n(p, &ctx, i * c, h, w, mcm[i])
+        };
+        let y_i: Vec<i64> = r[i].iter().zip(&mu).map(|(&a, &b)| a + b).collect();
+        if i < 3 { ctx.extend_from_slice(&y_i); }
+        parts[i] = y_i;
+    }
+    up_shuffle(&parts, c, h, w)
+}
 
 /// Dequantize entropy-decoded latent symbols by the per-position gain map.
 /// `y_q` and `gain_map` are row-major `[channels × positions]`, same length.
@@ -65,6 +153,23 @@ mod tests {
         assert_eq!(r.len(), 4);
     }
 
+    #[test]
+    fn shuffle_round_trips_and_uses_checkerboard_order() {
+        // [1ch, 4x4] ramp
+        let (c, h2, w2) = (1usize, 4usize, 4usize);
+        let y: Vec<i64> = (0..16).collect();
+        let parts = down_shuffle(&y, c, h2, w2);
+        // stage 0 owns (0,0) of each 2x2 block: indices 0, 2, 8, 10
+        assert_eq!(parts[0], vec![0, 2, 8, 10]);
+        // stage 1 owns (1,1): 5, 7, 13, 15  (the diagonal partner — decoded second)
+        assert_eq!(parts[1], vec![5, 7, 13, 15]);
+        // stage 2 owns (0,1): 1, 3, 9, 11 ; stage 3 owns (1,0): 4, 6, 12, 14
+        assert_eq!(parts[2], vec![1, 3, 9, 11]);
+        assert_eq!(parts[3], vec![4, 6, 12, 14]);
+        // up_shuffle is the exact inverse
+        assert_eq!(up_shuffle(&parts, c, h2 / 2, w2 / 2), y);
+    }
+
     #[test]
     fn sigma_index_thresholds_and_clamp() {
         // 1 channel, gain 0, rvs 0 → boundaries land on multiples of 128 offset by 64.

wai/wai-rs/tests/jpegai_reconstruct_calibration.rs

@@ -0,0 +1,63 @@
+//! Calibration for [`wai::jpegai_pipeline::reconstruct_yhat`] — the SGMM
+//! autoregressive spine — on REAL `test256.bin` data: rebuild the synthesis-input
+//! latent `y_hat` from the reference's dequantized residual (`resi_dq_y`) and hyper
+//! params (`psi_y`), and compare against the reference `yhat_y`.
+//!
+//! This closes the last structural gap in the decode chain: entropy → dequant →
+//! **reconstruct** → synthesis. The float reference loop reproduces `yhat_y` to
+//! ~139 dB (float round-off); our Q16 integer MCM stages carry the 16-bit PTQ floor
+//! of the MCM nets themselves (~61–74 dB per stage), so we assert well above that.
+//!
+//! env: WAI_JPEGAI_RECON_FIXTURE = recon_fixture.json {resi_dq_y, psi_y, yhat_y, c, h2, w2}
+//!      WAI_JPEGAI_MCM_WEIGHTS_DIR = dir with mcm_stage0..3_y.json
+#![cfg(feature = "neural_int")]
+
+use wai::jpegai_pipeline::reconstruct_yhat;
+use wai::jpegai_synthesis::Weights;
+
+fn q16(f: &[f64]) -> Vec<i64> { f.iter().map(|&x| (x * 65536.0).round() as i64).collect() }
+
+#[test]
+fn reconstruct_yhat_on_real_reference_data() {
+    let fx = match std::env::var("WAI_JPEGAI_RECON_FIXTURE") {
+        Ok(p) if std::path::Path::new(&p).exists() => p,
+        _ => { eprintln!("skip: set WAI_JPEGAI_RECON_FIXTURE"); return; }
+    };
+    let wdir = match std::env::var("WAI_JPEGAI_MCM_WEIGHTS_DIR") {
+        Ok(p) if std::path::Path::new(&p).exists() => p,
+        _ => { eprintln!("skip: set WAI_JPEGAI_MCM_WEIGHTS_DIR"); return; }
+    };
+    let v: serde_json::Value = serde_json::from_slice(&std::fs::read(&fx).unwrap()).unwrap();
+    let f64v = |k: &str| v[k].as_array().unwrap().iter().map(|x| x.as_f64().unwrap()).collect::<Vec<f64>>();
+    let (c, h2, w2) = (
+        v["c"].as_u64().unwrap() as usize,
+        v["h2"].as_u64().unwrap() as usize,
+        v["w2"].as_u64().unwrap() as usize,
+    );
+
+    let resi = q16(&f64v("resi_dq_y"));
+    let psi = q16(&f64v("psi_y"));
+    let yhat_ref = f64v("yhat_y");
+
+    let load = |i: usize| -> Weights {
+        let raw = std::fs::read(format!("{wdir}/mcm_stage{i}_y.json")).unwrap();
+        Weights::from_json(&serde_json::from_slice(&raw).unwrap())
+    };
+    let (w0, w1, w2w, w3) = (load(0), load(1), load(2), load(3));
+
+    let got = reconstruct_yhat(&resi, &psi, c, h2, w2, [&w0, &w1, &w2w, &w3]);
+    assert_eq!(got.len(), yhat_ref.len(), "y_hat size mismatch");
+
+    let (mut maxd, mut se, mut sig) = (0f64, 0f64, 0f64);
+    for i in 0..got.len() {
+        let g = got[i] as f64 / 65536.0;
+        let r = yhat_ref[i];
+        let d = (g - r).abs();
+        maxd = maxd.max(d);
+        se += d * d;
+        sig += r * r;
+    }
+    let snr = 10.0 * (sig / se).log10();
+    eprintln!("reconstruct_yhat on real data: max|Δ| = {maxd:.5}   SNR = {snr:.2} dB");
+    assert!(snr > 45.0, "SGMM reconstruction regressed (SNR {snr:.2} dB)");
+}