feat(compression): add-encoding-partition Phase 1 — API surface + scaffolding - #107
Merged
Conversation
…ffolding Implements the proposal merged today as PR #106. Phase 1 locks the API surface and ships the scaffolding (BlockSpec, validation, CompressionConfig.encoding_partition field, CompressionReport.blocks schema v2->v3, public exports). Phase 2 (the actual per-block dispatch in Compressor.apply) is filed as a follow-up; the current apply() raises NotImplementedError when partitioned configs arrive so downstream consumers don't silently get a single-encoding compression where a partitioned one was requested. What ships polygram/compression/partition.py (new, ~325 lines) BlockSpec dataclass — frozen, custom __hash__ (since encoding_kwargs is a dict). Per-family validator registry _BLOCK_SPEC_KWARG_VALIDATORS keeps the per-family logic out of __post_init__'s branch chain — adding Rung6+ is one new validator + one registry entry. PartitionCoverageError (subclasses ValueError). validate_partition_coverage(partition, n_features_input=...) — disjointness + completeness checks; error message names up to 10 offending feature ids. make_default_block(...) — 'default + heavy override' helper. polygram/compression/report.py SCHEMA_VERSION 2 -> 3. MAX_CLUSTERS_PER_BLOCK = 10_000 (global cluster-id namespace cap; Decision 2 in the design.md). BlockReport dataclass with per-block diagnostics. CompressionReport.blocks: tuple[BlockReport, ...] | None field. from_json defaults blocks=None for v2 payloads (back-compat). Equality + hash include blocks. _block_to_dict, _block_from_dict, _blocks_eq helpers added. polygram/compression/compressor.py Compressor.apply now refuses encoding_partition with a clear NotImplementedError naming the Phase 2 follow-up. Refusal happens before any I/O so the input SAE checkpoint isn't even read. polygram/config.py CompressionConfig.encoding_partition: tuple | None = None. __post_init__ type/membership/non-empty validation. The locked surface is tuple (NOT list — required for frozen-hashable contract). polygram/sae_import.py from_sae_lens(..., encoding_partition=...) kwarg accepted + type-validated. No-op routing (from_sae_lens builds Dictionaries, not compressed checkpoints; the partition is consumed at Compressor.apply time). polygram/compression/__init__.py New public exports: BlockSpec, BlockReport, PartitionCoverageError, make_default_block, validate_partition_coverage. Tests (33 new cases) tests/compression/test_add_encoding_partition.py §9.1 BlockSpec validation (11 cases) — encoding_class membership, per-family kwargs (Rung5, HEA_Rung2, MPSRung1 extras), empty feature_ids, duplicates, negative ids, empty block_id, hashability. §9.2 Partition coverage (5 cases) — disjoint+complete passes, overlap names duplicates, missing names holes, pure-extras names extras, negative n_features_input rejected. §9.5 make_default_block (4 cases) — covers-all-except-excluded, no-exclusions, refuses-empty-result, heavy-override pattern validates. §9.4 CompressionConfig.encoding_partition (5 cases) — default None, accepts tuple, rejects list, rejects non-BlockSpec, rejects empty tuple. §9.4 CompressionReport schema (8 cases) — SCHEMA_VERSION == 3, blocks default None, blocks round-trip, v2-payload back-compat load, v3 + blocks=None round-trip, blocks field-by-field equality, MAX_CLUSTERS_PER_BLOCK constant. Suite: 1037 -> 1070 (+33). Ruff clean. Coordination This change locks the API surface for sae-forge's add-block- structured-sae (Phase 0 of which explicitly cites these names). After this lands + polygram tags v0.13.0, sae-forge can bump its polygram>=0.13.0 pin and begin Phase 2 of its own change. The Phase 2 follow-up for THIS change (per-block dispatch in Compressor.apply + stitching, plus the BlockReport population from per-block compress sub-runs) is a separate impl PR. Filing it follows once Phase 1 is merged + tagged. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
…mple Addresses PR #107 review: - Compressor.apply's NotImplementedError now points at the openspec change directory (post-archive path) so downstream devs hitting the refusal can navigate to the Phase 1 design + the eventual Phase 2 impl PR directly. - partition.py module docstring gains a 'Usage example' block showing the canonical 'default + heavy override' pattern end to end (BlockSpec + make_default_block + validate_partition_coverage + CompressionConfig). Demonstrates the Phase 1 surface in runnable form. Non-functional polish; no behavioural change. 33/33 tests still pass. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
4 tasks
jascal
added a commit
that referenced
this pull request
May 21, 2026
…#108) Tags the polygram release that ships add-encoding-partition Phase 1 (#107) and archives the now-completed openspec change directory. Wave C unblocker: this release locks the API surface (BlockSpec, CompressionConfig.encoding_partition, CompressionReport.blocks, from_sae_lens kwarg) that sae-forge's add-block-structured-sae proposal explicitly cites as its Phase 0 prerequisite. Once this tag publishes, sae-forge can bump its polygram>=0.13.0 pin and begin Phase 2 of its own change. Phase 2 (the actual per-block dispatch in Compressor.apply) is a separate impl follow-up filed against this proposal — not in 0.13.0. Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
5 tasks
jascal
added a commit
that referenced
this pull request
May 21, 2026
#109) 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) <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Implements the proposal merged today as PR #106. Phase 1 locks the API surface and ships the scaffolding; Phase 2 (per-block dispatch in
Compressor.apply) is a separate follow-up impl PR.After this lands + polygram tags v0.13.0, sae-forge's
add-block-structured-saecan bump its pin and begin Phase 2 of its own change.What ships
polygram/compression/partition.py(new, ~325 lines)BlockSpecdataclass — frozen, custom__hash__(sinceencoding_kwargsis a dict). Per-family validator registry_BLOCK_SPEC_KWARG_VALIDATORSkeeps the per-family logic out of__post_init__'s branch chain — adding Rung6+ is one new validator + one registry entry.PartitionCoverageError(subclassesValueError).validate_partition_coverage(partition, n_features_input=...)— disjointness + completeness checks; error message names up to 10 offending feature ids.make_default_block(...)— 'default + heavy override' helper.polygram/compression/report.pySCHEMA_VERSION 2 → 3.MAX_CLUSTERS_PER_BLOCK = 10_000(global cluster-id namespace cap; Decision 2 in design.md).BlockReportdataclass with per-block diagnostics.CompressionReport.blocks: tuple[BlockReport, ...] | Nonefield.from_jsondefaultsblocks=Nonefor v2 payloads (back-compat).blocks._block_to_dict,_block_from_dict,_blocks_eqhelpers added.polygram/compression/compressor.pyCompressor.applynow refusesencoding_partitionwith a clearNotImplementedErrornaming the Phase 2 follow-up. Refusal happens before any I/O so the input SAE checkpoint isn't even read. This is load-bearing: without it, downstream consumers would silently get a single-encoding compression where a partitioned one was requested.polygram/config.pyCompressionConfig.encoding_partition: tuple | None = None.__post_init__type/membership/non-empty validation. The locked surface is tuple (NOT list — required for thefrozen=Truehashable contract).polygram/sae_import.pyfrom_sae_lens(..., encoding_partition=...)kwarg accepted + type-validated.Public exports
From
polygram.compression:BlockSpec,BlockReport,PartitionCoverageError,make_default_block,validate_partition_coverage.Tests (33 new cases)
SCHEMA_VERSION == 3, blocks default None, blocks round-trip, v2-payload back-compat load, v3 + blocks=None round-trip, blocks field-by-field equality,MAX_CLUSTERS_PER_BLOCKconstantSuite: 1037 → 1070 (+33). Ruff clean.
Phase 2 follow-up (separate change)
Per-block dispatch in
Compressor.apply:W_dec/W_enc/b_encto the block'sfeature_ids.CompressionConfigfrom the BlockSpec.CompressionReportwithblockspopulated.That work is significant (~150 LOC + the per-block test smoke matrix) and warrants its own PR.
Coordination with sae-forge
After this lands and polygram tags
v0.13.0:polygram>=0.13.0pin.add-block-structured-sae's Phase 0.1 task clears.auto_materialise.py,sweep.py,forge.py) can begin against the locked API surface.Test plan
🤖 Generated with Claude Code