f4de2c781c
- jpegai_decode::{decode, decode_with_stages} load the bundle's networks
(load_networks, under acc64) beside its scale path and run the hyper decoders,
context stages and synthesis through rec_y, rec_uv, syn_y and syn_uv;
decode_loaded(data, icci, norm, scale path, networks, observe) decodes with
both loaded once. The direct networks are reached only from
jpegai_decode::legacy (Front::Direct, Back::Direct, test-only, legacy_int).
- The loaders no longer panic: Weights::from_json, IntNet::from_json and
IcciNet::from_json return Result<_, BadMember>; a mistyped member is
model_mismatch. IntNet keeps only well-formed tensors, so a member §13.6.1 reads
that is mistyped is missing, as the lowering of the bytes refuses it.
- The neural sink lowers and loads the networks from their eight files' bytes
through the store of loaded sets (keyed by their listing, at acc64), parses only
the eICCI net besides, and reports a refusal (model_mismatch, field_range)
without falling back: a_refused_exact_decode_does_not_fall_back gains a
non-integer network weight and a mistyped eICCI field.
- The browser crate's decodeJpegAiFull reports a mistyped member.
- Goldens, unchanged (build host, release): jpegai_cross_platform's nine pinned
hashes, the Synthesis stage d56695726bb6ebed… among them; the decoded RGB
SHA-256 dbc494a5a670914e… (jpegai_reconstruction_binding with
WAI_JPEGAI_EXPECT_SHA256, jpegai_e2e_full: max|Δ| = 1 on 4.446%).
- Memory (jpegai_memory, reference kernels): with the scale path and the networks
loaded beforehand the decode holds 72,752,387 bytes, the direct decode
175,465,380; the networks take 89,727,913 to load and hold 55,530,900.
- Also green: clippy's disallowed-methods guard and jpegai_decode_symbols, the
sink's three dispatch tests, wai-web's tests and its wasm32 build; include/wai.h
unchanged.
- CHANGELOG: the M7b entry.
Co-Authored-By: Claude Opus 5.5 <[email protected]>
3187 lines
163 KiB
Rust
3187 lines
163 KiB
Rust
//! 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<u64> = 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::<Value>(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<u64>| 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<u8>, Vec<u8>) {
|
||
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<i64> = (0..c * out * k * k).map(|_| if r.below(4) == 0 { 0 } else { r.signed(wmax) }).collect();
|
||
let bias: Vec<i64> = (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<Result<IntImage, IntSynthError>> {
|
||
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<i64> = [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<Value>) -> Vec<u8> {
|
||
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<u8> {
|
||
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<i16>, m: Vec<i64>, b: Vec<i64> },
|
||
Conv { inc: u32, out: u32, k: u32, s: u32, p: u32, w: Vec<i16>, m: Vec<i64>, b: Vec<i64> },
|
||
Igdn { c: u32, g: Vec<i32>, b: Vec<i64> },
|
||
Relu,
|
||
Leaky { q: i64, f: u32 },
|
||
}
|
||
|
||
/// A `WHM1` network's bytes.
|
||
fn whm1_bytes(afrac: u32, gf: u32, shift: u32, layers: &[Whm]) -> Vec<u8> {
|
||
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<u8>, 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<Vec<i32>>,
|
||
offsets: Vec<i32>,
|
||
median_q: Vec<i64>,
|
||
rans: Vec<u8>,
|
||
}
|
||
|
||
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<wai_ops::Graph, wai_ops::Refusal>) -> 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<Whm>, 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<i16> = (0..c * out * k * k).map(|_| if r.below(4) == 0 { 0 } else { r.signed(wmax) as i16 }).collect();
|
||
let m: Vec<i64> = (0..out).map(|_| if r.below(32) == 0 { r.signed(1 << 40) } else { r.signed(1 << 12) }).collect();
|
||
let b: Vec<i64> = (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<i32> = (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<i64> = (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<i64> {
|
||
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<u8>, 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<i32> = (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<i32> = 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<Result<Vec<u8>, 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::<String, u64>::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<i32>| {
|
||
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<RefusalKind>| {
|
||
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<i16>, m: Vec<i64>, b: Vec<i64> },
|
||
Igdn { c: u32, g: Vec<i32>, b: Vec<i64> },
|
||
}
|
||
|
||
/// A `WAM1` network's bytes.
|
||
fn wam1_bytes(afrac: u32, gf: u32, shift: u32, layers: &[Wam]) -> Vec<u8> {
|
||
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<i64>,
|
||
cdfs: Vec<Vec<i32>>,
|
||
rans: Vec<u8>,
|
||
}
|
||
|
||
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<i64> {
|
||
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<u8>, 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<i32> = (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<i64> = (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<i16> = (0..c * out * k).map(|_| if r.below(4) == 0 { 0 } else { r.signed(wmax) as i16 }).collect();
|
||
let m: Vec<i64> = (0..out).map(|_| if r.below(32) == 0 { r.signed(1 << 40) } else { r.signed(1 << 12) }).collect();
|
||
let b: Vec<i64> = (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<i32> = (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<i64> = (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<i32> = (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<Result<Vec<i16>, 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::<String, u64>::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<i32>| {
|
||
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<RefusalKind>| {
|
||
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<i64>,
|
||
cdfs: Vec<Vec<i32>>,
|
||
cdf_lengths: Vec<usize>,
|
||
offsets: Vec<i32>,
|
||
z_cdfs: Vec<Vec<i32>>,
|
||
z_cdf_lengths: Vec<usize>,
|
||
z_offsets: Vec<i32>,
|
||
z_median_q: Vec<i64>,
|
||
m: usize,
|
||
sh: usize,
|
||
sw: usize,
|
||
n_z: usize,
|
||
zh: usize,
|
||
zw: usize,
|
||
z_rans: Vec<u8>,
|
||
rans: Vec<u8>,
|
||
}
|
||
|
||
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<i32> {
|
||
let len = 3 + r.below(4) as usize;
|
||
let mut cuts: Vec<i32> = (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<i16> = (0..c * out * k * k).map(|_| if r.below(4) == 0 { 0 } else { r.signed(wmax) as i16 }).collect();
|
||
let m: Vec<i64> = (0..out).map(|_| if r.below(32) == 0 { r.signed(1 << 40) } else { r.signed(1 << 12) }).collect();
|
||
let b: Vec<i64> = (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<Vec<u8>> {
|
||
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<u8>) {
|
||
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<u8>, Vec<u8>, 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<Vec<i32>> = (0..n_edges).map(|_| cdf_table(r)).collect();
|
||
let mut cdf_lengths: Vec<usize> = cdfs.iter().map(Vec::len).collect();
|
||
let offsets: Vec<i32> = (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<Vec<i32>> = (0..k).map(|_| cdf_table(r)).collect();
|
||
let mut z_cdf_lengths: Vec<usize> = 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<i32> = (0..k).map(|_| r.signed(3) as i32).collect();
|
||
let mut z_median_q: Vec<i64> = (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<i32> = 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<Vec<i64>> = (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<Vec<i64>> = 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<Result<Vec<u8>, 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::<String, u64>::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<i32>| {
|
||
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<RefusalKind>| {
|
||
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<u8> {
|
||
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<u8>; 4]> {
|
||
let mut o = 4;
|
||
let mut s: [Vec<u8>; 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<u8> {
|
||
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<u32> = 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<usize>, Vec<i32>, Result<Vec<u8>, String>);
|
||
|
||
fn mlic_outcome((t, r): (wai::int_mlicpp::Trace, Result<wai::int_mlicpp::IntMlicImage, wai::int_mlicpp::IntMlicError>)) -> 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<i16> = (0..c * out * k * k).map(|_| if r.below(4) == 0 { 0 } else { r.signed(wmax) as i16 }).collect();
|
||
let m: Vec<i64> = (0..out).map(|_| if r.below(128) == 0 { r.signed(1 << 40) } else { r.signed(1 << 12) }).collect();
|
||
let b: Vec<i64> = (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<u8>, Vec<u8>) {
|
||
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<u32> = [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<i32> = (0..c * c).map(|_| r.below(1 << 12) as i32).collect();
|
||
let b: Vec<i64> = (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<Vec<i32>> = (0..n_edges).map(|_| cdf_table(r)).collect();
|
||
let mut cdf_lengths: Vec<usize> = cdfs.iter().map(Vec::len).collect();
|
||
let offsets: Vec<i32> = (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<Vec<i32>> = (0..k).map(|_| cdf_table(r)).collect();
|
||
let mut z_cdf_lengths: Vec<usize> = 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<i32> = (0..k).map(|_| r.signed(3) as i32).collect();
|
||
let mut z_median_q: Vec<i64> = (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<i32> = 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::<Vec<u8>>();
|
||
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<Vec<i64>> = (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<i32>, 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<MlicOutcome> {
|
||
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::<String, u64>::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::<Vec<_>>().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<Whm>| 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<u8>, 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<u8>,
|
||
g_s: Vec<u8>,
|
||
c: usize,
|
||
lat: usize,
|
||
z0: Vec<i64>,
|
||
rows: Vec<Vec<i64>>,
|
||
}
|
||
|
||
/// 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<u32> {
|
||
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<i64> = (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<i64> = (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<usize> {
|
||
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<Vec<Vec<u8>>, (&'static str, usize, Option<usize>, 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::<String, u64>::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<Whm>| whm1_bytes(afrac, gf, shift, &l);
|
||
let case = |t: Vec<u8>, g_s: Vec<u8>, c: usize, rows: Vec<Vec<i64>>| 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<u8>,
|
||
lat: [Vec<u8>; 2],
|
||
grfs: [Vec<bool>; 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<u8> {
|
||
let mut v: Vec<u8> = 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<u8> = (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<i64> { (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<i64> {
|
||
let mut v: Vec<i64> = (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<i64> { (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<u8>; 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<Vec<i64>> = (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<i64> = (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<i64> {
|
||
let mut c = 0i64;
|
||
(0..63).map(|_| { c = (c + r.below(9) as i64).min(255); c }).collect()
|
||
};
|
||
let mut cdfs_y: Vec<i64> = (0..160).flat_map(|_| cdf_row(r)).collect();
|
||
let mut cdfs_uv: Vec<i64> = (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<i64> = (0..160).map(|_| gy0 + r.signed(100)).collect();
|
||
let gain_uv: Vec<i64> = (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<i64> { v.as_array().unwrap().iter().map(|x| x.as_i64().unwrap()).collect() };
|
||
let pmf: Vec<Vec<i64>> = 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::<i64>() != 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<u64> { 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<i64> = ["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<i64> { 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<u8>; 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<i32>,
|
||
mvs: Vec<(i16, i16)>,
|
||
rans: Vec<u8>,
|
||
i_frame: Option<Vec<u8>>,
|
||
}
|
||
|
||
/// 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<i32> {
|
||
let mut cuts: Vec<i32> = 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::<Vec<u8>>());
|
||
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<Vec<Vec<u8>>, &'static str>;
|
||
|
||
#[cfg(feature = "video_int")]
|
||
fn video_both(c: &VideoCase) -> (std::thread::Result<Frames>, std::thread::Result<Frames>) {
|
||
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<u8>),
|
||
I16(Vec<i16>),
|
||
}
|
||
|
||
#[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<u8>,
|
||
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<u8>)) -> Vec<u8> {
|
||
let Ok(w) = wai::staged::wir1::parse(p, kind) else { return p.to_vec() };
|
||
let mut streams: Vec<Vec<u8>> = (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<LayerCase> {
|
||
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<Units> = (0..n_units).map(|_| Units::random(r, kind, nb)).collect();
|
||
let targets: Vec<Units> = (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<String, Value>);
|
||
|
||
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<u64>, hot: bool) {
|
||
let n: u64 = dims.iter().product();
|
||
let a = [3i64, 255, 32767][r.below(3) as usize];
|
||
let w: Vec<i64> = (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<i64>) {
|
||
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<i64> = (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<i64> = (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<String> = 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<i64> {
|
||
(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<Vec<i64>, wai_ops::Refusal>, old: Option<Vec<i64>>, 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<Value> = (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<wai_ops::Graph, wai_ops::Refusal> {
|
||
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::<Vec<_>>()))).ok();
|
||
if let Some((hw, mw)) = &parsed {
|
||
let g2 = (|| -> Result<wai_ops::Graph, wai_ops::Refusal> {
|
||
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<i64> = (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::<Vec<i64>>()
|
||
}))
|
||
.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:?}");
|
||
}
|