From 22212badf2e58968da49e767f8ab37a738441be5 Mon Sep 17 00:00:00 2001 From: Transaction Science Date: Thu, 9 Jul 2026 07:01:28 -0400 Subject: [PATCH] =?UTF-8?q?wai-rs:=20hyper=5Fscale=5Fdecoder=20=E2=80=94?= =?UTF-8?q?=20floor=20(>>)=20not=20truncation;=20byte-exact=20vs=20the=20R?= =?UTF-8?q?EFERENCE=20decoder?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The fully-integer hyper_scale_decoder does its per-channel ÷2^n as an ARITHMETIC RIGHT SHIFT (floor), matching the real reference decoder. The exported ONNX represents it as integer `Div` (truncation toward zero), which differs from `>>` on negatives by 1 LSB — a ~1-LSB-wrong export. decode_hyper_scale used `/` (matching the ONNX); switch it to `>>` (matching the reference). Impact, on the real test256.bin hyper-latent (z[1,160,4,4] → i_s[1,160,16,16]): the old truncation matched only 30765/40960 of the reference i_s (the rest off by exactly +1, at the negative-remainder positions surfaced by the final abs); with `>>` it is 40960/40960 BYTE-EXACT vs the reference decoder. Retargets tests/jpegai_hyperscale_calibration.rs from the (truncating) onnx.reference golden to the real reference-decoder i_s — the correct target — validated byte-exact on real data. Structural lib test unaffected (shape/range/determinism). Co-Authored-By: Claude Opus 4.8 --- wai/wai-rs/src/jpegai_synthesis.rs | 13 +++- .../tests/jpegai_hyperscale_calibration.rs | 74 ++++++++----------- 2 files changed, 38 insertions(+), 49 deletions(-) diff --git a/wai/wai-rs/src/jpegai_synthesis.rs b/wai/wai-rs/src/jpegai_synthesis.rs index c693ce741..fad626c66 100644 --- a/wai/wai-rs/src/jpegai_synthesis.rs +++ b/wai/wai-rs/src/jpegai_synthesis.rs @@ -370,9 +370,14 @@ fn iconv(x: &[i64], cin: usize, h: usize, w: usize, wt: &[i64], b: &[i64], cout: /// 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). +/// are integer end to end); the per-channel ÷2^n is an ARITHMETIC RIGHT SHIFT (floor, +/// the reference-decoder semantics — the ONNX `Div` truncation is a ~1-LSB-wrong export). pub fn decode_hyper_scale(z: &[i64], cin: usize, h: usize, w: usize, net: &IntNet) -> (Vec, usize, usize, usize) { - // one QDQ block: clip → intconv → per-channel ÷2^exp (trunc toward 0) → [relu] + // one QDQ block: clip → intconv → per-channel ÷2^exp → [relu]. + // NB: the reference decoder does the power-of-two division as an ARITHMETIC RIGHT + // SHIFT (floor), not truncation-toward-zero — they differ for negatives. The + // exported ONNX represents it as `Div` (truncation), which is a ~1-LSB-wrong + // export; matching the real decoder (and the pipeline `i_s`) requires `>>`. let block = |x: &[i64], ci: usize, cname: &str, k: usize, pad: usize, ename: &str, relu_after: bool| -> (Vec, usize) { let xq: Vec = x.iter().map(|&v| it::clip(v, -128, 127)).collect(); let (cw, cd) = net.get(&format!("{cname}.weight")); @@ -382,10 +387,10 @@ pub fn decode_hyper_scale(z: &[i64], cin: usize, h: usize, w: usize, net: &IntNe 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 + let n = ev[o] as u32; // shift = 2^n for p in 0..plane { let idx = o * plane + p; - let mut r = y[idx] / div; // integer division, truncates toward zero + let mut r = y[idx] >> n; // arithmetic right shift = floor (reference semantics) if relu_after && r < 0 { r = 0; } y[idx] = r; } diff --git a/wai/wai-rs/tests/jpegai_hyperscale_calibration.rs b/wai/wai-rs/tests/jpegai_hyperscale_calibration.rs index 8b1c959f0..8febe622a 100644 --- a/wai/wai-rs/tests/jpegai_hyperscale_calibration.rs +++ b/wai/wai-rs/tests/jpegai_hyperscale_calibration.rs @@ -1,58 +1,42 @@ -//! 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. +//! Byte-exact calibration for the fully-integer `hyper_scale_decoder` +//! (`decode_hyper_scale`) against the REFERENCE DECODER on real test256.bin data: +//! run the real hyper-latent `z` through the net and match the reference scale-index +//! output `i_s` exactly. The per-channel ÷2^n is an arithmetic right shift (floor); +//! the exported ONNX's `Div` truncates, which is ~1-LSB wrong on negatives — this +//! test pins the *reference-decoder* semantics, not the ONNX. No ort/onnx 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'))" +//! env: WAI_JPEGAI_HYPERSCALE_IO = hyperscale_io_fixture.json +//! { z:[i64], i_s:[i64], z_shape:[1,C,h,w] } (dumped from the reference decoder) +//! WAI_JPEGAI_HYPERSCALE_WEIGHTS = hyperscale_y_weights.json (int weights) #![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") { +fn hyper_scale_byte_exact_vs_reference() { + let io = match std::env::var("WAI_JPEGAI_HYPERSCALE_IO") { + Ok(p) if std::path::Path::new(&p).exists() => p, + _ => { eprintln!("skip: set WAI_JPEGAI_HYPERSCALE_IO to the z→i_s fixture"); return; } + }; + let wj = 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 = (0..n).map(|i| ((i as i64 * 37) % 291) - 140).collect(); + let v: serde_json::Value = serde_json::from_slice(&std::fs::read(&io).unwrap()).unwrap(); + let z: Vec = v["z"].as_array().unwrap().iter().map(|x| x.as_i64().unwrap()).collect(); + let i_s: Vec = v["i_s"].as_array().unwrap().iter().map(|x| x.as_i64().unwrap()).collect(); + let zs: Vec = v["z_shape"].as_array().unwrap().iter().map(|x| x.as_u64().unwrap() as usize).collect(); + let (c, h, w) = (zs[1], zs[2], zs[3]); - let rv: serde_json::Value = serde_json::from_slice(&std::fs::read(&rjson).unwrap()).unwrap(); - let refv: Vec = rv["out"].as_array().unwrap().iter().map(|x| x.as_i64().unwrap()).collect(); - let shape: Vec = rv["shape"].as_array().unwrap().iter().map(|x| x.as_u64().unwrap() as usize).collect(); - eprintln!("reference shape = {shape:?}, {} values", refv.len()); + let net = IntNet::from_json(&serde_json::from_slice(&std::fs::read(&wj).unwrap()).unwrap()); + let (mine, co, ho, wo) = decode_hyper_scale(&z, c, h, w, &net); + eprintln!("decode_hyper_scale(real z[{c},{h},{w}]) -> [{co},{ho},{wo}] = {} vals", mine.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); + assert_eq!(mine.len(), i_s.len(), "output size mismatch"); + let mism = mine.iter().zip(&i_s).filter(|(a, b)| a != b).count(); + if let Some((i, (a, b))) = mine.iter().zip(&i_s).enumerate().find(|(_, (a, b))| a != b) { + eprintln!("first mismatch @ {i}: mine={a} ref={b}"); } - 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"); + eprintln!("hyper_scale vs reference i_s: {}/{} exact, {} mismatched", mine.len() - mism, mine.len(), mism); + assert_eq!(mism, 0, "hyper_scale must be BYTE-EXACT vs the reference decoder; {mism} mismatches"); } -- 2.54.0 (Apple Git-157)