feat(adr-185): add weight-loading capability to AETHER EmbeddingExtractor (§13.a)

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<f32>). 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
This commit is contained in:
ruv
2026-07-21 18:20:23 -07:00
parent e0cb7ed3b3
commit 65da488add
6 changed files with 377 additions and 0 deletions
+130
View File
@@ -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<Vec<f32>> {
let raw = fs::read_to_string(golden_dir().join("aether_input.json"))
.expect("read aether_input.json fixture");
let rows: Vec<Vec<f64>> = 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<f32> {
(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.");
}
}
}