2496c95410
RansEncoder::encode now codes a value outside its table's span through the bypass escape, exactly as RansDecoder reads it: the escape symbol, the count of 4-bit groups as a run of 15s closed by a group below 15, then the raw value's groups least significant first. It returns false only where no stream can code the value: a table shorter than three entries or longer than the slice, a frequency that is not positive, or value - offset outside i32. Every stream it writes still ends canonically. int_entropy::hist_to_cdf builds the encoder's table from a histogram in integer arithmetic alone: the escape is the last bin with frequency 1, every symbol gets at least 1, the total is exactly 2^precision, and the leftover goes by largest remainder with ties to the lower index. ternary::rans_encode_trits called encode inside debug_assert!, so a release build coded no symbol; the call now runs in every build. Co-Authored-By: Claude Opus 5.5 <[email protected]>
931 lines
36 KiB
Rust
931 lines
36 KiB
Rust
//! Balanced-ternary matrix-product-state transform — `wai.ternary.tt`
|
||
//! (quantum-*inspired*, classical, deterministic; extensions/ternary-tt).
|
||
//!
|
||
//! The honest lineage first. The ternary weight fabric (BitNet-b1.58 style,
|
||
//! {−1,0,+1}, five trits per byte) and the matrix-product-state (MPS / tensor
|
||
//! train) compression it rides on are established, working techniques — this
|
||
//! module does **not** reinvent them. What every existing ternary/MPS codec
|
||
//! shares, and what makes none of them WAI-conformant, is a **float scale**: a
|
||
//! per-core `α·(trit·x)` where `α` is `f32`. A float scale desyncs across
|
||
//! machines exactly as a float σ desyncs a learned entropy coder — so the decode
|
||
//! is reproducible only on the same build.
|
||
//!
|
||
//! The one change here is the whole contribution: **the scale is an integer**
|
||
//! (fixed-point, stored as a per-core shared exponent + per-column `u16`
|
||
//! mantissa, dequantized by an integer shift), and the entire *decode* path is
|
||
//! pure `i64`/`i128` with no float anywhere. So a ternary-MPS reconstruction is
|
||
//! byte-identical on every machine — a portable reconstruction hash you can put
|
||
//! in a receipt — the same `NeuralIntegerExact` discipline as `int_hyper`, now
|
||
//! for the quantum-inspired tensor-network transform. Encoder-side float (the
|
||
//! SVD, the optimal per-column scale) is fine and expected: like QAT it runs
|
||
//! once at encode and never at decode.
|
||
//!
|
||
//! This is a **lossy** codec whose decode is **deterministic**. It does not claim
|
||
//! to recover an arbitrary source losslessly; it claims that whatever it *does*
|
||
//! reconstruct is byte-identical everywhere, and its error is measured and
|
||
//! stated, never hidden ([`snr_db`]).
|
||
//!
|
||
//! Two rate knobs, both keeping the decode byte-exact:
|
||
//! - **bond dimension** `max_bond` — the MPS truncation (per layer);
|
||
//! - **residual layers** ([`TernaryStack`]) — RVQ over ternary-MPS: each layer
|
||
//! ternary-MPS-encodes the *residual* of the sum so far. Decode sums the
|
||
//! layers (a sum of deterministic integer decodes is still deterministic), so
|
||
//! fidelity climbs with layer count while the decode stays portable-exact.
|
||
//!
|
||
//! Structure of one layer: a left-canonical MPS — sites `0..n-1` are orthonormal
|
||
//! isometries (the `U` of a thin SVD), the **last** core absorbs the magnitude
|
||
//! (the remaining `S·Vᵀ`). Per-column integer scales keep the isometries
|
||
//! unit-scaled (no multiplicative magnitude drift); the last core carries the
|
||
//! amplitude. Decode contracts left-to-right in `i128`.
|
||
|
||
use crate::int_transform; // reused fixed-point idiom (i128 intermediates, >> shift)
|
||
|
||
/// Fixed-point fractional bits for core scales and the reconstruction
|
||
/// accumulator. Scales live at `2^SCALE_FRAC`; the i128 contraction rescales by
|
||
/// `>> SCALE_FRAC` after absorbing each core's per-column scale.
|
||
pub const SCALE_FRAC: u32 = 24;
|
||
|
||
/// Trits per packed byte. `3^5 = 243 < 256`, the near-entropy-optimal
|
||
/// balanced-ternary packing (1.58 bits/trit) the ternary fabric uses.
|
||
pub const TRITS_PER_BYTE: usize = 5;
|
||
|
||
// ---------------------------------------------------------------------------
|
||
// Balanced-ternary packing: trit ∈ {−1,0,+1} → digit ∈ {0,1,2}, base-3.
|
||
// ---------------------------------------------------------------------------
|
||
|
||
/// Pack balanced trits (`-1|0|+1`) five-to-a-byte (base-3). A non-`{-1,0,1}`
|
||
/// input is clamped to its sign — encoders never produce one, but the packer is
|
||
/// total. The final partial group is zero-padded; `unpack_trits` takes the exact
|
||
/// count so the padding is invisible.
|
||
pub fn pack_trits(trits: &[i8]) -> Vec<u8> {
|
||
let mut out = Vec::with_capacity(trits.len().div_ceil(TRITS_PER_BYTE));
|
||
for chunk in trits.chunks(TRITS_PER_BYTE) {
|
||
let mut byte: u16 = 0;
|
||
let mut w: u16 = 1;
|
||
for &t in chunk {
|
||
let digit = (t.clamp(-1, 1) + 1) as u16; // -1,0,1 -> 0,1,2
|
||
byte += digit * w;
|
||
w *= 3;
|
||
}
|
||
out.push(byte as u8);
|
||
}
|
||
out
|
||
}
|
||
|
||
/// Inverse of [`pack_trits`]; `n` is the exact trit count (drops padding).
|
||
pub fn unpack_trits(bytes: &[u8], n: usize) -> Vec<i8> {
|
||
let mut out = Vec::with_capacity(n);
|
||
for &b in bytes {
|
||
let mut v = b as u16;
|
||
for _ in 0..TRITS_PER_BYTE {
|
||
if out.len() == n {
|
||
break;
|
||
}
|
||
out.push((v % 3) as i8 - 1); // digit 0,1,2 -> -1,0,1
|
||
v /= 3;
|
||
}
|
||
}
|
||
out
|
||
}
|
||
|
||
// ---------------------------------------------------------------------------
|
||
// MPS types
|
||
// ---------------------------------------------------------------------------
|
||
|
||
/// One MPS core: a `[left, phys, right]` tensor of trits (`-1|0|+1`) plus a
|
||
/// compact per-column scale — a shared exponent and a `u16` mantissa per `right`
|
||
/// column. The fixed-point scale of column `j` is `(mant[j] << exp)`, recovered
|
||
/// by [`Self::scale_fx`] with a pure integer shift. `trits.len() ==
|
||
/// left*phys*right`, `scale_mant.len() == right`.
|
||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||
pub struct TernaryCore {
|
||
pub left: usize,
|
||
pub phys: usize,
|
||
pub right: usize,
|
||
pub trits: Vec<i8>,
|
||
pub scale_mant: Vec<u16>,
|
||
pub scale_exp: u8,
|
||
}
|
||
|
||
impl TernaryCore {
|
||
/// Dequantize column `j`'s scale to fixed-point `i64` at `2^SCALE_FRAC` — a
|
||
/// pure integer shift, byte-exact on every machine.
|
||
#[inline]
|
||
pub fn scale_fx(&self, j: usize) -> i64 {
|
||
(self.scale_mant[j] as i64) << self.scale_exp
|
||
}
|
||
|
||
/// Compact scale-table size in bytes: 2 B per column + 1 B shared exponent.
|
||
fn scale_bytes(&self) -> usize {
|
||
self.scale_mant.len() * 2 + 1
|
||
}
|
||
}
|
||
|
||
/// A ternary-MPS transform of a length-`n_signal` signal: an ordered chain of
|
||
/// cores (left bond of the first and right bond of the last are 1). Fully
|
||
/// integer — a portable, hashable object. One *layer* of a [`TernaryStack`].
|
||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||
pub struct TernaryMps {
|
||
pub n_sites: u8,
|
||
pub cores: Vec<TernaryCore>,
|
||
pub n_signal: usize,
|
||
}
|
||
|
||
impl TernaryMps {
|
||
/// Packed trit-payload size in bytes (the fundamental ternary density).
|
||
pub fn packed_len(&self) -> usize {
|
||
self.cores
|
||
.iter()
|
||
.map(|c| c.trits.len().div_ceil(TRITS_PER_BYTE))
|
||
.sum()
|
||
}
|
||
|
||
/// Total on-wire size in bytes: packed trits + the compact scale tables.
|
||
pub fn payload_len(&self) -> usize {
|
||
self.packed_len() + self.cores.iter().map(|c| c.scale_bytes()).sum::<usize>()
|
||
}
|
||
|
||
/// BLAKE3 over the canonical decode output — the portable reconstruction
|
||
/// hash. Two machines get this same hash iff every reconstructed sample
|
||
/// matches (the receiptable property).
|
||
pub fn reconstruction_hash(&self) -> [u8; 32] {
|
||
hash_signal(&self.decode(), self.n_signal)
|
||
}
|
||
}
|
||
|
||
fn hash_signal(sig: &[i64], n: usize) -> [u8; 32] {
|
||
let mut h = blake3::Hasher::new();
|
||
h.update(b"wai:ternary-tt-recon\x01");
|
||
h.update(&(n as u64).to_le_bytes());
|
||
for s in sig {
|
||
h.update(&s.to_le_bytes());
|
||
}
|
||
*h.finalize().as_bytes()
|
||
}
|
||
|
||
// ---------------------------------------------------------------------------
|
||
// Decode — pure integer, deterministic, byte-exact.
|
||
// ---------------------------------------------------------------------------
|
||
|
||
impl TernaryMps {
|
||
/// Contract the chain into the full signal. Pure `i128`/`i64`: trit tensors
|
||
/// contract as integers; each core's per-column integer scale folds in as a
|
||
/// `× scale >> SCALE_FRAC` step, keeping the accumulator at `2^SCALE_FRAC`.
|
||
/// No float anywhere — byte-identical on every machine.
|
||
pub fn decode(&self) -> Vec<i64> {
|
||
let c0 = &self.cores[0];
|
||
debug_assert_eq!(c0.left, 1);
|
||
let mut rows = c0.phys;
|
||
let mut bond = c0.right;
|
||
let mut acc = vec![0i128; rows * bond];
|
||
for s in 0..c0.phys {
|
||
for rp in 0..c0.right {
|
||
let t = c0.trits[s * c0.right + rp] as i128;
|
||
acc[s * bond + rp] = t * c0.scale_fx(rp) as i128; // at SCALE_FRAC
|
||
}
|
||
}
|
||
|
||
for core in &self.cores[1..] {
|
||
debug_assert_eq!(core.left, bond);
|
||
let (di, bn) = (core.phys, core.right);
|
||
let new_rows = rows * di;
|
||
let mut tmp = vec![0i128; new_rows * bn];
|
||
for row in 0..rows {
|
||
for s in 0..di {
|
||
let out_base = (row * di + s) * bn;
|
||
for r in 0..bond {
|
||
let a = acc[row * bond + r];
|
||
if a == 0 {
|
||
continue;
|
||
}
|
||
let tbase = (r * di + s) * bn;
|
||
for rp in 0..bn {
|
||
let t = core.trits[tbase + rp] as i128;
|
||
if t != 0 {
|
||
tmp[out_base + rp] += a * t;
|
||
}
|
||
}
|
||
}
|
||
}
|
||
}
|
||
for row in 0..new_rows {
|
||
for rp in 0..bn {
|
||
let sc = core.scale_fx(rp) as i128;
|
||
let v = &mut tmp[row * bn + rp];
|
||
*v = (*v * sc) >> SCALE_FRAC;
|
||
}
|
||
}
|
||
acc = tmp;
|
||
rows = new_rows;
|
||
bond = bn;
|
||
}
|
||
debug_assert_eq!(bond, 1);
|
||
debug_assert_eq!(rows, self.n_signal);
|
||
|
||
let half = 1i128 << (SCALE_FRAC - 1);
|
||
acc.iter()
|
||
.map(|&v| {
|
||
let r = if v >= 0 {
|
||
(v + half) >> SCALE_FRAC
|
||
} else {
|
||
-((-v + half) >> SCALE_FRAC)
|
||
};
|
||
r.clamp(i64::MIN as i128, i64::MAX as i128) as i64
|
||
})
|
||
.collect()
|
||
}
|
||
}
|
||
|
||
// ---------------------------------------------------------------------------
|
||
// Residual stack (RVQ over ternary-MPS) — the fidelity knob.
|
||
// ---------------------------------------------------------------------------
|
||
|
||
/// A stack of ternary-MPS layers whose decodes **sum** to the reconstruction.
|
||
/// Each layer encodes the residual left by the layers before it, so fidelity
|
||
/// climbs with depth. A sum of byte-exact deterministic decodes is itself
|
||
/// byte-exact deterministic — the receiptable property survives.
|
||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||
pub struct TernaryStack {
|
||
pub layers: Vec<TernaryMps>,
|
||
pub n_signal: usize,
|
||
}
|
||
|
||
impl TernaryStack {
|
||
/// Sum every layer's decode. Pure integer, deterministic.
|
||
pub fn decode(&self) -> Vec<i64> {
|
||
let mut out = vec![0i64; self.n_signal];
|
||
for l in &self.layers {
|
||
for (o, v) in out.iter_mut().zip(l.decode()) {
|
||
*o += v;
|
||
}
|
||
}
|
||
out
|
||
}
|
||
|
||
/// Total on-wire size: the sum of the layers' payloads.
|
||
pub fn payload_len(&self) -> usize {
|
||
self.layers.iter().map(|l| l.payload_len()).sum()
|
||
}
|
||
|
||
pub fn reconstruction_hash(&self) -> [u8; 32] {
|
||
hash_signal(&self.decode(), self.n_signal)
|
||
}
|
||
}
|
||
|
||
/// Encode `signal` as up to `n_layers` residual ternary-MPS layers. Stops early
|
||
/// once the residual is exactly zero (nothing left to code).
|
||
pub fn encode_residual(
|
||
signal: &[i64],
|
||
n_sites: u8,
|
||
phys: usize,
|
||
max_bond: usize,
|
||
n_layers: usize,
|
||
) -> TernaryStack {
|
||
let mut residual = signal.to_vec();
|
||
let mut layers = Vec::with_capacity(n_layers);
|
||
for _ in 0..n_layers.max(1) {
|
||
let mps = encode(&residual, n_sites, phys, max_bond);
|
||
let recon = mps.decode();
|
||
for (r, c) in residual.iter_mut().zip(&recon) {
|
||
*r -= c;
|
||
}
|
||
layers.push(mps);
|
||
if residual.iter().all(|&x| x == 0) {
|
||
break;
|
||
}
|
||
}
|
||
TernaryStack { layers, n_signal: signal.len() }
|
||
}
|
||
|
||
// ---------------------------------------------------------------------------
|
||
// Encode — float (SVD + optimal per-column scale), encoder-side only.
|
||
// ---------------------------------------------------------------------------
|
||
|
||
/// Encode a signal of length `phys^n_sites` into a single ternary-MPS layer with
|
||
/// bond dimension capped at `max_bond`. Float SVD + per-column ternary
|
||
/// quantization run here, at encode, once; the object decodes with no float.
|
||
/// `phys` is the local dimension per site — 2 is the natural binary reshape;
|
||
/// **3 is the qutrit/base-3 reshape** (a ternary *site* dimension, distinct from
|
||
/// the ternary *weights*), for signals whose length is a power of three.
|
||
pub fn encode(signal: &[i64], n_sites: u8, phys: usize, max_bond: usize) -> TernaryMps {
|
||
let n = signal.len();
|
||
assert_eq!(n, phys.pow(n_sites as u32), "signal length must be phys^n_sites");
|
||
assert!(phys >= 2 && max_bond >= 1);
|
||
|
||
let mut remainder: Vec<f64> = signal.iter().map(|&s| s as f64).collect();
|
||
let mut left = 1usize;
|
||
let mut cols = n;
|
||
let mut cores: Vec<TernaryCore> = Vec::with_capacity(n_sites as usize);
|
||
|
||
for site in 0..n_sites as usize {
|
||
let last = site == n_sites as usize - 1;
|
||
let m_rows = left * phys;
|
||
let m_cols = cols / phys;
|
||
|
||
if last {
|
||
debug_assert_eq!(m_cols, 1);
|
||
let (trits, mant, exp) = quantize_cols(&remainder, m_rows, 1);
|
||
cores.push(TernaryCore { left, phys, right: 1, trits, scale_mant: mant, scale_exp: exp });
|
||
break;
|
||
}
|
||
|
||
let (u, s, vt) = svd_thin(&remainder, m_rows, m_cols);
|
||
let k_full = s.len();
|
||
let chi = k_full.min(max_bond).max(1);
|
||
|
||
let mut core_f = vec![0f64; m_rows * chi];
|
||
for i in 0..m_rows {
|
||
for j in 0..chi {
|
||
core_f[i * chi + j] = if j < k_full { u[i * k_full + j] } else { 0.0 };
|
||
}
|
||
}
|
||
let (trits, mant, exp) = quantize_cols(&core_f, m_rows, chi);
|
||
cores.push(TernaryCore { left, phys, right: chi, trits, scale_mant: mant, scale_exp: exp });
|
||
|
||
let mut next = vec![0f64; chi * m_cols];
|
||
for j in 0..chi {
|
||
let sj = if j < k_full { s[j] } else { 0.0 };
|
||
for c in 0..m_cols {
|
||
next[j * m_cols + c] = sj * if j < k_full { vt[j * m_cols + c] } else { 0.0 };
|
||
}
|
||
}
|
||
remainder = next;
|
||
left = chi;
|
||
cols = m_cols;
|
||
}
|
||
|
||
TernaryMps { n_sites, cores, n_signal: n }
|
||
}
|
||
|
||
/// Ternary-quantize a `rows × cols` float matrix **column by column** (each
|
||
/// column gets its own L2-optimal ternary scale), then pack the per-column
|
||
/// scales into a shared exponent + `u16` mantissa. Returns row-major trits, the
|
||
/// mantissas, and the exponent. The scale quantization is deliberately part of
|
||
/// the codec (decode reads the *quantized* scale), so the measured error is
|
||
/// honest end-to-end.
|
||
fn quantize_cols(mat: &[f64], rows: usize, cols: usize) -> (Vec<i8>, Vec<u16>, u8) {
|
||
let mut trits = vec![0i8; rows * cols];
|
||
let mut scale_fx = vec![1i64; cols];
|
||
for j in 0..cols {
|
||
let absmax = (0..rows).fold(0f64, |m, i| m.max(mat[i * cols + j].abs()));
|
||
if absmax == 0.0 {
|
||
continue;
|
||
}
|
||
let mut best_s = absmax;
|
||
let mut best_err = f64::INFINITY;
|
||
for step in 1..=64 {
|
||
let s = absmax * (step as f64) / 64.0;
|
||
let mut err = 0.0;
|
||
for i in 0..rows {
|
||
let v = mat[i * cols + j];
|
||
let t = (v / s).round().clamp(-1.0, 1.0);
|
||
let d = v - t * s;
|
||
err += d * d;
|
||
}
|
||
if err < best_err {
|
||
best_err = err;
|
||
best_s = s;
|
||
}
|
||
}
|
||
for i in 0..rows {
|
||
trits[i * cols + j] = (mat[i * cols + j] / best_s).round().clamp(-1.0, 1.0) as i8;
|
||
}
|
||
scale_fx[j] = ((best_s * (1i64 << SCALE_FRAC) as f64).round() as i64).max(1);
|
||
}
|
||
// Shared exponent so the largest scale's mantissa fits in u16.
|
||
let maxs = scale_fx.iter().copied().max().unwrap_or(1).max(1);
|
||
let exp = if maxs < (1i64 << 16) {
|
||
0u8
|
||
} else {
|
||
(64 - maxs.leading_zeros() - 16) as u8
|
||
};
|
||
let mant: Vec<u16> = scale_fx
|
||
.iter()
|
||
.map(|&s| (s >> exp).clamp(1, u16::MAX as i64) as u16)
|
||
.collect();
|
||
(trits, mant, exp)
|
||
}
|
||
|
||
// ---------------------------------------------------------------------------
|
||
// Thin SVD — one-sided Jacobi (small matrices; encoder-side float).
|
||
// ---------------------------------------------------------------------------
|
||
|
||
/// Thin SVD `A(m×n) = U(m×k) · diag(s)(k) · Vt(k×n)`, `k = min(m,n)`, singular
|
||
/// values descending. One-sided Jacobi on the tall orientation; small matrices
|
||
/// only (MPS cores) — plenty accurate, no float dependency.
|
||
fn svd_thin(a: &[f64], m: usize, n: usize) -> (Vec<f64>, Vec<f64>, Vec<f64>) {
|
||
if m >= n {
|
||
one_sided_jacobi(a, m, n)
|
||
} else {
|
||
let at = transpose(a, m, n);
|
||
let (u2, s, vt2) = one_sided_jacobi(&at, n, m); // (n×k),(k),(k×m)
|
||
let k = s.len();
|
||
(transpose(&vt2, k, m), s, transpose(&u2, n, k))
|
||
}
|
||
}
|
||
|
||
fn transpose(a: &[f64], m: usize, n: usize) -> Vec<f64> {
|
||
let mut t = vec![0f64; m * n];
|
||
for i in 0..m {
|
||
for j in 0..n {
|
||
t[j * m + i] = a[i * n + j];
|
||
}
|
||
}
|
||
t
|
||
}
|
||
|
||
fn one_sided_jacobi(a: &[f64], m: usize, n: usize) -> (Vec<f64>, Vec<f64>, Vec<f64>) {
|
||
let mut w = a.to_vec(); // m×n → columns become U·Σ
|
||
let mut v = vec![0f64; n * n];
|
||
for i in 0..n {
|
||
v[i * n + i] = 1.0;
|
||
}
|
||
let eps = 1e-13;
|
||
for _sweep in 0..60 {
|
||
let mut off = 0.0f64;
|
||
for p in 0..n {
|
||
for q in (p + 1)..n {
|
||
let (mut app, mut aqq, mut apq) = (0.0, 0.0, 0.0);
|
||
for i in 0..m {
|
||
let wp = w[i * n + p];
|
||
let wq = w[i * n + q];
|
||
app += wp * wp;
|
||
aqq += wq * wq;
|
||
apq += wp * wq;
|
||
}
|
||
off += apq.abs();
|
||
if apq.abs() <= eps * (app.sqrt() * aqq.sqrt()).max(f64::MIN_POSITIVE) {
|
||
continue;
|
||
}
|
||
let tau = (aqq - app) / (2.0 * apq);
|
||
let t = tau.signum() / (tau.abs() + (1.0 + tau * tau).sqrt());
|
||
let c = 1.0 / (1.0 + t * t).sqrt();
|
||
let s = c * t;
|
||
for i in 0..m {
|
||
let wp = w[i * n + p];
|
||
let wq = w[i * n + q];
|
||
w[i * n + p] = c * wp - s * wq;
|
||
w[i * n + q] = s * wp + c * wq;
|
||
}
|
||
for i in 0..n {
|
||
let vp = v[i * n + p];
|
||
let vq = v[i * n + q];
|
||
v[i * n + p] = c * vp - s * vq;
|
||
v[i * n + q] = s * vp + c * vq;
|
||
}
|
||
}
|
||
}
|
||
if off <= eps {
|
||
break;
|
||
}
|
||
}
|
||
let mut sigma: Vec<(f64, usize)> = (0..n)
|
||
.map(|j| {
|
||
let nrm: f64 = (0..m).map(|i| w[i * n + j] * w[i * n + j]).sum();
|
||
(nrm.sqrt(), j)
|
||
})
|
||
.collect();
|
||
sigma.sort_by(|a, b| b.0.partial_cmp(&a.0).unwrap());
|
||
let k = n;
|
||
let mut u = vec![0f64; m * k];
|
||
let mut s_out = vec![0f64; k];
|
||
let mut vt = vec![0f64; k * n];
|
||
for (rank, &(sig, j)) in sigma.iter().enumerate() {
|
||
s_out[rank] = sig;
|
||
let inv = if sig > 1e-300 { 1.0 / sig } else { 0.0 };
|
||
for i in 0..m {
|
||
u[i * k + rank] = w[i * n + j] * inv;
|
||
}
|
||
for i in 0..n {
|
||
vt[rank * n + i] = v[i * n + j];
|
||
}
|
||
}
|
||
(u, s_out, vt)
|
||
}
|
||
|
||
// ---------------------------------------------------------------------------
|
||
// Honest error reporting.
|
||
// ---------------------------------------------------------------------------
|
||
|
||
/// Reconstruction SNR in dB of `recon` vs `orig` (higher = closer). The stated,
|
||
/// never-hidden loss of this lossy codec. `+inf` on an exact match.
|
||
pub fn snr_db(orig: &[i64], recon: &[i64]) -> f64 {
|
||
let mut sig = 0.0f64;
|
||
let mut err = 0.0f64;
|
||
for (&o, &r) in orig.iter().zip(recon) {
|
||
sig += (o as f64) * (o as f64);
|
||
let d = (o - r) as f64;
|
||
err += d * d;
|
||
}
|
||
if err == 0.0 {
|
||
return f64::INFINITY;
|
||
}
|
||
10.0 * (sig / err).log10()
|
||
}
|
||
|
||
// ---------------------------------------------------------------------------
|
||
// Entropy coding of the trit stream (rANS) — the density knob.
|
||
// ---------------------------------------------------------------------------
|
||
//
|
||
// The five-trits-per-byte packing is entropy-optimal only for *uniform* trits
|
||
// (1.585 bits each). Ternary-quantized MPS cores are somewhat **sparse** (~54%
|
||
// zeros on a realistic stream), so the true entropy is lower and rANS can shave
|
||
// the trit stream. MEASURED VERDICT: the gain is marginal — a 54/23/23 trit
|
||
// distribution has ≈1.45 bits/trit vs the packing's 1.6, ~9% headroom, most of
|
||
// which the header + rANS state overhead eats at these block sizes (it is even a
|
||
// small *loss* at low layer counts). Base-3 packing is already near-optimal;
|
||
// this path is kept for completeness and the honest measurement, not because it
|
||
// pays. The distribution is modeled once and stored as a tiny header, so decode
|
||
// rebuilds the exact CDF and recovers the trits losslessly — the reconstruction
|
||
// stays byte-exact and determinism is preserved either way.
|
||
|
||
const RANS_PRECISION: i32 = 16; // must match rans.rs PRECISION
|
||
|
||
/// A 3-symbol (`{-1,0,+1}` → `{0,1,2}`) CDF from the trit histogram, in rANS's
|
||
/// `n+2`-with-escape convention (freqs sum to `2^PRECISION − 1`, each real
|
||
/// symbol ≥ 1). Returns the CDF and the three stored freqs.
|
||
fn build_trit_cdf(trits: &[i8]) -> (Vec<i32>, [u32; 3]) {
|
||
let total = 1i32 << RANS_PRECISION;
|
||
let mut cnt = [0u64; 3];
|
||
for &t in trits {
|
||
cnt[(t.clamp(-1, 1) + 1) as usize] += 1;
|
||
}
|
||
let denom = cnt.iter().sum::<u64>().max(1) as f64;
|
||
let mut f = [0i32; 3];
|
||
for (k, &c) in cnt.iter().enumerate() {
|
||
f[k] = (((c as f64 / denom) * (total - 1) as f64).round() as i32).max(1);
|
||
}
|
||
let sum: i32 = f.iter().sum();
|
||
let amax = (0..3).max_by_key(|&i| f[i]).unwrap();
|
||
f[amax] += (total - 1) - sum; // real symbols now sum to total-1
|
||
let mut cdf = vec![0i32; 5]; // 3 real + escape + total
|
||
cdf[1] = f[0];
|
||
cdf[2] = f[0] + f[1];
|
||
cdf[3] = f[0] + f[1] + f[2]; // = total-1
|
||
cdf[4] = total; // escape freq 1
|
||
(cdf, [f[0] as u32, f[1] as u32, f[2] as u32])
|
||
}
|
||
|
||
/// rANS-encode a trit stream into a self-contained blob: `3×u16` freq header +
|
||
/// `u32` count + bitstream. Losslessly reversible by [`rans_decode_trits`].
|
||
pub fn rans_encode_trits(trits: &[i8]) -> Vec<u8> {
|
||
let (cdf, freqs) = build_trit_cdf(trits);
|
||
let mut enc = crate::rans::RansEncoder::new();
|
||
for &t in trits {
|
||
let v = (t.clamp(-1, 1) + 1) as i32;
|
||
let coded = enc.encode(v, &cdf, 5, 0);
|
||
debug_assert!(coded, "a trit is always in the table's span");
|
||
}
|
||
let body = enc.flush();
|
||
let mut out = Vec::with_capacity(10 + body.len());
|
||
for &fr in &freqs {
|
||
out.extend_from_slice(&(fr as u16).to_le_bytes());
|
||
}
|
||
out.extend_from_slice(&(trits.len() as u32).to_le_bytes());
|
||
out.extend_from_slice(&body);
|
||
out
|
||
}
|
||
|
||
/// Inverse of [`rans_encode_trits`] — recovers the exact trit stream.
|
||
pub fn rans_decode_trits(bytes: &[u8]) -> Vec<i8> {
|
||
let f0 = u16::from_le_bytes([bytes[0], bytes[1]]) as i32;
|
||
let f1 = u16::from_le_bytes([bytes[2], bytes[3]]) as i32;
|
||
let f2 = u16::from_le_bytes([bytes[4], bytes[5]]) as i32;
|
||
let n = u32::from_le_bytes([bytes[6], bytes[7], bytes[8], bytes[9]]) as usize;
|
||
let total = 1i32 << RANS_PRECISION;
|
||
let cdf = vec![0, f0, f0 + f1, f0 + f1 + f2, total];
|
||
let mut dec = crate::rans::RansDecoder::new(&bytes[10..]).unwrap();
|
||
(0..n).map(|_| (dec.decode(&cdf, 5, 0) - 1) as i8).collect()
|
||
}
|
||
|
||
impl TernaryStack {
|
||
/// Total ternary multiply-adds the decode performs — the portable,
|
||
/// recomputable work figure (Σ over layers and cores of the trit-tensor
|
||
/// size). A receipt's `work` field; a sink cannot inflate it.
|
||
pub fn trit_macs(&self) -> u64 {
|
||
self.layers
|
||
.iter()
|
||
.flat_map(|l| l.cores.iter())
|
||
.map(|c| c.trits.len() as u64)
|
||
.sum()
|
||
}
|
||
|
||
/// Concatenate every core's trits across every layer (encode order).
|
||
pub fn gather_trits(&self) -> Vec<i8> {
|
||
let mut v = Vec::new();
|
||
for l in &self.layers {
|
||
for c in &l.cores {
|
||
v.extend_from_slice(&c.trits);
|
||
}
|
||
}
|
||
v
|
||
}
|
||
|
||
/// On-wire size in bytes with the trit stream **rANS-entropy-coded** instead
|
||
/// of raw base-3 packed: `rANS(trits) + Σ scale tables`. The honest density
|
||
/// figure once sparsity is exploited.
|
||
pub fn entropy_payload_len(&self) -> usize {
|
||
let trits = self.gather_trits();
|
||
let ent = rans_encode_trits(&trits).len();
|
||
let scales: usize = self
|
||
.layers
|
||
.iter()
|
||
.flat_map(|l| l.cores.iter())
|
||
.map(|c| c.scale_bytes())
|
||
.sum();
|
||
ent + scales
|
||
}
|
||
}
|
||
|
||
// ---------------------------------------------------------------------------
|
||
// WQTT container — the canonical wire payload for `wai.ternary.tt`.
|
||
// ---------------------------------------------------------------------------
|
||
|
||
const WQTT_MAGIC: &[u8; 4] = b"WQTT";
|
||
|
||
/// A malformed `WQTT` container.
|
||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||
pub struct TernaryError(pub String);
|
||
|
||
impl std::fmt::Display for TernaryError {
|
||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||
write!(f, "malformed WQTT: {}", self.0)
|
||
}
|
||
}
|
||
impl std::error::Error for TernaryError {}
|
||
|
||
impl TernaryStack {
|
||
/// Serialize to a `WQTT` container — the canonical, self-describing wire
|
||
/// bytes: magic, `n_signal`, then each layer's cores (`left/phys/right`,
|
||
/// shared scale exponent, per-column `u16` mantissas, base-3-packed trits).
|
||
/// This byte string is the object a receipt's identity hashes.
|
||
pub fn to_wqtt(&self) -> Vec<u8> {
|
||
let mut o = Vec::new();
|
||
o.extend_from_slice(WQTT_MAGIC);
|
||
o.push(1u8); // version
|
||
o.extend_from_slice(&(self.n_signal as u32).to_le_bytes());
|
||
o.push(self.layers.len() as u8);
|
||
for l in &self.layers {
|
||
o.push(l.n_sites);
|
||
o.push(l.cores.len() as u8);
|
||
for c in &l.cores {
|
||
o.extend_from_slice(&(c.left as u16).to_le_bytes());
|
||
o.extend_from_slice(&(c.phys as u16).to_le_bytes());
|
||
o.extend_from_slice(&(c.right as u16).to_le_bytes());
|
||
o.push(c.scale_exp);
|
||
for &m in &c.scale_mant {
|
||
o.extend_from_slice(&m.to_le_bytes());
|
||
}
|
||
o.extend_from_slice(&pack_trits(&c.trits));
|
||
}
|
||
}
|
||
o
|
||
}
|
||
}
|
||
|
||
/// Parse a `WQTT` container back into a [`TernaryStack`]. Bounds-checked; a
|
||
/// truncated or mis-magicked buffer is a [`TernaryError`], never a panic.
|
||
pub fn from_wqtt(bytes: &[u8]) -> Result<TernaryStack, TernaryError> {
|
||
let err = |m: &str| TernaryError(m.to_string());
|
||
let mut p = 0usize;
|
||
let need = |p: usize, n: usize| -> Result<(), TernaryError> {
|
||
if p + n > bytes.len() { Err(TernaryError("truncated".into())) } else { Ok(()) }
|
||
};
|
||
need(p, 5)?;
|
||
if &bytes[0..4] != WQTT_MAGIC {
|
||
return Err(err("bad magic"));
|
||
}
|
||
p = 4;
|
||
let _ver = bytes[p];
|
||
p += 1;
|
||
need(p, 5)?;
|
||
let n_signal = u32::from_le_bytes([bytes[p], bytes[p + 1], bytes[p + 2], bytes[p + 3]]) as usize;
|
||
p += 4;
|
||
let n_layers = bytes[p] as usize;
|
||
p += 1;
|
||
|
||
let mut layers = Vec::with_capacity(n_layers);
|
||
for _ in 0..n_layers {
|
||
need(p, 2)?;
|
||
let n_sites = bytes[p];
|
||
let n_cores = bytes[p + 1] as usize;
|
||
p += 2;
|
||
let mut cores = Vec::with_capacity(n_cores);
|
||
for _ in 0..n_cores {
|
||
need(p, 7)?;
|
||
let left = u16::from_le_bytes([bytes[p], bytes[p + 1]]) as usize;
|
||
let phys = u16::from_le_bytes([bytes[p + 2], bytes[p + 3]]) as usize;
|
||
let right = u16::from_le_bytes([bytes[p + 4], bytes[p + 5]]) as usize;
|
||
let scale_exp = bytes[p + 6];
|
||
p += 7;
|
||
need(p, right * 2)?;
|
||
let mut scale_mant = Vec::with_capacity(right);
|
||
for _ in 0..right {
|
||
scale_mant.push(u16::from_le_bytes([bytes[p], bytes[p + 1]]));
|
||
p += 2;
|
||
}
|
||
let n_trits = left * phys * right;
|
||
let n_bytes = n_trits.div_ceil(TRITS_PER_BYTE);
|
||
need(p, n_bytes)?;
|
||
let trits = unpack_trits(&bytes[p..p + n_bytes], n_trits);
|
||
p += n_bytes;
|
||
cores.push(TernaryCore { left, phys, right, trits, scale_mant, scale_exp });
|
||
}
|
||
layers.push(TernaryMps { n_sites, cores, n_signal });
|
||
}
|
||
Ok(TernaryStack { layers, n_signal })
|
||
}
|
||
|
||
/// A real (not cosmetic) use of the shared fixed-point idiom: an integer matmul
|
||
/// in the same discipline, available to callers composing ternary cores with the
|
||
/// transformer stack.
|
||
#[doc(hidden)]
|
||
pub fn _fx_matmul(a: &[i64], b: &[i64], m: usize, k: usize, n: usize, shift: u32) -> Vec<i64> {
|
||
int_transform::matmul(a, b, m, k, n, shift)
|
||
}
|
||
|
||
#[cfg(test)]
|
||
mod tests {
|
||
use super::*;
|
||
|
||
fn sinusoid(n: usize) -> Vec<i64> {
|
||
(0..n)
|
||
.map(|i| {
|
||
let x = i as f64 / n as f64;
|
||
(1000.0 * (2.0 * std::f64::consts::PI * x).sin()).round() as i64
|
||
})
|
||
.collect()
|
||
}
|
||
|
||
/// A realistic-ish stream: three tones of differing amplitude/frequency plus
|
||
/// deterministic low-level noise — closer to an RF/telemetry block than a
|
||
/// single sinusoid. Deterministic (index-hashed noise), no RNG.
|
||
fn realistic(n: usize) -> Vec<i64> {
|
||
(0..n)
|
||
.map(|i| {
|
||
let x = i as f64 / n as f64;
|
||
let tau = 2.0 * std::f64::consts::PI;
|
||
let s = 900.0 * (tau * 3.0 * x).sin()
|
||
+ 400.0 * (tau * 11.0 * x).cos()
|
||
+ 180.0 * (tau * 27.0 * x).sin();
|
||
// deterministic pseudo-noise in ~[-40,40]
|
||
let h = (i as u64).wrapping_mul(2654435761) ^ 0x9E3779B97F4A7C15;
|
||
let noise = ((h >> 40) & 0x7f) as f64 - 40.0;
|
||
(s + noise).round() as i64
|
||
})
|
||
.collect()
|
||
}
|
||
|
||
#[test]
|
||
fn wqtt_container_round_trips() {
|
||
let sig = realistic(256);
|
||
let stack = encode_residual(&sig, 8, 2, 4, 3);
|
||
let bytes = stack.to_wqtt();
|
||
let back = from_wqtt(&bytes).expect("parse WQTT");
|
||
assert_eq!(back, stack, "WQTT round-trip must be exact");
|
||
assert_eq!(back.reconstruction_hash(), stack.reconstruction_hash());
|
||
// truncation and bad magic are errors, not panics
|
||
assert!(from_wqtt(&bytes[..bytes.len() - 3]).is_err());
|
||
assert!(from_wqtt(b"XXXX....").is_err());
|
||
}
|
||
|
||
#[test]
|
||
fn rans_trits_round_trips_losslessly() {
|
||
// Determinism survives entropy coding iff this is exact.
|
||
let sig = realistic(256);
|
||
let stack = encode_residual(&sig, 8, 2, 4, 3);
|
||
let trits = stack.gather_trits();
|
||
let blob = rans_encode_trits(&trits);
|
||
assert_eq!(rans_decode_trits(&blob), trits, "rANS trit round-trip must be lossless");
|
||
}
|
||
|
||
#[test]
|
||
fn pack_round_trips_all_patterns() {
|
||
let trits: Vec<i8> = (0..37).map(|i| (i % 3) as i8 - 1).collect();
|
||
let packed = pack_trits(&trits);
|
||
assert_eq!(packed.len(), 37usize.div_ceil(TRITS_PER_BYTE));
|
||
assert_eq!(unpack_trits(&packed, trits.len()), trits);
|
||
}
|
||
|
||
#[test]
|
||
fn decode_is_deterministic_byte_for_byte() {
|
||
let sig: Vec<i64> = (0..64).map(|i| ((i * 7 + 3) % 29) as i64 - 14).collect();
|
||
let mps = encode(&sig, 6, 2, 4);
|
||
assert_eq!(mps.reconstruction_hash(), mps.clone().reconstruction_hash());
|
||
assert_eq!(mps.decode(), mps.decode());
|
||
}
|
||
|
||
#[test]
|
||
fn residual_stack_is_deterministic() {
|
||
let sig = sinusoid(64);
|
||
let stack = encode_residual(&sig, 6, 2, 4, 4);
|
||
assert_eq!(stack.reconstruction_hash(), stack.clone().reconstruction_hash());
|
||
assert_eq!(stack.decode(), stack.decode());
|
||
}
|
||
|
||
#[test]
|
||
fn rate_distortion_curve_improves_with_layers() {
|
||
// The measurement the whole exercise is for: SNR vs on-wire bytes as
|
||
// residual layers accumulate. Honest end-to-end (decode reads quantized
|
||
// scales). Prints the curve; asserts fidelity is monotone-ish and that
|
||
// depth meaningfully beats a single layer.
|
||
let n = 64usize;
|
||
let sig = sinusoid(n);
|
||
let raw = n * 8;
|
||
println!("ternary-tt rate–distortion (n={n}, raw={raw}B, bond≤4):");
|
||
let mut prev = f64::NEG_INFINITY;
|
||
let mut snr_at = [0f64; 5];
|
||
for k in 1..=4 {
|
||
let stack = encode_residual(&sig, 6, 2, 4, k);
|
||
let snr = snr_db(&sig, &stack.decode());
|
||
let bytes = stack.payload_len();
|
||
println!(
|
||
" K={k}: snr={snr:5.1}dB payload={bytes:3}B ratio={:.1}×",
|
||
raw as f64 / bytes as f64
|
||
);
|
||
snr_at[k] = snr;
|
||
assert!(snr + 0.01 >= prev, "SNR must not regress with more layers");
|
||
prev = snr;
|
||
}
|
||
// Residual coding should buy real fidelity over a single layer.
|
||
assert!(
|
||
snr_at[4] > snr_at[1] + 3.0,
|
||
"4 layers should beat 1 by >3 dB (got {:.1} vs {:.1})",
|
||
snr_at[4],
|
||
snr_at[1]
|
||
);
|
||
}
|
||
|
||
#[test]
|
||
fn entropy_coded_rate_distortion_on_realistic_stream() {
|
||
// The decisive measurement: a 256-sample multi-tone+noise stream, RD with
|
||
// the trit stream rANS-entropy-coded vs raw base-3 packed. Reports trit
|
||
// sparsity (why rANS helps) and both payload sizes. Honest end-to-end.
|
||
let n = 256usize;
|
||
let sig = realistic(n);
|
||
let raw = n * 8;
|
||
// trit sparsity over a mid-depth stack
|
||
let probe = encode_residual(&sig, 8, 2, 4, 4);
|
||
let trits = probe.gather_trits();
|
||
let zeros = trits.iter().filter(|&&t| t == 0).count();
|
||
println!(
|
||
"ternary-tt entropy RD (n={n}, raw={raw}B, bond≤4): trit zeros {:.0}% of {}",
|
||
100.0 * zeros as f64 / trits.len().max(1) as f64,
|
||
trits.len()
|
||
);
|
||
let mut snr1 = 0.0;
|
||
let mut snrk = 0.0;
|
||
for k in 1..=6 {
|
||
let stack = encode_residual(&sig, 8, 2, 4, k);
|
||
let snr = snr_db(&sig, &stack.decode());
|
||
let raw_pl = stack.payload_len();
|
||
let ent_pl = stack.entropy_payload_len();
|
||
println!(
|
||
" K={k}: snr={snr:5.1}dB raw={raw_pl:4}B ({:.1}×) entropy={ent_pl:4}B ({:.1}×)",
|
||
raw as f64 / raw_pl as f64,
|
||
raw as f64 / ent_pl as f64,
|
||
);
|
||
if k == 1 {
|
||
snr1 = snr;
|
||
}
|
||
snrk = snr;
|
||
}
|
||
assert!(snrk > snr1 + 3.0, "depth should add fidelity");
|
||
// Honest claim: entropy coding is roughly BREAK-EVEN vs base-3 packing
|
||
// (marginal win at depth, small loss at low rate) — within 10% either way,
|
||
// never a blow-up. Not "rANS always wins" — the measurement says it barely
|
||
// moves, which is the finding.
|
||
let stack = encode_residual(&sig, 8, 2, 4, 4);
|
||
assert!(
|
||
stack.entropy_payload_len() <= stack.payload_len() * 11 / 10,
|
||
"rANS trits stay within 10% of raw-packed"
|
||
);
|
||
}
|
||
|
||
#[test]
|
||
fn qutrit_reshape_base3_decodes_deterministically() {
|
||
let sig: Vec<i64> = (0..27).map(|i| ((i * 11) % 41) as i64 - 20).collect();
|
||
let stack = encode_residual(&sig, 3, 3, 3, 3); // phys=3: base-3 reshape
|
||
assert_eq!(stack.n_signal, 27);
|
||
assert_eq!(stack.decode(), stack.decode());
|
||
assert_eq!(stack.layers[0].cores.last().unwrap().right, 1);
|
||
}
|
||
|
||
#[test]
|
||
fn empty_and_zero_signal_are_total() {
|
||
let sig = vec![0i64; 8];
|
||
let mps = encode(&sig, 3, 2, 2);
|
||
assert!(mps.decode().iter().all(|&s| s == 0));
|
||
}
|
||
|
||
#[test]
|
||
fn scale_dequant_is_pure_integer_shift() {
|
||
let sig = sinusoid(64);
|
||
let mps = encode(&sig, 6, 2, 4);
|
||
for c in &mps.cores {
|
||
for j in 0..c.right {
|
||
assert_eq!(c.scale_fx(j), (c.scale_mant[j] as i64) << c.scale_exp);
|
||
}
|
||
}
|
||
}
|
||
}
|