From 65da488add1823eb0dd28558cb8d0b28ac772ad3 Mon Sep 17 00:00:00 2001 From: ruv Date: Tue, 21 Jul 2026 18:20:23 -0700 Subject: [PATCH] =?UTF-8?q?feat(adr-185):=20add=20weight-loading=20capabil?= =?UTF-8?q?ity=20to=20AETHER=20EmbeddingExtractor=20(=C2=A713.a)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Closes the ADR-185 §13.a follow-up (the tractable, data-independent one of the three §6.7 gaps): the bound EmbeddingExtractor was random-Xavier-init only, with NO way to load real weights — so it was structurally untrained. This adds the *capability* to load weights whenever a trained checkpoint exists. It does NOT itself produce trained/SOTA embeddings and does NOT close §6.7 (still no trained checkpoint, no labeled data, no eval harness). Serialization: not greenfield — EmbeddingExtractor already had flatten_weights()/unflatten_weights() (flat Vec). wifi-densepose-aether is a deliberately dependency-free std-only leaf crate (ADR-185 §13), so rather than add safetensors/serde/bincode (which would undo the zero-dep property), the on-disk format is raw little-endian f32 with a 12-byte header (magic "AETHERW1" + u32 param count) — zero new deps. Native (v2/crates/wifi-densepose-aether/src/embedding.rs): - EmbeddingExtractor::save_weights(path) / load_weights(path). load_weights never panics: errors on unreadable file, short/oversized payload, bad magic, or a param-count mismatch (delegated to unflatten_weights). - Default (no weights) construction is UNCHANGED — still random init, clearly labeled untrained. Purely additive. Python binding (python/src/bindings/aether.rs, .pyi): - EmbeddingExtractor.load_weights(path) / save_weights(path) / param_count, GIL-released, ValueError on bad input. Tests (both Rust and Python, per the task): - Rust unit (embedding.rs): load_weights_actually_replaces_weights_and_ round_trips proves the loaded weights MOVE the embedding away from the random-init baseline (not a silent no-op), match the source extractor (round-trip), and are bit-identical after the file round-trip; plus a bad-magic/wrong-count rejection test. - Cross-language golden (aether_weights_parity.rs + test_aether.py): a shared deterministic weight formula (w[i]=k/65536-0.5, exact in f32+f64) written to the AETHER format; both native Rust and the Python binding load it and must produce the byte-identical embedding SHA-256 (tests/golden/aether_loaded_embedding.sha256) — proving the binding's load path is bit-identical to native, and (vs the random baseline) that the loaded weights are actually used. Verified: cargo test -p wifi-densepose-aether 98 passed, 0 failed cargo test --features aether --test aether_parity --test aether_weights_parity 3 passed, 0 failed maturin develop --features aether + pytest test_aether.py 13/13 pass default cargo build (no aether feature) clean --- python/src/bindings/aether.rs | 28 ++++ python/tests/aether_weights_parity.rs | 130 ++++++++++++++++ .../golden/aether_loaded_embedding.sha256 | 1 + python/tests/test_aether.py | 71 +++++++++ python/wifi_densepose/aether.pyi | 4 + .../wifi-densepose-aether/src/embedding.rs | 143 ++++++++++++++++++ 6 files changed, 377 insertions(+) create mode 100644 python/tests/aether_weights_parity.rs create mode 100644 python/tests/golden/aether_loaded_embedding.sha256 diff --git a/python/src/bindings/aether.rs b/python/src/bindings/aether.rs index f0a84fab..4a6461ee 100644 --- a/python/src/bindings/aether.rs +++ b/python/src/bindings/aether.rs @@ -29,6 +29,7 @@ //! ops touching no Python objects, so they run inside //! `py.allow_threads(|| ...)`. +use pyo3::exceptions::PyValueError; use pyo3::prelude::*; use wifi_densepose_aether::embedding::{ @@ -196,6 +197,33 @@ impl PyEmbeddingExtractor { self.embedding_dim } + /// Total trainable parameter count (transformer + projection). Equals the + /// number of `f32`s in a weight file for this architecture. + #[getter] + fn param_count(&self) -> usize { + self.inner.param_count() + } + + /// Load weights from `path` (a file written by `save_weights` or the Rust + /// `EmbeddingExtractor::save_weights`), replacing the current weights. + /// + /// By default an `EmbeddingExtractor` uses deterministic **random** init + /// (untrained); this is the additive path to load real weights once a + /// trained checkpoint exists (ADR-185 §13.a). Raises `ValueError` on a + /// missing/corrupt file or a param-count mismatch with this architecture. + /// GIL released during file I/O + deserialization. + fn load_weights(&mut self, py: Python<'_>, path: String) -> PyResult<()> { + py.allow_threads(|| self.inner.load_weights(&path)) + .map_err(PyValueError::new_err) + } + + /// Serialize the current weights to `path` (magic `AETHERW1` + `u32` count + /// + little-endian `f32` payload). GIL released. + fn save_weights(&self, py: Python<'_>, path: String) -> PyResult<()> { + py.allow_threads(|| self.inner.save_weights(&path)) + .map_err(|e| PyValueError::new_err(e.to_string())) + } + fn __repr__(&self) -> String { format!("EmbeddingExtractor(embedding_dim={})", self.embedding_dim) } diff --git a/python/tests/aether_weights_parity.rs b/python/tests/aether_weights_parity.rs new file mode 100644 index 00000000..75bc15cc --- /dev/null +++ b/python/tests/aether_weights_parity.rs @@ -0,0 +1,130 @@ +//! ADR-185 §13.a — weight-loading parity: native-Rust reference half. +//! +//! Proves the AETHER `load_weights` path produces a deterministic, non-random +//! embedding, and locks its SHA-256 into +//! `tests/golden/aether_loaded_embedding.sha256`. The pytest half +//! (`tests/test_aether.py`) writes a byte-identical weight file (same formula + +//! format) through the binding's `load_weights` and asserts the same hash — +//! together they prove the binding's weight-loading is bit-identical to native. +//! +//! Weight formula (shared with the pytest half): `w[i] = k/65536 - 0.5` where +//! `k = (i*1103515245 + 12345) mod 65536`. `k/65536` is a multiple of 2⁻¹⁶, +//! exactly representable in both f32 and f64, so both languages produce +//! byte-identical weights. +//! +//! File format: 8-byte magic `AETHERW1`, `u32` little-endian param count, then +//! that many little-endian `f32`. +//! +//! Regenerate (only on an intentional change): delete the .sha256 and re-run +//! `cargo test --features aether --test aether_weights_parity`. +#![cfg(feature = "aether")] + +use std::fs; +use std::path::PathBuf; + +use sha2::{Digest, Sha256}; +use wifi_densepose_aether::embedding::{EmbeddingConfig, EmbeddingExtractor}; +use wifi_densepose_aether::graph_transformer::TransformerConfig; + +fn golden_dir() -> PathBuf { + PathBuf::from(env!("CARGO_MANIFEST_DIR")) + .join("tests") + .join("golden") +} + +fn load_input() -> Vec> { + let raw = fs::read_to_string(golden_dir().join("aether_input.json")) + .expect("read aether_input.json fixture"); + let rows: Vec> = serde_json::from_str(&raw).expect("parse aether_input.json"); + rows.into_iter() + .map(|row| row.into_iter().map(|x| x as f32).collect()) + .collect() +} + +/// Same default construction as `aether_parity.rs` / the Python binding default. +fn new_extractor() -> EmbeddingExtractor { + let e_config = EmbeddingConfig { + d_model: 64, + d_proj: 128, + temperature: 0.07, + normalize: true, + }; + let t_config = TransformerConfig { + n_subcarriers: 56, + n_keypoints: 17, + d_model: 64, + n_heads: 4, + n_gnn_layers: 2, + }; + EmbeddingExtractor::new(t_config, e_config) +} + +fn formula_weights(n: usize) -> Vec { + (0..n) + .map(|i| { + let k = (i as u32).wrapping_mul(1_103_515_245).wrapping_add(12_345) % 65_536; + k as f32 / 65_536.0 - 0.5 + }) + .collect() +} + +fn write_weight_file(path: &PathBuf, weights: &[f32]) { + let mut buf = Vec::with_capacity(12 + weights.len() * 4); + buf.extend_from_slice(b"AETHERW1"); + buf.extend_from_slice(&(weights.len() as u32).to_le_bytes()); + for v in weights { + buf.extend_from_slice(&v.to_le_bytes()); + } + fs::write(path, buf).unwrap(); +} + +fn sha256_le(embedding: &[f32]) -> String { + let mut hasher = Sha256::new(); + for &x in embedding { + hasher.update(x.to_le_bytes()); + } + hasher + .finalize() + .iter() + .map(|b| format!("{b:02x}")) + .collect() +} + +#[test] +fn native_loaded_embedding_matches_committed_golden() { + let input = load_input(); + + let mut ext = new_extractor(); + let baseline = ext.extract(&input); // random Xavier init + + let weights = formula_weights(ext.param_count()); + let path = std::env::temp_dir().join(format!( + "aether_weights_parity_{}.bin", + std::process::id() + )); + write_weight_file(&path, &weights); + ext.load_weights(&path).expect("load_weights"); + fs::remove_file(&path).ok(); + + let loaded = ext.extract(&input); + // Loaded weights must actually take effect. + assert!( + baseline.iter().zip(&loaded).any(|(a, b)| (a - b).abs() > 1e-6), + "load_weights had no effect vs the random-init baseline" + ); + assert_eq!(loaded.len(), 128); + + let got = sha256_le(&loaded); + let sha_path = golden_dir().join("aether_loaded_embedding.sha256"); + match fs::read_to_string(&sha_path) { + Ok(expected) => assert_eq!( + got, + expected.trim(), + "native loaded-weights embedding hash drifted from committed golden" + ), + Err(_) => { + fs::write(&sha_path, &got).expect("write golden sha256"); + panic!("no committed golden found; wrote {got}. Re-run to verify parity."); + } + } +} diff --git a/python/tests/golden/aether_loaded_embedding.sha256 b/python/tests/golden/aether_loaded_embedding.sha256 new file mode 100644 index 00000000..88fb68d0 --- /dev/null +++ b/python/tests/golden/aether_loaded_embedding.sha256 @@ -0,0 +1 @@ +af7b87579b399a0586ed5b52f6e5b625709ea48affe52e856675203b602a5014 \ No newline at end of file diff --git a/python/tests/test_aether.py b/python/tests/test_aether.py index fc9a7c20..ff4bbba3 100644 --- a/python/tests/test_aether.py +++ b/python/tests/test_aether.py @@ -12,6 +12,7 @@ import hashlib import json import math import struct +import tempfile from pathlib import Path import pytest @@ -31,6 +32,19 @@ def build_extractor() -> aether.EmbeddingExtractor: return aether.EmbeddingExtractor(n_subcarriers=56, config=cfg) +def _formula_weights(n: int) -> list[float]: + # Byte-identical to aether_weights_parity.rs (k/65536 is exact in f32+f64). + return [((i * 1103515245 + 12345) % 65536) / 65536.0 - 0.5 for i in range(n)] + + +def _write_weight_file(path: Path, weights: list[float]) -> None: + # AETHER weight format: b"AETHERW1" + u32 count + LE f32 payload. + with open(path, "wb") as f: + f.write(b"AETHERW1") + f.write(struct.pack(" None: cfg = aether.AetherConfig(d_model=64, d_proj=128, temperature=0.07, normalize=True) assert cfg.d_model == 64 @@ -101,3 +115,60 @@ def test_base_wheel_import_error_message() -> None: # message is present in source so the base-wheel path stays honest. src = (Path(aether.__file__)).read_text() assert "pip install wifi-densepose[aether]" in src + + +# ─── Weight loading (ADR-185 §13.a) ────────────────────────────────── + +def test_load_weights_is_used_and_matches_native_golden() -> None: + ext = build_extractor() + baseline = ext.embed(load_input()) # random Xavier init + + weights = _formula_weights(ext.param_count) + with tempfile.TemporaryDirectory() as d: + wpath = Path(d) / "weights.bin" + _write_weight_file(wpath, weights) + ext.load_weights(str(wpath)) + + loaded = ext.embed(load_input()) + + # (1) The loaded weights actually take effect (not a silent no-op). + assert any(abs(a - b) > 1e-6 for a, b in zip(baseline, loaded)), ( + "load_weights had no effect — embedding still equals the random-init baseline" + ) + # (2) Bit-identical to the native-Rust reference that loaded the same weights. + packed = b"".join(struct.pack(" None: + ext = build_extractor() + inp = load_input() + with tempfile.TemporaryDirectory() as d: + wpath = Path(d) / "roundtrip.bin" + ext.save_weights(str(wpath)) # serialize current (random) weights + emb_before = ext.embed(inp) + ext2 = build_extractor() + ext2.load_weights(str(wpath)) # load them into a fresh extractor + assert ext2.embed(inp) == emb_before + + +def test_load_weights_rejects_bad_magic() -> None: + ext = build_extractor() + with tempfile.TemporaryDirectory() as d: + wpath = Path(d) / "bad.bin" + wpath.write_bytes(b"NOTAETHER" + b"\x00" * 8) + with pytest.raises(ValueError): + ext.load_weights(str(wpath)) + + +def test_load_weights_rejects_wrong_param_count() -> None: + ext = build_extractor() + with tempfile.TemporaryDirectory() as d: + wpath = Path(d) / "short.bin" + _write_weight_file(wpath, [0.1, 0.2, 0.3]) # far too few params + with pytest.raises(ValueError): + ext.load_weights(str(wpath)) diff --git a/python/wifi_densepose/aether.pyi b/python/wifi_densepose/aether.pyi index 3cce9094..9d4be719 100644 --- a/python/wifi_densepose/aether.pyi +++ b/python/wifi_densepose/aether.pyi @@ -44,6 +44,10 @@ class EmbeddingExtractor: def embed(self, csi_features: list[list[float]]) -> list[float]: ... @property def embedding_dim(self) -> int: ... + @property + def param_count(self) -> int: ... + def load_weights(self, path: str) -> None: ... + def save_weights(self, path: str) -> None: ... def __repr__(self) -> str: ... def info_nce_loss( diff --git a/v2/crates/wifi-densepose-aether/src/embedding.rs b/v2/crates/wifi-densepose-aether/src/embedding.rs index 2b2e43ae..5573b540 100644 --- a/v2/crates/wifi-densepose-aether/src/embedding.rs +++ b/v2/crates/wifi-densepose-aether/src/embedding.rs @@ -822,8 +822,73 @@ impl EmbeddingExtractor { self.projection = proj; Ok(()) } + + /// Serialize all weights (transformer + projection) to `path`. + /// + /// Format (little-endian, zero-dep so the std-only leaf crate stays + /// dependency-free): 8-byte magic `AETHERW1`, then a `u32` parameter + /// count, then that many `f32` values — exactly `flatten_weights()`. + /// + /// This is the counterpart of [`Self::load_weights`]. It enables loading a + /// *real trained* checkpoint once one exists (ADR-185 §13.a); it does not + /// itself make the default (random-init) extractor trained. + pub fn save_weights>(&self, path: P) -> std::io::Result<()> { + let weights = self.flatten_weights(); + let mut buf = Vec::with_capacity(WEIGHT_HEADER_LEN + weights.len() * 4); + buf.extend_from_slice(WEIGHT_MAGIC); + buf.extend_from_slice(&(weights.len() as u32).to_le_bytes()); + for v in &weights { + buf.extend_from_slice(&v.to_le_bytes()); + } + std::fs::write(path, buf) + } + + /// Load weights previously written by [`Self::save_weights`] (or any file in + /// that format) into this extractor, replacing the current (random-init or + /// prior) weights. + /// + /// Errors (never panics) on: unreadable file, a payload shorter than the + /// header, a wrong magic, a truncated/oversized payload, or a parameter + /// count that does not match this extractor's architecture (delegated to + /// [`Self::unflatten_weights`]). + pub fn load_weights>(&mut self, path: P) -> Result<(), String> { + let bytes = std::fs::read(path).map_err(|e| format!("failed to read weight file: {e}"))?; + if bytes.len() < WEIGHT_HEADER_LEN { + return Err(format!( + "weight file too short: {} bytes < {WEIGHT_HEADER_LEN}-byte header", + bytes.len() + )); + } + if &bytes[0..8] != WEIGHT_MAGIC { + return Err("bad magic: not an AETHER weight file (expected 'AETHERW1')".to_string()); + } + let count = u32::from_le_bytes([bytes[8], bytes[9], bytes[10], bytes[11]]) as usize; + let expected_len = WEIGHT_HEADER_LEN + count * 4; + if bytes.len() != expected_len { + return Err(format!( + "weight payload size mismatch: header declares {count} params ({expected_len} bytes), file is {} bytes", + bytes.len() + )); + } + let mut weights = Vec::with_capacity(count); + for i in 0..count { + let o = WEIGHT_HEADER_LEN + i * 4; + weights.push(f32::from_le_bytes([ + bytes[o], + bytes[o + 1], + bytes[o + 2], + bytes[o + 3], + ])); + } + self.unflatten_weights(&weights) + } } +/// Magic prefix for AETHER weight files (see [`EmbeddingExtractor::save_weights`]). +const WEIGHT_MAGIC: &[u8; 8] = b"AETHERW1"; +/// 8-byte magic + 4-byte `u32` param count. +const WEIGHT_HEADER_LEN: usize = 12; + // ── CSI feature statistics ───────────────────────────────────────────────── /// Compute mean and variance of all values in a CSI feature matrix. @@ -1219,6 +1284,84 @@ mod tests { } } + // ── Weight save/load (ADR-185 §13.a) ──────────────────────────────── + + /// Deterministic, non-random weight pattern. Values are `k/65536 - 0.5` + /// with `k ∈ [0, 65535]`, i.e. multiples of 2⁻¹⁶ — exactly representable + /// in both f32 and f64 so a cross-language (Rust ↔ Python) fixture using + /// the same formula produces byte-identical weights. + fn deterministic_weights(n: usize) -> Vec { + (0..n) + .map(|i| { + let k = (i as u32).wrapping_mul(1_103_515_245).wrapping_add(12_345) % 65_536; + k as f32 / 65_536.0 - 0.5 + }) + .collect() + } + + #[test] + fn load_weights_actually_replaces_weights_and_round_trips() { + let mut ext = EmbeddingExtractor::new(small_config(), small_embed_config()); + let csi = make_csi(4, 16, 42); + let baseline = ext.extract(&csi); // random Xavier init + + // Build a source extractor with deterministic non-default weights and + // serialize it. + let det = deterministic_weights(ext.param_count()); + let mut src = EmbeddingExtractor::new(small_config(), small_embed_config()); + src.unflatten_weights(&det).unwrap(); + let src_emb = src.extract(&csi); + + let path = std::env::temp_dir() + .join(format!("aether_wtest_{}_{:p}.bin", std::process::id(), &ext)); + src.save_weights(&path).unwrap(); + + // Load into the random-init extractor. + ext.load_weights(&path).unwrap(); + let loaded_emb = ext.extract(&csi); + + // (1) The loaded weights are ACTUALLY used — output moved away from the + // random-init baseline (proves load is not a silent no-op). + let differs = baseline + .iter() + .zip(&loaded_emb) + .any(|(a, b)| (a - b).abs() > 1e-6); + assert!( + differs, + "load_weights had no effect: embedding still equals the random-init baseline" + ); + + // (2) It matches the source extractor whose weights we saved (round-trip). + for (a, b) in src_emb.iter().zip(&loaded_emb) { + assert!((a - b).abs() < 1e-6, "loaded embedding != source: {a} vs {b}"); + } + + // (3) The weights are bit-identical after the file round-trip. + assert_eq!(ext.flatten_weights(), det); + + std::fs::remove_file(&path).ok(); + } + + #[test] + fn load_weights_rejects_bad_magic_and_wrong_count() { + let mut ext = EmbeddingExtractor::new(small_config(), small_embed_config()); + let base = std::env::temp_dir().join(format!("aether_wbad_{}.bin", std::process::id())); + + // Bad magic. + std::fs::write(&base, b"NOPEMAGIC\x00\x00\x00").unwrap(); + assert!(ext.load_weights(&base).is_err()); + + // Right magic, wrong param count for this architecture. + let mut bad = Vec::new(); + bad.extend_from_slice(WEIGHT_MAGIC); + bad.extend_from_slice(&3u32.to_le_bytes()); + bad.extend_from_slice(&[0u8; 12]); // 3 f32s — won't match param_count + std::fs::write(&base, &bad).unwrap(); + assert!(ext.load_weights(&base).is_err()); + + std::fs::remove_file(&base).ok(); + } + // ── FingerprintIndex tests ────────────────────────────────────────── #[test]