feat(compression): add-encoding-partition Phase 2 — per-block dispatch

[jascal]

Summary

Wires the per-block compress + stitch logic that Phase 1 (#107) deferred. Compressor.apply now actually runs per-block dispatch when an encoding_partition is supplied, instead of refusing with NotImplementedError.

After this lands + polygram tags v0.14.0, sae-forge's add-block-structured-sae Phase 2 (per-block dispatch in auto_materialise.py / sweep.py / forge.py) can begin against the now-functional partition surface.

What ships (~280 new lines)

Five new private helpers in polygram/compression/compressor.py:

helper	role
_partition_global_plan_into_blocks	Split global plan's clusters per block. Cross-block clusters are dropped (semantically invalid — two encodings can't merge a cluster spanning blocks).
_build_local_plan_for_block	Re-index cluster members/rep/zeroed to block-local indices for the strategy dispatch. Cluster ids stay GLOBAL.
_slice_state_to_block	Column-slice W_dec/W_enc/b_enc to block.feature_ids. b_dec invariant — shared by reference.
_stitch_block_into_state	Write per-block rewritten rows back into global state at block.feature_ids positions.
_build_block_report	Build the BlockReport with per-feature cluster_assignments (local cluster ids, -1 for unclustered).
_apply_partitioned	Top-level driver. Calls the helpers per block, returns (rewritten_state, merged_norms, block_reports).

Compressor.apply integration:

• Coverage validation fires before any I/O so a bad partition fails fast.
• Top-level plan rebuilt to drop cross-block clusters; downstream rank_ratio/post_A reflect the post-drop plan.
• CompressionReport.blocks populated.
• Defensive: when the dropped-cross-block plan is empty, the downstream rebuilt-Dictionary debugging artifact seeds from the lowest-fid features rather than crashing.

BlockReport.n_features_kept semantic aligned with top-level CompressionReport.n_features_kept: count of cluster representatives (= n_clusters), NOT total survivors.

Tests (9 new + 1 Phase 1 update)

test	covers
single-block partition matches top-level counts	sanity
2-block partition stitches correctly	per-block + top-level counts agree
2-block partition zeroes correct rows in output	end-to-end safetensors verification (with explicit n_fires for deterministic rep selection)
cross-block clusters are dropped	pair (3,4) bridging blocks → both features survive as singletons
coverage validation fires on incomplete partition	early failure at apply() time
coverage validation fires on overlapping partition	duplicate id named
block_report cluster_assignments are local to block	-1 sentinels for unclustered features
full CompressionReport with blocks round-trips	v3 schema
no partition uses single-encoding path	Phase 1 regression guard
test_compressor_apply_runs_partition_path_in_phase_2	Phase 1 → Phase 2 transition — the NotImplementedError block is gone

1070 → 1079 tests passing on full suite. 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 serialise as None. Top-level still computes (across the post-drop plan).
• merged_norms namespace — global cluster_ids preserved through local-plan construction; per-block sub_merged_norms is keyed on GLOBAL ids → clean union at top level. No namespace collisions.
• merge strategy — works the same as zero (strategy called per-block on each sub-state, merge_mode honored). Tests above cover zero explicitly; merge support is implicit through dispatch_strategy.

Coordination

After polygram tags v0.14.0, sae-forge can:

1. Bump polygram>=0.14.0.
2. Begin add-block-structured-sae Phase 2 — the --encoding-partition CLI flag + AutoMaterialiseSpec.encoding_partition field + per-row partition_label in ParetoFrontierRow. The downstream change's spec is already filed and locked.

Test plan

• 9 new Phase 2 tests pass locally.
• Phase 1 transition test updated and passes.
• Full polygram suite: 1070 → 1079 (no regressions).
• Ruff clean.
• CI verifies on Linux runners.

🤖 Generated with Claude Code