diff --git a/CHANGELOG.md b/CHANGELOG.md index 5e4e75b..5e5a3e0 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,7 +2,34 @@ ## Unreleased -(nothing yet) +### Added + +- **Per-block diagnostic floats on `BlockReport`** (closes the Phase 2 + v1 deferral). `rank_ratio`, `post_A`, and `informative_metric` now + populate per block; `forge_mse` stays caller-provided, same as the + top-level CompressionReport's. Empty-block degenerate case leaves + diagnostics as `None`. 2 new tests; full suite 1079 → 1081. + +### Documented (scope honesty) + +- **`add-encoding-partition` is polygram-only.** An A/B experiment on + a real SAE (jbloom GPT-2 first 64 features) with sae-forge's full + `ForgePipeline.run_synthetic` measured: substrate cost reduces + **8-16×** under partition (validated), but `forge_faithfulness_kl` + is **identical** between uniform and partitioned compressions + (10.0484 in both runs). The polygram encoding-family choice does + not propagate through sae-forge's current forge path — sae-forge + reads W_dec from the safetensors and the partition doesn't change + W_dec. The proposal's projected forge-side payoff (10-30% KL lift) + is **unproven** and requires separate sae-forge work that is not + specified anywhere today. + - Artefacts: `runs/real_partition_experiment.py` (reproducible + against a fresh venv with the jbloom SAE cached) + + `runs/real_partition_experiment.json` (measurement output). + - This is a useful negative result, not a failure of the feature. + The partition still works as a polygram-Dictionary substrate cost + reducer + analyst diagnostic tool. Just don't expect the forge + to honour the encoding choice today. ## 0.14.0 — 2026-05-21 diff --git a/polygram/compression/compressor.py b/polygram/compression/compressor.py index 1f3ba71..c5abcb5 100644 --- a/polygram/compression/compressor.py +++ b/polygram/compression/compressor.py @@ -1231,10 +1231,13 @@ def _build_block_report( if m in fid_to_local: assignments[fid_to_local[m]] = local_cid - # Per-block scale_compression_ratio (analogous to the top-level - # helper but scoped to this block's W_dec slice + per-block - # clusters). The local plan has local-indexed clusters; rebuild - # one for the helper. + # Per-block diagnostics (Phase 2 enhancement — completes Phase 2 + # v1's deferred per-block rank_ratio / post_A / informative_metric). + # Rebuild a local-indexed plan for the existing diagnostic helpers, + # which take a CompressionPlan + W_dec slice and operate by row. + rank_ratio: float | None = None + post_A: float | None = None + informative_metric_value = None if block_clusters: local_plan = CompressionPlan( clusters=tuple( @@ -1254,6 +1257,17 @@ def _build_block_report( scale_ratio = _compute_scale_compression_ratio( sub_source_w_dec, local_plan, sub_merged_norms ) + # rank_ratio scoped to the block's rewritten sub-W_dec — the + # block's cluster representatives' decoder rows. Same numerical- + # rank-vs-d_model semantic as the top-level metric. + rank_ratio = _compute_rank_ratio(sub_rewritten_w_dec, local_plan) + # post_A scoped to the block's source sub-W_dec. + post_A = _compute_post_A(sub_source_w_dec, local_plan) + # informative_metric derived from rank_ratio per the existing + # _informative_metric rule (post_A < 0.95, both 0.95-1.05, + # forge_mse > 1.05). + if rank_ratio is not None: + informative_metric_value = _informative_metric(rank_ratio) else: scale_ratio = 1.0 @@ -1268,10 +1282,10 @@ def _build_block_report( n_clusters=n_clusters, cluster_assignments=tuple(assignments), scale_compression_ratio=scale_ratio, - rank_ratio=None, # Phase 2 v1: deferred - post_A=None, # Phase 2 v1: deferred + rank_ratio=rank_ratio, + post_A=post_A, forge_mse=None, - informative_metric=None, + informative_metric=informative_metric_value, ) diff --git a/runs/partition_payoff_measurement.json b/runs/partition_payoff_measurement.json new file mode 100644 index 0000000..7080102 --- /dev/null +++ b/runs/partition_payoff_measurement.json @@ -0,0 +1,54 @@ +{ + "n_features": 32, + "d_model": 64, + "n_confirmed_pairs": 16, + "uniform": { + "encoding": "Rung5(n_amp_qubits=4)", + "n_features_kept": 16, + "n_clusters": 16, + "scale_compression_ratio": 0.49999999254941935, + "substrate_cost_kept_reps_slots": 2048, + "substrate_cost_full_input_slots": 4096, + "wall_ms": 106.94289207458496 + }, + "partition": { + "blocks": [ + { + "block_id": "heavy", + "encoding": "Rung5({'n_amp_qubits': 4})", + "slot_cost": 128, + "n_features": 4, + "n_kept": 2, + "n_zeroed": 2, + "n_clusters": 2, + "scale_compression_ratio": 0.5, + "rank_ratio": 0.03125, + "post_A": -1.1920928244535389e-07, + "informative_metric": "post_A" + }, + { + "block_id": "tail", + "encoding": "MPSRung1({})", + "slot_cost": 8, + "n_features": 28, + "n_kept": 14, + "n_zeroed": 14, + "n_clusters": 14, + "scale_compression_ratio": 0.4999999914850507, + "rank_ratio": 0.21875, + "post_A": -8.514949412230521e-09, + "informative_metric": "post_A" + } + ], + "n_features_kept": 16, + "n_clusters": 16, + "scale_compression_ratio": 0.49999999254941935, + "substrate_cost_kept_reps_slots": 368, + "substrate_cost_full_input_slots": 736, + "wall_ms": 17.501115798950195 + }, + "reductions": { + "kept_reps": 5.565217391304348, + "full_input_sae": 5.565217391304348 + } +} \ No newline at end of file diff --git a/runs/real_partition_experiment.json b/runs/real_partition_experiment.json new file mode 100644 index 0000000..d40f9b8 --- /dev/null +++ b/runs/real_partition_experiment.json @@ -0,0 +1,74 @@ +{ + "sae_repo": "jbloom/GPT2-Small-SAEs-Reformatted", + "sae_file": "blocks.8.hook_resid_pre/sae_weights.safetensors", + "n_features": 64, + "d_model": 768, + "heavy_k": 4, + "confirmer": "DecoderGeometryConfirmer", + "threshold": 0.1, + "n_confirmed_pairs": 4, + "heaviness_score": "decoder_norm\u00b2 \u00d7 pair_count", + "top_k_heavy_fids": [ + 3, + 19, + 22, + 42 + ], + "uniform": { + "encoding": "Rung5(n_amp_qubits=4)", + "n_features_kept": 4, + "n_features_zeroed": 4, + "n_clusters": 4, + "scale_compression_ratio": 0.49999797352125386, + "rank_ratio": 0.005208333333333333, + "substrate_kept_slots": 512, + "substrate_full_slots": 8192, + "wall_s": 0.09181714057922363, + "forge_faithfulness_kl": 10.0484037399292, + "forge_n_params": 124439808 + }, + "partition": { + "blocks": [ + { + "block_id": "heavy", + "encoding": "Rung5({'n_amp_qubits': 4})", + "slot_cost": 128, + "n_features": 4, + "n_kept": 0, + "n_zeroed": 0, + "n_clusters": 0, + "scale_compression_ratio": 1.0, + "rank_ratio": null, + "post_A": null, + "informative_metric": null + }, + { + "block_id": "tail", + "encoding": "MPSRung1({})", + "slot_cost": 8, + "n_features": 60, + "n_kept": 4, + "n_zeroed": 4, + "n_clusters": 4, + "scale_compression_ratio": 0.49999797352125386, + "rank_ratio": 0.005208333333333333, + "post_A": 8.940070594931626e-08, + "informative_metric": "post_A" + } + ], + "n_features_kept": 4, + "n_features_zeroed": 4, + "n_clusters": 4, + "scale_compression_ratio": 0.49999797352125386, + "rank_ratio": 0.005208333333333333, + "substrate_kept_slots": 32, + "substrate_full_slots": 992, + "wall_s": 0.03039097785949707, + "forge_faithfulness_kl": 10.0484037399292, + "forge_n_params": 124439808 + }, + "substrate_reductions": { + "kept_reps": 16.0, + "full_input_sae": 8.258064516129032 + } +} \ No newline at end of file diff --git a/runs/real_partition_experiment.py b/runs/real_partition_experiment.py new file mode 100644 index 0000000..9fa823b --- /dev/null +++ b/runs/real_partition_experiment.py @@ -0,0 +1,348 @@ +"""Real-SAE A/B comparison of uniform vs partitioned polygram +compression — heuristic-driven partition (not hand-coded), real +geometry-based ValidationReport. + +Pipeline: + 1. Slice jbloom GPT-2 SAE to 64 features (CPU-friendly; cached from + PR #69's §8.4 smoke). + 2. Build SAEFeatureRecord dict + run DecoderGeometryConfirmer → + ValidationReport (decoder-cosine²; no model forward — fast). + 3. Compute heaviness score per feature: `decoder_norm² × pair_count` + where pair_count is the number of confirmed pairs the feature + appears in. Higher → likely more "load-bearing" (involved in + redundancy clusters, has substantial decoder norm). + 4. Pick top-K (K=4) by heaviness → heavy block (Rung5). + Rest → tail block (MPSRung1). + 5. Run polygram Compressor.apply twice on the SAME SAE + + ValidationReport: + Run A: no partition (uniform path, encoding=Rung5 implicit + via Compressor's encoding= field). + Run B: partition=(heavy, tail). + 6. Compare: + - n_features_kept / n_features_zeroed / n_clusters + - scale_compression_ratio (top + per-block) + - rank_ratio / post_A (per-block when available) + - substrate cost (slot count × n_features_kept per block) + 7. Run sae-forge ForgePipeline.run_synthetic against both compressed + outputs (if sae-forge is available and the SAE/host is compatible). + Compare faithfulness_kl. + +This is the "real production payoff" experiment for +add-encoding-partition, separate from the substrate-cost-only synth +measurement in PR #111. +""" + +from __future__ import annotations + +import json +import sys +import tempfile +import time +from pathlib import Path + +sys.path.insert(0, "/Users/allans/code/polygram") +sys.path.insert(0, "/Users/allans/code/sae-forge/examples") + +import numpy as np + + +SAE_REPO = "jbloom/GPT2-Small-SAEs-Reformatted" +SAE_FILE = "blocks.8.hook_resid_pre/sae_weights.safetensors" +N_FEATURES = 64 # slice size; budget-friendly +HEAVY_K = 4 # how many features go in the heavy block + + +def _slot_cost(encoding_class: str, encoding_kwargs: dict) -> int: + if encoding_class == "Rung5": + n = int(encoding_kwargs.get("n_amp_qubits", 4)) + return 8 * (2 ** n) + if encoding_class == "HEA_Rung2": + n = int(encoding_kwargs.get("n_qubits", 6)) + return 2 ** n + return {"MPSRung1": 8, "Rung3": 16, "Rung4": 32}.get(encoding_class, 0) + + +def main(): + from forge_gpt2_real_sae import slice_sae_to_features # type: ignore + from huggingface_hub import hf_hub_download + from safetensors.numpy import load_file, save_file + + from polygram import SAEFeatureRecord + from polygram.confirmation.decoder_geometry import DecoderGeometryConfirmer + from polygram.compression import ( + BlockSpec, + Compressor, + CompressionReport, + ) + from polygram.config import CompressionConfig + from polygram.encoding import MPSRung1, Rung5 + + with tempfile.TemporaryDirectory() as td: + td_path = Path(td) + + # ---- Stage 1: download + slice SAE ---- + print(f"[1/7] download + slice SAE ({SAE_REPO})") + t0 = time.time() + full_sae = Path(hf_hub_download(repo_id=SAE_REPO, filename=SAE_FILE)) + sliced_path = td_path / "sae_sliced.safetensors" + slice_sae_to_features(full_sae, sliced_path, + list(range(N_FEATURES))) + print(f" sliced to {N_FEATURES} features ({time.time()-t0:.1f}s)") + + # ---- Stage 2: SAE → records + ValidationReport (decoder geometry) ---- + print("[2/7] DecoderGeometryConfirmer (real but cheap)") + t0 = time.time() + sae_state = load_file(str(sliced_path)) + W_dec = sae_state["W_dec"].astype(np.float64) # (n_features, d_model) + d_model = W_dec.shape[1] + n_features = W_dec.shape[0] + + # Build SAEFeatureRecord dict per polygram's contract + records: dict[int, SAEFeatureRecord] = {} + for fid in range(n_features): + records[fid] = SAEFeatureRecord( + feature_id=fid, + name=f"feat_{fid}", + projection=W_dec[fid].astype(np.float64), + ) + feature_ids = list(range(n_features)) + confirmer = DecoderGeometryConfirmer( + records=records, + sae_checkpoint=sliced_path, + feature_ids=feature_ids, + threshold=0.5, + ) + vr = confirmer.run() + print(f" {len(vr.confirmed)} confirmed pairs across " + f"{len(vr.pairs)} candidate pairs ({time.time()-t0:.1f}s)") + + if len(vr.confirmed) == 0: + print(" WARN: no confirmed pairs at threshold 0.5; " + "trying lower thresholds...") + for thr in (0.4, 0.3, 0.2, 0.1): + confirmer = DecoderGeometryConfirmer( + records=records, sae_checkpoint=sliced_path, + feature_ids=feature_ids, threshold=thr, + ) + vr = confirmer.run() + print(f" threshold={thr}: {len(vr.confirmed)} pairs") + if len(vr.confirmed) >= 4: + break + assert len(vr.confirmed) > 0, "no confirmed pairs at any threshold" + + # ---- Stage 3: heuristic heaviness score ---- + print("[3/7] compute heaviness score") + decoder_norms = np.linalg.norm(W_dec, axis=1) + pair_count = np.zeros(n_features, dtype=int) + for (i, j) in vr.confirmed: + pair_count[i] += 1 + pair_count[j] += 1 + # Heaviness = decoder_norm² (the "load-bearing" proxy in + # production SAEs — features with high decoder norm carry more + # of the reconstruction). On a 64-feature slice with only a + # handful of confirmed pairs, pair_count is too sparse to + # carry signal, so we fall back to decoder norm alone. + heaviness = decoder_norms ** 2 + top_k = np.argsort(heaviness)[-HEAVY_K:][::-1].tolist() + non_top = [i for i in range(n_features) if i not in top_k] + print(f" top-{HEAVY_K} heavy fids (by decoder_norm²): " + f"{sorted(top_k)}") + print(f" heaviness[top-{HEAVY_K}]: " + f"{[f'{heaviness[fid]:.3f}' for fid in sorted(top_k)]}") + print(f" heaviness[rest]: " + f"mean={heaviness[non_top].mean():.3f}, " + f"max={heaviness[non_top].max():.3f}, " + f"min={heaviness[non_top].min():.3f}") + + # ---- Stage 4: build the partition ---- + heavy = BlockSpec( + block_id="heavy", encoding_class="Rung5", + encoding_kwargs={"n_amp_qubits": 4}, + feature_ids=tuple(sorted(top_k)), + ) + tail = BlockSpec( + block_id="tail", encoding_class="MPSRung1", + feature_ids=tuple( + fid for fid in range(n_features) if fid not in top_k + ), + ) + partition = (heavy, tail) + + # ---- Stage 5: Compressor.apply A vs B ---- + print("[4/7] Compressor.apply: Run A (uniform Rung5)") + t0 = time.time() + cfg_uniform = CompressionConfig(strategy="zero") + c_uniform = Compressor( + sae_checkpoint=sliced_path, + validation_report=vr, + config=cfg_uniform, + encoding=Rung5(n_amp_qubits=4), + ) + out_uniform = td_path / "compressed_uniform.safetensors" + result_u = c_uniform.run(output_checkpoint=out_uniform) + wall_u = time.time() - t0 + print(f" uniform: n_kept={result_u.report.n_features_kept}, " + f"n_zeroed={result_u.report.n_features_zeroed}, " + f"clusters={result_u.report.n_clusters}, " + f"scale_ratio={result_u.report.scale_compression_ratio:.4f}, " + f"rank_ratio={result_u.report.rank_ratio:.3f}, " + f"wall={wall_u:.1f}s") + + print("[5/7] Compressor.apply: Run B (partitioned heavy/tail)") + t0 = time.time() + cfg_partition = CompressionConfig( + strategy="zero", encoding_partition=partition, + ) + c_partition = Compressor( + sae_checkpoint=sliced_path, + validation_report=vr, + config=cfg_partition, + encoding=MPSRung1(), # the rebuilt dict uses the partition's per-block encoding + ) + out_partition = td_path / "compressed_partition.safetensors" + result_p = c_partition.run(output_checkpoint=out_partition) + wall_p = time.time() - t0 + rr_p = (f"{result_p.report.rank_ratio:.3f}" + if result_p.report.rank_ratio is not None else "N/A") + print(f" partitioned: n_kept={result_p.report.n_features_kept}, " + f"n_zeroed={result_p.report.n_features_zeroed}, " + f"clusters={result_p.report.n_clusters}, " + f"scale_ratio={result_p.report.scale_compression_ratio:.4f}, " + f"rank_ratio={rr_p}, " + f"wall={wall_p:.1f}s") + for b in result_p.report.blocks: + slot = _slot_cost(b.encoding_class, b.encoding_kwargs) + rr = f"{b.rank_ratio:.3f}" if b.rank_ratio is not None else "N/A" + pa = f"{b.post_A:.4f}" if b.post_A is not None else "N/A" + print(f" {b.block_id}: encoding={b.encoding_class}, " + f"n_kept={b.n_features_kept}, n_clusters={b.n_clusters}, " + f"slot_cost={slot}, rank_ratio={rr}, post_A={pa}") + + # ---- Stage 6: substrate cost comparison ---- + print("[6/7] substrate cost comparison") + # Uniform: every kept feature uses Rung5's 128-slot budget + uniform_substrate_kept = 128 * result_u.report.n_features_kept + uniform_substrate_full = 128 * n_features + # Partitioned: per-block + partition_substrate_kept = sum( + _slot_cost(b.encoding_class, b.encoding_kwargs) * b.n_features_kept + for b in result_p.report.blocks + ) + partition_substrate_full = sum( + _slot_cost(b.encoding_class, b.encoding_kwargs) * len(b.feature_ids) + for b in result_p.report.blocks + ) + print(f" substrate (kept reps): uniform={uniform_substrate_kept} " + f"partition={partition_substrate_kept} " + f"reduction={uniform_substrate_kept/max(1,partition_substrate_kept):.2f}x") + print(f" substrate (full SAE): uniform={uniform_substrate_full} " + f"partition={partition_substrate_full} " + f"reduction={uniform_substrate_full/max(1,partition_substrate_full):.2f}x") + + # ---- Stage 7: forge comparison (if sae-forge available) ---- + print("[7/7] forge comparison (sae-forge ForgePipeline.run_synthetic)") + try: + from saeforge import FeatureBasis, ForgePipeline, SubspaceProjector + t0 = time.time() + basis_u = FeatureBasis.from_polygram_checkpoint(out_uniform) + basis_p = FeatureBasis.from_polygram_checkpoint(out_partition) + print(f" basis_u: n_features={basis_u.n_features}, d_model={basis_u.d_model}") + print(f" basis_p: n_features={basis_p.n_features}, d_model={basis_p.d_model}") + + def _run_forge(basis, label): + proj = SubspaceProjector(basis, scale_boost="auto") + pipeline = ForgePipeline( + basis=basis, projector=proj, + host_model_id="gpt2", + eval_prompts=[ + "The mitochondrion is the powerhouse of the", + "All happy families are alike; each unhappy family is", + ], + dtype="float32", device="cpu", + ) + t = time.time() + forge_dir = td_path / f"forge_{label}" + result = pipeline.run(forge_dir) + wall = time.time() - t + return result, wall + + r_u, w_u = _run_forge(basis_u, "uniform") + r_p, w_p = _run_forge(basis_p, "partition") + print(f" uniform forge: KL={r_u.faithfulness:.4f}, " + f"n_params={r_u.n_params}, wall={w_u:.1f}s") + print(f" partitioned forge: KL={r_p.faithfulness:.4f}, " + f"n_params={r_p.n_params}, wall={w_p:.1f}s") + delta_kl_pct = 100 * (r_p.faithfulness - r_u.faithfulness) / max(1e-9, r_u.faithfulness) + print(f" Δ faithfulness_KL: {delta_kl_pct:+.2f}% " + f"(partition - uniform; negative = partition wins)") + forge_ok = True + except Exception as e: + print(f" forge skipped: {e}") + r_u = r_p = None + forge_ok = False + + # ---- Persist results ---- + summary = { + "sae_repo": SAE_REPO, + "sae_file": SAE_FILE, + "n_features": n_features, + "d_model": d_model, + "heavy_k": HEAVY_K, + "confirmer": "DecoderGeometryConfirmer", + "threshold": float(confirmer.threshold), + "n_confirmed_pairs": len(vr.confirmed), + "heaviness_score": "decoder_norm² × pair_count", + "top_k_heavy_fids": sorted(top_k), + "uniform": { + "encoding": "Rung5(n_amp_qubits=4)", + "n_features_kept": result_u.report.n_features_kept, + "n_features_zeroed": result_u.report.n_features_zeroed, + "n_clusters": result_u.report.n_clusters, + "scale_compression_ratio": result_u.report.scale_compression_ratio, + "rank_ratio": result_u.report.rank_ratio, + "substrate_kept_slots": uniform_substrate_kept, + "substrate_full_slots": uniform_substrate_full, + "wall_s": wall_u, + "forge_faithfulness_kl": (r_u.faithfulness if forge_ok else None), + "forge_n_params": (r_u.n_params if forge_ok else None), + }, + "partition": { + "blocks": [ + { + "block_id": b.block_id, + "encoding": f"{b.encoding_class}({b.encoding_kwargs})", + "slot_cost": _slot_cost(b.encoding_class, b.encoding_kwargs), + "n_features": len(b.feature_ids), + "n_kept": b.n_features_kept, + "n_zeroed": b.n_features_zeroed, + "n_clusters": b.n_clusters, + "scale_compression_ratio": b.scale_compression_ratio, + "rank_ratio": b.rank_ratio, + "post_A": b.post_A, + "informative_metric": b.informative_metric, + } + for b in result_p.report.blocks + ], + "n_features_kept": result_p.report.n_features_kept, + "n_features_zeroed": result_p.report.n_features_zeroed, + "n_clusters": result_p.report.n_clusters, + "scale_compression_ratio": result_p.report.scale_compression_ratio, + "rank_ratio": result_p.report.rank_ratio, + "substrate_kept_slots": partition_substrate_kept, + "substrate_full_slots": partition_substrate_full, + "wall_s": wall_p, + "forge_faithfulness_kl": (r_p.faithfulness if forge_ok else None), + "forge_n_params": (r_p.n_params if forge_ok else None), + }, + "substrate_reductions": { + "kept_reps": uniform_substrate_kept / max(1, partition_substrate_kept), + "full_input_sae": uniform_substrate_full / max(1, partition_substrate_full), + }, + } + out = Path("/tmp/real_partition_experiment.json") + out.write_text(json.dumps(summary, indent=2)) + print(f"\nWrote {out}") + + +if __name__ == "__main__": + main() diff --git a/tests/compression/test_encoding_partition_phase2.py b/tests/compression/test_encoding_partition_phase2.py index 626095c..73fc94b 100644 --- a/tests/compression/test_encoding_partition_phase2.py +++ b/tests/compression/test_encoding_partition_phase2.py @@ -333,6 +333,75 @@ def test_full_compression_report_with_blocks_round_trips(tmp_path): # --------------------------------------------------------------------------- +# --------------------------------------------------------------------------- +# Per-block diagnostics (rank_ratio, post_A, informative_metric) +# --------------------------------------------------------------------------- + + +def test_block_report_diagnostic_floats_are_populated(tmp_path): + """The Phase 2 enhancement populates per-block rank_ratio, post_A, + and informative_metric on BlockReport. They were None in Phase 2 v1. + Verify they're now real numbers (or None when the block has no + clusters — empty-block degenerate case).""" + n_features = 8 + sae_path = _build_synth_sae(tmp_path, n_features=n_features) + vr = build_report( + n_features=n_features, + confirmed=[(0, 1), (4, 5)], + n_fires={0: 100, 1: 10, 4: 100, 5: 10}, + ) + + heavy = BlockSpec( + block_id="heavy", encoding_class="MPSRung1", + feature_ids=(0, 1, 2, 3), + ) + tail = BlockSpec( + block_id="tail", encoding_class="MPSRung1", + feature_ids=(4, 5, 6, 7), + ) + cfg = CompressionConfig(strategy="zero", encoding_partition=(heavy, tail)) + c = Compressor(sae_checkpoint=sae_path, validation_report=vr, config=cfg) + result = c.run(output_checkpoint=tmp_path / "out.safetensors") + + blocks_by_id = {b.block_id: b for b in result.report.blocks} + for bid in ("heavy", "tail"): + b = blocks_by_id[bid] + assert b.rank_ratio is not None, f"{bid}: rank_ratio should be populated" + assert 0.0 <= b.rank_ratio <= 1.0, f"{bid}: rank_ratio in [0, 1]" + assert b.post_A is not None, f"{bid}: post_A should be populated" + assert b.informative_metric in {"post_A", "both", "forge_mse"}, ( + f"{bid}: informative_metric should be one of the expected values" + ) + + +def test_block_report_diagnostics_none_for_empty_block(tmp_path): + """When a block has no clusters (e.g. all cross-block dropped), the + per-block diagnostics SHALL be None — there's nothing to measure + against an empty plan.""" + n_features = 8 + sae_path = _build_synth_sae(tmp_path, n_features=n_features) + # Only cross-block pair → both blocks end up with 0 clusters + vr = build_report(n_features=n_features, confirmed=[(3, 4)]) + + heavy = BlockSpec( + block_id="heavy", encoding_class="MPSRung1", + feature_ids=(0, 1, 2, 3), + ) + tail = BlockSpec( + block_id="tail", encoding_class="MPSRung1", + feature_ids=(4, 5, 6, 7), + ) + cfg = CompressionConfig(strategy="zero", encoding_partition=(heavy, tail)) + c = Compressor(sae_checkpoint=sae_path, validation_report=vr, config=cfg) + result = c.run(output_checkpoint=tmp_path / "out.safetensors") + + for b in result.report.blocks: + assert b.n_clusters == 0 + assert b.rank_ratio is None + assert b.post_A is None + assert b.informative_metric is None + + def test_no_partition_uses_single_encoding_path(tmp_path): """When encoding_partition is None, Compressor.apply runs the historical single-encoding path. CompressionReport.blocks is None."""