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
+28
View File
@@ -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)
}