// wai-conformance: int-kernels //! Each integer decoder through the integer operator set against its direct decode //! (`legacy_int`, test-only), on random parameter sets and payloads. The two must give //! the same output, or both refuse with the same code. The only inputs on which they may //! differ are the planned breaks the operator set makes (`extensions/integer-operators.md` //! §15.3), which are new refusals: each case is classified from the payload and the set //! alone, and a case that falls under a planned break must be refused by the decode //! through the operator set, whatever the direct decode does there. //! //! The decoders: `wai.neural.int_synth` (`int_synth/1`), `wai.splat.int_codec` //! (`whm1/1`), `wai.audio.int_codec` (`wam1/1`), `wai.neural.int_hyper` (`whm1/1`, its //! two networks; the keyframe of `wai.video.int_hyper` is the same decode) and //! `wai.neural.int_mlicpp` (`wmb1/1`, its three networks); and `wai.world.model`'s integer //! renderer (`world/1`, informative: `T` and `g_s`), against its direct renderer with the //! direct guard every other decoder ran; and `wai.image.jpegai`'s scale path and entropy //! decodes (`jpegai-int/1`), against its direct computation, and its hyper decoders, //! context stages and synthesis (the graphs of §13.6.2), against its direct networks; and //! `WIV1` and `WIR1` //! (`wai.video.int_motion` and the P-frames of `wai.video.int_hyper`, and the four //! `int_refine` capabilities), through the stage operators alone, which make no planned //! break (with `video_int` and `staged`). //! //! The same generators hold load step 8 to its oracle: the bounds pass, whose //! convolutions and inverse GDNs are in closed form, against the pass as it stood before //! (`wai_ops::bounds_oracle`, every term folded in 512 bits), on every graph they lower. //! //! The random cases run for a fixed count by default. `WAI_DIFF_SECS` runs each random //! test for that many seconds instead (the nightly run, `wai-ops-differential.yml`), and //! `WAI_DIFF_SEED` sets the first seed; a failing case names its seed. #![cfg(all(feature = "legacy_int", feature = "splat_int", feature = "audio_int", feature = "world_int"))] use std::panic::catch_unwind; use std::path::PathBuf; use std::time::{Duration, Instant}; use std::panic::AssertUnwindSafe; use serde_json::{json, Value}; use wai::int_hyper_synth::{IntModel, RefusalKind}; use wai::int_synth::{decode_from_spec, legacy, IntImage, IntSynthError}; fn wai() -> PathBuf { PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("..") } /// A splitmix64 sequence. struct Rng(u64); impl Rng { fn next(&mut self) -> u64 { self.0 = self.0.wrapping_add(0x9E37_79B9_7F4A_7C15); let mut z = self.0; z = (z ^ (z >> 30)).wrapping_mul(0xBF58_476D_1CE4_E5B9); z = (z ^ (z >> 27)).wrapping_mul(0x94D0_49BB_1331_11EB); z ^ (z >> 31) } fn below(&mut self, n: u64) -> u64 { self.next() % n } /// A value in `[−r, r]`. fn signed(&mut self, r: i64) -> i64 { (self.below(2 * r as u64 + 1) as i64) - r } } /// Run `case` on successive seeds: `n` of them, or for `WAI_DIFF_SECS` seconds. Each /// seed is `WAI_DIFF_SEED` (or a fixed default) mixed with the case's index. Returns how /// many ran. fn cases(name: &str, n: u64, mut case: impl FnMut(&mut Rng, u64)) -> u64 { let secs: Option = std::env::var("WAI_DIFF_SECS").ok().and_then(|s| s.parse().ok()); let seed0: u64 = std::env::var("WAI_DIFF_SEED").ok().and_then(|s| s.parse().ok()).unwrap_or(0x0D2C); let until = secs.map(|s| Instant::now() + Duration::from_secs(s)); let mut k = 0u64; loop { match until { Some(t) if Instant::now() >= t => break, None if k >= n => break, _ => {} } let seed = seed0 ^ k.wrapping_mul(0xD1B5_4A32_D192_ED03); case(&mut Rng(seed), seed); k += 1; } eprintln!("{name}: {k} cases from WAI_DIFF_SEED={seed0}"); k } // ------------------------------------------------------------------ int_synth (M1) const CAP: u128 = 1 << 26; /// The planned break of integer-operators §15.3 item 1 that a payload and set fall /// under, if any, read from the bytes alone. `"shift 64 to 127"` is the one that is not a /// refusal: such a shift now has its exact meaning, where the direct decode's result /// depended on the build. fn int_synth_planned_break(payload: &[u8], set: &[u8]) -> Option<&'static str> { if payload.len() < 12 { return None; } if payload[11] != 0 { return Some("a reserved byte that is not 0"); } let c = u16::from_le_bytes([payload[8], payload[9]]) as u128; if c == 0 { return Some("C of 0"); } let s_log2 = payload[10] as u32; if s_log2 >= 64 { return None; } if (1u128 << (2 * s_log2)).checked_mul(c).is_none_or(|n| n > CAP) { return Some("latents past 2^26"); } let Ok(root) = serde_json::from_slice::(set) else { return None }; let Some(layers) = root.get("layers").and_then(Value::as_array) else { return None }; if layers.is_empty() { return Some("no layers"); } let field = |l: &Value, k: &str, default: Option| l.get(k).map_or(default, Value::as_u64); let mut shift_64 = false; let mut side = 1u128 << s_log2; for l in layers { let (Some(in_ch), Some(out_ch), Some(k), Some(stride), Some(pad), Some(out_pad), Some(shift)) = ( field(l, "in_ch", None), field(l, "out_ch", None), field(l, "k", None), field(l, "stride", None), field(l, "pad", None), field(l, "out_pad", Some(0)), field(l, "shift", Some(0)), ) else { return None; }; if [in_ch, out_ch, k, stride, pad, out_pad].iter().any(|&v| v > u32::MAX as u64) { return Some("a field past 32 bits"); } if shift >= 128 { return Some("a shift of 128 or more"); } if k == 0 || stride == 0 { return Some("k or stride of 0"); } if out_ch == 0 { return Some("out_ch of 0"); } shift_64 |= shift >= 64; // The run-time plan's caps (integer-operators §3, §5.3): the uncropped sides and // output, then the output. let full = (side - 1) * stride as u128 + k as u128 + out_pad as u128; if full > CAP || out_ch as u128 * full * full > CAP { return Some("an uncropped output past 2^26"); } if full <= 2 * pad as u128 { return None; } side = full - 2 * pad as u128; if out_ch as u128 * side * side > CAP { return Some("a tensor past 2^26"); } } let lowered = wai_ops::lower::int_synth(set).ok()?; if wai::int_ops::load(lowered.get("synth")?).is_err() { return Some("a certificate that does not meet acc64"); } shift_64.then_some("shift 64 to 127") } /// A random `int_synth.json` of 1 to 3 layers and a payload for it: kernels 1 to 4, /// strides 1 to 3, paddings mostly below the kernel and output paddings mostly below the /// stride, shifts 0 to 24, weights of up to 2^4, 2^10 or 2^14 (one set in eight up to /// 2^20, where certificates start to fail), a quarter of them 0, and biases of up to /// 2^18; latent channels 1 to 3, sides 1 to 4, a quarter of the latents at the ends of /// their range. The payload's `H` and `W` are the output's sides, or (one case in eight) /// not; one payload in thirty-two has a reserved byte that is not 0. fn int_synth_case(r: &mut Rng) -> (Vec, Vec) { let c0 = 1 + r.below(3) as usize; let s_log2 = r.below(3) as u8; let (mut c, mut side) = (c0, 1i64 << s_log2); let n_layers = 1 + r.below(3) as usize; let big = r.below(8) == 0; let mut layers = Vec::new(); for i in 0..n_layers { let out = if i + 1 == n_layers && r.below(16) != 0 { 3 } else { 1 + r.below(4) as usize }; let k = 1 + r.below(4) as usize; let stride = 1 + r.below(3) as usize; let pad = if r.below(16) == 0 { r.below(k as u64 + 2) } else { r.below(k as u64) } as usize; let out_pad = if r.below(16) == 0 { r.below(stride as u64 + 2) } else { r.below(stride as u64) } as usize; let shift = r.below(25); let wmax = if big { 1 << 20 } else { [1i64 << 4, 1 << 10, 1 << 14][r.below(3) as usize] }; let weight: Vec = (0..c * out * k * k).map(|_| if r.below(4) == 0 { 0 } else { r.signed(wmax) }).collect(); let bias: Vec = (0..out).map(|_| r.signed(1 << 18)).collect(); layers.push(json!({ "in_ch": c, "out_ch": out, "k": k, "stride": stride, "pad": pad, "out_pad": out_pad, "weight": weight, "bias": bias, "shift": shift, "relu": r.below(2) == 1, })); side = ((side - 1) * stride as i64 + k as i64 + out_pad as i64 - 2 * pad as i64).max(1); c = out; } let (mut h, mut w) = (side as u32, side as u32); if r.below(8) == 0 { h += 1 + r.below(2) as u32; w = w.saturating_sub(r.below(2) as u32); } let mut p = Vec::new(); p.extend(h.to_le_bytes()); p.extend(w.to_le_bytes()); p.extend((c0 as u16).to_le_bytes()); p.extend([s_log2, if r.below(32) == 0 { 1 + r.below(255) as u8 } else { 0 }]); let n = c0 << (2 * s_log2); p.extend((0..n).map(|_| match r.below(8) { 0 => 0x80, 1 => 0x7F, _ => r.signed(127) as i8 as u8, })); (serde_json::to_vec(&json!({ "layers": layers })).unwrap(), p) } /// The direct decode, a panic (an overflow-checked build's arithmetic overflow) caught. fn direct(payload: &[u8], set: &[u8]) -> Option> { catch_unwind(|| legacy::decode_from_spec(payload, set)).ok() } /// Gate 4 of the migration: the decode through the operator set against the direct /// decode on random sets and payloads. Outside the planned breaks both decode to the /// same bytes or both refuse (one code, `payload_malformed`), and the direct decode never /// panics; under a planned break the decode through the operator set refuses. #[test] fn int_synth_random_sets_and_payloads_decode_alike() { let (mut alike, mut refused, mut breaks) = (0u64, 0u64, std::collections::BTreeMap::<&str, u64>::new()); let n = cases("int_synth", 3000, |r, seed| { let (set, pay) = int_synth_case(r); let new = decode_from_spec(&pay, &set); let why = || format!("seed {seed}\nset {}\npayload {:02x?}", String::from_utf8_lossy(&set), pay); match int_synth_planned_break(&pay, &set) { Some("shift 64 to 127") => *breaks.entry("shift 64 to 127").or_default() += 1, Some(b) => { assert!(new.is_err(), "{b}: the decode through the operator set must refuse it\n{}", why()); *breaks.entry(b).or_default() += 1; } None => match (&new, direct(&pay, &set)) { (Ok(a), Some(Ok(b))) => { assert!(*a == b, "the two decodes differ\n{}", why()); alike += 1; } (Err(_), Some(Err(_))) => refused += 1, (_, old) => panic!("new {new:?}, direct {old:?}\n{}", why()), }, } }); eprintln!("int_synth: {alike} decoded alike, {refused} refused by both, planned breaks {breaks:?}"); assert!(alike * 2 > n, "the generator must mostly produce decodable cases: {alike} of {n}"); assert!(n < 1000 || breaks.get("a certificate that does not meet acc64").is_some_and(|&b| b > 0), "no case reached a failing certificate"); } /// One layer of `int_synth.json`: every weight 1, the biases 10, 20 and 30 (or, for /// fewer channels, their first `out_ch`). fn layer(in_ch: u64, out_ch: u64, k: u64, stride: u64, pad: u64, out_pad: u64, shift: u64) -> Value { let bias: Vec = [10, 20, 30].into_iter().chain(std::iter::repeat(0)).take(out_ch as usize).collect(); json!({ "in_ch": in_ch, "out_ch": out_ch, "k": k, "stride": stride, "pad": pad, "out_pad": out_pad, "weight": vec![1; (in_ch * out_ch * k * k) as usize], "bias": bias, "shift": shift, }) } fn set(layers: Vec) -> Vec { serde_json::to_vec(&json!({ "layers": layers })).unwrap() } /// A payload for an `h`×`h` output: `C` channels of side 1, reserved byte `reserved`, the /// latents `lat`. fn payload(h: u32, c: u16, reserved: u8, lat: &[i8]) -> Vec { let mut p = Vec::new(); p.extend(h.to_le_bytes()); p.extend(h.to_le_bytes()); p.extend(c.to_le_bytes()); p.extend([0u8, reserved]); p.extend(lat.iter().map(|&v| v as u8)); p } /// The planned breaks of integer-operators §15.3 item 1, each on a payload and set the /// direct decode decodes (or, for a shift of 64 or more, decodes by build): the /// classifier names it, and the decode through the operator set refuses it, or, for a /// shift from 64 to 127, gives the shift its exact meaning. #[test] fn int_synth_planned_breaks_are_new_refusals() { let one = set(vec![layer(1, 3, 1, 1, 0, 0, 0)]); let five = payload(1, 1, 0, &[5]); let expect = |b: &str, pay: &[u8], s: &[u8], direct_decodes: bool| { assert_eq!(int_synth_planned_break(pay, s), Some(b), "{b}: the classifier"); if direct_decodes { assert!(matches!(direct(pay, s), Some(Ok(_))), "{b}: the direct decode decodes it"); } assert!(decode_from_spec(pay, s).is_err(), "{b}: the decode through the operator set refuses it"); }; // The base decodes alike: 5·1 + the bias. assert_eq!(int_synth_planned_break(&five, &one), None); assert_eq!(decode_from_spec(&five, &one).unwrap().rgb, [15, 25, 35]); assert_eq!(direct(&five, &one), Some(decode_from_spec(&five, &one))); expect("a reserved byte that is not 0", &payload(1, 1, 7, &[5]), &one, true); // C of 0: no latents; the set's one layer reads no channel, so the direct decode's // output is the bias. expect("C of 0", &payload(1, 0, 0, &[]), &set(vec![layer(0, 3, 1, 1, 0, 0, 0)]), true); // No layers: the direct decode presents the latents themselves, clamped, as the image. expect("no layers", &payload(1, 3, 0, &[5, -3, 100]), &set(vec![]), true); // A layer of no channels: the next reads none, so the output is its bias. expect("out_ch of 0", &five, &set(vec![layer(1, 0, 1, 1, 0, 0, 0), layer(0, 3, 1, 1, 0, 0, 0)]), true); expect("k or stride of 0", &five, &set(vec![layer(1, 3, 0, 1, 0, 1, 0)]), true); expect("k or stride of 0", &five, &set(vec![layer(1, 3, 1, 0, 0, 0, 0)]), true); expect("a field past 32 bits", &five, &set(vec![layer(1, 3, 1, 1 << 32, 0, 0, 0)]), true); // An uncropped output of 3·8193·8193 cropped to 1×1 that no tap reaches: the direct // decode forms no uncropped plane, and its output is the bias. expect("an uncropped output past 2^26", &five, &set(vec![layer(1, 3, 3, 1, 1 << 12, (1 << 13) - 2, 0)]), true); // The certificate: the falsification set that passes i64 by one unit decodes a // small latent directly. let overflow = std::fs::read(wai().join("ops-conformance/sets/overflow_acc64/int_synth.json")).unwrap(); expect("a certificate that does not meet acc64", &payload(1, 1, 0, &[1]), &overflow, true); // A shift of 128 or more; the direct decode panics or masks, by build. expect("a shift of 128 or more", &five, &set(vec![layer(1, 3, 1, 1, 0, 0, 128)]), false); // A shift from 64 to 127 has its exact meaning: every sum here is below 2^64. let s64 = set(vec![layer(1, 3, 1, 1, 0, 0, 64)]); assert_eq!(int_synth_planned_break(&five, &s64), Some("shift 64 to 127")); assert_eq!(decode_from_spec(&five, &s64).unwrap().rgb, [0, 0, 0]); } // ------------------------------------------------------------------ int_splat (M2) /// One layer of a `WHM1` network, as a set's bytes hold it (integer-payloads §9). #[derive(Clone)] enum Whm { ConvT { inc: u32, out: u32, k: u32, s: u32, p: u32, op: u32, w: Vec, m: Vec, b: Vec }, Conv { inc: u32, out: u32, k: u32, s: u32, p: u32, w: Vec, m: Vec, b: Vec }, Igdn { c: u32, g: Vec, b: Vec }, Relu, Leaky { q: i64, f: u32 }, } /// A `WHM1` network's bytes. fn whm1_bytes(afrac: u32, gf: u32, shift: u32, layers: &[Whm]) -> Vec { let mut v = b"WHM1".to_vec(); for x in [afrac, gf, shift, layers.len() as u32] { v.extend(x.to_le_bytes()); } let words = |v: &mut Vec, xs: &[u32]| xs.iter().for_each(|x| v.extend(x.to_le_bytes())); for l in layers { match l { Whm::ConvT { inc, out, k, s, p, op, w, m, b } => { v.push(0); words(&mut v, &[*inc, *out, *k, *s, *p, *op]); w.iter().for_each(|x| v.extend(x.to_le_bytes())); m.iter().chain(b).for_each(|x| v.extend(x.to_le_bytes())); } Whm::Conv { inc, out, k, s, p, w, m, b } => { v.push(1); words(&mut v, &[*inc, *out, *k, *s, *p]); w.iter().for_each(|x| v.extend(x.to_le_bytes())); m.iter().chain(b).for_each(|x| v.extend(x.to_le_bytes())); } Whm::Igdn { c, g, b } => { v.push(2); words(&mut v, &[*c]); g.iter().for_each(|x| v.extend(x.to_le_bytes())); b.iter().for_each(|x| v.extend(x.to_le_bytes())); } Whm::Relu => v.push(3), Whm::Leaky { q, f } => { v.push(4); v.extend(q.to_le_bytes()); v.extend(f.to_le_bytes()); } } } v } /// A `WIS1` cloud's typed fields, owned. #[derive(Clone, Debug)] struct Cloud { afrac: u32, c: usize, sh: usize, sw: usize, s: usize, n_attr: usize, cdfs: Vec>, offsets: Vec, median_q: Vec, rans: Vec, } impl Cloud { fn st(&self) -> wai::int_splat::SplatStream<'_> { wai::int_splat::SplatStream { afrac: self.afrac, c: self.c, sh: self.sh, sw: self.sw, s: self.s, n_attr: self.n_attr, cdfs: &self.cdfs, offsets: &self.offsets, median_q: &self.median_q, rans: &self.rans, } } } /// The planned break of integer-operators §15.3 that a splat set falls under, if any, /// read from the set's bytes alone: the set is refused at load by its lowering /// (`whm1/1` with the RGB8 step, loaded under `acc64`), for a rule the direct decode /// does not make. A set refused at load for a rule integer-payloads §7's plan also makes /// (the chain, a zero dimension, a shift a layer uses) is not a break: the plan, which /// both decodes run first, refuses it alike. No payload field changes meaning. fn int_splat_planned_break(set: &[u8]) -> Option<&'static str> { let net = wai_ops::lower::parse_whm1(set).ok()?; load_break(&net, wai_ops::lower::whm1(set, true)) } /// The planned break a parsed `WHM1` or `WAM1` network falls under, given its lowering /// (see [`int_splat_planned_break`]). fn load_break(net: &wai_ops::lower::Net, lowered: Result) -> Option<&'static str> { use wai_ops::lower::NetLayer as N; let loaded = lowered.as_ref().map_err(|r| *r).and_then(wai::int_ops::load); if loaded.is_ok() { return None; } let conv = net.layers.iter().any(|l| matches!(l, N::ConvT { .. } | N::Conv { .. })); let igdn = net.layers.iter().any(|l| matches!(l, N::Igdn { .. })); if !conv && !igdn { return Some("no layer of tag 0, 1 or 2"); } if (net.shift > 127 && !conv) || (net.gf > 127 && !igdn) { return Some("a shift or gf of 128 or more that no layer uses"); } if (1..=30).contains(&net.afrac) && net.shift <= 127 && net.gf <= 127 { let huge = |b: &[i64]| b.iter().any(|v| v.unsigned_abs() >= i64::MAX as u64); if net.layers.iter().any(|l| matches!(l, N::ConvT { bias, .. } | N::Conv { bias, .. } if huge(bias))) { return Some("a bias of magnitude 2^63 - 1 or more"); } } match loaded { Err(r) if lowered.is_ok() && r.at == wai_ops::Site::Load(8) => Some("a check range its producer cannot meet"), _ => None, } } /// A random layer stack for `int_splat_case`, from `c` input channels: transposed /// convolutions, convolutions, inverse GDNs, ReLUs and leaky ReLUs; kernels 1 to 4, /// strides 1 to 3, paddings mostly below the kernel. Weights of up to 2^4, 2^10 or 2^12 /// (one set in eight up to i16's range), a quarter 0; multipliers mostly up to 2^12 (one /// in 32 up to 2^40), biases mostly up to 2^20, with one in 64 at an end of i64; IGDN `γ` mostly /// non-negative and up to 2^12, `β` up to 2^28. Returns the layers and the output /// channels. fn splat_layers(r: &mut Rng, mut c: u32, n: usize, only_relu: bool) -> (Vec, u32) { let big = r.below(8) == 0; let mut layers = Vec::new(); for _ in 0..n { let kind = if only_relu { 3 } else { [0, 0, 0, 1, 2, 2, 3, 4][r.below(8) as usize] }; let l = match kind { 0 | 1 => { let out = 1 + r.below(4) as u32; let k = 1 + r.below(4) as u32; let s = 1 + r.below(3) as u32; let p = if r.below(16) == 0 { r.below(k as u64 + 3) } else { r.below(k as u64) } as u32; let wmax = if big { 32767 } else { [16, 1024, 4096][r.below(3) as usize] }; let w: Vec = (0..c * out * k * k).map(|_| if r.below(4) == 0 { 0 } else { r.signed(wmax) as i16 }).collect(); let m: Vec = (0..out).map(|_| if r.below(32) == 0 { r.signed(1 << 40) } else { r.signed(1 << 12) }).collect(); let b: Vec = (0..out) .map(|_| match r.below(64) { 0 => [i64::MAX, i64::MIN, i64::MIN + 1, i64::MAX - 1][r.below(4) as usize], 1 => r.signed(1 << 40), _ => r.signed(1 << 20), }) .collect(); let inc = c; c = out; if kind == 0 { let op = if r.below(16) == 0 { r.below(s as u64 + 1) } else { r.below(s as u64) } as u32; Whm::ConvT { inc, out, k, s, p, op, w, m, b } } else { Whm::Conv { inc, out, k, s, p, w, m, b } } } 2 => { let g: Vec = (0..c * c).map(|_| if r.below(16) == 0 { r.signed(1 << 12) as i32 } else { r.below(1 << 12) as i32 }).collect(); let b: Vec = (0..c).map(|_| if r.below(16) == 0 { r.signed(1 << 28) } else { r.below(1 << 28) as i64 }).collect(); Whm::Igdn { c, g, b } } 3 => Whm::Relu, _ => Whm::Leaky { q: if r.below(32) == 0 { r.signed(1 << 40) } else { r.signed(1 << 10) }, f: if r.below(32) == 0 { 64 } else { r.below(16) as u32 } }, }; layers.push(l); } (layers, c) } /// The output side of `layers` on a square input of side `n`, or `None` where the /// geometry leaves no output. fn splat_side(layers: &[Whm], mut n: i64) -> Option { for l in layers { n = match l { Whm::ConvT { k, s, p, op, .. } => (n - 1) * *s as i64 + *k as i64 + *op as i64 - 2 * *p as i64, Whm::Conv { k, s, p, .. } => { let np = n + 2 * *p as i64 - *k as i64; if np < 0 { return None } else { np / *s as i64 + 1 } } _ => n, }; if n < 1 { return None; } } Some(n) } /// A random `wai.splat.int_codec` set and cloud: a network of 1 to 4 layers (one set in /// 32 with no layers, one in 32 of ReLUs only) at `afrac` 12 mostly (one set in sixteen /// anywhere in 1..=30, one in sixteen at 0 or 31), a shift of 8 to 23, with a header `gf` /// or `shift` of 128 or more one set in 32 each; a latent of 1 to 3 channels of side 1 to 4 (one in sixteen /// not square), each channel with its own table of 3 to 6 entries, offset and median; the /// symbols mostly inside their table, one in eight escaped past it and one in 32 far /// past it; medians mostly small, one in eight at or past the 2^47 bound, and one in four /// placing a whole channel's latents at the input's `check` range or one past it. The /// declared attribute shape is the network's output (one case in eight off by one); the /// cloud's `afrac` is the set's (one in 32 not); one stream in sixteen is cut short, one /// in 32 has a word appended, one in 64 a byte changed. fn int_splat_case(r: &mut Rng) -> (Vec, Cloud) { let afrac = match r.below(16) { 0 => 1 + r.below(30) as u32, 1 => [0, 31][r.below(2) as usize], _ => 12, }; let gf = if r.below(32) == 0 { 128 + r.below(64) as u32 } else { 16 }; let shift = if r.below(32) == 0 { 128 + r.below(64) as u32 } else { 8 + r.below(16) as u32 }; let c0 = 1 + r.below(3) as u32; let n_layers = if r.below(32) == 0 { 0 } else { 1 + r.below(4) as usize }; let only_relu = r.below(32) == 0; let (layers, c_out) = splat_layers(r, c0, n_layers, only_relu); let set = whm1_bytes(afrac, gf, shift, &layers); let sh = 1 + r.below(4) as usize; let sw = if r.below(16) == 0 { 1 + r.below(4) as usize } else { sh }; let mut s = splat_side(&layers, sh as i64).unwrap_or(1) as usize; let mut n_attr = c_out as usize; if r.below(8) == 0 { if r.below(2) == 0 { s += 1 } else { n_attr += 1 } } let c = c0 as usize; let a = afrac.min(30); let mut cdfs = Vec::new(); let mut offsets = Vec::new(); let mut median_q = Vec::new(); for _ in 0..c { let len = 3 + r.below(4) as usize; let mut cuts: Vec = (0..len - 2).map(|_| 1 + r.below(65_535) as i32).collect(); cuts.sort_unstable(); cuts.dedup(); let mut t = vec![0]; t.extend(cuts); t.push(1 << 16); cdfs.push(t); offsets.push(r.signed(3) as i32); median_q.push(match r.below(8) { 0 => [1i64 << 47, -(1 << 47), (1 << 47) + 1, -(1 << 47) - 1][r.below(4) as usize], _ => r.signed(1 << (a + 4)), }); } let plane = sh * sw; let mut values: Vec = Vec::with_capacity(c * plane); for ch in 0..c { let span = cdfs[ch].len() as i64 - 2; for _ in 0..plane { let v = match r.below(32) { 0 => r.signed(1 << 20), 1..=4 => offsets[ch] as i64 + span + r.below(40) as i64 - 20, _ => offsets[ch] as i64 + r.below(span as u64) as i64, }; values.push(v as i32); } } // One case in four: a whole channel's latents at the input's check range, or one // past it, where the set lowers. if r.below(4) == 0 { if let Ok(g) = wai_ops::lower::whm1(&set, true) { let x0 = g.tensors[0].hi as i128; let ch = r.below(c as u64) as usize; let v = offsets[ch]; let target = [x0, x0 + 1, -x0, -x0 - 1][r.below(4) as usize]; let med = target - ((v as i128) << a); if med.unsigned_abs() <= 1 << 47 { median_q[ch] = med as i64; values[ch * plane..(ch + 1) * plane].iter_mut().for_each(|x| *x = v); } } } let mut enc = wai::rans::RansEncoder::new(); for ch in 0..c { for &v in &values[ch * plane..(ch + 1) * plane] { assert!(enc.encode(v, &cdfs[ch], cdfs[ch].len(), offsets[ch]), "a value the generator writes is codable"); } } let mut rans = enc.flush(); match r.below(64) { 0..=3 => { let cut = (4 * (1 + r.below(3)) as usize).min(rans.len()); rans.truncate(rans.len() - cut); } 4 | 5 => rans.extend([0, 0, 0, 0]), 6 => { let i = r.below(rans.len() as u64) as usize; rans[i] ^= 1 << r.below(8); } _ => {} } let cloud_afrac = if r.below(32) == 0 { afrac ^ 1 } else { afrac }; (set, Cloud { afrac: cloud_afrac, c, sh, sw, s, n_attr, cdfs, offsets, median_q, rans }) } /// The direct splat decode, a panic (an overflow-checked build's arithmetic overflow) /// caught. fn splat_direct(gs: &IntModel, cl: &Cloud) -> Option, RefusalKind>> { catch_unwind(AssertUnwindSafe(|| wai::int_splat::legacy::decode_typed(gs, &cl.st()).map(|d| d.attr).map_err(|r| r.kind))).ok() } /// Gate 4 of the migration, for `wai.splat.int_codec`: the decode through the operator /// set against the direct decode on random sets and clouds. Outside the planned breaks /// both decode to the same attribute bytes or both refuse with the same code, and the /// direct decode never panics; under a planned break the decode through the operator set /// refuses. #[test] fn int_splat_random_sets_and_clouds_decode_alike() { let (mut alike, mut refused, mut breaks) = (0u64, std::collections::BTreeMap::::new(), std::collections::BTreeMap::<&str, u64>::new()); let n = cases("int_splat", 3000, |r, seed| { let (set, cl) = int_splat_case(r); let why = || format!("seed {seed}\ncloud {cl:?}"); let gs = IntModel::from_bytes(&set).expect("the generator writes WHM1 that parses"); let new = wai::int_splat::ops::decode_typed(&gs, &cl.st()).map(|d| d.attr).map_err(|r| r.kind); match int_splat_planned_break(&set) { Some(b) => { assert!(new.is_err(), "{b}: the decode through the operator set must refuse it\n{}", why()); *breaks.entry(b).or_default() += 1; } None => match (&new, splat_direct(&gs, &cl)) { (Ok(a), Some(Ok(b))) => { assert!(*a == b, "the two decodes differ\n{}", why()); alike += 1; } (Err(a), Some(Err(b))) => { assert_eq!(*a, b, "the two decodes refuse with different codes\n{}", why()); *refused.entry(format!("{a:?}")).or_default() += 1; } (_, old) => panic!("new {new:?}, direct {old:?}\n{}", why()), }, } }); eprintln!("int_splat: {alike} decoded alike, refused by both {refused:?}, planned breaks {breaks:?}"); assert!(alike * 5 > n, "the generator must produce decodable cases: {alike} of {n}"); assert!(n < 1000 || refused.contains_key("FieldRange"), "no case reached a check range"); assert!(n < 1000 || breaks.len() >= 3, "too few kinds of planned break reached: {breaks:?}"); } /// The planned breaks of integer-operators §15.3 item 3 for `wai.splat.int_codec`, each /// on a set and cloud the direct decode decodes, or refuses with another code: the /// classifier names it, and the decode through the operator set refuses it /// (`model_mismatch`, before the stream). #[test] fn int_splat_planned_breaks_are_new_refusals() { // A cloud of 2 channels of side 2, every latent 2·2^12, each symbol the first of the // table `t` (offset 2). let cloud_with = |t: Vec| { let (c, n, s, n_attr) = (2usize, 2usize, 2usize, 2usize); let mut enc = wai::rans::RansEncoder::new(); for _ in 0..c * n * n { assert!(enc.encode(2, &t, 3, 2)); } Cloud { afrac: 12, c, sh: n, sw: n, s, n_attr, cdfs: vec![t; c], offsets: vec![2; c], median_q: vec![0; c], rans: enc.flush() } }; let conv1 = |inc: u32, out: u32, w: i16, b: i64| Whm::ConvT { inc, out, k: 1, s: 1, p: 0, op: 0, w: vec![w; (inc * out) as usize], m: vec![1; out as usize], b: vec![b; out as usize] }; let expect = |b: &str, set: &[u8], cl: &Cloud, direct: Option| { assert_eq!(int_splat_planned_break(set), Some(b), "{b}: the classifier"); let gs = IntModel::from_bytes(set).unwrap(); match direct { None => assert!(matches!(splat_direct(&gs, cl), Some(Ok(_))), "{b}: the direct decode decodes it"), Some(k) => assert_eq!(splat_direct(&gs, cl), Some(Err(k)), "{b}: the direct decode's code"), } let new = wai::int_splat::ops::decode_typed(&gs, &cl.st()).map(|_| ()).map_err(|r| r.kind); assert_eq!(new, Err(RefusalKind::ModelMismatch), "{b}: the decode through the operator set refuses it"); }; // The base decodes alike: one 1×1 transposed convolution, weight 1, shift 0. let base = whm1_bytes(12, 16, 0, &[conv1(2, 2, 1, 0)]); assert_eq!(int_splat_planned_break(&base), None); let gs = IntModel::from_bytes(&base).unwrap(); let cl = cloud_with(vec![0, 1 << 15, 1 << 16]); assert_eq!(splat_direct(&gs, &cl), Some(Ok(vec![255; 8])), "2·2^12 clamped to 2^12: 255"); assert_eq!(wai::int_splat::ops::decode_typed(&gs, &cl.st()).unwrap().attr, vec![255; 8]); // No layer of tag 0, 1 or 2: the direct decode presents the latents themselves. expect("no layer of tag 0, 1 or 2", &whm1_bytes(12, 16, 0, &[Whm::Relu]), &cl, None); expect("no layer of tag 0, 1 or 2", &whm1_bytes(12, 16, 0, &[]), &cl, None); // A header shift no layer uses (only an inverse GDN), or a gf no layer uses. let igdn = Whm::Igdn { c: 2, g: vec![0; 4], b: vec![1 << 24; 2] }; expect("a shift or gf of 128 or more that no layer uses", &whm1_bytes(12, 16, 200, &[igdn]), &cl, None); expect("a shift or gf of 128 or more that no layer uses", &whm1_bytes(12, 200, 0, &[conv1(2, 2, 1, 0)]), &cl, None); // A bias at the end of i64: the direct decode refuses every input at that layer // (field_range), after the stream; the lowering refuses the set before it, so a // stream that runs out is refused for the set, not the stream. expect("a bias of magnitude 2^63 - 1 or more", &whm1_bytes(12, 16, 0, &[conv1(2, 2, 1, i64::MAX)]), &cl, Some(RefusalKind::FieldRange)); let mut short = cloud_with(vec![0, 1, 1 << 16]); short.rans.truncate(8); expect("a bias of magnitude 2^63 - 1 or more", &whm1_bytes(12, 16, 0, &[conv1(2, 2, 1, i64::MIN)]), &short, Some(RefusalKind::StreamExhausted)); // A check range its producer cannot meet: layer 0 gives every value 2^30 (weights 0, // bias 2^30); layer 1, two input channels of weight 2^11 under the multiplier 2^40 and // shift 0, has X_1 = 2047. The direct decode refuses every input at layer 1 // (field_range); the load refuses the set (an empty check intersection, §8.1). let tight = Whm::ConvT { inc: 2, out: 2, k: 1, s: 1, p: 0, op: 0, w: vec![1 << 11; 4], m: vec![1 << 40; 2], b: vec![0; 2] }; let set = whm1_bytes(12, 16, 0, &[conv1(2, 2, 0, 1 << 30), tight]); expect("a check range its producer cannot meet", &set, &cl, Some(RefusalKind::FieldRange)); } // ------------------------------------------------------------------ int_audio (M3) /// One layer of a `WAM1` network, as a set's bytes hold it (integer-payloads §9). #[derive(Clone)] enum Wam { ConvT { inc: u32, out: u32, k: u32, s: u32, p: u32, op: u32, w: Vec, m: Vec, b: Vec }, Igdn { c: u32, g: Vec, b: Vec }, } /// A `WAM1` network's bytes. fn wam1_bytes(afrac: u32, gf: u32, shift: u32, layers: &[Wam]) -> Vec { let mut v = b"WAM1".to_vec(); for x in [afrac, gf, shift, layers.len() as u32] { v.extend(x.to_le_bytes()); } for l in layers { match l { Wam::ConvT { inc, out, k, s, p, op, w, m, b } => { v.push(0); [*inc, *out, *k, *s, *p, *op].iter().for_each(|x| v.extend(x.to_le_bytes())); w.iter().for_each(|x| v.extend(x.to_le_bytes())); m.iter().chain(b).for_each(|x| v.extend(x.to_le_bytes())); } Wam::Igdn { c, g, b } => { v.push(2); v.extend(c.to_le_bytes()); g.iter().for_each(|x| v.extend(x.to_le_bytes())); b.iter().for_each(|x| v.extend(x.to_le_bytes())); } } } v } /// A `WIA1` clip's typed fields, owned. #[derive(Clone, Debug)] struct Clip { afrac: u32, n_ch: usize, lat_len: usize, t_out: usize, offset: i32, median_q: Vec, cdfs: Vec>, rans: Vec, } impl Clip { fn st(&self) -> wai::int_audio::AudioStream<'_> { wai::int_audio::AudioStream { afrac: self.afrac, n_ch: self.n_ch, lat_len: self.lat_len, t_out: self.t_out, offset: self.offset, median_q: &self.median_q, cdfs: &self.cdfs, rans: &self.rans, } } } /// The planned break of integer-operators §15.3 that an audio set falls under, if any, /// read from the set's bytes alone: the set is refused at load by its lowering (`wam1/1` /// with the PCM step, loaded under `acc64`), for a rule the direct decode does not make /// (as for splat, [`int_splat_planned_break`]; a `WAM1` network has layers of tags 0 and /// 2 only, so one with no such layer has no layers). No payload field changes meaning. fn int_audio_planned_break(set: &[u8]) -> Option<&'static str> { let net = wai_ops::lower::parse_wam1(set).ok()?; load_break(&net, wai_ops::lower::wam1(set)) } /// The length a stack of layers gives on `n` latent samples, or `None` where one leaves /// no output. fn audio_len(layers: &[Wam], mut n: i64) -> Option { for l in layers { if let Wam::ConvT { k, s, p, op, .. } = l { n = (n - 1) * *s as i64 + *k as i64 + *op as i64 - 2 * *p as i64; if n < 1 { return None; } } } Some(n) } /// A random `wai.audio.int_codec` set and clip: a network of 1 to 4 layers (one set in /// 32 with none), one-dimensional transposed convolutions and inverse GDNs, the last a /// transposed convolution to one channel (one set in sixteen not), at `afrac` 12 mostly /// (one in sixteen anywhere in 1..=30, one in sixteen at 0 or 31), a shift of 8 to 23, /// and a header `gf` or `shift` of 128 or more one set in 32 each; kernels 1 to 5, /// strides 1 to 3, paddings mostly below the kernel; weights up to 2^4, 2^10 or 2^12 (one /// set in eight up to i16's range), multipliers mostly up to 2^12, biases mostly up to /// 2^20 with one in 64 at an end of i64; IGDN `γ` up to 2^12, `β` up to 2^28. A latent /// of 1 to 3 channels of 1 to 6 samples, one table length and offset for every channel; /// symbols mostly inside the table, some escaped and far past it; medians mostly small, /// one in eight at or past 2^47, and one case in four with a whole channel at the input's /// `check` range or one past it. `t_out` is one the latent is the shortest for (one case /// in eight not); the clip's `afrac` is the set's (one in 32 not); one stream in sixteen /// is cut short, one in 32 has a word appended, one in 64 a byte changed. fn int_audio_case(r: &mut Rng) -> (Vec, Clip) { let afrac = match r.below(16) { 0 => 1 + r.below(30) as u32, 1 => [0, 31][r.below(2) as usize], _ => 12, }; let gf = if r.below(32) == 0 { 128 + r.below(64) as u32 } else { 16 }; let shift = if r.below(32) == 0 { 128 + r.below(64) as u32 } else { 8 + r.below(16) as u32 }; let c0 = 1 + r.below(3) as u32; let n_layers = if r.below(32) == 0 { 0 } else { 1 + r.below(4) as usize }; let big = r.below(8) == 0; let (mut c, mut layers) = (c0, Vec::new()); for i in 0..n_layers { let last = i + 1 == n_layers; if !last && r.below(3) == 0 { let g: Vec = (0..c * c).map(|_| if r.below(16) == 0 { r.signed(1 << 12) as i32 } else { r.below(1 << 12) as i32 }).collect(); let b: Vec = (0..c).map(|_| if r.below(16) == 0 { r.signed(1 << 28) } else { r.below(1 << 28) as i64 }).collect(); layers.push(Wam::Igdn { c, g, b }); continue; } let out = if last && r.below(16) != 0 { 1 } else { 1 + r.below(3) as u32 }; let k = 1 + r.below(5) as u32; let s = 1 + r.below(3) as u32; let p = if r.below(16) == 0 { r.below(k as u64 + 3) } else { r.below(k as u64) } as u32; let op = if r.below(16) == 0 { r.below(s as u64 + 1) } else { r.below(s as u64) } as u32; let wmax = if big { 32767 } else { [16, 1024, 4096][r.below(3) as usize] }; let w: Vec = (0..c * out * k).map(|_| if r.below(4) == 0 { 0 } else { r.signed(wmax) as i16 }).collect(); let m: Vec = (0..out).map(|_| if r.below(32) == 0 { r.signed(1 << 40) } else { r.signed(1 << 12) }).collect(); let b: Vec = (0..out) .map(|_| match r.below(64) { 0 => [i64::MAX, i64::MIN, i64::MIN + 1, i64::MAX - 1][r.below(4) as usize], 1 => r.signed(1 << 40), _ => r.signed(1 << 20), }) .collect(); layers.push(Wam::ConvT { inc: c, out, k, s, p, op, w, m, b }); c = out; } let set = wam1_bytes(afrac, gf, shift, &layers); let lat_len = 1 + r.below(6) as usize; let full = audio_len(&layers, lat_len as i64).unwrap_or(1).max(1) as usize; let shorter = if lat_len > 1 { audio_len(&layers, lat_len as i64 - 1).unwrap_or(0).max(0) as usize } else { 0 }; let mut t_out = if full > shorter { shorter + 1 + r.below((full - shorter) as u64) as usize } else { full }; if r.below(8) == 0 { t_out = if r.below(2) == 0 { full + 1 } else { shorter.max(1) }; } let n_ch = c0 as usize; let a = afrac.min(30); let len = 3 + r.below(4) as usize; let mut cuts: Vec = (0..len - 2).map(|_| 1 + r.below(65_535) as i32).collect(); cuts.sort_unstable(); cuts.dedup(); let mut t = vec![0]; t.extend(cuts); t.push(1 << 16); let offset = r.signed(3) as i32; let span = t.len() as i64 - 2; let mut median_q: Vec = (0..n_ch) .map(|_| match r.below(8) { 0 => [1i64 << 47, -(1 << 47), (1 << 47) + 1, -(1 << 47) - 1][r.below(4) as usize], _ => r.signed(1 << (a + 4)), }) .collect(); let mut values: Vec = (0..n_ch * lat_len) .map(|_| match r.below(32) { 0 => r.signed(1 << 20), 1..=4 => offset as i64 + span + r.below(40) as i64 - 20, _ => offset as i64 + r.below(span as u64) as i64, } as i32) .collect(); if r.below(4) == 0 { if let Ok(g) = wai_ops::lower::wam1(&set) { let x0 = g.tensors[0].hi as i128; let ch = r.below(n_ch as u64) as usize; let target = [x0, x0 + 1, -x0, -x0 - 1][r.below(4) as usize]; let med = target - ((offset as i128) << a); if med.unsigned_abs() <= 1 << 47 { median_q[ch] = med as i64; values[ch * lat_len..(ch + 1) * lat_len].iter_mut().for_each(|x| *x = offset); } } } let mut enc = wai::rans::RansEncoder::new(); for &v in &values { assert!(enc.encode(v, &t, t.len(), offset), "a value the generator writes is codable"); } let mut rans = enc.flush(); match r.below(64) { 0..=3 => { let cut = (4 * (1 + r.below(3)) as usize).min(rans.len()); rans.truncate(rans.len() - cut); } 4 | 5 => rans.extend([0, 0, 0, 0]), 6 => { let i = r.below(rans.len() as u64) as usize; rans[i] ^= 1 << r.below(8); } _ => {} } let clip_afrac = if r.below(32) == 0 { afrac ^ 1 } else { afrac }; (set, Clip { afrac: clip_afrac, n_ch, lat_len, t_out, offset, median_q, cdfs: vec![t; n_ch], rans }) } /// The direct audio decode, a panic (an overflow-checked build's arithmetic overflow) /// caught. fn audio_direct(m: &wai::int_audio::AudioModel, cl: &Clip) -> Option, RefusalKind>> { catch_unwind(AssertUnwindSafe(|| wai::int_audio::legacy::decode_typed(m, &cl.st()).map_err(|r| r.kind))).ok() } /// Gate 4 of the migration, for `wai.audio.int_codec`: the decode through the operator /// set against the direct decode on random sets and clips. Outside the planned breaks /// both decode to the same samples or both refuse with the same code, and the direct /// decode never panics; under a planned break the decode through the operator set /// refuses. #[test] fn int_audio_random_sets_and_clips_decode_alike() { let (mut alike, mut refused, mut breaks) = (0u64, std::collections::BTreeMap::::new(), std::collections::BTreeMap::<&str, u64>::new()); let n = cases("int_audio", 3000, |r, seed| { let (set, cl) = int_audio_case(r); let why = || format!("seed {seed}\nclip {cl:?}"); let m = wai::int_audio::AudioModel::from_bytes(&set).expect("the generator writes WAM1 that parses"); let new = wai::int_audio::ops::decode_typed(&m, &cl.st()).map_err(|r| r.kind); match int_audio_planned_break(&set) { Some(b) => { assert!(new.is_err(), "{b}: the decode through the operator set must refuse it\n{}", why()); *breaks.entry(b).or_default() += 1; } None => match (&new, audio_direct(&m, &cl)) { (Ok(a), Some(Ok(b))) => { assert!(*a == b, "the two decodes differ\n{}", why()); alike += 1; } (Err(a), Some(Err(b))) => { assert_eq!(*a, b, "the two decodes refuse with different codes\n{}", why()); *refused.entry(format!("{a:?}")).or_default() += 1; } (_, old) => panic!("new {new:?}, direct {old:?}\n{}", why()), }, } }); eprintln!("int_audio: {alike} decoded alike, refused by both {refused:?}, planned breaks {breaks:?}"); assert!(alike * 5 > n, "the generator must produce decodable cases: {alike} of {n}"); assert!(n < 1000 || refused.contains_key("FieldRange"), "no case reached a check range"); assert!(n < 1000 || breaks.len() >= 3, "too few kinds of planned break reached: {breaks:?}"); } /// The planned breaks of integer-operators §15.3 item 3 for `wai.audio.int_codec`, each /// on a set and clip the direct decode decodes, or refuses with another code: the /// classifier names it, and the decode through the operator set refuses it /// (`model_mismatch`, before the stream). #[test] fn int_audio_planned_breaks_are_new_refusals() { use wai::int_audio::AudioModel; // A clip of one channel of 4 latents, each 2·2^12, presenting 4 samples; with // `freq1`, each symbol costs 16 bits, so a stream cut to its seed runs out. let clip_with = |t: Vec| { let mut enc = wai::rans::RansEncoder::new(); for _ in 0..4 { assert!(enc.encode(2, &t, 3, 2)); } Clip { afrac: 12, n_ch: 1, lat_len: 4, t_out: 4, offset: 2, median_q: vec![0], cdfs: vec![t], rans: enc.flush() } }; let conv1 = |w: i16, m: i64, b: i64| Wam::ConvT { inc: 1, out: 1, k: 1, s: 1, p: 0, op: 0, w: vec![w], m: vec![m], b: vec![b] }; let expect = |b: &str, set: &[u8], cl: &Clip, direct: Option| { assert_eq!(int_audio_planned_break(set), Some(b), "{b}: the classifier"); let m = AudioModel::from_bytes(set).unwrap(); match direct { None => assert!(matches!(audio_direct(&m, cl), Some(Ok(_))), "{b}: the direct decode decodes it"), Some(k) => assert_eq!(audio_direct(&m, cl), Some(Err(k)), "{b}: the direct decode's code"), } let new = wai::int_audio::ops::decode_typed(&m, &cl.st()).map(|_| ()).map_err(|r| r.kind); assert_eq!(new, Err(RefusalKind::ModelMismatch), "{b}: the decode through the operator set refuses it"); }; // The base decodes alike: one 1×1 transposed convolution, weight 1, shift 0, so each // sample is (2·2^12·32767 + 2^11) >> 12 = 65534, clamped to 32767. let cl = clip_with(vec![0, 1 << 15, 1 << 16]); let base = wam1_bytes(12, 16, 0, &[conv1(1, 1, 0)]); assert_eq!(int_audio_planned_break(&base), None); let m = AudioModel::from_bytes(&base).unwrap(); assert_eq!(audio_direct(&m, &cl), Some(Ok(vec![32767; 4]))); assert_eq!(wai::int_audio::ops::decode_typed(&m, &cl.st()).unwrap(), vec![32767; 4]); // No layers: the direct decode presents the latents themselves. expect("no layer of tag 0, 1 or 2", &wam1_bytes(12, 16, 0, &[]), &cl, None); // A header shift no layer uses (only an inverse GDN), or a gf no layer uses. let igdn = Wam::Igdn { c: 1, g: vec![0], b: vec![1 << 24] }; expect("a shift or gf of 128 or more that no layer uses", &wam1_bytes(12, 16, 200, &[igdn]), &cl, None); expect("a shift or gf of 128 or more that no layer uses", &wam1_bytes(12, 200, 0, &[conv1(1, 1, 0)]), &cl, None); // A bias at the end of i64: refused at that layer, after the stream, directly; at // load through the lowering, so a stream that runs out is refused for the set. expect("a bias of magnitude 2^63 - 1 or more", &wam1_bytes(12, 16, 0, &[conv1(1, 1, i64::MAX)]), &cl, Some(RefusalKind::FieldRange)); let mut short = clip_with(vec![0, 1, 1 << 16]); short.rans.truncate(8); expect("a bias of magnitude 2^63 - 1 or more", &wam1_bytes(12, 16, 0, &[conv1(1, 1, i64::MIN)]), &short, Some(RefusalKind::StreamExhausted)); // A check range its producer cannot meet: every value 2^30 into a layer whose range is // ±4095 (weight 2^11, multiplier 2^40, shift 0); and every sample 2^62, past the PCM // step's range. The direct decode refuses each at that step (field_range). expect("a check range its producer cannot meet", &wam1_bytes(12, 16, 0, &[conv1(0, 1, 1 << 30), conv1(1 << 11, 1 << 40, 0)]), &cl, Some(RefusalKind::FieldRange)); expect("a check range its producer cannot meet", &wam1_bytes(12, 16, 0, &[conv1(0, 1, 1 << 62)]), &cl, Some(RefusalKind::FieldRange)); } // ------------------------------------------------------------------ int_hyper (M4) /// A conditional codec's entropy tables (`entropy.json`) and a `WIH1` payload's typed /// fields, owned. #[derive(Clone, Debug)] struct Cond { edges: Vec, cdfs: Vec>, cdf_lengths: Vec, offsets: Vec, z_cdfs: Vec>, z_cdf_lengths: Vec, z_offsets: Vec, z_median_q: Vec, m: usize, sh: usize, sw: usize, n_z: usize, zh: usize, zw: usize, z_rans: Vec, rans: Vec, } impl Cond { fn ent(&self) -> wai::int_hyper_synth::CondEntropy { wai::int_hyper_synth::CondEntropy { edges: self.edges.clone(), cdfs: self.cdfs.clone(), cdf_lengths: self.cdf_lengths.clone(), offsets: self.offsets.clone(), z_cdfs: self.z_cdfs.clone(), z_cdf_lengths: self.z_cdf_lengths.clone(), z_offsets: self.z_offsets.clone(), z_median_q: self.z_median_q.clone(), } } fn st(&self) -> wai::int_hyper_synth::CondStream<'_> { wai::int_hyper_synth::CondStream { m: self.m, sh: self.sh, sw: self.sw, n_z: self.n_z, zh: self.zh, zw: self.zw, z_rans: &self.z_rans, rans: &self.rans } } } /// The planned break of integer-operators §15.3 that a conditional codec's set falls /// under, if any, read from its two networks' bytes alone: either network is refused at /// load by its lowering (`h_s` by `whm1/1`, `g_s` by `whm1/1` with the RGB8 step, each /// loaded under `acc64`) for a rule the direct decode does not make (see /// [`int_splat_planned_break`]). The entropy tables are not lowered, and no payload field /// changes meaning. fn int_hyper_planned_break(g_s: &[u8], h_s: &[u8]) -> Option<&'static str> { let (h, g) = (wai_ops::lower::parse_whm1(h_s).ok()?, wai_ops::lower::parse_whm1(g_s).ok()?); load_break(&h, wai_ops::lower::whm1(h_s, false)).or_else(|| load_break(&g, wai_ops::lower::whm1(g_s, true))) } /// A random CDF table of 3 to 6 entries. fn cdf_table(r: &mut Rng) -> Vec { let len = 3 + r.below(4) as usize; let mut cuts: Vec = (0..len - 2).map(|_| 1 + r.below(65_535) as i32).collect(); cuts.sort_unstable(); cuts.dedup(); let mut t = vec![0]; t.extend(cuts); t.push(1 << 16); t } /// A last layer of a random network: a convolution (stride 1) or a transposed convolution /// from `c` channels to `out`, with `splat_layers`' parameter draws; one in 32 a ReLU. fn last_layer(r: &mut Rng, c: u32, out: u32) -> Whm { if r.below(32) == 0 { return Whm::Relu; } let k = 1 + r.below(3) as u32; let p = if r.below(16) == 0 { r.below(k as u64 + 2) } else { (k / 2) as u64 } as u32; let wmax = if r.below(8) == 0 { 32767 } else { [16, 1024, 4096][r.below(3) as usize] }; let w: Vec = (0..c * out * k * k).map(|_| if r.below(4) == 0 { 0 } else { r.signed(wmax) as i16 }).collect(); let m: Vec = (0..out).map(|_| if r.below(32) == 0 { r.signed(1 << 40) } else { r.signed(1 << 12) }).collect(); let b: Vec = (0..out) .map(|_| match r.below(64) { 0 => [i64::MAX, i64::MIN, i64::MIN + 1, i64::MAX - 1][r.below(4) as usize], 1 => r.signed(1 << 40), _ => r.signed(1 << 20), }) .collect(); if r.below(2) == 0 { Whm::Conv { inc: c, out, k, s: 1, p, w, m, b } } else { let s = 1 + r.below(2) as u32; Whm::ConvT { inc: c, out, k, s, p: p.min(k.saturating_sub(1)), op: r.below(s as u64) as u32, w, m, b } } } /// `n` values coded against one table each, with `RansEncoder`: `tables[i]` and its length /// and offset for value `i`. `None` where a value is not codable against its table. fn rans_of(values: &[(i32, &[i32], usize, i32)]) -> Option> { let mut enc = wai::rans::RansEncoder::new(); for &(v, t, len, off) in values { if !enc.encode(v, t, len, off) { return None; } } Some(enc.flush()) } /// A stream damaged one case in sixteen (cut short), one in 32 (a word appended) or one in /// 64 (a byte changed). fn damage(r: &mut Rng, s: &mut Vec) { match r.below(64) { 0..=3 => { let cut = (4 * (1 + r.below(3)) as usize).min(s.len()); s.truncate(s.len() - cut); } 4 | 5 => s.extend([0, 0, 0, 0]), 6 if !s.is_empty() => { let i = r.below(s.len() as u64) as usize; s[i] ^= 1 << r.below(8); } _ => {} } } /// A random `wai.neural.int_hyper` set and `WIH1` payload. The set: `afrac` 12 mostly (one /// in sixteen anywhere in 1..=30, one in sixteen at 0 or 31, and `h_s`'s not `g_s`'s one /// in 32); each network's header `gf` or `shift` of 128 or more one set in 32 each; `h_s` /// from `n_z` (1 to 3) channels, up to one of `splat_layers`' layers then a last layer to /// `2·m` channels (one in sixteen to another count); `g_s` from `m` (1 to 3) channels, up /// to two layers then a last layer to three channels (one in sixteen not); the tables: 1 /// to 4 ascending scale edges (one set in 32 none), a latent table, length and offset per /// edge, and a hyper-latent table, offset and median per hyper-latent channel (one set in /// 32 a channel short), one length in 32 past its row; medians mostly small, one in eight /// at or past the 2^47 bound or near i64's ends. The payload: a hyper-latent of side 2 or /// 3 (one in sixteen not square), its symbols mostly inside their table and some escaped, /// one case in four with a whole channel at `h_s`'s input `check` range or one past it; /// the latent's shape `h_s`'s output (one case in eight off by one), each symbol coded /// against the table the direct decode's `σ` buckets it to (where `h_s` runs directly); /// each stream damaged as [`damage`] says. fn int_hyper_case(r: &mut Rng) -> (Vec, Vec, Cond) { let afrac = match r.below(16) { 0 => 1 + r.below(30) as u32, 1 => [0, 31][r.below(2) as usize], _ => 12, }; let h_afrac = if r.below(32) == 0 { afrac ^ 1 } else { afrac }; let header = |r: &mut Rng| { let gf = if r.below(32) == 0 { 128 + r.below(64) as u32 } else { 16 }; let shift = if r.below(32) == 0 { 128 + r.below(64) as u32 } else { 8 + r.below(16) as u32 }; (gf, shift) }; let (n_z, m) = (1 + r.below(3) as u32, 1 + r.below(3) as u32); let (n, relu_only) = (r.below(2) as usize, r.below(32) == 0); let (mut hl, c) = splat_layers(r, n_z, n, relu_only); let out = if r.below(16) == 0 { 1 + r.below(6) as u32 } else { 2 * m }; hl.push(last_layer(r, c, out)); let (gf, shift) = header(r); let h_s = whm1_bytes(h_afrac, gf, shift, &hl); let (n, relu_only) = (r.below(3) as usize, r.below(32) == 0); let (mut gl, c) = splat_layers(r, m, n, relu_only); let out = if r.below(16) == 0 { 1 + r.below(4) as u32 } else { 3 }; gl.push(last_layer(r, c, out)); let (gf, shift) = header(r); let g_s = whm1_bytes(afrac, gf, shift, &gl); let a = afrac.min(30); let n_edges = if r.below(32) == 0 { 0 } else { 1 + r.below(4) as usize }; let mut edges = Vec::new(); let mut e = 1 + r.below(1 << (a + 2)) as i64; for _ in 0..n_edges { edges.push(e); e += 1 + r.below(1 << (a + 2)) as i64; } let cdfs: Vec> = (0..n_edges).map(|_| cdf_table(r)).collect(); let mut cdf_lengths: Vec = cdfs.iter().map(Vec::len).collect(); let offsets: Vec = (0..n_edges).map(|_| r.signed(3) as i32).collect(); let k = if r.below(32) == 0 { n_z as usize - 1 } else { n_z as usize }; let z_cdfs: Vec> = (0..k).map(|_| cdf_table(r)).collect(); let mut z_cdf_lengths: Vec = z_cdfs.iter().map(Vec::len).collect(); if r.below(32) == 0 { match (n_edges, k) { (1.., _) if r.below(2) == 0 => cdf_lengths[0] += 1, (_, 1..) => z_cdf_lengths[0] += 1, _ => {} } } let z_offsets: Vec = (0..k).map(|_| r.signed(3) as i32).collect(); let mut z_median_q: Vec = (0..k) .map(|_| match r.below(8) { 0 => [1i64 << 47, -(1 << 47), 1 << 60, -(1 << 60), i64::MAX, i64::MIN][r.below(6) as usize], _ => r.signed(1 << (a + 4)), }) .collect(); let zh = 2 + r.below(2) as usize; let zw = if r.below(16) == 0 { 1 + r.below(3) as usize } else { zh }; let (nz, plane) = (n_z as usize, zh * zw); // The hyper-latent's symbol values, per channel against its own table (channel 0's // where a channel has none). let zt = |ch: usize| ch.min(k.saturating_sub(1)); let mut zv: Vec = Vec::with_capacity(nz * plane); for ch in 0..nz { let (off, span) = if k == 0 { (0, 1) } else { (z_offsets[zt(ch)] as i64, z_cdfs[zt(ch)].len() as i64 - 2) }; for _ in 0..plane { zv.push(match r.below(32) { 0 => r.signed(1 << 20), 1..=4 => off + span + r.below(40) as i64 - 20, _ => off + r.below(span.max(1) as u64) as i64, } as i32); } } // One case in four: a whole hyper-latent channel at h_s's input check range, or one // past it, where h_s lowers. if r.below(4) == 0 && k > 0 { if let Ok(g) = wai_ops::lower::whm1(&h_s, false) { let x0 = g.tensors[0].hi as i128; let ch = r.below(k as u64) as usize; let v = z_offsets[ch]; let target = [x0, x0 + 1, -x0, -x0 - 1][r.below(4) as usize]; let med = target - ((v as i128) << a); if let Ok(med) = i64::try_from(med) { z_median_q[ch] = med; zv[ch * plane..(ch + 1) * plane].iter_mut().for_each(|x| *x = v); } } } let z_rans = if k == 0 { vec![0, 0, 0, 128, 0, 0, 0, 0] } else { let vals: Vec<(i32, &[i32], usize, i32)> = (0..nz * plane).map(|i| { let t = zt(i / plane); (zv[i], z_cdfs[t].as_slice(), z_cdf_lengths[t], z_offsets[t]) }).collect(); rans_of(&vals).unwrap_or_else(|| vec![0, 0, 0, 128, 0, 0, 0, 0]) }; // The latent: h_s's output shape, and each symbol's table from the direct decode's σ. let hm = IntModel::from_bytes(&h_s).expect("the generator writes WHM1 that parses"); let side = |n: usize| splat_side(&hl, n as i64).unwrap_or(1) as usize; let (mut sh, sw) = (side(zh), side(zw)); if r.below(8) == 0 { sh += 1; } let mu = m as usize; let z_hat: Option> = (0..nz * plane) .map(|i| { let med = *z_median_q.get(zt(i / plane))?; (zv[i] as i64).checked_mul(1 << a)?.checked_add(med) }) .collect(); let sigma: Option> = z_hat.and_then(|z| { let (o, oc, oh, ow) = wai::int_hyper_synth::legacy::run_net_checked(&hm, &z, nz, zh, zw).ok()?; ((oc, oh, ow) == (2 * mu, sh, sw)).then(|| o[..mu * sh * sw].to_vec()) }); let ny = mu * sh * sw; let rans = match (&sigma, n_edges) { (Some(s), 1..) => { let mut vals = Vec::with_capacity(ny); for p in 0..ny { let b = wai::int_entropy::det_build_index(s[p], &edges); let span = cdfs[b].len() as i64 - 2; let v = match r.below(32) { 0 => r.signed(1 << 20), 1..=3 => offsets[b] as i64 + span + r.below(40) as i64 - 20, _ => offsets[b] as i64 + r.below(span.max(1) as u64) as i64, } as i32; vals.push((v, cdfs[b].as_slice(), cdf_lengths[b], offsets[b])); } rans_of(&vals).unwrap_or_else(|| (0..8 + 4 * r.below(8)).map(|_| r.below(256) as u8).collect()) } _ => (0..8 + 4 * r.below(8)).map(|_| r.below(256) as u8).collect(), }; let mut cond = Cond { edges, cdfs, cdf_lengths, offsets, z_cdfs, z_cdf_lengths, z_offsets, z_median_q, m: mu, sh, sw, n_z: nz, zh, zw, z_rans, rans, }; damage(r, &mut cond.z_rans); damage(r, &mut cond.rans); (g_s, h_s, cond) } /// The direct conditional decode, a panic (an overflow-checked build's arithmetic /// overflow) caught. fn hyper_direct(g: &IntModel, h: &IntModel, c: &Cond) -> Option, RefusalKind>> { let e = c.ent(); catch_unwind(AssertUnwindSafe(|| wai::int_hyper_synth::legacy::decode_conditional_checked(g, h, &e, &c.st()).map(|d| d.0).map_err(|r| r.kind))).ok() } /// Gate 4 of the migration, for `wai.neural.int_hyper` (and so the keyframe of /// `wai.video.int_hyper`): the decode through the operator set against the direct decode /// on random sets and payloads. Outside the planned breaks both decode to the same image /// or both refuse with the same code, and the direct decode never panics; under a planned /// break the decode through the operator set refuses. #[test] fn int_hyper_random_sets_and_payloads_decode_alike() { let (mut alike, mut refused, mut breaks) = (0u64, std::collections::BTreeMap::::new(), std::collections::BTreeMap::<&str, u64>::new()); let n = cases("int_hyper", 3000, |r, seed| { let (g_s, h_s, c) = int_hyper_case(r); let why = || format!("seed {seed}\npayload {c:?}"); let (g, h) = (IntModel::from_bytes(&g_s).unwrap(), IntModel::from_bytes(&h_s).unwrap()); let new = wai::int_hyper_synth::ops::decode(&g, &h, &c.ent(), &c.st()).map(|d| d.0).map_err(|r| r.kind); match int_hyper_planned_break(&g_s, &h_s) { Some(b) => { assert!(new.is_err(), "{b}: the decode through the operator set must refuse it\n{}", why()); *breaks.entry(b).or_default() += 1; } None => match (&new, hyper_direct(&g, &h, &c)) { (Ok(a), Some(Ok(b))) => { assert!(*a == b, "the two decodes differ\n{}", why()); alike += 1; } (Err(a), Some(Err(b))) => { assert_eq!(*a, b, "the two decodes refuse with different codes\n{}", why()); *refused.entry(format!("{a:?}")).or_default() += 1; } (_, old) => panic!("new {new:?}, direct {old:?}\n{}", why()), }, } }); eprintln!("int_hyper: {alike} decoded alike, refused by both {refused:?}, planned breaks {breaks:?}"); // Two networks and two streams compound the generator's refusal rates, so fewer cases // decode than for one network; an eighth is still hundreds per push. assert!(alike * 8 > n, "the generator must produce decodable cases: {alike} of {n}"); assert!(n < 1000 || refused.contains_key("FieldRange"), "no case reached a check range"); assert!(n < 1000 || ["StreamExhausted", "StreamNoncanonical", "GeometryMismatch", "ModelMismatch"].iter().all(|k| refused.contains_key(*k)), "a refusal kind was not reached: {refused:?}"); assert!(n < 1000 || breaks.len() >= 3, "too few kinds of planned break reached: {breaks:?}"); } /// The planned breaks of integer-operators §15.3 item 3 for `wai.neural.int_hyper`, each /// on a set and payload the direct decode decodes, or refuses with another code: the /// classifier names it, and the decode through the operator set refuses it /// (`model_mismatch`, before either stream). #[test] fn int_hyper_planned_breaks_are_new_refusals() { // A payload of two hyper-latent channels of side 2, every symbol the first of the // table `t` (offset 2), so every hyper-latent is 2·2^12; a latent of one channel of // side 2, every symbol the first of the one scale edge's table (offset 2). let payload = |t: Vec| { let z: Vec<(i32, &[i32], usize, i32)> = (0..8).map(|_| (2, t.as_slice(), t.len(), 2)).collect(); let y = [0, 1 << 15, 1 << 16]; let yv: Vec<(i32, &[i32], usize, i32)> = (0..4).map(|_| (2, y.as_slice(), 3, 2)).collect(); Cond { edges: vec![1 << 12], cdfs: vec![y.to_vec()], cdf_lengths: vec![3], offsets: vec![2], z_cdfs: vec![t.clone(); 2], z_cdf_lengths: vec![t.len(); 2], z_offsets: vec![2; 2], z_median_q: vec![0; 2], m: 1, sh: 2, sw: 2, n_z: 2, zh: 2, zw: 2, z_rans: rans_of(&z).unwrap(), rans: rans_of(&yv).unwrap(), } }; let cl = payload(vec![0, 1 << 15, 1 << 16]); // h_s: σ and μ are the two hyper-latent channels; g_s: the latent to each of three channels. let ident = Whm::Conv { inc: 2, out: 2, k: 1, s: 1, p: 0, w: vec![1, 0, 0, 1], m: vec![1; 2], b: vec![0; 2] }; let to3 = |w: i16, b: i64| Whm::ConvT { inc: 1, out: 3, k: 1, s: 1, p: 0, op: 0, w: vec![w; 3], m: vec![1; 3], b: vec![b; 3] }; let (h0, g0) = (whm1_bytes(12, 16, 0, &[ident.clone()]), whm1_bytes(12, 16, 0, &[to3(1, 0)])); let decode = |g_s: &[u8], h_s: &[u8], c: &Cond| { let (g, h) = (IntModel::from_bytes(g_s).unwrap(), IntModel::from_bytes(h_s).unwrap()); (wai::int_hyper_synth::ops::decode(&g, &h, &c.ent(), &c.st()).map(|d| d.0).map_err(|r| r.kind), hyper_direct(&g, &h, c)) }; // The base decodes alike: every latent 2·2^12 + 2·2^12, clamped to 2^12: 255. assert_eq!(int_hyper_planned_break(&g0, &h0), None); let (new, old) = decode(&g0, &h0, &cl); assert_eq!((new, old), (Ok(vec![255; 12]), Some(Ok(vec![255; 12])))); let expect = |b: &str, g_s: &[u8], h_s: &[u8], c: &Cond, direct: Option| { assert_eq!(int_hyper_planned_break(g_s, h_s), Some(b), "{b}: the classifier"); let (new, old) = decode(g_s, h_s, c); match direct { None => assert!(matches!(old, Some(Ok(_))), "{b}: the direct decode decodes it: {old:?}"), Some(k) => assert_eq!(old, Some(Err(k)), "{b}: the direct decode's code"), } assert_eq!(new, Err(RefusalKind::ModelMismatch), "{b}: the decode through the operator set refuses it"); }; // No layer of tag 0, 1 or 2 in h_s: σ and μ are the hyper-latent's channels directly. expect("no layer of tag 0, 1 or 2", &g0, &whm1_bytes(12, 16, 0, &[Whm::Relu]), &cl, None); // A header shift no layer of h_s uses (only an inverse GDN that leaves the values as // they are), and a gf no layer of g_s uses. let igdn = Whm::Igdn { c: 2, g: vec![0; 4], b: vec![1 << 24; 2] }; expect("a shift or gf of 128 or more that no layer uses", &g0, &whm1_bytes(12, 16, 200, &[igdn]), &cl, None); expect("a shift or gf of 128 or more that no layer uses", &whm1_bytes(12, 200, 0, &[to3(1, 0)]), &h0, &cl, None); // A bias at the end of i64 in g_s: refused directly as g_s runs (field_range), after // both streams; at load through the lowering. In h_s, with a hyper-latent stream that // runs out: refused directly for the stream, through the lowering for the set. expect("a bias of magnitude 2^63 - 1 or more", &whm1_bytes(12, 16, 0, &[to3(1, i64::MAX)]), &h0, &cl, Some(RefusalKind::FieldRange)); let mut short = payload(vec![0, 1, 1 << 16]); short.z_rans.truncate(8); let h_min = Whm::Conv { inc: 2, out: 2, k: 1, s: 1, p: 0, w: vec![1, 0, 0, 1], m: vec![1; 2], b: vec![i64::MIN; 2] }; expect("a bias of magnitude 2^63 - 1 or more", &g0, &whm1_bytes(12, 16, 0, &[h_min]), &short, Some(RefusalKind::StreamExhausted)); // A check range its producer cannot meet, in g_s: its first layer gives every value // 2^30 (weight 0, bias 2^30), into a layer of weight 2^11 under the multiplier 2^40 // and shift 0, whose range is ±4095. Refused directly at that layer (field_range). let first = Whm::ConvT { inc: 1, out: 1, k: 1, s: 1, p: 0, op: 0, w: vec![0], m: vec![1], b: vec![1 << 30] }; let tight = Whm::ConvT { inc: 1, out: 3, k: 1, s: 1, p: 0, op: 0, w: vec![1 << 11; 3], m: vec![1 << 40; 3], b: vec![0; 3] }; expect("a check range its producer cannot meet", &whm1_bytes(12, 16, 0, &[first, tight]), &h0, &cl, Some(RefusalKind::FieldRange)); } // ------------------------------------------------------------------ int_mlicpp (M5) /// A `WMB1` bundle's bytes: the three `WHM1` networks, then the tables (with `afrac`). fn wmb1_bytes(h_s: &[u8], entropy: &[u8], g_s: &[u8], tables: &Value) -> Vec { let t = serde_json::to_vec(tables).unwrap(); let mut v = b"WMB1".to_vec(); for s in [h_s, entropy, g_s, t.as_slice()] { v.extend((s.len() as u32).to_le_bytes()); v.extend(s); } v } /// A `WMB1` bundle's four sections. fn wmb1_sections(b: &[u8]) -> Option<[Vec; 4]> { let mut o = 4; let mut s: [Vec; 4] = Default::default(); for x in &mut s { let n = u32::from_le_bytes(b.get(o..o + 4)?.try_into().ok()?) as usize; *x = b.get(o + 4..o + 4 + n)?.to_vec(); o += 4 + n; } Some(s) } /// A `WIM1` payload's bytes: `H`, `W`; `Cz, Zh, Zw, Cy, Yh, Yw, K`; the two streams. fn wim1_bytes(h: u32, w: u32, dims: [u16; 7], z: &[u8], y: &[u8]) -> Vec { let mut p = b"WIM1".to_vec(); p.extend(h.to_le_bytes()); p.extend(w.to_le_bytes()); dims.iter().for_each(|d| p.extend(d.to_le_bytes())); p.extend(0u16.to_le_bytes()); p.extend((z.len() as u32).to_le_bytes()); p.extend((y.len() as u32).to_le_bytes()); p.extend(z); p.extend(y); p } /// The planned break of integer-operators §15.3 that a `wai.neural.int_mlicpp` case falls /// under, if any, read from the bundle and the payload's header alone. A bundle that does /// not parse, or a header or bundle that the checks both decodes make first refuse /// (`int_mlicpp::ops::check`: the header, the `z` tables, each network's geometry on the /// header's shapes), is not a break: both refuse it alike. Otherwise the decode through the /// operator set loads the three networks (`h_s`, `entropy`, `g_s` with the RGB8 step, by /// `whm1/1` under `acc64`) and plans each on the payload's shapes before either stream, /// and refuses (`model_mismatch`) a bundle that: /// - falls under one of the breaks of the other `WHM1` decoders ([`load_break`]); /// - has a layer that does not fit the channels the layer before it gives, a header shift /// or `gf` of 128 or more that a layer uses, or a leaky ReLU's shift of 64 or more, which /// the direct decode refuses (`model_mismatch` too) only when that network runs, after /// the streams and networks before it, so an earlier refusal of the payload comes first /// there; /// - has a network whose input is not the payload's channel count (`Cz`, `3·Cy`, `Cy`), /// which the direct decode likewise refuses when that network runs; or a convolution /// whose padded side exceeds `2^26`, which the operator set's run-time plan caps /// (integer-operators §5.2) and the direct decode's geometry does not. fn int_mlicpp_planned_break(bundle: &[u8], payload: &[u8]) -> Option<&'static str> { use wai_ops::lower::NetLayer as N; let m = wai::int_mlicpp::IntMlicModel::from_bundle(bundle).ok()?; let hd = wai::int_mlicpp::ops::check(&m, payload).ok()?; let secs = wmb1_sections(bundle)?; for (k, rgb8) in [(0, false), (1, false), (2, true)] { let net = wai_ops::lower::parse_whm1(&secs[k]).ok()?; if let Some(b) = load_break(&net, wai_ops::lower::whm1(&secs[k], rgb8)) { return Some(b); } let conv = net.layers.iter().any(|l| matches!(l, N::ConvT { .. } | N::Conv { .. })); let igdn = net.layers.iter().any(|l| matches!(l, N::Igdn { .. })); if (net.shift > 127 && conv) || (net.gf > 127 && igdn) { return Some("a shift or gf of 128 or more that a layer uses"); } if net.layers.iter().any(|l| matches!(l, N::Leaky { slope_frac, .. } if *slope_frac > 63)) { return Some("a leaky ReLU's shift of 64 or more"); } let mut c: Option = None; for l in &net.layers { let (inc, out) = match l { N::ConvT { inc, out, .. } | N::Conv { inc, out, .. } => (*inc, *out), N::Igdn { c, .. } => (*c, *c), _ => continue, }; if c.is_some_and(|c| c != inc) { return Some("a layer that does not fit the channels before it"); } c = Some(out); } } let n = wai::int_mlicpp::ops::load(&m).ok()?; let (cy, yh, yw) = (hd.cy as u64, hd.yh as u64, hd.yw as u64); for (g, input) in [(n.h_s.graph(), (hd.cz as u64, hd.zh as u64, hd.zw as u64)), (n.entropy.graph(), (3 * cy, yh, yw)), (n.g_s.graph(), (cy, yh, yw))] { match wai_ops::plan::run_plan(g, &[input], &[]) { Err(r) if r.code == wai_ops::Code::ModelMismatch && r.at == wai_ops::Site::Input(0) => { return Some("a network whose input is not the payload's channel count"); } Err(r) if r.code == wai_ops::Code::TooLarge => return Some("a convolution's padded side past 2^26"), _ => {} } } None } /// A refusal of either `int_mlicpp` decode, as both name it: the code (integer-payloads §8), /// then which check: for a network as it runs, the network (the direct decode names the /// layer's arithmetic, the operator set the `check` range); for any other, the decode's /// own message, which the two share. fn mlic_label(e: &wai::int_mlicpp::IntMlicError) -> String { use wai::int_mlicpp::IntMlicError as E; let (code, m) = match e { E::Payload(m) => ("payload_malformed", m.as_str()), E::Model(m) => ("model_mismatch", m.as_str()), }; let m = match m.find(" refused") { Some(i) if ["h_s", "entropy", "g_s"].contains(&&m[..i]) => &m[..i], _ => m, }; format!("{code}: {m}") } /// One `int_mlicpp` decode's outcome: the `y` symbols it read (buckets and symbols), and /// the image or its refusal's label. type MlicOutcome = (Vec, Vec, Result, String>); fn mlic_outcome((t, r): (wai::int_mlicpp::Trace, Result)) -> MlicOutcome { (t.buckets, t.symbols, r.map(|i| i.rgb).map_err(|e| mlic_label(&e))) } /// A layer of an entropy network: a stride-1 convolution from `c` channels to `out` that /// keeps the size (kernel 1 or 3, padding `k / 2`; one in 32 another padding); weights of /// up to 2^4, 2^8 or 2^10 (one in sixteen up to i16's range), a quarter 0; multipliers up /// to 2^12 (one in 128 up to 2^40); biases mostly up to 2^20, one in 256 at an end of i64. fn same_size_layer(r: &mut Rng, c: u32, out: u32) -> Whm { let k = [1u32, 3][r.below(2) as usize]; let p = if r.below(32) == 0 { r.below(3) as u32 } else { k / 2 }; let wmax = if r.below(16) == 0 { 32767 } else { [16, 256, 1024][r.below(3) as usize] }; let w: Vec = (0..c * out * k * k).map(|_| if r.below(4) == 0 { 0 } else { r.signed(wmax) as i16 }).collect(); let m: Vec = (0..out).map(|_| if r.below(128) == 0 { r.signed(1 << 40) } else { r.signed(1 << 12) }).collect(); let b: Vec = (0..out) .map(|_| match r.below(256) { 0 => [i64::MAX, i64::MIN, i64::MIN + 1, i64::MAX - 1][r.below(4) as usize], 1 => r.signed(1 << 40), _ => r.signed(1 << 20), }) .collect(); Whm::Conv { inc: c, out, k, s: 1, p, w, m, b } } /// A random `wai.neural.int_mlicpp` bundle and `WIM1` payload. The bundle: `afrac` 12 /// mostly (one in sixteen anywhere in 1..=30, one in 32 at 0 or 31, and one network's not /// the others' one in 64); each network's header `shift` 12 to 23, and its `gf` or `shift` /// of 128 or more one in 64 each; `h_s` from `Cz` (1 to 3) channels, up to one of `splat_layers`' layers then a last /// layer to `2·Cy` (one in 32 another count); `entropy` from `3·Cy`, up to one layer that /// keeps the size (a leaky ReLU, a ReLU, an inverse GDN, or [`same_size_layer`]) and up to /// one convolution and leaky ReLU more, then [`same_size_layer`] to `2·Cy` (one in 32 /// another count); `g_s` from `Cy` (1, 2 or 4), up to two layers then a last layer to three /// channels (one in 32 not); the tables as [`int_hyper_case`]'s, with `afrac` (one set in /// 64 with no scale edge, a hyper-latent channel short, or a table's length past its row). /// The payload: `K` a divisor of `Cy` (one in 64 not); a hyper-latent of side 2 or 3 (one /// in sixteen not square), its symbols as [`int_hyper_case`]'s (medians at an extreme one /// in sixteen), one case in eight with a whole channel at `h_s`'s input `check` range or /// one past it; the latent's size `h_s`'s /// output and the image's `g_s`'s (each one in sixteen off by one); the `y` stream coded in /// the decode's order (each slice's anchors, then the rest), each symbol against the table /// the direct networks' `σ` buckets it to, mostly inside its table, some escaped and some /// far, until a network refuses (the rest of the stream is then random); each stream /// damaged as [`damage`] says. fn int_mlicpp_case(r: &mut Rng) -> (Vec, Vec) { use wai::int_hyper_synth::legacy::run_net_checked; let afrac = match r.below(32) { 0 | 1 => 1 + r.below(30) as u32, 2 => [0, 31][r.below(2) as usize], _ => 12, }; let odd = r.below(64) == 0; let header = |r: &mut Rng| { let gf = if r.below(64) == 0 { 128 + r.below(64) as u32 } else { 16 }; let shift = if r.below(64) == 0 { 128 + r.below(64) as u32 } else { 12 + r.below(12) as u32 }; (gf, shift) }; let cz = 1 + r.below(3) as u32; let cy = [1u32, 2, 2, 4][r.below(4) as usize]; let divisors: Vec = [1u32, 2, 4].into_iter().filter(|d| cy % d == 0).collect(); let kk = if r.below(64) == 0 { 3 } else { divisors[r.below(divisors.len() as u64) as usize] }; // h_s: [Cz, Zh, Zw] → [2·Cy, Yh, Yw]. let (n, relu_only) = (r.below(2) as usize, r.below(32) == 0); let (mut hl, c) = splat_layers(r, cz, n, relu_only); let out = if r.below(32) == 0 { 1 + r.below(6) as u32 } else { 2 * cy }; hl.push(last_layer(r, c, out)); let (gf, shift) = header(r); let h_s = whm1_bytes(if odd { afrac ^ 1 } else { afrac }, gf, shift, &hl); // entropy: [3·Cy, Yh, Yw] → [2·Cy, Yh, Yw], keeping the size. let mut el = Vec::new(); let mut c = 3 * cy; match r.below(8) { 0 | 1 => el.push(Whm::Leaky { q: if r.below(32) == 0 { r.signed(1 << 40) } else { r.signed(1 << 10) }, f: if r.below(32) == 0 { 64 } else { r.below(16) as u32 } }), 2 => el.push(Whm::Relu), 3 => { let g: Vec = (0..c * c).map(|_| r.below(1 << 12) as i32).collect(); let b: Vec = (0..c).map(|_| r.below(1 << 28) as i64).collect(); el.push(Whm::Igdn { c, g, b }); } 4 | 5 => { let mid = 1 + r.below(6) as u32; el.push(same_size_layer(r, c, mid)); c = mid; } _ => {} } if r.below(2) == 0 { let mid = 1 + r.below(6) as u32; el.push(same_size_layer(r, c, mid)); el.push(Whm::Leaky { q: r.signed(1 << 10), f: r.below(16) as u32 }); c = mid; } let out = if r.below(32) == 0 { 1 + r.below(6) as u32 } else { 2 * cy }; el.push(same_size_layer(r, c, out)); let (gf, shift) = header(r); let entropy = whm1_bytes(afrac, gf, shift, &el); // g_s: [Cy, Yh, Yw] → [3, H, W]. let (n, relu_only) = (r.below(3) as usize, r.below(32) == 0); let (mut gl, c) = splat_layers(r, cy, n, relu_only); let out = if r.below(32) == 0 { 1 + r.below(4) as u32 } else { 3 }; gl.push(last_layer(r, c, out)); let (gf, shift) = header(r); let g_s = whm1_bytes(afrac, gf, shift, &gl); // The tables. let a = afrac.min(30); let n_edges = if r.below(64) == 0 { 0 } else { 1 + r.below(4) as usize }; let mut edges = Vec::new(); let mut e = 1 + r.below(1 << (a + 2)) as i64; for _ in 0..n_edges { edges.push(e); e += 1 + r.below(1 << (a + 2)) as i64; } let cdfs: Vec> = (0..n_edges).map(|_| cdf_table(r)).collect(); let mut cdf_lengths: Vec = cdfs.iter().map(Vec::len).collect(); let offsets: Vec = (0..n_edges).map(|_| r.signed(3) as i32).collect(); let k = if r.below(64) == 0 { cz as usize - 1 } else { cz as usize }; let z_cdfs: Vec> = (0..k).map(|_| cdf_table(r)).collect(); let mut z_cdf_lengths: Vec = z_cdfs.iter().map(Vec::len).collect(); if r.below(64) == 0 { match (n_edges, k) { (1.., _) if r.below(2) == 0 => cdf_lengths[0] += 1, (_, 1..) => z_cdf_lengths[0] += 1, _ => {} } } let z_offsets: Vec = (0..k).map(|_| r.signed(3) as i32).collect(); let mut z_median_q: Vec = (0..k) .map(|_| match r.below(16) { 0 => [1i64 << 47, -(1 << 47), 1 << 60, -(1 << 60), i64::MAX, i64::MIN][r.below(6) as usize], _ => r.signed(1 << (a + 4)), }) .collect(); // The hyper-latent. let zh = 2 + r.below(2) as usize; let zw = if r.below(16) == 0 { 1 + r.below(3) as usize } else { zh }; let (nz, plane) = (cz as usize, zh * zw); let zt = |ch: usize| ch.min(k.saturating_sub(1)); let mut zv: Vec = Vec::with_capacity(nz * plane); for ch in 0..nz { let (off, span) = if k == 0 { (0, 1) } else { (z_offsets[zt(ch)] as i64, z_cdfs[zt(ch)].len() as i64 - 2) }; for _ in 0..plane { zv.push(match r.below(32) { 0 => r.signed(1 << 20), 1..=4 => off + span + r.below(40) as i64 - 20, _ => off + r.below(span.max(1) as u64) as i64, } as i32); } } if r.below(8) == 0 && k > 0 { if let Ok(g) = wai_ops::lower::whm1(&h_s, false) { let x0 = g.tensors[0].hi as i128; let ch = r.below(k as u64) as usize; let v = z_offsets[ch]; let target = [x0, x0 + 1, -x0, -x0 - 1][r.below(4) as usize]; if let Ok(med) = i64::try_from(target - ((v as i128) << a)) { z_median_q[ch] = med; zv[ch * plane..(ch + 1) * plane].iter_mut().for_each(|x| *x = v); } } } let random = |r: &mut Rng| (0..8 + 4 * r.below(8)).map(|_| r.below(256) as u8).collect::>(); let mut z_rans = if k == 0 { random(r) } else { let vals: Vec<(i32, &[i32], usize, i32)> = (0..nz * plane).map(|i| { let t = zt(i / plane); (zv[i], z_cdfs[t].as_slice(), z_cdf_lengths[t], z_offsets[t]) }).collect(); rans_of(&vals).unwrap_or_else(|| random(r)) }; // The sizes, then the latent, coded in the decode's order against the direct networks. let side = |l: &[Whm], n: usize| splat_side(l, n as i64).unwrap_or(1) as usize; let (mut yh, yw) = (side(&hl, zh), side(&hl, zw)); if r.below(16) == 0 { yh += 1; } let (mut hh, ww) = (side(&gl, yh), side(&gl, yw)); if r.below(16) == 0 { hh += 1; } let (hm, em) = (IntModel::from_bytes(&h_s).unwrap(), IntModel::from_bytes(&entropy).unwrap()); let cyu = cy as usize; let ny = cyu * yh * yw; let z_hat: Option> = (0..nz * plane).map(|i| (zv[i] as i64).checked_mul(1 << a)?.checked_add(*z_median_q.get(zt(i / plane))?)).collect(); let hyper = z_hat.and_then(|z| { let (o, oc, oh, ow) = run_net_checked(&hm, &z, nz, zh, zw).ok()?; ((oc, oh, ow) == (2 * cyu, yh, yw)).then_some(o) }); let mut vals: Vec<(i32, Vec, usize, i32)> = Vec::new(); if let (Some(hyper), 1..) = (hyper, n_edges) { let mut ctx = vec![0i64; 3 * ny]; ctx[ny..].copy_from_slice(&hyper); let width = cyu / kk as usize; 'all: for j in 0..kk as usize { for anchor in [true, false] { let Ok((params, ..)) = run_net_checked(&em, &ctx, 3 * cyu, yh, yw) else { break 'all }; if params.len() != 2 * ny { break 'all; } for ch in j * width..(j + 1) * width { for (row, col) in (0..yh).flat_map(|y| (0..yw).map(move |x| (y, x))) { if ((row + col) % 2 == 0) != anchor { continue; } let e = ch * yh * yw + row * yw + col; let b = wai::int_entropy::det_build_index(params[e], &edges); let span = cdfs[b].len() as i64 - 2; let v = match r.below(32) { 0 => r.signed(1 << 20), 1..=3 => offsets[b] as i64 + span + r.below(40) as i64 - 20, _ => offsets[b] as i64 + r.below(span.max(1) as u64) as i64, } as i32; let Some(y) = wai::int_entropy::dequantize(v, a, params[ny + e]) else { break 'all }; vals.push((v, cdfs[b].clone(), cdf_lengths[b], offsets[b])); ctx[e] = y; } } } } } let refs: Vec<(i32, &[i32], usize, i32)> = vals.iter().map(|(v, t, l, o)| (*v, t.as_slice(), *l, *o)).collect(); let mut y_rans = if refs.is_empty() { random(r) } else { rans_of(&refs).unwrap_or_else(|| random(r)) }; damage(r, &mut z_rans); damage(r, &mut y_rans); let tables = json!({ "afrac": afrac, "scale_edges": edges, "cdfs": cdfs, "cdf_lengths": cdf_lengths, "offsets": offsets, "z_cdfs": z_cdfs, "z_cdf_lengths": z_cdf_lengths, "z_offsets": z_offsets, "z_median_q": z_median_q, }); let bundle = wmb1_bytes(&h_s, &entropy, &g_s, &tables); let dims = [cz as u16, zh as u16, zw as u16, cy as u16, yh as u16, yw as u16, kk as u16]; (bundle, wim1_bytes(hh as u32, ww as u32, dims, &z_rans, &y_rans)) } /// The direct `int_mlicpp` decode, a panic (an overflow-checked build's arithmetic overflow) /// caught. fn mlic_direct(m: &wai::int_mlicpp::IntMlicModel, p: &[u8]) -> Option { catch_unwind(AssertUnwindSafe(|| mlic_outcome(wai::int_mlicpp::legacy::decode_traced(m, p)))).ok() } /// Gate 4 of the migration, for `wai.neural.int_mlicpp`: the decode through the operator /// set against the direct decode on random bundles and payloads. Outside the planned /// breaks both read the same `y` symbols, decode to the same image or both refuse with the /// same code at the same check (a network's ranges, or the decode's own), and the direct /// decode never panics; under a planned break the decode through the operator set refuses /// (`model_mismatch`). A bundle that does not parse is refused by both alike before /// either decode starts. #[test] fn int_mlicpp_random_sets_and_payloads_decode_alike() { use std::collections::BTreeMap; let (mut alike, mut unparsed, mut refused, mut breaks) = (0u64, 0u64, BTreeMap::::new(), BTreeMap::<&str, u64>::new()); let n = cases("int_mlicpp", 3000, |r, seed| { let (bundle, p) = int_mlicpp_case(r); let why = || format!("seed {seed}\nbundle {bundle:?}\npayload {p:?}"); let Ok(m) = wai::int_mlicpp::IntMlicModel::from_bundle(&bundle) else { unparsed += 1; return; }; let new = mlic_outcome(wai::int_mlicpp::ops::decode_traced(&m, &p)); match int_mlicpp_planned_break(&bundle, &p) { Some(b) => { assert!(matches!(&new.2, Err(e) if e.starts_with("model_mismatch")), "{b}: the decode through the operator set must refuse it: {:?}\n{}", new.2, why()); *breaks.entry(b).or_default() += 1; } None => match mlic_direct(&m, &p) { Some(old) => { assert!(new == old, "the two decodes differ: {:?} / {:?}\n{}", new.2.as_ref().err(), old.2.as_ref().err(), why()); match &new.2 { Ok(_) => alike += 1, Err(e) => { let (code, msg) = e.split_once(": ").unwrap_or((e, "")); let kind = format!("{code}: {}", msg.split_whitespace().take(2).collect::>().join(" ")); *refused.entry(kind).or_default() += 1; } } } None => panic!("the direct decode panicked\n{}", why()), }, } }); eprintln!("int_mlicpp: {alike} decoded alike, {unparsed} bundles refused by both before decoding, refused by both {refused:?}, planned breaks {breaks:?}"); // Three networks, two streams and a bundle compound the generator's refusal rates. assert!(alike * 10 > n, "the generator must produce decodable cases: {alike} of {n}"); let reached = |k: &str| refused.keys().any(|x| x.contains(k)); assert!(n < 1000 || ["h_s", "entropy", "g_s", "stream", "model_mismatch", "payload_malformed"].iter().all(|k| reached(k)), "a refusal was not reached: {refused:?}"); assert!(n < 1000 || breaks.len() >= 3, "too few kinds of planned break reached: {breaks:?}"); } /// The planned breaks of integer-operators §15.3 for `wai.neural.int_mlicpp`, each on a /// bundle and payload the direct decode decodes, or refuses later or with another code: /// the classifier names it, and the decode through the operator set refuses it /// (`model_mismatch`, before either stream). #[test] fn int_mlicpp_planned_breaks_are_new_refusals() { use wai::int_mlicpp::{legacy, ops, IntMlicModel}; // Cy = 1, K = 1, a latent of side 2; Cz = 2, a hyper-latent of side 2; an image of // side 2. h_s: the hyper-latent's two channels as they are; entropy: σ and μ from them; // g_s: the latent to each of three channels. Every z symbol is the first of its table // (offset 2), so every hyper-latent is 2·2^12; every y symbol the first of the one // scale edge's table (offset 2), so every latent is 2·2^12 + 2·2^12, clamped to 2^12: // 255. let t3 = vec![0, 1 << 15, 1 << 16]; let tables = |z_short: bool| { let nz = if z_short { 1 } else { 2 }; json!({ "afrac": 12, "scale_edges": [1 << 12], "cdfs": [t3], "cdf_lengths": [3], "offsets": [2], "z_cdfs": vec![t3.clone(); nz], "z_cdf_lengths": vec![3; nz], "z_offsets": vec![2; nz], "z_median_q": vec![0; nz], }) }; let zs: Vec<(i32, &[i32], usize, i32)> = (0..8).map(|_| (2, t3.as_slice(), 3, 2)).collect(); let ys: Vec<(i32, &[i32], usize, i32)> = (0..4).map(|_| (2, t3.as_slice(), 3, 2)).collect(); let (z, y) = (rans_of(&zs).unwrap(), rans_of(&ys).unwrap()); let payload = |z: &[u8], y: &[u8]| wim1_bytes(2, 2, [2, 2, 2, 1, 2, 2, 1], z, y); let ok = payload(&z, &y); let cut = |s: &[u8]| s[..s.len() - 4].to_vec(); let ident = Whm::Conv { inc: 2, out: 2, k: 1, s: 1, p: 0, w: vec![1, 0, 0, 1], m: vec![1; 2], b: vec![0; 2] }; let pick = Whm::Conv { inc: 3, out: 2, k: 1, s: 1, p: 0, w: vec![0, 1, 0, 0, 0, 1], m: vec![1; 2], b: vec![0; 2] }; let to3 = |w: i16, b: i64| Whm::ConvT { inc: 1, out: 3, k: 1, s: 1, p: 0, op: 0, w: vec![w; 3], m: vec![1; 3], b: vec![b; 3] }; let net = |shift: u32, gf: u32, l: Vec| whm1_bytes(12, gf, shift, &l); let (h0, e0, g0) = (net(0, 16, vec![ident.clone()]), net(0, 16, vec![pick.clone()]), net(0, 16, vec![to3(1, 0)])); let b = |h: &[u8], e: &[u8], g: &[u8]| wmb1_bytes(h, e, g, &tables(false)); let decode = |bundle: &[u8], p: &[u8]| { let m = IntMlicModel::from_bundle(bundle).unwrap(); let new = ops::decode(&m, p).map(|i| i.rgb).map_err(|e| mlic_label(&e)); let old = catch_unwind(AssertUnwindSafe(|| legacy::decode(&m, p).map(|i| i.rgb).map_err(|e| mlic_label(&e)))).ok(); (new, old) }; // The base decodes alike. assert_eq!(int_mlicpp_planned_break(&b(&h0, &e0, &g0), &ok), None); let (new, old) = decode(&b(&h0, &e0, &g0), &ok); assert_eq!((new, old), (Ok(vec![255; 12]), Some(Ok(vec![255; 12])))); let expect = |kind: &str, bundle: Vec, p: &[u8], direct: Option<&str>| { assert_eq!(int_mlicpp_planned_break(&bundle, p), Some(kind), "{kind}: the classifier"); let (new, old) = decode(&bundle, p); match direct { None => assert!(matches!(old, Some(Ok(_))), "{kind}: the direct decode decodes it: {old:?}"), Some(code) => assert!(matches!(&old, Some(Err(e)) if e.starts_with(code)), "{kind}: the direct decode's code: {old:?}"), } assert!(matches!(&new, Err(e) if e.starts_with("model_mismatch")), "{kind}: the decode through the operator set refuses it before the streams: {new:?}"); }; // The breaks of the other WHM1 decoders, in each of the three networks. expect("no layer of tag 0, 1 or 2", b(&net(0, 16, vec![Whm::Relu]), &e0, &g0), &ok, None); let igdn = Whm::Igdn { c: 2, g: vec![0; 4], b: vec![1 << 24; 2] }; expect("a shift or gf of 128 or more that no layer uses", b(&net(200, 16, vec![igdn]), &e0, &g0), &ok, None); expect("a shift or gf of 128 or more that no layer uses", b(&h0, &e0, &net(0, 200, vec![to3(1, 0)])), &ok, None); expect("a bias of magnitude 2^63 - 1 or more", b(&h0, &e0, &net(0, 16, vec![to3(1, i64::MAX)])), &ok, Some("payload_malformed: g_s")); let first = Whm::ConvT { inc: 1, out: 1, k: 1, s: 1, p: 0, op: 0, w: vec![0], m: vec![1], b: vec![1 << 30] }; let tight = Whm::ConvT { inc: 1, out: 3, k: 1, s: 1, p: 0, op: 0, w: vec![1 << 11; 3], m: vec![1 << 40; 3], b: vec![0; 3] }; expect("a check range its producer cannot meet", b(&h0, &e0, &net(0, 16, vec![first, tight])), &ok, Some("payload_malformed: g_s")); // The rules the direct decode makes only when a network runs: the same code where the // streams are whole, and the stream's where the streams run out first. let chain = |inc: u32| vec![to3(1, 0), Whm::ConvT { inc, out: 3, k: 1, s: 1, p: 0, op: 0, w: vec![1; inc as usize * 3], m: vec![1; 3], b: vec![0; 3] }]; expect("a layer that does not fit the channels before it", b(&h0, &e0, &net(0, 16, chain(2))), &ok, Some("model_mismatch: g_s")); expect("a layer that does not fit the channels before it", b(&h0, &e0, &net(0, 16, chain(2))), &payload(&z, &cut(&y)), Some("payload_malformed: y stream")); expect("a shift or gf of 128 or more that a layer uses", b(&h0, &e0, &net(200, 16, vec![to3(1, 0)])), &ok, Some("model_mismatch: g_s")); expect("a shift or gf of 128 or more that a layer uses", b(&h0, &e0, &net(200, 16, vec![to3(1, 0)])), &payload(&z, &cut(&y)), Some("payload_malformed: y stream")); let leaky64 = net(0, 16, vec![pick.clone(), Whm::Leaky { q: 1, f: 64 }]); expect("a leaky ReLU's shift of 64 or more", b(&h0, &leaky64, &g0), &ok, Some("model_mismatch: entropy")); expect("a leaky ReLU's shift of 64 or more", b(&h0, &leaky64, &g0), &payload(&cut(&z), &y), Some("payload_malformed: z stream")); let reads3 = net(0, 16, vec![Whm::Conv { inc: 3, out: 2, k: 1, s: 1, p: 0, w: vec![1; 6], m: vec![1; 2], b: vec![0; 2] }]); expect("a network whose input is not the payload's channel count", b(&reads3, &e0, &g0), &ok, Some("model_mismatch: h_s")); expect("a network whose input is not the payload's channel count", b(&reads3, &e0, &g0), &payload(&cut(&z), &y), Some("payload_malformed: z stream")); // A convolution padded by 2^26 a side and strided by 2^27: the direct decode's geometry // gives the output, of side 2, and decodes; the operator set's run-time plan caps the // padded side. let far = Whm::Conv { inc: 1, out: 1, k: 1, s: 1 << 27, p: 1 << 26, w: vec![1], m: vec![1], b: vec![0] }; expect("a convolution's padded side past 2^26", b(&h0, &e0, &net(0, 16, vec![far, to3(1, 0)])), &ok, None); // Not breaks: a bundle with too few z tables, and a network the shapes refuse, are // refused alike before either decode reads a stream. let short = wmb1_bytes(&h0, &e0, &g0, &tables(true)); assert_eq!(int_mlicpp_planned_break(&short, &ok), None); let (new, old) = decode(&short, &ok); assert!(matches!((&new, &old), (Err(a), Some(Err(b))) if a == b && a.starts_with("model_mismatch")), "{new:?} {old:?}"); } // ------------------------------------------------------------------ the world renderer (M6) /// A random set of `wai.world.model`'s integer renderer, and a clip for it. #[derive(Clone, Debug)] struct WorldCase { t: Vec, g_s: Vec, c: usize, lat: usize, z0: Vec, rows: Vec>, } /// The input channel count of a network's first layer of tag 0, 1 or 2. fn in_channels(n: &wai_ops::lower::Net) -> Option { use wai_ops::lower::NetLayer as N; n.layers.iter().find_map(|l| match l { N::ConvT { inc, .. } | N::Conv { inc, .. } => Some(*inc), N::Igdn { c, .. } => Some(*c), _ => None, }) } /// The planned break of integer-operators §15.3 item 2 a world case falls under, if any, /// read from the set and the clip alone. The direct renderer checked nothing before it ran, /// so each check the renderer makes before its first frame is new: the set refused at load /// (`world/1`, loaded under `acc64`) for the `afrac` rule, for `A < 1`, for one of a /// `WHM1` network's rules (§13.2, named as [`load_break`] names them) or for layers that do /// not chain; a row that does not carry `A` values; and `T` that does not give the latent's /// shape back or `g_s` that does not plan on it. A value outside a `check` range is not a /// break: the direct guard refuses it alike. fn world_planned_break(cs: &WorldCase) -> Option<&'static str> { use wai_ops::plan::run_plan; let (tn, gn) = (wai_ops::lower::parse_whm1(&cs.t).ok()?, wai_ops::lower::parse_whm1(&cs.g_s).ok()?); let loaded = wai_ops::lower::world(&cs.t, &cs.g_s, cs.c as u32).and_then(|l| { let (t, g) = (wai::int_ops::load(l.get("T").unwrap())?, wai::int_ops::load(l.get("g_s").unwrap())?); Ok((t, g)) }); let (t, g) = match loaded { Ok(tg) => tg, Err(_) => { if [tn.afrac, gn.afrac].iter().any(|a| !(1..=30).contains(a)) || tn.afrac != gn.afrac { return Some("afrac outside 1..=30, or unequal"); } if let Some(b) = load_break(&tn, wai_ops::lower::whm1(&cs.t, false)).or_else(|| load_break(&gn, wai_ops::lower::whm1(&cs.g_s, true))) { return Some(b); } if in_channels(&tn).is_some_and(|cin| cin <= cs.c as u32) { return Some("A < 1"); } return Some("layers that do not chain, or a shift a layer cannot take"); } }; let a = t.graph().scalars.len(); if cs.rows.iter().any(|r| r.len() != a) { return Some("a row that does not carry A values"); } let shape = [(cs.c as u64, cs.lat as u64, cs.lat as u64)]; let fits = run_plan(t.graph(), &shape, &vec![0; a]).is_ok_and(|hw| { let o = t.graph().outputs[0] as usize; (t.graph().tensors[o].channels as u64, hw[o].0, hw[o].1) == shape[0] }); if !fits || run_plan(g.graph(), &shape, &[]).is_err() { return Some("T or g_s that does not fit the latent"); } None } /// [`same_size_layer`], its multipliers (when `unit` is set) about `2^shift` over the /// layer's terms and largest weight, so that its output keeps about its input's /// magnitude, of either sign. A random transition of such layers leaves the latent mostly /// inside `T`'s clamp, where a difference in `T`'s arithmetic shows. fn world_layer_conv(r: &mut Rng, c: u32, out: u32, shift: u32, unit: bool) -> Whm { let mut l = same_size_layer(r, c, out); if unit && let Whm::Conv { k, w, m, .. } = &mut l { let wmax = w.iter().map(|v| v.unsigned_abs() as i64).max().unwrap_or(1).max(1); let terms = (c * *k * *k) as i64; let g = ((1i64 << shift.min(62)) / (terms * wmax)).max(1); for v in m.iter_mut() { *v = if r.below(2) == 0 { g } else { -g }; } } l } /// A random world set and clip. `T`: `afrac` 12 mostly (one in sixteen anywhere in /// 1..=30, one in 32 at 0 or 31, and `g_s`'s unequal one in 32), a header `shift` of 8 to /// 23 (128 or more one in 64) and `gf` 16 (128 or more one in 64); `C` of 1 to 3 and `A` /// of 1 to 4 (0 one in 32); a first layer from `C + A` channels ([`world_layer_conv`], /// near unit gain three cases in four), up to one leaky ReLU (its slope below 1 three /// cases in four), ReLU or inverse GDN, and a last layer back to `C` (one in sixteen to /// `C + 1`, and one in sixteen strided by 2). `g_s`: 1 to 3 of [`splat_layers`]' layers /// from `C` channels (none one in 32, ReLUs only one in 32). The clip: a latent of side 1 /// to 4; `z_0` within `±2^(afrac+1)` mostly, one case in four with a value at `g_s`'s /// input range or one past it and one in sixteen at an end of `i64`; 0 to 4 rows of moves /// `−1, 0, 1` at `2^afrac`, one row in eight with a value at `T`'s `x_0` range or one past /// it, one in 32 with a value up to `2^40`, and one in sixteen without `A` values. fn world_case(r: &mut Rng) -> WorldCase { let afrac = match r.below(32) { 0 | 1 => 1 + r.below(30) as u32, 2 => [0, 31][r.below(2) as usize], _ => 12, }; let g_afrac = if r.below(32) == 0 { afrac ^ 1 } else { afrac }; let shift = if r.below(64) == 0 { 128 + r.below(64) as u32 } else { 8 + r.below(16) as u32 }; let gf = if r.below(64) == 0 { 128 + r.below(64) as u32 } else { 16 }; let c = 1 + r.below(3) as u32; let a = if r.below(32) == 0 { 0 } else { 1 + r.below(4) as u32 }; let h1 = 1 + r.below(4) as u32; let unit = r.below(4) != 0; let mut tl = vec![world_layer_conv(r, c + a, h1, shift, unit)]; match r.below(4) { 0 => { let f = r.below(16) as u32; tl.push(Whm::Leaky { q: if unit { r.below(1 << f) as i64 } else { r.signed(1 << 10) }, f }); } 1 => tl.push(Whm::Relu), 2 => tl.push(Whm::Igdn { c: h1, g: (0..h1 * h1).map(|_| r.below(1 << 12) as i32).collect(), b: (0..h1).map(|_| (1 << 24) + r.below(1 << 20) as i64).collect(), }), _ => {} } let out = if r.below(16) == 0 { c + 1 } else { c }; let mut last = world_layer_conv(r, h1, out, shift, unit); if r.below(16) == 0 && let Whm::Conv { s, .. } = &mut last { *s = 2; } tl.push(last); let t = whm1_bytes(afrac, gf, shift, &tl); let n = if r.below(32) == 0 { 0 } else { 1 + r.below(3) as usize }; let only_relu = r.below(32) == 0; let (gl, _) = splat_layers(r, c, n, only_relu); let g_s = whm1_bytes(g_afrac, 16, 8 + r.below(16) as u32, &gl); let (c, lat) = (c as usize, 1 + r.below(4) as usize); let one = 1i64 << afrac.min(30); let ranges = wai_ops::lower::world(&t, &g_s, c as u32).ok().map(|l| { let (tg, gg) = (l.get("T").unwrap(), l.get("g_s").unwrap()); (tg.tensors[1 + tg.scalars.len()].hi, gg.tensors[0].hi) }); let edge = |r: &mut Rng, x: i64| [x, x.saturating_add(1), -x, -x.saturating_add(1)][r.below(4) as usize]; let mut z0: Vec = (0..c * lat * lat).map(|_| r.signed(2 * one)).collect(); let k = r.below(z0.len() as u64) as usize; match (r.below(16), ranges) { (0..=3, Some((_, xg))) => z0[k] = edge(r, xg), (4, _) => z0[k] = [i64::MIN, i64::MAX][r.below(2) as usize], _ => {} } let mut rows = Vec::new(); for _ in 0..r.below(5) { let n = match r.below(16) { 0 if a > 0 => a as usize - 1, 1 => a as usize + 1, _ => a as usize, }; let mut row: Vec = (0..n).map(|_| r.signed(1) * one).collect(); if !row.is_empty() { let p = r.below(row.len() as u64) as usize; match (r.below(32), ranges) { (0..=3, Some((xt, _))) => row[p] = edge(r, xt), (4, _) => row[p] = r.signed(1 << 40), _ => {} } } rows.push(row); } WorldCase { t, g_s, c, lat, z0, rows } } /// Where a world graph refused, as the layer whose input check it is: for `T` (`A` fills, /// the concatenation, then the layers), the output of operator `A + ℓ`; for `g_s`, its /// input for layer 0 and the output of operator `ℓ − 1` for layer `ℓ`. `None` for a /// refusal that is no layer's check. fn world_layer(net: &str, a: usize, at: wai_ops::Site) -> Option { match (net, at) { ("T", wai_ops::Site::Op(i)) if i as usize >= a => Some(i as usize - a), ("g_s", wai_ops::Site::Input(0)) => Some(0), ("g_s", wai_ops::Site::Op(i)) => Some(i as usize + 1), _ => None, } } /// The renderer through the operator set, as `World::decode_planes` runs it after its /// checks, one graph at a time, so that a refusal names its network, step and layer. fn world_traced(w: &wai::int_world::ops::World, cs: &WorldCase) -> Result>, (&'static str, usize, Option, wai_ops::Code)> { use wai_ops::Tensor; let a = w.planes(); let g = |z: &Tensor, step: usize| { let out = wai::int_ops::run1(w.g_s(), z.clone(), &[]).map_err(|r| ("g_s", step, world_layer("g_s", a, r.at), r.code))?; let plane = (out.h * out.w) as usize; let mut hwc = vec![0u8; out.data.len()]; for (k, &v) in out.data.iter().enumerate() { hwc[(k % plane) * out.c as usize + k / plane] = v as u8; } Ok(hwc) }; let mut z = Tensor::new(cs.c as u32, cs.lat as u32, cs.lat as u32, cs.z0.clone()); let mut frames = vec![g(&z, 0)?]; for (k, row) in cs.rows.iter().enumerate() { z = wai::int_ops::run1(w.t(), z, row).map_err(|r| ("T", k + 1, world_layer("T", a, r.at), r.code))?; frames.push(g(&z, k + 1)?); } Ok(frames) } /// Gate 4 of the migration, for `wai.world.model`'s integer renderer: the renderer through /// the operator set against the direct renderer with the direct guard /// (`int_world::legacy::checked`) on random sets and clips. Outside the planned breaks /// both render the same frames, which are also the unguarded direct renderer's, through /// the same latents after every transition, or both /// refuse a value (`field_range`), in the same network, at the same step and layer, and the /// guard never panics; under a planned break the renderer through the operator set refuses /// before its first frame. #[test] fn int_world_random_sets_and_clips_render_alike() { use std::collections::BTreeMap; use wai::int_world::{legacy, ops}; let (mut alike, mut refused, mut breaks) = (0u64, BTreeMap::::new(), BTreeMap::<&str, u64>::new()); let n = cases("int_world", 3000, |r, seed| { let cs = world_case(r); let why = || format!("seed {seed}\ncase {cs:?}"); let (t, g) = (IntModel::from_bytes(&cs.t).expect("the generator writes WHM1 that parses"), IntModel::from_bytes(&cs.g_s).unwrap()); let new = ops::decode_planes(&t, &g, &cs.z0, &cs.rows, cs.c, cs.lat).map(|c| c.frames); if let Some(b) = world_planned_break(&cs) { assert!(new.is_err(), "{b}: the renderer through the operator set must refuse it\n{}", why()); *breaks.entry(b).or_default() += 1; return; } let world = ops::load(&t, &g, cs.c).unwrap_or_else(|e| panic!("no planned break, yet the set is refused: {e}\n{}", why())); let traced = world_traced(&world, &cs); let direct = catch_unwind(AssertUnwindSafe(|| legacy::checked::decode_planes(&t, &g, &cs.z0, &cs.rows, cs.c, cs.lat).map(|c| c.frames))) .unwrap_or_else(|_| panic!("the direct guard panicked\n{}", why())); match (&new, &traced, &direct) { (Ok(a), Ok(b), Ok(d)) => { assert!(a == d && b == d, "the two renderers' frames differ\n{}", why()); let plain = legacy::decode_planes(&t, &g, &cs.z0, &cs.rows, cs.c, cs.lat).frames; assert!(plain == *d, "the guard changed the direct renderer's frames\n{}", why()); // The latents too: a random g_s mostly saturates the RGB8 step, which can // hide a difference in T. let (mut zo, mut zd) = (cs.z0.clone(), cs.z0.clone()); for (k, row) in cs.rows.iter().enumerate() { zo = world.step(&zo, row, cs.lat).unwrap_or_else(|e| panic!("transition {}: {e}\n{}", k + 1, why())); zd = legacy::step_planes(&t, &zd, row, cs.c, cs.lat); assert!(zo == zd, "the latents after transition {} differ\n{}", k + 1, why()); } alike += 1; } (Err(a), Err((net, step, layer, code)), Err(d)) => { assert_eq!((a.kind, *code), (RefusalKind::FieldRange, wai_ops::Code::FieldRange), "the renderer refused other than a value\n{}", why()); assert!(d.out_of_range, "the guard refused other than a value: {d:?}\n{}", why()); assert_eq!((*net, *step, *layer), (d.net, d.step, Some(d.layer)), "the refusals' network, step and layer\n{}", why()); *refused.entry(format!("{net} layer {}", d.layer)).or_default() += 1; } _ => panic!("through the operator set {:?}, traced {:?}, directly {:?}\n{}", new.as_ref().err(), traced.as_ref().err(), direct.as_ref().err(), why()), } }); eprintln!("int_world: {alike} rendered alike, refused by both {refused:?}, planned breaks {breaks:?}"); assert!(alike * 8 > n, "the generator must produce renderable cases: {alike} of {n}"); assert!(n < 1000 || (refused.keys().any(|k| k.starts_with("T")) && refused.keys().any(|k| k.starts_with("g_s"))), "a network's check range was not reached: {refused:?}"); assert!(n < 1000 || breaks.len() >= 5, "too few kinds of planned break reached: {breaks:?}"); } /// The planned breaks of integer-operators §15.3 item 2 for the world renderer, each on a /// set and clip the direct renderer rendered from (or panicked on, or whose guard refused a /// value): the classifier names it, and the renderer through the operator set refuses it /// before its first frame (`model_mismatch`). #[test] fn int_world_planned_breaks_are_new_refusals() { use wai::int_world::{legacy, ops}; // C = 1, A = 1, a latent of side 2: T passes z through (weights 1 for z, 0 for the plane), // g_s gives three channels of the latent, the RGB8 step then maps 2^12 to 255. let pass = |b: i64| Whm::Conv { inc: 2, out: 1, k: 1, s: 1, p: 0, w: vec![1, 0], m: vec![1], b: vec![b] }; let to3 = |w: i16, b: i64| Whm::ConvT { inc: 1, out: 3, k: 1, s: 1, p: 0, op: 0, w: vec![w; 3], m: vec![1; 3], b: vec![b; 3] }; let net = |afrac: u32, shift: u32, gf: u32, l: Vec| whm1_bytes(afrac, gf, shift, &l); let case = |t: Vec, g_s: Vec, c: usize, rows: Vec>| WorldCase { t, g_s, c, lat: 2, z0: vec![1 << 12; c * 4], rows }; let base = case(net(12, 0, 16, vec![pass(0)]), net(12, 0, 16, vec![to3(1, 0)]), 1, vec![vec![0], vec![1 << 12]]); let render = |cs: &WorldCase| { let (t, g) = (IntModel::from_bytes(&cs.t).unwrap(), IntModel::from_bytes(&cs.g_s).unwrap()); let new = ops::decode_planes(&t, &g, &cs.z0, &cs.rows, cs.c, cs.lat).map(|c| c.frames).map_err(|e| e.kind); let old = catch_unwind(AssertUnwindSafe(|| legacy::decode_planes(&t, &g, &cs.z0, &cs.rows, cs.c, cs.lat).frames)).ok(); let guard = catch_unwind(AssertUnwindSafe(|| legacy::checked::decode_planes(&t, &g, &cs.z0, &cs.rows, cs.c, cs.lat).map(|c| c.frames.len()))).ok(); (new, old, guard) }; assert_eq!(world_planned_break(&base), None); let (new, old, _) = render(&base); assert_eq!(new, Ok(vec![vec![255; 12]; 3]), "the base renders through the operator set"); assert_eq!(old, Some(vec![vec![255; 12]; 3]), "the base renders directly"); let expect = |kind: &str, cs: WorldCase| { assert_eq!(world_planned_break(&cs), Some(kind), "{kind}: the classifier"); let (new, old, guard) = render(&cs); assert_eq!(new.map(|_| ()), Err(RefusalKind::ModelMismatch), "{kind}: the renderer through the operator set refuses it"); let d = match (old, guard) { (Some(f), Some(Ok(_))) => format!("rendered {} frames, and so did its guard", f.len()), (Some(f), Some(Err(g))) => format!("rendered {} frames; its guard refused at {} layer {} (step {})", f.len(), g.net, g.layer, g.step), (Some(f), None) => format!("rendered {} frames; its guard panicked", f.len()), (None, _) => "panicked".to_string(), }; println!("{kind}: refused model_mismatch; the direct renderer {d}"); }; expect("afrac outside 1..=30, or unequal", case(net(13, 0, 16, vec![pass(0)]), net(12, 0, 16, vec![to3(1, 0)]), 1, vec![vec![0]])); expect("afrac outside 1..=30, or unequal", case(net(31, 0, 16, vec![pass(0)]), net(31, 0, 16, vec![to3(1, 0)]), 1, vec![vec![0]])); expect("A < 1", case(net(12, 0, 16, vec![pass(0)]), net(12, 0, 16, vec![to3(1, 0)]), 2, vec![])); expect("no layer of tag 0, 1 or 2", case(net(12, 0, 16, vec![pass(0)]), net(12, 0, 16, vec![Whm::Relu]), 1, vec![vec![0]])); expect("a shift or gf of 128 or more that no layer uses", case(net(12, 0, 16, vec![pass(0)]), net(12, 0, 200, vec![to3(1, 0)]), 1, vec![vec![0]])); expect("a bias of magnitude 2^63 - 1 or more", case(net(12, 0, 16, vec![pass(i64::MAX)]), net(12, 0, 16, vec![to3(1, 0)]), 1, vec![vec![0]])); let first = Whm::Conv { inc: 2, out: 1, k: 1, s: 1, p: 0, w: vec![0, 0], m: vec![1], b: vec![1 << 30] }; let tight = Whm::Conv { inc: 1, out: 1, k: 1, s: 1, p: 0, w: vec![1 << 11], m: vec![1 << 40], b: vec![0] }; expect("a check range its producer cannot meet", case(net(12, 0, 16, vec![first, tight]), net(12, 0, 16, vec![to3(1, 0)]), 1, vec![vec![0]])); expect("layers that do not chain, or a shift a layer cannot take", case(net(12, 200, 16, vec![pass(0)]), net(12, 0, 16, vec![to3(1, 0)]), 1, vec![vec![0]])); expect("a row that does not carry A values", case(net(12, 0, 16, vec![pass(0)]), net(12, 0, 16, vec![to3(1, 0)]), 1, vec![vec![0, 0]])); let halve = Whm::Conv { inc: 2, out: 1, k: 1, s: 2, p: 0, w: vec![1, 0], m: vec![1], b: vec![0] }; expect("T or g_s that does not fit the latent", case(net(12, 0, 16, vec![halve]), net(12, 0, 16, vec![to3(1, 0)]), 1, vec![vec![0]])); let wide = Whm::Conv { inc: 2, out: 2, k: 1, s: 1, p: 0, w: vec![1, 0, 1, 0], m: vec![1; 2], b: vec![0; 2] }; expect("T or g_s that does not fit the latent", case(net(12, 0, 16, vec![wide]), net(12, 0, 16, vec![to3(1, 0)]), 1, vec![vec![0]])); let big = Whm::Conv { inc: 1, out: 1, k: 3, s: 1, p: 0, w: vec![1; 9], m: vec![1], b: vec![0] }; expect("T or g_s that does not fit the latent", case(net(12, 0, 16, vec![pass(0)]), net(12, 0, 16, vec![big, to3(1, 0)]), 1, vec![vec![0]])); } // ------------------------------------------------------------------ load step 8 /// What the bounds pass made of the graphs compared: refused, or the profiles met. #[derive(Default, Debug)] struct Verdicts { graphs: u64, refused: u64, acc32: u64, acc64_only: u64, neither: u64, } /// Load step 8's pass on `g` against its oracle: equal passes, node for node (every /// intermediate, result, saturation mark, parameter condition, `K_c`, width and flag, /// and every tensor's interval), or the same refusal; and the same first failure under /// each profile. fn same_pass(g: &wai_ops::Graph, v: &mut Verdicts, why: &dyn Fn() -> String) { use wai_ops::bounds::{self, Profile}; let sizes = wai_ops::plan::load_plan(g).unwrap_or_else(|r| panic!("a lowered graph plans: {r:?}\n{}", why())); let (new, old) = (bounds::pass(g, sizes.as_deref()), wai_ops::bounds_oracle::pass(g, sizes.as_deref())); assert!(new == old, "the bounds pass and its oracle differ\n{}", why()); v.graphs += 1; match (&new, &old) { (Ok(a), Ok(b)) => { for p in Profile::ALL { assert_eq!(a.first_failure(g, p), b.first_failure(g, p), "{}", why()); } match (a.meets(g, Profile::Acc32), a.meets(g, Profile::Acc64)) { (true, _) => v.acc32 += 1, (false, true) => v.acc64_only += 1, _ => v.neither += 1, } } _ => v.refused += 1, } } /// The bounds pass against its oracle on every graph the generators above lower: /// `int_synth/1`'s graph, the splat set's `g_s`, the audio synthesis, the conditional /// codec's `h_s` and `g_s`, the MLIC++ bundle's three networks, and the world /// renderer's `T` and `g_s`. Every verdict is reached: a refusal by the pass /// itself (a `check` range its producer cannot meet), `acc32`, `acc64` only, and neither. #[test] fn the_bounds_pass_equals_its_oracle_on_random_sets() { let mut v = Verdicts::default(); let n = cases("bounds pass", 3000, |r, seed| { let why = || format!("seed {seed}"); let (set, _) = int_synth_case(r); if let Ok(l) = wai_ops::lower::int_synth(&set) { for (_, g) in &l.roles { same_pass(g, &mut v, &why); } } let (set, _) = int_splat_case(r); if let Ok(g) = wai_ops::lower::whm1(&set, true) { same_pass(&g, &mut v, &why); } let (set, _) = int_audio_case(r); if let Ok(g) = wai_ops::lower::wam1(&set) { same_pass(&g, &mut v, &why); } let (g_s, h_s, _) = int_hyper_case(r); for (set, rgb8) in [(&h_s, false), (&g_s, true)] { if let Ok(g) = wai_ops::lower::whm1(set, rgb8) { same_pass(&g, &mut v, &why); } } let (bundle, _) = int_mlicpp_case(r); if let Some(secs) = wmb1_sections(&bundle) { for (k, rgb8) in [(0, false), (1, false), (2, true)] { if let Ok(g) = wai_ops::lower::whm1(&secs[k], rgb8) { same_pass(&g, &mut v, &why); } } } let cs = world_case(r); if let Ok(l) = wai_ops::lower::world(&cs.t, &cs.g_s, cs.c as u32) { for (_, g) in &l.roles { same_pass(g, &mut v, &why); } } }); eprintln!("bounds pass: {n} cases, {v:?}"); assert!(v.graphs > n * 2, "the generators must mostly lower: {} graphs from {n} cases", v.graphs); assert!(n < 1000 || (v.refused > 0 && v.acc32 > 0 && v.acc64_only > 0 && v.neither > 0), "a verdict was not reached: {v:?}"); } // ------------------------------------------------------------- jpegai (M7a) // // `wai.image.jpegai`'s scale path and entropy decodes (integer-operators §13.6.1, lowered // by `jpegai-int/1`) against the direct computation (`jpegai_ops::legacy`): random bundles // (the members §13.6.1 reads) and random entropy payloads. The direct computation refuses // nothing: where the operator set refuses a stream (`stream_exhausted`) the direct // decode panics (overflow checks on) or computes on from a window that wrapped (off). /// A random bundle's members as JSON, and the three entropy payloads with each plane's /// GRFS flags. struct JCase { norm: Value, hs: [Value; 2], soz: Vec, lat: [Vec; 2], grfs: [Vec; 2], } impl std::fmt::Debug for JCase { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { write!(f, "pmf rows {}, payloads {} / {} / {}", self.norm["pmf"].as_array().map_or(0, Vec::len), self.soz.len(), self.lat[0].len(), self.lat[1].len()) } } /// Random bytes, or one of the committed test codestream's payloads, damaged or not. fn jpegai_payload(r: &mut Rng, committed: &[u8]) -> Vec { let mut v: Vec = match r.below(4) { 0 => committed.to_vec(), 1 => (0..r.below(16)).map(|_| r.below(256) as u8).collect(), _ => (0..[64, 512, 4096, 9000][r.below(4) as usize]).map(|_| r.below(256) as u8).collect(), }; match r.below(6) { 0 if !v.is_empty() => { let n = r.below(v.len() as u64) as usize; v.truncate(n) } 1 => { let mut pre: Vec = (0..[8u64, 256, 4096][r.below(3) as usize]).map(|_| r.below(256) as u8).collect(); pre.extend_from_slice(&v); v = pre } 2 if !v.is_empty() => { let i = r.below(v.len() as u64) as usize; v[i] ^= 1 << r.below(8) } _ => {} } v } /// One plane's scale network: `c1`, `c2` channels after blocks 1 and 2, `cz` in. fn jpegai_scale(r: &mut Rng, cz: u64, c1: u64, c2: u64) -> Value { let big = r.below(16) == 0; let w = |n: u64, r: &mut Rng| -> Vec { (0..n).map(|_| if big && r.below(64) == 0 { r.signed(1 << 20) } else { r.signed(3) }).collect() }; // Shifts up to 8, then, in one network in about eight, one shift of 31 to 40 in each // block; rarer, one of 64 to 127 (past acc32's 63) or of 128 (a load check). let mode = r.below(64); let shifts = |n: u64, r: &mut Rng| -> Vec { let mut v: Vec = (0..n).map(|_| r.below(9) as i64).collect(); let i = r.below(n) as usize; match mode { 0 => v[i] = [64, 100, 127][r.below(3) as usize], 1 => v[i] = 128, 2..=9 => v[i] = 31 + r.below(10) as i64, _ => {} } v }; let bias = |n: u64, r: &mut Rng| -> Vec { (0..n).map(|_| r.signed(500)).collect() }; let c3 = 16 * cz; let mut v = json!({ "conv1.weight": { "dims": [c1, cz, 1, 1], "v": w(c1 * cz, r) }, "conv1.bias": { "dims": [c1], "v": bias(c1, r) }, "onnx::Pow_85": { "dims": [1, c1, 1, 1], "v": shifts(c1, r) }, "depthwise.weight": { "dims": [c2, c1, 3, 3], "v": w(c2 * c1 * 9, r) }, "depthwise.bias": { "dims": [c2], "v": bias(c2, r) }, "onnx::Pow_91": { "dims": [1, c2, 1, 1], "v": shifts(c2, r) }, "pointwise.weight": { "dims": [c3, c2, 1, 1], "v": w(c3 * c2, r) }, "pointwise.bias": { "dims": [c3], "v": bias(c3, r) }, "onnx::Pow_97": { "dims": [1, c3, 1, 1], "v": shifts(c3, r) }, }); if r.below(64) == 0 { v["depthwise.weight"]["dims"] = json!([c2, c1, 1, 1]); } v } /// A random bundle and payloads. Most bundles pass every load check; about one in five /// breaks one, or meets acc64 and not acc32, or has a shift from 64 to 127. fn jpegai_case(r: &mut Rng, payloads: &[Vec; 3]) -> JCase { let rows = if r.below(16) == 0 { 33 + r.below(8) } else { 32 }; let w = [1u64, 2, 3, 4, 8, 16, 64][r.below(7) as usize]; let mut pmf: Vec> = (0..rows).map(|_| { let mut row = vec![0i64; w as usize]; for _ in 0..256 { row[r.below(w) as usize] += 1; } row }).collect(); let mut bounds: Vec = (0..32).map(|_| [0i64, 1, 2, 4, 30, 127, 128, 200, 255][r.below(9) as usize]).collect(); let cdf_row = |r: &mut Rng| -> Vec { let mut c = 0i64; (0..63).map(|_| { c = (c + r.below(9) as i64).min(255); c }).collect() }; let mut cdfs_y: Vec = (0..160).flat_map(|_| cdf_row(r)).collect(); let mut cdfs_uv: Vec = (0..96).flat_map(|_| cdf_row(r)).collect(); let g0 = |r: &mut Rng| [-5000i64, -1000, 0, 300, 380, 400, 1000, 3000][r.below(8) as usize]; let (gy0, guv0) = (g0(r), g0(r)); let mut gain_y: Vec = (0..160).map(|_| gy0 + r.signed(100)).collect(); let gain_uv: Vec = (0..96).map(|_| guv0 + r.signed(100)).collect(); let mut per_channel = 16; match r.below(64) { 0 => { pmf[0][0] += 1 } 1 => { pmf.truncate(31) } 2 => { bounds[r.below(32) as usize] = [256, -1][r.below(2) as usize] } 3 => { let i = r.below(160 * 63 - 1) as usize; cdfs_y[i] = 256 } 4 => { let row = r.below(96) as usize * 63; cdfs_uv[row] = cdfs_uv[row + 1] + 1 } 5 => { gain_y.pop(); } 6 => { per_channel = 15 } _ => {} } let (c1, c2) = (1 + r.below(3), 1 + r.below(3)); let norm = json!({ "pmf": pmf, "bounds": bounds, "per_channel": per_channel, "cdfs_y": cdfs_y, "cdfs_uv": cdfs_uv, "gain_y": gain_y, "gain_uv": gain_uv }); let (u1, u2) = (1 + r.below(3), 1 + r.below(3)); let hs = [jpegai_scale(r, 160, c1, c2), jpegai_scale(r, 96, u1, u2)]; let soz = jpegai_payload(r, &payloads[0]); let lat = [jpegai_payload(r, &payloads[1]), jpegai_payload(r, &payloads[2])]; let grfs = [(0..160).map(|_| r.below(2) == 0).collect(), (0..96).map(|_| r.below(2) == 0).collect()]; JCase { norm, hs, soz, lat, grfs } } /// The planned break of integer-operators §15.3 item 4 a bundle falls under, read from /// its members alone: a load check of §13.6.1 it fails, or scale graphs that do not meet /// `acc32`, the capability's profile (a sum past `i32`, or a shift past 63, where the /// direct computation's result depended on the build). Each is a refusal. fn jpegai_planned_break(c: &JCase) -> Option<&'static str> { let n = &c.norm; let ints = |v: &Value| -> Vec { v.as_array().unwrap().iter().map(|x| x.as_i64().unwrap()).collect() }; let pmf: Vec> = n["pmf"].as_array().unwrap().iter().map(ints).collect(); if pmf.len() < 32 || pmf.iter().any(|r| r.len() != pmf[0].len() || r.iter().any(|&p| p < 0) || r.iter().sum::() != 256) { return Some("pmf"); } let bounds = ints(&n["bounds"]); if bounds.len() < 32 || bounds.iter().any(|&b| !(0..=255).contains(&b)) { return Some("bounds"); } if n["per_channel"] != 16 { return Some("per_channel"); } for (k, rows) in [("cdfs_y", 160usize), ("cdfs_uv", 96)] { let c = ints(&n[k]); if c.len() != rows * 63 || c.iter().any(|&v| !(0..=255).contains(&v)) || c.chunks(63).any(|r| r.windows(2).any(|p| p[0] > p[1])) { return Some("cdfs"); } } if ints(&n["gain_y"]).len() != 160 || ints(&n["gain_uv"]).len() != 96 { return Some("gains"); } for (hs, cz) in c.hs.iter().zip([160u64, 96]) { let d = |k: &str| -> Vec { hs[k]["dims"].as_array().unwrap().iter().map(|x| x.as_u64().unwrap()).collect() }; let (c1, c2) = (d("conv1.weight")[0], d("depthwise.weight")[0]); if d("depthwise.weight") != [c2, c1, 3, 3] || d("conv1.weight") != [c1, cz, 1, 1] { return Some("dims"); } let shifts: Vec = ["onnx::Pow_85", "onnx::Pow_91", "onnx::Pow_97"].iter().flat_map(|k| ints(&hs[*k]["v"])).collect(); if shifts.iter().any(|&s| !(0..=127).contains(&s)) { return Some("shift"); } } let b = |v: &Value| serde_json::to_vec(v).unwrap(); let g = wai_ops::lower::jpegai_scale(&b(n), &b(&c.hs[0]), &b(&c.hs[1])).expect("a bundle that passes the load checks lowers"); if g.roles.iter().any(|(_, g)| wai_ops::graph::load(&g.to_bytes(), wai_ops::Profile::Acc32).is_err()) { assert!(g.roles.iter().all(|(_, g)| wai_ops::graph::load(&g.to_bytes(), wai_ops::Profile::Acc64).is_ok()), "the generator's sums fit i64"); return Some("acc32"); } None } /// The tables as the decoders parse them. fn jpegai_parsed(c: &JCase) -> (wai::jpegai_decode::NormTables, [wai::jpegai_synthesis::IntNet; 2]) { let ints = |k: &str| -> Vec { c.norm[k].as_array().unwrap().iter().map(|x| x.as_i64().unwrap()).collect() }; let n = wai::jpegai_decode::NormTables { pmf: c.norm["pmf"].as_array().unwrap().iter().map(|r| r.as_array().unwrap().iter().map(|x| x.as_i64().unwrap()).collect()).collect(), bounds: ints("bounds"), per_channel: c.norm["per_channel"].as_u64().unwrap() as usize, cdfs_y: ints("cdfs_y"), cdfs_uv: ints("cdfs_uv"), gain_y: ints("gain_y"), gain_uv: ints("gain_uv"), lsbs_t0: vec![], lsbs_t1: vec![], lef_thr: [0; 3], lef_mag: [0.0; 3], }; (n, [wai::jpegai_synthesis::IntNet::from_json(&c.hs[0]).unwrap(), wai::jpegai_synthesis::IntNet::from_json(&c.hs[1]).unwrap()]) } /// The committed test codestream's SOZ, SORP and SORS payloads. fn jpegai_committed_payloads() -> [Vec; 3] { use wai::jpegai_bitstream::{find_segment, parse_container, Marker}; let data = std::fs::read(wai().join("wai-rs/tests/vectors/jpegai/test256.bin")).unwrap(); let segs = parse_container(&data); [Marker::Soz, Marker::Sorp, Marker::Sors].map(|m| { let s = find_segment(&segs, m).unwrap(); data[s.start..s.start + s.len].to_vec() }) } /// Random bundles and payloads: outside the planned breaks, both planes' integers through /// the operator set equal the direct computation's, or the operator set refuses /// (`stream_exhausted`) where the direct computation panics with overflow checks on (or, /// with them off, panics or computes on from a wrapped window, counted). Under a planned /// break the bundle is refused (`model_mismatch`), from its bytes and from its parsed /// form alike. #[test] fn jpegai_random_bundles_and_streams_decode_alike() { use wai::jpegai_ops::{self, Plane}; let payloads = jpegai_committed_payloads(); let (mut alike, mut refused, mut computed_on) = (0u64, 0u64, 0u64); let mut breaks = std::collections::BTreeMap::<&str, u64>::new(); let n = cases("jpegai", 1500, |r, seed| { let c = jpegai_case(r, &payloads); let why = || format!("seed {seed}\n{c:?}"); let b = |v: &Value| serde_json::to_vec(v).unwrap(); let from_bytes = jpegai_ops::load_bytes(&b(&c.norm), &b(&c.hs[0]), &b(&c.hs[1])); let (norm, hs) = jpegai_parsed(&c); let loaded = jpegai_ops::load(&norm, &hs[0], &hs[1]); assert_eq!(from_bytes.as_ref().err().map(|r| r.code), loaded.as_ref().err().map(|r| r.code), "the bytes and the parsed form load alike\n{}", why()); if let Some(k) = jpegai_planned_break(&c) { *breaks.entry(k).or_default() += 1; assert_eq!(loaded.err().map(|r| r.code), Some(wai_ops::Code::ModelMismatch), "{k}: the bundle must be refused\n{}", why()); return; } let sp = loaded.unwrap_or_else(|r| panic!("a bundle under no planned break is refused: {r:?}\n{}", why())); for (k, p) in [Plane::Luma, Plane::Chroma].into_iter().enumerate() { let gain = if k == 0 { &norm.gain_y } else { &norm.gain_uv }; let new = jpegai_ops::plane(&sp, p, &c.soz, &c.lat[k], &c.grfs[k]); let old = catch_unwind(AssertUnwindSafe(|| jpegai_ops::legacy::plane(&norm, &hs[k], gain, p, &c.soz, &c.lat[k], &c.grfs[k]))).ok(); match (new, old) { (Ok(a), Some(b)) => { assert!(a == b, "{p:?}: the two decodes differ\n{}", why()); alike += 1 } (Err(r), None) => { assert_eq!(r.code, wai_ops::Code::StreamExhausted, "{p:?}\n{}", why()); refused += 1 } (Err(r), Some(_)) => { assert_eq!(r.code, wai_ops::Code::StreamExhausted, "{p:?}\n{}", why()); assert!(!cfg!(debug_assertions), "{p:?}: refused, and the checked build's direct decode did not panic\n{}", why()); computed_on += 1 } (Ok(_), None) => panic!("{p:?}: the direct decode panicked where the operator set decodes\n{}", why()), } } }); eprintln!("jpegai: {n} cases; planes decoded alike {alike}, refused where the direct decode panicked {refused}, refused where the unchecked build computed on {computed_on}; planned breaks {breaks:?}"); assert!(alike * 8 > n, "the generator must produce decodable planes: {alike} of {n} cases"); assert!(refused > 0, "no stream was refused"); assert!(n < 1000 || breaks.len() >= 5, "too few kinds of planned break reached: {breaks:?}"); } // -------------------------------------------------------- WIV1 and WIR1 (M8, no graph) // // `wai.video.int_motion` (and the P-frames of `wai.video.int_hyper`) and the four // `int_refine` capabilities decode through the stage operators alone // (integer-operators §13.5). They make no planned break: the decode through the operator // set and the direct decode give the same frames or units, or the same refusal, on // every input, and neither panics. /// A random `WIV1` clip, its fields as `int_video::VideoStream` holds them, and for /// keyframe 1 its I-frame. #[cfg(feature = "video_int")] #[derive(Debug)] struct VideoCase { n_frames: usize, h: usize, w: usize, block: usize, q: i64, offset: i32, mv_precision: u8, cdf: Vec, mvs: Vec<(i16, i16)>, rans: Vec, i_frame: Option>, } /// A table of `len` entries from 0 to 65536, strictly increasing: `len − 1` positive /// frequencies of random sizes. #[cfg(feature = "video_int")] fn table_of(r: &mut Rng, len: usize) -> Vec { let mut cuts: Vec = Vec::new(); while cuts.len() < len - 2 { let c = 1 + r.below(65_535) as i32; if !cuts.contains(&c) { cuts.push(c); } } cuts.sort_unstable(); let mut t = vec![0]; t.extend(cuts); t.push(1 << 16); t } /// A random clip: 1 to 4 frames of 1 to 12 × 1 to 12 pixels in blocks of 1 to 6 (so the /// block grid's last row and column are often partial); whole-pel or half-pel; `q` /// mostly small, one clip in sixteen 65535 and one in sixteen anywhere in its range; a /// table of 3 to 6 entries (one in eight up to 64) at an offset mostly small, one in /// sixteen at either end of `i32`; vectors mostly within two frames of the pixel, one in /// sixteen anywhere in `i16`; keyframe 1 one clip in four, with an I-frame of the clip's /// size (one in 32 a byte short). The residuals are coded against the table with the /// crate's encoder: mostly inside the table, one in eight escaped, one in 64 near the /// ends of `i32`. One clip in 64 breaks a field (`q`, the precision, the table, the /// vector count). The stream is damaged as [`damage`] says. #[cfg(feature = "video_int")] fn int_video_case(r: &mut Rng) -> VideoCase { let n_frames = 1 + r.below(4) as usize; let (h, w, block) = (1 + r.below(12) as usize, 1 + r.below(12) as usize, 1 + r.below(6) as usize); let mv_precision = if n_frames == 1 { 1 } else { 1 + r.below(2) as u8 }; let q = match r.below(16) { 0 => 65_535, 1 => 1 + r.below(65_535) as i64, _ => 1 + r.below(12) as i64, }; let len = if r.below(8) == 0 { 3 + r.below(62) as usize } else { 3 + r.below(4) as usize }; let cdf = table_of(r, len); let offset = match r.below(16) { 0 => i32::MIN + r.below(8) as i32, 1 => i32::MAX - r.below(8) as i32, _ => r.signed(40) as i32, }; let nb = h.div_ceil(block) * w.div_ceil(block); let far = r.below(16) == 0; let mvs: Vec<(i16, i16)> = (0..nb * (n_frames - 1)) .map(|_| if far { (r.signed(32_767) as i16, r.signed(32_767) as i16) } else { (r.signed(2 * h as i64 + 2) as i16, r.signed(2 * w as i64 + 2) as i16) }) .collect(); let i_frame = (r.below(4) == 0).then(|| (0..h * w * 3 - (r.below(32) == 0) as usize).map(|_| r.below(256) as u8).collect::>()); let first = i_frame.is_some() as usize; let n2 = cdf.len() as i64 - 2; let mut enc = wai::rans::RansEncoder::new(); for _ in 0..(n_frames - first) * h * w * 3 { let v = match r.below(64) { 0 => [i32::MIN as i64, i32::MAX as i64, i32::MIN as i64 + 1, i32::MAX as i64 - 1][r.below(4) as usize], 1..=8 => { let far = 1u64 << r.below(24); offset as i64 + if r.below(2) == 0 { -1 - r.below(far) as i64 } else { n2 + r.below(far) as i64 } } _ => offset as i64 + r.below(n2 as u64) as i64, }; let v = v.clamp(i32::MIN as i64, i32::MAX as i64) as i32; if !enc.encode(v, &cdf, cdf.len(), offset) { let inside = (offset as i64 + r.below(n2 as u64) as i64).clamp(i32::MIN as i64, i32::MAX as i64) as i32; assert!(enc.encode(inside, &cdf, cdf.len(), offset) || offset as i64 + n2 > i32::MAX as i64, "a value inside the table codes"); } } let mut rans = enc.flush(); damage(r, &mut rans); let mut c = VideoCase { n_frames, h, w, block, q, offset, mv_precision, cdf, mvs, rans, i_frame }; match r.below(64) { 0 => c.q = [0, 65_536, -1][r.below(3) as usize], 1 => c.mv_precision = [0, 3][r.below(2) as usize], 2 => c.cdf[1] = 0, 3 => { c.mvs.pop(); } _ => {} } c } /// The case's clip through the stage operators and directly, each under /// `catch_unwind`. #[cfg(feature = "video_int")] type Frames = Result>, &'static str>; #[cfg(feature = "video_int")] fn video_both(c: &VideoCase) -> (std::thread::Result, std::thread::Result) { use wai::int_video::{legacy, ops, VideoStream}; let s = VideoStream { n_frames: c.n_frames, h: c.h, w: c.w, block: c.block, q: c.q, offset: c.offset, mv_precision: c.mv_precision, cdf: &c.cdf, mvs: &c.mvs, rans: &c.rans, }; let through = catch_unwind(AssertUnwindSafe(|| match &c.i_frame { None => ops::decode_typed(&s).map(|c| c.frames), Some(f) => ops::decode_pframes_typed(f, &s).map(|c| c.frames), })); let direct = catch_unwind(AssertUnwindSafe(|| match &c.i_frame { None => legacy::decode_typed(&s).map(|c| c.frames), Some(f) => legacy::decode_pframes_typed(f, &s).map(|c| c.frames), })); (through, direct) } /// The case's clip with E9 in its 32-bit form (`int_video::ops::decode_typed_i32`). #[cfg(all(feature = "video_int", feature = "i32_forms"))] fn video_i32(c: &VideoCase) -> Frames { use wai::int_video::{ops, VideoStream}; let s = VideoStream { n_frames: c.n_frames, h: c.h, w: c.w, block: c.block, q: c.q, offset: c.offset, mv_precision: c.mv_precision, cdf: &c.cdf, mvs: &c.mvs, rans: &c.rans, }; match &c.i_frame { None => ops::decode_typed_i32(&s).map(|c| c.frames), Some(f) => ops::decode_pframes_typed_i32(f, &s).map(|c| c.frames), } } /// Random clips: the decode through the stage operators (`int_video::ops`) and the /// direct decode (`int_video::legacy`) give the same frames or the same refusal, and /// neither panics; with `i32_forms`, so does the decode with E9 in its 32-bit form. #[cfg(feature = "video_int")] #[test] fn int_video_random_clips_decode_alike() { let mut outcomes = std::collections::BTreeMap::<&str, u64>::new(); let n = cases("int_video", 3000, |r, seed| { let c = int_video_case(r); let (through, direct) = video_both(&c); let (Ok(through), Ok(direct)) = (through, direct) else { panic!("a decode panicked, seed {seed}\n{c:?}") }; if through != direct { panic!("the two decodes differ, seed {seed}: {:?} through the operator set, {:?} directly\n{c:?}", through.as_ref().err(), direct.as_ref().err()); } #[cfg(feature = "i32_forms")] assert!(video_i32(&c) == through, "the 32-bit form differs from the exact form, seed {seed}\n{c:?}"); *outcomes.entry(through.err().unwrap_or("decoded alike")).or_default() += 1; }); eprintln!("int_video: {n} cases; {outcomes:?}"); let alike = outcomes.get("decoded alike").copied().unwrap_or(0); assert!(alike * 4 > n, "the generator must produce decodable clips: {alike} of {n}"); assert!(outcomes.len() >= 4, "too few kinds of outcome: {outcomes:?}"); } /// A unit's current output, or a layer's target, as a buffer of its element type. #[cfg(feature = "staged")] #[derive(Debug, Clone)] enum Units { U8(Vec), I16(Vec), } #[cfg(feature = "staged")] impl Units { fn below(&self) -> wai::staged::wir1::Below<'_> { match self { Units::U8(v) => wai::staged::wir1::Below::U8(v), Units::I16(v) => wai::staged::wir1::Below::I16(v), } } fn random(r: &mut Rng, kind: wai::staged::wir1::UnitKind, n: usize) -> Units { let ends = r.below(8) == 0; match kind { wai::staged::wir1::UnitKind::Pcm => Units::I16((0..n).map(|_| if ends && r.below(2) == 0 { [i16::MIN, i16::MAX][r.below(2) as usize] } else { r.next() as i16 }).collect()), _ => Units::U8((0..n).map(|_| if ends && r.below(2) == 0 { [0, 255][r.below(2) as usize] } else { r.next() as u8 }).collect()), } } } /// A random `WIR1` layer and the current outputs its units are applied to. #[cfg(feature = "staged")] #[derive(Debug)] struct LayerCase { kind: wai::staged::wir1::UnitKind, payload: Vec, belows: Vec<(Units, (u32, u32, u32))>, } /// `p` (a layer of `kind`) with unit `u`'s stream changed by `f`, its length field /// rewritten; `p` unchanged where the layer does not parse. #[cfg(feature = "staged")] fn restream(kind: wai::staged::wir1::UnitKind, p: &[u8], u: usize, f: impl FnOnce(&mut Vec)) -> Vec { let Ok(w) = wai::staged::wir1::parse(p, kind) else { return p.to_vec() }; let mut streams: Vec> = (0..w.n_units as usize).map(|k| w.stream(p, k).to_vec()).collect(); f(&mut streams[u]); let at = 44 + 4 * w.cdf_len * w.n_tables(); let mut out = p[..at].to_vec(); for s in &streams { out.extend((s.len() as u32).to_le_bytes()); } for s in &streams { out.extend(s); } out } /// A random layer: a unit kind and a predictor it allows; 1 to 3 units; RGB8 units of 1 /// to 10 × 1 to 10, attribute units of 1 to 3 planes of 1 to 5 × 1 to 5, PCM units of 1 /// to 3 channels; up2 from the half-size picture; rate at 1 to 6 output samples per 1 to /// 6 source samples, one layer in eight at rates up to 2^20, from 1 to 40 source /// samples. The current outputs and the targets random, one in eight with many elements /// at the range's ends; `q` mostly 1 to 8, one layer in sixteen up to 65535; a table of /// at most 2 to 16 symbols (one layer in four), so that most residuals escape, or the /// encoder's span; one table or one per unit. Coded by the reference encoder /// (`wir1::encode`). One layer in eight has a unit's stream damaged as [`damage`] says, /// and one in sixteen applies a unit to a current output of another shape or type. #[cfg(feature = "staged")] fn wir1_case(r: &mut Rng) -> Option { use wai::staged::wir1::{encode, Predictor, RefineParams, UnitKind}; let kind = [UnitKind::Rgb8, UnitKind::Attr, UnitKind::Pcm][r.below(3) as usize]; let predictor = match (kind, r.below(2)) { (UnitKind::Rgb8, 1) => Predictor::Up2, (UnitKind::Pcm, 1) => Predictor::Rate, _ => Predictor::Identity, }; let mut rate = (0u32, 0u32); let below_shape: (u32, u32, u32); let dims = match (kind, predictor) { (UnitKind::Rgb8, _) => { let d = (1 + r.below(10) as u32, 1 + r.below(10) as u32, 3); below_shape = if predictor == Predictor::Up2 { (d.0.div_ceil(2), d.1.div_ceil(2), 3) } else { d }; d } (UnitKind::Attr, _) => { let d = (1 + r.below(3) as u32, 1 + r.below(5) as u32, 1 + r.below(5) as u32); below_shape = d; d } (UnitKind::Pcm, Predictor::Rate) => { let big = r.below(8) == 0; let pick = |r: &mut Rng| if big { 1 + r.below(1 << 20) as u32 } else { 1 + r.below(6) as u32 }; rate = (pick(r), pick(r)); let (n, c) = (1 + r.below(40) as u32, 1 + r.below(3) as u32); let d0 = n as u64 * rate.0 as u64 / rate.1 as u64; if d0 == 0 || d0 > 2000 { return None; } below_shape = (n, 1, c); (d0 as u32, 1, c) } _ => { let d = (1 + r.below(60) as u32, 1, 1 + r.below(3) as u32); below_shape = d; d } }; let n_units = 1 + r.below(3) as usize; let nb = (below_shape.0 * below_shape.1 * below_shape.2) as usize; let nt = (dims.0 * dims.1 * dims.2) as usize; let belows: Vec = (0..n_units).map(|_| Units::random(r, kind, nb)).collect(); let targets: Vec = (0..n_units).map(|_| Units::random(r, kind, nt)).collect(); let q = if r.below(16) == 0 { 1 + r.below(65_535) as u32 } else { 1 + r.below(8) as u32 }; let max_symbols = if r.below(4) == 0 { 2 + r.below(15) as u32 } else { 4095 }; let p = RefineParams { q, predictor, rate, table_per_unit: r.below(2) == 0, max_symbols }; let bs: Vec<_> = belows.iter().map(Units::below).collect(); let ts: Vec<_> = targets.iter().map(Units::below).collect(); let (mut payload, _) = encode(kind, &bs, below_shape, &ts, dims, &p).ok()?; if r.below(8) == 0 { let u = r.below(n_units as u64) as usize; payload = restream(kind, &payload, u, |s| damage(r, s)); } let mut belows: Vec<(Units, (u32, u32, u32))> = belows.into_iter().map(|b| (b, below_shape)).collect(); if r.below(16) == 0 { let u = r.below(n_units as u64) as usize; belows[u] = match r.below(3) { 0 => (belows[u].0.clone(), (below_shape.0 + 1, below_shape.1, below_shape.2)), 1 => (match &belows[u].0 { Units::U8(v) => Units::I16(v.iter().map(|&x| x as i16).collect()), Units::I16(v) => Units::U8(v.iter().map(|&x| x as u8).collect()) }, below_shape), _ => (match &belows[u].0 { Units::U8(v) => Units::U8(v[1..].to_vec()), Units::I16(v) => Units::I16(v[1..].to_vec()) }, below_shape), }; } Some(LayerCase { kind, payload, belows }) } /// Random layers: every unit through the stage operators (`wir1::ops`) and directly /// (`wir1::legacy`) gives the same elements or the same refusal, the streams' checks /// alike, and nothing panics; with `i32_forms`, so does every unit with E9 in its 32-bit /// form. #[cfg(feature = "staged")] #[test] fn wir1_random_layers_apply_alike() { use wai::staged::wir1::{legacy, ops, parse}; let mut outcomes = std::collections::BTreeMap::<&str, u64>::new(); let n = cases("wir1", 3000, |r, seed| { let Some(c) = wir1_case(r) else { *outcomes.entry("no layer").or_default() += 1; return; }; // A damaged stream can leave a length the parser refuses (under 8 bytes); then // neither decode runs. let w = match parse(&c.payload, c.kind) { Ok(w) => w, Err(e) => { assert_eq!(e.code(), "stream_length", "the encoder's layer, seed {seed}"); *outcomes.entry("refused at parse").or_default() += 1; return; } }; let checks = catch_unwind(AssertUnwindSafe(|| (ops::check_streams(&w, &c.payload), legacy::check_streams(&w, &c.payload)))); let Ok((a, b)) = checks else { panic!("a streams' check panicked, seed {seed}\n{c:?}") }; assert_eq!(a, b, "the streams' checks differ, seed {seed}\n{c:?}"); for (u, (below, shape)) in c.belows.iter().enumerate() { let both = catch_unwind(AssertUnwindSafe(|| { (ops::apply_unit(&w, &c.payload, u as u32, below.below(), *shape), legacy::apply_unit(&w, &c.payload, u as u32, below.below(), *shape)) })); let Ok((through, direct)) = both else { panic!("a decode panicked, seed {seed}, unit {u}\n{c:?}") }; if through != direct { panic!("the two decodes differ, seed {seed}, unit {u}: {:?} through the operator set, {:?} directly", through.as_ref().err(), direct.as_ref().err()); } #[cfg(feature = "i32_forms")] assert!(ops::apply_unit_i32(&w, &c.payload, u as u32, below.below(), *shape) == through, "the 32-bit form differs from the exact form, seed {seed}, unit {u}"); *outcomes.entry(through.err().map_or("applied alike", |e| e.code())).or_default() += 1; } }); eprintln!("wir1: {n} cases; {outcomes:?}"); let alike = outcomes.get("applied alike").copied().unwrap_or(0); assert!(alike * 2 > n, "the generator must produce decodable units: {alike} of {n}"); assert!(outcomes.len() >= 4, "too few kinds of outcome: {outcomes:?}"); } // ------------------------------------------------- jpegai's networks (M7b, §13.6.2) // // `wai.image.jpegai`'s hyper decoders, context stages and synthesis (integer-operators // §13.6.2, lowered by `jpegai-int/1` to `rec_y`, `rec_uv`, `syn_y` and `syn_uv`) against // the direct networks (`jpegai_synthesis`, `jpegai_pipeline`), on random networks with // the shapes the direct networks read and random inputs. The direct networks refuse // nothing. Where the operator set refuses, the refusal is a planned break: a network file // that fails a load check (`model_mismatch`, from the files alone), or a value past `A32` // (`field_range`) at a checked tensor, which the direct networks computed on. /// A random weight file in the export's form, member by member. struct NetFile(serde_json::Map); impl NetFile { fn new() -> NetFile { NetFile(serde_json::Map::new()) } /// `m.weight` of `dims` with `taps` products per output, its multiplier scaled so that /// the convolution keeps its input's magnitude, and `m.bias` of `bias` values when /// given. `hot` makes one convolution in about 64 grow its input past `A32`. fn conv(&mut self, r: &mut Rng, m: &str, dims: [u64; 4], taps: u64, bias: Option, hot: bool) { let n: u64 = dims.iter().product(); let a = [3i64, 255, 32767][r.below(3) as usize]; let w: Vec = (0..n).map(|_| r.signed(a)).collect(); // Most shifts 31; some others, from 24 to 40, with the multiplier scaled to match. let shift: u64 = if r.below(4) == 0 { 24 + r.below(17) } else { 31 }; let root = (taps as f64).sqrt().max(1.0); let mut mult = ((1u128 << shift) as f64 / (a as f64 * root)) as i64; mult = (mult as f64 * [0.5, 1.0, 1.5, 2.0][r.below(4) as usize]) as i64; if hot && r.below(64) == 0 { mult = mult.saturating_mul(1 << 12); } self.0.insert(format!("{m}.weight"), json!({ "kind": "conv", "dims": dims, "mult": mult.max(1), "shift": shift, "w": w })); if let Some(b) = bias { self.q16(&format!("{m}.bias"), (0..b).map(|_| r.signed(1 << 17)).collect()); } } fn q16(&mut self, name: &str, v: Vec) { self.0.insert(name.into(), json!({ "kind": "q16", "dims": [v.len()], "v": v })); } fn res(&mut self, r: &mut Rng, m: &str, c: u64, hot: bool) { self.conv(r, &format!("{m}.conv1"), [c, c, 3, 3], 9 * c, Some(c), hot); self.conv(r, &format!("{m}.conv2"), [c, c, 3, 3], 9 * c, Some(c), hot); } fn gate(&mut self, r: &mut Rng, m: &str, c: u64, hot: bool) { self.conv(r, &format!("{m}.conv"), [c, c / 4, 3, 3], 9 * c / 4, None, hot); self.conv(r, &format!("{m}.conv2"), [c, c, 1, 1], c, None, hot); } fn cab(&mut self, r: &mut Rng, c: u64, hot: bool) { self.res(r, "CAB.residual_trunk.0", c, hot); self.res(r, "CAB.residual_trunk.1", c, hot); self.conv(r, "CAB.subscale", [c, c, 3, 3], 9 * c, Some(c), hot); self.res(r, "CAB.residual_mask2.0", c, hot); self.res(r, "CAB.residual_mask3.0", c, hot); self.conv(r, "CAB.upscale", [c, c, 3, 3], 4 * c, Some(c), hot); } fn tab(&mut self, r: &mut Rng, m: &str, c: u64, hid: u64, hot: bool) { let g: Vec = (0..c).map(|_| r.below(1 << 17) as i64).collect(); self.q16(&format!("{m}.prep_data.norm1.weight"), g); self.q16(&format!("{m}.prep_data.norm1.bias"), (0..c).map(|_| r.signed(1 << 15)).collect()); self.conv(r, &format!("{m}.prep_data.conv1"), [3 * c, c, 1, 1], c, None, hot); self.conv(r, &format!("{m}.prep_data.conv2"), [3 * c, 1, 3, 3], 9, None, hot); self.q16(&format!("{m}.attn.temperature"), (0..4).map(|_| (1 << 16) + r.below(1 << 19) as i64).collect()); self.conv(r, &format!("{m}.attn.project_out"), [c, c, 1, 1], c, None, hot); let g: Vec = (0..c).map(|_| r.signed(1 << 17)).collect(); self.q16(&format!("{m}.ffn.norm1.weight"), g); self.q16(&format!("{m}.ffn.norm1.bias"), (0..c).map(|_| r.signed(1 << 15)).collect()); self.conv(r, &format!("{m}.ffn.project_in"), [2 * hid, c, 1, 1], c, None, hot); self.conv(r, &format!("{m}.ffn.dwconv"), [2 * hid, 1, 3, 3], 9, None, hot); self.conv(r, &format!("{m}.ffn.project_out"), [c, hid, 1, 1], hid, None, hot); } fn value(self) -> Value { Value::Object(self.0) } } /// A hyper decoder of internal width `w` over `cz` channels. fn ptq_hyper(r: &mut Rng, cz: u64, w: u64, hot: bool) -> Value { let mut f = NetFile::new(); f.conv(r, "conv1", [w, cz, 1, 1], cz, None, hot); f.conv(r, "conv2", [w, w, 4, 4], 4 * w, Some(w), hot); f.conv(r, "conv3", [w, w, 3, 3], 9 * w, Some(w), hot); f.conv(r, "conv4", [4 * cz, w, 3, 3], 9 * w, Some(4 * cz), hot); f.value() } /// Context stage `i` at the widths the direct stages read (160 throughout). fn ptq_mcm(r: &mut Rng, i: u64, hot: bool) -> Value { let mut f = NetFile::new(); if i > 0 { f.conv(r, "conv.0", [160, 160 * i, 1, 1], 160 * i, None, hot); f.conv(r, "conv.1", [160, 32, 3, 3], 288, Some(160), hot); } f.conv(r, "fusion_pred_net.conv1", [160, if i == 0 { 160 } else { 320 }, 1, 1], 320, None, hot); f.conv(r, "fusion_pred_net.conv2", [160, 160, 1, 1], 160, None, hot); f.conv(r, "fusion_pred_net.conv3", [160, 160, 1, 1], 160, None, hot); f.value() } /// A synthesis network at the widths the direct networks read: luma 128 wide (its /// feed-forward hidden width `hid`), chroma 64. fn ptq_synth(r: &mut Rng, chroma: bool, hid: u64, hot: bool) -> Value { let mut f = NetFile::new(); if !chroma { let w = 128; f.res(r, "first_stage.0.0", 160, hot); f.conv(r, "first_stage.1", [160, w, 3, 3], 4 * 160, Some(w), hot); f.gate(r, "first_stage.2", w, hot); f.conv(r, "conv2_t", [w, w, 3, 3], 4 * w, Some(w), hot); f.cab(r, w, hot); f.gate(r, "iact2", w, hot); f.conv(r, "conv3_t", [4 * w, w, 1, 1], w, Some(4 * w), hot); f.conv(r, "TAM.ds_conv", [w, w, 3, 3], 9 * w, Some(w), hot); f.tab(r, "TAM.TABs.0", w, hid, hot); f.tab(r, "TAM.TABs.1", w, hid, hot); f.conv(r, "TAM.us_conv", [w, w, 3, 3], 4 * w, Some(w), hot); f.gate(r, "iact3", w, hot); f.conv(r, "conv4_t", [w, 1, 3, 3], 4 * w, Some(1), hot); } else { let w = 64; f.conv(r, "first_stage.conv1", [48, 256, 3, 3], 9 * 256, Some(48), hot); f.conv(r, "conv2_t", [144, w, 3, 3], 4 * 144, Some(w), hot); f.cab(r, w, hot); f.gate(r, "iact2", w, hot); f.conv(r, "conv3_t", [4 * w, w, 1, 1], w, Some(4 * w), hot); f.tab(r, "TAM.TABs.0", w, hid, hot); f.tab(r, "TAM.TABs.1", w, hid, hot); f.gate(r, "iact3", w, hot); f.conv(r, "conv4_t", [w, 8, 3, 3], 4 * w, Some(8), hot); } f.value() } /// The direct networks' weights, parsed as the decoder parses them. fn ptq_weights(v: &Value) -> wai::jpegai_synthesis::Weights { wai::jpegai_synthesis::Weights::from_json(v).unwrap() } /// Break one load check of §13.6.2 in a network file; the name of the check, or `None` /// (one call in six). fn ptq_break(r: &mut Rng, v: &mut Value) -> Option<&'static str> { let o = v.as_object_mut().unwrap(); let convs: Vec = o.iter().filter(|(_, e)| e["kind"] == "conv").map(|(k, _)| k.clone()).collect(); let k = convs[r.below(convs.len() as u64) as usize].clone(); Some(match r.below(6) { 0 => { o.remove(&k); "a member missing" } 1 => { o[&k]["shift"] = json!(128); "a shift of 128" } 2 => { o[&k]["w"].as_array_mut().unwrap().pop(); "a weight short" } 3 => { let d = o[&k]["dims"][0].as_u64().unwrap(); o[&k]["dims"][0] = json!(d + 1); "dims" } 4 => { let b = k.replace(".weight", ".bias"); match o.get_mut(&b) { Some(e) => { e["v"].as_array_mut().unwrap().push(json!(0)); "a bias long" } None => return None, } } _ => return None, }) } /// A random Q16 tensor of `n` values, of magnitude about `2^bits`. fn q16_values(r: &mut Rng, n: usize, bits: u32) -> Vec { (0..n).map(|_| r.signed(1 << bits)).collect() } /// The outcome of the networks through the operator set: equal to the direct networks', /// or a refusal that is a planned break (checked here). fn ptq_outcome(role: &str, new: Result, wai_ops::Refusal>, old: Option>, g: &wai_ops::Loaded, why: &dyn Fn() -> String) -> &'static str { match (new, old) { (Ok(a), Some(b)) => { assert!(a == b, "{role}: the two decodes differ\n{}", why()); "alike" } (Err(r), _) => { // A value past A32, at an input or at a checked tensor's producer. assert_eq!(r.code, wai_ops::Code::FieldRange, "{role}: refused with {:?}\n{}", r, why()); let checked = match r.at { wai_ops::Site::Input(k) => g.graph().tensors[g.graph().inputs[k as usize].tensor as usize].policy, wai_ops::Site::Op(k) => g.graph().tensors[g.graph().ops[k as usize].outs[0] as usize].policy, other => panic!("{role}: refused at {other:?}\n{}", why()), }; assert_eq!(checked, wai_ops::graph::Policy::Check, "{role}: refused at a tensor that is not checked\n{}", why()); "past A32" } (Ok(_), None) => panic!("{role}: the direct network panicked where the operator set decodes\n{}", why()), } } /// Random hyper decoders and context stages, with random hyper-latents and residuals: /// `ŷ` through `rec_y` and `rec_uv` is the direct networks', or the operator set refuses /// a value past `A32` (where the direct networks computed on, or panicked with overflow /// checks on); a network file that fails a load check is refused, from its bytes and its /// parsed form alike. #[test] fn jpegai_random_networks_reconstruct_alike() { use wai::jpegai_ops::Plane; use wai::jpegai_pipeline::{reconstruct_yhat, reconstruct_yhat_uv}; use wai::jpegai_synthesis::{decode_hyper, q16_from_symbols}; let mut outcomes = std::collections::BTreeMap::<&str, u64>::new(); let n = cases("jpegai_rec", 64, |r, seed| { let hot = r.below(2) == 0; let chroma = r.below(2) == 0; let (plane, cz) = if chroma { (Plane::Chroma, 96u64) } else { (Plane::Luma, 160) }; let width = 1 + r.below(4); let mut hyp = ptq_hyper(r, cz, width, hot); let mut mcm: Vec = (0..4).map(|i| ptq_mcm(r, i, hot)).collect(); let broken = if r.below(8) == 0 { if chroma || r.below(2) == 0 { ptq_break(r, &mut hyp) } else { let i = r.below(4) as usize; ptq_break(r, &mut mcm[i]) } } else { None }; let why = || format!("seed {seed}: {plane:?}, broken {broken:?}"); let b = |v: &Value| serde_json::to_vec(v).unwrap(); // This plane's graph alone: `rec_y` from the hyper decoder and the four stages, // or `rec_uv` from the hyper decoder. let hyp_net = wai_ops::lower::jpegai_nets::parse_ptq(&b(&hyp)); let mcm_nets: Vec<_> = mcm.iter().map(|v| wai_ops::lower::jpegai_nets::parse_ptq(&b(v))).collect(); let graph = (|| -> Result { let mut h = hyp_net.clone()?; if chroma { wai_ops::lower::jpegai_nets::rec(&mut h, None, true) } else { let mut m: [wai_ops::lower::jpegai_nets::PtqNet; 4] = [mcm_nets[0].clone()?, mcm_nets[1].clone()?, mcm_nets[2].clone()?, mcm_nets[3].clone()?]; wai_ops::lower::jpegai_nets::rec(&mut h, Some(&mut m), false) } })() .and_then(|g| wai_ops::graph::load(&g.to_bytes(), wai_ops::Profile::Acc64)); // The same, from the networks as this crate parses them, where its parser reads them. let parsed = catch_unwind(AssertUnwindSafe(|| (ptq_weights(&hyp), mcm.iter().map(ptq_weights).collect::>()))).ok(); if let Some((hw, mw)) = &parsed { let g2 = (|| -> Result { let mut h = hw.ptq()?; if chroma { wai_ops::lower::jpegai_nets::rec(&mut h, None, true) } else { let mut m = [mw[0].ptq()?, mw[1].ptq()?, mw[2].ptq()?, mw[3].ptq()?]; wai_ops::lower::jpegai_nets::rec(&mut h, Some(&mut m), false) } })() .and_then(|g| wai_ops::graph::load(&g.to_bytes(), wai_ops::Profile::Acc64)); assert_eq!(g2.as_ref().map(|g| g.graph().clone()).map_err(|e| e.code), graph.as_ref().map(|g| g.graph().clone()).map_err(|e| e.code), "the parsed networks lower as their bytes\n{}", why()); } if let Some(k) = broken { assert_eq!(graph.err().map(|e| e.code), Some(wai_ops::Code::ModelMismatch), "{k}: refused\n{}", why()); *outcomes.entry(k).or_default() += 1; return; } let g = graph.unwrap_or_else(|e| panic!("networks under no planned break are refused: {e:?}\n{}", why())); let (hw, mw) = parsed.expect("well-formed networks parse"); let c = cz as usize; let (zh, zw) = ([1usize, 2, 4][r.below(3) as usize], [1usize, 2, 4][r.below(3) as usize]); let z: Vec = (0..c * zh * zw).map(|_| r.below(64) as i64 - 31).collect(); let bits = [12u32, 20, 26][r.below(3) as usize]; let resi = q16_values(r, c * 16 * zh * zw, bits); let new = wai_ops::run_with(&g, vec![wai_ops::Tensor::new(cz as u32, zh as u32, zw as u32, z.clone()), wai_ops::Tensor::new(cz as u32, 4 * zh as u32, 4 * zw as u32, resi.clone())], &[], &mut wai_ops::Sequence::new(), &wai::int_ops::Host) .map(|mut o| o.swap_remove(0).data); let old = catch_unwind(AssertUnwindSafe(|| { let (psi, ..) = decode_hyper(&q16_from_symbols(&z), c, zh, zw, &hw); if chroma { reconstruct_yhat_uv(&resi, &psi, c, 4 * zh, 4 * zw) } else { reconstruct_yhat(&resi, &psi, c, 4 * zh, 4 * zw, [&mw[0], &mw[1], &mw[2], &mw[3]]) } })) .ok(); *outcomes.entry(ptq_outcome(if chroma { "rec_uv" } else { "rec_y" }, new, old, &g, &why)).or_default() += 1; }); eprintln!("jpegai networks, rec: {n} cases; outcomes {outcomes:?}"); assert!(outcomes.get("alike").copied().unwrap_or(0) * 2 > n, "the generator must produce decodable cases: {outcomes:?}"); } /// Random synthesis networks at the direct networks' widths, with random latents: the /// planes through `syn_y` and `syn_uv` are the direct synthesis's, or the operator set /// refuses a value past `A32`. These run at the whole picture's size, so the fixed count /// is small; the nightly run gives them their time. #[test] fn jpegai_random_synthesis_networks_decode_alike() { use wai::jpegai_synthesis::{decode_uv_f32, decode_y_f32}; let mut outcomes = std::collections::BTreeMap::<&str, u64>::new(); let n = cases("jpegai_syn", if cfg!(debug_assertions) { 2 } else { 8 }, |r, seed| { let chroma = r.below(2) == 1; let hot = r.below(2) == 0; let hid = [32u64, 64, 256][r.below(3) as usize]; let v = ptq_synth(r, chroma, hid, hot); let why = || format!("seed {seed}: chroma {chroma}, hid {hid}, hot {hot}"); let bytes = serde_json::to_vec(&v).unwrap(); let mut net = wai_ops::lower::jpegai_nets::parse_ptq(&bytes).unwrap(); let g = if chroma { wai_ops::lower::jpegai_nets::syn_uv(&mut net) } else { wai_ops::lower::jpegai_nets::syn_y(&mut net) } .and_then(|g| wai_ops::graph::load(&g.to_bytes(), wai_ops::Profile::Acc64)) .unwrap_or_else(|e| panic!("well-formed networks are refused: {e:?}\n{}", why())); let w = ptq_weights(&v); let bits = [12u32, 18, 22][r.below(3) as usize]; let ly = q16_values(r, 160 * 256, bits); let luv = q16_values(r, 96 * 256, 18); let mut inputs = vec![wai_ops::Tensor::new(160, 16, 16, ly.clone())]; if chroma { inputs.push(wai_ops::Tensor::new(96, 16, 16, luv.clone())); } let new = wai_ops::run_with(&g, inputs, &[], &mut wai_ops::Sequence::new(), &wai::int_ops::Host).map(|mut o| o.swap_remove(0).data); let old = catch_unwind(AssertUnwindSafe(|| { let f = if chroma { let mut cat = ly.clone(); cat.extend_from_slice(&luv); decode_uv_f32(&cat, &w) } else { decode_y_f32(&ly, &w) }; // Back to the Q16 integers the direct networks' denormalisation gave. f.iter().map(|&x| (x as f64 * 65536.0) as i64).collect::>() })) .ok(); *outcomes.entry(ptq_outcome(if chroma { "syn_uv" } else { "syn_y" }, new, old, &g, &why)).or_default() += 1; }); eprintln!("jpegai networks, synthesis: {n} cases; outcomes {outcomes:?}"); }