From 74af44d683ac87cb4fdbf8f0992da5d9f5941ac1 Mon Sep 17 00:00:00 2001 From: jascal Date: Thu, 21 May 2026 09:00:55 -0400 Subject: [PATCH 1/2] =?UTF-8?q?feat(compression):=20add-encoding-partition?= =?UTF-8?q?=20Phase=201=20=E2=80=94=20API=20surface=20+=20scaffolding?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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) --- CHANGELOG.md | 38 +- .../specs/pareto-compression/spec.md | 191 ++++++++ polygram/compression/__init__.py | 12 + polygram/compression/compressor.py | 21 + polygram/compression/partition.py | 325 +++++++++++++ polygram/compression/report.py | 163 ++++++- polygram/config.py | 36 ++ polygram/sae_import.py | 24 + .../test_add_encoding_partition.py | 438 ++++++++++++++++++ 9 files changed, 1246 insertions(+), 2 deletions(-) create mode 100644 openspec/changes/add-encoding-partition/specs/pareto-compression/spec.md create mode 100644 polygram/compression/partition.py create mode 100644 tests/compression/test_add_encoding_partition.py diff --git a/CHANGELOG.md b/CHANGELOG.md index eae0ea3..17daa1b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,7 +2,43 @@ ## Unreleased -(nothing yet) +### Added + +- **`add-encoding-partition` Phase 1** — locks the API surface for + per-block heterogeneous encoding (the gating prerequisite for + sae-forge's `add-block-structured-sae`). 33 new tests, full + suite green (1037 → 1070). + - `BlockSpec` dataclass + per-family validator registry + (`_BLOCK_SPEC_KWARG_VALIDATORS`) for clean Rung6+ extensibility. + Hashable (frozen + custom `__hash__` since `encoding_kwargs` is + a dict). `block_id`, `encoding_class`, `encoding_kwargs`, + `learn_axis_assignment`, `feature_ids`. + - `PartitionCoverageError` (subclasses `ValueError`) + + `validate_partition_coverage(partition, n_features_input=...)` + helper. Names offending feature ids (capped at first 10). + - `make_default_block(...)` convenience constructor for the + "default + heavy override" partition pattern. + - `CompressionConfig.encoding_partition: tuple[BlockSpec, ...] | None` + field. Type/membership/non-empty validation at config + construction. Default `None` preserves the single-encoding path + byte-equivalently. + - `CompressionReport.blocks: tuple[BlockReport, ...] | None` + + new `BlockReport` dataclass mirroring per-block diagnostics. + `CompressionReport.SCHEMA_VERSION` bumped 2 → 3 with back-compat + loader for v2 payloads (defaults `blocks=None`). + - `MAX_CLUSTERS_PER_BLOCK = 10_000` module-level constant for the + global cluster-id namespace. + - `from_sae_lens(..., encoding_partition=...)` kwarg accepted + + type-validated. + - Public exports from `polygram.compression`: `BlockSpec`, + `BlockReport`, `PartitionCoverageError`, `make_default_block`, + `validate_partition_coverage`. + + Phase 2 (the actual per-block dispatch in `Compressor.apply`) is + filed as a follow-up impl. The current `Compressor.apply` raises + `NotImplementedError` with a clear pointer when a partitioned + config is supplied — downstream consumers don't silently get a + single-encoding compression where a partitioned one was requested. ## 0.12.0 — 2026-05-20 diff --git a/openspec/changes/add-encoding-partition/specs/pareto-compression/spec.md b/openspec/changes/add-encoding-partition/specs/pareto-compression/spec.md new file mode 100644 index 0000000..36da1bb --- /dev/null +++ b/openspec/changes/add-encoding-partition/specs/pareto-compression/spec.md @@ -0,0 +1,191 @@ +# pareto-compression Specification (delta) + +## ADDED Requirements + +### Requirement: BlockSpec dataclass + +`BlockSpec` SHALL be a frozen dataclass with five fields: + +| Field | Type | Semantics | +|-------|------|-----------| +| `block_id` | `str` | Non-empty human-readable id; surfaces in `BlockReport.block_id` | +| `encoding_class` | `str` | One of `{"MPSRung1", "Rung3", "Rung4", "Rung5", "HEA_Rung2"}` | +| `encoding_kwargs` | `dict[str, Any]` | Per-family kwargs per the validator registry | +| `learn_axis_assignment` | `bool` | Per-block axis-assignment policy | +| `feature_ids` | `tuple[int, ...]` | Non-empty, non-negative, no duplicates within block | + +`__post_init__` SHALL validate via a module-level +`_BLOCK_SPEC_KWARG_VALIDATORS: dict[str, Callable[[dict], None]]` +registry, not an inline branch chain. Future encoding families add a +new validator + one registry entry. + +`__hash__` SHALL hash a `frozenset` of `encoding_kwargs.items()` +(dict is itself unhashable). This makes `BlockSpec` hashable as +required by downstream consumers' cache-key contracts (e.g. +sae-forge's `compute_cache_key`). + +#### Scenario: Valid BlockSpec with Rung5 + n_amp_qubits + +- **WHEN** `BlockSpec(block_id="heavy", encoding_class="Rung5", encoding_kwargs={"n_amp_qubits": 4}, feature_ids=(0, 1, 2, 3))` is constructed +- **THEN** validation passes and the instance is hashable + +#### Scenario: Rung5 without n_amp_qubits is rejected + +- **WHEN** `BlockSpec(encoding_class="Rung5", encoding_kwargs={}, ...)` is constructed +- **THEN** `ValueError` is raised naming the missing `n_amp_qubits` kwarg + +#### Scenario: MPSRung1 with non-empty kwargs is rejected + +- **WHEN** `BlockSpec(encoding_class="MPSRung1", encoding_kwargs={"foo": 1}, ...)` is constructed +- **THEN** `ValueError` is raised noting the encoding accepts no kwargs + +#### Scenario: feature_ids with duplicates rejected + +- **WHEN** `BlockSpec(feature_ids=(1, 2, 1), ...)` is constructed +- **THEN** `ValueError` names the duplicate id + +#### Scenario: feature_ids with negative ids rejected + +- **WHEN** `BlockSpec(feature_ids=(0, -1), ...)` is constructed +- **THEN** `ValueError` notes feature_ids must be non-negative + +### Requirement: PartitionCoverageError + validate_partition_coverage + +`PartitionCoverageError` SHALL subclass `ValueError`. +`validate_partition_coverage(partition, *, n_features_input)` SHALL +check: + +- **Disjointness**: no feature id in more than one block. Violation + raises `PartitionCoverageError` naming up to 10 duplicate ids with + their containing block_ids. +- **Completeness**: the union of all blocks' `feature_ids` equals + `set(range(n_features_input))`. Missing ids raise an `incomplete` + error; extras raise an `extra` error; both name up to 10 ids. + +Called by `Compressor.apply` immediately after loading the input +SAE; NOT called from `BlockSpec.__post_init__` (because +`n_features_input` is unknown at BlockSpec construction). + +#### Scenario: Disjoint complete partition passes + +- **WHEN** partition `(BlockSpec(feature_ids=(0,1,2,3)), BlockSpec(feature_ids=(4,5,6,7)))` is validated against `n_features_input=8` +- **THEN** no exception is raised + +#### Scenario: Overlapping partition raises naming duplicates + +- **WHEN** two blocks share feature_id 3 +- **THEN** `PartitionCoverageError` is raised mentioning `3` and both blocks' `block_id`s + +#### Scenario: Incomplete partition raises naming missing ids + +- **WHEN** blocks cover `{0,1,2}` against `n_features_input=4` (missing 3) +- **THEN** `PartitionCoverageError` is raised naming the missing id + +### Requirement: CompressionConfig.encoding_partition + +`CompressionConfig` SHALL accept an optional +`encoding_partition: tuple[BlockSpec, ...] | None = None` field. + +When `None` (default), `Compressor.apply` runs the historical +single-encoding path byte-equivalently — no behavioural drift for +existing callers. + +When set, `__post_init__` SHALL validate: +- Type is `tuple` (NOT list — required for `CompressionConfig`'s + frozen-hashable contract). +- Every element is a `BlockSpec` instance. +- Tuple is non-empty. + +Coverage validation against `n_features_input` is deferred to +`Compressor.apply` (where `n_features_input` becomes known). + +#### Scenario: Default encoding_partition is None + +- **WHEN** `CompressionConfig()` is constructed without arguments +- **THEN** `cfg.encoding_partition is None` + +#### Scenario: Partition as list is rejected + +- **WHEN** `CompressionConfig(encoding_partition=[BlockSpec(...)])` (list, not tuple) is constructed +- **THEN** `TypeError` is raised noting "tuple of BlockSpec" + +#### Scenario: Non-BlockSpec member rejected + +- **WHEN** `CompressionConfig(encoding_partition=("not-a-blockspec",))` is constructed +- **THEN** `TypeError` is raised naming `BlockSpec` + +#### Scenario: Empty partition tuple rejected + +- **WHEN** `CompressionConfig(encoding_partition=())` is constructed +- **THEN** `ValueError` is raised noting non-empty requirement + +### Requirement: CompressionReport.blocks + schema v2 → v3 + +`CompressionReport` SHALL gain a +`blocks: tuple[BlockReport, ...] | None = None` field. `BlockReport` +mirrors the top-level diagnostic fields per block (plus the +encoding/feature-id metadata): + +| Field | Type | Default | +|-------|------|---------| +| `block_id` | `str` | required | +| `encoding_class` | `str` | required | +| `encoding_kwargs` | `dict` | required | +| `learn_axis_assignment` | `bool` | required | +| `feature_ids` | `tuple[int, ...]` | required | +| `n_features_kept` | `int` | required | +| `n_features_zeroed` | `int` | required | +| `n_clusters` | `int` | required | +| `cluster_assignments` | `tuple[int, ...] \| None` | `None` | +| `scale_compression_ratio` | `float` | `1.0` | +| `rank_ratio` | `float \| None` | `None` | +| `post_A` | `float \| None` | `None` | +| `forge_mse` | `float \| None` | `None` | +| `informative_metric` | `Literal[...] \| None` | `None` | + +`CompressionReport.SCHEMA_VERSION` SHALL be bumped from `2` to `3`. +`from_json` SHALL accept v2 payloads (no `blocks` key) and default +`blocks=None`. Equality + hash SHALL include `blocks`. + +The module SHALL export `MAX_CLUSTERS_PER_BLOCK = 10_000` as the +constant for the global cluster-id namespace. + +#### Scenario: SCHEMA_VERSION bumped to 3 + +- **WHEN** `polygram.compression.report.SCHEMA_VERSION` is read +- **THEN** the value is `3` + +#### Scenario: v3 report with blocks round-trips through JSON + +- **WHEN** a `CompressionReport` with two populated `BlockReport`s is serialised via `to_json` and reconstructed via `from_json` +- **THEN** the round-tripped instance equals the original (NaN-aware on float fields, exact-equal on the rest) + +#### Scenario: v2 payload loads with blocks=None + +- **GIVEN** a JSON payload with `schema_version=2` and no `blocks` key +- **WHEN** `CompressionReport.from_json(payload)` is called +- **THEN** the call succeeds; `r.schema_version == 2`, `r.blocks is None` + +#### Scenario: v3 report with blocks=None round-trips + +- **WHEN** a `CompressionReport` with `schema_version=3` and `blocks=None` (the unpartitioned-compress case) round-trips via `to_json` + `from_json` +- **THEN** the result has `blocks is None` and equals the original + +### Requirement: Compressor.apply refuses encoding_partition (Phase 1) + +`Compressor.apply` SHALL detect `self.config.encoding_partition is +not None` immediately upon entry and raise `NotImplementedError` +naming `add-encoding-partition` Phase 2 as the load-bearing follow-up. + +This refusal is the Phase 1 contract: the API surface is locked +(`BlockSpec` + `CompressionConfig.encoding_partition` + +`CompressionReport.blocks`) but the actual per-block dispatch is +deferred to Phase 2. Without the refusal, downstream consumers +would silently get a single-encoding compression instead of the +partitioned one they requested. + +#### Scenario: Compressor.apply refuses partitioned config + +- **GIVEN** a `Compressor` constructed with a `CompressionConfig` whose `encoding_partition` is a non-empty tuple +- **WHEN** `compressor.apply(output_checkpoint=...)` is called +- **THEN** `NotImplementedError` is raised mentioning "Phase 2 follow-up" diff --git a/polygram/compression/__init__.py b/polygram/compression/__init__.py index 637e0b5..b3989db 100644 --- a/polygram/compression/__init__.py +++ b/polygram/compression/__init__.py @@ -24,6 +24,12 @@ from polygram.compression.compressor import Compressor from polygram.compression.epoch import EpochCompressor +from polygram.compression.partition import ( + BlockSpec, + PartitionCoverageError, + make_default_block, + validate_partition_coverage, +) from polygram.compression.epoch_report import ( EpochIteration, EpochReport, @@ -39,6 +45,7 @@ SlotPopulation, ) from polygram.compression.report import ( + BlockReport, ClusterPlan, CompressionPlan, CompressionReport, @@ -46,6 +53,8 @@ ) __all__ = [ + "BlockReport", + "BlockSpec", "ClusterPlan", "Compressor", "CompressionPlan", @@ -58,10 +67,13 @@ "Panel", "ParetoOutcome", "ParetoReport", + "PartitionCoverageError", "RegrowPlan", "RegrowReport", "RegrowResult", "RegrowStrategy", "Regrower", "SlotPopulation", + "make_default_block", + "validate_partition_coverage", ] diff --git a/polygram/compression/compressor.py b/polygram/compression/compressor.py index 104b053..3b9d4d6 100644 --- a/polygram/compression/compressor.py +++ b/polygram/compression/compressor.py @@ -833,6 +833,27 @@ 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, " + "which this PR ships, locks the API surface + builds " + "the BlockSpec / CompressionReport.blocks scaffolding " + "but does not yet implement per-block compress + " + "stitch). Set encoding_partition=None for the v1 " + "single-encoding path." + ) + if output_checkpoint is None: raise ValueError( "Compressor.apply: output_checkpoint is required" diff --git a/polygram/compression/partition.py b/polygram/compression/partition.py new file mode 100644 index 0000000..97af4e6 --- /dev/null +++ b/polygram/compression/partition.py @@ -0,0 +1,325 @@ +"""Per-block heterogeneous encoding partition for +:class:`polygram.compression.Compressor`. + +Defines :class:`BlockSpec` (a single block's encoding + axis-assignment ++ feature-id list) and :func:`validate_partition_coverage` (the +disjointness + completeness check). Used by +:class:`polygram.config.CompressionConfig`'s ``encoding_partition`` +field to drive per-block compression in +:meth:`polygram.compression.Compressor.apply`. + +See ``openspec/changes/add-encoding-partition/proposal.md`` for the +full design. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any, Callable, Literal + + +SUPPORTED_ENCODING_CLASSES: frozenset[str] = frozenset( + {"MPSRung1", "Rung3", "Rung4", "Rung5", "HEA_Rung2"} +) + + +class PartitionCoverageError(ValueError): + """Raised when an ``encoding_partition`` fails the disjointness or + completeness coverage check against the input feature count. + + Subclasses :class:`ValueError` so existing + ``except ValueError`` blocks still catch partition-validation + failures; the dedicated class lets callers distinguish coverage + errors from other ``ValueError``s when desired. + """ + + +def _validate_empty_kwargs(kwargs: dict[str, Any]) -> None: + if kwargs: + raise ValueError( + f"BlockSpec: encoding_class accepts no kwargs; " + f"got {sorted(kwargs)}" + ) + + +def _validate_rung5_kwargs(kwargs: dict[str, Any]) -> None: + if "n_amp_qubits" not in kwargs: + raise ValueError( + "BlockSpec: encoding_class='Rung5' requires " + "encoding_kwargs['n_amp_qubits']: int >= 1" + ) + v = kwargs["n_amp_qubits"] + if not isinstance(v, int) or isinstance(v, bool) or v < 1: + raise ValueError( + f"BlockSpec: encoding_kwargs['n_amp_qubits'] must be " + f"int >= 1; got {v!r}" + ) + extras = set(kwargs) - {"n_amp_qubits"} + if extras: + raise ValueError( + f"BlockSpec: encoding_class='Rung5' got unexpected " + f"kwargs: {sorted(extras)}" + ) + + +def _validate_hea_rung2_kwargs(kwargs: dict[str, Any]) -> None: + if "n_qubits" not in kwargs: + raise ValueError( + "BlockSpec: encoding_class='HEA_Rung2' requires " + "encoding_kwargs['n_qubits']: int >= 1" + ) + v = kwargs["n_qubits"] + if not isinstance(v, int) or isinstance(v, bool) or v < 1: + raise ValueError( + f"BlockSpec: encoding_kwargs['n_qubits'] must be " + f"int >= 1; got {v!r}" + ) + extras = set(kwargs) - {"n_qubits"} + if extras: + raise ValueError( + f"BlockSpec: encoding_class='HEA_Rung2' got unexpected " + f"kwargs: {sorted(extras)}" + ) + + +# Per-family validator registry. Adding a future family +# (Rung6, HEA_Rung3, …) is one new validator function + one +# registry entry, not a new branch in BlockSpec.__post_init__. +# See `openspec/changes/add-encoding-partition/design.md` Decision 2b. +_BLOCK_SPEC_KWARG_VALIDATORS: dict[str, Callable[[dict[str, Any]], None]] = { + "MPSRung1": _validate_empty_kwargs, + "Rung3": _validate_empty_kwargs, + "Rung4": _validate_empty_kwargs, + "Rung5": _validate_rung5_kwargs, + "HEA_Rung2": _validate_hea_rung2_kwargs, +} + + +@dataclass(frozen=True) +class BlockSpec: + """One block of a per-block heterogeneous encoding partition. + + Each block defines a subset of feature ids that SHALL be compressed + with the same encoding family + kwargs + axis-assignment policy. + A :class:`polygram.config.CompressionConfig`'s + ``encoding_partition`` is a tuple of these. + + Fields: + block_id: human-readable id (e.g. ``"heavy"``, ``"tail"``). + Surfaces in the ``CompressionReport.blocks`` per-block + report so analysts can correlate blocks with their + partition-manifest entry. + encoding_class: one of + ``{"MPSRung1", "Rung3", "Rung4", "Rung5", "HEA_Rung2"}``. + encoding_kwargs: per-family kwargs: + - ``Rung5`` requires ``{"n_amp_qubits": int >= 1}``. + - ``HEA_Rung2`` requires ``{"n_qubits": int >= 1}``. + - Others MUST be empty. + learn_axis_assignment: per-block axis-assignment policy + (matches the top-level ``CompressionConfig`` field of + the same name). + feature_ids: tuple of non-negative ints, non-empty, no + duplicates within this block. Disjointness across blocks + + completeness vs n_features_input are checked separately + at ``Compressor.apply`` time via + :func:`validate_partition_coverage`. + """ + + block_id: str + encoding_class: str + encoding_kwargs: dict[str, Any] = field(default_factory=dict) + learn_axis_assignment: bool = False + feature_ids: tuple[int, ...] = () + + def __post_init__(self) -> None: + if not isinstance(self.block_id, str) or not self.block_id: + raise ValueError( + f"BlockSpec: block_id must be a non-empty str; got " + f"{self.block_id!r}" + ) + if self.encoding_class not in SUPPORTED_ENCODING_CLASSES: + raise ValueError( + f"BlockSpec: encoding_class must be one of " + f"{sorted(SUPPORTED_ENCODING_CLASSES)}; got " + f"{self.encoding_class!r}" + ) + if not isinstance(self.encoding_kwargs, dict): + raise TypeError( + f"BlockSpec: encoding_kwargs must be a dict; got " + f"{type(self.encoding_kwargs).__name__}" + ) + # Per-family validator (Decision 2b — registry, not elif chain). + validator = _BLOCK_SPEC_KWARG_VALIDATORS[self.encoding_class] + validator(self.encoding_kwargs) + if not isinstance(self.learn_axis_assignment, bool): + raise TypeError( + f"BlockSpec: learn_axis_assignment must be bool; got " + f"{type(self.learn_axis_assignment).__name__}" + ) + if not isinstance(self.feature_ids, tuple): + raise TypeError( + f"BlockSpec: feature_ids must be a tuple; got " + f"{type(self.feature_ids).__name__}" + ) + if not self.feature_ids: + raise ValueError("BlockSpec: feature_ids must be non-empty") + seen: set[int] = set() + for fid in self.feature_ids: + if not isinstance(fid, int) or isinstance(fid, bool): + raise TypeError( + f"BlockSpec: feature_ids must contain int; got " + f"{type(fid).__name__}" + ) + if fid < 0: + raise ValueError( + f"BlockSpec: feature_ids must be non-negative; " + f"got {fid}" + ) + if fid in seen: + raise ValueError( + f"BlockSpec: feature_ids contains duplicate id {fid} " + f"in block {self.block_id!r}" + ) + seen.add(fid) + + # `frozen=True` provides default __eq__ / __hash__, but `dict` is + # unhashable so __hash__ would crash on a populated encoding_kwargs. + # Override to hash a frozenset of items. + def __hash__(self) -> int: + return hash(( + self.block_id, + self.encoding_class, + frozenset(self.encoding_kwargs.items()), + self.learn_axis_assignment, + self.feature_ids, + )) + + +def validate_partition_coverage( + partition: tuple[BlockSpec, ...], + *, + n_features_input: int, +) -> None: + """Verify a partition's disjointness + completeness against the + input feature count. + + Disjointness: no feature id appears in more than one block. + Completeness: the union of all blocks' ``feature_ids`` is exactly + ``set(range(n_features_input))``. + + On violation raises :class:`PartitionCoverageError` naming the + offending ids (capped at first 10) so analysts can edit the + manifest without scanning the whole list. + + Called by :meth:`polygram.compression.Compressor.apply` immediately + after loading the input SAE (when ``n_features_input`` becomes + known). Not called from ``BlockSpec.__post_init__`` because a + BlockSpec is built before the input feature count is known. + """ + if n_features_input < 0: + raise ValueError( + f"validate_partition_coverage: n_features_input must be " + f"non-negative; got {n_features_input}" + ) + + seen: dict[int, str] = {} # fid -> first-seen block_id + duplicates: list[tuple[int, str, str]] = [] # (fid, block_a, block_b) + for block in partition: + for fid in block.feature_ids: + if fid in seen: + duplicates.append((fid, seen[fid], block.block_id)) + else: + seen[fid] = block.block_id + + if duplicates: + dupe_strs = [ + f"{fid} (in blocks {a!r} and {b!r})" + for fid, a, b in duplicates[:10] + ] + more = "" + if len(duplicates) > 10: + more = f" ... and {len(duplicates) - 10} more" + raise PartitionCoverageError( + f"encoding_partition has overlapping feature_ids: " + f"{', '.join(dupe_strs)}{more}" + ) + + expected = set(range(n_features_input)) + actual = set(seen.keys()) + missing = sorted(expected - actual) + extras = sorted(actual - expected) + + if missing: + truncated = missing[:10] + more = f" ... and {len(missing) - 10} more" if len(missing) > 10 else "" + raise PartitionCoverageError( + f"encoding_partition is incomplete: feature_ids " + f"{truncated}{more} are not covered by any block " + f"(n_features_input={n_features_input})" + ) + if extras: + truncated = extras[:10] + more = f" ... and {len(extras) - 10} more" if len(extras) > 10 else "" + raise PartitionCoverageError( + f"encoding_partition has extra feature_ids " + f"{truncated}{more} outside the input range " + f"[0, {n_features_input})" + ) + + +def make_default_block( + *, + encoding_class: Literal["MPSRung1", "Rung3", "Rung4", "Rung5", "HEA_Rung2"], + n_features_input: int, + encoding_kwargs: dict[str, Any] | None = None, + learn_axis_assignment: bool = False, + excluded_feature_ids: set[int] | frozenset[int] = frozenset(), + block_id: str = "default", +) -> BlockSpec: + """Build a :class:`BlockSpec` covering all features in + ``range(n_features_input)`` minus ``excluded_feature_ids``. + + Useful for "default + heavy override" partition patterns where + a small set of heavy features get a custom block and the rest + falls through to a default. See + ``openspec/changes/add-encoding-partition/design.md`` Decision 2c. + + Example: + >>> heavy = BlockSpec( + ... block_id="heavy", encoding_class="Rung5", + ... encoding_kwargs={"n_amp_qubits": 4}, + ... feature_ids=(0, 1, 2, 3), + ... ) + >>> tail = make_default_block( + ... encoding_class="MPSRung1", + ... n_features_input=128, + ... excluded_feature_ids={0, 1, 2, 3}, + ... ) + >>> partition = (heavy, tail) + """ + excluded = frozenset(int(x) for x in excluded_feature_ids) + feature_ids = tuple( + fid for fid in range(int(n_features_input)) if fid not in excluded + ) + if not feature_ids: + raise ValueError( + f"make_default_block: excluded_feature_ids covers every " + f"feature in range({n_features_input}); the resulting " + f"block would be empty" + ) + return BlockSpec( + block_id=block_id, + encoding_class=encoding_class, + encoding_kwargs=encoding_kwargs if encoding_kwargs is not None else {}, + learn_axis_assignment=learn_axis_assignment, + feature_ids=feature_ids, + ) + + +__all__ = [ + "BlockSpec", + "PartitionCoverageError", + "SUPPORTED_ENCODING_CLASSES", + "make_default_block", + "validate_partition_coverage", +] diff --git a/polygram/compression/report.py b/polygram/compression/report.py index 3f6bb24..666b980 100644 --- a/polygram/compression/report.py +++ b/polygram/compression/report.py @@ -39,7 +39,13 @@ from polygram.dictionary import Dictionary -SCHEMA_VERSION = 2 +SCHEMA_VERSION = 3 +# Per-block global-cluster-id namespace cap. Each block's local +# cluster ids in [0, n_clusters_in_block) get encoded as +# `block_index * MAX_CLUSTERS_PER_BLOCK + local_id` in the top-level +# CompressionReport's per-feature cluster ids. See +# `openspec/changes/add-encoding-partition/design.md` Decision 2. +MAX_CLUSTERS_PER_BLOCK = 10_000 @dataclass(frozen=True) @@ -90,6 +96,35 @@ def n_features_kept(self) -> int: return len(self.clusters) +@dataclass(frozen=True, eq=False) +class BlockReport: + """Per-block diagnostic record emitted by `Compressor.apply` when + a `CompressionConfig.encoding_partition` is set. + + Mirrors the per-block subset of the top-level + :class:`CompressionReport` fields (n_features_kept, + n_features_zeroed, n_clusters, plus the diagnostic floats), keyed + by the block's analyst-supplied ``block_id``. + + Added by ``add-encoding-partition``. v3 schema. + """ + + block_id: str + encoding_class: str + encoding_kwargs: dict + learn_axis_assignment: bool + feature_ids: tuple[int, ...] + n_features_kept: int + n_features_zeroed: int + n_clusters: int + cluster_assignments: tuple[int, ...] | None = None + scale_compression_ratio: float = 1.0 + rank_ratio: float | None = None + post_A: float | None = None + forge_mse: float | None = None + informative_metric: Literal["post_A", "both", "forge_mse"] | None = None + + @dataclass(frozen=True, eq=False) class CompressionReport: """Post-`apply()` artifact carrying provenance + the applied plan.""" @@ -111,6 +146,13 @@ class CompressionReport: post_A: float | None = None forge_mse: float | None = None informative_metric: Literal["post_A", "both", "forge_mse"] | None = None + # v3: per-block heterogeneous-encoding diagnostics. When the + # `Compressor` was driven by a `CompressionConfig.encoding_partition`, + # `blocks` carries one BlockReport per partition block. When the + # compression used a single top-level encoding (the historical path), + # `blocks` is None. Schema bump from 2 to 3; from_json defaults + # `blocks=None` for v2 payloads. + blocks: tuple[BlockReport, ...] | None = None # ---- JSON ------------------------------------------------------ @@ -185,6 +227,14 @@ def from_json(cls, source: str | os.PathLike) -> "CompressionReport": else None ) informative = payload.get("informative_metric") + + # v3 (add-encoding-partition): `blocks` is a list of BlockReport + # dicts when populated, or absent/null in v2 payloads. + blocks_raw = payload.get("blocks") + blocks: tuple[BlockReport, ...] | None = None + if blocks_raw is not None: + blocks = tuple(_block_from_dict(b) for b in blocks_raw) + return cls( schema_version=int(payload["schema_version"]), source_checkpoint=str(payload["source_checkpoint"]), @@ -209,6 +259,7 @@ def from_json(cls, source: str | os.PathLike) -> "CompressionReport": post_A=post_A, forge_mse=forge_mse, informative_metric=informative, + blocks=blocks, ) # ---- Equality -------------------------------------------------- @@ -239,9 +290,17 @@ def __eq__(self, other: object) -> bool: and floats_eq(self.post_A, other.post_A) and floats_eq(self.forge_mse, other.forge_mse) and self.informative_metric == other.informative_metric + and _blocks_eq(self.blocks, other.blocks) ) def __hash__(self) -> int: + # blocks contains dict (encoding_kwargs) which is unhashable; + # hash by structural identifiers (block_id, encoding_class) + # of each block instead of the full BlockReport tuple. + blocks_hash = ( + tuple((b.block_id, b.encoding_class) for b in self.blocks) + if self.blocks is not None else None + ) return hash(( self.schema_version, self.source_checkpoint_sha256, @@ -251,6 +310,7 @@ def __hash__(self) -> int: self.rank_ratio, self.post_A, self.forge_mse, + blocks_hash, )) # ---- Internal -------------------------------------------------- @@ -292,6 +352,12 @@ def _serialize(self) -> str: float(self.forge_mse) if self.forge_mse is not None else None ), "informative_metric": self.informative_metric, + # v3 (add-encoding-partition): list of per-block dicts, or + # null when the compression used a single top-level encoding. + "blocks": ( + [_block_to_dict(b) for b in self.blocks] + if self.blocks is not None else None + ), } return json.dumps(payload, sort_keys=True, separators=(",", ":")) @@ -342,3 +408,98 @@ def _opt_float(v: Any) -> float | None: cluster_norm_std=_opt_float(raw.get("cluster_norm_std")), merged_norm=_opt_float(raw.get("merged_norm")), ) + + +def _block_to_dict(b: BlockReport) -> dict[str, Any]: + """Serialise a :class:`BlockReport` to a JSON-roundtrippable dict. + Mirrors `_cluster_to_dict`'s shape for the new v3 partition fields. + """ + return { + "block_id": str(b.block_id), + "encoding_class": str(b.encoding_class), + "encoding_kwargs": dict(b.encoding_kwargs), + "learn_axis_assignment": bool(b.learn_axis_assignment), + "feature_ids": [int(f) for f in b.feature_ids], + "n_features_kept": int(b.n_features_kept), + "n_features_zeroed": int(b.n_features_zeroed), + "n_clusters": int(b.n_clusters), + "cluster_assignments": ( + list(b.cluster_assignments) + if b.cluster_assignments is not None else None + ), + "scale_compression_ratio": float(b.scale_compression_ratio), + "rank_ratio": ( + float(b.rank_ratio) if b.rank_ratio is not None else None + ), + "post_A": ( + float(b.post_A) if b.post_A is not None else None + ), + "forge_mse": ( + float(b.forge_mse) if b.forge_mse is not None else None + ), + "informative_metric": b.informative_metric, + } + + +def _block_from_dict(raw: dict[str, Any]) -> BlockReport: + def _opt_float(v: Any) -> float | None: + return None if v is None else float(v) + + ca = raw.get("cluster_assignments") + return BlockReport( + block_id=str(raw["block_id"]), + encoding_class=str(raw["encoding_class"]), + encoding_kwargs=dict(raw.get("encoding_kwargs") or {}), + learn_axis_assignment=bool(raw.get("learn_axis_assignment", False)), + feature_ids=tuple(int(f) for f in raw["feature_ids"]), + n_features_kept=int(raw["n_features_kept"]), + n_features_zeroed=int(raw["n_features_zeroed"]), + n_clusters=int(raw["n_clusters"]), + cluster_assignments=( + tuple(int(x) for x in ca) if ca is not None else None + ), + scale_compression_ratio=float(raw.get("scale_compression_ratio", 1.0)), + rank_ratio=_opt_float(raw.get("rank_ratio")), + post_A=_opt_float(raw.get("post_A")), + forge_mse=_opt_float(raw.get("forge_mse")), + informative_metric=raw.get("informative_metric"), + ) + + +def _blocks_eq( + a: "tuple[BlockReport, ...] | None", + b: "tuple[BlockReport, ...] | None", +) -> bool: + """NaN-aware element-wise equality for the optional `blocks` field. + Either both `None` (single-encoding compression) or same-length + tuples whose corresponding BlockReports compare equal field-by-field. + """ + if a is None and b is None: + return True + if a is None or b is None: + return False + if len(a) != len(b): + return False + for x, y in zip(a, b): + if not _block_report_eq(x, y): + return False + return True + + +def _block_report_eq(x: BlockReport, y: BlockReport) -> bool: + return ( + x.block_id == y.block_id + and x.encoding_class == y.encoding_class + and x.encoding_kwargs == y.encoding_kwargs + and x.learn_axis_assignment == y.learn_axis_assignment + and x.feature_ids == y.feature_ids + and x.n_features_kept == y.n_features_kept + and x.n_features_zeroed == y.n_features_zeroed + and x.n_clusters == y.n_clusters + and x.cluster_assignments == y.cluster_assignments + and floats_eq(x.scale_compression_ratio, y.scale_compression_ratio) + and floats_eq(x.rank_ratio, y.rank_ratio) + and floats_eq(x.post_A, y.post_A) + and floats_eq(x.forge_mse, y.forge_mse) + and x.informative_metric == y.informative_metric + ) diff --git a/polygram/config.py b/polygram/config.py index 08b910a..bdcfdd5 100644 --- a/polygram/config.py +++ b/polygram/config.py @@ -309,6 +309,18 @@ class CompressionConfig(_ConfigMixin): confirmer: str | None = None target_n_features_kept: int | None = None score_field: str = "polygram_overlap" + # Per-block heterogeneous encoding partition. When non-None, + # `Compressor.apply` SHALL dispatch per-block (each block's + # features get compressed with its own encoding family + + # kwargs + axis-assignment policy). Coverage validation + # (disjointness + completeness vs n_features_input) runs at + # `Compressor.apply` time, not here. See + # ``openspec/changes/add-encoding-partition/proposal.md``. + # + # `tuple[BlockSpec, ...]` (not list) — required for + # CompressionConfig's frozen=True hash + downstream + # cache-key contracts (e.g. sae-forge's compute_cache_key). + encoding_partition: "tuple | None" = None def __post_init__(self) -> None: if self.strategy not in _SUPPORTED_STRATEGIES: @@ -342,6 +354,30 @@ def __post_init__(self) -> None: f"CompressionConfig: score_field must be one of " f"{_SUPPORTED_SCORE_FIELDS}; got {self.score_field!r}" ) + if self.encoding_partition is not None: + # Lazy import to keep the no-partition call path + # independent of partition.py (which itself is torch- + # free; the lazy import is for module-load-order + # cleanliness against the wider polygram surface). + from polygram.compression.partition import BlockSpec + if not isinstance(self.encoding_partition, tuple): + raise TypeError( + f"CompressionConfig: encoding_partition must be a " + f"tuple of BlockSpec; got " + f"{type(self.encoding_partition).__name__}" + ) + for i, block in enumerate(self.encoding_partition): + if not isinstance(block, BlockSpec): + raise TypeError( + f"CompressionConfig: encoding_partition[{i}] " + f"must be a BlockSpec; got " + f"{type(block).__name__}" + ) + if not self.encoding_partition: + raise ValueError( + "CompressionConfig: encoding_partition must be " + "either None or a non-empty tuple of BlockSpec" + ) # --------------------------------------------------------------------------- diff --git a/polygram/sae_import.py b/polygram/sae_import.py index 514fd55..512dd41 100644 --- a/polygram/sae_import.py +++ b/polygram/sae_import.py @@ -630,6 +630,7 @@ def from_sae_lens( assign_amp_knobs: bool | None = None, assign_phase_knobs: bool | None = None, learn_axis_assignment: "bool | object | None" = None, + encoding_partition: "tuple | None" = None, ) -> tuple["Dictionary | ClusteredDictionary", SelectionReport]: """Build a `Dictionary` from an explicit subset of SAE features. @@ -681,6 +682,29 @@ def from_sae_lens( cfg = config if config is not None else SAEImportConfig() resolved_profile = _resolve_profile(profile, cfg) + # `encoding_partition` (add-encoding-partition Phase 1) — accepted + # to lock the API surface that sae-forge's `add-block-structured-sae` + # consumes. Phase 2 (the actual per-block dispatch) lives in + # `Compressor.apply`; calling sites that supply a partition AND + # invoke `Compressor.apply` get a clean `NotImplementedError` + # there. `from_sae_lens` itself builds Dictionaries (not + # compressed checkpoints), so accepting the kwarg here is a no-op + # routing decision: the partition is stored on + # `CompressionConfig` by callers and consumed at compress time. + if encoding_partition is not None: + from polygram.compression.partition import BlockSpec + if not isinstance(encoding_partition, tuple): + raise TypeError( + f"from_sae_lens: encoding_partition must be a tuple of " + f"BlockSpec; got {type(encoding_partition).__name__}" + ) + for i, block in enumerate(encoding_partition): + if not isinstance(block, BlockSpec): + raise TypeError( + f"from_sae_lens: encoding_partition[{i}] must be " + f"a BlockSpec; got {type(block).__name__}" + ) + if assign_gamma is None: assign_gamma = cfg.assign_gamma if assign_amp_knobs is None: diff --git a/tests/compression/test_add_encoding_partition.py b/tests/compression/test_add_encoding_partition.py new file mode 100644 index 0000000..3d360a4 --- /dev/null +++ b/tests/compression/test_add_encoding_partition.py @@ -0,0 +1,438 @@ +"""Tests for `add-encoding-partition` Phase 1 — BlockSpec, partition +coverage validation, CompressionConfig.encoding_partition, and the +CompressionReport.blocks schema bump. + +Phase 1 of this change locks the API surface and ships the +scaffolding; the actual per-block dispatch in `Compressor.apply` is +a Phase 2 follow-up. This test file pins the contract that Phase 2 +implementations will need to honour. +""" + +from __future__ import annotations + +import json + +import pytest + +from polygram.compression import ( + BlockReport, + BlockSpec, + CompressionReport, + PartitionCoverageError, + make_default_block, + validate_partition_coverage, +) +from polygram.compression.report import ( + CompressionPlan, + SCHEMA_VERSION, + MAX_CLUSTERS_PER_BLOCK, +) +from polygram.config import CompressionConfig + + +# --------------------------------------------------------------------------- +# §9.1 — BlockSpec validation +# --------------------------------------------------------------------------- + + +def test_block_spec_accepts_mps_rung1_with_empty_kwargs(): + b = BlockSpec(block_id="b", encoding_class="MPSRung1", feature_ids=(0, 1, 2)) + assert b.encoding_class == "MPSRung1" + assert b.feature_ids == (0, 1, 2) + + +def test_block_spec_rejects_unknown_encoding_class(): + with pytest.raises(ValueError, match="encoding_class"): + BlockSpec(block_id="b", encoding_class="MadeUp", feature_ids=(0,)) + + +def test_block_spec_rejects_missing_n_amp_qubits_for_rung5(): + with pytest.raises(ValueError, match="n_amp_qubits"): + BlockSpec(block_id="b", encoding_class="Rung5", + encoding_kwargs={}, feature_ids=(0,)) + + +def test_block_spec_rejects_negative_n_amp_qubits(): + with pytest.raises(ValueError, match="n_amp_qubits"): + BlockSpec(block_id="b", encoding_class="Rung5", + encoding_kwargs={"n_amp_qubits": -1}, feature_ids=(0,)) + + +def test_block_spec_rejects_extra_kwargs_for_mps_rung1(): + with pytest.raises(ValueError, match="no kwargs"): + BlockSpec(block_id="b", encoding_class="MPSRung1", + encoding_kwargs={"foo": 1}, feature_ids=(0,)) + + +def test_block_spec_rejects_missing_n_qubits_for_hea_rung2(): + with pytest.raises(ValueError, match="n_qubits"): + BlockSpec(block_id="b", encoding_class="HEA_Rung2", + encoding_kwargs={}, feature_ids=(0,)) + + +def test_block_spec_rejects_empty_feature_ids(): + with pytest.raises(ValueError, match="non-empty"): + BlockSpec(block_id="b", encoding_class="MPSRung1", + feature_ids=()) + + +def test_block_spec_rejects_duplicate_feature_ids_within_block(): + with pytest.raises(ValueError, match="duplicate"): + BlockSpec(block_id="b", encoding_class="MPSRung1", + feature_ids=(1, 2, 1)) + + +def test_block_spec_rejects_negative_feature_id(): + with pytest.raises(ValueError, match="non-negative"): + BlockSpec(block_id="b", encoding_class="MPSRung1", + feature_ids=(0, -1)) + + +def test_block_spec_rejects_empty_block_id(): + with pytest.raises(ValueError, match="block_id"): + BlockSpec(block_id="", encoding_class="MPSRung1", feature_ids=(0,)) + + +def test_block_spec_is_hashable(): + """frozen=True + custom __hash__ → BlockSpec is hashable even + though encoding_kwargs is a dict. Required for + CompressionConfig's downstream cache-key contract.""" + b = BlockSpec(block_id="heavy", encoding_class="Rung5", + encoding_kwargs={"n_amp_qubits": 4}, + feature_ids=(0, 1, 2, 3)) + h = hash(b) + assert isinstance(h, int) + # Re-hash with the same fields → same hash. + b2 = BlockSpec(block_id="heavy", encoding_class="Rung5", + encoding_kwargs={"n_amp_qubits": 4}, + feature_ids=(0, 1, 2, 3)) + assert hash(b) == hash(b2) + + +# --------------------------------------------------------------------------- +# §9.2 — Partition coverage +# --------------------------------------------------------------------------- + + +def test_partition_coverage_disjoint_complete_passes(): + a = BlockSpec(block_id="a", encoding_class="MPSRung1", + feature_ids=(0, 1, 2, 3)) + b = BlockSpec(block_id="b", encoding_class="MPSRung1", + feature_ids=(4, 5, 6, 7)) + # Should not raise + validate_partition_coverage((a, b), n_features_input=8) + + +def test_partition_coverage_overlap_raises_naming_duplicates(): + 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)) # 3 overlaps + with pytest.raises(PartitionCoverageError) as exc_info: + validate_partition_coverage((a, b), n_features_input=6) + assert "3" in str(exc_info.value) + assert "'a'" in str(exc_info.value) + assert "'b'" in str(exc_info.value) + + +def test_partition_coverage_missing_raises_naming_holes(): + a = BlockSpec(block_id="a", encoding_class="MPSRung1", + feature_ids=(0, 1, 2)) # missing 3 + with pytest.raises(PartitionCoverageError, match="incomplete"): + validate_partition_coverage((a,), n_features_input=4) + + +def test_partition_coverage_extras_only_raises_naming_extras(): + """Pure-extras case (no missing): blocks cover [0..7] inclusive but + n_features_input=4 → 4,5,6,7 are extras.""" + a = BlockSpec(block_id="a", encoding_class="MPSRung1", + feature_ids=(0, 1, 2, 3, 4, 5, 6, 7)) + with pytest.raises(PartitionCoverageError, match="extra"): + validate_partition_coverage((a,), n_features_input=4) + + +def test_partition_coverage_rejects_negative_n_features_input(): + a = BlockSpec(block_id="a", encoding_class="MPSRung1", + feature_ids=(0,)) + with pytest.raises(ValueError, match="non-negative"): + validate_partition_coverage((a,), n_features_input=-1) + + +# --------------------------------------------------------------------------- +# §9.5 — make_default_block helper +# --------------------------------------------------------------------------- + + +def test_make_default_block_covers_all_except_excluded(): + block = make_default_block( + encoding_class="MPSRung1", + n_features_input=16, + excluded_feature_ids={0, 1, 2, 3}, + ) + assert block.feature_ids == tuple(range(4, 16)) + assert block.encoding_class == "MPSRung1" + assert block.block_id == "default" + + +def test_make_default_block_with_no_exclusions_covers_all(): + block = make_default_block( + encoding_class="MPSRung1", + n_features_input=8, + ) + assert block.feature_ids == (0, 1, 2, 3, 4, 5, 6, 7) + + +def test_make_default_block_refuses_empty_result(): + """If every feature is excluded, the resulting block would be empty, + which violates BlockSpec.__post_init__'s non-empty requirement. + make_default_block catches this earlier with a clearer message.""" + with pytest.raises(ValueError, match="empty"): + make_default_block( + encoding_class="MPSRung1", + n_features_input=4, + excluded_feature_ids={0, 1, 2, 3}, + ) + + +def test_make_default_block_with_heavy_override_partition_works(): + """Realistic 'default + heavy override' pattern: a small heavy + block + a default tail computed by the helper, validated together.""" + heavy = BlockSpec( + block_id="heavy", encoding_class="Rung5", + encoding_kwargs={"n_amp_qubits": 4}, + feature_ids=(0, 1, 2, 3), + ) + tail = make_default_block( + encoding_class="MPSRung1", + n_features_input=16, + excluded_feature_ids={0, 1, 2, 3}, + ) + # Coverage validates cleanly + validate_partition_coverage((heavy, tail), n_features_input=16) + + +# --------------------------------------------------------------------------- +# §9.4 — CompressionConfig.encoding_partition validation +# --------------------------------------------------------------------------- + + +def test_compression_config_default_partition_is_none(): + cfg = CompressionConfig() + assert cfg.encoding_partition is None + + +def test_compression_config_accepts_partition_tuple(): + a = BlockSpec(block_id="a", encoding_class="MPSRung1", + feature_ids=(0,)) + cfg = CompressionConfig(encoding_partition=(a,)) + assert cfg.encoding_partition == (a,) + + +def test_compression_config_rejects_partition_as_list(): + """The locked API surface is `tuple[BlockSpec, ...]`, not list, + for hashability. Passing a list raises TypeError.""" + a = BlockSpec(block_id="a", encoding_class="MPSRung1", + feature_ids=(0,)) + with pytest.raises(TypeError, match="tuple of BlockSpec"): + CompressionConfig(encoding_partition=[a]) + + +def test_compression_config_rejects_non_blockspec_in_partition(): + with pytest.raises(TypeError, match="BlockSpec"): + CompressionConfig(encoding_partition=("not-a-blockspec",)) + + +def test_compression_config_rejects_empty_partition_tuple(): + with pytest.raises(ValueError, match="non-empty"): + CompressionConfig(encoding_partition=()) + + +# --------------------------------------------------------------------------- +# §9.4 — CompressionReport schema v2 → v3 round-trip with blocks +# --------------------------------------------------------------------------- + + +def _empty_plan() -> CompressionPlan: + return CompressionPlan(clusters=(), feature_ids=()) + + +def _make_report_with_blocks() -> CompressionReport: + blocks = ( + BlockReport( + block_id="heavy", + encoding_class="Rung5", + encoding_kwargs={"n_amp_qubits": 4}, + learn_axis_assignment=True, + feature_ids=(0, 1, 2, 3), + n_features_kept=2, + n_features_zeroed=2, + n_clusters=2, + cluster_assignments=(0, 0, 1, 1), + scale_compression_ratio=0.85, + rank_ratio=0.95, + post_A=0.97, + forge_mse=0.001, + informative_metric="both", + ), + BlockReport( + block_id="tail", + encoding_class="MPSRung1", + encoding_kwargs={}, + learn_axis_assignment=False, + feature_ids=(4, 5, 6, 7), + n_features_kept=4, + n_features_zeroed=0, + n_clusters=4, + cluster_assignments=(0, 1, 2, 3), + ), + ) + return CompressionReport( + schema_version=SCHEMA_VERSION, + source_checkpoint="/x/sae.safetensors", + source_checkpoint_sha256="a" * 64, + output_checkpoint="/x/sae.compressed.safetensors", + output_checkpoint_sha256="b" * 64, + validation_report_dictionary_name="d", + validation_report_schema_version=1, + strategy="merge", + plan=_empty_plan(), + n_features_zeroed=2, + n_features_kept=6, + n_clusters=6, + scale_compression_ratio=0.92, + blocks=blocks, + ) + + +def test_compression_report_schema_version_is_3(): + assert SCHEMA_VERSION == 3 + + +def test_compression_report_blocks_default_is_none(): + r = CompressionReport( + schema_version=SCHEMA_VERSION, + source_checkpoint="/x", source_checkpoint_sha256="a" * 64, + output_checkpoint="/y", output_checkpoint_sha256="b" * 64, + validation_report_dictionary_name="d", + validation_report_schema_version=1, + strategy="merge", plan=_empty_plan(), + n_features_zeroed=0, n_features_kept=0, n_clusters=0, + ) + assert r.blocks is None + + +def test_compression_report_blocks_roundtrip_via_json(): + r = _make_report_with_blocks() + rt = CompressionReport.from_json(r.to_json()) + assert rt.blocks is not None + assert len(rt.blocks) == 2 + assert rt == r + + +def test_compression_report_v2_payload_loads_with_blocks_none(): + """A v2-schema payload (no `blocks` key) SHALL load without error, + defaulting blocks=None — preserves back-compat with the v2 contract + that pre-add-encoding-partition consumers wrote.""" + v2_payload = { + "schema_version": 2, + "source_checkpoint": "/x", + "source_checkpoint_sha256": "a" * 64, + "output_checkpoint": "/y", + "output_checkpoint_sha256": "b" * 64, + "validation_report_dictionary_name": "d", + "validation_report_schema_version": 1, + "strategy": "merge", + "feature_ids": [], + "clusters": [], + "n_features_zeroed": 0, + "n_features_kept": 0, + "n_clusters": 0, + "scale_compression_ratio": 1.0, + "rank_ratio": None, + "post_A": None, + "forge_mse": None, + "informative_metric": None, + # NOTE: no "blocks" key + } + r = CompressionReport.from_json(json.dumps(v2_payload)) + assert r.schema_version == 2 + assert r.blocks is None + + +def test_compression_report_v3_payload_blocks_none_round_trip(): + """A v3-schema payload with blocks=None (the unpartitioned compress + case) round-trips cleanly — blocks is None on both sides.""" + r = CompressionReport( + schema_version=SCHEMA_VERSION, + source_checkpoint="/x", source_checkpoint_sha256="a" * 64, + output_checkpoint="/y", output_checkpoint_sha256="b" * 64, + validation_report_dictionary_name="d", + validation_report_schema_version=1, + strategy="merge", plan=_empty_plan(), + n_features_zeroed=0, n_features_kept=0, n_clusters=0, + blocks=None, + ) + rt = CompressionReport.from_json(r.to_json()) + assert rt.blocks is None + assert rt == r + + +def test_compression_report_blocks_equality_field_by_field(): + """Two reports with identical blocks compare equal; differing + on any block field breaks equality.""" + r1 = _make_report_with_blocks() + r2 = _make_report_with_blocks() + assert r1 == r2 + # Mutate one field on the first block's BlockReport + blocks_mod = list(r2.blocks) + block0 = blocks_mod[0] + import dataclasses + blocks_mod[0] = dataclasses.replace(block0, n_features_kept=99) + r3 = dataclasses.replace(r2, blocks=tuple(blocks_mod)) + assert r1 != r3 + + +# --------------------------------------------------------------------------- +# §9.7 — Compressor.apply refusal (Phase 1 stub; Phase 2 follow-up) +# --------------------------------------------------------------------------- + + +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, + ) + with pytest.raises(NotImplementedError, match="Phase 2 follow-up"): + compressor.apply(output_checkpoint=dst) + + +# --------------------------------------------------------------------------- +# §9.6 — Global cluster-id namespace constant +# --------------------------------------------------------------------------- + + +def test_max_clusters_per_block_constant_present(): + """The per-block global cluster-id namespace constant SHALL be + exported so Phase 2 implementations + downstream consumers can + reference the load-bearing cap. See Decision 2 in design.md.""" + assert MAX_CLUSTERS_PER_BLOCK == 10_000 From 4f0936204e335cba79772248ba29ecca305d99e2 Mon Sep 17 00:00:00 2001 From: jascal Date: Thu, 21 May 2026 09:08:07 -0400 Subject: [PATCH 2/2] =?UTF-8?q?polish(partition):=20review=20nits=20?= =?UTF-8?q?=E2=80=94=20actionable=20error=20pointer=20+=20usage=20example?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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) --- polygram/compression/compressor.py | 10 +++++-- polygram/compression/partition.py | 47 ++++++++++++++++++++++++++++++ 2 files changed, 54 insertions(+), 3 deletions(-) diff --git a/polygram/compression/compressor.py b/polygram/compression/compressor.py index 3b9d4d6..0642023 100644 --- a/polygram/compression/compressor.py +++ b/polygram/compression/compressor.py @@ -847,11 +847,15 @@ def apply( raise NotImplementedError( "Compressor.apply: encoding_partition support is a " "Phase 2 follow-up of add-encoding-partition (Phase 1, " - "which this PR ships, locks the API surface + builds " + "currently shipped, locks the API surface + builds " "the BlockSpec / CompressionReport.blocks scaffolding " "but does not yet implement per-block compress + " - "stitch). Set encoding_partition=None for the v1 " - "single-encoding path." + "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." ) if output_checkpoint is None: diff --git a/polygram/compression/partition.py b/polygram/compression/partition.py index 97af4e6..6887032 100644 --- a/polygram/compression/partition.py +++ b/polygram/compression/partition.py @@ -10,6 +10,53 @@ See ``openspec/changes/add-encoding-partition/proposal.md`` for the full design. + +Usage example — "default + heavy override" pattern +================================================== + +The common analyst workflow: a small set of heavy features deserves a +high-capacity encoding (Rung5 with learn-axis-assignment), the long +tail gets the cheap default (MPSRung1):: + + from polygram.compression import ( + BlockSpec, + CompressionConfig, + make_default_block, + validate_partition_coverage, + ) + + # Heavy features (analyst-supplied; e.g. top-K by firing rate) + heavy = BlockSpec( + block_id="heavy", + encoding_class="Rung5", + encoding_kwargs={"n_amp_qubits": 4}, + learn_axis_assignment=True, + feature_ids=(0, 1, 2, 3), + ) + + # Tail defaults to MPSRung1 over the remaining ids + tail = make_default_block( + encoding_class="MPSRung1", + n_features_input=128, # total feature count of the input SAE + excluded_feature_ids={0, 1, 2, 3}, + ) + + partition = (heavy, tail) + + # Validate before passing to CompressionConfig — covers disjointness + # + completeness against the input SAE's feature count. + validate_partition_coverage(partition, n_features_input=128) + + config = CompressionConfig(encoding_partition=partition) + # Phase 1: Compressor.apply will refuse with NotImplementedError + # because the per-block dispatch is the Phase 2 follow-up. Phase 2 + # makes the same call site work end-to-end. + +Phase 1 (this module's current state) ships everything except the +actual per-block compress + stitch loop in +:meth:`Compressor.apply`. Phase 2 adds that loop and populates +:attr:`CompressionReport.blocks` with one +:class:`polygram.compression.BlockReport` per partition block. """ from __future__ import annotations