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.
  */

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::<Vec<u8>>();
+    let i64v = |k: &str| v[k].as_array().unwrap().iter().map(|x| x.as_i64().unwrap()).collect::<Vec<i64>>();
+    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<i64> = 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");
+    }
+}