fix: correct EMA RVQ exports and verify trained-model parity

.github/workflows/ci.yml

@@ -4,7 +4,6 @@ on:
   push:
     branches: [main]
   pull_request:
-    branches: [main]
 
 concurrency:
   group: ci-${{ github.ref }}

astro-site/src/content/docs/deployment.md

@@ -70,6 +70,43 @@ Use `load_codec()` for the common interface across families. The lower-level exa
 
 The manifest is the source of truth. The runtime reads it first, then resolves the encoder, optional decoder, and codebook files from the names listed there.
 
+## RVQ Codebook Export and Training Parity
+
+When exporting `vq.get_weights()`, declare `rvq_num_levels`, `rvq_use_ema`, and
+`rvq_kmeans_init` from the trained model's configuration. Plain RVQ has one
+codebook matrix per level. EMA RVQ also stores count vectors and embedding-sum
+matrices, plus an optional k-means flag; these are training state, not additional
+codebooks. Export rejects weights that do not match the declared layout. The
+standalone NPZ/header writers accept codebook matrices only; first call
+`extract_codebooks(weights, num_levels=..., use_ema=..., kmeans_init=...)` for EMA.
+
+Saved-run repackaging uses the run's `config.json` and loads positional weight
+arrays in numeric order. Preserve original packages and use `golden repackage
+<id> --output-dir <separate-directory>` when repairing an existing export.
+
+With sample inputs, `reference_vectors.npz` contains both deployed-path vectors
+and independent `source_latents`, `source_indices`, `source_quantized_latents`,
+and `source_reconstructions` from the restored training quantizer. Validation
+requires identical indices and quantized latents for identical source latents;
+the float encoder and decoder comparisons use `atol=rtol=1e-5`. This isolates codebook
+correctness from encoder conversion and precision differences. Sample-data
+reconstructions also use the discrete trained path rather than bypassing VQ.
+Training references run on CPU with full float32 operations, so GPU TF32 rounding
+cannot change the reference or fail the float32 conversion gate.
+For INT8-only decoder packages, the training comparison uses the exported float
+Keras decoder companion; the selected INT8 decoder is checked against its own
+deployed reference vectors. Float tolerances are not applied to INT8 outputs.
+
+Strict release validation rejects older RVQ references without these source
+arrays. Re-export them from the original training state; a self-consistent
+deployed reference alone does not establish training parity. Repackaging also
+recomputes model-card precision measurements instead of copying historical
+reports that may have used different codebooks.
+
+Correcting the exported level count changes packet interpretation. Keep existing
+packets paired with their original bundle and bind new packets to the corrected
+bundle revision; replacing tables inside an existing package is not a migration.
+
 ## Encoder Preprocessing and INT8 Parity
 
 Both RVQ encoders accept one normalized frame: PPG is `(1, 1, 320, 1)` at 64 Hz and ECG is `(1, 1, 512, 1)` at 256 Hz. Resample and frame the raw single-channel recording first, then normalize each frame independently:

compressionkit/experiments/repackage.py

@@ -138,7 +138,9 @@ def repackage_rvq_golden(
 
     import keras
 
+    from compressionkit.export.codebook import extract_codebooks, load_rvq_weights
     from compressionkit.export.deploy import export_for_deployment
+    from compressionkit.export.quantization import refresh_rvq_encoder_precision_reports
     from compressionkit.export.stimulus import export_stimulus_npz, generate_normalized_stimulus
 
     output_dir = output_dir or (run_dir / "deploy")
@@ -155,8 +157,18 @@ def repackage_rvq_golden(
 
     encoder = keras.models.load_model(encoder_path)
     decoder = keras.models.load_model(decoder_path)
-    rvq_npz = np.load(rvq_weights_path)
-    rvq_weights = [rvq_npz[k] for k in sorted(rvq_npz.files)]
+    rvq_weights = load_rvq_weights(rvq_weights_path)
+    # Use the saved run contract, not the current registry/config defaults.
+    saved_model = json.loads((run_dir / "config.json").read_text())["model"]
+    rvq_num_levels = saved_model["num_levels"]
+    rvq_use_ema = saved_model.get("use_ema", False)
+    rvq_kmeans_init = saved_model.get("kmeans_init", False)
+    codebooks = extract_codebooks(
+        rvq_weights, num_levels=rvq_num_levels, use_ema=rvq_use_ema, kmeans_init=rvq_kmeans_init
+    )
+    sizes = saved_model.get("codebook_sizes") or [saved_model["latent_width"]] * rvq_num_levels
+    if [cb.shape for cb in codebooks] != [(size, saved_model["embedding_dim"]) for size in sizes]:
+        raise ValueError("Saved RVQ codebook shapes do not match the training config")
 
     input_shape = encoder.input_shape
     if len(input_shape) in (3, 4):
@@ -181,20 +193,6 @@ def repackage_rvq_golden(
         "license": "other",
         "preprocessing_contract": preprocessing_contract,
     }
-    precision_paths = [
-        run_dir.parent / "precision-study" / f"{experiment.experiment_id}.json",
-        run_dir.parent / "rvq_encoder_precision_comparison.json",
-    ]
-    for precision_path in precision_paths:
-        if not precision_path.is_file():
-            continue
-        precision_payload = json.loads(precision_path.read_text())
-        variants = precision_payload.get("experiments", {}).get(experiment.experiment_id, {}).get("variants", {})
-        if isinstance(variants, dict):
-            model_card_info["encoder_precision_report"] = {
-                name: payload.get("report", {}) for name, payload in variants.items() if isinstance(payload, dict)
-            }
-            break
     if scorecard_payload is not None:
         model_card_info["scorecard_summary"] = scorecard_payload
 
@@ -206,28 +204,28 @@ def repackage_rvq_golden(
         sample_rate=experiment.sample_rate,
         epsilon=preprocessing_contract["normalization"]["epsilon"],
     )
-    latents = encoder.predict(sample_inputs, verbose=0)
-    sample_reconstructions = decoder.predict(latents, verbose=0)
-    if isinstance(sample_reconstructions, dict):
-        sample_reconstructions = sample_reconstructions.get("reconstruction", sample_reconstructions.get("output"))
 
     artifacts = export_for_deployment(
         encoder=encoder,
         decoder=decoder,
         rvq_weights=rvq_weights,
+        rvq_num_levels=rvq_num_levels,
+        rvq_use_ema=rvq_use_ema,
+        rvq_kmeans_init=rvq_kmeans_init,
         rep_dataset=rep_dataset,
         quantization_validation_dataset=quantization_validation_dataset,
         output_dir=output_dir,
         sample_inputs=sample_inputs,
         sample_targets=sample_inputs,
-        sample_reconstructions=np.asarray(sample_reconstructions, dtype=np.float32),
         model_name=experiment.run_name,
         export_decoder_float32=True,
         export_decoder_int8=export_decoder_int8,
         model_card_info=model_card_info,
     )
     quantization_report = json.loads((output_dir / "quantization_report.json").read_text())
     report_metrics = quantization_report["metrics"]
+    # Historical reports may have used different (including faulty) codebooks.
+    refresh_rvq_encoder_precision_reports(output_dir, quantization_validation_dataset)
     export_stimulus_npz(
         modality=experiment.modality,
         output_path=output_dir / "sample_stimulus.npz",

compressionkit/export/codebook.py

@@ -15,25 +15,77 @@
 logger = logging.getLogger(__name__)
 
 
-def extract_codebooks(rvq_weights: list[np.ndarray]) -> list[np.ndarray]:
+def extract_codebooks(
+    rvq_weights: list[np.ndarray],
+    *,
+    num_levels: int | None = None,
+    use_ema: bool = False,
+    kmeans_init: bool = False,
+) -> list[np.ndarray]:
     """Extract per-level codebook matrices from RVQ weight list.
 
-    The weight list from ``rvq.get_weights()`` contains one embedding matrix
-    per RVQ level in order, each with shape ``(num_embeddings, embedding_dim)``.
+    Plain RVQ stores one matrix per level. EMA RVQ stores a codebook,
+    count vector, and embedding-sum matrix per level, followed by an optional
+    scalar k-means flag. The layout must be declared; matrix rank alone
+    cannot distinguish codebooks from EMA accumulators.
 
     Args:
         rvq_weights: List of weight arrays from ``rvq.get_weights()``.
+        num_levels: Expected number of levels. Required for EMA checkpoints.
+        use_ema: Whether the source is an EMA RVQ checkpoint.
+        kmeans_init: Whether the EMA checkpoint includes its warm-start flag.
 
     Returns:
         List of codebook arrays, one per level.
+
+    Raises:
+        ValueError: The weights do not match the declared layout.
     """
-    codebooks = []
-    for w in rvq_weights:
-        if w.ndim == 2:
-            codebooks.append(w)
+    if use_ema and num_levels is None:
+        raise ValueError("EMA codebook extraction requires num_levels")
+    if kmeans_init and not use_ema:
+        raise ValueError("kmeans_init requires an EMA checkpoint")
+    levels = len(rvq_weights) if num_levels is None else num_levels
+    if levels < 1:
+        raise ValueError("At least one RVQ level is required")
+    stride = 3 if use_ema else 1
+    expected = stride * levels + int(kmeans_init)
+    if len(rvq_weights) != expected:
+        raise ValueError(f"Expected {expected} RVQ weight arrays for {levels} levels, got {len(rvq_weights)}")
+    codebooks: list[np.ndarray] = []
+    for level in range(levels):
+        cb = np.asarray(rvq_weights[stride * level])
+        if cb.ndim != 2 or min(cb.shape) < 1 or not np.issubdtype(cb.dtype, np.floating):
+            raise ValueError(f"Level {level}: expected a nonempty floating-point codebook matrix")
+        if not np.isfinite(cb).all():
+            raise ValueError(f"Level {level}: non-finite codebook values")
+        if codebooks and cb.shape[1] != codebooks[0].shape[1]:
+            raise ValueError("RVQ codebooks must share an embedding dimension")
+        if use_ema:
+            count, total = (np.asarray(w) for w in rvq_weights[stride * level + 1 : stride * level + 3])
+            if count.shape != (cb.shape[0],) or total.shape != cb.shape:
+                raise ValueError(f"Level {level}: EMA count/sum shapes do not match the codebook")
+        codebooks.append(cb)
+    if kmeans_init:
+        flag = np.asarray(rvq_weights[-1])
+        if flag.shape != () or flag.item() not in (0, 1):
+            raise ValueError("Expected a scalar 0/1 kmeans_done flag")
     return codebooks
 
 
+def load_rvq_weights(path: str | Path) -> list[np.ndarray]:
+    """Read a positional RVQ training archive in numeric weight order.
+
+    Reject missing or unexpected keys instead of silently reordering or
+    skipping state. Lexical sorting would place ``arr_10`` before ``arr_2``.
+    """
+    with np.load(path, allow_pickle=False) as archive:
+        keys = [f"arr_{i}" for i in range(len(archive.files))]
+        if set(archive.files) != set(keys):
+            raise ValueError("RVQ training archive must contain contiguous arr_0 ... arr_N keys")
+        return [archive[key] for key in keys]
+
+
 def _format_c_array(
     name: str,
     data: np.ndarray,
@@ -48,7 +100,10 @@ def _format_c_array(
     chunk_size = 8
     for i in range(0, len(flat), chunk_size):
         chunk = flat[i : i + chunk_size]
-        vals = ", ".join(f"{v:.8g}f" for v in chunk)
+        # Nine significant digits round-trip float32. C requires a decimal
+        # point or exponent before the suffix (``0f`` is not a valid literal).
+        literals = [f"{v:.9g}" for v in chunk]
+        vals = ", ".join(f"{v if '.' in v or 'e' in v else v + '.0'}f" for v in literals)
         comma = "," if i + chunk_size < len(flat) else ""
         rows.append(f"    {vals}{comma}")
     rows.append("};")
@@ -62,7 +117,8 @@ def export_codebooks_npz(
     """Save RVQ codebook tables to a NumPy ``.npz`` archive.
 
     Args:
-        rvq_weights: Weight list from ``rvq.get_weights()``.
+        rvq_weights: Codebook matrices only; use ``extract_codebooks`` first
+            for an EMA checkpoint.
         output_path: Path for the ``.npz`` file.
 
     Returns:
@@ -87,7 +143,8 @@ def export_codebooks_header(
     """Export RVQ codebook tables as a C header file.
 
     Args:
-        rvq_weights: Weight list from ``rvq.get_weights()``.
+        rvq_weights: Codebook matrices only; use ``extract_codebooks`` first
+            for an EMA checkpoint.
         output_path: Path for the ``.h`` file.
         prefix: Prefix for C array names (e.g. ``rvq_codebook_l0``).
         guard: Include guard macro name; auto-generated if ``None``.
@@ -99,6 +156,9 @@ def export_codebooks_header(
     output_path.parent.mkdir(parents=True, exist_ok=True)
     codebooks = extract_codebooks(rvq_weights)
 
+    if any(cb.shape != codebooks[0].shape for cb in codebooks):
+        raise ValueError("The C header requires equally sized RVQ codebooks")
+
     if guard is None:
         guard = output_path.stem.upper().replace("-", "_").replace(".", "_") + "_H"
 
@@ -128,4 +188,5 @@ def export_codebooks_header(
     "export_codebooks_header",
     "export_codebooks_npz",
     "extract_codebooks",
+    "load_rvq_weights",
 ]

compressionkit/export/deploy.py

@@ -159,6 +159,9 @@ def export_for_deployment(
     decoder: keras.Model,
     rvq_weights: list[np.ndarray],
     *,
+    rvq_num_levels: int | None = None,
+    rvq_use_ema: bool = False,
+    rvq_kmeans_init: bool = False,
     rep_dataset: np.ndarray,
     quantization_validation_dataset: np.ndarray | None = None,
     output_dir: str | Path,
@@ -184,13 +187,17 @@ def export_for_deployment(
         encoder: Trained Keras encoder model.
         decoder: Trained Keras decoder model.
         rvq_weights: Weight list from ``rvq.get_weights()``.
+        rvq_num_levels: Expected trained RVQ level count; required for EMA.
+        rvq_use_ema: Whether weights use the EMA codebook/count/sum layout.
+        rvq_kmeans_init: Whether EMA weights include a k-means warm-start flag.
         rep_dataset: Representative input dataset for TFLite calibration.
         quantization_validation_dataset: Disjoint preprocessed frames for
             post-export INT8-vs-float32 parity validation. Required for INT8.
         output_dir: Directory to write all deployment artifacts.
         sample_inputs: Optional sample input frames for validation.
         sample_targets: Optional sample target frames for validation.
-        sample_reconstructions: Optional model reconstructions of targets.
+        sample_reconstructions: Optional reconstructions when inputs are absent.
+            With sample inputs, recomputed through the trained discrete RVQ path.
         quantization: Quantization mode for TFLite conversion.
         io_type: I/O type string for TFLite conversion.
         codebook_prefix: Prefix for codebook C array names.
@@ -207,6 +214,29 @@ def export_for_deployment(
         ``DeploymentArtifacts`` with paths to all exported files.
     """
     output_dir = Path(output_dir)
+    codebooks = extract_codebooks(
+        rvq_weights, num_levels=rvq_num_levels, use_ema=rvq_use_ema, kmeans_init=rvq_kmeans_init
+    )
+    if any(cb.shape != codebooks[0].shape for cb in codebooks):
+        raise ValueError("Deployment currently requires equally sized RVQ codebooks")
+    if codebooks[0].shape[1] != encoder.output_shape[-1] or codebooks[0].shape[1] != decoder.input_shape[-1]:
+        raise ValueError("RVQ embedding dimension does not match encoder output and decoder input")
+    source_reference = None
+    if sample_inputs is not None:
+        from compressionkit.export.rvq_reference import build_rvq_reference
+
+        source_reference = build_rvq_reference(
+            encoder,
+            decoder,
+            rvq_weights,
+            sample_inputs,
+            num_levels=len(codebooks),
+            use_ema=rvq_use_ema,
+            kmeans_init=rvq_kmeans_init,
+        )
+        # Sample reconstructions must represent the discrete trained path,
+        # not decoder(encoder(x)) with the quantizer bypassed.
+        sample_reconstructions = source_reference["source_reconstructions"]
     if quantization == "INT8" and quantization_validation_dataset is None:
         raise ValueError("INT8 export requires a disjoint quantization_validation_dataset")
     output_dir.mkdir(parents=True, exist_ok=True)
@@ -324,11 +354,11 @@ def export_for_deployment(
     # 3. Codebook tables
     logger.info("Exporting codebook tables...")
     artifacts.codebook_npz = export_codebooks_npz(
-        rvq_weights,
+        codebooks,
         output_dir / "codebook.npz",
     )
     artifacts.codebook_header = export_codebooks_header(
-        rvq_weights,
+        codebooks,
         output_dir / "codebook.h",
         prefix=codebook_prefix,
     )
@@ -363,7 +393,6 @@ def export_for_deployment(
         artifacts.scorecard = write_scorecard_artifact(output_dir, scorecard_summary)
 
     # 5. Canonical runtime hydration spec
-    codebooks = extract_codebooks(rvq_weights)
     preprocessing_contract = dict(model_card_info.get("preprocessing_contract", {})) if model_card_info else {}
     codec_spec = {
         "family": "rvq",
@@ -530,6 +559,7 @@ def export_for_deployment(
                 "indices": encoded_indices,
                 "reconstructions": decoded_frames,
             }
+            reference_payload.update(source_reference)
             if sample_targets is not None:
                 reference_payload["targets"] = np.asarray(sample_targets, dtype=np.float32)
             np.savez_compressed(artifacts.reference_vectors, **reference_payload)
@@ -540,6 +570,16 @@ def export_for_deployment(
             )
             artifacts.reference_vectors = Path()
 
+    if artifacts.reference_vectors != Path():
+        from compressionkit.export.validate import _validate_reference_vectors
+
+        parity_errors: list[str] = []
+        _validate_reference_vectors(
+            output_dir, "rvq", parity_errors, [], max_vectors=len(sample_inputs), require_training_reference=True
+        )
+        if parity_errors:
+            raise ValueError("RVQ trained-model parity failed: " + "; ".join(parity_errors))
+
     # 7b. Compare the deployed INT8 encoder against its float32 companion on
     # a disjoint, real preprocessed validation partition.
     if quantization == "INT8":