mirror of
https://github.com/ruvnet/RuVector.git
synced 2026-05-25 06:36:37 +00:00
feat: add ruvector-dither crate and integrate thermorust+dither into exo
ruvector-dither (new crate):
- GoldenRatioDither: additive φ-sequence with best 1-D equidistribution
- PiDither: cyclic 256-entry π-byte table for deterministic weight dithering
- quantize_dithered / quantize_slice_dithered: drop-in pre-quantization offset
- quantize_to_code: integer-code variant for packed-weight use
- ChannelDither: per-channel pool seeded by (layer_id, channel_id) pairs
- DitherSource trait for generic dither composition
- 15 unit tests + 3 doctests; 4 Criterion benchmark groups
exo-backend-classical integration:
- ThermoLayer (thermo_layer.rs): Ising motif coherence gate using thermorust
- Runs Metropolis steps on clamped activations
- Returns ThermoSignal { lambda, magnetisation, dissipation_j, energy_after }
- λ-signal = −ΔE/|E₀|: positive means pattern is settling toward coherence
- DitheredQuantizer (dither_quantizer.rs): wraps ruvector-dither for exo tensors
- GoldenRatio or Pi kind, per-layer seeding, reset support
- Supports 3/5/7/8-bit quantization with ε-LSB dither amplitude
- 8 new unit tests across both modules; all 74 existing tests still pass
https://claude.ai/code/session_019Lt11HYsW1265X7jB7haoC
This commit is contained in:
parent
3b5048c84a
commit
e9230450d6
14 changed files with 949 additions and 2 deletions
7
Cargo.lock
generated
7
Cargo.lock
generated
|
|
@ -8374,6 +8374,13 @@ dependencies = [
|
|||
"wasm-bindgen-test",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "ruvector-dither"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"criterion 0.5.1",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "ruvector-domain-expansion"
|
||||
version = "2.0.5"
|
||||
|
|
|
|||
|
|
@ -112,6 +112,7 @@ members = [
|
|||
"examples/rvf-kernel-optimized",
|
||||
"examples/verified-applications",
|
||||
"crates/thermorust",
|
||||
"crates/ruvector-dither",
|
||||
]
|
||||
resolver = "2"
|
||||
|
||||
|
|
|
|||
28
crates/ruvector-dither/Cargo.toml
Normal file
28
crates/ruvector-dither/Cargo.toml
Normal file
|
|
@ -0,0 +1,28 @@
|
|||
[package]
|
||||
name = "ruvector-dither"
|
||||
version = "0.1.0"
|
||||
edition = "2021"
|
||||
license = "MIT OR Apache-2.0"
|
||||
authors = ["rUv <ruv@ruv.io>"]
|
||||
repository = "https://github.com/ruvnet/ruvector"
|
||||
homepage = "https://ruv.io"
|
||||
documentation = "https://docs.rs/ruvector-dither"
|
||||
description = "Deterministic low-discrepancy dithering for low-bit quantization: golden-ratio and π-digit sequences for blue-noise error shaping"
|
||||
keywords = ["quantization", "dither", "golden-ratio", "inference", "wasm"]
|
||||
categories = ["science", "algorithms", "no-std"]
|
||||
readme = "README.md"
|
||||
|
||||
[features]
|
||||
default = []
|
||||
# Enable no_std mode (requires an allocator)
|
||||
no_std = []
|
||||
|
||||
[dependencies]
|
||||
# No runtime deps — fully no_std compatible
|
||||
|
||||
[dev-dependencies]
|
||||
criterion = { version = "0.5", features = ["html_reports"] }
|
||||
|
||||
[[bench]]
|
||||
name = "dither_bench"
|
||||
harness = false
|
||||
61
crates/ruvector-dither/benches/dither_bench.rs
Normal file
61
crates/ruvector-dither/benches/dither_bench.rs
Normal file
|
|
@ -0,0 +1,61 @@
|
|||
use criterion::{black_box, criterion_group, criterion_main, BenchmarkId, Criterion};
|
||||
use ruvector_dither::{
|
||||
channel::ChannelDither, GoldenRatioDither, PiDither, quantize_dithered,
|
||||
quantize_slice_dithered,
|
||||
};
|
||||
|
||||
fn bench_single_quantize(c: &mut Criterion) {
|
||||
let mut group = c.benchmark_group("quantize_dithered_single");
|
||||
for bits in [5u32, 7, 8] {
|
||||
group.bench_with_input(BenchmarkId::from_parameter(bits), &bits, |b, &bits| {
|
||||
let mut d = GoldenRatioDither::new(0.0);
|
||||
b.iter(|| quantize_dithered(black_box(0.314_f32), bits, 0.5, &mut d));
|
||||
});
|
||||
}
|
||||
group.finish();
|
||||
}
|
||||
|
||||
fn bench_slice_quantize(c: &mut Criterion) {
|
||||
let mut group = c.benchmark_group("quantize_slice");
|
||||
for n in [64usize, 256, 1024] {
|
||||
group.bench_with_input(BenchmarkId::from_parameter(n), &n, |b, &n| {
|
||||
let input: Vec<f32> = (0..n).map(|i| (i as f32 / n as f32) * 2.0 - 1.0).collect();
|
||||
b.iter(|| {
|
||||
let mut buf = input.clone();
|
||||
let mut d = GoldenRatioDither::new(0.0);
|
||||
quantize_slice_dithered(black_box(&mut buf), 8, 0.5, &mut d);
|
||||
black_box(buf)
|
||||
});
|
||||
});
|
||||
}
|
||||
group.finish();
|
||||
}
|
||||
|
||||
fn bench_pi_dither(c: &mut Criterion) {
|
||||
c.bench_function("pi_dither_1k", |b| {
|
||||
let mut d = PiDither::new(0);
|
||||
let mut buf: Vec<f32> = vec![0.5; 1024];
|
||||
b.iter(|| {
|
||||
quantize_slice_dithered(black_box(&mut buf), 7, 0.5, &mut d);
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
fn bench_channel_dither(c: &mut Criterion) {
|
||||
c.bench_function("channel_dither_256activations_32ch", |b| {
|
||||
let mut cd = ChannelDither::new(0, 32, 8, 0.5);
|
||||
let mut acts: Vec<f32> = vec![0.314; 256];
|
||||
b.iter(|| {
|
||||
cd.quantize_batch(black_box(&mut acts));
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
criterion_group!(
|
||||
benches,
|
||||
bench_single_quantize,
|
||||
bench_slice_quantize,
|
||||
bench_pi_dither,
|
||||
bench_channel_dither
|
||||
);
|
||||
criterion_main!(benches);
|
||||
80
crates/ruvector-dither/src/channel.rs
Normal file
80
crates/ruvector-dither/src/channel.rs
Normal file
|
|
@ -0,0 +1,80 @@
|
|||
//! Per-channel and per-layer dither management.
|
||||
//!
|
||||
//! `ChannelDither` bundles one `GoldenRatioDither` state per channel,
|
||||
//! seeded from `(layer_id, channel_id)` pairs so every channel is
|
||||
//! structurally decorrelated without any RNG.
|
||||
|
||||
use crate::{DitherSource, GoldenRatioDither};
|
||||
|
||||
/// Per-channel dither pool seeded from `(layer_id, channel_id)` pairs.
|
||||
///
|
||||
/// Allocates one `GoldenRatioDither` per channel; each is independently
|
||||
/// advanced, so channels cannot constructively interfere.
|
||||
pub struct ChannelDither {
|
||||
channels: Vec<GoldenRatioDither>,
|
||||
bits: u32,
|
||||
eps: f32,
|
||||
}
|
||||
|
||||
impl ChannelDither {
|
||||
/// Build a pool of `n_channels` dithers for `layer_id` / `bits` / `eps`.
|
||||
pub fn new(layer_id: u32, n_channels: usize, bits: u32, eps: f32) -> Self {
|
||||
let channels = (0..n_channels)
|
||||
.map(|ch| GoldenRatioDither::from_ids(layer_id, ch as u32))
|
||||
.collect();
|
||||
Self { channels, bits, eps }
|
||||
}
|
||||
|
||||
/// Quantize `activations` in-place. Each column (channel dimension) uses
|
||||
/// its own independent dither state.
|
||||
///
|
||||
/// `activations` is a flat row-major tensor of shape `[batch, channels]`.
|
||||
/// If the slice is not a multiple of `n_channels`, the remainder is
|
||||
/// processed using channel 0.
|
||||
pub fn quantize_batch(&mut self, activations: &mut [f32]) {
|
||||
let nc = self.channels.len();
|
||||
let qmax = ((1u32 << (self.bits - 1)) - 1) as f32;
|
||||
let lsb = 1.0 / qmax;
|
||||
for (i, x) in activations.iter_mut().enumerate() {
|
||||
let ch = i % nc;
|
||||
let d = self.channels[ch].next(self.eps * lsb);
|
||||
*x = ((*x + d) * qmax).round().clamp(-qmax, qmax) / qmax;
|
||||
}
|
||||
}
|
||||
|
||||
/// Number of channels in this pool.
|
||||
pub fn n_channels(&self) -> usize {
|
||||
self.channels.len()
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn channel_dither_correct_count() {
|
||||
let cd = ChannelDither::new(0, 16, 8, 0.5);
|
||||
assert_eq!(cd.n_channels(), 16);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn channel_dither_in_bounds() {
|
||||
let mut cd = ChannelDither::new(1, 8, 5, 0.5);
|
||||
let mut acts: Vec<f32> = (0..64).map(|i| (i as f32 / 63.0) * 2.0 - 1.0).collect();
|
||||
cd.quantize_batch(&mut acts);
|
||||
for v in acts {
|
||||
assert!(v >= -1.0 && v <= 1.0, "out of bounds: {v}");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn different_layers_produce_different_outputs() {
|
||||
let input: Vec<f32> = vec![0.5; 16];
|
||||
let mut buf0 = input.clone();
|
||||
let mut buf1 = input.clone();
|
||||
ChannelDither::new(0, 8, 8, 0.5).quantize_batch(&mut buf0);
|
||||
ChannelDither::new(99, 8, 8, 0.5).quantize_batch(&mut buf1);
|
||||
assert_ne!(buf0, buf1, "different layer_ids must yield different dithered outputs");
|
||||
}
|
||||
}
|
||||
95
crates/ruvector-dither/src/golden.rs
Normal file
95
crates/ruvector-dither/src/golden.rs
Normal file
|
|
@ -0,0 +1,95 @@
|
|||
//! Golden-ratio quasi-random dither sequence.
|
||||
//!
|
||||
//! State update: `state = frac(state + φ)` where φ = (√5−1)/2 ≈ 0.618…
|
||||
//!
|
||||
//! This is the 1-D Halton sequence in base φ — it has the best possible
|
||||
//! equidistribution for a 1-D low-discrepancy sequence.
|
||||
|
||||
use crate::DitherSource;
|
||||
|
||||
/// Additive golden-ratio dither with zero-mean output in `[-0.5, 0.5]`.
|
||||
///
|
||||
/// The sequence has period 1 (irrational) so it never exactly repeats.
|
||||
/// Two instances with different seeds stay decorrelated.
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct GoldenRatioDither {
|
||||
state: f32,
|
||||
}
|
||||
|
||||
/// φ = (√5 − 1) / 2
|
||||
const PHI: f32 = 0.618_033_98_f32;
|
||||
|
||||
impl GoldenRatioDither {
|
||||
/// Create a new sequence seeded at `initial_state` ∈ [0, 1).
|
||||
///
|
||||
/// For per-layer / per-channel decorrelation, seed with
|
||||
/// `frac(layer_id × φ + channel_id × φ²)`.
|
||||
#[inline]
|
||||
pub fn new(initial_state: f32) -> Self {
|
||||
Self { state: initial_state.abs().fract() }
|
||||
}
|
||||
|
||||
/// Construct from a `(layer_id, channel_id)` pair for structural decorrelation.
|
||||
#[inline]
|
||||
pub fn from_ids(layer_id: u32, channel_id: u32) -> Self {
|
||||
let s = ((layer_id as f32) * PHI + (channel_id as f32) * PHI * PHI).fract();
|
||||
Self { state: s }
|
||||
}
|
||||
|
||||
/// Current state (useful for serialisation / checkpointing).
|
||||
#[inline]
|
||||
pub fn state(&self) -> f32 {
|
||||
self.state
|
||||
}
|
||||
}
|
||||
|
||||
impl DitherSource for GoldenRatioDither {
|
||||
/// Advance and return next value in `[-0.5, 0.5]`.
|
||||
#[inline]
|
||||
fn next_unit(&mut self) -> f32 {
|
||||
self.state = (self.state + PHI).fract();
|
||||
self.state - 0.5
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::DitherSource;
|
||||
|
||||
#[test]
|
||||
fn output_is_in_range() {
|
||||
let mut d = GoldenRatioDither::new(0.0);
|
||||
for _ in 0..10_000 {
|
||||
let v = d.next_unit();
|
||||
assert!(v >= -0.5 && v <= 0.5, "out of range: {v}");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn mean_is_near_zero() {
|
||||
let mut d = GoldenRatioDither::new(0.0);
|
||||
let n = 100_000;
|
||||
let mean: f32 = (0..n).map(|_| d.next_unit()).sum::<f32>() / n as f32;
|
||||
assert!(mean.abs() < 0.01, "mean too large: {mean}");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn from_ids_decorrelates() {
|
||||
let mut d0 = GoldenRatioDither::from_ids(0, 0);
|
||||
let mut d1 = GoldenRatioDither::from_ids(1, 7);
|
||||
// Confirm they start at different states
|
||||
let v0 = d0.next_unit();
|
||||
let v1 = d1.next_unit();
|
||||
assert!((v0 - v1).abs() > 1e-4, "distinct seeds should produce distinct first values");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn deterministic_across_calls() {
|
||||
let mut d1 = GoldenRatioDither::new(0.123);
|
||||
let mut d2 = GoldenRatioDither::new(0.123);
|
||||
for _ in 0..1000 {
|
||||
assert_eq!(d1.next_unit(), d2.next_unit());
|
||||
}
|
||||
}
|
||||
}
|
||||
63
crates/ruvector-dither/src/lib.rs
Normal file
63
crates/ruvector-dither/src/lib.rs
Normal file
|
|
@ -0,0 +1,63 @@
|
|||
//! # ruvector-dither
|
||||
//!
|
||||
//! Deterministic, low-discrepancy **pre-quantization dithering** for low-bit
|
||||
//! inference on tiny devices (WASM, Seed, STM32).
|
||||
//!
|
||||
//! ## Why dither?
|
||||
//!
|
||||
//! Quantizers at 3 / 5 / 7 bits can align with power-of-two boundaries and
|
||||
//! produce idle tones / limit cycles — sticky activations and periodic errors
|
||||
//! that degrade accuracy. A sub-LSB pre-quantization offset:
|
||||
//!
|
||||
//! - Decorrelates the signal from grid boundaries.
|
||||
//! - Pushes quantization error toward high frequencies (blue-noise-like),
|
||||
//! which average out downstream.
|
||||
//! - Uses **no RNG** — outputs are deterministic, reproducible across
|
||||
//! platforms (WASM / x86 / ARM), and cache-friendly.
|
||||
//!
|
||||
//! ## Sequences
|
||||
//!
|
||||
//! | Type | State update | Properties |
|
||||
//! |------|-------------|------------|
|
||||
//! | [`GoldenRatioDither`] | frac(state + φ) | Best 1-D equidistribution |
|
||||
//! | [`PiDither`] | table of π bytes | Reproducible, period = 256 |
|
||||
//!
|
||||
//! ## Quick start
|
||||
//!
|
||||
//! ```
|
||||
//! use ruvector_dither::{GoldenRatioDither, PiDither, quantize_dithered};
|
||||
//!
|
||||
//! // Quantize with golden-ratio dither, 8-bit, ε = 0.5 LSB
|
||||
//! let mut gr = GoldenRatioDither::new(0.0);
|
||||
//! let q = quantize_dithered(0.314, 8, 0.5, &mut gr);
|
||||
//! assert!(q >= -1.0 && q <= 1.0);
|
||||
//!
|
||||
//! // Quantize with π-digit dither
|
||||
//! let mut pi = PiDither::new(0);
|
||||
//! let q2 = quantize_dithered(0.271, 5, 0.5, &mut pi);
|
||||
//! assert!(q2 >= -1.0 && q2 <= 1.0);
|
||||
//! ```
|
||||
|
||||
#![cfg_attr(feature = "no_std", no_std)]
|
||||
|
||||
pub mod golden;
|
||||
pub mod pi;
|
||||
pub mod quantize;
|
||||
pub mod channel;
|
||||
|
||||
pub use golden::GoldenRatioDither;
|
||||
pub use pi::PiDither;
|
||||
pub use quantize::{quantize_dithered, quantize_slice_dithered};
|
||||
pub use channel::ChannelDither;
|
||||
|
||||
/// Trait implemented by any deterministic dither source.
|
||||
pub trait DitherSource {
|
||||
/// Advance the sequence and return the next zero-mean offset in `[-0.5, +0.5]`.
|
||||
fn next_unit(&mut self) -> f32;
|
||||
|
||||
/// Scale output to ε × LSB amplitude.
|
||||
#[inline]
|
||||
fn next(&mut self, eps_lsb: f32) -> f32 {
|
||||
self.next_unit() * eps_lsb
|
||||
}
|
||||
}
|
||||
106
crates/ruvector-dither/src/pi.rs
Normal file
106
crates/ruvector-dither/src/pi.rs
Normal file
|
|
@ -0,0 +1,106 @@
|
|||
//! π-digit dither: cyclic table of the first 256 digits of π scaled to [-0.5, 0.5].
|
||||
//!
|
||||
//! Period = 256. Each entry is an independent offset making the sequence
|
||||
//! suitable for small buffers where you want exact reproducibility from a
|
||||
//! named tensor / layer rather than a stateful RNG.
|
||||
|
||||
use crate::DitherSource;
|
||||
|
||||
/// First 256 bytes of π (hex digits 3.243F6A8885A308D3…).
|
||||
///
|
||||
/// Each byte spans [0, 255]; we map to [-0.5, 0.5] by `(b as f32 / 255.0) - 0.5`.
|
||||
#[rustfmt::skip]
|
||||
const PI_BYTES: [u8; 256] = [
|
||||
0x32, 0x43, 0xF6, 0xA8, 0x88, 0x5A, 0x30, 0x8D, 0x31, 0x31, 0x98, 0xA2,
|
||||
0xE0, 0x37, 0x07, 0x34, 0x4A, 0x40, 0x93, 0x82, 0x22, 0x99, 0xF3, 0x1D,
|
||||
0x00, 0x82, 0xEF, 0xA9, 0x8E, 0xC4, 0xE6, 0xC8, 0x94, 0x52, 0x21, 0xE6,
|
||||
0x38, 0xD0, 0x13, 0x77, 0xBE, 0x54, 0x66, 0xCF, 0x34, 0xE9, 0x0C, 0x6C,
|
||||
0xC0, 0xAC, 0x29, 0xB7, 0xC9, 0x7C, 0x50, 0xDD, 0x3F, 0x84, 0xD5, 0xB5,
|
||||
0xB5, 0x47, 0x09, 0x17, 0x92, 0x16, 0xD5, 0xD9, 0x89, 0x79, 0xFB, 0x1B,
|
||||
0xD1, 0x31, 0x0B, 0xA6, 0x98, 0xDF, 0xB5, 0xAC, 0x2F, 0xFD, 0x72, 0xDB,
|
||||
0xD0, 0x1A, 0xDF, 0xB7, 0xB8, 0xE1, 0xAF, 0xED, 0x6A, 0x26, 0x7E, 0x96,
|
||||
0xBA, 0x7C, 0x90, 0x45, 0xF1, 0x2C, 0x7F, 0x99, 0x24, 0xA1, 0x99, 0x47,
|
||||
0xB3, 0x91, 0x6C, 0xF7, 0x08, 0x01, 0xF2, 0xE2, 0x85, 0x8E, 0xFC, 0x16,
|
||||
0x63, 0x69, 0x20, 0xD8, 0x71, 0x57, 0x4E, 0x69, 0xA4, 0x58, 0xFE, 0xA3,
|
||||
0xF4, 0x93, 0x3D, 0x7E, 0x0D, 0x95, 0x74, 0x8F, 0x72, 0x8E, 0xB6, 0x58,
|
||||
0x71, 0x8B, 0xCD, 0x58, 0x82, 0x15, 0x4A, 0xEE, 0x7B, 0x54, 0xA4, 0x1D,
|
||||
0xC2, 0x5A, 0x59, 0xB5, 0x9C, 0x30, 0xD5, 0x39, 0x2A, 0xF2, 0x60, 0x13,
|
||||
0xC5, 0xD1, 0xB0, 0x23, 0x28, 0x60, 0x85, 0xF0, 0xCA, 0x41, 0x79, 0x18,
|
||||
0xB8, 0xDB, 0x38, 0xEF, 0x8E, 0x79, 0xDC, 0xB0, 0x60, 0x3A, 0x18, 0x0E,
|
||||
0x6C, 0x9E, 0xD0, 0xE8, 0x9D, 0x44, 0x8F, 0x39, 0xF9, 0x93, 0xDB, 0x07,
|
||||
0x3A, 0xA3, 0x45, 0x22, 0x7E, 0xD8, 0xAC, 0x87, 0x2F, 0x85, 0x5D, 0x28,
|
||||
0x55, 0xB0, 0x89, 0x73, 0x36, 0xF3, 0xEB, 0xCD, 0xF6, 0x00, 0x4A, 0xDB,
|
||||
0x36, 0x47, 0xDB, 0xF7, 0x82, 0x48, 0xDB, 0xF3, 0xD3, 0x7C, 0x45, 0x10,
|
||||
0xC6, 0x7A, 0x70, 0xAA, 0x56, 0x78, 0x5A, 0xC6, 0x37, 0x10, 0xA2, 0x44,
|
||||
0x32, 0x34, 0xFE, 0x08,
|
||||
];
|
||||
|
||||
/// Cyclic π-digit dither. Period = 256; index wraps with bitwise AND.
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct PiDither {
|
||||
idx: u8,
|
||||
}
|
||||
|
||||
impl PiDither {
|
||||
/// Create a new instance starting at `offset` (0–255).
|
||||
#[inline]
|
||||
pub fn new(offset: u8) -> Self {
|
||||
Self { idx: offset }
|
||||
}
|
||||
|
||||
/// Construct from a tensor/layer identifier for structural reproducibility.
|
||||
#[inline]
|
||||
pub fn from_tensor_id(tensor_id: u32) -> Self {
|
||||
// Mix bits so different tensor IDs get distinct offsets
|
||||
let mixed = tensor_id.wrapping_mul(0x9E37_79B9).wrapping_add(tensor_id >> 16);
|
||||
Self { idx: (mixed & 0xFF) as u8 }
|
||||
}
|
||||
}
|
||||
|
||||
impl DitherSource for PiDither {
|
||||
/// Advance and return next value in `[-0.5, 0.5]`.
|
||||
#[inline]
|
||||
fn next_unit(&mut self) -> f32 {
|
||||
let b = PI_BYTES[self.idx as usize];
|
||||
self.idx = self.idx.wrapping_add(1);
|
||||
(b as f32 / 255.0) - 0.5
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::DitherSource;
|
||||
|
||||
#[test]
|
||||
fn output_is_in_range() {
|
||||
let mut d = PiDither::new(0);
|
||||
for _ in 0..256 * 4 {
|
||||
let v = d.next_unit();
|
||||
assert!(v >= -0.5 && v <= 0.5, "out of range: {v}");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn period_is_256() {
|
||||
let mut d = PiDither::new(0);
|
||||
let first: Vec<f32> = (0..256).map(|_| d.next_unit()).collect();
|
||||
let second: Vec<f32> = (0..256).map(|_| d.next_unit()).collect();
|
||||
assert_eq!(first, second);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn mean_is_near_zero() {
|
||||
let mut d = PiDither::new(0);
|
||||
let sum: f32 = (0..256).map(|_| d.next_unit()).sum();
|
||||
let mean = sum / 256.0;
|
||||
assert!(mean.abs() < 0.05, "π-digit mean too large: {mean}");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn from_tensor_id_gives_distinct_offsets() {
|
||||
let d0 = PiDither::from_tensor_id(0);
|
||||
let d1 = PiDither::from_tensor_id(1);
|
||||
assert_ne!(d0.idx, d1.idx);
|
||||
}
|
||||
}
|
||||
130
crates/ruvector-dither/src/quantize.rs
Normal file
130
crates/ruvector-dither/src/quantize.rs
Normal file
|
|
@ -0,0 +1,130 @@
|
|||
//! Drop-in quantization helpers that apply dither before rounding.
|
||||
|
||||
use crate::DitherSource;
|
||||
|
||||
/// Quantize a single value with deterministic dither.
|
||||
///
|
||||
/// # Arguments
|
||||
/// - `x` – input activation in `[-1.0, 1.0]`
|
||||
/// - `bits` – quantizer bit-width (e.g. 3, 5, 7, 8)
|
||||
/// - `eps` – dither amplitude in LSB units (0.0 = no dither, 0.5 = half-LSB recommended)
|
||||
/// - `source` – stateful dither sequence
|
||||
///
|
||||
/// Returns the quantized value in `[-1.0, 1.0]`.
|
||||
///
|
||||
/// # Example
|
||||
/// ```
|
||||
/// use ruvector_dither::{GoldenRatioDither, quantize_dithered};
|
||||
/// let mut d = GoldenRatioDither::new(0.0);
|
||||
/// let q = quantize_dithered(0.314, 8, 0.5, &mut d);
|
||||
/// assert!(q >= -1.0 && q <= 1.0);
|
||||
/// ```
|
||||
#[inline]
|
||||
pub fn quantize_dithered(x: f32, bits: u32, eps: f32, source: &mut impl DitherSource) -> f32 {
|
||||
debug_assert!(bits >= 1 && bits <= 31, "bits must be in [1, 31]");
|
||||
let qmax = ((1u32 << (bits - 1)) - 1) as f32;
|
||||
let lsb = 1.0 / qmax;
|
||||
let dither = source.next(eps * lsb);
|
||||
let shifted = (x + dither) * qmax;
|
||||
let rounded = shifted.round().clamp(-qmax, qmax);
|
||||
rounded / qmax
|
||||
}
|
||||
|
||||
/// Quantize a slice in-place with deterministic dither.
|
||||
///
|
||||
/// Each element gets an independent dither sample from `source`.
|
||||
///
|
||||
/// # Example
|
||||
/// ```
|
||||
/// use ruvector_dither::{GoldenRatioDither, quantize_slice_dithered};
|
||||
/// let mut vals = vec![0.1_f32, 0.5, -0.3, 0.9, -0.8];
|
||||
/// let mut d = GoldenRatioDither::new(0.0);
|
||||
/// quantize_slice_dithered(&mut vals, 5, 0.5, &mut d);
|
||||
/// for &v in &vals {
|
||||
/// assert!(v >= -1.0 && v <= 1.0);
|
||||
/// }
|
||||
/// ```
|
||||
pub fn quantize_slice_dithered(
|
||||
xs: &mut [f32],
|
||||
bits: u32,
|
||||
eps: f32,
|
||||
source: &mut impl DitherSource,
|
||||
) {
|
||||
let qmax = ((1u32 << (bits - 1)) - 1) as f32;
|
||||
let lsb = 1.0 / qmax;
|
||||
for x in xs.iter_mut() {
|
||||
let dither = source.next(eps * lsb);
|
||||
let shifted = (*x + dither) * qmax;
|
||||
*x = shifted.round().clamp(-qmax, qmax) / qmax;
|
||||
}
|
||||
}
|
||||
|
||||
/// Quantize to a raw integer code (signed, in `[-(2^(bits-1)), 2^(bits-1)-1]`).
|
||||
///
|
||||
/// Useful when you need the integer representation rather than a re-scaled float.
|
||||
#[inline]
|
||||
pub fn quantize_to_code(x: f32, bits: u32, eps: f32, source: &mut impl DitherSource) -> i32 {
|
||||
debug_assert!(bits >= 1 && bits <= 31);
|
||||
let qmax = ((1u32 << (bits - 1)) - 1) as f32;
|
||||
let lsb = 1.0 / qmax;
|
||||
let dither = source.next(eps * lsb);
|
||||
((x + dither) * qmax).round().clamp(-qmax, qmax) as i32
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::{GoldenRatioDither, PiDither};
|
||||
|
||||
#[test]
|
||||
fn output_in_unit_range() {
|
||||
let mut d = GoldenRatioDither::new(0.0);
|
||||
for bits in [3u32, 5, 7, 8] {
|
||||
for &x in &[-1.0_f32, -0.5, 0.0, 0.5, 1.0] {
|
||||
let q = quantize_dithered(x, bits, 0.5, &mut d);
|
||||
assert!(q >= -1.0 && q <= 1.0, "bits={bits}, x={x}, q={q}");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn dither_reduces_idle_tones() {
|
||||
// A constant signal at exactly 0.5 * LSB without dither quantizes
|
||||
// to the same code every time (idle tone). With dither the code
|
||||
// alternates, so the variance of codes should be > 0.
|
||||
let bits = 5u32;
|
||||
let qmax = ((1u32 << (bits - 1)) - 1) as f32;
|
||||
let lsb = 1.0 / qmax;
|
||||
let x = 0.5 * lsb; // exactly half an LSB
|
||||
|
||||
let mut codes_with: Vec<i32> = Vec::with_capacity(256);
|
||||
let mut d = GoldenRatioDither::new(0.0);
|
||||
for _ in 0..256 {
|
||||
codes_with.push(quantize_to_code(x, bits, 0.5, &mut d));
|
||||
}
|
||||
let unique: std::collections::HashSet<i32> = codes_with.iter().copied().collect();
|
||||
assert!(unique.len() > 1, "dithered signal must produce >1 unique code");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn slice_quantize_in_bounds() {
|
||||
let mut vals: Vec<f32> = (-50..=50).map(|i| i as f32 * 0.02).collect();
|
||||
let mut pi = PiDither::new(0);
|
||||
quantize_slice_dithered(&mut vals, 7, 0.5, &mut pi);
|
||||
for v in vals {
|
||||
assert!(v >= -1.0 && v <= 1.0, "out of range: {v}");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn deterministic_with_same_seed() {
|
||||
let input = vec![0.1_f32, 0.4, -0.7, 0.9];
|
||||
let quantize = |input: &[f32]| {
|
||||
let mut buf = input.to_vec();
|
||||
let mut d = GoldenRatioDither::new(0.5);
|
||||
quantize_slice_dithered(&mut buf, 8, 0.5, &mut d);
|
||||
buf
|
||||
};
|
||||
assert_eq!(quantize(&input), quantize(&input));
|
||||
}
|
||||
}
|
||||
29
examples/exo-ai-2025/Cargo.lock
generated
29
examples/exo-ai-2025/Cargo.lock
generated
|
|
@ -517,7 +517,7 @@ dependencies = [
|
|||
"clap",
|
||||
"criterion-plot",
|
||||
"is-terminal",
|
||||
"itertools",
|
||||
"itertools 0.10.5",
|
||||
"num-traits",
|
||||
"once_cell",
|
||||
"oorandom",
|
||||
|
|
@ -538,7 +538,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
|
|||
checksum = "6b50826342786a51a89e2da3a28f1c32b06e387201bc2d19791f622c673706b1"
|
||||
dependencies = [
|
||||
"cast",
|
||||
"itertools",
|
||||
"itertools 0.10.5",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
|
|
@ -795,11 +795,14 @@ dependencies = [
|
|||
"exo-manifold",
|
||||
"exo-temporal 0.1.0",
|
||||
"parking_lot",
|
||||
"rand 0.8.5",
|
||||
"ruvector-core",
|
||||
"ruvector-dither",
|
||||
"ruvector-domain-expansion",
|
||||
"ruvector-graph",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"thermorust",
|
||||
"thiserror 2.0.17",
|
||||
"uuid",
|
||||
]
|
||||
|
|
@ -1305,6 +1308,15 @@ dependencies = [
|
|||
"either",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "itertools"
|
||||
version = "0.12.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "ba291022dbbd398a455acf126c1e341954079855bc60dfdda641363bd6922569"
|
||||
dependencies = [
|
||||
"either",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "itoa"
|
||||
version = "1.0.15"
|
||||
|
|
@ -2233,6 +2245,10 @@ dependencies = [
|
|||
"uuid",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "ruvector-dither"
|
||||
version = "0.1.0"
|
||||
|
||||
[[package]]
|
||||
name = "ruvector-domain-expansion"
|
||||
version = "2.0.5"
|
||||
|
|
@ -2530,6 +2546,15 @@ version = "0.2.0"
|
|||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "7b2093cf4c8eb1e67749a6762251bc9cd836b6fc171623bd0a9d324d37af2417"
|
||||
|
||||
[[package]]
|
||||
name = "thermorust"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"itertools 0.12.1",
|
||||
"rand 0.8.5",
|
||||
"rand_distr",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "thiserror"
|
||||
version = "1.0.69"
|
||||
|
|
|
|||
|
|
@ -24,6 +24,9 @@ exo-exotic = { path = "../exo-exotic" }
|
|||
ruvector-core = { version = "0.1", features = ["simd"] }
|
||||
ruvector-graph = "0.1"
|
||||
ruvector-domain-expansion = { path = "../../../../crates/ruvector-domain-expansion", features = ["rvf"] }
|
||||
thermorust = { path = "../../../../crates/thermorust" }
|
||||
ruvector-dither = { path = "../../../../crates/ruvector-dither" }
|
||||
rand = { version = "0.8", features = ["small_rng"] }
|
||||
|
||||
# Utility dependencies
|
||||
serde = { version = "1.0", features = ["derive"] }
|
||||
|
|
|
|||
|
|
@ -0,0 +1,161 @@
|
|||
//! DitheredQuantizer: deterministic low-bit quantization for exo activations.
|
||||
//!
|
||||
//! Wraps `ruvector-dither` to provide drop-in dithered quantization for
|
||||
//! exo-backend-classical activation and weight tensors.
|
||||
//!
|
||||
//! Dithering breaks power-of-two resonances that cause idle tones / sticky
|
||||
//! activations in 3/5/7-bit inference — without any RNG overhead.
|
||||
//!
|
||||
//! # Quick start
|
||||
//!
|
||||
//! ```
|
||||
//! use exo_backend_classical::dither_quantizer::{DitheredQuantizer, DitherKind};
|
||||
//!
|
||||
//! // 8-bit, golden-ratio dither, layer 0, 16 channels, ε = 0.5 LSB
|
||||
//! let mut q = DitheredQuantizer::new(DitherKind::GoldenRatio, 0, 16, 8, 0.5);
|
||||
//!
|
||||
//! let mut activations = vec![0.3_f32, -0.7, 0.5, 0.1];
|
||||
//! q.quantize(&mut activations);
|
||||
//! assert!(activations.iter().all(|&v| v >= -1.0 && v <= 1.0));
|
||||
//! ```
|
||||
|
||||
use ruvector_dither::{channel::ChannelDither, quantize_slice_dithered, PiDither};
|
||||
|
||||
/// Which deterministic dither sequence to use.
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub enum DitherKind {
|
||||
/// Golden-ratio quasi-random sequence (best equidistribution, no period).
|
||||
GoldenRatio,
|
||||
/// π-digit cyclic sequence (period = 256; ideal for weight pack-time use).
|
||||
Pi,
|
||||
}
|
||||
|
||||
enum Source {
|
||||
Golden(ChannelDither),
|
||||
Pi(PiDither),
|
||||
}
|
||||
|
||||
/// Dithered quantizer for exo activation / weight tensors.
|
||||
pub struct DitheredQuantizer {
|
||||
source: Source,
|
||||
bits: u32,
|
||||
eps: f32,
|
||||
}
|
||||
|
||||
impl DitheredQuantizer {
|
||||
/// Create a new quantizer.
|
||||
///
|
||||
/// - `kind` – dither sequence type
|
||||
/// - `layer_id` – identifies this layer (seeds per-channel states)
|
||||
/// - `n_channels` – number of independent channels (ignored for Pi)
|
||||
/// - `bits` – quantizer bit-width (3–8)
|
||||
/// - `eps` – dither amplitude in LSB units (0.5 recommended)
|
||||
pub fn new(kind: DitherKind, layer_id: u32, n_channels: usize, bits: u32, eps: f32) -> Self {
|
||||
let source = match kind {
|
||||
DitherKind::GoldenRatio => {
|
||||
Source::Golden(ChannelDither::new(layer_id, n_channels, bits, eps))
|
||||
}
|
||||
DitherKind::Pi => {
|
||||
Source::Pi(PiDither::from_tensor_id(layer_id))
|
||||
}
|
||||
};
|
||||
Self { source, bits, eps }
|
||||
}
|
||||
|
||||
/// Quantize `activations` in-place.
|
||||
///
|
||||
/// Each element is rounded to the nearest representable value in
|
||||
/// `[-1.0, 1.0]` at `bits`-bit precision with dither applied.
|
||||
pub fn quantize(&mut self, activations: &mut [f32]) {
|
||||
match &mut self.source {
|
||||
Source::Golden(cd) => cd.quantize_batch(activations),
|
||||
Source::Pi(pd) => quantize_slice_dithered(activations, self.bits, self.eps, pd),
|
||||
}
|
||||
}
|
||||
|
||||
/// Reset the dither state to the initial seed (useful for reproducible tests).
|
||||
pub fn reset(&mut self, layer_id: u32, n_channels: usize) {
|
||||
match &mut self.source {
|
||||
Source::Golden(cd) => {
|
||||
*cd = ChannelDither::new(layer_id, n_channels, self.bits, self.eps);
|
||||
}
|
||||
Source::Pi(pd) => {
|
||||
*pd = PiDither::from_tensor_id(layer_id);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Bit-width used by this quantizer.
|
||||
pub fn bits(&self) -> u32 {
|
||||
self.bits
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn golden_quantizer_in_bounds() {
|
||||
let mut q = DitheredQuantizer::new(DitherKind::GoldenRatio, 0, 8, 8, 0.5);
|
||||
let mut acts: Vec<f32> = (0..64).map(|i| (i as f32 / 63.0) * 2.0 - 1.0).collect();
|
||||
q.quantize(&mut acts);
|
||||
for v in &acts {
|
||||
assert!(*v >= -1.0 && *v <= 1.0, "out of bounds: {v}");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn pi_quantizer_in_bounds() {
|
||||
let mut q = DitheredQuantizer::new(DitherKind::Pi, 42, 1, 5, 0.5);
|
||||
let mut acts = vec![0.3_f32, -0.7, 0.5, 0.1, -1.0, 1.0];
|
||||
q.quantize(&mut acts);
|
||||
for v in &acts {
|
||||
assert!(*v >= -1.0 && *v <= 1.0, "out of bounds: {v}");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn different_layers_different_output() {
|
||||
let input: Vec<f32> = vec![0.5; 16];
|
||||
|
||||
let quantize = |layer: u32| {
|
||||
let mut buf = input.clone();
|
||||
let mut q = DitheredQuantizer::new(DitherKind::GoldenRatio, layer, 8, 8, 0.5);
|
||||
q.quantize(&mut buf);
|
||||
buf
|
||||
};
|
||||
assert_ne!(quantize(0), quantize(1));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn deterministic_after_reset() {
|
||||
let input: Vec<f32> = vec![0.3, -0.4, 0.7, -0.1, 0.9];
|
||||
let mut q = DitheredQuantizer::new(DitherKind::GoldenRatio, 7, 4, 8, 0.5);
|
||||
|
||||
let mut buf1 = input.clone();
|
||||
q.quantize(&mut buf1);
|
||||
|
||||
q.reset(7, 4);
|
||||
let mut buf2 = input.clone();
|
||||
q.quantize(&mut buf2);
|
||||
|
||||
assert_eq!(buf1, buf2, "reset must restore deterministic output");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn three_bit_quantization() {
|
||||
let mut q = DitheredQuantizer::new(DitherKind::Pi, 0, 1, 3, 0.5);
|
||||
let mut acts = vec![-0.9_f32, -0.5, 0.0, 0.5, 0.9];
|
||||
q.quantize(&mut acts);
|
||||
for v in &acts {
|
||||
assert!(*v >= -1.0 && *v <= 1.0);
|
||||
}
|
||||
// 3-bit: qmax = 3, only multiples of 1/3 are valid
|
||||
let step = 1.0 / 3.0;
|
||||
for v in &acts {
|
||||
let rem = (v / step).round() * step - v;
|
||||
assert!(rem.abs() < 1e-5, "3-bit output should be on grid: {v}");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -6,8 +6,10 @@
|
|||
|
||||
#![warn(missing_docs)]
|
||||
|
||||
pub mod dither_quantizer;
|
||||
pub mod domain_bridge;
|
||||
pub mod graph;
|
||||
pub mod thermo_layer;
|
||||
pub mod transfer_orchestrator;
|
||||
pub mod vector;
|
||||
|
||||
|
|
|
|||
|
|
@ -0,0 +1,185 @@
|
|||
//! ThermoLayer: thermodynamic coherence gate for exo-backend-classical.
|
||||
//!
|
||||
//! Wraps a `thermorust` Ising motif and treats the energy drop ΔE as a
|
||||
//! **coherence λ-signal**: a large negative ΔE means the activation pattern
|
||||
//! is "settling" (becoming coherent); a near-zero ΔE means it is already
|
||||
//! at a local minimum or chaotically fluctuating at high temperature.
|
||||
//!
|
||||
//! The λ-signal can be used to gate min-cut operations or to weight
|
||||
//! confidence scores in the ruvector-attn-mincut pipeline.
|
||||
//!
|
||||
//! # Integration sketch
|
||||
//! ```no_run
|
||||
//! use exo_backend_classical::thermo_layer::{ThermoLayer, ThermoConfig};
|
||||
//!
|
||||
//! let cfg = ThermoConfig { n: 16, beta: 3.0, steps_per_call: 20, ..Default::default() };
|
||||
//! let mut layer = ThermoLayer::new(cfg);
|
||||
//!
|
||||
//! // Activations from an attention layer (length must equal `n`).
|
||||
//! let mut acts = vec![0.5_f32; 16];
|
||||
//! let signal = layer.run(&mut acts, 20);
|
||||
//! println!("λ = {:.4}, dissipation = {:.3e} J", signal.lambda, signal.dissipation_j);
|
||||
//! ```
|
||||
|
||||
use rand::SeedableRng;
|
||||
use thermorust::{
|
||||
dynamics::{Params, step_discrete},
|
||||
energy::{Couplings, EnergyModel, Ising},
|
||||
metrics::magnetisation,
|
||||
State,
|
||||
};
|
||||
|
||||
/// Configuration for a `ThermoLayer`.
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct ThermoConfig {
|
||||
/// Number of units in the Ising motif (must match activation vector length).
|
||||
pub n: usize,
|
||||
/// Inverse temperature β = 1/(kT). Higher = colder, more deterministic.
|
||||
pub beta: f32,
|
||||
/// Ferromagnetic coupling strength J for ring topology.
|
||||
pub coupling: f32,
|
||||
/// Metropolis steps executed per `run()` call.
|
||||
pub steps_per_call: usize,
|
||||
/// Landauer cost in Joules per accepted irreversible flip.
|
||||
pub irreversible_cost: f64,
|
||||
/// RNG seed (fixed → fully deterministic).
|
||||
pub seed: u64,
|
||||
}
|
||||
|
||||
impl Default for ThermoConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
n: 16,
|
||||
beta: 3.0,
|
||||
coupling: 0.2,
|
||||
steps_per_call: 20,
|
||||
irreversible_cost: 2.87e-21, // kT ln2 at 300 K
|
||||
seed: 0,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Thermodynamic coherence signal returned by `ThermoLayer::run`.
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct ThermoSignal {
|
||||
/// λ-signal: −ΔE / |E_initial| (positive = energy decreased = more coherent).
|
||||
pub lambda: f32,
|
||||
/// Magnetisation m ∈ [−1, 1] after update.
|
||||
pub magnetisation: f32,
|
||||
/// Cumulative Joules dissipated since layer creation.
|
||||
pub dissipation_j: f64,
|
||||
/// Energy after the update step.
|
||||
pub energy_after: f32,
|
||||
}
|
||||
|
||||
/// Ising-motif thermodynamic gate.
|
||||
pub struct ThermoLayer {
|
||||
model: Ising,
|
||||
state: State,
|
||||
params: Params,
|
||||
rng: rand::rngs::SmallRng,
|
||||
}
|
||||
|
||||
impl ThermoLayer {
|
||||
/// Create a new `ThermoLayer` from `cfg`.
|
||||
pub fn new(cfg: ThermoConfig) -> Self {
|
||||
let couplings = Couplings::ferromagnetic_ring(cfg.n, cfg.coupling);
|
||||
let model = Ising::new(couplings);
|
||||
let state = State::ones(cfg.n);
|
||||
let params = Params {
|
||||
beta: cfg.beta,
|
||||
eta: 0.05,
|
||||
irreversible_cost: cfg.irreversible_cost,
|
||||
clamp_mask: vec![false; cfg.n],
|
||||
};
|
||||
let rng = rand::rngs::SmallRng::seed_from_u64(cfg.seed);
|
||||
Self { model, state, params, rng }
|
||||
}
|
||||
|
||||
/// Apply activations as external fields, run MH steps, return coherence signal.
|
||||
///
|
||||
/// The activation vector is **modified in place** by the thermodynamic
|
||||
/// relaxation: each element is replaced by the Ising spin value after
|
||||
/// `steps_per_call` Metropolis updates. Values are clamped to {-1, +1}.
|
||||
pub fn run(&mut self, activations: &mut [f32], steps: usize) -> ThermoSignal {
|
||||
let n = self.state.len().min(activations.len());
|
||||
|
||||
// Clamp inputs to ±1 and load as spin state.
|
||||
for i in 0..n {
|
||||
self.state.x[i] = activations[i].clamp(-1.0, 1.0).signum();
|
||||
}
|
||||
|
||||
let e_before = self.model.energy(&self.state);
|
||||
|
||||
// Run Metropolis steps.
|
||||
for _ in 0..steps {
|
||||
step_discrete(&self.model, &mut self.state, &self.params, &mut self.rng);
|
||||
}
|
||||
|
||||
let e_after = self.model.energy(&self.state);
|
||||
let d_e = e_after - e_before;
|
||||
let lambda = if e_before.abs() > 1e-9 {
|
||||
-d_e / e_before.abs()
|
||||
} else {
|
||||
0.0
|
||||
};
|
||||
|
||||
// Write relaxed spins back to the caller's buffer.
|
||||
for i in 0..n {
|
||||
activations[i] = self.state.x[i];
|
||||
}
|
||||
|
||||
ThermoSignal {
|
||||
lambda,
|
||||
magnetisation: magnetisation(&self.state),
|
||||
dissipation_j: self.state.dissipated_j,
|
||||
energy_after: e_after,
|
||||
}
|
||||
}
|
||||
|
||||
/// Reset the spin state to all +1.
|
||||
pub fn reset(&mut self) {
|
||||
for xi in &mut self.state.x {
|
||||
*xi = 1.0;
|
||||
}
|
||||
self.state.dissipated_j = 0.0;
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn thermo_layer_runs_without_panic() {
|
||||
let cfg = ThermoConfig { n: 8, steps_per_call: 10, ..Default::default() };
|
||||
let mut layer = ThermoLayer::new(cfg);
|
||||
let mut acts = vec![1.0_f32; 8];
|
||||
let sig = layer.run(&mut acts, 10);
|
||||
assert!(sig.lambda.is_finite());
|
||||
assert!(sig.magnetisation >= -1.0 && sig.magnetisation <= 1.0);
|
||||
assert!(sig.dissipation_j >= 0.0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn activations_are_binarised() {
|
||||
let cfg = ThermoConfig { n: 4, steps_per_call: 0, ..Default::default() };
|
||||
let mut layer = ThermoLayer::new(cfg);
|
||||
let mut acts = vec![0.7_f32, -0.3, 0.1, -0.9];
|
||||
layer.run(&mut acts, 0);
|
||||
for a in &acts {
|
||||
assert!((*a - 1.0).abs() < 1e-6 || (*a + 1.0).abs() < 1e-6, "not ±1: {a}");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn lambda_finite_after_many_steps() {
|
||||
let cfg = ThermoConfig { n: 16, beta: 5.0, ..Default::default() };
|
||||
let mut layer = ThermoLayer::new(cfg);
|
||||
for _ in 0..10 {
|
||||
let mut acts = vec![1.0_f32; 16];
|
||||
let sig = layer.run(&mut acts, 50);
|
||||
assert!(sig.lambda.is_finite());
|
||||
}
|
||||
}
|
||||
}
|
||||
Loading…
Add table
Add a link
Reference in a new issue