open-standards/wai/wai-rs/tests/jpegai_nets_ops.rs

// wai-conformance: int-kernels int-kernels-amd
//! `wai.image.jpegai`'s other integer networks through the integer operator set: the
//! hyper decoder, the four context stages and the synthesis, as `jpegai-int/1` lowers
//! them (`extensions/integer-operators.md` §13.6.2) to the graphs `rec_y`, `rec_uv`,
//! `syn_y` and `syn_uv`, and as this crate loads them (`jpegai_ops::Networks`).
//!
//! The committed bundle (`ops-conformance/sets/jpegai`, synthetic, at small widths) has
//! its lowered graphs and certificate under `ops-conformance/certificates/`, where
//! `wai-verify` and the independent checker (`tools/wai_ops_check.py`) reproduce them
//! byte for byte. The `jpegai/v1` bundle of the models root is not committed; its
//! certificate is, by digest, and the tests that read it take its files and the test
//! codestream from `WAI_JPEGAI_WEIGHTS_DIR`, `WAI_JPEGAI_NORM` and
//! `WAI_JPEGAI_CODESTREAM` (`tools/wai_jpegai_fetch_e2e_inputs.sh`), with
//! `WAI_REQUIRE_SOTA_VECTORS` deciding whether their absence is a skip or a failure.
#![cfg(all(feature = "neural_int", feature = "legacy_int"))]

#[path = "support/require_vectors.rs"]
mod require_vectors;

use std::collections::BTreeMap;
use std::path::{Path, PathBuf};

use serde_json::Value;
use wai::jpegai_decode::{DecoderWeights, NormTables, PlaneWeights};
use wai::jpegai_ops::{self, Networks, Plane};
use wai::jpegai_synthesis::{IntNet, Weights};
use wai_ops::lower::jpegai_nets::NET_MEMBERS;
use wai_ops::{Profile, Tensor};

const CAP: &str = "wai.image.jpegai";
const SYNTHETIC: &str = "1bc0daae512797a8817f2e77e57d009ace0d22acf4fe6945af8f4e0ee95244b9";
const V1: &str = "3e243dc11566922f5cb7f4c628afd440b18d56a08b99b07cc33d5c8522aeec37";
const PRODUCE: &str = "tools/wai_jpegai_fetch_e2e_inputs.sh <dir> (pinned by tools/jpegai-e2e-v1.sha256)";
const ROLES: [&str; 4] = ["rec_y", "rec_uv", "syn_y", "syn_uv"];

fn wai() -> PathBuf {
    PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("..")
}

fn read(p: &Path) -> Vec<u8> {
    std::fs::read(p).unwrap_or_else(|e| panic!("{}: {e}", p.display()))
}

fn json(b: &[u8]) -> Value {
    serde_json::from_slice(b).unwrap()
}

fn sha256(b: &[u8]) -> String {
    wai_ops::sha256::sha256(b).iter().map(|x| format!("{x:02x}")).collect()
}

fn i64s(v: &Value) -> Vec<i64> {
    v.as_array().unwrap().iter().map(|x| x.as_i64().unwrap()).collect()
}

/// A bundle's members as bytes, by their names in the set (`norm.json`, `weights/…`).
struct Members(BTreeMap<String, Vec<u8>>);

impl Members {
    fn synthetic() -> Members {
        let d = wai().join("ops-conformance/sets/jpegai");
        let mut m = BTreeMap::new();
        for n in ["norm.json", "weights/hs_y.json", "weights/hs_uv.json"].into_iter().chain(NET_MEMBERS) {
            m.insert(n.to_string(), read(&d.join(n)));
        }
        Members(m)
    }
    /// The `jpegai/v1` bundle: its twelve members, from the variables the whole-decode
    /// tests read (a directory of `<name>.json`, and `norm.json`).
    fn pinned(test: &str) -> Option<Members> {
        let [dir, norm] = require_vectors::env_paths(test, ["WAI_JPEGAI_WEIGHTS_DIR", "WAI_JPEGAI_NORM"], PRODUCE)?;
        let d = PathBuf::from(dir);
        let mut m = BTreeMap::new();
        m.insert("norm.json".to_string(), read(Path::new(&norm)));
        for n in ["hs_y", "hs_uv", "hyp_y", "hyp_uv", "synth_y", "synth_uv", "mcm0_y", "mcm1_y", "mcm2_y", "mcm3_y", "icci"] {
            m.insert(format!("weights/{n}.json"), read(&d.join(format!("{n}.json"))));
        }
        Some(Members(m))
    }
    fn get(&self, n: &str) -> &[u8] {
        &self.0[n]
    }
    fn bytes(&self) -> wai_ops::lower::JpegaiBytes<'_> {
        wai_ops::lower::JpegaiBytes::from_members(|n| self.0.get(n).map(Vec::as_slice)).expect("every member the lowering reads")
    }
    fn nets(&self) -> [&[u8]; 8] {
        NET_MEMBERS.map(|n| self.get(n))
    }
    fn digest(&self) -> String {
        let m: BTreeMap<String, &[u8]> = self.0.iter().map(|(k, v)| (k.clone(), v.as_slice())).collect();
        wai::int_set_cache::SetKey::listing_digest(&m).unwrap().iter().map(|x| format!("{x:02x}")).collect()
    }
    /// The networks as this crate parses them (`Weights::from_json`).
    fn weights(&self, n: &str) -> Weights {
        Weights::from_json(&json(self.get(n))).unwrap()
    }
}

/// Gate 1, on one bundle: the lowering of its bytes writes `hs_y`, `hs_uv`, then the four
/// network graphs; the networks as this crate parses them, converted
/// (`jpegai_ops::load_networks`), are the graphs `wai-ops` lowers from the bytes; the
/// committed certificate `digest` is those six graphs', with the verdict `acc64`, and
/// `acc32` for the scale path alone (§11's `role_profiles`).
fn lowers_to_its_certificate(name: &str, m: &Members, digest: &str) {
    let lowered = wai_ops::lower::jpegai(&m.bytes()).unwrap_or_else(|r| panic!("{name}: does not lower: {r:?}"));
    let roles: Vec<&str> = lowered.roles.iter().map(|(r, _)| *r).collect();
    assert_eq!(roles, ["hs_y", "hs_uv", "rec_y", "rec_uv", "syn_y", "syn_uv"], "{name}: the roles, in §13.0's order");
    let w = |n: &str| m.weights(n);
    let dw = DecoderWeights {
        luma: PlaneWeights { hyper_scale: IntNet::from_json(&json(m.get("weights/hs_y.json"))).unwrap(), hyper_decoder: w("weights/hyp_y.json"), synthesis: w("weights/synth_y.json") },
        chroma: PlaneWeights { hyper_scale: IntNet::from_json(&json(m.get("weights/hs_uv.json"))).unwrap(), hyper_decoder: w("weights/hyp_uv.json"), synthesis: w("weights/synth_uv.json") },
        mcm_y: [w("weights/mcm0_y.json"), w("weights/mcm1_y.json"), w("weights/mcm2_y.json"), w("weights/mcm3_y.json")],
        icci: wai::jpegai_icci::IcciNet { luma: empty_path(), chroma: empty_path(), conv_last_v: empty_conv(), scale_v: 0.0 },
    };
    let parsed = jpegai_ops::nets_of(&dw).unwrap_or_else(|r| panic!("{name}: does not convert: {r:?}"));
    assert!(parsed == wai_ops::lower::jpegai_nets::parse_nets(m.nets()).unwrap(), "{name}: the converted networks are the ones wai-ops parses from the bytes");
    let nets = jpegai_ops::load_networks(&dw).unwrap_or_else(|r| panic!("{name}: refused under acc64: {r:?}"));
    let from_bytes = jpegai_ops::load_networks_bytes(m.nets()).unwrap();
    let cert = json(&read(&wai().join(format!("ops-conformance/certificates/{CAP}/{digest}.json"))));
    assert_eq!(cert["set_sha256"].as_str(), Some(digest), "{name}");
    for (i, role) in ROLES.iter().enumerate() {
        let g = lowered.get(role).unwrap();
        assert!(*nets.graph(role).unwrap().graph() == *g, "{name}: {role} as loaded from the parsed networks");
        assert!(*from_bytes.graph(role).unwrap().graph() == *g, "{name}: {role} as loaded from the bytes");
        assert!(!nets.graph(role).unwrap().meets(Profile::Acc32), "{name}: {role} meets acc64 and not acc32");
        let c = &cert["graphs"][i + 2];
        assert_eq!(c["name"], *role, "{name}");
        assert_eq!(c["graph_sha256"].as_str(), Some(sha256(&g.to_bytes()).as_str()), "{name}: the certificate is {role}'s");
        let lo = (i32::MIN as i64).to_string();
        let hi = (i32::MAX as i64).to_string();
        for x in c["inputs"].as_array().unwrap() {
            let range = (x["lo"].as_str().unwrap(), x["hi"].as_str().unwrap());
            assert!(x["policy"] == "check" && (range == ("-31", "32") || range == (lo.as_str(), hi.as_str())), "{name}: {role}'s input ranges: {range:?}");
        }
    }
    for (i, role) in ["hs_y", "hs_uv"].iter().enumerate() {
        let c = &cert["graphs"][i];
        assert_eq!(c["name"], *role);
        assert!(c["nodes"].as_array().unwrap().iter().all(|n| n["fits"].as_array().unwrap().iter().any(|p| p == "acc32")), "{name}: every node of {role} fits acc32");
    }
    let v = &cert["verdict"];
    assert_eq!((v["acc64"].as_bool(), v["acc32"].as_bool()), (Some(true), Some(false)), "{name}: the set meets acc64 and not acc32");
    assert_eq!(v["acc32_first_failure"]["graph"], "rec_y", "{name}: the first acc32 failure lies past the scale path");
}

fn empty_conv() -> wai::jpegai_icci::Conv {
    wai::jpegai_icci::Conv { w: vec![], b: vec![], cin: 0, cout: 0, kh: 0, kw: 0, ph: 0, pw: 0 }
}
fn empty_path() -> wai::jpegai_icci::IcciPath {
    wai::jpegai_icci::IcciPath { conv_first: empty_conv(), blocks: vec![], conv_last: empty_conv(), scale: 0.0 }
}

/// The committed synthetic bundle: its four network graphs are the committed ones, and
/// its set digest is the listing of its eleven members.
#[test]
fn the_committed_bundle_lowers_to_its_committed_networks() {
    let m = Members::synthetic();
    assert_eq!(m.digest(), SYNTHETIC, "the set's digest");
    let dir = wai().join(format!("ops-conformance/certificates/{CAP}/{SYNTHETIC}"));
    let lowered = wai_ops::lower::jpegai(&m.bytes()).unwrap();
    for role in ROLES {
        assert!(lowered.get(role).unwrap().to_bytes() == read(&dir.join(format!("{role}.whg1"))), "the lowering of the bytes differs from the committed {role}.whg1");
    }
    lowers_to_its_certificate("synthetic", &m, SYNTHETIC);
}

/// The `jpegai/v1` bundle: its networks lower to its committed certificate's graphs.
#[test]
fn the_pinned_bundle_lowers_its_networks_to_its_certificate() {
    let Some(m) = Members::pinned("jpegai_nets_ops::the_pinned_bundle_lowers_its_networks_to_its_certificate") else {
        return;
    };
    assert_eq!(m.digest(), V1, "the set's digest");
    lowers_to_its_certificate("jpegai/v1", &m, V1);
    eprintln!("jpegai/v1: rec_y, rec_uv, syn_y and syn_uv lower to the committed certificate {V1} and meet acc64");
}

/// The whole `jpegai/v1` bundle and the test codestream, as the decoders parse them.
struct Whole {
    data: Vec<u8>,
    weights: DecoderWeights,
    norm: NormTables,
}

impl Whole {
    fn pinned(test: &str) -> Option<Whole> {
        let [cs, dir, nf] = require_vectors::env_paths(test, ["WAI_JPEGAI_CODESTREAM", "WAI_JPEGAI_WEIGHTS_DIR", "WAI_JPEGAI_NORM"], PRODUCE)?;
        let d = PathBuf::from(dir);
        let j = |n: &str| json(&read(&d.join(n)));
        let w = |n: &str| Weights::from_json(&j(n)).unwrap();
        let net = |n: &str| IntNet::from_json(&j(n)).unwrap();
        let weights = DecoderWeights {
            luma: PlaneWeights { hyper_scale: net("hs_y.json"), hyper_decoder: w("hyp_y.json"), synthesis: w("synth_y.json") },
            chroma: PlaneWeights { hyper_scale: net("hs_uv.json"), hyper_decoder: w("hyp_uv.json"), synthesis: w("synth_uv.json") },
            mcm_y: [w("mcm0_y.json"), w("mcm1_y.json"), w("mcm2_y.json"), w("mcm3_y.json")],
            icci: wai::jpegai_icci::IcciNet::from_json(&j("icci.json")).unwrap(),
        };
        let nv = json(&read(Path::new(&nf)));
        let three = |k: &str| -> [f64; 3] { std::array::from_fn(|i| nv[k][i].as_f64().unwrap()) };
        let norm = NormTables {
            pmf: nv["pmf"].as_array().unwrap().iter().map(i64s).collect(),
            bounds: i64s(&nv["bounds"]),
            per_channel: nv["per_channel"].as_u64().unwrap() as usize,
            cdfs_y: i64s(&nv["cdfs_y"]),
            cdfs_uv: i64s(&nv["cdfs_uv"]),
            gain_y: i64s(&nv["gain_y"]),
            gain_uv: i64s(&nv["gain_uv"]),
            lsbs_t0: i64s(&nv["lsbs_t0"]).iter().map(|&x| x as i32).collect(),
            lsbs_t1: i64s(&nv["lsbs_t1"]).iter().map(|&x| x as i32).collect(),
            lef_thr: three("lef_thr0").map(|x| x as i64),
            lef_mag: three("lef_mag0").map(|x| x as f32),
        };
        Some(Whole { data: read(Path::new(&cs)), weights, norm })
    }
}

/// What the networks of §13.6.2 read from the test codestream, and what the direct
/// decode makes of it: for each plane `z`, the dequantised residual in Q16 (the input
/// the float stages hand `rec_*`), the direct decode's `ŷ`, and the LSBS-refined latent
/// in Q16 (the input of `syn_*`); then the direct synthesis's planes, as the float
/// values the decode reports.
struct Inputs {
    z: [Vec<i64>; 2],
    resi: [Vec<i64>; 2],
    yhat: [Vec<i64>; 2],
    latent: [Vec<i64>; 2],
    syn_y: Vec<f32>,
    syn_uv: Vec<f32>,
}

fn q16(f: &[f32]) -> Vec<i64> {
    f.iter().map(|&x| wai::float_repro::round_half_away(x as f64 * 65536.0) as i64).collect()
}

/// The decode's own steps up to each network (`jpegai_decode`'s, restated), with the
/// direct networks (`jpegai_synthesis`, `jpegai_pipeline`).
fn inputs(wh: &Whole) -> Inputs {
    use wai::jpegai_bitstream::{find_segment, parse_container, parse_picture_header, Marker};
    use wai::jpegai_pipeline::{dequantize_resi, gain_unit_scaler_vec, lsbs_refine, reconstruct_yhat, reconstruct_yhat_uv};
    use wai::jpegai_synthesis::{decode_hyper, decode_uv_f32, decode_y_f32, q16_from_symbols};
    let sp = wai::jpegai_decode::load(&wh.weights, &wh.norm).unwrap();
    let segs = parse_container(&wh.data);
    let pay = |m| {
        let s = find_segment(&segs, m).unwrap();
        wh.data[s.start..s.start + s.len].to_vec()
    };
    let ph = parse_picture_header(&wh.data, find_segment(&segs, Marker::Pih).unwrap().start_bit());
    let mut out: Vec<(Vec<i64>, Vec<i64>, Vec<i64>, Vec<i64>)> = vec![];
    for (plane, lat_m, pw, gain) in [(Plane::Luma, Marker::Sorp, &wh.weights.luma, &wh.norm.gain_y), (Plane::Chroma, Marker::Sors, &wh.weights.chroma, &wh.norm.gain_uv)] {
        let c = plane.channels();
        let mut grfs = ph.components[plane as usize].grfs_channel_flags.clone();
        grfs.resize(c, false);
        let ints = jpegai_ops::plane(&sp, plane, &pay(Marker::Soz), &pay(lat_m), &grfs).unwrap();
        let y_q: Vec<i16> = ints.y_q.iter().map(|&q| q as i16).collect();
        let resi = dequantize_resi(&y_q, &ints.likely, &grfs, &gain_unit_scaler_vec(gain, 0), c);
        let resi_q = q16(&resi);
        let (psi, ..) = decode_hyper(&q16_from_symbols(&ints.z), c, 4, 4, &pw.hyper_decoder);
        let yhat = if plane == Plane::Luma {
            let m = &wh.weights.mcm_y;
            reconstruct_yhat(&resi_q, &psi, c, 16, 16, [&m[0], &m[1], &m[2], &m[3]])
        } else {
            reconstruct_yhat_uv(&resi_q, &psi, c, 16, 16)
        };
        let yf: Vec<f32> = yhat.iter().map(|&v| v as f64 as f32 / 65536.0).collect();
        let lat = lsbs_refine(&yf, &resi, &ints.likely, &wh.norm.lsbs_t0, &wh.norm.lsbs_t1, c, 16, 16);
        out.push((ints.z, resi_q, yhat, q16(&lat)));
    }
    let [(zy, ry, yy, ly), (zuv, ruv, yuv, luv)]: [_; 2] = out.try_into().unwrap();
    let syn_y = decode_y_f32(&ly, &wh.weights.luma.synthesis);
    let mut cat = ly.clone();
    cat.extend_from_slice(&luv);
    let syn_uv = decode_uv_f32(&cat, &wh.weights.chroma.synthesis);
    Inputs { z: [zy, zuv], resi: [ry, ruv], yhat: [yy, yuv], latent: [ly, luv], syn_y, syn_uv }
}

fn floats(q: &[i64]) -> Vec<f32> {
    q.iter().map(|&v| v as f32 / 65536.0).collect()
}

/// Gate 2, the networks: on the test codestream, each plane's `ŷ` through `rec_y` and
/// `rec_uv` is the direct hyper decoder's and context stages', every value, and the
/// planes through `syn_y` and `syn_uv` are the direct synthesis's, bit for bit as the
/// float values the decode reports.
#[test]
fn the_networks_decode_through_the_lowering_as_directly() {
    let Some(wh) = Whole::pinned("jpegai_nets_ops::the_networks_decode_through_the_lowering_as_directly") else {
        return;
    };
    let i = inputs(&wh);
    let nets = jpegai_ops::load_networks(&wh.weights).unwrap();
    for (k, plane) in [Plane::Luma, Plane::Chroma].into_iter().enumerate() {
        let yhat = jpegai_ops::reconstruct(&nets, plane, i.z[k].clone(), i.resi[k].clone()).unwrap();
        assert!(yhat == i.yhat[k], "{plane:?}: ŷ through the lowering differs from the direct hyper decoder and context stages");
        eprintln!("{plane:?}: ŷ ({} values) equals the direct decode's", yhat.len());
    }
    let y = floats(&jpegai_ops::synthesize_y(&nets, i.latent[0].clone()).unwrap());
    assert!(y.iter().map(|v| v.to_bits()).eq(i.syn_y.iter().map(|v| v.to_bits())), "Y through syn_y differs from the direct synthesis");
    let uv = floats(&jpegai_ops::synthesize_uv(&nets, i.latent[0].clone(), i.latent[1].clone()).unwrap());
    assert!(uv.iter().map(|v| v.to_bits()).eq(i.syn_uv.iter().map(|v| v.to_bits())), "U, V through syn_uv differ from the direct synthesis");
    eprintln!("Y ({}) and U, V ({}) through the synthesis graphs equal the direct synthesis's", y.len(), uv.len());
}

/// The pinned decode's stage hashes (`tests/jpegai_cross_platform.rs`' `PINNED`): the
/// float values after each stage, and the RGB bytes, of the `jpegai/v1` decode.
const PINNED: [(&str, &str); 8] = [
    ("latent_y", "5889f3806308496c64850aee04f3706178f92d1e0edcda1e6975fe256779061e"),
    ("latent_uv", "be89a76f726ef362f267aba3dee8f61594cc512576a3864f52f665cc5b6ab77b"),
    ("synthesis", "d56695726bb6ebed403199f60a79a5698734eef1e5ca2813859d490c0882d2a4"),
    ("efl", "1783ea77bd4f6b0ac6a599137577a18db2a86753aa3eb3d79ca987e1cf25792e"),
    ("icci", "3566c51c4042a749d7826df812663c4e3ff7a9e807cbd0c3bb0545c9dbd79f2b"),
    ("efn", "09dc833c7a6b91ce08da20adb8bbc55e2eeedba7d8a2dc4b133ab0915b267de7"),
    ("lef", "e2b20c72ac6a59dd12067f75562c82715061f6d7a3444ca075fc0b11e3f56468"),
    ("rgb", "dbc494a5a670914edf55c07c48afc0dd348cd05d3a82fff293bf384b2bb69ea3"),
];

/// Gate 2, the decode: the whole decode with every integer stage through the operator
/// set (the scale path, the entropy decodes and the four network graphs) gives the
/// direct decode's float values after every stage, bit for bit, and its RGB bytes; and
/// those are the pinned decode's: the `Synthesis` stage hash and the RGB SHA-256 are
/// unchanged.
#[test]
fn the_whole_decode_through_the_networks_is_the_pinned_decode() {
    let Some(wh) = Whole::pinned("jpegai_nets_ops::the_whole_decode_through_the_networks_is_the_pinned_decode") else {
        return;
    };
    let bits = |planes: &[&[f32]]| -> String { sha256(&planes.iter().flat_map(|p| p.iter().flat_map(|v| v.to_bits().to_le_bytes())).collect::<Vec<u8>>()) };
    let mut direct = vec![];
    let a = wai::jpegai_decode::legacy::decode_with_stages(&wh.data, &wh.weights, &wh.norm, &mut |s, p| direct.push((s.name(), bits(p)))).expect("the direct decode");
    direct.push(("rgb", sha256(&a.data)));
    let sp = wai::jpegai_decode::load(&wh.weights, &wh.norm).unwrap();
    let nets = jpegai_ops::load_networks(&wh.weights).unwrap();
    let mut ops = vec![];
    let b = wai::jpegai_decode::decode_loaded(&wh.data, &wh.weights.icci, &wh.norm, &sp, &nets, &mut |s, p| ops.push((s.name(), bits(p)))).expect("the decode through the operator set");
    // The decode a sink makes, which loads the bundle itself, is the same.
    let c = wai::jpegai_decode::decode(&wh.data, &wh.weights, &wh.norm).expect("the decode");
    assert!(c.data == b.data, "jpegai_decode::decode");
    ops.push(("rgb", sha256(&b.data)));
    assert_eq!(ops, direct, "the stages through the networks' graphs");
    let pinned: Vec<(&str, String)> = PINNED.iter().map(|(n, h)| (*n, h.to_string())).collect();
    assert_eq!(ops, pinned, "the decode through the networks' graphs is not the pinned decode");
    eprintln!("every stage hash ({}) through the networks' graphs is the pinned decode's, the synthesis {} and the RGB {} among them", ops.len(), PINNED[2].1, PINNED[7].1);
}

/// The accumulator class, measured (§13.6.2's table): the four graphs run on the test
/// codestream's inputs through the interpreter alone, with every value each operator
/// forms held to its certificate's interval (`wai_ops::bounds::Within`), and the largest
/// magnitude each kind of sum reaches, per graph, printed beside the certificate's bound.
/// Nothing a decode forms leaves its interval, and no sum passes `2^63`.
#[test]
fn the_networks_keep_to_their_certificate_on_the_test_codestream() {
    let Some(wh) = Whole::pinned("jpegai_nets_ops::the_networks_keep_to_their_certificate_on_the_test_codestream") else {
        return;
    };
    let i = inputs(&wh);
    let nets = jpegai_ops::load_networks(&wh.weights).unwrap();
    let runs: [(&str, Vec<Tensor>); 4] = [
        ("rec_y", vec![Tensor::new(160, 4, 4, i.z[0].clone()), Tensor::new(160, 16, 16, i.resi[0].clone())]),
        ("rec_uv", vec![Tensor::new(96, 4, 4, i.z[1].clone()), Tensor::new(96, 16, 16, i.resi[1].clone())]),
        ("syn_y", vec![Tensor::new(160, 16, 16, i.latent[0].clone())]),
        ("syn_uv", vec![Tensor::new(160, 16, 16, i.latent[0].clone()), Tensor::new(96, 16, 16, i.latent[1].clone())]),
    ];
    const SUMS: [&str; 9] = ["conv.acc", "convT.acc", "rq.t", "rq.p", "mul.p", "gram.acc", "apply.acc", "ln.T", "l2.T"];
    for (role, x) in runs {
        let g = nets.graph(role).unwrap();
        let sizes = wai_ops::plan::load_plan(g.graph()).unwrap();
        let b = wai_ops::bounds::pass(g.graph(), sizes.as_deref()).unwrap();
        let mut w = wai_ops::bounds::Within::new(&b);
        wai_ops::interp::run_observed(g, &x, &[], &mut wai_ops::Sequence::new(), &mut w).unwrap_or_else(|r| panic!("{role}: {r:?}"));
        assert!(w.outside.is_empty(), "{role}: values outside the certificate: {:?}", w.outside);
        assert!(w.unrecorded.is_empty(), "{role}: keys the certificate does not record: {:?}", w.unrecorded);
        // The largest magnitude seen under each sum, over the graph's nodes, with the
        // certificate's bound for the same key over the same nodes.
        for key in SUMS {
            let mut seen: Option<(u128, u32)> = None;
            let mut bound: u128 = 0;
            for (k, n) in b.nodes.iter().enumerate() {
                if let Some((lo, hi)) = w.extreme(k as u32, key) {
                    let m = lo.unsigned_abs().max(hi.unsigned_abs());
                    if seen.is_none_or(|(s, _)| m > s) {
                        seen = Some((m, k as u32));
                    }
                    if let Some(iv) = n.inter.iter().find(|x| x.key == key) {
                        bound = bound.max(iv.iv.bits() as u128);
                    }
                }
            }
            if let Some((m, at)) = seen {
                let bits = 128 - m.leading_zeros();
                assert!(m < 1u128 << 63 || key.starts_with("rq.p") || key.starts_with("ln.") || key.starts_with("l2."), "{role}: {key} reaches {m} at node {at}");
                eprintln!("{role:7} {key:10} largest |value| {m:>22} ({bits:>2} bits) at node {at:>3}; certified within {bound} bits");
            }
        }
        // Every activation the graph writes (each producer result), and its inputs.
        let mut act: u128 = 0;
        for k in 0..b.nodes.len() as u32 {
            if let Some((lo, hi)) = w.extreme(k, "out.0") {
                act = act.max(lo.unsigned_abs()).max(hi.unsigned_abs());
            }
        }
        let inp = x.iter().flat_map(|t| t.data.iter()).map(|v| v.unsigned_abs() as u128).max().unwrap_or(0);
        assert!(act < 1 << 31 && inp < 1 << 31, "{role}: an activation or an input reaches 2^31");
        eprintln!("{role:7} largest |input| {inp} ({} bits), largest |activation| {act} ({} bits); every one within i32", 128 - inp.leading_zeros(), 128 - act.leading_zeros());
        eprintln!("{role:7} {} values held to the certificate", w.seen);
    }
}

/// Gate 5, the fast kernels: under every kernel policy, each of the four graphs gives the
/// interpreter's values, on the committed synthetic networks with inputs over their
/// ranges at small sizes (in the fast-kernel steps) and, where the bundle is supplied, on
/// the test codestream's inputs. (`jpegai_ops`' own test counts the convolutions that
/// ran on the fast kernel.)
#[test]
fn every_kernel_policy_runs_the_networks_alike() {
    use wai::int_kernels::{with_policy, Policy};
    static POLICIES: std::sync::LazyLock<Vec<Policy>> = std::sync::LazyLock::new(wai::int_kernels::equivalence_policies);
    let m = Members::synthetic();
    let syn = jpegai_ops::load_networks_bytes(m.nets()).unwrap();
    let mut s: u64 = 0x0D2C_007B;
    let mut range = |n: usize, lo: i64, hi: i64| -> Vec<i64> {
        (0..n).map(|_| {
            s = s.wrapping_mul(6364136223846793005).wrapping_add(1442695040888963407);
            lo + ((s >> 11) % (hi - lo + 1) as u64) as i64
        }).collect()
    };
    let mut cases: Vec<(&Networks, &str, Vec<Tensor>)> = vec![];
    for k in 0..3 {
        let a = 1i64 << (16 + 4 * k);
        cases.push((&syn, "rec_y", vec![Tensor::new(160, 1, 1, range(160, -31, 32)), Tensor::new(160, 4, 4, range(160 * 16, -a, a))]));
        cases.push((&syn, "rec_uv", vec![Tensor::new(96, 2, 1, range(192, -31, 32)), Tensor::new(96, 8, 4, range(96 * 32, -a, a))]));
        cases.push((&syn, "syn_y", vec![Tensor::new(160, 2, 2, range(640, -a, a))]));
        cases.push((&syn, "syn_uv", vec![Tensor::new(160, 2, 2, range(640, -a, a)), Tensor::new(96, 2, 2, range(384, -a, a))]));
    }
    let pinned = Whole::pinned("jpegai_nets_ops::every_kernel_policy_runs_the_networks_alike");
    let real = pinned.as_ref().map(|wh| (inputs(wh), jpegai_ops::load_networks(&wh.weights).unwrap()));
    if let Some((i, n)) = &real {
        cases.push((n, "rec_y", vec![Tensor::new(160, 4, 4, i.z[0].clone()), Tensor::new(160, 16, 16, i.resi[0].clone())]));
        cases.push((n, "rec_uv", vec![Tensor::new(96, 4, 4, i.z[1].clone()), Tensor::new(96, 16, 16, i.resi[1].clone())]));
        cases.push((n, "syn_y", vec![Tensor::new(160, 16, 16, i.latent[0].clone())]));
        cases.push((n, "syn_uv", vec![Tensor::new(160, 16, 16, i.latent[0].clone()), Tensor::new(96, 16, 16, i.latent[1].clone())]));
    }
    let (mut ran, mut refused) = (0, 0);
    for (n, role, x) in &cases {
        let g = n.graph(role).unwrap();
        let run = || wai_ops::run_with(g, x.clone(), &[], &mut wai_ops::Sequence::new(), &wai::int_ops::Host).map(|mut o| o.swap_remove(0).data);
        let want = with_policy(POLICIES[0], run);
        for p in &POLICIES[1..] {
            assert!(with_policy(*p, run) == want, "{role} {p:?}: the outputs differ from the interpreter alone");
        }
        ran += 1;
        refused += want.is_err() as usize;
    }
    eprintln!("{ran} runs of the four graphs alike under every policy ({refused} refused alike){}", if real.is_some() { ", the test codestream's among them" } else { "" });
}

/// One network file of the committed synthetic bundle changed by `edit` (the file's JSON
/// object, its members by name), as bytes; and whether this crate's parser
/// (`Weights::from_json`) reads it, which it does only where every member of kind `conv`
/// or `q16` has its fields of their types (and otherwise refuses it).
fn edited_net(file: &str, edit: impl Fn(&mut serde_json::Map<String, Value>)) -> (Vec<u8>, bool) {
    let m = Members::synthetic();
    let mut v = json(m.get(file));
    edit(v.as_object_mut().unwrap());
    let typed = v.as_object().unwrap().values().all(|e| match e.get("kind").and_then(Value::as_str) {
        Some("conv") => e["w"].as_array().is_some_and(|a| a.iter().all(Value::is_i64)) && e["dims"].as_array().is_some_and(|a| a.iter().all(Value::is_u64)) && e["mult"].is_i64() && e["shift"].is_u64(),
        Some("q16") => e["v"].as_array().is_some_and(|a| a.iter().all(Value::is_i64)),
        _ => true,
    });
    (serde_json::to_vec(&v).unwrap(), typed)
}

/// The eight network files of the committed bundle with `file` replaced by `b`.
fn nets_with<'a>(m: &'a Members, file: &str, b: &'a [u8]) -> [&'a [u8]; 8] {
    NET_MEMBERS.map(|n| if n == file { b } else { m.get(n) })
}

/// Gate 3, the bundle: each load check of §13.6.2, broken one change at a time in the
/// committed bundle's network files, is refused `model_mismatch` by the load from the
/// files' bytes, and by this crate's parser or the load from the networks as it parses
/// them. The direct decoder made none of these checks (integer-operators §15.3 item 4):
/// its loaders panicked on a mistyped member, and it computed with whatever shapes, shifts
/// and values the files held.
#[test]
fn every_load_check_of_the_networks_refuses_model_mismatch() {
    use serde_json::json as j;
    let m = Members::synthetic();
    let conv = |o: &mut serde_json::Map<String, Value>, k: &str| -> serde_json::Map<String, Value> { o[k].as_object().unwrap().clone() };
    let cases: Vec<(&str, &str, (Vec<u8>, bool))> = vec![
        ("a weight member missing", "weights/synth_y.json", edited_net("weights/synth_y.json", |o| { o.remove("iact2.conv2.weight"); })),
        ("a layer norm's γ missing", "weights/synth_uv.json", edited_net("weights/synth_uv.json", |o| { o.remove("TAM.TABs.1.ffn.norm1.weight"); })),
        ("a weight member of kind q16", "weights/hyp_y.json", edited_net("weights/hyp_y.json", |o| { let mut c = conv(o, "conv3.weight"); c.insert("kind".into(), j!("q16")); c.insert("v".into(), j!([1])); o.insert("conv3.weight".into(), Value::Object(c)); })),
        ("weights of other dims", "weights/hyp_uv.json", edited_net("weights/hyp_uv.json", |o| { o["conv1.weight"]["dims"] = j!([2, 95, 1, 1]); })),
        ("one weight too few", "weights/mcm1_y.json", edited_net("weights/mcm1_y.json", |o| { o["conv.0.weight"]["w"].as_array_mut().unwrap().pop(); })),
        ("a kernel that is not the topology's", "weights/synth_y.json", edited_net("weights/synth_y.json", |o| { o["TAM.ds_conv.weight"]["dims"] = j!([4, 4, 1, 1]); o["TAM.ds_conv.weight"]["w"] = j!(vec![1; 16]); })),
        ("a gate convolution whose groups are not 4", "weights/synth_y.json", edited_net("weights/synth_y.json", |o| { o["iact3.conv.weight"]["dims"] = j!([4, 2, 3, 3]); o["iact3.conv.weight"]["w"] = j!(vec![1; 72]); })),
        ("a context convolution whose groups are not 5", "weights/mcm2_y.json", edited_net("weights/mcm2_y.json", |o| { o["conv.1.weight"]["dims"] = j!([5, 5, 3, 3]); o["conv.1.weight"]["w"] = j!(vec![1; 225]); })),
        ("a bias one value short", "weights/synth_uv.json", edited_net("weights/synth_uv.json", |o| { o["conv3_t.bias"]["v"].as_array_mut().unwrap().pop(); })),
        ("a shift of 128", "weights/mcm3_y.json", edited_net("weights/mcm3_y.json", |o| { o["fusion_pred_net.conv2.weight"]["shift"] = j!(128); })),
        ("a shift of -1", "weights/mcm0_y.json", edited_net("weights/mcm0_y.json", |o| { o["fusion_pred_net.conv1.weight"]["shift"] = j!(-1); })),
        ("a weight that is not an integer", "weights/hyp_y.json", edited_net("weights/hyp_y.json", |o| { o["conv2.weight"]["w"][0] = j!(0.5); })),
        ("a multiplier past i64", "weights/synth_y.json", edited_net("weights/synth_y.json", |o| { o["conv2_t.weight"]["mult"] = j!(9223372036854775808u64); })),
        ("three temperatures", "weights/synth_y.json", edited_net("weights/synth_y.json", |o| { o["TAM.TABs.0.attn.temperature"]["v"].as_array_mut().unwrap().pop(); })),
        ("ψ of other than 4C channels", "weights/hyp_uv.json", edited_net("weights/hyp_uv.json", |o| { o["conv4.weight"]["dims"] = j!([380, 2, 3, 3]); o["conv4.weight"]["w"] = j!(vec![1; 380 * 18]); o["conv4.bias"]["v"] = j!(vec![0; 380]); })),
        ("a luma synthesis of two planes", "weights/synth_y.json", edited_net("weights/synth_y.json", |o| { o["conv4_t.weight"]["dims"] = j!([4, 2, 3, 3]); o["conv4_t.weight"]["w"] = j!(vec![1; 72]); o["conv4_t.bias"]["v"] = j!([0, 0]); })),
        ("an odd feed-forward width", "weights/synth_uv.json", edited_net("weights/synth_uv.json", |o| {
            o["TAM.TABs.0.ffn.project_in.weight"]["dims"] = j!([3, 4, 1, 1]); o["TAM.TABs.0.ffn.project_in.weight"]["w"] = j!(vec![1; 12]);
            o["TAM.TABs.0.ffn.dwconv.weight"]["dims"] = j!([3, 1, 3, 3]); o["TAM.TABs.0.ffn.dwconv.weight"]["w"] = j!(vec![1; 27]);
        })),
        ("a member that is not an object of its kind", "weights/mcm1_y.json", edited_net("weights/mcm1_y.json", |o| { o["conv.1.bias"] = j!({ "kind": "q16", "v": "1" }); })),
        // Past acc64: one weight of 2^40 makes a sum that no i64 holds under A32.
        ("a sum past i64", "weights/synth_y.json", edited_net("weights/synth_y.json", |o| { o["CAB.subscale.weight"]["w"][0] = j!(1i64 << 40); })),
    ];
    let mut parsed_too = 0;
    for (what, file, (b, typed)) in &cases {
        let r = jpegai_ops::load_networks_bytes(nets_with(&m, file, b));
        assert_eq!(r.as_ref().err().map(|r| r.code), Some(wai_ops::Code::ModelMismatch), "{what}: the load from the bytes");
        let parsed = Weights::from_json(&json(b));
        assert_eq!(parsed.is_ok(), *typed, "{what}: this crate's parser reads exactly the typed files");
        if !*typed {
            parsed_too += 1;
        } else {
            let mut dw = synthetic_weights(&m);
            let w = parsed.unwrap();
            match *file {
                "weights/hyp_y.json" => dw.luma.hyper_decoder = w,
                "weights/hyp_uv.json" => dw.chroma.hyper_decoder = w,
                "weights/synth_y.json" => dw.luma.synthesis = w,
                "weights/synth_uv.json" => dw.chroma.synthesis = w,
                f => dw.mcm_y[f.as_bytes()["weights/mcm".len()] as usize - b'0' as usize] = w,
            }
            let r = jpegai_ops::load_networks(&dw);
            assert_eq!(r.as_ref().err().map(|r| r.code), Some(wai_ops::Code::ModelMismatch), "{what}: the load from the parsed networks");
            parsed_too += 1;
        }
    }
    // Two members of one name: JSON this crate's parser reads as one member; the bytes are
    // refused.
    let file = "weights/hyp_y.json";
    let text = String::from_utf8(m.get(file).to_vec()).unwrap();
    let dup = text.replacen('{', "{\"conv3.bias\":{\"kind\":\"q16\",\"v\":[0,0]},", 1).into_bytes();
    assert_eq!(jpegai_ops::load_networks_bytes(nets_with(&m, file, &dup)).err().map(|r| r.code), Some(wai_ops::Code::ModelMismatch), "two members of one name");
    // The committed bundle itself loads, from either form.
    assert!(jpegai_ops::load_networks_bytes(m.nets()).is_ok() && jpegai_ops::load_networks(&synthetic_weights(&m)).is_ok());
    assert_eq!(parsed_too, cases.len(), "every case is refused by the parser or the load from its result");
    eprintln!("{} load checks refused model_mismatch from the bytes and from the parsed networks, and two members of one name", cases.len());
}

/// The committed synthetic bundle's networks as this crate parses them (with empty scale
/// networks and eICCI, which the networks' load does not read).
fn synthetic_weights(m: &Members) -> DecoderWeights {
    let w = |n: &str| m.weights(n);
    DecoderWeights {
        luma: PlaneWeights { hyper_scale: IntNet::from_json(&json(m.get("weights/hs_y.json"))).unwrap(), hyper_decoder: w("weights/hyp_y.json"), synthesis: w("weights/synth_y.json") },
        chroma: PlaneWeights { hyper_scale: IntNet::from_json(&json(m.get("weights/hs_uv.json"))).unwrap(), hyper_decoder: w("weights/hyp_uv.json"), synthesis: w("weights/synth_uv.json") },
        mcm_y: [w("weights/mcm0_y.json"), w("weights/mcm1_y.json"), w("weights/mcm2_y.json"), w("weights/mcm3_y.json")],
        icci: wai::jpegai_icci::IcciNet { luma: empty_path(), chroma: empty_path(), conv_last_v: empty_conv(), scale_v: 0.0 },
    }
}

/// Gate 3, the ranges: a dequantised residual or a latent one past `A32`, at either end,
/// is refused `field_range` at its input, before any operator runs, by every graph that
/// reads it (on the committed networks); and, on the test codestream, a residual at the
/// range's end where the context stages' mean keeps the sum within it decodes, through
/// `rec_y` and `rec_uv`, to the direct decode's values, while one past it is refused
/// where the direct decode computed on (integer-operators §15.3 item 4).
#[test]
fn a_value_past_its_range_is_refused_field_range() {
    use wai_ops::{Code, Site};
    let m = Members::synthetic();
    let syn = jpegai_ops::load_networks_bytes(m.nets()).unwrap();
    let past = [i32::MAX as i64 + 1, i32::MIN as i64 - 1];
    let fill = |n: usize, at: usize, v: i64| -> Vec<i64> { (0..n).map(|i| if i == at { v } else { 0 }).collect() };
    let mut refused = 0;
    for v in past {
        for (role, inputs, k) in [
            ("rec_y", vec![Tensor::new(160, 1, 1, vec![0; 160]), Tensor::new(160, 4, 4, fill(160 * 16, 7, v))], 1u32),
            ("rec_uv", vec![Tensor::new(96, 1, 1, vec![0; 96]), Tensor::new(96, 4, 4, fill(96 * 16, 95, v))], 1),
            ("syn_y", vec![Tensor::new(160, 2, 2, fill(640, 3, v))], 0),
            ("syn_uv", vec![Tensor::new(160, 2, 2, vec![0; 640]), Tensor::new(96, 2, 2, fill(384, 0, v))], 1),
        ] {
            let r = wai_ops::run_with(syn.graph(role).unwrap(), inputs, &[], &mut wai_ops::Sequence::new(), &wai::int_ops::Host);
            assert_eq!(r.err(), Some(wai_ops::Refusal::new(Code::FieldRange, Site::Input(k))), "{role}: input {k} holding {v}");
            refused += 1;
        }
    }
    eprintln!("{refused} inputs one past A32 refused at their input");
    let Some(wh) = Whole::pinned("jpegai_nets_ops::a_value_past_its_range_is_refused_field_range") else {
        return;
    };
    let i = inputs(&wh);
    let nets = jpegai_ops::load_networks(&wh.weights).unwrap();
    for (k, plane) in [Plane::Luma, Plane::Chroma].into_iter().enumerate() {
        // The mean each position gets: ŷ − r. A luma position is taken from the last
        // sub-lattice, `(1, 0)`, which no context stage reads, so the other means stay.
        let mu: Vec<i64> = i.yhat[k].iter().zip(&i.resi[k]).map(|(y, r)| y - r).collect();
        let last = |p: usize| plane == Plane::Chroma || ((p % 256) / 16 % 2, p % 16 % 2) == (1, 0);
        let lo_at = (0..mu.len()).find(|&p| last(p) && mu[p] >= 0).unwrap();
        let hi_at = (0..mu.len()).find(|&p| last(p) && mu[p] <= 0).unwrap();
        for (at, edge, one_past) in [(lo_at, i32::MIN as i64, i32::MIN as i64 - 1), (hi_at, i32::MAX as i64, i32::MAX as i64 + 1)] {
            let mut r = i.resi[k].clone();
            r[at] = edge;
            let got = jpegai_ops::reconstruct(&nets, plane, i.z[k].clone(), r.clone()).unwrap_or_else(|e| panic!("{plane:?}: a residual of {edge} where the mean is {}: {e:?}", mu[at]));
            let want: Vec<i64> = r.iter().zip(&mu).map(|(a, b)| a + b).collect();
            assert!(got == want, "{plane:?}: a residual at {edge} decodes as the direct decode would");
            r[at] = one_past;
            assert_eq!(jpegai_ops::reconstruct(&nets, plane, i.z[k].clone(), r).err(), Some(wai_ops::Refusal::new(Code::FieldRange, Site::Input(1))), "{plane:?}: a residual of {one_past}");
        }
    }
    eprintln!("on the test codestream: a residual at either end of A32 decodes as directly, one past it is refused field_range");
}