8386f0f45a
The duplicate ports neural/jpegai_hyper.rs and neural/jpegai_mans.rs, which the float decode path ran its front half with, are deleted (O10): that path runs the shared driver since the switch. The direct pieces of the exact decoder leave production, test-only (legacy_int): the crate-root mANS port (jpegai_mans), the hyper-scale network (jpegai_synthesis::decode_hyper_scale and its convolution) and the pipeline's skip_mask, rvs_additive, rvs_likely and sigma_index. The calibration tests that read them require legacy_int. jpegai_mans_golden runs on E8 and the scale graph alone; jpegai_e2e (the float chain against the reference PNG) takes its front half from the driver, still checking the dumped RVS additions, masks and residuals; the byte-equality crate drops the neural port and keeps E8 and the direct port. Co-Authored-By: Claude Opus 5.5 <[email protected]>
396 lines
19 KiB
Rust
396 lines
19 KiB
Rust
//! Integer Gaussian-conditional entropy — the integer-entropy contract's
|
||
//! reusable core (SPEC §7, *Replicate-soundness*).
|
||
//!
|
||
//! A learned conditional codec entropy-decodes each latent against a CDF chosen
|
||
//! by that element's scale `σ`. In the float path ([`crate::neural`]'s
|
||
//! `hyperprior::build_index`) `σ` is a float from an ONNX hyper-decoder and the
|
||
//! bucket is a float comparison, so ~1e-5 of cross-platform float drift in `σ`
|
||
//! flips a boundary bucket and desyncs the **entire** stream. That single fact
|
||
//! is why the conditional-entropy family is
|
||
//! [`crate::codecs::ReplicateClass::NeuralFloat`]: even with an integer
|
||
//! arithmetic coder downstream, a float σ source makes the decode reproducible
|
||
//! only on the same build.
|
||
//!
|
||
//! This module is the integer replacement for that bucketing step: `σ` carried
|
||
//! as a fixed-point integer, bucketed by integer comparison. It is deterministic
|
||
//! **by construction** — given the same integer `σ` it returns the same bucket
|
||
//! on every machine, because there is no float operation on the path. The
|
||
//! existing integer rANS is the arithmetic coder these indices feed; the integer
|
||
//! synthesis ([`crate::int_synth`]) is the other half.
|
||
//!
|
||
//! ## What graduates a codec to `NeuralIntegerExact`
|
||
//!
|
||
//! Adopting integer bucketing is necessary but not sufficient. The desync lives
|
||
//! in the *source* of `σ`, not only in the comparison. A codec is cross-hardware
|
||
//! replicate-sound when **all three** hold:
|
||
//!
|
||
//! 1. `σ` is produced by an **integer** hyper-decoder, so the σ bytes are
|
||
//! identical on every machine (the reference integer source is JPEG AI's scale
|
||
//! path through the integer operator set, `jpegai_ops`, whose `scale_index`
|
||
//! (integer-operators §5E E5) is the integer-log-domain special case of
|
||
//! [`det_build_index`]);
|
||
//! 2. the scale→bucket step uses this module (no float comparison);
|
||
//! 3. the synthesis is integer ([`crate::int_synth`]).
|
||
//!
|
||
//! With a **float** ONNX hyper-decoder, [`quantize_scale`] snaps σ to the grid
|
||
//! and removes drift *inside a bucket*, but cannot guarantee the boundary case —
|
||
//! so such a codec stays `NeuralFloat`. Pure i64, no ONNX: like `int_synth`,
|
||
//! this builds and is checkable for bit-exactness without `ort`, and can compile
|
||
//! to WASM for an in-browser deterministic decode.
|
||
|
||
/// Precompute a CompressAI-style ascending `scale_table` as fixed-point integer
|
||
/// bucket edges (Q`frac`: `frac` fractional bits, grid step `2^-frac`). Compute
|
||
/// once per model; pass the result to [`det_build_index`].
|
||
pub fn fixed_scale_edges(scale_table: &[f32], frac: u32) -> Vec<i64> {
|
||
let s = (1u64 << frac) as f64;
|
||
scale_table.iter().map(|&e| (e as f64 * s).floor() as i64).collect()
|
||
}
|
||
|
||
/// Snap a float `σ` to the fixed-point grid (Q`frac`), deterministic floor.
|
||
///
|
||
/// For a codec whose σ source is still a float ONNX hyper-decoder: this removes
|
||
/// drift smaller than the grid step *within* a bucket, but is **not** a hard
|
||
/// guarantee at a bucket edge (a σ within one grid step of an edge can still
|
||
/// land differently across platforms). A guaranteed-replicate codec must source
|
||
/// σ from an integer hyper-decoder, where the integer σ is exact and this snap
|
||
/// is unnecessary.
|
||
pub fn quantize_scale(sigma: f32, frac: u32) -> i64 {
|
||
let s = (1u64 << frac) as f64;
|
||
((sigma.max(0.0) as f64) * s).floor() as i64
|
||
}
|
||
|
||
/// Integer `build_index`: bucket a fixed-point `σ` to a scale-table index, with
|
||
/// the same monotone semantics as the float `hyperprior::build_index` (index
|
||
/// rises with σ, clamped to `[0, len-1]`), but pure integer — **identical on
|
||
/// every machine** given the same `sigma_q`. `edges_q` comes from
|
||
/// [`fixed_scale_edges`]; `sigma_q` from an integer hyper-decoder (exact) or
|
||
/// [`quantize_scale`] (float source).
|
||
pub fn det_build_index(sigma_q: i64, edges_q: &[i64]) -> usize {
|
||
let n = edges_q.len();
|
||
debug_assert!(n >= 1, "scale table must be non-empty");
|
||
let sigma_q = sigma_q.max(edges_q[0]); // never below the smallest tabulated scale
|
||
let mut idx = n - 1;
|
||
for &e in &edges_q[..n - 1] {
|
||
if sigma_q <= e {
|
||
idx -= 1;
|
||
}
|
||
}
|
||
idx
|
||
}
|
||
|
||
/// Rebuild a latent element from its decoded symbol: `sym · 2^frac + base`, where
|
||
/// `base` is the element's integer mean (a conditional codec's μ) or its channel's
|
||
/// median (a factorized prior's), both at `2^frac`.
|
||
///
|
||
/// `None` when `frac > 62`, or when the exact value does not fit `i64`. The symbol
|
||
/// is read from the stream, and the bypass escape lets a stream carry any 32-bit
|
||
/// value, so an untrusted stream can ask for a value `i64` cannot hold; the caller
|
||
/// then refuses the stream. A plain `sym as i64 * one + base` would panic there in
|
||
/// an overflow-checked build and wrap without the checks, so the two builds would
|
||
/// decode the same bytes differently. Here both refuse. The arithmetic is `i128`,
|
||
/// where no step can overflow (`|sym · 2^62| < 2^94`); the Kani harness
|
||
/// `kani_dequantize_is_exact_or_refused` checks that for every input.
|
||
pub fn dequantize(sym: i32, frac: u32, base: i64) -> Option<i64> {
|
||
if frac > 62 {
|
||
return None;
|
||
}
|
||
i64::try_from(((sym as i128) << frac) + base as i128).ok()
|
||
}
|
||
|
||
/// The quantised CDF an encoder writes for a histogram, integer only: one table of
|
||
/// `hist.len() + 2` entries over the `hist.len()` symbols and the escape, the rANS
|
||
/// decoder's layout (`extensions/integer-payloads.md` §2; the escape is the last
|
||
/// symbol, `cdf_len − 2`). Every symbol, the escape included, gets frequency at least
|
||
/// 1, and the frequencies sum to exactly `2^precision`: the escape gets 1, each
|
||
/// symbol `1 + ⌊R·h_i / H⌋` with `R = 2^precision − hist.len() − 1` and `H` the
|
||
/// histogram's total, and the `R − Σ⌊·⌋` left over go one each to the symbols with
|
||
/// the largest remainders `R·h_i mod H`, ties to the lower index. An all-zero
|
||
/// histogram counts every symbol once. So the table is strictly increasing, from 0 to
|
||
/// `2^precision`, and the same histogram gives the same table on every machine.
|
||
///
|
||
/// `None` when the histogram is empty, when `precision` is outside `1..=16` (the
|
||
/// decoder's precision is 16), or when the symbols and the escape cannot each have
|
||
/// frequency 1 (`hist.len() + 1 > 2^precision`).
|
||
pub fn hist_to_cdf(hist: &[u64], precision: u32) -> Option<Vec<i32>> {
|
||
if hist.is_empty() || !(1..=16).contains(&precision) {
|
||
return None;
|
||
}
|
||
let total: u128 = 1u128 << precision;
|
||
let n = hist.len() as u128;
|
||
if n + 1 > total {
|
||
return None;
|
||
}
|
||
let r = total - n - 1;
|
||
let all_zero = hist.iter().all(|&h| h == 0);
|
||
let weight = |h: u64| -> u128 { if all_zero { 1 } else { h as u128 } };
|
||
let h_total: u128 = hist.iter().map(|&h| weight(h)).sum();
|
||
// Room for the escape too, so pushing it below never regrows the vector: the fit
|
||
// holds one frequency per symbol and the escape, never twice that.
|
||
let mut freq: Vec<u128> = Vec::with_capacity(hist.len() + 1);
|
||
freq.extend(hist.iter().map(|&h| 1 + r * weight(h) / h_total));
|
||
let given: u128 = freq.iter().map(|f| f - 1).sum();
|
||
let mut order: Vec<usize> = (0..hist.len()).collect();
|
||
// Largest remainder first; equal remainders to the lower index (a stable sort).
|
||
order.sort_by_key(|&i| std::cmp::Reverse(r * weight(hist[i]) % h_total));
|
||
for &i in order.iter().take((r - given) as usize) {
|
||
freq[i] += 1;
|
||
}
|
||
freq.push(1); // the escape
|
||
let mut cdf = Vec::with_capacity(freq.len() + 1);
|
||
let mut acc: u128 = 0;
|
||
cdf.push(0i32);
|
||
for f in freq {
|
||
acc += f;
|
||
cdf.push(acc as i32);
|
||
}
|
||
debug_assert_eq!(acc, total);
|
||
Some(cdf)
|
||
}
|
||
|
||
// ─────────────────────────────────────────────────────────────────────
|
||
// Kani proof harnesses — the soundness hinge, checked rather than asserted
|
||
// ─────────────────────────────────────────────────────────────────────
|
||
//
|
||
// SPEC §7 makes this function load-bearing: a bucket that flips selects a
|
||
// different CDF and desyncs the remainder of the stream. Two things therefore
|
||
// have to hold, and both are proved here rather than argued in prose.
|
||
//
|
||
// 1. The index addresses a real scale-table entry. An out-of-range bucket is
|
||
// either an out-of-bounds panic downstream or, worse, a silently wrong CDF.
|
||
// 2. The bucketing is *monotone* in σ. The doc comment claims "index rises
|
||
// with σ"; a non-monotone bucketing would mis-order the table without ever
|
||
// panicking, which no amount of round-trip testing on well-formed vectors
|
||
// would reliably surface.
|
||
//
|
||
// Note the second is proved for an ARBITRARY edge table, not only an ascending
|
||
// one — the property follows from counting, not from the table being sorted, so
|
||
// the proof does not depend on a precondition a caller might not maintain.
|
||
|
||
/// The bucket always addresses a real entry of the scale table.
|
||
#[cfg(kani)]
|
||
#[kani::proof]
|
||
#[kani::unwind(6)]
|
||
fn kani_det_build_index_in_range() {
|
||
let edges: [i64; 4] = kani::any();
|
||
let sigma: i64 = kani::any();
|
||
let idx = det_build_index(sigma, &edges);
|
||
kani::assert(idx < edges.len(), "bucket index addresses a real scale-table entry");
|
||
kani::cover!(true, "the bucketing is reachable");
|
||
}
|
||
|
||
/// **Monotonicity**: a larger σ never selects a smaller bucket.
|
||
///
|
||
/// This is the semantic property the integer bucketing exists to guarantee. It
|
||
/// holds by counting — a larger σ satisfies `σ <= e` for no more edges than a
|
||
/// smaller one does — so it is proved without assuming the table is ascending.
|
||
#[cfg(kani)]
|
||
#[kani::proof]
|
||
#[kani::unwind(6)]
|
||
fn kani_det_build_index_is_monotone_in_sigma() {
|
||
let edges: [i64; 4] = kani::any();
|
||
let a: i64 = kani::any();
|
||
let b: i64 = kani::any();
|
||
kani::assume(a <= b);
|
||
kani::assert(
|
||
det_build_index(a, &edges) <= det_build_index(b, &edges),
|
||
"a larger sigma never selects a smaller bucket",
|
||
);
|
||
kani::cover!(true, "both bucketings are reachable");
|
||
}
|
||
|
||
/// The single-entry table degenerates correctly: one bucket, always index 0.
|
||
#[cfg(kani)]
|
||
#[kani::proof]
|
||
#[kani::unwind(4)]
|
||
fn kani_det_build_index_singleton_table() {
|
||
let edges: [i64; 1] = kani::any();
|
||
let sigma: i64 = kani::any();
|
||
kani::assert(det_build_index(sigma, &edges) == 0, "a one-entry table has one bucket");
|
||
}
|
||
|
||
/// `dequantize` over every symbol, scale and base: it never panics, and it returns
|
||
/// the exact `sym · 2^frac + base` precisely when that value fits `i64` (and
|
||
/// `frac ≤ 62`), `None` otherwise. The expected value is formed by multiplication
|
||
/// by a power of two, not by the shift the function uses, so the harness states the
|
||
/// specification rather than restating the code.
|
||
#[cfg(kani)]
|
||
#[kani::proof]
|
||
fn kani_dequantize_is_exact_or_refused() {
|
||
let sym: i32 = kani::any();
|
||
let frac: u32 = kani::any();
|
||
let base: i64 = kani::any();
|
||
let got = dequantize(sym, frac, base);
|
||
if frac <= 62 {
|
||
let exact = sym as i128 * (1i128 << frac) + base as i128;
|
||
let fits = exact >= i64::MIN as i128 && exact <= i64::MAX as i128;
|
||
kani::assert(got.is_some() == fits, "refused exactly when the value leaves i64");
|
||
if let Some(v) = got {
|
||
kani::assert(v as i128 == exact, "the value returned is the exact one");
|
||
}
|
||
} else {
|
||
kani::assert(got.is_none(), "a scale above 2^62 is refused");
|
||
}
|
||
kani::cover!(got.is_some(), "an in-range value is returned");
|
||
kani::cover!(got.is_none() && frac <= 62, "an out-of-range value is refused");
|
||
}
|
||
|
||
#[cfg(test)]
|
||
mod tests {
|
||
use super::*;
|
||
|
||
/// The float `build_index` the conditional family uses today — copied here
|
||
/// (it is `pub(crate)` in the ORT-gated `neural` module) to demonstrate the
|
||
/// hazard this module removes. Byte-for-byte the same logic as
|
||
/// `neural::hyperprior::build_index`.
|
||
fn float_build_index(sigma: f32, scale_table: &[f32]) -> usize {
|
||
let n = scale_table.len();
|
||
let sigma = sigma.max(scale_table[0]);
|
||
let mut idx = n - 1;
|
||
for &s in &scale_table[..n - 1] {
|
||
if sigma <= s {
|
||
idx -= 1;
|
||
}
|
||
}
|
||
idx
|
||
}
|
||
|
||
/// CompressAI's default 64 log-spaced scale buckets, 0.11 … 256.
|
||
fn scale_table() -> Vec<f32> {
|
||
let (lo, hi, n) = (0.11f32, 256.0f32, 64usize);
|
||
let (ll, lh) = (lo.ln(), hi.ln());
|
||
(0..n).map(|i| (ll + (lh - ll) * i as f32 / (n - 1) as f32).exp()).collect()
|
||
}
|
||
|
||
const FRAC: u32 = 16; // Q16: grid step ~1.5e-5
|
||
|
||
/// `hist_to_cdf`'s table: total exactly 2^16, every frequency at least 1,
|
||
/// strictly increasing, the escape last with frequency 1, the same table for the
|
||
/// same histogram, and frequencies ordered as the counts are. Includes histograms
|
||
/// far larger than the total, all-zero and one-symbol histograms, and a ladder
|
||
/// where the remainders tie.
|
||
#[test]
|
||
fn hist_to_cdf_properties() {
|
||
let cases: Vec<Vec<u64>> = vec![
|
||
vec![0],
|
||
vec![5],
|
||
vec![0, 0, 0],
|
||
vec![1, 1, 1, 1, 1, 1, 1],
|
||
vec![1_000_000_000_000, 1, 0, 3, u64::MAX / 4],
|
||
(0..4095u64).map(|i| (i * 2654435761) % 1000).collect(),
|
||
(0..300u64).map(|i| if i == 150 { 1 << 40 } else { 0 }).collect(),
|
||
];
|
||
for hist in &cases {
|
||
let cdf = hist_to_cdf(hist, 16).expect("a table");
|
||
assert_eq!(cdf.len(), hist.len() + 2);
|
||
assert_eq!(cdf[0], 0);
|
||
assert_eq!(*cdf.last().unwrap(), 1 << 16, "the total is exactly 2^16");
|
||
assert!(cdf.windows(2).all(|w| w[1] > w[0]), "every frequency is at least 1");
|
||
assert_eq!(cdf[cdf.len() - 1] - cdf[cdf.len() - 2], 1, "the escape is the last bin, frequency 1");
|
||
assert_eq!(hist_to_cdf(hist, 16).unwrap(), cdf, "deterministic");
|
||
let f: Vec<i32> = cdf.windows(2).map(|w| w[1] - w[0]).collect();
|
||
for i in 0..hist.len() {
|
||
for j in 0..hist.len() {
|
||
if hist[i] > hist[j] {
|
||
assert!(f[i] >= f[j], "a larger count never gets a smaller frequency");
|
||
}
|
||
}
|
||
}
|
||
}
|
||
// Equal counts: the leftover goes to the lower indices.
|
||
let cdf = hist_to_cdf(&[1, 1, 1], 16).unwrap();
|
||
let f: Vec<i32> = cdf.windows(2).map(|w| w[1] - w[0]).collect();
|
||
assert_eq!(f, vec![21845, 21845, 21845, 1]);
|
||
assert_eq!(hist_to_cdf(&[1, 1, 1, 1], 16).unwrap()[1..].iter().zip(&cdf_ladder()).filter(|(a, b)| a != b).count(), 0);
|
||
// Refusals: empty, a precision the decoder does not use, too many symbols.
|
||
assert!(hist_to_cdf(&[], 16).is_none());
|
||
assert!(hist_to_cdf(&[1], 0).is_none());
|
||
assert!(hist_to_cdf(&[1], 17).is_none());
|
||
assert!(hist_to_cdf(&vec![1; 65536], 16).is_none());
|
||
assert!(hist_to_cdf(&vec![1; 65535], 16).is_some());
|
||
}
|
||
|
||
/// Four equal counts: R = 65531, 16382 each with 3 left over, to indices 0, 1, 2.
|
||
fn cdf_ladder() -> Vec<i32> {
|
||
vec![16384, 32768, 49152, 65535, 65536]
|
||
}
|
||
|
||
#[test]
|
||
fn integer_bucketing_matches_float_in_bucket_interiors() {
|
||
let st = scale_table();
|
||
let edges = fixed_scale_edges(&st, FRAC);
|
||
// Sample the geometric midpoint of every adjacent bucket pair — points
|
||
// far (≫ one grid step) from any edge, where quantization is exact.
|
||
for w in st.windows(2) {
|
||
let mid = (w[0] * w[1]).sqrt();
|
||
assert_eq!(
|
||
det_build_index(quantize_scale(mid, FRAC), &edges),
|
||
float_build_index(mid, &st),
|
||
"interior σ={mid} must bucket identically to the float path",
|
||
);
|
||
}
|
||
}
|
||
|
||
#[test]
|
||
fn float_path_desyncs_on_cross_platform_drift() {
|
||
// THE hazard the contract names: two platforms' float hyper-decoders
|
||
// produce σ for the same element differing by ~1e-5 (real ONNX drift).
|
||
// At a bucket edge the float bucketing DISAGREES → the stream desyncs.
|
||
let st = scale_table();
|
||
let edge = st[40];
|
||
let drift = edge * 1e-5;
|
||
let sigma_a = edge - drift; // platform A
|
||
let sigma_b = edge + drift; // platform B
|
||
assert_ne!(
|
||
float_build_index(sigma_a, &st),
|
||
float_build_index(sigma_b, &st),
|
||
"float build_index should flip the bucket across the edge (the desync)",
|
||
);
|
||
}
|
||
|
||
#[test]
|
||
fn integer_source_makes_the_bucket_platform_invariant() {
|
||
// The fix: an INTEGER hyper-decoder emits the identical integer σ on
|
||
// every platform (integer arithmetic has no drift), so there is no
|
||
// per-platform σ divergence to flip the bucket. Model that single
|
||
// deterministic σ_q and bucket it: the result is the same on A and B.
|
||
let st = scale_table();
|
||
let edges = fixed_scale_edges(&st, FRAC);
|
||
let sigma_q = 37_000_i64; // whatever the integer hyper-decoder emitted
|
||
let on_a = det_build_index(sigma_q, &edges);
|
||
let on_b = det_build_index(sigma_q, &edges); // identical integer input on B
|
||
assert_eq!(on_a, on_b, "integer bucketing is platform-invariant by construction");
|
||
// And it is monotone: a larger integer σ never buckets lower.
|
||
let mut prev = 0;
|
||
for q in (edges[0]..*edges.last().unwrap()).step_by(997) {
|
||
let idx = det_build_index(q, &edges);
|
||
assert!(idx >= prev, "bucketing must be monotone in σ");
|
||
prev = idx;
|
||
}
|
||
}
|
||
|
||
#[test]
|
||
fn clamps_below_the_smallest_scale() {
|
||
let st = scale_table();
|
||
let edges = fixed_scale_edges(&st, FRAC);
|
||
assert_eq!(det_build_index(i64::MIN / 2, &edges), 0);
|
||
assert_eq!(det_build_index(i64::MAX / 2, &edges), edges.len() - 1);
|
||
}
|
||
|
||
#[test]
|
||
fn dequantize_is_the_fixed_point_latent_or_a_refusal() {
|
||
// In range it is exactly the expression the conditional decoders use.
|
||
for (sym, base) in [(0, 0), (5, -3000), (-7, 12_345), (i32::MAX, 0), (i32::MIN, 0)] {
|
||
assert_eq!(dequantize(sym, 12, base), Some(sym as i64 * (1 << 12) + base));
|
||
}
|
||
// At the edge of i64: the largest value that fits, and one past it.
|
||
assert_eq!(dequantize(1, 62, i64::MAX - (1 << 62)), Some(i64::MAX));
|
||
assert_eq!(dequantize(1, 62, i64::MAX - (1 << 62) + 1), None);
|
||
assert_eq!(dequantize(-2, 62, 0), Some(i64::MIN));
|
||
assert_eq!(dequantize(-2, 62, -1), None);
|
||
// A stream-supplied symbol that no latent at this scale can be.
|
||
assert_eq!(dequantize(i32::MAX, 40, 0), None);
|
||
// Scales beyond 2^62 are refused rather than shifted out of range.
|
||
assert_eq!(dequantize(0, 63, 0), None);
|
||
assert_eq!(dequantize(1, u32::MAX, 0), None);
|
||
}
|
||
}
|