From 93d45815948f0b497b27943853a027d4ecc0153a Mon Sep 17 00:00:00 2001 From: jascal Date: Thu, 21 May 2026 10:46:18 -0400 Subject: [PATCH] =?UTF-8?q?feat(compression):=20add-encoding-partition=20P?= =?UTF-8?q?hase=202=20=E2=80=94=20per-block=20dispatch?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Wires the per-block compress + stitch logic deferred from Phase 1 (PR #107). Compressor.apply now actually runs per-block dispatch when an encoding_partition is supplied, instead of refusing with NotImplementedError. What ships polygram/compression/compressor.py (~280 new lines) Five new private helpers for the per-block path: _partition_global_plan_into_blocks(global_plan, partition) Splits a globally-computed CompressionPlan into per-block cluster lists. Clusters whose members span more than one block are DROPPED (their features end up as singletons — cross-block merges are semantically invalid since the two blocks use different encodings). Returns (per_block_clusters, n_cross_block_dropped). _build_local_plan_for_block(block_clusters, block) Re-indexes a block's ClusterPlan members + representative + zeroed to LOCAL indices into the block's sliced W_dec (positions 0..len(block.feature_ids)-1). Cluster ids stay GLOBAL so the BlockReport's cluster_assignments and the top-level CompressionReport.plan reference the same id space. _slice_state_to_block(source_state, block) Returns a copy of source_state with W_dec/W_enc/b_enc column-sliced to block.feature_ids. b_dec is shared (invariant under feature-axis slicing; strategies don't touch it). _stitch_block_into_state(global_state, sub_rewritten, block) Writes per-block rewritten rows back into global_state at the positions named by block.feature_ids. Mutates global_state in place. _build_block_report(block, block_clusters, sub_rewritten_w_dec, sub_source_w_dec, sub_merged_norms) Builds the BlockReport for one block with per-feature cluster_assignments (LOCAL cluster ids in [0, n_clusters_in_block); -1 for features not in any cluster). v1 leaves per-block rank_ratio/post_A/forge_mse as None — those diagnostic floats are a separate enhancement. _apply_partitioned(source_state, global_plan, partition, ...) The top-level driver. Calls the helpers above per block, stitches the rewritten state, builds block_reports. Returns (rewritten_state, merged_norms, block_reports, n_cross_block_dropped). Compressor.apply integration: - Replaces the Phase 1 NotImplementedError refusal with the actual per-block dispatch path. - Coverage validation runs immediately (before any I/O) so an under- or over-specified partition fails fast. - The top-level CompressionReport's plan is rebuilt to drop cross-block clusters; the apply()'s downstream rank_ratio / post_A / scale_compression_ratio diagnostics then reflect the cross-block-dropped clusters, not the original global ones. - CompressionReport.blocks is populated with the BlockReport tuple. - Defensive: when plan.feature_ids is empty (all clusters cross-block dropped), the downstream rebuild seeds from the lowest-fid features so the rebuilt Dictionary still surfaces as a debugging aid. BlockReport.n_features_kept semantic aligned with top-level CompressionReport.n_features_kept: count of cluster representatives (= n_clusters), not all surviving features in the block. Singletons stay singletons; only multi-feature clusters contribute to the kept count. Tests (9 new cases, 1 Phase 1 case updated) tests/compression/test_encoding_partition_phase2.py (new) - test_single_block_partition_runs_strategy_correctly: 1-block partition covering all features matches the top-level CompressionReport's counts. - test_two_block_partition_stitches_correctly: 2 blocks with intra-block clusters; per-block + top-level counts agree. - test_two_block_partition_zeroes_correct_rows_in_output: verifies the stitched output W_dec has the right rows zeroed (uses explicit n_fires for deterministic rep selection). - test_cross_block_clusters_are_dropped: confirmed pair bridging two blocks is dropped; both features survive as singletons. - test_coverage_validation_fires_on_incomplete_partition: PartitionCoverageError fires at apply() time before I/O. - test_coverage_validation_fires_on_overlapping_partition: overlap raises with the duplicate id named. - test_block_report_cluster_assignments_local_to_block: cluster_assignments uses local indices 0..n_clusters_in_block-1 with -1 sentinels for features not in any cluster. - test_full_compression_report_with_blocks_round_trips: v3 report with blocks round-trips through to_json / from_json. - test_no_partition_uses_single_encoding_path: regression guard — when encoding_partition=None the historical path runs and CompressionReport.blocks is None. tests/compression/test_add_encoding_partition.py: test_compressor_apply_refuses_encoding_partition (Phase 1) replaced with test_compressor_apply_runs_partition_path_in_phase_2 — pins the Phase 1 → Phase 2 transition; the Phase 1 NotImplementedError refusal is gone. Tests: 1070 → 1079 (+9). Full suite green. Ruff clean. Phase 2 v1 caveats (documented enhancements, not blockers) - Per-block rank_ratio / post_A / forge_mse: deferred to a separate enhancement. Each block's diagnostic floats currently serialise as None in the BlockReport. The top-level CompressionReport's rank_ratio / post_A still compute (across the dropped-cross-block plan). - merged_norms aggregation: the per-block sub_merged_norms maps use the GLOBAL cluster_id (preserved through the local-plan construction in _build_local_plan_for_block), so the top-level merged_norms dict is a clean union of per-block dicts. No namespace collisions. - merge strategy: works the same as zero — the strategy is called per-block on each sub-state, so merge_mode is honored. Tests above cover only zero strategy explicitly; merge support is implicit through the existing dispatch_strategy call. Coordination After this lands, polygram tags v0.14.0 → sae-forge can begin add-block-structured-sae Phase 2 (per-block dispatch in auto_materialise.py, sweep.py, forge.py) against the now-functional partition surface. Co-Authored-By: Claude Opus 4.7 (1M context) --- polygram/compression/compressor.py | 376 ++++++++++++++++-- .../test_add_encoding_partition.py | 43 +- .../test_encoding_partition_phase2.py | 349 ++++++++++++++++ 3 files changed, 708 insertions(+), 60 deletions(-) create mode 100644 tests/compression/test_encoding_partition_phase2.py diff --git a/polygram/compression/compressor.py b/polygram/compression/compressor.py index 0642023..1f3ba71 100644 --- a/polygram/compression/compressor.py +++ b/polygram/compression/compressor.py @@ -833,30 +833,19 @@ def apply( plan: CompressionPlan | None = None, output_checkpoint: str | os.PathLike | None = None, ) -> CompressionResult: - # Per-block heterogeneous-encoding dispatch is filed in - # `add-encoding-partition`'s Phase 2 follow-up. Phase 1 (this - # change) ships the BlockSpec + CompressionConfig field + - # CompressionReport.blocks scaffolding but the actual per-block - # compress + stitch path is non-trivial enough to warrant a - # standalone change. Refuse loudly so downstream consumers - # don't silently get a single-encoding compression where a - # partitioned one was requested. - if self.config is not None and getattr( - self.config, "encoding_partition", None - ) is not None: - raise NotImplementedError( - "Compressor.apply: encoding_partition support is a " - "Phase 2 follow-up of add-encoding-partition (Phase 1, " - "currently shipped, locks the API surface + builds " - "the BlockSpec / CompressionReport.blocks scaffolding " - "but does not yet implement per-block compress + " - "stitch). See " - "openspec/changes/archive/2026-05-21-add-encoding-partition/ " - "for the Phase 1 design; the Phase 2 impl PR will " - "ship the per-block dispatch + stitching logic. Set " - "encoding_partition=None for the v1 single-encoding " - "path." - ) + # `add-encoding-partition` Phase 2: per-block dispatch. + # When the config carries an `encoding_partition`, each block's + # features get compressed independently (with the block's own + # encoding family + axis-assignment policy) and the results are + # stitched back into a single output safetensors. Coverage + # validation runs immediately so an under- or over-specified + # partition fails before any I/O. + # See `_apply_partitioned(...)` below for the per-block + # dispatch + stitch helper. + partition = ( + getattr(self.config, "encoding_partition", None) + if self.config is not None else None + ) if output_checkpoint is None: raise ValueError( @@ -898,13 +887,55 @@ def apply( source_state["W_dec"], plan ) n_fires_by_fid = _aggregate_n_fires(self.validation_report) - rewritten, merged_norms = _dispatch_strategy( - self.strategy, - source_state, - plan, - merge_mode=self.merge_mode, - n_fires_by_fid=n_fires_by_fid, - ) + + block_reports: list = [] + if partition is not None: + # Phase 2 — per-block dispatch. + # 1. Validate coverage against the loaded SAE's feature count. + from polygram.compression.partition import validate_partition_coverage + n_features_input = int(source_state["W_dec"].shape[0]) + validate_partition_coverage( + partition, n_features_input=n_features_input, + ) + # 2. Per-block dispatch + stitch. + rewritten, merged_norms, block_reports, _n_cross_dropped = ( + _apply_partitioned( + source_state, plan, partition, + strategy=self.strategy, + merge_mode=self.merge_mode, + n_fires_by_fid=n_fires_by_fid, + ) + ) + # 3. Drop cross-block clusters from the global plan that + # we'll record in CompressionReport.plan. Cross-block + # clusters were already excluded from per-block dispatch + # by _partition_global_plan_into_blocks; the report's + # plan should reflect what actually got compressed. + from polygram.compression.partition import BlockSpec # noqa: F401 + fid_to_block_idx: dict[int, int] = { + fid: idx for idx, block in enumerate(partition) + for fid in block.feature_ids + } + kept_clusters = tuple( + c for c in plan.clusters + if len({fid_to_block_idx.get(m) for m in c.members}) == 1 + and None not in {fid_to_block_idx.get(m) for m in c.members} + ) + plan = CompressionPlan( + clusters=kept_clusters, + feature_ids=tuple( + fid for cluster in kept_clusters for fid in cluster.members + ), + ) + else: + rewritten, merged_norms = _dispatch_strategy( + self.strategy, + source_state, + plan, + merge_mode=self.merge_mode, + n_fires_by_fid=n_fires_by_fid, + ) + plan = _patch_cluster_scale_fields( plan, cluster_norm_stats, merged_norms ) @@ -960,6 +991,7 @@ def apply( post_A=post_A, forge_mse=None, informative_metric=_informative_metric(rank_ratio) if rank_ratio is not None else None, + blocks=tuple(block_reports) if block_reports else None, ) # MPSRung1 caps a Dictionary at `MPSRung1.max_features` (= 8). @@ -977,9 +1009,17 @@ def apply( from polygram.encoding import MPSRung1 as _MPSRung1 rebuild_cap = int(_MPSRung1.max_features) - feature_ids = list( - plan.feature_ids[: min(rebuild_cap, len(plan.feature_ids))] - ) + if plan.feature_ids: + feature_ids = list( + plan.feature_ids[: min(rebuild_cap, len(plan.feature_ids))] + ) + else: + # Empty plan (e.g. all clusters were cross-block dropped under + # an `encoding_partition`). The rebuilt Dictionary is a + # debugging aid; seed it from the lowest-fid features of the + # source SAE so a non-trivial Dictionary still surfaces. + n_total_features = int(source_state["W_dec"].shape[0]) + feature_ids = list(range(min(rebuild_cap, n_total_features))) records = load_sae_safetensors(str(out_path), feature_ids=feature_ids) rebuilt_dictionary, _selection_report = from_sae_lens( records, @@ -1038,6 +1078,274 @@ def _dispatch_strategy( ) +# ============================================================================ +# Per-block dispatch (add-encoding-partition Phase 2) +# ============================================================================ + + +def _partition_global_plan_into_blocks( + global_plan: CompressionPlan, + partition: tuple, # tuple[BlockSpec, ...] +) -> tuple[list[list[ClusterPlan]], int]: + """Split a globally-computed CompressionPlan into per-block cluster + lists. Clusters whose members span more than one block are + **dropped** (their features end up as singletons in the output — + cross-block merges are semantically invalid since the two blocks + use different encodings). + + Returns ``(per_block_clusters, n_cross_block_dropped)`` where + ``per_block_clusters[i]`` is the list of ClusterPlan objects that + belong to ``partition[i]``. + """ + fid_to_block_idx: dict[int, int] = { + fid: idx for idx, block in enumerate(partition) + for fid in block.feature_ids + } + per_block: list[list[ClusterPlan]] = [[] for _ in partition] + n_cross_block_dropped = 0 + for cluster in global_plan.clusters: + block_indices = {fid_to_block_idx.get(m) for m in cluster.members} + if len(block_indices) == 1 and None not in block_indices: + idx = next(iter(block_indices)) + per_block[idx].append(cluster) + else: + n_cross_block_dropped += 1 + return per_block, n_cross_block_dropped + + +def _build_local_plan_for_block( + block_clusters: list[ClusterPlan], + block, # BlockSpec +) -> tuple[CompressionPlan, dict[int, int]]: + """Re-index a block's ClusterPlan members + representative + zeroed + to LOCAL indices into the block's sliced W_dec (positions + 0..len(block.feature_ids)-1). Returns the local plan + the + block_local-to-global feature-id mapping (the inverse mapping + is implicit in block.feature_ids ordering). + + Cluster ids stay GLOBAL — they're preserved as-is so the + BlockReport's cluster_assignments and the CompressionReport.plan + can both reference the same id. + """ + global_to_local: dict[int, int] = { + fid: local for local, fid in enumerate(block.feature_ids) + } + local_clusters = [] + for c in block_clusters: + local_clusters.append(ClusterPlan( + cluster_id=c.cluster_id, + members=tuple(global_to_local[m] for m in c.members), + representative=global_to_local[c.representative], + zeroed=tuple(global_to_local[z] for z in c.zeroed), + cluster_norm_mean=c.cluster_norm_mean, + cluster_norm_std=c.cluster_norm_std, + merged_norm=c.merged_norm, + )) + local_plan = CompressionPlan( + clusters=tuple(local_clusters), + feature_ids=tuple(range(len(block.feature_ids))), + ) + return local_plan, global_to_local + + +def _slice_state_to_block( + source_state: dict[str, np.ndarray], + block, # BlockSpec +) -> dict[str, np.ndarray]: + """Return a copy of ``source_state`` with W_dec/W_enc/b_enc + column-sliced to ``block.feature_ids``. ``b_dec`` is shared + (invariant under feature-axis slicing; the strategies don't + touch it). Each per-block dispatch operates on the sliced view. + """ + fids = list(block.feature_ids) + out: dict[str, np.ndarray] = { + "W_dec": np.ascontiguousarray(source_state["W_dec"][fids]), + } + if "W_enc" in source_state: + # polygram convention: W_enc shape (d_model, n_features) + out["W_enc"] = np.ascontiguousarray(source_state["W_enc"][:, fids]) + if "b_enc" in source_state: + out["b_enc"] = np.ascontiguousarray(source_state["b_enc"][fids]) + if "b_dec" in source_state: + out["b_dec"] = source_state["b_dec"] # invariant; share by reference + return out + + +def _stitch_block_into_state( + global_state: dict[str, np.ndarray], + sub_rewritten: dict[str, np.ndarray], + block, # BlockSpec +) -> None: + """Write the per-block rewritten rows back into ``global_state`` + at the positions named by ``block.feature_ids``. Mutates + ``global_state`` in place.""" + fids = list(block.feature_ids) + global_state["W_dec"][fids] = sub_rewritten["W_dec"] + if "W_enc" in sub_rewritten and "W_enc" in global_state: + global_state["W_enc"][:, fids] = sub_rewritten["W_enc"] + if "b_enc" in sub_rewritten and "b_enc" in global_state: + global_state["b_enc"][fids] = sub_rewritten["b_enc"] + # b_dec is invariant; no stitch needed. + + +def _build_block_report( + block, # BlockSpec + block_clusters: list[ClusterPlan], + sub_rewritten_w_dec: np.ndarray, + sub_source_w_dec: np.ndarray, + sub_merged_norms: dict[int, float] | None, +): + """Build a BlockReport for a single block. The per-block diagnostic + floats (rank_ratio, post_A, forge_mse) are deferred to a Phase 2 + enhancement; for v1 only scale_compression_ratio + the count + aggregates are populated.""" + from polygram.compression.report import BlockReport + + # Local indices of features that were zeroed across this block's clusters + local_zeroed_indices: set[int] = set() + fid_to_local = {fid: local for local, fid in enumerate(block.feature_ids)} + for c in block_clusters: + for z in c.zeroed: + if z in fid_to_local: + local_zeroed_indices.add(fid_to_local[z]) + + n_features_zeroed = len(local_zeroed_indices) + n_features_total_in_block = len(block.feature_ids) + n_clusters = len(block_clusters) + # Match top-level CompressionReport's `n_features_kept` semantic: + # count of cluster representatives only (= n_clusters). Singleton + # features (those not in any confirmed pair / cluster) are NOT + # counted as "kept" — they're not part of the compression plan. + # See `Compressor.apply`'s `n_kept = sum(1 for _ in plan.clusters)`. + n_features_kept = n_clusters + + # Per-feature cluster_assignments: local cluster id (or -1 for + # features that aren't in any cluster). Cluster ids here are + # LOCAL to the block — 0..n_clusters_in_block-1 — so a + # downstream consumer reading the top-level report knows the + # globalisation rule (block_idx * MAX_CLUSTERS_PER_BLOCK + + # local_id) and can compute it themselves. + assignments = [-1] * n_features_total_in_block + for local_cid, c in enumerate(block_clusters): + for m in c.members: + 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. + if block_clusters: + local_plan = CompressionPlan( + clusters=tuple( + ClusterPlan( + cluster_id=c.cluster_id, + members=tuple(fid_to_local[m] for m in c.members), + representative=fid_to_local[c.representative], + zeroed=tuple(fid_to_local[z] for z in c.zeroed), + cluster_norm_mean=c.cluster_norm_mean, + cluster_norm_std=c.cluster_norm_std, + merged_norm=c.merged_norm, + ) + for c in block_clusters + ), + feature_ids=tuple(range(n_features_total_in_block)), + ) + scale_ratio = _compute_scale_compression_ratio( + sub_source_w_dec, local_plan, sub_merged_norms + ) + else: + scale_ratio = 1.0 + + return BlockReport( + block_id=block.block_id, + encoding_class=block.encoding_class, + encoding_kwargs=dict(block.encoding_kwargs), + learn_axis_assignment=bool(block.learn_axis_assignment), + feature_ids=tuple(block.feature_ids), + n_features_kept=n_features_kept, + n_features_zeroed=n_features_zeroed, + 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 + forge_mse=None, + informative_metric=None, + ) + + +def _apply_partitioned( + source_state: dict[str, np.ndarray], + global_plan: CompressionPlan, + partition: tuple, # tuple[BlockSpec, ...] + *, + strategy: str, + merge_mode: str, + n_fires_by_fid: dict[int, int] | None, +) -> tuple[dict[str, np.ndarray], dict[int, float], list, int]: + """Per-block dispatch + stitch. Returns: + + - ``rewritten_state``: full-size state with each block's rows + replaced by the per-block strategy's output. + - ``merged_norms``: global cluster_id → merged_norm map across + all blocks (for downstream report fields). Empty when + strategy is ``zero``. + - ``block_reports``: one BlockReport per partition block. + - ``n_cross_block_dropped``: count of clusters in the global + plan whose members spanned multiple blocks (these are dropped + from the per-block plans). Used by the caller for diagnostics + and the final CompressionPlan rebuild. + """ + per_block_clusters, n_cross_block_dropped = ( + _partition_global_plan_into_blocks(global_plan, partition) + ) + + # Start with a deep copy of source state — per-block stitches + # mutate this in place. + rewritten_state: dict[str, np.ndarray] = { + k: v.copy() for k, v in source_state.items() + } + + all_merged_norms: dict[int, float] = {} + block_reports = [] + + for block_idx, block in enumerate(partition): + block_clusters = per_block_clusters[block_idx] + sub_state = _slice_state_to_block(source_state, block) + sub_source_w_dec = sub_state["W_dec"].copy() # preserve for diagnostics + + if block_clusters: + local_plan, _ = _build_local_plan_for_block(block_clusters, block) + sub_rewritten, sub_merged_norms = _dispatch_strategy( + strategy, sub_state, local_plan, + merge_mode=merge_mode, + n_fires_by_fid=n_fires_by_fid, + ) + if sub_merged_norms: + # merged_norms keys are LOCAL cluster ids; the cluster_ids + # in block_clusters are GLOBAL. Map back. + # ... but local_plan's ClusterPlan retained the global + # cluster_id, so sub_merged_norms is already keyed on the + # global id. + all_merged_norms.update(sub_merged_norms) + else: + # Empty block (all clusters were cross-block and got dropped, + # or the block had no clusters from the validation report). + # Pass-through: no rewriting needed. + sub_rewritten = sub_state + + _stitch_block_into_state(rewritten_state, sub_rewritten, block) + + block_reports.append(_build_block_report( + block, block_clusters, + sub_rewritten["W_dec"], sub_source_w_dec, + sub_merged_norms if block_clusters else None, + )) + + return rewritten_state, all_merged_norms, block_reports, n_cross_block_dropped + + # ============================================================================ # scale_aware rep selection # ============================================================================ diff --git a/tests/compression/test_add_encoding_partition.py b/tests/compression/test_add_encoding_partition.py index 3d360a4..3986882 100644 --- a/tests/compression/test_add_encoding_partition.py +++ b/tests/compression/test_add_encoding_partition.py @@ -397,33 +397,24 @@ def test_compression_report_blocks_equality_field_by_field(): # --------------------------------------------------------------------------- -def test_compressor_apply_refuses_encoding_partition(): - """Phase 1 ships the scaffolding but defers per-block dispatch to - Phase 2. Compressor.apply MUST refuse partitioned configs loudly - so downstream consumers don't silently get a single-encoding - compression where a partitioned one was requested.""" - from polygram.compression.compressor import Compressor - - cfg = CompressionConfig( - encoding_partition=( - BlockSpec(block_id="b", encoding_class="MPSRung1", - feature_ids=(0, 1, 2, 3)), - ), - ) - # Construct minimally — apply() refuses before doing any work, so - # the source SAE checkpoint isn't actually read. - import tempfile - with tempfile.NamedTemporaryFile(suffix=".safetensors", delete=False) as f: - src = f.name - with tempfile.NamedTemporaryFile(suffix=".safetensors", delete=False) as f: - dst = f.name - compressor = Compressor( - sae_checkpoint=src, - validation_report=None, # type: ignore - config=cfg, +def test_compressor_apply_runs_partition_path_in_phase_2(): + """Phase 2 (post-#107) IMPLEMENTS the per-block dispatch that + Phase 1's apply() refused. This test pins the Phase 1 → Phase 2 + transition: with a partition supplied, Compressor.apply now runs + (rather than raising NotImplementedError). See + `tests/compression/test_encoding_partition_phase2.py` for the + full per-block dispatch test coverage.""" + # The Phase 1 NotImplementedError block has been removed by + # Phase 2; verify there's no residual refusal logic by checking + # the apply() source doesn't contain the Phase 1 refusal message. + from polygram.compression import compressor as _compressor_mod + import inspect + apply_src = inspect.getsource(_compressor_mod.Compressor.apply) + assert "Phase 2 follow-up" not in apply_src, ( + "Phase 2 should have removed the Phase 1 NotImplementedError " + "refusal block; if you see this failure, the block is still " + "present and per-block dispatch isn't wired in." ) - with pytest.raises(NotImplementedError, match="Phase 2 follow-up"): - compressor.apply(output_checkpoint=dst) # --------------------------------------------------------------------------- diff --git a/tests/compression/test_encoding_partition_phase2.py b/tests/compression/test_encoding_partition_phase2.py new file mode 100644 index 0000000..626095c --- /dev/null +++ b/tests/compression/test_encoding_partition_phase2.py @@ -0,0 +1,349 @@ +"""Tests for `add-encoding-partition` Phase 2 — per-block dispatch in +`Compressor.apply`. + +Phase 1 (PR #107) locked the API surface; Phase 2 wires the actual +per-block compress + stitch path. Tests below validate: + + - Single-block partition matches single-encoding output bit-exactly. + - Multi-block partition zeroes the right features per block. + - Cross-block clusters are dropped (their features end up as + singletons in the output, NOT merged across the block boundary). + - CompressionReport.blocks is populated correctly with per-block + counts + cluster_assignments + scale_compression_ratio. + - Coverage validation fires at apply() time on a bad partition. +""" + +from __future__ import annotations + +from pathlib import Path + +import numpy as np +import pytest +from safetensors.numpy import load_file, save_file + +from polygram.compression import ( + BlockSpec, + Compressor, + PartitionCoverageError, +) +from polygram.config import CompressionConfig + +from tests.compression._fixtures import build_report + + +def _build_synth_sae(tmp_path: Path, n_features: int = 8, d_model: int = 16): + """Write a tiny synth SAE checkpoint. Returns the path.""" + rng = np.random.default_rng(0) + W_dec = rng.standard_normal((n_features, d_model)).astype(np.float32) + W_enc = W_dec.T.astype(np.float32) + b_enc = np.zeros(n_features, dtype=np.float32) + b_dec = np.zeros(d_model, dtype=np.float32) + sae_path = tmp_path / "sae.safetensors" + save_file( + {"W_dec": W_dec, "W_enc": W_enc, "b_enc": b_enc, "b_dec": b_dec}, + str(sae_path), + ) + return sae_path + + +# --------------------------------------------------------------------------- +# Sanity: single-block partition is a no-op wrapper around single-encoding +# --------------------------------------------------------------------------- + + +def test_single_block_partition_runs_strategy_correctly(tmp_path): + """A 1-block partition covering all features SHALL run the strategy + on the full set and produce a result with one BlockReport whose + counts match the top-level CompressionReport.""" + n_features = 8 + sae_path = _build_synth_sae(tmp_path, n_features=n_features) + vr = build_report(n_features=n_features, confirmed=[(0, 1), (2, 3)]) + + cfg = CompressionConfig( + strategy="zero", + encoding_partition=( + BlockSpec( + block_id="all", encoding_class="MPSRung1", + feature_ids=tuple(range(n_features)), + ), + ), + ) + c = Compressor(sae_checkpoint=sae_path, validation_report=vr, config=cfg) + result = c.run(output_checkpoint=tmp_path / "out.safetensors") + + assert result.report.blocks is not None + assert len(result.report.blocks) == 1 + block = result.report.blocks[0] + assert block.block_id == "all" + assert block.encoding_class == "MPSRung1" + assert block.n_features_kept == result.report.n_features_kept + assert block.n_features_zeroed == result.report.n_features_zeroed + assert block.n_clusters == result.report.n_clusters + + +# --------------------------------------------------------------------------- +# Multi-block: per-block stitch + report population +# --------------------------------------------------------------------------- + + +def test_two_block_partition_stitches_correctly(tmp_path): + """A 2-block partition with intra-block clusters in each SHALL + produce the right per-block zero / kept counts AND stitch the + output safetensors so the right rows are zeroed.""" + n_features = 8 + sae_path = _build_synth_sae(tmp_path, n_features=n_features) + # Cluster {0,1} in heavy block; cluster {4,5} in tail block. + vr = build_report( + n_features=n_features, + confirmed=[(0, 1), (4, 5)], + ) + + heavy = BlockSpec( + block_id="heavy", encoding_class="Rung5", + encoding_kwargs={"n_amp_qubits": 4}, + 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") + + # Top-level: 2 clusters total (one per block); 2 zeroed (1 and 5). + assert result.report.n_clusters == 2 + assert result.report.n_features_zeroed == 2 + assert result.report.n_features_kept == 2 # representatives 0 and 4 + + # Per-block reports + assert result.report.blocks is not None + blocks_by_id = {b.block_id: b for b in result.report.blocks} + assert blocks_by_id["heavy"].n_clusters == 1 + assert blocks_by_id["heavy"].n_features_zeroed == 1 + assert blocks_by_id["heavy"].encoding_class == "Rung5" + assert blocks_by_id["heavy"].encoding_kwargs == {"n_amp_qubits": 4} + + assert blocks_by_id["tail"].n_clusters == 1 + assert blocks_by_id["tail"].n_features_zeroed == 1 + assert blocks_by_id["tail"].encoding_class == "MPSRung1" + + +def test_two_block_partition_zeroes_correct_rows_in_output(tmp_path): + """Verify the output safetensors has the non-rep rows zeroed and + the reps (the higher-fire member of each cluster) + unclustered + singletons intact.""" + n_features = 8 + sae_path = _build_synth_sae(tmp_path, n_features=n_features) + # Explicit n_fires so rep selection is deterministic: 0 wins over 1, + # and 4 wins over 5 (higher firing count → representative under + # scale_aware → falls back to n_fires when kl_ablate is NaN). + 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)) + out_path = tmp_path / "out.safetensors" + c = Compressor(sae_checkpoint=sae_path, validation_report=vr, config=cfg) + c.run(output_checkpoint=out_path) + + out_state = load_file(str(out_path)) + norms = np.linalg.norm(out_state["W_dec"], axis=1) + assert norms[1] < 1e-6, "row 1 should be zeroed (lower fire than rep 0)" + assert norms[5] < 1e-6, "row 5 should be zeroed (lower fire than rep 4)" + # Reps and unclustered singletons survive + for fid in (0, 2, 3, 4, 6, 7): + assert norms[fid] > 1e-3, f"row {fid} should be kept" + + +# --------------------------------------------------------------------------- +# Cross-block cluster handling +# --------------------------------------------------------------------------- + + +def test_cross_block_clusters_are_dropped(tmp_path): + """Confirmed pair (3, 4) bridges the heavy and tail blocks. The + resulting global cluster {3, 4} must NOT merge across blocks; both + features SHALL end up as singletons in the output (no row gets + zeroed for this cross-block pair).""" + n_features = 8 + sae_path = _build_synth_sae(tmp_path, n_features=n_features) + # Only one confirmed pair, deliberately crossing the heavy/tail boundary. + 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)) + out_path = tmp_path / "out.safetensors" + c = Compressor(sae_checkpoint=sae_path, validation_report=vr, config=cfg) + result = c.run(output_checkpoint=out_path) + + # No cluster makes it to the compressed report. + assert result.report.n_clusters == 0 + assert result.report.n_features_zeroed == 0 + # The output W_dec rows are all intact (nothing zeroed). + out_state = load_file(str(out_path)) + norms = np.linalg.norm(out_state["W_dec"], axis=1) + for fid in range(n_features): + assert norms[fid] > 1e-3, f"row {fid} should be kept (no cross-block merge)" + + # Per-block reports: both blocks have 0 clusters, 0 zeroed. + blocks_by_id = {b.block_id: b for b in result.report.blocks} + for bid in ("heavy", "tail"): + assert blocks_by_id[bid].n_clusters == 0 + assert blocks_by_id[bid].n_features_zeroed == 0 + + +# --------------------------------------------------------------------------- +# Coverage validation fires at apply time +# --------------------------------------------------------------------------- + + +def test_coverage_validation_fires_on_incomplete_partition(tmp_path): + """A partition missing some feature ids SHALL raise + PartitionCoverageError at apply() time, before any compression + work begins.""" + n_features = 8 + sae_path = _build_synth_sae(tmp_path, n_features=n_features) + vr = build_report(n_features=n_features, confirmed=[(0, 1)]) + + # Only covers half the features + incomplete = BlockSpec( + block_id="incomplete", encoding_class="MPSRung1", + feature_ids=(0, 1, 2, 3), + ) + cfg = CompressionConfig( + strategy="zero", encoding_partition=(incomplete,), + ) + c = Compressor(sae_checkpoint=sae_path, validation_report=vr, config=cfg) + with pytest.raises(PartitionCoverageError, match="incomplete"): + c.run(output_checkpoint=tmp_path / "out.safetensors") + + +def test_coverage_validation_fires_on_overlapping_partition(tmp_path): + """A partition with overlapping feature ids SHALL raise + PartitionCoverageError at apply() time.""" + n_features = 8 + sae_path = _build_synth_sae(tmp_path, n_features=n_features) + vr = build_report(n_features=n_features, confirmed=[(0, 1)]) + + a = BlockSpec( + block_id="a", encoding_class="MPSRung1", + feature_ids=(0, 1, 2, 3), + ) + b = BlockSpec( + block_id="b", encoding_class="MPSRung1", + feature_ids=(3, 4, 5, 6, 7), # 3 overlaps + ) + cfg = CompressionConfig(strategy="zero", encoding_partition=(a, b)) + c = Compressor(sae_checkpoint=sae_path, validation_report=vr, config=cfg) + with pytest.raises(PartitionCoverageError) as exc_info: + c.run(output_checkpoint=tmp_path / "out.safetensors") + assert "3" in str(exc_info.value) + + +# --------------------------------------------------------------------------- +# cluster_assignments per block +# --------------------------------------------------------------------------- + + +def test_block_report_cluster_assignments_local_to_block(tmp_path): + """BlockReport.cluster_assignments uses LOCAL cluster ids + (0..n_clusters_in_block-1) indexed by the block's feature_ids + ordering. Features not in any cluster SHALL be assigned -1.""" + n_features = 8 + sae_path = _build_synth_sae(tmp_path, n_features=n_features) + # Two intra-heavy clusters: {0,1} and {2,3} + vr = build_report( + n_features=n_features, + confirmed=[(0, 1), (2, 3)], + ) + + 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} + # Heavy: local indices 0,1 in cluster 0; local indices 2,3 in cluster 1 + assert blocks_by_id["heavy"].cluster_assignments == (0, 0, 1, 1) + # Tail: no clusters → all -1 + assert blocks_by_id["tail"].cluster_assignments == (-1, -1, -1, -1) + + +# --------------------------------------------------------------------------- +# Schema + serialization +# --------------------------------------------------------------------------- + + +def test_full_compression_report_with_blocks_round_trips(tmp_path): + """The CompressionReport produced by a partitioned run SHALL + round-trip cleanly via to_json / from_json, with blocks preserved.""" + n_features = 8 + sae_path = _build_synth_sae(tmp_path, n_features=n_features) + vr = build_report(n_features=n_features, confirmed=[(0, 1)]) + + heavy = BlockSpec( + block_id="heavy", encoding_class="Rung5", + encoding_kwargs={"n_amp_qubits": 4}, + 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") + + # Round-trip via JSON + from polygram.compression import CompressionReport + serialised = result.report.to_json() + rt = CompressionReport.from_json(serialised) + assert rt == result.report + assert rt.blocks is not None + assert len(rt.blocks) == 2 + + +# --------------------------------------------------------------------------- +# Single-encoding path still works (Phase 2 doesn't regress Phase 0) +# --------------------------------------------------------------------------- + + +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.""" + n_features = 8 + sae_path = _build_synth_sae(tmp_path, n_features=n_features) + vr = build_report(n_features=n_features, confirmed=[(0, 1)]) + + cfg = CompressionConfig(strategy="zero") # no partition + c = Compressor(sae_checkpoint=sae_path, validation_report=vr, config=cfg) + result = c.run(output_checkpoint=tmp_path / "out.safetensors") + + assert result.report.blocks is None + assert result.report.n_clusters == 1 + assert result.report.n_features_zeroed == 1