Transaction-Science/open-standards
Watch
1
0
Fork
You've already forked open-standards
 0
Code Issues Pull requests Projects Releases 15 Packages Wiki Activity Actions

wai-rs: hyper_scale_decoder floor(>>) not truncation — byte-exact vs the reference decoder #4

Merged
dcharlot merged 1 commit from hyperscale-floorfix into main 2026-07-12 12:22:07 -04:00
Conversation 0 Commits 1 Files changed 2 +37 -48
2 changed files with 37 additions and 48 deletions
Download patch file Download diff file Expand all files Collapse all files

13
wai/wai-rs/src/jpegai_synthesis.rs
View file

@ -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<i64>, 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<i64>, usize) {
let xq: Vec<i64> = 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;
}

74
wai/wai-rs/tests/jpegai_hyperscale_calibration.rs
View file

@ -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<i64> = (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<i64> = v["z"].as_array().unwrap().iter().map(|x| x.as_i64().unwrap()).collect();
let i_s: Vec<i64> = v["i_s"].as_array().unwrap().iter().map(|x| x.as_i64().unwrap()).collect();
let zs: Vec<usize> = 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<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 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");
}