wai-rs: jpegai_pipeline::reconstruct_yhat — the SGMM autoregressive spine #7
2026-07-12
wai-rs: jpegai_pipeline::reconstruct_yhat — the SGMM autoregressive spine (real data 76.8 dB)
Rebuilds the synthesis-input latent y_hat from the dequantized residual + the hyper
params — the last large structural piece of the JPEG-AI decode chain. Ports the
reference context.pred loop:
p = chunk(psi, 4) # per-stage hyper params, ascending
r = down_shuffle(resi) -> (k00,k11,k01,k10) # 2x2 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)
Reuses the byte-exact MCM stages (mcm_stage0 / mcm_stage_n). Two subtleties pinned
from the reference source (not the ONNX): the checkerboard order is (0,0),(1,1),
(0,1),(1,0) — the diagonal pair first — and the context concatenation is ASCENDING
(y_0 first), despite the ONNX input being labelled y[1-0]; stage 0 has no context.
Validated on real test256.bin data: reconstruct_yhat(resi_dq_y, psi_y) vs the
reference yhat_y = 76.83 dB SNR (max|Δ| = 0.005), carrying the MCM nets' 16-bit PTQ
floor (the float reference loop matches yhat_y to ~139 dB). Unit test pins the
shuffle round-trip + checkerboard ordering.
Co-Authored-By: Claude Opus 4.8 <[email protected]>