wai/wai-rs/src/jpegai_mans.rs

@@ -0,0 +1,239 @@
+//! Byte-exact **mANS** (masked-ANS) entropy decoder for JPEG-AI — a faithful port
+//! of the ISO reference `cpp_exts/mans/decompressor.cpp`. Two interleaved ANS states,
+//! MSB-first reads from a rolling 64-bit window (`read`/`flush`), sentinel-bit init,
+//! and a group-of-4 inbound→flush→outbound schedule. Pure integer ⇒ byte-identical
+//! on every machine. Backs the entropy stage of `wai.neural.jpeg_ai`.
+//!
+//! `decode_sgm` (main y-latent): per position a distribution `sigma_idx` selects one
+//! of the 32 precomputed transition tables; `mask=false` ⇒ the position is skipped
+//! (value 0) and does not consume the stream. The state threads across calls, so a
+//! single [`MansDecoder`] decodes a whole stream over many `decode_sgm` invocations.
+
+/// Low 56..63-bit masks for `flush()` (BIT_MASK_HIGH in the reference).
+const BIT_MASK_HIGH: [u64; 8] = [
+    0x00ff_ffff_ffff_ffff, 0x01ff_ffff_ffff_ffff, 0x03ff_ffff_ffff_ffff, 0x07ff_ffff_ffff_ffff,
+    0x0fff_ffff_ffff_ffff, 0x1fff_ffff_ffff_ffff, 0x3fff_ffff_ffff_ffff, 0x7fff_ffff_ffff_ffff,
+];
+
+/// The backward-consumed 64-bit bit window + the two interleaved ANS states.
+#[derive(Clone, Copy)]
+struct BitStream {
+    ptr_end: usize, // index into `memory` (one past the current window's low byte block)
+    head: u64,
+    bit_pos: u32,
+    state1: u8,
+    state2: u8,
+}
+
+impl BitStream {
+    fn le64(mem: &[u8], at: usize) -> u64 {
+        let mut b = [0u8; 8];
+        b.copy_from_slice(&mem[at..at + 8]);
+        u64::from_le_bytes(b)
+    }
+    /// Read `step` bits (MSB-first) out of the window.
+    fn read(&mut self, step: u32) -> u64 {
+        self.bit_pos -= step;
+        let result = self.head >> self.bit_pos;
+        self.head ^= result << self.bit_pos;
+        result
+    }
+    /// Refill the window from earlier bytes (moving backward).
+    fn flush(&mut self, mem: &[u8]) {
+        self.ptr_end -= (7 ^ (self.bit_pos >> 3)) as usize;
+        self.head = Self::le64(mem, self.ptr_end) & BIT_MASK_HIGH[(self.bit_pos & 7) as usize];
+        self.bit_pos |= 56;
+    }
+    fn init_backward(&mut self, mem: &[u8], ptr: usize) {
+        self.ptr_end = ptr - 8;
+        self.head = Self::le64(mem, self.ptr_end);
+        self.bit_pos = 63 ^ self.head.leading_zeros();
+        self.head ^= 1u64 << self.bit_pos;
+    }
+    fn init_with_states(&mut self, mem: &[u8], ptr: usize) {
+        self.init_backward(mem, ptr);
+        self.state1 = self.read(8) as u8;
+        self.state2 = self.read(8) as u8;
+        self.flush(mem);
+    }
+}
+
+const MAX_Z: usize = 63;
+
+/// A byte-exact mANS decoder over one memory buffer (single-thread streams).
+pub struct MansDecoder {
+    memory: Vec<u8>,          // 8 zero bytes prepended, then the stream bytes
+    stream: BitStream,        // persists across decode calls (state threads)
+    trans_y: Vec<u32>,        // num_distributions * 256 flat transition table (sgm)
+    bounds_y: Vec<u8>,        // per-distribution outbound bound (sgm)
+    z_pre: [u16; 512],        // transitionsZPreprocess (built once, for the fac path)
+    trans_z: [u32; 256],      // per-row Z transition table (rebuilt per channel)
+}
+
+impl MansDecoder {
+    /// `mem` = the raw stream bytes; `offsets` = per-thread end offsets
+    /// (`cumsum(thread_sizes)`). Only single-thread (`offsets.len()==1`) is supported.
+    pub fn new(mem: &[u8], offsets: &[u32]) -> Self {
+        assert_eq!(offsets.len(), 1, "only single-thread streams supported");
+        let cap = mem.len();
+        let mut memory = vec![0u8; cap + 8];
+        memory[8..8 + cap].copy_from_slice(mem);
+        let thread_off = offsets[0] as usize + 8;
+        let mut stream = BitStream { ptr_end: 0, head: 0, bit_pos: 0, state1: 0, state2: 0 };
+        stream.init_with_states(&memory, thread_off);
+        // transitionsZPreprocess: for i in 511..1, nBits grows so (i<<nBits) reaches [256,512)
+        let mut z_pre = [0u16; 512];
+        let mut nbits: u32 = 0;
+        for i in (1..512usize).rev() {
+            if (i << nbits) < 256 { nbits += 1; }
+            z_pre[i] = (((i << nbits) ^ 256) as u16) | ((nbits as u16) << 8);
+        }
+        MansDecoder { memory, stream, trans_y: Vec::new(), bounds_y: Vec::new(), z_pre, trans_z: [0u32; 256] }
+    }
+
+    /// Install the 32 sgm transition tables (`num_dist × 256` u32) + per-dist bounds.
+    pub fn set_sgm_transitions(&mut self, transitions: &[u32], bounds: &[u8]) {
+        self.trans_y = transitions.to_vec();
+        self.bounds_y = bounds.to_vec();
+    }
+
+    #[inline]
+    fn t_y(&self, dist: usize, state: u8) -> u32 {
+        self.trans_y[dist * 256 + state as usize]
+    }
+
+    #[inline]
+    fn decode_y_inbound(&mut self, which: u8, sigma: &[u8], out: &mut [i16], index: usize) {
+        let transition = self.t_y(sigma[index] as usize, if which == 1 { self.stream.state1 } else { self.stream.state2 });
+        let ns = ((transition >> 16) as u64 | self.stream.read(transition >> 24)) as u8;
+        if which == 1 { self.stream.state1 = ns; } else { self.stream.state2 = ns; }
+        out[index] = transition as i16; // low 16 bits, signed = the symbol
+    }
+
+    #[inline]
+    fn decode_y_outbound(&mut self, sigma: &[u8], out: &mut [i16], index: usize) {
+        let bound = self.bounds_y[sigma[index] as usize] as i32;
+        if out[index] as i32 + bound == 0 {
+            let step = if self.stream.read(1) != 0 { 3 } else { 16 };
+            let value_temp = self.stream.read(step) as i32;
+            let sign = value_temp & 1;
+            out[index] = (((value_temp >> 1) + (bound - sign)) ^ (-sign)) as i16;
+            self.stream.flush(&self.memory);
+        }
+    }
+
+    /// Decode one span: `out[i] = symbol` where `mask[i]`, else left untouched (0).
+    /// State threads across calls. Faithful port of `decodeSingleThread`.
+    pub fn decode_sgm(&mut self, sigma: &[u8], out: &mut [i16], mask: &[bool]) {
+        let length = sigma.len();
+        let mut index = 0;
+        // main groups of 4: state1/state2/state1/state2, flush between in/out bound
+        while index + 3 < length {
+            if mask[index] { self.decode_y_inbound(1, sigma, out, index); }
+            if mask[index + 1] { self.decode_y_inbound(2, sigma, out, index + 1); }
+            if mask[index + 2] { self.decode_y_inbound(1, sigma, out, index + 2); }
+            if mask[index + 3] { self.decode_y_inbound(2, sigma, out, index + 3); }
+            self.stream.flush(&self.memory);
+            if mask[index] { self.decode_y_outbound(sigma, out, index); }
+            if mask[index + 1] { self.decode_y_outbound(sigma, out, index + 1); }
+            if mask[index + 2] { self.decode_y_outbound(sigma, out, index + 2); }
+            if mask[index + 3] { self.decode_y_outbound(sigma, out, index + 3); }
+            index += 4;
+        }
+        if length & 2 != 0 {
+            if mask[index] {
+                self.decode_y_inbound(1, sigma, out, index);
+                self.decode_y_outbound(sigma, out, index);
+            }
+            if mask[index + 1] {
+                self.decode_y_inbound(2, sigma, out, index + 1);
+                self.decode_y_outbound(sigma, out, index + 1);
+            }
+            index += 2;
+            self.stream.flush(&self.memory);
+        }
+        if length & 1 != 0 {
+            if mask[index] {
+                self.decode_y_inbound(1, sigma, out, index);
+                self.decode_y_outbound(sigma, out, index);
+            }
+            self.stream.flush(&self.memory);
+        }
+    }
+
+    // ---- fac (factorized z hyper-latent) path ----
+
+    /// Build the per-row Z transition table from a channel's cumulative `cdfs`
+    /// (`MAX_Z` entries). Faithful port of `setZDistributions`.
+    fn set_z_distributions(&mut self, cdfs: &[u8]) {
+        let mut curr = 0usize;
+        for i in 0..MAX_Z {
+            let c = cdfs[i] as usize;
+            let mut k = c - curr; // = count for symbol i
+            for j in curr..c {
+                self.trans_z[j] = ((self.z_pre[k] as u32) << 16) | i as u32;
+                k += 1;
+            }
+            curr = c;
+        }
+        self.trans_z[255] = ((self.z_pre[1] as u32) << 16) | MAX_Z as u32;
+    }
+
+    #[inline]
+    fn decode_z_inbound(&mut self, which: u8, out: &mut [u8], index: usize) {
+        let transition = self.trans_z[if which == 1 { self.stream.state1 } else { self.stream.state2 } as usize];
+        let ns = ((transition >> 16) as u64 | self.stream.read(transition >> 24)) as u8;
+        if which == 1 { self.stream.state1 = ns; } else { self.stream.state2 = ns; }
+        out[index] = transition as u8; // low 8 bits = symbol (0..=MAX_Z)
+    }
+
+    #[inline]
+    fn decode_z_outbound(&mut self, out: &mut [u8], index: usize) {
+        if out[index] as usize == MAX_Z {
+            out[index] = self.stream.read(6) as u8;
+        }
+    }
+
+    fn decode_row(&mut self, cdfs: &[u8], out: &mut [u8]) {
+        self.set_z_distributions(cdfs);
+        let length = out.len();
+        let mut index = 0;
+        while index + 3 < length {
+            self.decode_z_inbound(1, out, index);
+            self.decode_z_inbound(2, out, index + 1);
+            self.decode_z_inbound(1, out, index + 2);
+            self.decode_z_inbound(2, out, index + 3);
+            self.decode_z_outbound(out, index);
+            self.decode_z_outbound(out, index + 1);
+            self.decode_z_outbound(out, index + 2);
+            self.decode_z_outbound(out, index + 3);
+            self.stream.flush(&self.memory);
+            index += 4;
+        }
+        if length & 2 != 0 {
+            self.decode_z_inbound(1, out, index);
+            self.decode_z_outbound(out, index);
+            self.decode_z_inbound(2, out, index + 1);
+            self.decode_z_outbound(out, index + 1);
+            index += 2;
+        }
+        if length & 1 != 0 {
+            self.decode_z_inbound(1, out, index);
+            self.decode_z_outbound(out, index);
+        }
+        self.stream.flush(&self.memory);
+    }
+
+    /// Decode the factorized z hyper-latent: `channels` rows, `size_per_channel`
+    /// symbols each, each row using its own `MAX_Z`-entry cumulative CDF (`cdfs` is
+    /// `channels × MAX_Z`). Faithful port of `decode_factorize` (single-thread).
+    pub fn decode_factorize(&mut self, cdfs: &[u8], out: &mut [u8], channels: usize, size_per_channel: usize) {
+        for i in 0..channels {
+            let row = &cdfs[i * MAX_Z..(i + 1) * MAX_Z];
+            let (a, b) = (i * size_per_channel, (i + 1) * size_per_channel);
+            let mut buf = vec![0u8; size_per_channel];
+            self.decode_row(row, &mut buf);
+            out[a..b].copy_from_slice(&buf);
+        }
+    }
+}

wai/wai-rs/src/jpegai_synthesis.rs

@@ -1,8 +1,18 @@
-//! Whole-net **byte-exact integer JPEG-AI synthesis** (model_y, tools_0) — threads
-//! the real Restormer topology (traced from the ISO reference ONNX) through the
-//! integer ops in [`crate::int_transformer`] + [`crate::int_hyper_synth`], loading
-//! 16-bit-PTQ'd weights the sink supplies. Pure i64 ⇒ byte-identical on every
-//! machine; backs `wai.neural.jpeg_ai` / `intent = replicate`.
+//! Whole-net **byte-exact integer JPEG-AI synthesis** (model_y + model_uv, tools_0)
+//! — threads the real Restormer topology (traced from the ISO reference ONNX)
+//! through the integer ops in [`crate::int_transformer`] + [`crate::int_hyper_synth`],
+//! loading 16-bit-PTQ'd weights the sink supplies. Pure i64 ⇒ byte-identical on
+//! every machine; backs `wai.neural.jpeg_ai` / `intent = replicate`.
+//!
+//! [`decode_y`] reconstructs the luma plane [1×256×256] from the 160-channel latent;
+//! [`decode_uv`] reconstructs the chroma planes [2×256×256] from the 256-channel
+//! (y‖uv) latent. Both validated vs the float ONNX at ~58 dB PSNR (16-bit PTQ).
+//! [`decode_hyper`] is the entropy-model hyper-prior decoder (common_modules) —
+//! a plain conv chain turning the hyper-latent into the SGMM params (~69.6 dB).
+//! [`decode_hyper_scale`] is the fully-integer scale-index decoder (int symbols →
+//! int scale indices); it is BYTE-EXACT vs the spec reference (max|Δ| = 0), no PTQ.
+//! [`mcm_stage0`]/[`mcm_stage_n`] are the 4-stage multi-stage context model (MCM) —
+//! small conv fusion nets refining the entropy context per decode stage (~61–74 dB).
 //!
 //! Topology (16×16-block trace, 256×256 image): latent[160,16,16]
 //!   first_stage(res-block + ⇑2 + gated) → conv2_t(⇑2) → CAB(dual-branch trunk×mask
@@ -218,6 +228,208 @@ pub fn decode_y(latent: &[i64], wt: &Weights) -> Vec<u8> {
     }).collect()
 }
 
+/// Decode the **U,V chroma planes** [2×256×256] from the 256-channel 16×16 latent
+/// (the concatenated y+uv latent-space features the sink supplies as input `y`).
+/// Pure-integer, deterministic; reproduces the float ONNX chroma synthesis
+/// (model_uv, tools_0).
+///
+/// model_uv is a **structurally distinct** net from model_y (same middle, different
+/// ends):
+///   • front-end: `conv1(256→48)` ‖ `Slice(input, last 96 ch)` → `Concat(144)` →
+///     `conv2_t(144→64, ⇑2)` — no first_stage residual/upsample/gate;
+///   • TAM: the two TABs run **at the full DepthToSpace resolution** — no ds/us conv;
+///   • tail: `conv4_t(64→8, ⇑2)` → `DepthToSpace(CRD,2)` → denorm → **2-channel UV**.
+/// Widths are 64 (vs model_y's 128). CAB / iact gating / TAB / denorm are identical.
+pub fn decode_uv(latent: &[i64], wt: &Weights) -> Vec<u8> {
+    let eps = (1e-5 * (1u128 << 32) as f64) as u128; // LN eps = 1e-5
+    let (h0, w0) = (16usize, 16usize);
+    let cin = 256usize;
+    let hw = h0 * w0;
+    // first_stage: conv1 (256→48); Concat with Slice(input, starts=-96, axis=1) —
+    // the last 96 channels of the input — → 144 ch; conv2_t (144→64, ⇑2)
+    let (c1, _, _) = wt.conv("first_stage.conv1", latent, cin, h0, w0, 3, 1, 1, 1);
+    let mut cat = Vec::with_capacity((48 + 96) * hw);
+    cat.extend_from_slice(&c1); // 48 ch (each hw contiguous)
+    cat.extend_from_slice(&latent[(cin - 96) * hw..]); // last 96 ch of the input
+    let (xc, mut h, mut w) = wt.convt("conv2_t", &cat, 144, h0, w0, 3, 2, 1, 1); // →[64,32,32]
+    let c = 64usize;
+    // CAB: trunk (2 res blocks) × sigmoid(mask branch: subscale ⇓2 → 2 res → upscale ⇑2)
+    //   + conv2_t residual — identical shape to model_y, width 64
+    let trunk = {
+        let t0 = res_block("CAB.residual_trunk.0", &xc, c, h, w, wt);
+        res_block("CAB.residual_trunk.1", &t0, c, h, w, wt)
+    };
+    let mask = {
+        let (s0, sh, sw) = wt.conv("CAB.subscale", &xc, c, h, w, 3, 2, 1, 1); // ⇓2
+        let m2 = res_block("CAB.residual_mask2.0", &s0, c, sh, sw, wt);
+        let m3 = res_block("CAB.residual_mask3.0", &m2, c, sh, sw, wt);
+        let (u, _, _) = wt.convt("CAB.upscale", &m3, c, sh, sw, 3, 2, 1, 1); // ⇑2
+        u.iter().map(|&v| it::sigmoid(v)).collect::<Vec<_>>()
+    };
+    let mut x = eadd(&xc, &emul(&trunk, &mask));
+    // iact2 gated: ReLU6 feeds the conv path, gate uses RAW x
+    let xcl = relu6_of(&x);
+    let (g, _, _) = wt.conv("iact2.conv", &xcl, c, h, w, 3, 1, 1, 4);
+    let (g, _, _) = wt.conv("iact2.conv2", &g, c, h, w, 1, 1, 0, 1);
+    x = gate(&x, &g);
+    // conv3_t (1×1) 64→256; DepthToSpace ⇑2 → [64,64,64]
+    let (x3, _, _) = wt.conv("conv3_t", &x, c, h, w, 1, 1, 0, 1);
+    let (xd, cd, nh, nw) = depth_to_space_crd(&x3, 256, h, w, 2);
+    (h, w) = (nh, nw);
+    // TAM: TAB0 → TAB1 directly at full resolution (model_uv has no ds/us conv)
+    let mut xt = tab("TAM.TABs.0", &xd, cd, 4, h, w, wt, eps);
+    xt = tab("TAM.TABs.1", &xt, cd, 4, h, w, wt, eps);
+    // iact3 gated: ReLU6 feeds the conv path, gate uses RAW xt
+    let xcl = relu6_of(&xt);
+    let (g, _, _) = wt.conv("iact3.conv", &xcl, c, h, w, 3, 1, 1, 4);
+    let (g, _, _) = wt.conv("iact3.conv2", &g, c, h, w, 1, 1, 0, 1);
+    x = gate(&xt, &g);
+    // conv4_t ⇑2 64→8 → [8,128,128]; DepthToSpace ⇑2 → [2,256,256]
+    let (y, yh, yw) = wt.convt("conv4_t", &x, c, h, w, 3, 2, 1, 1);
+    let (yd, cyd, dh, dw) = depth_to_space_crd(&y, 8, yh, yw, 2);
+    debug_assert_eq!((cyd, dh, dw), (2, 256, 256));
+    // output denormalize: clip(·128 + 127.5, 0, 255) → u8 (2 channels)
+    let half = S / 2;
+    yd.iter().map(|&v| {
+        let d = (v * 128 + 127 * S + half).clamp(0, 255 * S);
+        ((d + half) >> F).clamp(0, 255) as u8
+    }).collect()
+}
+
+/// **hyper_decoder** (common_modules) — the entropy-model hyper-prior decoder that
+/// turns the hyper-latent `z` [160×h×w] (Q16) into the 640-channel SGMM entropy
+/// parameters `p` [640×2h×2w] (Q16). A plain conv chain (no transformer):
+/// `conv1(1×1) → conv2_t(4×4, ⇑2) → ReLU6 → conv3(3×3) → ReLU6 → conv4(→640)`.
+/// The output stays in fixed-point (it feeds the arithmetic decoder, not the
+/// display), so there is no pixel denorm. Pure i64 ⇒ byte-identical everywhere.
+pub fn decode_hyper(z: &[i64], cin: usize, h: usize, w: usize, wt: &Weights) -> (Vec<i64>, usize, usize) {
+    let (x, _, _) = wt.conv("conv1", z, cin, h, w, 1, 1, 0, 1); // 160→160, 1×1 (no bias)
+    let (x, h, w) = wt.convt("conv2", &x, 160, h, w, 4, 2, 1, 0); // 160→160, 4×4 ⇑2
+    let x = relu6_of(&x);
+    let (x, _, _) = wt.conv("conv3", &x, 160, h, w, 3, 1, 1, 1); // 160→160, 3×3
+    let x = relu6_of(&x);
+    wt.conv("conv4", &x, 160, h, w, 3, 1, 1, 1) // 160→640, 3×3 → p[640,2h,2w]
+}
+
+/// A fully-integer weight store (int32 conv weights/biases + integer per-channel
+/// power-of-two exponents) for [`decode_hyper_scale`]. Unlike [`Weights`] there is
+/// NO quantization — the hyper_scale_decoder net is exact integer arithmetic, so
+/// the weights are loaded verbatim as i64.
+pub struct IntNet {
+    t: HashMap<String, (Vec<i64>, Vec<usize>)>,
+}
+impl IntNet {
+    pub fn from_json(v: &serde_json::Value) -> Self {
+        let mut t = HashMap::new();
+        for (k, e) in v.as_object().expect("intnet root object") {
+            let vals = e["v"].as_array().unwrap().iter().map(|x| x.as_i64().unwrap()).collect();
+            let dims = e["dims"].as_array().unwrap().iter().map(|x| x.as_u64().unwrap() as usize).collect();
+            t.insert(k.clone(), (vals, dims));
+        }
+        IntNet { t }
+    }
+    fn get(&self, n: &str) -> &(Vec<i64>, Vec<usize>) {
+        self.t.get(n).unwrap_or_else(|| panic!("missing int tensor {n}"))
+    }
+}
+
+/// Exact integer conv, groups=1, i64 accumulation (no requant): output channel `o`,
+/// pixel `p` = `bias[o] + Σ_i Σ_ky Σ_kx w[o,i,ky,kx]·x[i, p+ky-pad, p+kx-pad]`.
+fn iconv(x: &[i64], cin: usize, h: usize, w: usize, wt: &[i64], b: &[i64], cout: usize, k: usize, stride: usize, pad: usize) -> (Vec<i64>, usize, usize) {
+    let ho = (h + 2 * pad - k) / stride + 1;
+    let wo = (w + 2 * pad - k) / stride + 1;
+    let mut out = vec![0i64; cout * ho * wo];
+    for o in 0..cout {
+        let bias = b.get(o).copied().unwrap_or(0);
+        for oy in 0..ho {
+            for ox in 0..wo {
+                let mut acc = bias;
+                for i in 0..cin {
+                    for ky in 0..k {
+                        let iy = oy * stride + ky;
+                        if iy < pad || iy - pad >= h { continue; }
+                        let sy = iy - pad;
+                        for kx in 0..k {
+                            let ix = ox * stride + kx;
+                            if ix < pad || ix - pad >= w { continue; }
+                            let sx = ix - pad;
+                            let wv = wt[((o * cin + i) * k + ky) * k + kx];
+                            acc += wv * x[(i * h + sy) * w + sx];
+                        }
+                    }
+                }
+                out[(o * ho + oy) * wo + ox] = acc;
+            }
+        }
+    }
+    (out, ho, wo)
+}
+
+/// **hyper_scale_decoder** (common_modules) — the scale-index hyper decoder. A
+/// FULLY-INTEGER net: int symbols `z`[160×h×w] → int scale indices `I_s`[160×4h×4w].
+/// Three QDQ-style blocks (`clip(−128,127) → intconv → ÷2^n(per-channel, trunc) →
+/// ReLU`) over conv1(1×1) / depthwise(3×3) / pointwise(160→2560, 1×1), then
+/// `DepthToSpace(CRD, 4)` → `abs` → `clip(0, 3967)`. Byte-exact (no PTQ — the ops
+/// are integer end to end); division truncates toward zero (NOT arithmetic shift).
+pub fn decode_hyper_scale(z: &[i64], cin: usize, h: usize, w: usize, net: &IntNet) -> (Vec<i64>, usize, usize, usize) {
+    // one QDQ block: clip → intconv → per-channel ÷2^exp (trunc toward 0) → [relu]
+    let block = |x: &[i64], ci: usize, cname: &str, k: usize, pad: usize, ename: &str, relu_after: bool| -> (Vec<i64>, usize) {
+        let xq: Vec<i64> = x.iter().map(|&v| it::clip(v, -128, 127)).collect();
+        let (cw, cd) = net.get(&format!("{cname}.weight"));
+        let (cb, _) = net.get(&format!("{cname}.bias"));
+        let cout = cd[0];
+        let (mut y, _, _) = iconv(&xq, ci, h, w, cw, cb, cout, k, 1, pad);
+        let (ev, _) = net.get(ename); // per-out-channel exponent
+        let plane = h * w;
+        for o in 0..cout {
+            let div = 1i64 << ev[o]; // 2^n
+            for p in 0..plane {
+                let idx = o * plane + p;
+                let mut r = y[idx] / div; // integer division, truncates toward zero
+                if relu_after && r < 0 { r = 0; }
+                y[idx] = r;
+            }
+        }
+        (y, cout)
+    };
+    let (x1, _) = block(z, cin, "conv1", 1, 0, "onnx::Pow_85", true);
+    let (x2, _) = block(&x1, 160, "depthwise", 3, 1, "onnx::Pow_91", true);
+    let (x3, c3) = block(&x2, 160, "pointwise", 1, 0, "onnx::Pow_97", false);
+    // DepthToSpace CRD block 4: 2560 → 160 ch, spatial ×4
+    let (ds, co, ho, wo) = depth_to_space_crd(&x3, c3, h, w, 4);
+    // I_s = clip(abs(·), 0, 3967)
+    let is = ds.iter().map(|&v| v.abs().clamp(0, 3967)).collect();
+    (is, co, ho, wo)
+}
+
+/// **MCM stage 0** (common_modules/MCM) — the multi-stage context model's first
+/// stage: `p0 → conv1 → ReLU → conv2 → ReLU → conv3 → y[0]` (bias-free 1×1 convs,
+/// width 160). Q16 fixed-point; `wt` = this stage's own PTQ weights.
+pub fn mcm_stage0(p0: &[i64], h: usize, w: usize, wt: &Weights) -> Vec<i64> {
+    let (x, _, _) = wt.conv("fusion_pred_net.conv1", p0, 160, h, w, 1, 1, 0, 1);
+    let x = relu(&x);
+    let (x, _, _) = wt.conv("fusion_pred_net.conv2", &x, 160, h, w, 1, 1, 0, 1);
+    let x = relu(&x);
+    wt.conv("fusion_pred_net.conv3", &x, 160, h, w, 1, 1, 0, 1).0
+}
+
+/// **MCM stage i≥1** — process the accumulated context then fuse with prior `p`:
+/// `cy = conv.1(conv.0(ctx))` (conv.0 1×1 `ctx_ch→160`, conv.1 3×3 grouped-5 +bias),
+/// then `Concat(cy, p)` [→320] → `fusion.conv1 → ReLU → conv2 → ReLU → conv3 → y[i]`.
+/// `ctx` is the caller-assembled context `concat(y[i-1..0])` (160/320/480 ch).
+pub fn mcm_stage_n(p: &[i64], ctx: &[i64], ctx_ch: usize, h: usize, w: usize, wt: &Weights) -> Vec<i64> {
+    let (c0, _, _) = wt.conv("conv.0", ctx, ctx_ch, h, w, 1, 1, 0, 1); // ctx_ch→160, 1×1
+    let (cy, _, _) = wt.conv("conv.1", &c0, 160, h, w, 3, 1, 1, 5);    // 3×3 grouped-5, +bias
+    let mut cat = Vec::with_capacity(320 * h * w);
+    cat.extend_from_slice(&cy); // 160 ch
+    cat.extend_from_slice(p);   // 160 ch → 320
+    let (x, _, _) = wt.conv("fusion_pred_net.conv1", &cat, 320, h, w, 1, 1, 0, 1); // 320→160
+    let x = relu(&x);
+    let (x, _, _) = wt.conv("fusion_pred_net.conv2", &x, 160, h, w, 1, 1, 0, 1);
+    let x = relu(&x);
+    wt.conv("fusion_pred_net.conv3", &x, 160, h, w, 1, 1, 0, 1).0
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -242,4 +454,99 @@ mod tests {
         let (mn, mx) = (*y.iter().min().unwrap(), *y.iter().max().unwrap());
         eprintln!("decode_y: {} px, value range [{}, {}]", y.len(), mn, mx);
     }
+
+    /// Structural: with the real 16-bit-PTQ model_uv weights, the chroma net threads
+    /// end to end, is deterministic, and emits a 2×256×256 UV output. (PSNR-vs-float
+    /// is `tests/jpegai_uv_calibration.rs`, which measures ~58.8 dB.)
+    #[test]
+    fn decode_uv_runs_deterministically_on_real_weights() {
+        let path = match std::env::var("WAI_JPEGAI_UV_WEIGHTS") {
+            Ok(p) => p,
+            Err(_) => { eprintln!("skip: set WAI_JPEGAI_UV_WEIGHTS to the exported weights JSON"); return; }
+        };
+        if !std::path::Path::new(&path).exists() { eprintln!("skip: weights not present at {path}"); return; }
+        let v: serde_json::Value = serde_json::from_slice(&std::fs::read(&path).unwrap()).unwrap();
+        let wt = Weights::from_json(&v);
+        // deterministic synthetic 256-ch latent [256,16,16] in Q16 (a ramp)
+        let latent: Vec<i64> = (0..256 * 16 * 16).map(|i| ((i as i64 % 97) - 48) * (S / 64)).collect();
+        let uv = decode_uv(&latent, &wt);
+        assert_eq!(uv.len(), 2 * 256 * 256, "UV output must be 2×256×256");
+        assert_eq!(uv, decode_uv(&latent, &wt), "decode_uv must be deterministic");
+        let (mn, mx) = (*uv.iter().min().unwrap(), *uv.iter().max().unwrap());
+        eprintln!("decode_uv: {} samples (2ch), value range [{}, {}]", uv.len(), mn, mx);
+    }
+
+    /// Structural: with the real 16-bit-PTQ hyper_decoder weights, the entropy
+    /// hyper-prior net threads end to end, is deterministic, and emits the
+    /// 640-channel params at 2× the hyper-latent resolution. (SNR-vs-float is
+    /// `tests/jpegai_hyper_calibration.rs`, which measures ~69.6 dB.)
+    #[test]
+    fn decode_hyper_runs_deterministically_on_real_weights() {
+        let path = match std::env::var("WAI_JPEGAI_HYPERDEC_WEIGHTS") {
+            Ok(p) => p,
+            Err(_) => { eprintln!("skip: set WAI_JPEGAI_HYPERDEC_WEIGHTS to the exported weights JSON"); return; }
+        };
+        if !std::path::Path::new(&path).exists() { eprintln!("skip: weights not present at {path}"); return; }
+        let v: serde_json::Value = serde_json::from_slice(&std::fs::read(&path).unwrap()).unwrap();
+        let wt = Weights::from_json(&v);
+        let (c, h, w) = (160usize, 8usize, 8usize);
+        // integer-valued hyper-latent (entropy-decoded symbols) in Q16
+        let z: Vec<i64> = (0..c * h * w).map(|i| ((i as i64 % 11) - 5) * S).collect();
+        let (p, ph, pw) = decode_hyper(&z, c, h, w, &wt);
+        assert_eq!((ph, pw), (16, 16), "hyper output must be 2× (⇑2)");
+        assert_eq!(p.len(), 640 * 16 * 16, "hyper output must be 640×16×16");
+        assert_eq!(decode_hyper(&z, c, h, w, &wt).0, p, "decode_hyper must be deterministic");
+        eprintln!("decode_hyper: {} params at {}×{}", p.len(), ph, pw);
+    }
+
+    /// Structural: with the real integer hyper_scale_decoder weights, the fully-integer
+    /// scale-index net threads end to end, is deterministic, and emits `160×4h×4w` scale
+    /// indices in [0, 3967]. (Byte-exact-vs-reference is
+    /// `tests/jpegai_hyperscale_calibration.rs`, which asserts max|Δ| = 0.)
+    #[test]
+    fn decode_hyper_scale_runs_deterministically_on_real_weights() {
+        let path = match std::env::var("WAI_JPEGAI_HYPERSCALE_WEIGHTS") {
+            Ok(p) => p,
+            Err(_) => { eprintln!("skip: set WAI_JPEGAI_HYPERSCALE_WEIGHTS to the integer weights JSON"); return; }
+        };
+        if !std::path::Path::new(&path).exists() { eprintln!("skip: weights not present at {path}"); return; }
+        let v: serde_json::Value = serde_json::from_slice(&std::fs::read(&path).unwrap()).unwrap();
+        let net = IntNet::from_json(&v);
+        let (c, h, w) = (160usize, 8usize, 8usize);
+        let z: Vec<i64> = (0..c * h * w).map(|i| ((i as i64 * 37) % 291) - 140).collect();
+        let (is, co, ho, wo) = decode_hyper_scale(&z, c, h, w, &net);
+        assert_eq!((co, ho, wo), (160, 32, 32), "hyper_scale output must be 160×4h×4w");
+        assert_eq!(is.len(), 160 * 32 * 32);
+        assert!(is.iter().all(|&v| (0..=3967).contains(&v)), "I_s must be clipped to [0,3967]");
+        assert_eq!(decode_hyper_scale(&z, c, h, w, &net).0, is, "decode_hyper_scale must be deterministic");
+        eprintln!("decode_hyper_scale: {} scale indices at {co}×{ho}×{wo}", is.len());
+    }
+
+    /// Structural: the 4 MCM context stages thread on the real per-stage weights,
+    /// are deterministic, and emit 160-channel context features at the input
+    /// resolution. (SNR-vs-float is `tests/jpegai_mcm_calibration.rs`, ~61–74 dB.)
+    #[test]
+    fn mcm_stages_run_deterministically_on_real_weights() {
+        let dir = match std::env::var("WAI_JPEGAI_MCM_WEIGHTS_DIR") {
+            Ok(p) => p,
+            Err(_) => { eprintln!("skip: set WAI_JPEGAI_MCM_WEIGHTS_DIR to the MCM weights dir"); return; }
+        };
+        let load = |i: usize| {
+            let p = format!("{dir}/mcm_stage{i}_y.json");
+            if !std::path::Path::new(&p).exists() { return None; }
+            let v: serde_json::Value = serde_json::from_slice(&std::fs::read(&p).unwrap()).unwrap();
+            Some(Weights::from_json(&v))
+        };
+        let (Some(w0), Some(w1)) = (load(0), load(1)) else { eprintln!("skip: MCM weights not present"); return; };
+        let (h, w) = (16usize, 16usize);
+        let p: Vec<i64> = (0..160 * h * w).map(|i| ((i as i64 % 71) - 35) * (S / 64)).collect();
+        let y0 = mcm_stage0(&p, h, w, &w0);
+        assert_eq!(y0.len(), 160 * h * w);
+        assert_eq!(mcm_stage0(&p, h, w, &w0), y0, "mcm_stage0 must be deterministic");
+        // stage1: context = y0 (160 ch)
+        let y1 = mcm_stage_n(&p, &y0, 160, h, w, &w1);
+        assert_eq!(y1.len(), 160 * h * w);
+        assert_eq!(mcm_stage_n(&p, &y0, 160, h, w, &w1), y1, "mcm_stage_n must be deterministic");
+        eprintln!("MCM: y0 {} + y1 {} context features (160ch @ {h}×{w})", y0.len(), y1.len());
+    }
 }

wai/wai-rs/src/lib.rs

@@ -73,6 +73,12 @@ pub mod int_transformer;
 #[cfg(feature = "neural_int")]
 pub mod jpegai_synthesis;
 
+// Byte-exact mANS entropy decoder for JPEG-AI (port of the ISO reference
+// decompressor.cpp) — two interleaved ANS states, sgm (y-latent) + fac (z-latent)
+// paths. Pure integer ⇒ byte-identical everywhere.
+#[cfg(feature = "neural_int")]
+pub mod jpegai_mans;
+
 // Deterministic integer generative synthesis (`wai.generate.*`) — the
 // `intent=create` analog of the integer-exact codecs: a generator whose entire
 // pipeline is integer/fixed-point, pinned to (model, prompt, seed, steps), so the

wai/wai-rs/tests/jpegai_hyper_calibration.rs

@@ -0,0 +1,66 @@
+//! Calibration: integer `jpegai_synthesis::decode_hyper` vs the float reference —
+//! run the SAME hyper-latent `z` through `hyper_decoder.onnx` (ort) and through
+//! `decode_hyper`, compare the 640-channel entropy params `p`. The output is not
+//! pixels but fixed-point SGMM params, so we use a scale-invariant SNR
+//! (20·log10(rms(signal)/rms(error))) plus max abs error. Needs `neural` (ort).
+//!
+//! env: WAI_JPEGAI_HYPERDEC = path to common_modules/hyper_decoder.onnx,
+//!      WAI_JPEGAI_HYPERDEC_WEIGHTS = path to the exported 16-bit weights JSON.
+#![cfg(feature = "neural")]
+
+use ort::session::{builder::GraphOptimizationLevel, Session};
+use ort::value::Value;
+use wai::jpegai_synthesis::{decode_hyper, Weights};
+
+fn ort_run(path: &str, x: Vec<f32>, c: usize, h: usize, w: usize) -> (Vec<f32>, Vec<usize>) {
+    let mut s = Session::builder().unwrap()
+        .with_intra_threads(1).unwrap().with_inter_threads(1).unwrap()
+        .with_optimization_level(GraphOptimizationLevel::Level3).unwrap()
+        .commit_from_file(path).unwrap();
+    let name = s.inputs().first().unwrap().name().to_string();
+    let arr = ndarray::Array4::from_shape_vec((1, c, h, w), x).unwrap();
+    let v = Value::from_array(arr).unwrap();
+    let out = s.run(ort::inputs![name.as_str() => v]).unwrap();
+    let view = out[0].try_extract_array::<f32>().unwrap();
+    (view.iter().copied().collect(), view.shape().to_vec())
+}
+
+#[test]
+fn hyper_decoder_int_vs_float_snr() {
+    let onnx = match std::env::var("WAI_JPEGAI_HYPERDEC") {
+        Ok(p) if std::path::Path::new(&p).exists() => p,
+        _ => { eprintln!("skip: set WAI_JPEGAI_HYPERDEC to hyper_decoder.onnx"); return; }
+    };
+    let wjson = match std::env::var("WAI_JPEGAI_HYPERDEC_WEIGHTS") {
+        Ok(p) if std::path::Path::new(&p).exists() => p,
+        _ => { eprintln!("skip: set WAI_JPEGAI_HYPERDEC_WEIGHTS to the weights JSON"); return; }
+    };
+    let (c, h, w) = (160usize, 8usize, 8usize);
+    let n = c * h * w;
+    // integer-valued hyper-latent (entropy-decoded symbols), small magnitude
+    let lf: Vec<f32> = (0..n).map(|i| ((i % 11) as f32) - 5.0).collect();
+
+    let (pf, shape) = ort_run(&onnx, lf.clone(), c, h, w);
+    let (fmn, fmx) = pf.iter().fold((f32::MAX, f32::MIN), |(a, b), &v| (a.min(v), b.max(v)));
+    eprintln!("float hyper out shape = {:?}, len {}, range [{:.3}, {:.3}]", shape, pf.len(), fmn, fmx);
+
+    let li: Vec<i64> = lf.iter().map(|&x| (x * 65536.0).round() as i64).collect();
+    let v: serde_json::Value = serde_json::from_slice(&std::fs::read(&wjson).unwrap()).unwrap();
+    let wt = Weights::from_json(&v);
+    let (pi, ph, pw) = decode_hyper(&li, c, h, w, &wt);
+    eprintln!("int hyper out = {} vals at [{},{}]", pi.len(), ph, pw);
+
+    assert_eq!(pi.len(), pf.len(), "output size mismatch (640×2h×2w expected)");
+    let (mut maxd, mut se, mut sig) = (0f64, 0f64, 0f64);
+    for i in 0..pi.len() {
+        let f = pf[i] as f64;
+        let q = pi[i] as f64 / 65536.0; // dequantize Q16
+        let d = (q - f).abs();
+        maxd = maxd.max(d);
+        se += d * d;
+        sig += f * f;
+    }
+    let snr = 10.0 * (sig / se).log10();
+    eprintln!("INT vs FLOAT hyper_decoder:  max|Δ| = {:.4}   SNR = {:.2} dB", maxd, snr);
+    assert!(snr > 45.0, "hyper_decoder SNR {snr:.2} dB too low (expected high — pure conv chain)");
+}

wai/wai-rs/tests/jpegai_hyperscale_calibration.rs

@@ -0,0 +1,58 @@
+//! Calibration: integer `jpegai_synthesis::decode_hyper_scale` vs the spec-compliant
+//! reference. The hyper_scale_decoder net is FULLY INTEGER (int64 symbols in → int32
+//! scale indices out) — but ONNX Runtime rejects it (`Clip` has no int64 kernel), so
+//! the reference is produced by the pure-Python `onnx.reference.ReferenceEvaluator`
+//! and saved to a JSON vector; this test replays the SAME deterministic `z` through
+//! `decode_hyper_scale` and asserts BYTE-EXACT equality (max|Δ| = 0). No ort needed.
+//!
+//! env: WAI_JPEGAI_HYPERSCALE_WEIGHTS = exported integer weights JSON,
+//!      WAI_JPEGAI_HYPERSCALE_REF     = reference-output JSON {shape, out:[i64]}.
+//!
+//! Generate the reference vector (matching this test's deterministic `z`) with the
+//! pure-Python ONNX reference evaluator:
+//!   python -c "import numpy as np,json; from onnx.reference import ReferenceEvaluator; \
+//!     n=160*8*8; z=np.array([((i*37)%291)-140 for i in range(n)],dtype=np.int64).reshape(1,160,8,8); \
+//!     s=ReferenceEvaluator('hyper_scale_decoder.onnx'); o=s.run(None,{s.input_names[0]:z})[0]; \
+//!     json.dump({'shape':list(o.shape),'out':o.astype('int64').flatten().tolist()}, open('ref.json','w'))"
+#![cfg(feature = "neural_int")]
+
+use wai::jpegai_synthesis::{decode_hyper_scale, IntNet};
+
+#[test]
+fn hyper_scale_decoder_byte_exact_vs_reference() {
+    let wjson = match std::env::var("WAI_JPEGAI_HYPERSCALE_WEIGHTS") {
+        Ok(p) if std::path::Path::new(&p).exists() => p,
+        _ => { eprintln!("skip: set WAI_JPEGAI_HYPERSCALE_WEIGHTS to the weights JSON"); return; }
+    };
+    let rjson = match std::env::var("WAI_JPEGAI_HYPERSCALE_REF") {
+        Ok(p) if std::path::Path::new(&p).exists() => p,
+        _ => { eprintln!("skip: set WAI_JPEGAI_HYPERSCALE_REF to the reference JSON"); return; }
+    };
+    let (c, h, w) = (160usize, 8usize, 8usize);
+    let n = c * h * w;
+    // SAME deterministic int symbols as the reference generator
+    let zi: Vec<i64> = (0..n).map(|i| ((i as i64 * 37) % 291) - 140).collect();
+
+    let rv: serde_json::Value = serde_json::from_slice(&std::fs::read(&rjson).unwrap()).unwrap();
+    let refv: Vec<i64> = rv["out"].as_array().unwrap().iter().map(|x| x.as_i64().unwrap()).collect();
+    let shape: Vec<usize> = rv["shape"].as_array().unwrap().iter().map(|x| x.as_u64().unwrap() as usize).collect();
+    eprintln!("reference shape = {shape:?}, {} values", refv.len());
+
+    let wv: serde_json::Value = serde_json::from_slice(&std::fs::read(&wjson).unwrap()).unwrap();
+    let net = IntNet::from_json(&wv);
+    let (mine, co, ho, wo) = decode_hyper_scale(&zi, c, h, w, &net);
+    eprintln!("int hyper_scale out = {} vals at {co}×{ho}×{wo}", mine.len());
+
+    assert_eq!(mine.len(), refv.len(), "output size mismatch (160×32×32 expected)");
+    let mut maxd = 0i64;
+    let mut mism = 0usize;
+    let mut first = None;
+    for i in 0..mine.len() {
+        let d = (mine[i] - refv[i]).abs();
+        if d != 0 { mism += 1; if first.is_none() { first = Some((i, mine[i], refv[i])); } }
+        maxd = maxd.max(d);
+    }
+    eprintln!("INT vs REF hyper_scale_decoder:  max|Δ| = {maxd}   mismatched = {mism}/{}", mine.len());
+    if let Some((i, m, r)) = first { eprintln!("first mismatch @ {i}: mine={m} ref={r}"); }
+    assert_eq!(maxd, 0, "hyper_scale_decoder must be BYTE-EXACT (fully-integer net); {mism} mismatches");
+}

wai/wai-rs/tests/jpegai_mans_calibration.rs

@@ -0,0 +1,76 @@
+//! Byte-exact calibration for the mANS entropy decoder: replay the real JPEG-AI
+//! `SORP` stream (main y-latent) through `MansDecoder::decode_sgm` and require an
+//! EXACT match to the reference `expected` (the golden was produced by the ISO
+//! reference `ans.so`, and this Rust is a faithful port of `decompressor.cpp`).
+//! No ort needed — the fixture carries the transition tables + stream + goldens.
+//!
+//! env: WAI_JPEGAI_MANS_FIXTURE = path to mans_sorp_fixture.json
+//!   { dt:[u32; 32*256], bounds:[u8;32], mem:[u8], offsets:[u32],
+//!     calls:[{sigma:[u8], mask:[bool], expected:[i16]}] }
+#![cfg(feature = "neural_int")]
+
+use wai::jpegai_mans::MansDecoder;
+
+#[test]
+fn mans_sgm_byte_exact_vs_reference() {
+    let path = match std::env::var("WAI_JPEGAI_MANS_FIXTURE") {
+        Ok(p) if std::path::Path::new(&p).exists() => p,
+        _ => { eprintln!("skip: set WAI_JPEGAI_MANS_FIXTURE to mans_sorp_fixture.json"); return; }
+    };
+    let v: serde_json::Value = serde_json::from_slice(&std::fs::read(&path).unwrap()).unwrap();
+    let arr_u32 = |k: &str| v[k].as_array().unwrap().iter().map(|x| x.as_u64().unwrap() as u32).collect::<Vec<_>>();
+    let dt = arr_u32("dt");
+    let bounds: Vec<u8> = v["bounds"].as_array().unwrap().iter().map(|x| x.as_u64().unwrap() as u8).collect();
+    let mem: Vec<u8> = v["mem"].as_array().unwrap().iter().map(|x| x.as_u64().unwrap() as u8).collect();
+    let offsets = arr_u32("offsets");
+
+    let mut dec = MansDecoder::new(&mem, &offsets);
+    dec.set_sgm_transitions(&dt, &bounds);
+
+    let mut total = 0usize;
+    let mut mism = 0usize;
+    let mut first: Option<(usize, usize, i16, i16)> = None;
+    for (ci, call) in v["calls"].as_array().unwrap().iter().enumerate() {
+        let sigma: Vec<u8> = call["sigma"].as_array().unwrap().iter().map(|x| x.as_u64().unwrap() as u8).collect();
+        let mask: Vec<bool> = call["mask"].as_array().unwrap().iter().map(|x| x.as_bool().unwrap()).collect();
+        let exp: Vec<i16> = call["expected"].as_array().unwrap().iter().map(|x| x.as_i64().unwrap() as i16).collect();
+        let mut out = vec![0i16; sigma.len()];
+        dec.decode_sgm(&sigma, &mut out, &mask);
+        for i in 0..out.len() {
+            total += 1;
+            if out[i] != exp[i] {
+                mism += 1;
+                if first.is_none() { first = Some((ci, i, out[i], exp[i])); }
+            }
+        }
+    }
+    eprintln!("mANS decode_sgm vs reference: {}/{} exact, {} mismatched", total - mism, total, mism);
+    if let Some((ci, i, got, exp)) = first { eprintln!("first mismatch: call {ci} pos {i}: got {got} exp {exp}"); }
+    assert_eq!(mism, 0, "mANS decoder must be BYTE-EXACT vs the ISO reference; {mism} mismatches");
+}
+
+/// Byte-exact `fac` (factorized z hyper-latent) path vs an oracle roundtrip vector.
+/// env: WAI_JPEGAI_MANS_FAC = path to mans_fac_fixture.json
+///   { cdfs:[u8; channels*63], mem:[u8], offsets:[u32], channels, spc, expected:[u8] }
+#[test]
+fn mans_fac_byte_exact_vs_reference() {
+    let path = match std::env::var("WAI_JPEGAI_MANS_FAC") {
+        Ok(p) if std::path::Path::new(&p).exists() => p,
+        _ => { eprintln!("skip: set WAI_JPEGAI_MANS_FAC to mans_fac_fixture.json"); return; }
+    };
+    let v: serde_json::Value = serde_json::from_slice(&std::fs::read(&path).unwrap()).unwrap();
+    let u8v = |k: &str| v[k].as_array().unwrap().iter().map(|x| x.as_u64().unwrap() as u8).collect::<Vec<u8>>();
+    let cdfs = u8v("cdfs");
+    let mem = u8v("mem");
+    let offsets: Vec<u32> = v["offsets"].as_array().unwrap().iter().map(|x| x.as_u64().unwrap() as u32).collect();
+    let channels = v["channels"].as_u64().unwrap() as usize;
+    let spc = v["spc"].as_u64().unwrap() as usize;
+    let expected = u8v("expected");
+
+    let mut dec = MansDecoder::new(&mem, &offsets);
+    let mut out = vec![0u8; channels * spc];
+    dec.decode_factorize(&cdfs, &mut out, channels, spc);
+    let mism = out.iter().zip(&expected).filter(|(a, b)| a != b).count();
+    eprintln!("mANS decode_factorize vs reference: {}/{} exact, {} mismatched", out.len() - mism, out.len(), mism);
+    assert_eq!(mism, 0, "mANS fac decoder must be BYTE-EXACT; {mism} mismatches");
+}

wai/wai-rs/tests/jpegai_mcm_calibration.rs

@@ -0,0 +1,85 @@
+//! Calibration: integer MCM stage decoders (`mcm_stage0` / `mcm_stage_n`) vs the
+//! float reference — run the SAME inputs through each `MCM/stage{i}.onnx` (ort) and
+//! through the Rust stage, compare the fixed-point context features `y[i]` (SNR,
+//! since the output is params not pixels). Each stage validated independently.
+//! Needs `neural` (ort).
+//!
+//! env: WAI_JPEGAI_MCM_DIR         = dir with stage0..3.onnx,
+//!      WAI_JPEGAI_MCM_WEIGHTS_DIR = dir with mcm_stage0..3_y.json
+#![cfg(feature = "neural")]
+
+use ort::session::{builder::GraphOptimizationLevel, Session};
+use ort::value::Value;
+use wai::jpegai_synthesis::{mcm_stage0, mcm_stage_n, Weights};
+
+fn arr(c: usize, h: usize, w: usize, seed: i64) -> (Vec<f32>, Vec<i64>) {
+    let n = c * h * w;
+    let f: Vec<f32> = (0..n).map(|i| ((((i as i64 + seed) * 31) % 197 - 98) as f32) / 64.0).collect();
+    let q: Vec<i64> = f.iter().map(|&x| (x * 65536.0).round() as i64).collect();
+    (f, q)
+}
+
+fn run_stage(path: &str, named: Vec<(&str, Vec<f32>, usize)>, h: usize, w: usize) -> Vec<f32> {
+    let mut s = Session::builder().unwrap()
+        .with_intra_threads(1).unwrap().with_inter_threads(1).unwrap()
+        .with_optimization_level(GraphOptimizationLevel::Level3).unwrap()
+        .commit_from_file(path).unwrap();
+    let inputs: Vec<(String, Value)> = named.into_iter().map(|(name, data, c)| {
+        let a = ndarray::Array4::from_shape_vec((1, c, h, w), data).unwrap();
+        (name.to_string(), Value::from_array(a).unwrap().into_dyn())
+    }).collect();
+    let out = s.run(inputs).unwrap();
+    out[0].try_extract_array::<f32>().unwrap().iter().copied().collect()
+}
+
+fn snr(int_q16: &[i64], flt: &[f32]) -> (f64, f64) {
+    let (mut maxd, mut se, mut sig) = (0f64, 0f64, 0f64);
+    for i in 0..int_q16.len() {
+        let q = int_q16[i] as f64 / 65536.0;
+        let f = flt[i] as f64;
+        let d = (q - f).abs();
+        maxd = maxd.max(d); se += d * d; sig += f * f;
+    }
+    (maxd, if se > 0.0 { 10.0 * (sig / se).log10() } else { f64::INFINITY })
+}
+
+#[test]
+fn mcm_stages_int_vs_float_snr() {
+    let dir = match std::env::var("WAI_JPEGAI_MCM_DIR") {
+        Ok(p) if std::path::Path::new(&p).exists() => p,
+        _ => { eprintln!("skip: set WAI_JPEGAI_MCM_DIR to the MCM onnx dir"); 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 to the weights dir"); return; }
+    };
+    let (h, w) = (16usize, 16usize);
+    let load = |i: usize| -> Weights {
+        let v: serde_json::Value = serde_json::from_slice(&std::fs::read(format!("{wdir}/mcm_stage{i}_y.json")).unwrap()).unwrap();
+        Weights::from_json(&v)
+    };
+    let mut worst = f64::INFINITY;
+
+    // stage 0: single input p0
+    {
+        let (pf, pq) = arr(160, h, w, 1);
+        let rf = run_stage(&format!("{dir}/stage0.onnx"), vec![("p0", pf, 160)], h, w);
+        let mine = mcm_stage0(&pq, h, w, &load(0));
+        let (md, sn) = snr(&mine, &rf);
+        eprintln!("MCM stage0:  max|Δ| = {md:.4}  SNR = {sn:.2} dB  (len {})", mine.len());
+        worst = worst.min(sn);
+    }
+    // stages 1..3: inputs (p_i, context of width 160/320/480)
+    for (i, (pn, cn, cc)) in [("p1", "y[0]", 160usize), ("p2", "y[1-0]", 320), ("p3", "y[2-0]", 480)].into_iter().enumerate() {
+        let stage = i + 1;
+        let (pf, pq) = arr(160, h, w, 7 * stage as i64 + 1);
+        let (cf, cq) = arr(cc, h, w, 13 * stage as i64 + 2);
+        let rf = run_stage(&format!("{dir}/stage{stage}.onnx"), vec![(pn, pf, 160), (cn, cf, cc)], h, w);
+        let mine = mcm_stage_n(&pq, &cq, cc, h, w, &load(stage));
+        let (md, sn) = snr(&mine, &rf);
+        eprintln!("MCM stage{stage}:  max|Δ| = {md:.4}  SNR = {sn:.2} dB  (len {})", mine.len());
+        worst = worst.min(sn);
+    }
+    eprintln!("MCM worst-stage SNR = {worst:.2} dB");
+    assert!(worst > 45.0, "an MCM stage regressed (worst SNR {worst:.2} dB)");
+}

wai/wai-rs/tests/jpegai_uv_calibration.rs

@@ -0,0 +1,63 @@
+//! Calibration: my integer `jpegai_synthesis::decode_uv` vs the float reference —
+//! run the SAME 256-channel latent through `synthesis16.onnx` (model_uv, ort) and
+//! through `decode_uv`, compare the UV planes (max abs diff + PSNR). This is the
+//! numerical-faithfulness pass for the chroma decoder. Needs `neural` (ort).
+//!
+//! env: WAI_JPEGAI_SYNTH16_UV = path to model_uv/synthesis16.onnx,
+//!      WAI_JPEGAI_UV_WEIGHTS  = path to the exported 16-bit uv weights JSON.
+#![cfg(feature = "neural")]
+
+use ort::session::{builder::GraphOptimizationLevel, Session};
+use ort::value::Value;
+use wai::jpegai_synthesis::{decode_uv, Weights};
+
+fn ort_synth(path: &str, x: Vec<f32>) -> (Vec<f32>, Vec<usize>) {
+    let mut s = Session::builder().unwrap()
+        .with_intra_threads(1).unwrap().with_inter_threads(1).unwrap()
+        .with_optimization_level(GraphOptimizationLevel::Level3).unwrap()
+        .commit_from_file(path).unwrap();
+    let name = s.inputs().first().unwrap().name().to_string();
+    let arr = ndarray::Array4::from_shape_vec((1, 256, 16, 16), x).unwrap();
+    let v = Value::from_array(arr).unwrap();
+    let out = s.run(ort::inputs![name.as_str() => v]).unwrap();
+    let view = out[0].try_extract_array::<f32>().unwrap();
+    (view.iter().copied().collect(), view.shape().to_vec())
+}
+
+#[test]
+fn synthesis_uv_int_vs_float_psnr() {
+    let onnx = match std::env::var("WAI_JPEGAI_SYNTH16_UV") {
+        Ok(p) if std::path::Path::new(&p).exists() => p,
+        _ => { eprintln!("skip: set WAI_JPEGAI_SYNTH16_UV to model_uv/synthesis16.onnx"); return; }
+    };
+    let wjson = match std::env::var("WAI_JPEGAI_UV_WEIGHTS") {
+        Ok(p) if std::path::Path::new(&p).exists() => p,
+        _ => { eprintln!("skip: set WAI_JPEGAI_UV_WEIGHTS to the weights JSON"); return; }
+    };
+    // deterministic small latent [256,16,16] (typical post-LSBS magnitude)
+    let n = 256 * 16 * 16;
+    let lf: Vec<f32> = (0..n).map(|i| (((i % 97) as f32) - 48.0) / 96.0).collect();
+
+    let (uvf, shape) = ort_synth(&onnx, lf.clone());
+    let (fmn, fmx) = uvf.iter().fold((f32::MAX, f32::MIN), |(a, b), &v| (a.min(v), b.max(v)));
+    eprintln!("float uv-synth out shape = {:?}, len {}, range [{:.3}, {:.3}]", shape, uvf.len(), fmn, fmx);
+    eprintln!("float sample [0..6] = {:?}", &uvf[..6]);
+
+    let li: Vec<i64> = lf.iter().map(|&x| (x * 65536.0).round() as i64).collect();
+    let v: serde_json::Value = serde_json::from_slice(&std::fs::read(&wjson).unwrap()).unwrap();
+    let wt = Weights::from_json(&v);
+    let uvi = decode_uv(&li, &wt);
+
+    assert_eq!(uvi.len(), uvf.len(), "output size mismatch (2×256×256 expected)");
+    let (mut maxd, mut mse) = (0f64, 0f64);
+    for i in 0..uvi.len() {
+        let f = uvf[i].clamp(0.0, 255.0) as f64;
+        let d = (uvi[i] as f64 - f).abs();
+        maxd = maxd.max(d);
+        mse += d * d;
+    }
+    mse /= uvi.len() as f64;
+    let psnr = if mse > 0.0 { 20.0 * (255.0f64).log10() - 10.0 * mse.log10() } else { f64::INFINITY };
+    eprintln!("INT vs FLOAT uv-synthesis:  max|Δ| = {:.2}   MSE = {:.2}   PSNR = {:.2} dB", maxd, mse, psnr);
+    assert!(psnr > 45.0, "uv-synthesis PSNR {psnr:.2} dB too low (expected ~55-58)");
+}