diff --git a/wai/wai-rs/include/wai.h b/wai/wai-rs/include/wai.h index d16c9b09f..dcd0b42de 100644 --- a/wai/wai-rs/include/wai.h +++ b/wai/wai-rs/include/wai.h @@ -40,6 +40,27 @@ #define WAI_ERR_PANIC 4 +/** + * Probability output scale of [`softmax`]: results are at `2^PROB_FRAC` + * (they sum to ~`2^PROB_FRAC`). + */ +#define PROB_FRAC 16 + +/** + * Fractional bits of the Q16 fixed-point format. + */ +#define F 16 + +/** + * 1.0 in Q16. + */ +#define S (1 << F) + +/** + * Number of action planes appended to the latent for the transition (dx, dy, btn). + */ +#define A 3 + #define MARKER_SOC 65408 #define MARKER_EOC 65409 @@ -224,6 +245,11 @@ #define PULSE25 3 +/** + * `CausalOp::mutation_kind` for an action input (move/look/buttons + free payload). + */ +#define MUT_ACTION 1 + /** * Q32.32 fixed-point value. */ diff --git a/wai/wai-rs/tests/jpegai_e2e_hyperlatent.rs b/wai/wai-rs/tests/jpegai_e2e_hyperlatent.rs new file mode 100644 index 000000000..70d6847d4 --- /dev/null +++ b/wai/wai-rs/tests/jpegai_e2e_hyperlatent.rs @@ -0,0 +1,65 @@ +//! End-to-end integration: recover the hyper-latent `z` **from the real codestream**. +//! +//! This is the first stage chained straight off the bitstream rather than a fixture: +//! parse the container, take the SOZ substream, and run the byte-exact mANS `fac` +//! (factorized) decoder over it with the normative per-channel CDFs. The decoded +//! symbols are offset by 31 (the mid-point of the 63-symbol alphabet) to give signed +//! `z`, which must equal the reference `z` that `decode_hyper_scale` consumes. +//! +//! SOZ carries both planes back to back: luma (160 ch) then chroma (96 ch), 16 samples +//! per channel, decoded by one MansDecoder whose state threads across the calls. +//! +//! env: WAI_JPEGAI_SOZ_FIXTURE = e2e_soz_fixture.json +//! { codestream:[u8], cdfs_y:[u8;160*63], cdfs_uv:[u8;96*63], per_channel, +//! z_y_expected:[i64;2560], z_uv_expected:[i64;1536] } +#![cfg(feature = "neural_int")] + +use wai::jpegai_bitstream::{find_segment, parse_container, Marker}; +use wai::jpegai_mans::MansDecoder; + +/// The normative offset applied to factorized symbols: the alphabet's mid-point. +const Z_OFFSET: i64 = 31; + +#[test] +fn hyper_latent_decodes_from_the_real_codestream() { + let path = match std::env::var("WAI_JPEGAI_SOZ_FIXTURE") { + Ok(p) if std::path::Path::new(&p).exists() => p, + _ => { eprintln!("skip: set WAI_JPEGAI_SOZ_FIXTURE"); return; } + }; + let v: serde_json::Value = serde_json::from_slice(&std::fs::read(&path).unwrap()).unwrap(); + let u8v = |k: &str| v[k].as_array().unwrap().iter().map(|x| x.as_u64().unwrap() as u8).collect::>(); + let i64v = |k: &str| v[k].as_array().unwrap().iter().map(|x| x.as_i64().unwrap()).collect::>(); + let data = u8v("codestream"); + let per_channel = v["per_channel"].as_u64().unwrap() as usize; + + // 1. the container gives us the SOZ substream's bytes + let segs = parse_container(&data); + let soz = find_segment(&segs, Marker::Soz).expect("SOZ present"); + let mem = &data[soz.start..soz.start + soz.len]; + eprintln!("SOZ: {} bytes at {}", mem.len(), soz.start); + + // 2. one mANS decoder over that substream; state threads across both planes + let mut dec = MansDecoder::new(mem, &[mem.len() as u32]); + + for (plane, cdf_key, exp_key, channels) in + [("y", "cdfs_y", "z_y_expected", 160usize), ("uv", "cdfs_uv", "z_uv_expected", 96)] + { + let cdfs = u8v(cdf_key); + let expected = i64v(exp_key); + let mut symbols = vec![0u8; channels * per_channel]; + dec.decode_factorize(&cdfs, &mut symbols, channels, per_channel); + + // 3. symbols → signed z + let z: Vec = symbols.iter().map(|&s| s as i64 - Z_OFFSET).collect(); + assert_eq!(z.len(), expected.len(), "{plane}: z size"); + let mism = z.iter().zip(&expected).filter(|(a, b)| a != b).count(); + if let Some((i, (a, b))) = z.iter().zip(&expected).enumerate().find(|(_, (a, b))| a != b) { + eprintln!("{plane}: first mismatch @ {i}: got {a} exp {b}"); + } + let nz = z.iter().filter(|&&x| x != 0).count(); + eprintln!("z[{plane}]: {}/{} exact ({} non-zero, range [{}, {}])", + z.len() - mism, z.len(), nz, + z.iter().min().unwrap(), z.iter().max().unwrap()); + assert_eq!(mism, 0, "{plane}: hyper-latent must decode byte-exact from the codestream"); + } +}