open-standards/wai/wai-rs/tests/jpegai_hyper_calibration.rs
Transaction Science e6453a9566 wai: gate the direct-network tests on legacy_int in their files too
Gate 7 made these seven tests require `legacy_int` in Cargo.toml, but their files
were still gated on `neural_int` or `neural` alone, which `check-tests-in-ci.py`
reports: a file's `#![cfg]` must name what its `[[test]]` requires. Each file is
now gated on `all(..., feature = "legacy_int")`, as `jpegai_ops.rs` is.

Three test comments named the design's appendix and a plan item. They now name
the text's sections (§11's `role_profiles`, §13.6.2's table).

Co-Authored-By: Claude Opus 5.5 <[email protected]>
2026-10-18 02:41:38 -04:00

66 lines
3.3 KiB
Rust
Raw Permalink Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

//! 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(all(feature = "neural", feature = "legacy_int"))]
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).unwrap();
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)");
}