cbc4caf7b0
Gate 5's fault (a grouped convolution finished with its neighbour channel's requantisation on the fast kernel) left the unit test green: every committed vector of rec_y, rec_uv, syn_y and syn_uv refused past A32, because the synthetic networks' multipliers grew each convolution's output by up to 2^6 per product, so no decode of those graphs ran its operators end to end. - The generator now sets each synthetic convolution's multiplier so that a weight of 32767 times its products keeps the input's magnitude (and halves the biases): every rec_y, rec_uv, syn_y and syn_uv vector of graphs/lower_jpegai.jsonl now decodes, and only the vectors one past an input's range refuse. The set's digest is 1bc0daae…; its graphs, certificate and vectors are regenerated, and both certifiers reproduce them (wai-verify vectors verify; wai_ops_check.py ops-conformance), with the tests naming the new digest. - jpegai_ops::fast_kernels::the_jpegai_networks_run_their_convolutions_on_the_fast_kernel requires a completed decode of each of the four graphs, not of any one. Co-Authored-By: Claude Opus 5.5 <[email protected]>
509 lines
28 KiB
Rust
509 lines
28 KiB
Rust
//! Every committed parameter set, lowered (§13) and decoded through the interpreter
|
|
//! and the stage operators, reproduces its committed output digest.
|
|
//!
|
|
//! The drivers here are test drivers, written from the payload extensions: they are
|
|
//! what a capability's driver does around its graphs (§13.5). They check the lowerings
|
|
//! and the operators against outputs that were fixed before either existed.
|
|
//!
|
|
//! Every graph a driver runs goes through `graph::load` of its bytes under `acc64`,
|
|
//! the profile every lowering of §13 declares, as a sink's does: load step 8 runs the
|
|
//! bounds pass and refuses a set whose certificate does not meet it.
|
|
|
|
use std::path::{Path, PathBuf};
|
|
|
|
use wai_ops::{Profile, bounds, graph, plan};
|
|
use wai_ops::interp::{Sequence, Tensor, run_observed};
|
|
use wai_ops::json::{self, Value, hex};
|
|
use wai_ops::lower;
|
|
use wai_ops::sha256::sha256;
|
|
use wai_ops::stage::{self, rans64};
|
|
|
|
fn wai() -> PathBuf {
|
|
PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("..")
|
|
}
|
|
fn read(p: &str) -> Vec<u8> {
|
|
std::fs::read(wai().join(p)).unwrap_or_else(|e| panic!("{p}: {e}"))
|
|
}
|
|
fn jread(p: &str) -> Value {
|
|
json::parse(&read(p)).unwrap()
|
|
}
|
|
fn sha(b: &[u8]) -> String {
|
|
hex(&sha256(b))
|
|
}
|
|
fn hwc(t: &Tensor) -> Vec<u8> {
|
|
let (c, h, w) = (t.c as usize, t.h as usize, t.w as usize);
|
|
let mut o = vec![0u8; c * h * w];
|
|
for cc in 0..c {
|
|
for p in 0..h * w {
|
|
o[p * c + cc] = u8::try_from(t.data[cc * h * w + p]).unwrap();
|
|
}
|
|
}
|
|
o
|
|
}
|
|
/// One decode of a lowered graph with one data input: loaded from its bytes, as a sink
|
|
/// loads it, then run, observed: every value each operator forms under a named
|
|
/// intermediate (§8.5, Table 8.5-A), and every producer result, must lie in the interval
|
|
/// the graph's bounds pass records for it.
|
|
fn one(g: &wai_ops::Graph, x: Tensor, sc: &[i64]) -> Tensor {
|
|
let g = graph::load(&g.to_bytes(), Profile::Acc64).expect("a lowered graph loads");
|
|
let sizes = plan::load_plan(&g).unwrap();
|
|
let b = bounds::pass(&g, sizes.as_deref()).unwrap();
|
|
let mut w = bounds::Within::new(&b);
|
|
let out = run_observed(&g, &[x], sc, &mut Sequence::new(), &mut w).unwrap().remove(0);
|
|
assert!(w.outside.is_empty(), "values outside their certified intervals: {:?}", w.outside);
|
|
assert!(w.unrecorded.is_empty(), "values under keys the certificate does not record: {:?}", w.unrecorded);
|
|
assert!(w.seen > 0);
|
|
out
|
|
}
|
|
fn u32le(b: &[u8], o: usize) -> u32 {
|
|
u32::from_le_bytes(b[o..o + 4].try_into().unwrap())
|
|
}
|
|
fn i64s(v: &Value) -> Vec<i64> {
|
|
v.as_arr().unwrap().iter().map(|x| x.as_i64().unwrap()).collect()
|
|
}
|
|
fn sidecar(p: &str, name: &str) -> String {
|
|
jread(p).get(name).unwrap().get("sha256").unwrap().as_str().unwrap().to_string()
|
|
}
|
|
|
|
#[test]
|
|
fn int_synth_sets_reproduce_their_digests() {
|
|
for name in ["gradient", "bloom", "facade", "photo"] {
|
|
let set = lower::int_synth(&read(&format!("neural-conformance/vectors/{name}/int_synth.json"))).unwrap();
|
|
let g = set.get("synth").unwrap();
|
|
let p = read(&format!("neural-conformance/vectors/{name}/latent.bin"));
|
|
let (h, w, c, s) = (u32le(&p, 0), u32le(&p, 4), u16::from_le_bytes([p[8], p[9]]) as u32, 1u32 << p[10]);
|
|
let lat: Vec<i64> = p[12..].iter().map(|&b| b as i8 as i64).collect();
|
|
let out = one(g, Tensor::new(c, s, s, lat), &[]);
|
|
assert_eq!((out.c, out.h, out.w), (3, h, w), "{name}: output shape");
|
|
assert_eq!(sha(&hwc(&out)), sidecar("neural-conformance/sha256.json", name), "{name}");
|
|
}
|
|
}
|
|
|
|
fn world_clip(dir: &str, json_name: &str, t: &str, gs: &str, two: bool) -> Vec<String> {
|
|
const ACTIONS: [(i64, i64); 5] = [(0, 0), (1, 0), (-1, 0), (0, 1), (0, -1)];
|
|
let j = jread(&format!("{dir}/{json_name}"));
|
|
let c = j.get("c").unwrap().as_u32().unwrap();
|
|
let lat = j.get("lat").unwrap().as_u32().unwrap();
|
|
let afrac = j.get("afrac").unwrap().as_u32().unwrap();
|
|
let set = lower::world(&read(&format!("{dir}/{t}")), &read(&format!("{dir}/{gs}")), c).unwrap();
|
|
let (tg, gg) = (set.get("T").unwrap(), set.get("g_s").unwrap());
|
|
let q = |d: i64| stage::dequantize(d, afrac, 0).unwrap();
|
|
let rows: Vec<Vec<i64>> = if two {
|
|
j.get("pairs").unwrap().as_arr().unwrap().iter().map(|p| {
|
|
let a = i64s(p);
|
|
let ((xa, ya), (xb, yb)) = (ACTIONS[a[0] as usize], ACTIONS[a[1] as usize]);
|
|
vec![q(xa), q(ya), q(xb), q(yb)]
|
|
}).collect()
|
|
} else {
|
|
i64s(j.get("actions").unwrap()).iter().map(|&a| {
|
|
let (dx, dy) = ACTIONS[a as usize];
|
|
vec![q(dx), q(dy), 0]
|
|
}).collect()
|
|
};
|
|
let mut z = Tensor::new(c, lat, lat, i64s(j.get("z0").unwrap()));
|
|
let mut frames = vec![sha(&hwc(&one(gg, z.clone(), &[])))];
|
|
for r in &rows {
|
|
z = one(tg, z, r);
|
|
frames.push(sha(&hwc(&one(gg, z.clone(), &[]))));
|
|
}
|
|
frames
|
|
}
|
|
|
|
#[test]
|
|
fn world_sets_reproduce_every_frame() {
|
|
for (dir, js, t, gs, two, key) in [
|
|
("wai-rs/tests/vectors/int_world", "world.json", "world_T.bin", "world_gs.bin", false, "world"),
|
|
("wai-rs/tests/vectors/int_world2", "world2.json", "world2_T.bin", "world2_gs.bin", true, "world2"),
|
|
] {
|
|
let got = world_clip(dir, js, t, gs, two);
|
|
let want: Vec<String> = jread(&format!("{dir}/sha256.json")).get(key).unwrap().get("frames_sha256").unwrap()
|
|
.as_arr().unwrap().iter().map(|v| v.as_str().unwrap().to_string()).collect();
|
|
assert_eq!(got, want, "{dir}");
|
|
}
|
|
}
|
|
|
|
/// A `t4097`/`lenient`/`canonical` stream of `n` symbols, each against `tables[pick(i)]`.
|
|
fn decode_stream(stream: &[u8], tables: &[rans64::Table], reads: &[usize], modes: rans64::Modes) -> Vec<i64> {
|
|
rans64::decode(stream, tables, reads, modes).unwrap()
|
|
}
|
|
|
|
const WIH1_MODES: rans64::Modes = rans64::Modes { table: rans64::TableRule::T4097, escape: rans64::Escape::Lenient, end: rans64::End::Canonical };
|
|
|
|
fn table(cdf: &[i64], len: usize, off: i64) -> rans64::Table {
|
|
rans64::Table { cdf: cdf[..len].to_vec(), offset: off }
|
|
}
|
|
|
|
#[test]
|
|
fn int_hyper_tiny_reproduces_its_digest() {
|
|
let d = "wai-rs/tests/vectors/int_hyper";
|
|
let set = lower::int_hyper(&read(&format!("{d}/tiny_h_s.bin")), &read(&format!("{d}/tiny_g_s.bin"))).unwrap();
|
|
let afrac = lower::parse_whm1(&read(&format!("{d}/tiny_g_s.bin"))).unwrap().afrac;
|
|
let e = jread(&format!("{d}/tiny_entropy.json"));
|
|
let p = read(&format!("{d}/tiny_bitstream.wih1"));
|
|
assert_eq!(&p[0..4], b"WIH1");
|
|
let f = |k: usize| u32le(&p, 8 + 4 * k);
|
|
let (m, sh, sw, nz, zh, zw) = (f(0), f(1), f(2), f(3), f(4), f(5));
|
|
let lz = u32le(&p, 32) as usize;
|
|
let zs = &p[36..36 + lz];
|
|
let ly = u32le(&p, 36 + lz) as usize;
|
|
let ys = &p[40 + lz..40 + lz + ly];
|
|
let rows = |k: &str| -> Vec<Vec<i64>> { e.get(k).unwrap().as_arr().unwrap().iter().map(i64s).collect() };
|
|
let (zc, zl, zo, zm) = (rows("z_cdfs"), i64s(e.get("z_cdf_lengths").unwrap()), i64s(e.get("z_offsets").unwrap()), i64s(e.get("z_median_q").unwrap()));
|
|
let ztabs: Vec<_> = (0..nz as usize).map(|k| table(&zc[k], zl[k] as usize, zo[k])).collect();
|
|
let reads: Vec<usize> = (0..nz as usize).flat_map(|k| std::iter::repeat_n(k, (zh * zw) as usize)).collect();
|
|
let zsym = decode_stream(zs, &ztabs, &reads, WIH1_MODES);
|
|
let z: Vec<i64> = zsym.iter().enumerate().map(|(i, &s)| stage::dequantize(s, afrac, zm[i / (zh * zw) as usize]).unwrap()).collect();
|
|
let hyper = one(set.get("h_s").unwrap(), Tensor::new(nz, zh, zw, z), &[]);
|
|
assert_eq!((hyper.c, hyper.h, hyper.w), (2 * m, sh, sw));
|
|
let n = (m * sh * sw) as usize;
|
|
let (sigma, mu) = hyper.data.split_at(n);
|
|
let (yc, yl, yo, edges) = (rows("cdfs"), i64s(e.get("cdf_lengths").unwrap()), i64s(e.get("offsets").unwrap()), i64s(e.get("scale_edges").unwrap()));
|
|
let ytabs: Vec<_> = (0..edges.len()).map(|b| table(&yc[b], yl[b] as usize, yo[b])).collect();
|
|
let yreads: Vec<usize> = sigma.iter().map(|&s| stage::bucket_ceil(s, &edges).unwrap()).collect();
|
|
let ysym = decode_stream(ys, &ytabs, &yreads, WIH1_MODES);
|
|
let y: Vec<i64> = ysym.iter().zip(mu).map(|(&r, &u)| stage::dequantize(r, afrac, u).unwrap()).collect();
|
|
let rgb = one(set.get("g_s").unwrap(), Tensor::new(m, sh, sw, y), &[]);
|
|
assert_eq!(sha(&hwc(&rgb)), sidecar(&format!("{d}/sha256.json"), "tiny"));
|
|
}
|
|
|
|
#[test]
|
|
fn splat_reproduces_its_digest() {
|
|
let d = "wai-rs/tests/vectors/int_splat";
|
|
let set = lower::splat(&read(&format!("{d}/splat_gs.bin"))).unwrap();
|
|
let p = read(&format!("{d}/splat.wis1"));
|
|
assert_eq!(&p[0..4], b"WIS1");
|
|
let afrac = p[5] as u32;
|
|
let f = |k: usize| u32le(&p, 8 + 4 * k);
|
|
let (c, sh, sw, s, attr) = (f(0), f(1), f(2), f(3), f(4));
|
|
let mut o = 32 + 8 * attr as usize;
|
|
let lens: Vec<usize> = (0..c as usize).map(|k| u16::from_le_bytes([p[o + 2 * k], p[o + 2 * k + 1]]) as usize).collect();
|
|
o += 2 * c as usize;
|
|
let offs: Vec<i64> = (0..c as usize).map(|k| i32::from_le_bytes(p[o + 4 * k..o + 4 * k + 4].try_into().unwrap()) as i64).collect();
|
|
o += 4 * c as usize;
|
|
let meds: Vec<i64> = (0..c as usize).map(|k| i64::from_le_bytes(p[o + 8 * k..o + 8 * k + 8].try_into().unwrap())).collect();
|
|
o += 8 * c as usize;
|
|
let mut tabs = Vec::new();
|
|
for k in 0..c as usize {
|
|
let cdf: Vec<i64> = (0..lens[k]).map(|j| u32le(&p, o + 4 * j) as i64).collect();
|
|
o += 4 * lens[k];
|
|
tabs.push(rans64::Table { cdf, offset: offs[k] });
|
|
}
|
|
let l = u32le(&p, o) as usize;
|
|
let st = &p[o + 4..o + 4 + l];
|
|
let per = (sh * sw) as usize;
|
|
let reads: Vec<usize> = (0..c as usize).flat_map(|k| std::iter::repeat_n(k, per)).collect();
|
|
let sym = decode_stream(st, &tabs, &reads, WIH1_MODES);
|
|
let y: Vec<i64> = sym.iter().enumerate().map(|(i, &v)| stage::dequantize(v, afrac, meds[i / per]).unwrap()).collect();
|
|
let out = one(set.get("g_s").unwrap(), Tensor::new(c, sh, sw, y), &[]);
|
|
assert_eq!((out.c, out.h, out.w), (attr, s, s));
|
|
let bytes: Vec<u8> = out.data.iter().map(|&v| u8::try_from(v).unwrap()).collect();
|
|
assert_eq!(sha(&bytes), sidecar(&format!("{d}/sha256.json"), "splat"));
|
|
}
|
|
|
|
#[test]
|
|
fn audio_reproduces_its_digest() {
|
|
let d = "wai-rs/tests/vectors/int_audio";
|
|
let g = lower::wam1(&read(&format!("{d}/audio_model.bin"))).unwrap();
|
|
let p = read(&format!("{d}/audio_bitstream.wia1"));
|
|
assert_eq!(&p[0..4], b"WIA1");
|
|
let afrac = p[5] as u32;
|
|
let (n_ch, lat_len, t_out) = (u32le(&p, 12), u32le(&p, 16), u32le(&p, 20));
|
|
let offset = i32::from_le_bytes(p[24..28].try_into().unwrap()) as i64;
|
|
let cdf_len = u16::from_le_bytes([p[28], p[29]]) as usize;
|
|
let mut o = 32;
|
|
let meds: Vec<i64> = (0..n_ch as usize).map(|k| i64::from_le_bytes(p[o + 8 * k..o + 8 * k + 8].try_into().unwrap())).collect();
|
|
o += 8 * n_ch as usize;
|
|
let mut tabs = Vec::new();
|
|
for _ in 0..n_ch {
|
|
tabs.push(rans64::Table { cdf: (0..cdf_len).map(|j| u32le(&p, o + 4 * j) as i64).collect(), offset });
|
|
o += 4 * cdf_len;
|
|
}
|
|
let l = u32le(&p, o) as usize;
|
|
let st = &p[o + 4..o + 4 + l];
|
|
let per = lat_len as usize;
|
|
let reads: Vec<usize> = (0..n_ch as usize).flat_map(|k| std::iter::repeat_n(k, per)).collect();
|
|
let sym = decode_stream(st, &tabs, &reads, WIH1_MODES);
|
|
let y: Vec<i64> = sym.iter().enumerate().map(|(i, &v)| stage::dequantize(v, afrac, meds[i / per]).unwrap()).collect();
|
|
let out = one(&g, Tensor::new(n_ch, 1, lat_len, y), &[t_out as i64]);
|
|
assert_eq!((out.c, out.h, out.w), (1, 1, t_out));
|
|
let pcm: Vec<u8> = out.data.iter().flat_map(|&v| i16::try_from(v).unwrap().to_le_bytes()).collect();
|
|
assert_eq!(sha(&pcm), sidecar(&format!("{d}/sha256.json"), "audio"));
|
|
}
|
|
|
|
#[test]
|
|
fn int_mlicpp_reproduces_its_digest() {
|
|
let d = "wai-rs/tests/vectors/int_mlicpp";
|
|
let bundle = read(&format!("{d}/model.wmb"));
|
|
let set = lower::wmb1(&bundle).unwrap();
|
|
let exp = jread(&format!("{d}/expected.json"));
|
|
// The tables section.
|
|
let mut o = 4;
|
|
let mut secs = Vec::new();
|
|
for _ in 0..4 {
|
|
let n = u32le(&bundle, o) as usize;
|
|
secs.push(&bundle[o + 4..o + 4 + n]);
|
|
o += 4 + n;
|
|
}
|
|
let t = json::parse(secs[3]).unwrap();
|
|
let afrac = t.get("afrac").unwrap().as_u32().unwrap();
|
|
let rows = |k: &str| -> Vec<Vec<i64>> { t.get(k).unwrap().as_arr().unwrap().iter().map(i64s).collect() };
|
|
// The payload: a WAI1 envelope holds its manifest, then the payload, each after
|
|
// its u32 length.
|
|
let env = read(&format!("{d}/stream.wai"));
|
|
let ml = u32le(&env, 4) as usize;
|
|
let pl = u32le(&env, 8 + ml) as usize;
|
|
let p = &env[12 + ml..12 + ml + pl];
|
|
assert_eq!(sha(p), exp.get("payload_sha256").unwrap().as_str().unwrap());
|
|
assert_eq!(&p[0..4], b"WIM1");
|
|
let h16 = |k: usize| u16::from_le_bytes([p[k], p[k + 1]]) as u32;
|
|
let (cz, zh, zw, cy, yh, yw, kk) = (h16(12), h16(14), h16(16), h16(18), h16(20), h16(22), h16(24));
|
|
let (lz, ly) = (u32le(p, 28) as usize, u32le(p, 32) as usize);
|
|
let (zs, ys) = (&p[36..36 + lz], &p[36 + lz..36 + lz + ly]);
|
|
let modes = rans64::Modes { table: rans64::TableRule::T2, escape: rans64::Escape::Lenient, end: rans64::End::FinalOnly };
|
|
let (zc, zl, zo, zm) = (rows("z_cdfs"), i64s(t.get("z_cdf_lengths").unwrap()), i64s(t.get("z_offsets").unwrap()), i64s(t.get("z_median_q").unwrap()));
|
|
let ztabs: Vec<_> = (0..cz as usize).map(|k| table(&zc[k], zl[k] as usize, zo[k])).collect();
|
|
let per = (zh * zw) as usize;
|
|
let reads: Vec<usize> = (0..cz as usize).flat_map(|k| std::iter::repeat_n(k, per)).collect();
|
|
let zsym = rans64::decode(zs, &ztabs, &reads, modes).unwrap();
|
|
let z: Vec<i64> = zsym.iter().enumerate().map(|(i, &s)| stage::dequantize(s, afrac, zm[i / per]).unwrap()).collect();
|
|
let hyper = one(set.get("h_s").unwrap(), Tensor::new(cz, zh, zw, z), &[]);
|
|
// The y stream: every slice's two passes from one decoder, in order.
|
|
let (yc, yl, yo, edges) = (rows("cdfs"), i64s(t.get("cdf_lengths").unwrap()), i64s(t.get("offsets").unwrap()), i64s(t.get("scale_edges").unwrap()));
|
|
let ytabs: Vec<_> = (0..edges.len()).map(|b| table(&yc[b], yl[b] as usize, yo[b])).collect();
|
|
let mut dec = rans64::Rans::new(ys).unwrap();
|
|
let plane = (yh * yw) as usize;
|
|
let mut y = vec![0i64; cy as usize * plane];
|
|
let per_slice = (cy / kk) as usize;
|
|
for j in 0..kk as usize {
|
|
for anchor in [true, false] {
|
|
let mut inp = y.clone();
|
|
inp.extend_from_slice(&hyper.data);
|
|
let sm = one(set.get("entropy").unwrap(), Tensor::new(3 * cy, yh, yw, inp), &[]);
|
|
let n = cy as usize * plane;
|
|
for ch in j * per_slice..(j + 1) * per_slice {
|
|
for idx in stage::coding_order(1, yh, yw, Some(anchor)) {
|
|
let e = ch * plane + idx as usize;
|
|
let b = stage::bucket_ceil(sm.data[e], &edges).unwrap();
|
|
let r = dec.symbol(&ytabs[b], rans64::Escape::Lenient).unwrap();
|
|
y[e] = stage::dequantize(r, afrac, sm.data[n + e]).unwrap();
|
|
}
|
|
}
|
|
}
|
|
}
|
|
assert!(!dec.ran_out());
|
|
let rgb = one(set.get("g_s").unwrap(), Tensor::new(cy, yh, yw, y), &[]);
|
|
assert_eq!(sha(&hwc(&rgb)), exp.get("rgb_sha256").unwrap().as_str().unwrap());
|
|
}
|
|
|
|
/// `jpegai-int/1` of a bundle already parsed writes the graphs the lowering of its bytes
|
|
/// writes, and loads them as `graph::load` of those bytes does: on the committed set, on
|
|
/// the bundle `WAI_OPS_JPEGAI_BUNDLE` names when it is set, and on each committed set
|
|
/// with one member changed so that a load check fails (both refuse).
|
|
#[test]
|
|
fn jpegai_parsed_lowering_is_the_lowering_of_its_bytes() {
|
|
let mut sets = vec![(read("ops-conformance/sets/jpegai/norm.json"), read("ops-conformance/sets/jpegai/weights/hs_y.json"), read("ops-conformance/sets/jpegai/weights/hs_uv.json"))];
|
|
if let Ok(dir) = std::env::var("WAI_OPS_JPEGAI_BUNDLE") {
|
|
let rd = |n: &str| std::fs::read(Path::new(&dir).join(n)).unwrap();
|
|
sets.push((rd("norm.json"), rd("hs_y.json"), rd("hs_uv.json")));
|
|
}
|
|
let committed = ["hs_y", "hs_uv"].map(|r| read(&format!("ops-conformance/certificates/wai.image.jpegai/1bc0daae512797a8817f2e77e57d009ace0d22acf4fe6945af8f4e0ee95244b9/{r}.whg1")));
|
|
for (i, (n, y, uv)) in sets.iter().enumerate() {
|
|
let bytes = lower::jpegai_scale(n, y, uv).unwrap();
|
|
if i == 0 {
|
|
assert_eq!(bytes.get("hs_y").unwrap().to_bytes(), committed[0], "the committed hs_y");
|
|
assert_eq!(bytes.get("hs_uv").unwrap().to_bytes(), committed[1], "the committed hs_uv");
|
|
}
|
|
let net = lower::parse_jpegai(n, y, uv).unwrap();
|
|
assert_eq!(lower::jpegai_net(&net).unwrap(), bytes, "set {i}: the parsed bundle lowers as its bytes do");
|
|
let loaded = lower::jpegai_net_loaded(net, Profile::Acc32).unwrap();
|
|
for (l, role) in [(&loaded.hs_y, "hs_y"), (&loaded.hs_uv, "hs_uv")] {
|
|
let g = bytes.get(role).unwrap();
|
|
assert_eq!(**l, *g, "set {i}: {role} loaded from the parsed bundle");
|
|
assert!(graph::load(&g.to_bytes(), Profile::Acc32).is_ok() && l.meets(Profile::Acc32));
|
|
}
|
|
}
|
|
// A bound of 256, and a per-channel shift of 128: both forms refuse, `model_mismatch`.
|
|
let (n, y, uv) = &sets[0];
|
|
let text = |b: &[u8]| String::from_utf8(b.to_vec()).unwrap();
|
|
let (nt, yt) = (text(n), text(y));
|
|
assert!(nt.contains("\"bounds\":[") && yt.contains("\"onnx::Pow_85\":{\"dims\":[1,"), "the committed set is written compactly");
|
|
let shift = |t: &str| {
|
|
let at = t.find("\"onnx::Pow_85\"").unwrap();
|
|
let v = at + t[at..].find("\"v\":[").unwrap() + 5;
|
|
format!("{}128,{}", &t[..v], &t[v..].splitn(2, ',').nth(1).unwrap())
|
|
};
|
|
for (norm, hs_y) in [(nt.replacen("\"bounds\":[", "\"bounds\":[256,", 1), yt.clone()), (nt.clone(), shift(&yt))] {
|
|
let (norm, hs_y) = (norm.into_bytes(), hs_y.into_bytes());
|
|
let a = lower::jpegai_scale(&norm, &hs_y, uv).unwrap_err();
|
|
let net = lower::parse_jpegai(&norm, &hs_y, uv).unwrap();
|
|
assert_eq!(lower::jpegai_net(&net).unwrap_err(), a);
|
|
assert_eq!(lower::jpegai_net_loaded(net, Profile::Acc32).err().unwrap(), a);
|
|
assert_eq!(a.code, wai_ops::Code::ModelMismatch);
|
|
}
|
|
}
|
|
|
|
/// The jpegai scale path on the real bundle, where one is supplied: lowered, and run on
|
|
/// every hyper-latent value in range, against the legacy decoder's arithmetic restated.
|
|
/// The bundle is not committed, so the test is ignored unless asked for:
|
|
/// `WAI_OPS_JPEGAI_BUNDLE` names a directory holding `norm.json`, `hs_y.json` and
|
|
/// `hs_uv.json`.
|
|
#[test]
|
|
#[ignore = "needs a jpegai bundle that is not committed: set WAI_OPS_JPEGAI_BUNDLE and run with --ignored"]
|
|
fn jpegai_scale_path_on_a_supplied_bundle() {
|
|
let dir = std::env::var("WAI_OPS_JPEGAI_BUNDLE").expect("WAI_OPS_JPEGAI_BUNDLE names the bundle");
|
|
let rd = |n: &str| std::fs::read(Path::new(&dir).join(n)).unwrap();
|
|
let set = lower::jpegai_scale(&rd("norm.json"), &rd("hs_y.json"), &rd("hs_uv.json")).unwrap();
|
|
for (role, file, cz) in [("hs_y", "hs_y.json", 160u32), ("hs_uv", "hs_uv.json", 96)] {
|
|
let g = set.get(role).unwrap();
|
|
let net = json::parse(&rd(file)).unwrap();
|
|
let norm = json::parse(&rd("norm.json")).unwrap();
|
|
let gains = i64s(norm.get(if role == "hs_y" { "gain_y" } else { "gain_uv" }).unwrap());
|
|
let z: Vec<i64> = (0..cz as i64 * 16).map(|i| (i * 7919 + 13) % 64 - 31).collect();
|
|
let got = one(g, Tensor::new(cz, 4, 4, z.clone()), &[]);
|
|
let want = legacy_scale(&net, &z, cz as usize, &gains);
|
|
assert_eq!(got.data, want, "{role}");
|
|
eprintln!("{role}: {} values equal", want.len());
|
|
}
|
|
}
|
|
|
|
/// The hyper-scale arithmetic as the legacy decoder writes it: clip, integer
|
|
/// convolution with its bias, an arithmetic right shift per channel, ReLU, depth to
|
|
/// space (CRD), abs, clip to 3967, then the gain.
|
|
fn legacy_scale(net: &Value, z: &[i64], cz: usize, gains: &[i64]) -> Vec<i64> {
|
|
let t = |k: &str| i64s(net.get(k).unwrap().get("v").unwrap());
|
|
let block = |x: &[i64], cin: usize, name: &str, k: usize, pad: usize, sh: &str, relu: bool| -> (Vec<i64>, usize) {
|
|
let (w, b, s) = (t(&format!("{name}.weight")), t(&format!("{name}.bias")), t(sh));
|
|
let cout = b.len();
|
|
let mut y = vec![0i64; cout * 16];
|
|
for o in 0..cout {
|
|
for oy in 0..4 {
|
|
for ox in 0..4 {
|
|
let mut acc = b[o];
|
|
for i in 0..cin {
|
|
for ky in 0..k {
|
|
for kx in 0..k {
|
|
let (iy, ix) = (oy as i64 + ky as i64 - pad as i64, ox as i64 + kx as i64 - pad as i64);
|
|
if (0..4).contains(&iy) && (0..4).contains(&ix) {
|
|
acc += w[((o * cin + i) * k + ky) * k + kx] * x[i * 16 + (iy * 4 + ix) as usize].clamp(-128, 127);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
let mut r = acc >> s[o];
|
|
if relu && r < 0 {
|
|
r = 0;
|
|
}
|
|
y[o * 16 + oy * 4 + ox] = r;
|
|
}
|
|
}
|
|
}
|
|
(y, cout)
|
|
};
|
|
let (x1, c1) = block(z, cz, "conv1", 1, 0, "onnx::Pow_85", true);
|
|
let (x2, c2) = block(&x1, c1, "depthwise", 3, 1, "onnx::Pow_91", true);
|
|
let (x3, _) = block(&x2, c2, "pointwise", 1, 0, "onnx::Pow_97", false);
|
|
let mut out = vec![0i64; cz * 256];
|
|
for c in 0..cz {
|
|
for y in 0..16 {
|
|
for x in 0..16 {
|
|
let src = (c * 16 + (y % 4) * 4 + (x % 4)) * 16 + (y / 4) * 4 + (x / 4);
|
|
out[(c * 16 + y) * 16 + x] = x3[src].abs().clamp(0, 3967) + gains[c];
|
|
}
|
|
}
|
|
}
|
|
out
|
|
}
|
|
|
|
/// §13.6.1 end to end against the legacy decoder's intermediates on a real codestream:
|
|
/// the latent tables, then for each plane the hyper-latent (`fac`), the scale graph,
|
|
/// `coded`, `block_stat`, `band_lookup`, `scale_index` and the latent (`sgm`). Neither
|
|
/// the bundle nor the codestream is committed: `WAI_OPS_JPEGAI_BUNDLE` names the bundle
|
|
/// and `WAI_OPS_JPEGAI_DUMP` a JSON file of the legacy decoder's values for one
|
|
/// codestream (its segment payloads and every intermediate above).
|
|
#[test]
|
|
#[ignore = "needs a bundle and a legacy dump that are not committed: set WAI_OPS_JPEGAI_BUNDLE and WAI_OPS_JPEGAI_DUMP and run with --ignored"]
|
|
fn jpegai_entropy_path_matches_the_legacy_decoder() {
|
|
let dir = std::env::var("WAI_OPS_JPEGAI_BUNDLE").expect("WAI_OPS_JPEGAI_BUNDLE names the bundle");
|
|
let dump = std::env::var("WAI_OPS_JPEGAI_DUMP").expect("WAI_OPS_JPEGAI_DUMP names the legacy dump");
|
|
let rd = |n: &str| std::fs::read(Path::new(&dir).join(n)).unwrap();
|
|
let norm = json::parse(&rd("norm.json")).unwrap();
|
|
let d = json::parse(&std::fs::read(&dump).unwrap()).unwrap();
|
|
let set = lower::jpegai_scale(&rd("norm.json"), &rd("hs_y.json"), &rd("hs_uv.json")).unwrap();
|
|
let pmf: Vec<Vec<i64>> = norm.get("pmf").unwrap().as_arr().unwrap().iter().map(i64s).collect();
|
|
let tables = stage::mans::build_latent_tables(&pmf).unwrap();
|
|
let want: Vec<Vec<i64>> = d.get("tables").unwrap().as_arr().unwrap().iter().map(i64s).collect();
|
|
for (s, (t, w)) in tables.iter().zip(&want).enumerate() {
|
|
assert_eq!(t.iter().map(|&v| v as i64).collect::<Vec<_>>(), *w, "table {s}");
|
|
}
|
|
let bounds = i64s(norm.get("bounds").unwrap());
|
|
let seg = |k: &str| -> Vec<u8> { i64s(d.get("segments").unwrap().get(k).unwrap()).iter().map(|&b| b as u8).collect() };
|
|
let rows_of = |k: &str| -> Vec<Vec<i64>> { i64s(norm.get(k).unwrap()).chunks(63).map(|r| r.to_vec()).collect() };
|
|
let (ry, ruv) = (rows_of("cdfs_y"), rows_of("cdfs_uv"));
|
|
for (plane, role, c, skip, gain_k, lat) in [("y", "hs_y", 160usize, 0usize, "gain_y", "SORP"), ("uv", "hs_uv", 96, 160, "gain_uv", "SORS")] {
|
|
let p = d.get("planes").unwrap().get(plane).unwrap();
|
|
let gains = i64s(norm.get(gain_k).unwrap());
|
|
let mut rows: Vec<Vec<i64>> = ry[..skip].to_vec();
|
|
rows.extend_from_slice(if plane == "y" { &ry[..c] } else { &ruv[..c] });
|
|
let lens = vec![16usize; skip + c];
|
|
let syms = stage::mans::decode_fac(&seg("SOZ"), &rows, &lens).unwrap();
|
|
let z: Vec<i64> = syms[skip..].iter().flatten().map(|&s| s - 31).collect();
|
|
assert_eq!(z, i64s(p.get("z").unwrap()), "{plane}: z");
|
|
let v = one(set.get(role).unwrap(), Tensor::new(c as u32, 4, 4, z), &[]);
|
|
let is: Vec<i64> = v.data.iter().enumerate().map(|(i, &x)| x - gains[i / 256]).collect();
|
|
assert_eq!(is, i64s(p.get("i_s").unwrap()), "{plane}: I_s");
|
|
let vt = Tensor::new(c as u32, 16, 16, v.data.clone());
|
|
let likely = stage::block_stat(&vt, 8, 6, 0, 4351).unwrap();
|
|
assert_eq!(likely.data, i64s(p.get("likely").unwrap()), "{plane}: likely");
|
|
let grfs: Vec<bool> = p.get("grfs").unwrap().as_arr().unwrap().iter().map(|b| b.as_bool().unwrap()).collect();
|
|
let table = vec![vec![-31, 20, 122, 181], vec![-80, -31, 71, 130]];
|
|
let rvs: Vec<i64> = likely.data.iter().enumerate().map(|(i, &l)| stage::band_lookup(l, &[1004, 1353, 1938], &table, if grfs[i / 256] { 0 } else { 1 }).unwrap()).collect();
|
|
assert_eq!(rvs, i64s(p.get("rvs").unwrap()), "{plane}: rvs");
|
|
let sigma: Vec<i64> = v.data.iter().zip(&rvs).map(|(&x, &r)| stage::scale_index(x + r, 7, 32).unwrap()).collect();
|
|
assert_eq!(sigma, i64s(p.get("sigma").unwrap()), "{plane}: sigma");
|
|
let mask: Vec<bool> = v.data.iter().map(|&x| stage::coded(x, 382)).collect();
|
|
let want_mask: Vec<bool> = p.get("mask").unwrap().as_arr().unwrap().iter().map(|b| b.as_bool().unwrap()).collect();
|
|
assert_eq!(mask, want_mask, "{plane}: mask");
|
|
let spans: Vec<stage::mans::Span> = (0..c).map(|ch| stage::mans::Span {
|
|
mask: mask[ch * 256..(ch + 1) * 256].to_vec(),
|
|
sigma: sigma[ch * 256..(ch + 1) * 256].iter().map(|&s| s as usize).collect(),
|
|
}).collect();
|
|
let yq: Vec<i64> = stage::mans::decode_sgm(&seg(lat), &pmf, &bounds[..32], &spans).unwrap().concat();
|
|
assert_eq!(yq, i64s(p.get("y_q").unwrap()), "{plane}: y_q");
|
|
eprintln!("{plane}: z, I_s, likely, rvs, sigma, mask and y_q equal the legacy decoder's");
|
|
}
|
|
}
|
|
|
|
/// E8 against streams the reference encoder wrote, and the decode tables it built:
|
|
/// `WAI_OPS_MANS_GOLDENS` names the fixture (`jpegai_mans_goldens.json`, not
|
|
/// committed).
|
|
#[test]
|
|
#[ignore = "needs a fixture that is not committed: set WAI_OPS_MANS_GOLDENS and run with --ignored"]
|
|
fn mans_matches_reference_encoder_streams() {
|
|
let f = std::env::var("WAI_OPS_MANS_GOLDENS").expect("WAI_OPS_MANS_GOLDENS names the fixture");
|
|
let v = json::parse(&std::fs::read(f).unwrap()).unwrap();
|
|
let pmf: Vec<Vec<i64>> = v.get("pmf").unwrap().as_arr().unwrap().iter().map(i64s).collect();
|
|
let built = stage::mans::build_latent_tables(&pmf).unwrap();
|
|
let want: Vec<Vec<i64>> = v.get("decode_transitions").unwrap().as_arr().unwrap().iter().map(i64s).collect();
|
|
for (t, w) in built.iter().zip(&want) {
|
|
assert_eq!(t.iter().map(|&x| x as i64).collect::<Vec<_>>(), *w);
|
|
}
|
|
let bounds = i64s(v.get("bounds").unwrap());
|
|
let bytes = |x: &Value| -> Vec<u8> { i64s(x).iter().map(|&b| b as u8).collect() };
|
|
let y = v.get("y").unwrap();
|
|
let mask: Vec<bool> = y.get("masks").unwrap().as_arr().unwrap().iter().map(|b| b.as_bool().unwrap()).collect();
|
|
let sigma: Vec<usize> = i64s(y.get("sigma").unwrap()).iter().map(|&s| s as usize).collect();
|
|
let got = stage::mans::decode_sgm(&bytes(y.get("memory").unwrap()), &pmf, &bounds, &[stage::mans::Span { mask: mask.clone(), sigma }]).unwrap();
|
|
let exp = i64s(y.get("expected").unwrap());
|
|
for i in 0..exp.len() {
|
|
if mask[i] {
|
|
assert_eq!(got[0][i], exp[i] as i16 as i64, "y symbol {i}");
|
|
}
|
|
}
|
|
let z = v.get("z").unwrap();
|
|
let rows: Vec<Vec<i64>> = z.get("cdfs").unwrap().as_arr().unwrap().iter().map(i64s).collect();
|
|
let exp: Vec<Vec<i64>> = z.get("expected").unwrap().as_arr().unwrap().iter().map(i64s).collect();
|
|
let lens: Vec<usize> = exp.iter().map(Vec::len).collect();
|
|
let got = stage::mans::decode_fac(&bytes(z.get("memory").unwrap()), &rows, &lens).unwrap();
|
|
assert_eq!(got, exp);
|
|
eprintln!("mans: tables, the sgm stream and the fac stream equal the reference encoder's");
|
|
}
|