open-standards/wai/wai-ops/tests/goldens.rs

//! 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() {
    int_mlicpp_digest_of("wai-rs/tests/vectors/int_mlicpp");
}

/// `natural`, the `wai.neural.int_mlicpp` set trained on openly licensed natural pictures
/// (`int_mlicpp_natural/natural_set.json`), decoded by the same stage operators and graphs.
#[test]
fn trained_int_mlicpp_reproduces_its_digest() {
    int_mlicpp_digest_of("wai-rs/tests/vectors/int_mlicpp_natural");
}

fn int_mlicpp_digest_of(d: &str) {
    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");
}