collapse.py

"""Learned collapsing of adjacent tokens into single units.

The earlier n-gram experiment *added* bigram rows next to the unigrams, so a
Georgian word cut into six pieces became six unigrams plus five bigrams — eleven
items instead of six. That raised fragmentation instead of lowering it, which is
why well-tokenised languages lost ground: an English sentence of four tokens
gained three bigram rows and half its pooled mass went to them.

Collapsing is a soft *replacement*: a learned merged row enters the numerator
while the neighbouring unigram rows lose mass. That changes the direction toward
units the tokenizer should have produced and targets tokens-per-word directly —
the quantity whose rank correlation with held-out Recall@1 is -0.607.

For adjacent pair i, a merge probability p_i gates a learned merged row M_i:

    numerator   = sum_i (1 - (p_{i-1} + p_i)/2) * E[t_i]  +  sum_i p_i * M_i
    denominator = n - sum_i p_i

At p = 0 this is exactly the plain mean, so the control is recovered by
construction. Increasing p rotates the pooled direction from neighbouring
unigrams toward the learned merged row; overlapping candidates share the
reduction instead of needing a discrete matching. Under cosine scoring the
positive scalar denominator cancels, so the mechanism is the relative numerator
composition, not the smaller denominator.

The merge score carries a per-language bias, so a language whose words are
already single tokens can shut merging off globally rather than paying for a
mechanism it does not need — the failure mode of the ungated n-gram arm.
"""

from __future__ import annotations

from collections import defaultdict
from dataclasses import dataclass
import heapq
import math
import statistics
from typing import Any

import torch
from torch import nn

_MIX_A = 2654435761
_MIX_B = 40503
_SKETCH_A = (6364136223846793005, 3202034522624059733, 3935559000370003845)
_SKETCH_B = (1442695040888963407, 2691343689449507681, 4768777513237032717)
_SKETCH_SIGN_A = (2862933555777941757, 7046029254386353131, 3935559000370003845)
_MAX_WEIGHTED_SOURCE_BYTES = 256 * 1024 * 1024


def _weighted_index_add_(
    target: torch.Tensor,
    index: torch.Tensor,
    source: torch.Tensor,
    weight: torch.Tensor,
    *,
    max_chunk_bytes: int = _MAX_WEIGHTED_SOURCE_BYTES,
) -> torch.Tensor:
    """Index-add weighted rows without materialising an unbounded N x dim product."""
    if source.ndim != 2 or weight.ndim != 1 or source.shape[0] != weight.shape[0]:
        raise ValueError("weighted index-add needs aligned matrix rows and scalar weights")
    if index.ndim != 1 or index.shape[0] != source.shape[0]:
        raise ValueError("weighted index-add needs one destination per source row")
    if max_chunk_bytes <= 0:
        raise ValueError("weighted index-add chunk budget must be positive")
    bytes_per_row = max(1, source.shape[1] * source.element_size())
    rows_per_chunk = max(1, max_chunk_bytes // bytes_per_row)
    for start in range(0, source.shape[0], rows_per_chunk):
        stop = min(source.shape[0], start + rows_per_chunk)
        target.index_add_(
            0,
            index[start:stop],
            source[start:stop] * weight[start:stop].unsqueeze(1),
        )
    return target


@dataclass(frozen=True)
class PromotedVocabulary:
    """Deterministic result of bounded exact-pair promotion."""

    pair_keys: tuple[int, ...]
    records: tuple[dict[str, Any], ...]
    candidate_count: int
    reservation_memberships: tuple[tuple[int, int], ...]


class ContrastiveVocabularyDiscovery:
    """Bounded candidate table with signed CountSketch admission.

    The sketch receives ``-p * dL/dp`` after each backward pass. It transfers
    only the largest absolute pair/language aggregates from a forward. A fixed
    signed CountSketch retains cumulative admission evidence even for keys not
    currently in the exact table; the bounded exact table stores utility,
    probability and support only from admission onward.
    """

    admission_mode = "utility"

    def __init__(
        self,
        *,
        n_languages: int,
        candidate_capacity: int,
        top_per_forward: int,
    ) -> None:
        if n_languages <= 0:
            raise ValueError("n_languages must be positive")
        if candidate_capacity <= 0:
            raise ValueError("candidate_capacity must be positive")
        if top_per_forward <= 0:
            raise ValueError("top_per_forward must be positive")
        self.n_languages = int(n_languages)
        self.candidate_capacity = int(candidate_capacity)
        self.top_per_forward = int(top_per_forward)
        # composite pair/language key -> [signed utility, p sum, captured count]
        self.records: dict[int, list[float | int]] = {}
        self._sketch_width = max(64, self.candidate_capacity)
        self._utility_sketch = torch.zeros(
            (len(_SKETCH_A), self._sketch_width), dtype=torch.float64
        )
        # Lazy min-heap over retained rows.  Updates append a new rank and leave
        # the previous entry stale; periodic rebuilding gives the metadata a
        # fixed bound without sorting the entire sketch on every forward.
        self._heap: list[tuple[float, int, int, int]] = []
        self.observed_occurrences = 0
        self.transferred_records = 0
        self.unique_records_before_transfer = 0
        self.dropped_by_top_per_forward = 0
        self.retained_updates = 0
        self.admitted_into_free_slots = 0
        self.capacity_replacements = 0
        self.capacity_rejections = 0

    @staticmethod
    def _rank(key: int, row: list[float | int]) -> tuple[float, int, int]:
        return float(row[0]), int(row[2]), -int(key)

    def _push_heap(self, key: int) -> None:
        heapq.heappush(self._heap, (*self._rank(key, self.records[key]), int(key)))

    def _discard_stale_heap_entries(self) -> None:
        while self._heap:
            utility, count, negative_key, key = self._heap[0]
            row = self.records.get(key)
            if row is not None and (utility, count, negative_key) == self._rank(key, row):
                return
            heapq.heappop(self._heap)

    def _rebuild_heap(self) -> None:
        self._heap = [(*self._rank(key, row), key) for key, row in self.records.items()]
        heapq.heapify(self._heap)

    def _update_utility_sketch(
        self, keys: torch.Tensor, values: torch.Tensor
    ) -> torch.Tensor:
        """Update fixed memory and return the median signed estimate per key."""

        keys = keys.to(device="cpu", dtype=torch.long)
        values = values.to(device="cpu", dtype=torch.float64)
        estimates = []
        for depth, (multiplier, offset, sign_multiplier) in enumerate(zip(
            _SKETCH_A, _SKETCH_B, _SKETCH_SIGN_A, strict=True
        )):
            bucket = torch.remainder(keys * multiplier + offset, self._sketch_width)
            sign = torch.where(
                torch.bitwise_and(keys * sign_multiplier + offset, 1) == 0,
                1.0,
                -1.0,
            ).to(torch.float64)
            self._utility_sketch[depth].index_add_(0, bucket, values * sign)
            estimates.append(self._utility_sketch[depth, bucket] * sign)
        return torch.stack(estimates).median(dim=0).values

    @torch.no_grad()
    def observe(
        self,
        pair_key: torch.Tensor,
        language: torch.Tensor,
        probability: torch.Tensor,
        gradient: torch.Tensor,
    ) -> None:
        if not pair_key.numel():
            return
        utility = -probability.detach().float() * gradient.detach().float()
        finite = torch.isfinite(utility)
        if not bool(finite.any()):
            return
        pair_key = pair_key.detach()[finite].to(torch.long)
        language = language.detach()[finite].to(torch.long)
        probability = probability.detach()[finite].float()
        utility = utility[finite]
        self.observed_occurrences += int(pair_key.numel())

        composite = pair_key * self.n_languages + language
        unique, inverse = torch.unique(composite, return_inverse=True)
        utility_sum = torch.zeros(unique.numel(), device=utility.device)
        probability_sum = torch.zeros(unique.numel(), device=utility.device)
        count = torch.zeros(unique.numel(), dtype=torch.long, device=utility.device)
        utility_sum.index_add_(0, inverse, utility)
        probability_sum.index_add_(0, inverse, probability)
        count.index_add_(0, inverse, torch.ones_like(inverse))

        self.unique_records_before_transfer += int(unique.numel())
        keep = min(self.top_per_forward, unique.numel())
        if keep < unique.numel():
            self.dropped_by_top_per_forward += int(unique.numel() - keep)
            chosen = torch.topk(utility_sum.abs(), keep, sorted=False).indices
            unique = unique[chosen]
            utility_sum = utility_sum[chosen]
            probability_sum = probability_sum[chosen]
            count = count[chosen]

        unique_cpu = unique.cpu()
        utility_cpu = utility_sum.cpu()
        admission_estimates = self._update_utility_sketch(unique_cpu, utility_cpu)
        keys = unique_cpu.tolist()
        utilities = utility_cpu.tolist()
        probabilities = probability_sum.cpu().tolist()
        counts = count.cpu().tolist()
        self.transferred_records += len(keys)
        incoming = {
            int(key): [float(value), float(p_sum), int(n)]
            for key, value, p_sum, n in zip(
                keys, utilities, probabilities, counts, strict=True
            )
        }
        admission = {
            int(key): float(estimate)
            for key, estimate in zip(keys, admission_estimates.tolist(), strict=True)
        }

        # First update retained rows, then consider unseen rows in descending
        # batch rank.  This makes replacement deterministic and independent of
        # the device order returned by top-k.
        retained_keys = sorted(key for key in incoming if key in self.records)
        self.retained_updates += len(retained_keys)
        for key in retained_keys:
            value, p_sum, n = incoming.pop(key)
            row = self.records[key]
            row[0] = float(row[0]) + float(value)
            row[1] = float(row[1]) + float(p_sum)
            row[2] = int(row[2]) + int(n)
            self._push_heap(key)

        unseen = sorted(
            incoming.items(),
            key=lambda item: (admission[item[0]], int(item[1][2]), -item[0]),
            reverse=True,
        )
        for key, row in unseen:
            if len(self.records) < self.candidate_capacity:
                self.records[key] = row
                self._push_heap(key)
                self.admitted_into_free_slots += 1
                continue
            self._discard_stale_heap_entries()
            if not self._heap:
                self._rebuild_heap()
            admission_rank = (admission[key], int(row[2]), -key)
            if admission_rank <= self._heap[0][:3]:
                self.capacity_rejections += 1
                continue
            victim = heapq.heappop(self._heap)[3]
            del self.records[victim]
            self.records[key] = row
            self._push_heap(key)
            self.capacity_replacements += 1

        # The retained dictionary never exceeds candidate_capacity.  Lazy heap
        # metadata is also bounded: after a call it has at most 2x as many
        # entries as the registered candidate table.
        if len(self._heap) > 2 * self.candidate_capacity:
            self._rebuild_heap()

    def _prune(self) -> None:
        ranked = sorted(
            self.records.items(),
            key=lambda item: self._rank(item[0], item[1]),
            reverse=True,
        )[: self.candidate_capacity]
        self.records = dict(ranked)
        self._rebuild_heap()

    def snapshot_records(self) -> tuple[dict[str, float | int], ...]:
        """Return the complete bounded pair/language table deterministically."""

        self._prune()
        rows = []
        for composite, row in sorted(self.records.items()):
            pair_key, language = divmod(int(composite), self.n_languages)
            utility, probability_sum, count = float(row[0]), float(row[1]), int(row[2])
            rows.append({
                "composite_key": int(composite),
                "pair_key": pair_key,
                "language_index": language,
                "utility": utility,
                "probability_sum": probability_sum,
                "captured_support": count,
                "mean_probability": probability_sum / count if count else 0.0,
            })
        return tuple(rows)

    def state_dict(self) -> dict[str, Any]:
        """Return the complete bounded discovery state for an exact resume."""

        self._prune()
        return {
            "version": 2,
            "n_languages": self.n_languages,
            "candidate_capacity": self.candidate_capacity,
            "top_per_forward": self.top_per_forward,
            "sketch_width": self._sketch_width,
            "records": {
                int(key): [float(row[0]), float(row[1]), int(row[2])]
                for key, row in sorted(self.records.items())
            },
            "utility_sketch": self._utility_sketch.clone(),
            "observed_occurrences": self.observed_occurrences,
            "transferred_records": self.transferred_records,
            "admission_pressure": self.admission_pressure(),
        }

    def load_state_dict(self, raw: dict[str, Any]) -> None:
        """Restore a state produced by :meth:`state_dict` after validation."""

        version = int(raw.get("version", -1))
        if version not in (1, 2):
            raise ValueError("unsupported discovery state version")
        if int(raw.get("n_languages", -1)) != self.n_languages:
            raise ValueError("discovery state language count does not match")
        if int(raw.get("candidate_capacity", -1)) != self.candidate_capacity:
            raise ValueError("discovery state candidate capacity does not match")
        if int(raw.get("top_per_forward", -1)) != self.top_per_forward:
            raise ValueError("discovery state top-per-forward does not match")
        if int(raw.get("sketch_width", -1)) != self._sketch_width:
            raise ValueError("discovery state sketch width does not match")
        sketch = raw.get("utility_sketch")
        if (
            not torch.is_tensor(sketch)
            or sketch.dtype != torch.float64
            or tuple(sketch.shape) != tuple(self._utility_sketch.shape)
            or not bool(torch.isfinite(sketch).all())
        ):
            raise ValueError("discovery state utility sketch is invalid")
        raw_records = raw.get("records")
        if not isinstance(raw_records, dict) or len(raw_records) > self.candidate_capacity:
            raise ValueError("discovery state records are invalid")
        records: dict[int, list[float | int]] = {}
        for raw_key, raw_row in raw_records.items():
            key = int(raw_key)
            if key < 0 or not isinstance(raw_row, (list, tuple)) or len(raw_row) != 3:
                raise ValueError("discovery state record is invalid")
            utility, probability, count = float(raw_row[0]), float(raw_row[1]), int(raw_row[2])
            if not (math.isfinite(utility) and math.isfinite(probability)) or count < 0:
                raise ValueError("discovery state record values are invalid")
            records[key] = [utility, probability, count]
        observed = int(raw.get("observed_occurrences", -1))
        transferred = int(raw.get("transferred_records", -1))
        if observed < 0 or transferred < 0:
            raise ValueError("discovery state counters are invalid")
        counter_names = tuple(self.admission_pressure())
        raw_pressure = raw.get("admission_pressure", {})
        if version == 2 and (
            not isinstance(raw_pressure, dict)
            or set(raw_pressure) != set(counter_names)
        ):
            raise ValueError("discovery state admission counters are invalid")
        pressure = {
            name: int(raw_pressure.get(name, 0)) for name in counter_names
        }
        if any(value < 0 for value in pressure.values()):
            raise ValueError("discovery state admission counters are invalid")

        self.records = records
        self._utility_sketch.copy_(sketch)
        self.observed_occurrences = observed
        self.transferred_records = transferred
        for name, value in pressure.items():
            setattr(self, name, value)
        self._rebuild_heap()

    def admission_pressure(self) -> dict[str, int]:
        return {
            "unique_records_before_transfer": self.unique_records_before_transfer,
            "dropped_by_top_per_forward": self.dropped_by_top_per_forward,
            "retained_updates": self.retained_updates,
            "admitted_into_free_slots": self.admitted_into_free_slots,
            "capacity_replacements": self.capacity_replacements,
            "capacity_rejections": self.capacity_rejections,
        }


    def promote(
        self,
        *,
        budget: int,
        min_support: int,
        per_language_quota: int,
        per_language_min_support: int,
        min_utility: float = 0.0,
    ) -> PromotedVocabulary:
        """Select utility-thresholded exact rows under a safety L0 ceiling."""

        if budget <= 0:
            raise ValueError("budget must be positive")
        if min_support <= 0 or per_language_min_support <= 0:
            raise ValueError("support thresholds must be positive")
        if per_language_quota < 0:
            raise ValueError("per_language_quota must be non-negative")
        if not math.isfinite(min_utility) or min_utility < 0:
            raise ValueError("min_utility must be finite and non-negative")
        self._prune()

        by_pair: dict[int, dict[str, Any]] = defaultdict(
            lambda: {"utility": 0.0, "p_sum": 0.0, "count": 0, "languages": {}}
        )
        for composite, row in self.records.items():
            pair_key, language = divmod(composite, self.n_languages)
            utility, p_sum, count = float(row[0]), float(row[1]), int(row[2])
            target = by_pair[pair_key]
            target["utility"] += utility
            target["p_sum"] += p_sum
            target["count"] += count
            target["languages"][language] = {
                "utility": utility,
                "p_sum": p_sum,
                "count": count,
            }

        selected: set[int] = set()
        reservation_languages: dict[int, set[int]] = defaultdict(set)
        if per_language_quota:
            candidates_by_language: list[list[tuple[float, int, int, int]]] = [
                [] for _ in range(self.n_languages)
            ]
            for pair_key, row in by_pair.items():
                # A language reservation protects multilingual coverage, but
                # it cannot override the registered global U(q) > 0
                # eligibility constraint.  Build the per-language lists in
                # one pass over observed records; scanning every pair once per
                # marker is prohibitive for the 1,914-language run.
                if row["utility"] <= 0 or row["utility"] < min_utility:
                    continue
                for language, language_row in row["languages"].items():
                    if (language_row["utility"] > 0
                            and language_row["count"] >= per_language_min_support):
                        candidates_by_language[language].append((
                            float(language_row["utility"]),
                            int(language_row["count"]),
                            -int(pair_key),
                            int(pair_key),
                        ))
            for language, candidates in enumerate(candidates_by_language):
                candidates.sort(reverse=True)
                for candidate in candidates[:per_language_quota]:
                    pair_key = candidate[3]
                    selected.add(pair_key)
                    reservation_languages[pair_key].add(language)

        if len(selected) > budget:
            selected = set(sorted(
                selected,
                key=lambda key: (
                    float(by_pair[key]["utility"]),
                    int(by_pair[key]["count"]),
                    -key,
                ),
                reverse=True,
            )[:budget])
            reservation_languages = {
                key: languages for key, languages in reservation_languages.items()
                if key in selected
            }

        global_candidates = [
            key for key, row in by_pair.items()
            if (
                row["utility"] > 0
                and row["utility"] >= min_utility
                and row["count"] >= min_support
                and key not in selected
            )
        ]
        global_candidates.sort(
            key=lambda key: (
                float(by_pair[key]["utility"]),
                int(by_pair[key]["count"]),
                -key,
            ),
            reverse=True,
        )
        selected.update(global_candidates[: max(0, budget - len(selected))])

        pair_keys = tuple(sorted(selected))
        rows = []
        for pair_key in pair_keys:
            row = by_pair[pair_key]
            reserved = sorted(reservation_languages.get(pair_key, ()))
            rows.append({
                "pair_key": pair_key,
                "utility": float(row["utility"]),
                "captured_support": int(row["count"]),
                "mean_probability": (
                    float(row["p_sum"]) / int(row["count"]) if row["count"] else 0.0
                ),
                "language_count": len(row["languages"]),
                "reservation_language_count": len(reserved),
                "max_reservation_support": max(
                    (int(row["languages"][language]["count"]) for language in reserved),
                    default=0,
                ),
                "global_support_eligible": int(row["count"]) >= min_support,
            })
        return PromotedVocabulary(
            pair_keys=pair_keys,
            records=tuple(rows),
            candidate_count=len(by_pair),
            reservation_memberships=tuple(sorted(
                (pair_key, language)
                for pair_key, languages in reservation_languages.items()
                for language in languages
            )),
        )

    def audit_promotion(
        self,
        promoted: PromotedVocabulary,
        *,
        budget: int,
        min_support: int,
        per_language_quota: int,
        per_language_min_support: int,
        min_utility: float = 0.0,
        reservation_audit_quota: int | None = None,
        trainable_language_indices: tuple[int, ...] | None = None,
    ) -> dict[str, Any]:
        """Compare a quota selection with the same discovery under quota zero."""

        if reservation_audit_quota is None:
            reservation_audit_quota = per_language_quota
        if reservation_audit_quota < 0:
            raise ValueError("reservation_audit_quota must be non-negative")
        if trainable_language_indices is None:
            trainable_language_indices = tuple(range(self.n_languages))
        else:
            trainable_language_indices = tuple(
                int(index) for index in trainable_language_indices
            )
        if (
            not trainable_language_indices
            or len(set(trainable_language_indices)) != len(trainable_language_indices)
            or any(
                index < 0 or index >= self.n_languages
                for index in trainable_language_indices
            )
        ):
            raise ValueError("trainable language indices are invalid")
        trainable_set = set(trainable_language_indices)
        global_only = (
            promoted
            if per_language_quota == 0
            else self.promote(
                budget=budget,
                min_support=min_support,
                per_language_quota=0,
                per_language_min_support=per_language_min_support,
                min_utility=min_utility,
            )
        )
        reservation_selection = (
            promoted
            if per_language_quota == reservation_audit_quota
            else self.promote(
                budget=budget,
                min_support=min_support,
                per_language_quota=reservation_audit_quota,
                per_language_min_support=per_language_min_support,
                min_utility=min_utility,
            )
        )
        all_candidate_counts = [0] * self.n_languages
        aggregate_utility: dict[int, float] = defaultdict(float)
        for composite, row in self.records.items():
            pair_key, language = divmod(int(composite), self.n_languages)
            if language not in trainable_set:
                raise ValueError(
                    "discovery contains a nontrainable language candidate"
                )
            all_candidate_counts[language] += 1
            aggregate_utility[pair_key] += float(row[0])
        all_eligible_counts = [0] * self.n_languages
        for composite, row in self.records.items():
            pair_key, language = divmod(int(composite), self.n_languages)
            if (
                aggregate_utility[pair_key] > 0
                and aggregate_utility[pair_key] >= min_utility
                and float(row[0]) > 0
                and int(row[2]) >= per_language_min_support
            ):
                all_eligible_counts[language] += 1
        candidate_counts = [
            all_candidate_counts[index] for index in trainable_language_indices
        ]
        eligible_counts = [
            all_eligible_counts[index] for index in trainable_language_indices
        ]
        all_reservation_counts = [0] * self.n_languages
        for _, language in reservation_selection.reservation_memberships:
            if language not in trainable_set:
                raise ValueError("reservation selected a nontrainable language")
            all_reservation_counts[language] += 1
        reservation_counts = [
            all_reservation_counts[index] for index in trainable_language_indices
        ]

        reservation_utility = math.fsum(
            float(record["utility"]) for record in reservation_selection.records
        )
        global_utility = math.fsum(
            float(record["utility"]) for record in global_only.records
        )
        utility_displaced = global_utility - reservation_utility
        membership_capacity = reservation_audit_quota * len(trainable_language_indices)
        reservation_keys = set(reservation_selection.pair_keys)
        global_keys = set(global_only.pair_keys)
        return {
            "deployed_selector": {
                "per_language_quota": per_language_quota,
                "matches_global_only": promoted.pair_keys == global_only.pair_keys,
                "reservation_audit_quota": reservation_audit_quota,
            },
            "candidate_table": {
                "capacity": self.candidate_capacity,
                "pair_language_records": len(self.records),
                "saturation": len(self.records) / self.candidate_capacity,
                "candidate_pairs": promoted.candidate_count,
                "initializer_languages_total": self.n_languages,
                "languages_total": len(trainable_language_indices),
                "trainable_language_indices": list(trainable_language_indices),
                "languages_with_candidates": sum(
                    count > 0 for count in candidate_counts
                ),
                "records_per_language": candidate_counts,
                "reservation_eligible_records_per_language": eligible_counts,
                "min_records_per_language": min(candidate_counts),
                "median_records_per_language": statistics.median(candidate_counts),
                "max_records_per_language": max(candidate_counts),
                "admission_pressure": {
                    **{
                        "unique_records_before_transfer": (
                            self.unique_records_before_transfer
                        ),
                        "dropped_by_top_per_forward": (
                            self.dropped_by_top_per_forward
                        ),
                        "transferred_records": self.transferred_records,
                    },
                    "retained_updates": self.retained_updates,
                    "admitted_into_free_slots": self.admitted_into_free_slots,
                    "capacity_replacements": self.capacity_replacements,
                    "capacity_rejections": self.capacity_rejections,
                },
            },
            "reservation": {
                "quota_per_language": reservation_audit_quota,
                "membership_capacity": membership_capacity,
                "memberships": len(reservation_selection.reservation_memberships),
                "membership_occupancy": (
                    len(reservation_selection.reservation_memberships)
                    / membership_capacity
                    if membership_capacity
                    else 0.0
                ),
                "unique_reserved_rows": len({
                    pair_key
                    for pair_key, _ in reservation_selection.reservation_memberships
                }),
                "languages_with_reservations": sum(
                    count > 0 for count in reservation_counts
                ),
                "memberships_per_language": reservation_counts,
            },
            "global_only_counterfactual": {
                "global_selected_rows": len(global_only.pair_keys),
                "global_selected_utility": global_utility,
                "reservation_selected_rows": len(reservation_selection.pair_keys),
                "reservation_selected_utility": reservation_utility,
                "utility_displaced_by_reservations": utility_displaced,
                "relative_utility_displaced_by_reservations": (
                    utility_displaced / global_utility if global_utility else None
                ),
                "reservation_only_rows": len(reservation_keys - global_keys),
                "global_only_rows": len(global_keys - reservation_keys),
            },
        }

    def select_external(
        self,
        pair_keys: tuple[int, ...] | list[int],
        *,
        minimum_support: int,
    ) -> PromotedVocabulary:
        """Validate and materialize a hash-pinned external promotion decision."""

        if minimum_support <= 0:
            raise ValueError("external-selection minimum support must be positive")
        keys = tuple(int(key) for key in pair_keys)
        if not keys or tuple(sorted(set(keys))) != keys:
            raise ValueError("external pair keys must be non-empty, sorted and unique")
        self._prune()
        by_pair: dict[int, dict[str, Any]] = defaultdict(
            lambda: {"utility": 0.0, "p_sum": 0.0, "count": 0, "languages": {}}
        )
        for composite, raw in self.records.items():
            pair_key, language = divmod(composite, self.n_languages)
            utility, p_sum, count = float(raw[0]), float(raw[1]), int(raw[2])
            row = by_pair[pair_key]
            row["utility"] += utility
            row["p_sum"] += p_sum
            row["count"] += count
            row["languages"][language] = {
                "utility": utility, "p_sum": p_sum, "count": count,
            }
        missing = [key for key in keys if key not in by_pair]
        if missing:
            raise ValueError(f"external pair keys absent from discovery: {missing[:8]}")
        ineligible = [
            key for key in keys
            if by_pair[key]["utility"] <= 0 or by_pair[key]["count"] < minimum_support
        ]
        if ineligible:
            raise ValueError(
                f"external pair keys fail utility/support eligibility: {ineligible[:8]}"
            )
        rows = []
        for pair_key in keys:
            row = by_pair[pair_key]
            rows.append({
                "pair_key": pair_key,
                "utility": float(row["utility"]),
                "captured_support": int(row["count"]),
                "mean_probability": (
                    float(row["p_sum"]) / int(row["count"]) if row["count"] else 0.0
                ),
                "language_count": len(row["languages"]),
                "reservation_language_count": 0,
                "max_reservation_support": 0,
                "global_support_eligible": int(row["count"]) >= minimum_support,
            })
        return PromotedVocabulary(
            pair_keys=keys,
            records=tuple(rows),
            candidate_count=len(by_pair),
            reservation_memberships=(),
        )


class FrequencyVocabularyDiscovery:
    """Independent fixed-memory heavy hitters ranked only by occurrence count.

    Unlike :class:`ContrastiveVocabularyDiscovery`, this observer never reads
    the value or sign of ``-p * dL/dp``.  A Count-Min sketch sees every pair
    occurrence, while a bounded exact key table retains the largest sketch
    estimates.  The table is pair-level rather than pair/language-level so a
    globally frequent pair does not spend multiple candidate slots merely
    because it occurs in several languages.

    The sketch has the same depth and width rule as the utility collector.  Its
    estimates can over-count because of sketch collisions; they cannot
    under-count.  This is an independent fixed-memory frequency selector, not
    an exact unbounded corpus histogram.
    """

    admission_mode = "frequency"

    def __init__(
        self,
        *,
        n_languages: int,
        candidate_capacity: int,
        top_per_forward: int,
    ) -> None:
        if n_languages <= 0:
            raise ValueError("n_languages must be positive")
        if candidate_capacity <= 0:
            raise ValueError("candidate_capacity must be positive")
        if top_per_forward <= 0:
            raise ValueError("top_per_forward must be positive")
        self.n_languages = int(n_languages)
        self.candidate_capacity = int(candidate_capacity)
        self.top_per_forward = int(top_per_forward)
        # pair key -> [unused utility, unused probability sum, support estimate]
        self.records: dict[int, list[float | int]] = {}
        self._sketch_width = max(64, self.candidate_capacity)
        self._frequency_sketch = torch.zeros(
            (len(_SKETCH_A), self._sketch_width), dtype=torch.long
        )
        self._heap: list[tuple[int, int, int]] = []
        self.observed_occurrences = 0
        self.transferred_records = 0
        self.unique_records_before_transfer = 0
        self.dropped_by_top_per_forward = 0
        self.retained_updates = 0
        self.admitted_into_free_slots = 0
        self.capacity_replacements = 0
        self.capacity_rejections = 0

    @staticmethod
    def _rank(key: int, row: list[float | int]) -> tuple[int, int]:
        return int(row[2]), -int(key)

    def _push_heap(self, key: int) -> None:
        heapq.heappush(self._heap, (*self._rank(key, self.records[key]), int(key)))

    def _discard_stale_heap_entries(self) -> None:
        while self._heap:
            support, negative_key, key = self._heap[0]
            row = self.records.get(key)
            if row is not None and (support, negative_key) == self._rank(key, row):
                return
            heapq.heappop(self._heap)

    def _rebuild_heap(self) -> None:
        self._heap = [
            (*self._rank(key, row), key) for key, row in self.records.items()
        ]
        heapq.heapify(self._heap)

    def _update_frequency_sketch(
        self, keys: torch.Tensor, counts: torch.Tensor
    ) -> torch.Tensor:
        """Update Count-Min state for every key and return its current estimate."""

        keys = keys.to(dtype=torch.long)
        counts = counts.to(device=keys.device, dtype=torch.long)
        if self._frequency_sketch.device != keys.device:
            self._frequency_sketch = self._frequency_sketch.to(keys.device)
        estimates = []
        for depth, (multiplier, offset) in enumerate(
            zip(_SKETCH_A, _SKETCH_B, strict=True)
        ):
            bucket = torch.remainder(
                keys * multiplier + offset, self._sketch_width
            )
            self._frequency_sketch[depth].index_add_(0, bucket, counts)
            estimates.append(self._frequency_sketch[depth, bucket])
        return torch.stack(estimates).amin(dim=0)

    @staticmethod
    def _largest_with_deterministic_boundary(
        keys: torch.Tensor, estimates: torch.Tensor, keep: int
    ) -> torch.Tensor:
        """Keep the largest estimates, resolving the cutoff by smaller key."""

        if keep >= keys.numel():
            return torch.arange(keys.numel(), device=keys.device)
        boundary = torch.topk(estimates, keep, sorted=False).values.min()
        above = torch.nonzero(estimates > boundary, as_tuple=False).squeeze(1)
        ties = torch.nonzero(estimates == boundary, as_tuple=False).squeeze(1)
        remaining = keep - int(above.numel())
        if remaining <= 0:
            return above[:keep]
        tie_choice = torch.topk(
            -keys.index_select(0, ties), remaining, sorted=False
        ).indices
        return torch.cat((above, ties.index_select(0, tie_choice)))

    @torch.no_grad()
    def observe(
        self,
        pair_key: torch.Tensor,
        language: torch.Tensor,
        probability: torch.Tensor,
        gradient: torch.Tensor,
    ) -> None:
        del language, probability, gradient
        if not pair_key.numel():
            return
        pair_key = pair_key.detach().to(torch.long)
        self.observed_occurrences += int(pair_key.numel())
        unique, counts = torch.unique(pair_key, return_counts=True)
        estimates = self._update_frequency_sketch(unique, counts)

        self.unique_records_before_transfer += int(unique.numel())
        keep = min(self.top_per_forward, unique.numel())
        if keep < unique.numel():
            self.dropped_by_top_per_forward += int(unique.numel() - keep)
            chosen = self._largest_with_deterministic_boundary(
                unique, estimates, keep
            )
            unique = unique.index_select(0, chosen)
            estimates = estimates.index_select(0, chosen)

        keys = unique.cpu().tolist()
        supports = estimates.cpu().tolist()
        self.transferred_records += len(keys)
        incoming = {int(key): int(support) for key, support in zip(keys, supports)}

        retained_keys = sorted(key for key in incoming if key in self.records)
        self.retained_updates += len(retained_keys)
        for key in retained_keys:
            support = incoming.pop(key)
            self.records[key][2] = max(int(self.records[key][2]), support)
            self._push_heap(key)

        unseen = sorted(
            incoming.items(), key=lambda item: (item[1], -item[0]), reverse=True
        )
        for key, support in unseen:
            row: list[float | int] = [0.0, 0.0, int(support)]
            if len(self.records) < self.candidate_capacity:
                self.records[key] = row
                self._push_heap(key)
                self.admitted_into_free_slots += 1
                continue
            self._discard_stale_heap_entries()
            if not self._heap:
                self._rebuild_heap()
            admission_rank = (int(support), -int(key))
            if admission_rank <= self._heap[0][:2]:
                self.capacity_rejections += 1
                continue
            victim = heapq.heappop(self._heap)[2]
            del self.records[victim]
            self.records[key] = row
            self._push_heap(key)
            self.capacity_replacements += 1

        if len(self._heap) > 2 * self.candidate_capacity:
            self._rebuild_heap()

    def _prune(self) -> None:
        ranked = sorted(
            self.records.items(),
            key=lambda item: self._rank(item[0], item[1]),
            reverse=True,
        )[: self.candidate_capacity]
        self.records = dict(ranked)
        self._rebuild_heap()

    def snapshot_records(self) -> tuple[dict[str, float | int], ...]:
        """Return one deterministic row per retained global pair key."""

        self._prune()
        return tuple(
            {
                "composite_key": int(pair_key),
                "pair_key": int(pair_key),
                "language_index": -1,
                "utility": 0.0,
                "probability_sum": 0.0,
                "captured_support": int(row[2]),
                "mean_probability": 0.0,
            }
            for pair_key, row in sorted(self.records.items())
        )

    def state_dict(self) -> dict[str, Any]:
        """Return the complete fixed-memory frequency state for exact resume."""

        self._prune()
        return {
            "version": 1,
            "protocol": "independent-frequency-count-min-v1",
            "n_languages": self.n_languages,
            "candidate_capacity": self.candidate_capacity,
            "top_per_forward": self.top_per_forward,
            "sketch_width": self._sketch_width,
            "records": {
                int(key): [0.0, 0.0, int(row[2])]
                for key, row in sorted(self.records.items())
            },
            "frequency_sketch": self._frequency_sketch.detach().cpu().clone(),
            "observed_occurrences": self.observed_occurrences,
            "transferred_records": self.transferred_records,
            "admission_pressure": self.admission_pressure(),
        }

    def load_state_dict(self, raw: dict[str, Any]) -> None:
        """Restore a state produced by :meth:`state_dict` after validation."""

        if (
            int(raw.get("version", -1)) != 1
            or raw.get("protocol") != "independent-frequency-count-min-v1"
        ):
            raise ValueError("unsupported frequency discovery state")
        if int(raw.get("n_languages", -1)) != self.n_languages:
            raise ValueError("frequency discovery language count does not match")
        if int(raw.get("candidate_capacity", -1)) != self.candidate_capacity:
            raise ValueError("frequency discovery candidate capacity does not match")
        if int(raw.get("top_per_forward", -1)) != self.top_per_forward:
            raise ValueError("frequency discovery top-per-forward does not match")
        if int(raw.get("sketch_width", -1)) != self._sketch_width:
            raise ValueError("frequency discovery sketch width does not match")
        sketch = raw.get("frequency_sketch")
        if (
            not torch.is_tensor(sketch)
            or sketch.dtype != torch.long
            or tuple(sketch.shape) != tuple(self._frequency_sketch.shape)
            or bool((sketch < 0).any())
        ):
            raise ValueError("frequency discovery sketch is invalid")
        raw_records = raw.get("records")
        if not isinstance(raw_records, dict) or len(raw_records) > self.candidate_capacity:
            raise ValueError("frequency discovery records are invalid")
        records: dict[int, list[float | int]] = {}
        for raw_key, raw_row in raw_records.items():
            key = int(raw_key)
            if key < 0 or not isinstance(raw_row, (list, tuple)) or len(raw_row) != 3:
                raise ValueError("frequency discovery record is invalid")
            support = int(raw_row[2])
            if support < 0:
                raise ValueError("frequency discovery support is invalid")
            records[key] = [0.0, 0.0, support]
        observed = int(raw.get("observed_occurrences", -1))
        transferred = int(raw.get("transferred_records", -1))
        if observed < 0 or transferred < 0:
            raise ValueError("frequency discovery counters are invalid")
        counter_names = tuple(self.admission_pressure())
        pressure = raw.get("admission_pressure")
        if not isinstance(pressure, dict) or set(pressure) != set(counter_names):
            raise ValueError("frequency discovery admission counters are invalid")
        pressure = {name: int(pressure[name]) for name in counter_names}
        if any(value < 0 for value in pressure.values()):
            raise ValueError("frequency discovery admission counters are invalid")

        self.records = records
        self._frequency_sketch = sketch.detach().cpu().clone()
        self.observed_occurrences = observed
        self.transferred_records = transferred
        for name, value in pressure.items():
            setattr(self, name, value)
        self._rebuild_heap()

    def admission_pressure(self) -> dict[str, int]:
        return {
            "unique_records_before_transfer": self.unique_records_before_transfer,
            "dropped_by_top_per_forward": self.dropped_by_top_per_forward,
            "retained_updates": self.retained_updates,
            "admitted_into_free_slots": self.admitted_into_free_slots,
            "capacity_replacements": self.capacity_replacements,
            "capacity_rejections": self.capacity_rejections,
        }



class CollapseTable(nn.Module):
    def __init__(
        self,
        buckets: int,
        dim: int,
        n_languages: int,
        bias_init: float = -2.0,
        *,
        pair_key_base: int | None = None,
        pair_keys: list[int] | tuple[int, ...] | torch.Tensor | None = None,
        pair_slots: list[int] | tuple[int, ...] | torch.Tensor | None = None,
        reclaimed_token_ids: list[int] | tuple[int, ...] | torch.Tensor | None = None,
        residual_buckets: int = 0,
    ) -> None:
        super().__init__()
        self.buckets = int(buckets)
        self.residual_buckets = int(residual_buckets)
        if self.residual_buckets < 0:
            raise ValueError("residual_buckets must be non-negative")
        self.pair_key_base = int(pair_key_base or 0)
        reclaimed = torch.as_tensor(
            reclaimed_token_ids if reclaimed_token_ids is not None else [], dtype=torch.long
        )
        self.merged = nn.Parameter(torch.zeros(
            0 if reclaimed.numel() else self.buckets, dim
        ))
        self.score = nn.Parameter(torch.zeros(self.buckets))
        self.residual_merged = nn.Parameter(torch.zeros(self.residual_buckets, dim))
        self.residual_score = nn.Parameter(torch.zeros(self.residual_buckets))
        self.language_bias = nn.Parameter(torch.full((n_languages,), float(bias_init)))
        keys = torch.as_tensor(pair_keys if pair_keys is not None else [], dtype=torch.long)
        slots = torch.as_tensor(pair_slots if pair_slots is not None else [], dtype=torch.long)
        if keys.numel() != slots.numel():
            raise ValueError("pair_keys and pair_slots must have the same length")
        if keys.numel() and not bool((keys[1:] > keys[:-1]).all()):
            raise ValueError("pair_keys must be strictly increasing")
        if slots.numel() and (int(slots.min()) < 0 or int(slots.max()) >= self.buckets):
            raise ValueError("pair_slots must index the collapse table")
        if reclaimed.numel() and reclaimed.numel() != keys.numel():
            raise ValueError("reclaimed_token_ids must align one-to-one with pair_keys")
        if reclaimed.unique().numel() != reclaimed.numel():
            raise ValueError("reclaimed_token_ids must be unique")
        self.register_buffer("pair_keys", keys, persistent=True)
        self.register_buffer("pair_slots", slots, persistent=True)
        self.register_buffer("reclaimed_token_ids", reclaimed, persistent=True)
        self.discovery: (
            ContrastiveVocabularyDiscovery | FrequencyVocabularyDiscovery | None
        ) = None
        self.diagnostic_discovery: (
            ContrastiveVocabularyDiscovery | FrequencyVocabularyDiscovery | None
        ) = None
        self._vocabulary_telemetry: Any | None = None
        self._loss_path_scale: float | None = None
        self._loss_path_observer: Any | None = None

    def extra_repr(self) -> str:
        return (f"buckets={self.buckets}, explicit_pairs={self.pair_keys.numel()}, "
                f"residual_buckets={self.residual_buckets}, "
                f"reclaimed_rows={self.reclaimed_token_ids.numel()}")

    def set_vocabulary_telemetry(self, window: Any | None) -> None:
        """Attach detached run telemetry without making it module state."""

        self._vocabulary_telemetry = window

    def begin_loss_path_pass(self, scale: float, observer: Any) -> None:
        """Scale every soft pair gate and capture occurrence gradients.

        This is intentionally transient (not module state): an audit pass may
        read the graph but must never alter a checkpoint or discovery table.
        """

        if self._loss_path_scale is not None or self._loss_path_observer is not None:
            raise RuntimeError("a loss-path pass is already active")
        if float(scale) not in (0.0, 0.5, 1.0):
            raise ValueError("registered loss-path scale must be 0, 0.5, or 1")
        if observer is None or not callable(observer):
            raise TypeError("loss-path observer must be callable")
        self._loss_path_scale = float(scale)
        self._loss_path_observer = observer

    def end_loss_path_pass(self) -> None:
        if self._loss_path_scale is None or self._loss_path_observer is None:
            raise RuntimeError("no loss-path pass is active")
        self._loss_path_scale = None
        self._loss_path_observer = None

    def abort_loss_path_pass(self) -> None:
        self._loss_path_scale = None
        self._loss_path_observer = None

    def _hash(self, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
        return ((a * _MIX_A + b * _MIX_B).abs()) % self.buckets

    def _hash_residual(self, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
        if self.residual_buckets <= 0:
            raise ValueError("residual hash requires a positive residual capacity")
        return ((a * _MIX_A + b * _MIX_B).abs()) % self.residual_buckets

    def _exact_key(self, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
        if self.pair_key_base <= 0:
            raise ValueError("pair_key_base is required for exact-pair vocabulary")
        return a.to(torch.long) * self.pair_key_base + b.to(torch.long)

    def configure_discovery(
        self,
        *,
        candidate_capacity: int,
        top_per_forward: int,
        admission_mode: str = "utility",
    ) -> None:
        if self.pair_keys.numel() and not self.residual_buckets:
            raise ValueError("cannot discover after explicit vocabulary promotion")
        discovery_class = {
            "utility": ContrastiveVocabularyDiscovery,
            "frequency": FrequencyVocabularyDiscovery,
        }.get(str(admission_mode))
        if discovery_class is None:
            raise ValueError("discovery admission_mode must be utility or frequency")
        self.discovery = discovery_class(
            n_languages=self.language_bias.numel(),
            candidate_capacity=candidate_capacity,
            top_per_forward=top_per_forward,
        )

    def configure_diagnostic_discovery(
        self,
        *,
        candidate_capacity: int,
        top_per_forward: int,
        admission_mode: str = "utility",
    ) -> None:
        """Attach a resettable observer which cannot affect promotion state."""

        if self.pair_keys.numel() and not self.residual_buckets:
            raise ValueError("cannot diagnose discovery after explicit promotion")
        discovery_class = {
            "utility": ContrastiveVocabularyDiscovery,
            "frequency": FrequencyVocabularyDiscovery,
        }.get(str(admission_mode))
        if discovery_class is None:
            raise ValueError("discovery admission_mode must be utility or frequency")
        self.diagnostic_discovery = discovery_class(
            n_languages=self.language_bias.numel(),
            candidate_capacity=candidate_capacity,
            top_per_forward=top_per_forward,
        )

    def _explicit_lookup(self, key: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        if not self.pair_keys.numel():
            return torch.zeros_like(key, dtype=torch.bool), torch.zeros_like(key)
        positions = torch.searchsorted(self.pair_keys, key)
        safe = positions.clamp(max=self.pair_keys.numel() - 1)
        found = (positions < self.pair_keys.numel()) & (self.pair_keys[safe] == key)
        return found, self.pair_slots[safe]

    @staticmethod
    def _fold_rows(value: torch.Tensor, rows: int) -> torch.Tensor:
        if value.ndim == 0:
            return value.detach().clone()
        if rows <= 0 or value.shape[0] < rows:
            raise ValueError("residual folding requires no more rows than its source")
        destination = torch.arange(value.shape[0], device=value.device) % rows
        folded = torch.zeros(
            (rows, *value.shape[1:]), dtype=value.dtype, device=value.device
        )
        folded.index_add_(0, destination, value)
        counts = torch.bincount(destination, minlength=rows).to(value.dtype)
        return folded / counts.reshape(rows, *([1] * (value.ndim - 1)))

    @staticmethod
    def _row_state(
        state: dict[Any, Any],
        source: nn.Parameter,
        rows: torch.Tensor,
    ) -> dict[Any, Any]:
        if not state:
            raise ValueError("growth requires initialized optimizer state")
        result: dict[Any, Any] = {}
        for name, value in state.items():
            if not torch.is_tensor(value):
                result[name] = value
            elif value.ndim == 0:
                result[name] = value.detach().clone()
            elif tuple(value.shape) == tuple(source.shape):
                result[name] = value.index_select(0, rows.to(value.device)).detach().clone()
            else:
                raise ValueError(f"unknown optimizer row-state shape: {name}")
        return result

    @staticmethod
    def _concat_row_states(
        left: dict[Any, Any] | None,
        right: dict[Any, Any],
        order: torch.Tensor,
    ) -> dict[Any, Any]:
        if left is None:
            return {
                name: (
                    value.index_select(0, order.to(value.device))
                    if torch.is_tensor(value) and value.ndim
                    else value.detach().clone()
                    if torch.is_tensor(value)
                    else value
                )
                for name, value in right.items()
            }
        if set(left) != set(right):
            raise ValueError("optimizer state keys differ across growth tiers")
        result: dict[Any, Any] = {}
        for name in left:
            old_value, new_value = left[name], right[name]
            if torch.is_tensor(old_value) != torch.is_tensor(new_value):
                raise ValueError(f"optimizer state type differs: {name}")
            if not torch.is_tensor(old_value):
                if old_value != new_value:
                    raise ValueError(f"optimizer scalar state differs: {name}")
                result[name] = old_value
            elif old_value.ndim == 0:
                if not torch.equal(old_value, new_value):
                    raise ValueError(f"optimizer step differs: {name}")
                result[name] = old_value.detach().clone()
            else:
                combined = torch.cat((old_value, new_value))
                result[name] = combined.index_select(
                    0, order.to(combined.device)
                ).detach().clone()
        return result

    def grow_exact_vocabulary(
        self,
        pair_keys: tuple[int, ...] | list[int],
        *,
        optimizer: torch.optim.Optimizer,
        residual_buckets: int,
    ) -> dict[str, int]:
        """Append exact rows while retaining a residual proposal channel."""

        additions = torch.as_tensor(
            pair_keys, dtype=torch.long, device=self.merged.device
        )
        if not additions.numel():
            raise ValueError("growth requires at least one exact pair")
        if additions.numel() > 1 and not bool((additions[1:] > additions[:-1]).all()):
            raise ValueError("growth pair keys must be strictly increasing")
        if self.reclaimed_token_ids.numel():
            raise ValueError("reclaimed vocabulary cannot attach a residual tier")
        if int(residual_buckets) <= 0:
            raise ValueError("growth residual capacity must be positive")
        if self.residual_buckets and int(residual_buckets) != self.residual_buckets:
            raise ValueError("growth cannot resize an existing residual tier")
        if self.pair_keys.numel():
            positions = torch.searchsorted(self.pair_keys, additions)
            safe = positions.clamp(max=self.pair_keys.numel() - 1)
            if bool(
                ((positions < self.pair_keys.numel()) & (self.pair_keys[safe] == additions)).any()
            ):
                raise ValueError("growth pair already exists in the exact vocabulary")

        old_merged, old_score = self.merged, self.score
        old_residual_merged, old_residual_score = (
            self.residual_merged,
            self.residual_score,
        )
        first = not self.residual_buckets
        tracked = (old_merged, old_score, old_residual_merged, old_residual_score)
        occurrences = {
            parameter: sum(
                candidate is parameter
                for group in optimizer.param_groups
                for candidate in group["params"]
            )
            for parameter in tracked
        }
        residual_occurrences = {0, 1} if first else {1}
        if (
            occurrences[old_merged] != 1
            or occurrences[old_score] != 1
            or occurrences[old_residual_merged] not in residual_occurrences
            or occurrences[old_residual_score] not in residual_occurrences
        ):
            raise ValueError("growth parameters must occur once in the optimizer")

        source_merged = old_merged if first else old_residual_merged
        source_score = old_score if first else old_residual_score
        source_index = (
            self._hash(
                torch.div(additions, self.pair_key_base, rounding_mode="floor"),
                additions.remainder(self.pair_key_base),
            )
            if first
            else self._hash_residual(
                torch.div(additions, self.pair_key_base, rounding_mode="floor"),
                additions.remainder(self.pair_key_base),
            )
        )
        source_merged_state = optimizer.state.get(source_merged, {})
        source_score_state = optimizer.state.get(source_score, {})
        added_merged_state = self._row_state(
            source_merged_state, source_merged, source_index
        )
        added_score_state = self._row_state(
            source_score_state, source_score, source_index
        )

        old_keys = self.pair_keys
        old_slots = self.pair_slots
        combined_keys = torch.cat((old_keys, additions))
        order = torch.argsort(combined_keys, stable=True)
        combined_merged = torch.cat(
            (
                old_merged.index_select(0, old_slots),
                source_merged.index_select(0, source_index),
            )
        ).index_select(0, order)
        combined_score = torch.cat(
            (
                old_score.index_select(0, old_slots),
                source_score.index_select(0, source_index),
            )
        ).index_select(0, order)
        old_merged_state = (
            None
            if not old_keys.numel()
            else self._row_state(
                optimizer.state.get(old_merged, {}), old_merged, old_slots
            )
        )
        old_score_state = (
            None
            if not old_keys.numel()
            else self._row_state(
                optimizer.state.get(old_score, {}), old_score, old_slots
            )
        )
        new_merged_state = self._concat_row_states(
            old_merged_state, added_merged_state, order
        )
        new_score_state = self._concat_row_states(
            old_score_state, added_score_state, order
        )

        if first:
            if source_merged.shape[0] < int(residual_buckets):
                raise ValueError("residual capacity exceeds the proposal table")
            residual_merged_value = self._fold_rows(
                source_merged.detach(), int(residual_buckets)
            )
            residual_score_value = self._fold_rows(
                source_score.detach(), int(residual_buckets)
            )
            new_residual_merged_state = {
                name: (
                    self._fold_rows(value, int(residual_buckets))
                    if torch.is_tensor(value) and value.ndim
                    else value.detach().clone()
                    if torch.is_tensor(value)
                    else value
                )
                for name, value in source_merged_state.items()
            }
            new_residual_score_state = {
                name: (
                    self._fold_rows(value, int(residual_buckets))
                    if torch.is_tensor(value) and value.ndim
                    else value.detach().clone()
                    if torch.is_tensor(value)
                    else value
                )
                for name, value in source_score_state.items()
            }
        else:
            residual_merged_value = old_residual_merged.detach().clone()
            residual_score_value = old_residual_score.detach().clone()
            new_residual_merged_state = optimizer.state[old_residual_merged]
            new_residual_score_state = optimizer.state[old_residual_score]

        new_merged = nn.Parameter(
            combined_merged.detach().clone(), requires_grad=old_merged.requires_grad
        )
        new_score = nn.Parameter(
            combined_score.detach().clone(), requires_grad=old_score.requires_grad
        )
        new_residual_merged = nn.Parameter(
            residual_merged_value,
            requires_grad=source_merged.requires_grad,
        )
        new_residual_score = nn.Parameter(
            residual_score_value,
            requires_grad=source_score.requires_grad,
        )
        replacements = {
            old_merged: new_merged,
            old_score: new_score,
            old_residual_merged: new_residual_merged,
            old_residual_score: new_residual_score,
        }
        for group in optimizer.param_groups:
            replaced = []
            for parameter in group["params"]:
                replaced.append(replacements.get(parameter, parameter))
                # Resume compatibility deliberately removes zero-sized tensors
                # from AdamW.  The first growth transition creates the real
                # residual tier, so attach each new parameter beside the source
                # parameter whose folded value and moments initialize it.
                if first and parameter is old_score:
                    # Match ``CollapseTable`` registration order exactly.
                    # Optimizer checkpoint loading binds state positionally, so
                    # putting residual_merged beside merged would swap the
                    # score vector and residual matrix after resume.
                    if occurrences[old_residual_merged] == 0:
                        replaced.append(new_residual_merged)
                    if occurrences[old_residual_score] == 0:
                        replaced.append(new_residual_score)
            group["params"] = replaced
        for parameter in tracked:
            optimizer.state.pop(parameter, None)
        optimizer.state[new_merged] = new_merged_state
        optimizer.state[new_score] = new_score_state
        optimizer.state[new_residual_merged] = new_residual_merged_state
        optimizer.state[new_residual_score] = new_residual_score_state

        self.merged = new_merged
        self.score = new_score
        self.residual_merged = new_residual_merged
        self.residual_score = new_residual_score
        self.pair_keys = combined_keys.index_select(0, order)
        self.pair_slots = torch.arange(
            self.pair_keys.numel(), dtype=torch.long, device=self.pair_keys.device
        )
        old_exact_rows = int(old_keys.numel())
        self.buckets = int(self.pair_keys.numel())
        self.residual_buckets = int(residual_buckets)
        self.discovery = None
        self.diagnostic_discovery = None
        return {
            "old_exact_rows": old_exact_rows,
            "added_exact_rows": int(additions.numel()),
            "active_exact_rows": int(self.pair_keys.numel()),
            "residual_rows": self.residual_buckets,
        }

    def promote_exact_vocabulary(
        self,
        pair_keys: tuple[int, ...] | list[int],
        *,
        optimizer: torch.optim.Optimizer | None = None,
    ) -> dict[str, int]:
        """Map selected exact pairs to collision-free live parameter rows."""

        keys = torch.as_tensor(pair_keys, dtype=torch.long, device=self.merged.device)
        if keys.numel() and not bool((keys[1:] > keys[:-1]).all()):
            raise ValueError("promoted pair keys must be strictly increasing")
        left = torch.div(keys, self.pair_key_base, rounding_mode="floor")
        right = keys.remainder(self.pair_key_base)
        source = self._hash(left, right)

        # Preserve each original hash row for one pair. Colliding extra pairs
        # receive unused rows initialized, including Adam moments, from the same
        # source bucket.
        reserved = set(int(value) for value in source.cpu().tolist())
        free = iter(index for index in range(self.buckets) if index not in reserved)
        seen: set[int] = set()
        slots = []
        copied = []
        for src in source.cpu().tolist():
            src = int(src)
            if src not in seen:
                slot = src
                seen.add(src)
            else:
                try:
                    slot = next(free)
                except StopIteration as error:
                    raise ValueError("not enough collapse rows for collision-free promotion") from error
                copied.append((src, slot))
            slots.append(slot)

        with torch.no_grad():
            for src, dst in copied:
                self.merged[dst].copy_(self.merged[src])
                self.score[dst].copy_(self.score[src])
            if optimizer is not None:
                for parameter in (self.merged, self.score):
                    state = optimizer.state.get(parameter, {})
                    for value in state.values():
                        if (torch.is_tensor(value) and value.ndim
                                and value.shape[0] == self.buckets):
                            for src, dst in copied:
                                value[dst].copy_(value[src])

        self.pair_keys = keys
        self.pair_slots = torch.tensor(slots, dtype=torch.long, device=self.merged.device)
        self.discovery = None
        self.diagnostic_discovery = None
        return {
            "active_pairs": int(keys.numel()),
            "unique_source_buckets": len(reserved),
            "promoted_hash_collisions": len(copied),
        }

    def promote_exact_vocabulary_to_reclaimed(
        self,
        pair_keys: tuple[int, ...] | list[int],
        reclaimed_token_ids: tuple[int, ...] | list[int] | torch.Tensor,
        embedding_weight: nn.Parameter,
        *,
        optimizer: torch.optim.Optimizer,
    ) -> dict[str, int]:
        """Move a promoted exact vocabulary into existing base-table rows.

        The transition is elementwise-equivalent to ``promote_exact_vocabulary``
        at the promotion boundary. AdamW moments follow the copied vectors, and
        the obsolete full merged parameter/state is removed before training
        continues.
        """

        keys = torch.as_tensor(pair_keys, dtype=torch.long, device=self.merged.device)
        reclaimed = torch.as_tensor(
            reclaimed_token_ids, dtype=torch.long, device=self.merged.device
        )
        if not keys.numel() or keys.numel() != reclaimed.numel():
            raise ValueError("direct reclaimed promotion needs one base row per exact pair")
        if not bool((keys[1:] > keys[:-1]).all()):
            raise ValueError("promoted pair keys must be strictly increasing")
        if (reclaimed.unique().numel() != reclaimed.numel()
                or int(reclaimed.min()) < 0
                or int(reclaimed.max()) >= embedding_weight.shape[0]):
            raise ValueError("invalid direct-promotion reclaimed token ids")
        if self.pair_keys.numel() or self.reclaimed_token_ids.numel():
            raise ValueError("direct reclaimed promotion can run only once")
        if optimizer is None:
            raise ValueError("direct reclaimed promotion requires the live optimizer")

        left = torch.div(keys, self.pair_key_base, rounding_mode="floor")
        right = keys.remainder(self.pair_key_base)
        source = self._hash(left, right)
        unique_sources = int(source.unique().numel())
        old_merged, old_score = self.merged, self.score
        if tuple(old_merged.shape) != (self.buckets, embedding_weight.shape[1]):
            raise ValueError("direct promotion source table shape mismatch")

        occurrences = {
            "merged": sum(parameter is old_merged for group in optimizer.param_groups
                          for parameter in group["params"]),
            "score": sum(parameter is old_score for group in optimizer.param_groups
                         for parameter in group["params"]),
            "embedding": sum(parameter is embedding_weight for group in optimizer.param_groups
                             for parameter in group["params"]),
        }
        if occurrences != {"merged": 1, "score": 1, "embedding": 1}:
            raise ValueError(f"unexpected optimizer parameter membership: {occurrences}")
        merged_state = optimizer.state.get(old_merged)
        score_state = optimizer.state.get(old_score)
        embedding_state = optimizer.state.get(embedding_weight)
        if not merged_state or not score_state or not embedding_state:
            raise ValueError("AdamW state must exist before direct promotion")
        if set(merged_state) != set(score_state) or set(merged_state) != set(embedding_state):
            raise ValueError("optimizer state keys differ across promoted parameters")

        new_score_state: dict[object, object] = {}
        with torch.no_grad():
            embedding_weight.index_copy_(
                0, reclaimed.to(embedding_weight.device),
                old_merged.index_select(0, source).to(embedding_weight.device),
            )
            for state_key in merged_state:
                merged_value = merged_state[state_key]
                score_value = score_state[state_key]
                embedding_value = embedding_state[state_key]
                if not torch.is_tensor(merged_value):
                    if not (merged_value == score_value == embedding_value):
                        raise ValueError(f"optimizer scalar state mismatch: {state_key}")
                    new_score_state[state_key] = score_value
                    continue
                if merged_value.ndim == 0:
                    if not (torch.equal(merged_value, score_value)
                            and torch.equal(merged_value, embedding_value)):
                        raise ValueError(f"optimizer step mismatch: {state_key}")
                    new_score_state[state_key] = score_value.detach().clone()
                    continue
                if (tuple(merged_value.shape) != tuple(old_merged.shape)
                        or tuple(embedding_value.shape) != tuple(embedding_weight.shape)
                        or tuple(score_value.shape) != tuple(old_score.shape)):
                    raise ValueError(f"unknown optimizer row-state shape: {state_key}")
                embedding_value.index_copy_(
                    0, reclaimed.to(embedding_value.device),
                    merged_value.index_select(0, source).to(embedding_value.device),
                )
                new_score_state[state_key] = score_value.index_select(
                    0, source.to(score_value.device)
                ).detach().clone()

        new_score = nn.Parameter(
            old_score.index_select(0, source).detach().clone(),
            requires_grad=old_score.requires_grad,
        )
        for group in optimizer.param_groups:
            replacement = []
            for parameter in group["params"]:
                if parameter is old_merged:
                    continue
                replacement.append(new_score if parameter is old_score else parameter)
            group["params"] = replacement
        optimizer.state.pop(old_merged)
        optimizer.state.pop(old_score)
        optimizer.state[new_score] = new_score_state

        self.merged = nn.Parameter(
            torch.empty((0, embedding_weight.shape[1]), dtype=old_merged.dtype,
                        device=old_merged.device),
            requires_grad=False,
        )
        self.score = new_score
        self.pair_keys = keys
        self.pair_slots = torch.arange(keys.numel(), device=keys.device)
        self.reclaimed_token_ids = reclaimed
        self.buckets = int(keys.numel())
        self.discovery = None
        self.diagnostic_discovery = None
        return {
            "active_pairs": int(keys.numel()),
            "unique_source_buckets": unique_sources,
            "promoted_hash_collisions": int(keys.numel()) - unique_sources,
            "reclaimed_rows": int(reclaimed.numel()),
            "released_merged_rows": int(old_merged.shape[0]),
        }

    def compact_exact_vocabulary(self) -> dict[str, int]:
        """Drop inactive discovery rows after training, before final save."""

        active = int(self.pair_keys.numel())
        if not active:
            raise ValueError("cannot compact an empty exact vocabulary")
        old_buckets = self.buckets
        with torch.no_grad():
            merged = self.merged[self.pair_slots].detach().clone()
            score = self.score[self.pair_slots].detach().clone()
        self.merged = nn.Parameter(merged, requires_grad=self.merged.requires_grad)
        self.score = nn.Parameter(score, requires_grad=self.score.requires_grad)
        self.pair_slots = torch.arange(active, device=self.pair_keys.device)
        self.buckets = active
        return {"old_rows": old_buckets, "active_rows": active, "dropped_rows": old_buckets - active}

    def pool(
        self,
        vectors: torch.Tensor,     # (N, dim) token vectors, sentence-ordered
        content: torch.Tensor,     # (N,) token ids
        segment: torch.Tensor,     # (N,) sentence index
        language: torch.Tensor,    # (B,) language per sentence
        n_sentences: int,
        embedding_weight: torch.Tensor | None = None,
    ) -> torch.Tensor:
        n = content.numel()
        dim = vectors.shape[1]
        weight = torch.ones(n, dtype=vectors.dtype, device=vectors.device)
        numerator = torch.zeros(n_sentences, dim, dtype=vectors.dtype, device=vectors.device)
        denominator = torch.zeros(n_sentences, dtype=vectors.dtype, device=vectors.device)
        denominator.index_add_(0, segment, weight)
        telemetry_observed = False

        if n > 1:
            left, right = content[:-1], content[1:]
            inside = segment[:-1] == segment[1:]
            if bool(inside.any()):
                pair_left = left[inside]
                pair_right = right[inside]
                pair_segment = segment[:-1][inside]
                pair_language = language[pair_segment]
                exact_key = self._exact_key(pair_left, pair_right)
                all_positions = torch.nonzero(inside, as_tuple=False).squeeze(1)
                discovery_probability = None
                discovery_key = None
                discovery_language = None
                if self.pair_keys.numel() and self.residual_buckets:
                    found, exact_index = self._explicit_lookup(exact_key)
                    missed = ~found
                    exact_index = exact_index[found]
                    residual_index = self._hash_residual(
                        pair_left[missed], pair_right[missed]
                    )
                    exact_probability = torch.sigmoid(
                        self.score[exact_index]
                        + self.language_bias[language[pair_segment[found]]]
                    )
                    residual_probability = torch.sigmoid(
                        self.residual_score[residual_index]
                        + self.language_bias[language[pair_segment[missed]]]
                    )
                    p = torch.cat((exact_probability, residual_probability))
                    pair_segment = torch.cat((pair_segment[found], pair_segment[missed]))
                    pair_positions = torch.cat((all_positions[found], all_positions[missed]))
                    merged = torch.cat(
                        (
                            self.merged[exact_index],
                            self.residual_merged[residual_index],
                        )
                    )
                    telemetry_index = torch.cat(
                        (exact_index, residual_index + self.buckets)
                    )
                    telemetry_exact_hits = torch.cat(
                        (
                            torch.ones_like(exact_index, dtype=torch.bool),
                            torch.zeros_like(residual_index, dtype=torch.bool),
                        )
                    )
                    discovery_probability = residual_probability
                    discovery_key = exact_key[missed]
                    discovery_language = pair_language[missed]
                    audit_key = torch.cat((exact_key[found], exact_key[missed]))
                elif self.pair_keys.numel():
                    found, index = self._explicit_lookup(exact_key)
                    pair_segment = pair_segment[found]
                    pair_language = pair_language[found]
                    exact_key = exact_key[found]
                    index = index[found]
                    pair_positions = all_positions[found]
                    p = torch.sigmoid(
                        self.score[index]
                        + self.language_bias[language[pair_segment]]
                    )
                    if self.reclaimed_token_ids.numel():
                        if embedding_weight is None:
                            raise ValueError(
                                "reclaimed collapse requires the base embedding table"
                            )
                        merged = embedding_weight[
                            self.reclaimed_token_ids.index_select(0, index)
                        ]
                    else:
                        merged = self.merged[index]
                    telemetry_index = index
                    telemetry_exact_hits = torch.ones_like(index, dtype=torch.bool)
                    discovery_probability = p
                    discovery_key = exact_key
                    discovery_language = pair_language
                    audit_key = exact_key
                else:
                    index = self._hash(pair_left, pair_right)
                    pair_positions = all_positions
                    p = torch.sigmoid(
                        self.score[index]
                        + self.language_bias[language[pair_segment]]
                    )
                    merged = self.merged[index]
                    telemetry_index = index
                    telemetry_exact_hits = torch.zeros_like(index, dtype=torch.bool)
                    discovery_probability = p
                    discovery_key = exact_key
                    discovery_language = pair_language
                    audit_key = exact_key
                if self._loss_path_scale is not None:
                    if self._loss_path_observer is None:
                        raise RuntimeError("loss-path scale has no observer")
                    original_probability = p.detach()
                    p = p * self._loss_path_scale
                    saved_audit_key = audit_key.detach()
                    saved_audit_probability = original_probability
                    audit_observer = self._loss_path_observer

                    def observe_loss_path(gradient):
                        audit_observer(
                            saved_audit_key, saved_audit_probability, gradient
                        )

                    if p.requires_grad:
                        p.register_hook(observe_loss_path)
                if not p.numel():
                    _weighted_index_add_(numerator, segment, vectors, weight)
                    if self._vocabulary_telemetry is not None:
                        self._vocabulary_telemetry.observe_q15(
                            base_tokens=n,
                            probabilities=torch.empty(
                                0, dtype=torch.float32, device=content.device
                            ),
                            hash_indices=telemetry_index,
                            effective_tokens=denominator.sum(),
                            exact_hits=telemetry_exact_hits,
                        )
                        telemetry_observed = True
                    return numerator / denominator.clamp(min=1e-3).unsqueeze(1)
                discoveries = tuple(
                    observer for observer in (
                        self.discovery, self.diagnostic_discovery
                    ) if observer is not None
                )
                if (
                    discoveries
                    and self.training
                    and discovery_probability is not None
                    and discovery_probability.requires_grad
                    and discovery_probability.numel()
                ):
                    saved_key = discovery_key.detach()
                    saved_language = discovery_language.detach()
                    saved_probability = discovery_probability.detach()
                    def observe_all(gradient):
                        for observer in discoveries:
                            observer.observe(
                                saved_key, saved_language, saved_probability, gradient
                            )
                    discovery_probability.register_hook(observe_all)
                # each participant of a merge gives up half a slot
                half = 0.5 * p
                weight = weight.index_add(0, pair_positions, -half)
                weight = weight.index_add(0, pair_positions + 1, -half)
                if self.reclaimed_token_ids.numel():
                    if embedding_weight is None:
                        raise ValueError("reclaimed collapse requires the base embedding table")
                    merged = embedding_weight[
                        self.reclaimed_token_ids.index_select(0, index)
                    ]
                _weighted_index_add_(numerator, pair_segment, merged, p)
                denominator.index_add_(0, pair_segment, -p)

                if self._vocabulary_telemetry is not None:
                    self._vocabulary_telemetry.observe_q15(
                        base_tokens=n,
                        probabilities=p,
                        hash_indices=telemetry_index,
                        effective_tokens=denominator.sum(),
                        exact_hits=telemetry_exact_hits,
                    )
                    telemetry_observed = True

        _weighted_index_add_(numerator, segment, vectors, weight)
        if not telemetry_observed and self._vocabulary_telemetry is not None:
            self._vocabulary_telemetry.observe_q15(
                base_tokens=n,
                probabilities=torch.empty(
                    0, dtype=torch.float32, device=content.device
                ),
                hash_indices=torch.empty(
                    0, dtype=torch.long, device=content.device
                ),
                effective_tokens=denominator.sum(),
                exact_hits=torch.empty(
                    0, dtype=torch.bool, device=content.device
                ),
            )
        return numerator / denominator.clamp(min=1e-3).unsqueeze(1)

    def mean_merge_probability(self, content, segment, language) -> torch.Tensor:
        if content.numel() < 2:
            return torch.zeros((), device=content.device)
        inside = segment[:-1] == segment[1:]
        if not bool(inside.any()):
            return torch.zeros((), device=content.device)
        left, right = content[:-1][inside], content[1:][inside]
        pair_segment = segment[:-1][inside]
        if self.pair_keys.numel():
            found, index = self._explicit_lookup(self._exact_key(left, right))
            index = index[found]
            pair_segment = pair_segment[found]
        else:
            index = self._hash(left, right)
        if not index.numel():
            return torch.zeros((), device=content.device)
        p = torch.sigmoid(self.score[index] + self.language_bias[language[pair_segment]])
        return p.mean()


def ablate_language_bias(module: nn.Module, mode: str = "none") -> None:
    """Apply an evaluation-only necessity ablation to a loaded collapse model.

    ``mean`` replaces every language bias with their learned mean. It preserves
    the overall merge tendency while removing language-specific routing; every
    other learned parameter remains unchanged. Because the full model was
    trained with language routing active, this does not replace a separately
    trained shared-bias control for a training-time causal claim.
    """

    if mode == "none":
        return
    if mode != "mean":
        raise ValueError(f"Unknown collapse language-bias ablation: {mode!r}")
    tables = [child for child in module.modules() if isinstance(child, CollapseTable)]
    if len(tables) != 1:
        raise ValueError(f"Expected exactly one CollapseTable, found {len(tables)}")
    with torch.no_grad():
        bias = tables[0].language_bias
        bias.fill_(bias.float().mean().to(dtype=bias.dtype, device=bias.device))