"""Shape and evidence audit for parallel multi-head representation routing.""" from __future__ import annotations from dataclasses import asdict, dataclass from datetime import datetime import hashlib import json import math import re from typing import Any MAX_HEADS = 32 MAX_SEQUENCE = 128 MAX_DIMENSION = 128 MAX_OPERATIONS = 4_000_000 MIN_RESOLUTION = 1e-300 MAX_RESOLUTION = 1e-3 IDENTIFIER = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._:@/-]{0,127}$") SHA256_ID = re.compile(r"^[a-z][a-z0-9-]*@sha256:[0-9a-f]{64}$") def _identifier(name: str, value: object) -> str: if type(value) is not str or IDENTIFIER.fullmatch(value) is None: raise ValueError(f"{name} must be a bounded identifier") return value def _timestamp(name: str, value: object) -> str: if type(value) is not str or len(value) > 64: raise ValueError(f"{name} must be a bounded ISO-8601 timestamp") try: parsed = datetime.fromisoformat(value.replace("Z", "+00:00")) except ValueError as exc: raise ValueError(f"{name} must be an ISO-8601 timestamp") from exc if parsed.tzinfo is None: raise ValueError(f"{name} must include an offset") return value def _integer(name: str, value: object, minimum: int, maximum: int) -> int: if type(value) is not int or not minimum <= value <= maximum: raise ValueError(f"{name} must be an integer in [{minimum}, {maximum}]") return value def _finite(name: str, value: object, resolution: float, limit: float) -> float: if type(value) not in (int, float) or type(value) is bool: raise TypeError(f"{name} must be numeric and not bool") checked = float(value) if not math.isfinite(checked): raise ValueError(f"{name} must be finite") if abs(checked) > limit: raise ValueError(f"{name} exceeds its contracted magnitude") if checked != 0.0 and abs(checked) < resolution: raise ValueError(f"{name} is below the contracted numeric resolution") return checked def _canonical(value: Any) -> Any: if isinstance(value, dict): return {key: _canonical(item) for key, item in sorted(value.items())} if isinstance(value, tuple): return [_canonical(item) for item in value] return value def _content_id(prefix: str, payload: dict[str, Any]) -> str: encoded = json.dumps( _canonical(payload), sort_keys=True, separators=(",", ":"), allow_nan=False ).encode("utf-8") return f"{prefix}@sha256:{hashlib.sha256(encoded).hexdigest()}" def _matrix( name: str, value: object, rows: int, columns: int, resolution: float, limit: float, ) -> tuple[tuple[float, ...], ...]: if type(value) is not tuple or len(value) != rows: raise TypeError(f"{name} must be an immutable tuple with declared height") checked_rows: list[tuple[float, ...]] = [] for row in value: if type(row) is not tuple or len(row) != columns: raise TypeError(f"{name} rows must be immutable tuples with declared width") checked_rows.append(tuple(_finite(name, item, resolution, limit) for item in row)) return tuple(checked_rows) @dataclass(frozen=True) class HeadProjection: head_id: str projection_version: str query_projection: tuple[tuple[float, ...], ...] key_projection: tuple[tuple[float, ...], ...] value_projection: tuple[tuple[float, ...], ...] def __post_init__(self) -> None: _identifier("head_id", self.head_id) _identifier("projection_version", self.projection_version) for field in ("query_projection", "key_projection", "value_projection"): if type(getattr(self, field)) is not tuple: raise TypeError(f"{field} must be an immutable tuple") @dataclass(frozen=True) class MultiHeadContract: contract_version: str model_id: str model_version: str layer_id: str input_semantics: str output_semantics: str model_dimension: int head_dimension: int value_dimension: int maximum_tokens: int maximum_operations: int mask_semantics: str scaling_semantics: str softmax_semantics: str specialization_claim_policy: str minimum_numeric_resolution: float maximum_absolute_input: float maximum_absolute_logit: float heads: tuple[HeadProjection, ...] output_projection: tuple[tuple[float, ...], ...] scope: str model_owner: str audit_owner: str def __post_init__(self) -> None: for field in ( "contract_version", "model_id", "model_version", "layer_id", "input_semantics", "output_semantics", "scope", "model_owner", "audit_owner", ): _identifier(field, getattr(self, field)) for field in ("model_dimension", "head_dimension", "value_dimension"): _integer(field, getattr(self, field), 1, MAX_DIMENSION) _integer("maximum_tokens", self.maximum_tokens, 1, MAX_SEQUENCE) _integer("maximum_operations", self.maximum_operations, 1, MAX_OPERATIONS) if self.mask_semantics != "bidirectional-all-token-pairs-v1": raise ValueError("unsupported mask_semantics") if self.scaling_semantics != "divide-dot-by-sqrt-head-dimension-v1": raise ValueError("unsupported scaling_semantics") if self.softmax_semantics != "rowwise-max-shift-binary64-v1": raise ValueError("unsupported softmax_semantics") if self.specialization_claim_policy != "requires-task-ablation-evidence-v1": raise ValueError("unsupported specialization_claim_policy") resolution = _finite( "minimum_numeric_resolution", self.minimum_numeric_resolution, MIN_RESOLUTION, MAX_RESOLUTION, ) if not MIN_RESOLUTION <= resolution <= MAX_RESOLUTION: raise ValueError("minimum_numeric_resolution is outside supported bounds") input_limit = _finite( "maximum_absolute_input", self.maximum_absolute_input, resolution, 1 / resolution ) logit_limit = _finite( "maximum_absolute_logit", self.maximum_absolute_logit, resolution, 100.0 ) if input_limit < resolution or logit_limit < resolution: raise ValueError("numeric limits must be at least the declared resolution") if type(self.heads) is not tuple or not 1 <= len(self.heads) <= MAX_HEADS: raise TypeError("heads must be a non-empty bounded immutable tuple") checked_heads: list[HeadProjection] = [] for head in self.heads: if type(head) is not HeadProjection: raise TypeError("heads must contain concrete HeadProjection records") head.__post_init__() checked_heads.append( HeadProjection( head_id=head.head_id, projection_version=head.projection_version, query_projection=_matrix( "query_projection", head.query_projection, self.model_dimension, self.head_dimension, resolution, input_limit ), key_projection=_matrix( "key_projection", head.key_projection, self.model_dimension, self.head_dimension, resolution, input_limit ), value_projection=_matrix( "value_projection", head.value_projection, self.model_dimension, self.value_dimension, resolution, input_limit ), ) ) head_ids = [head.head_id for head in checked_heads] if len(head_ids) != len(set(head_ids)): raise ValueError("head IDs must be unique") checked_output = _matrix( "output_projection", self.output_projection, len(checked_heads) * self.value_dimension, self.model_dimension, resolution, input_limit, ) object.__setattr__(self, "heads", tuple(checked_heads)) object.__setattr__(self, "output_projection", checked_output) @property def content_id(self) -> str: return _content_id("multi-head-contract", asdict(self)) @property def projection_content_id(self) -> str: return _content_id( "multi-head-projections", { "model_id": self.model_id, "model_version": self.model_version, "layer_id": self.layer_id, "heads": tuple(asdict(head) for head in self.heads), "output_projection": self.output_projection, }, ) @dataclass(frozen=True) class RepresentationEvidence: evidence_id: str evidence_content_id: str contract_content_id: str projection_content_id: str source_id: str source_version: str observed_at: str scope: str token_ids: tuple[str, ...] representations: tuple[tuple[float, ...], ...] def __post_init__(self) -> None: for field in ("evidence_id", "source_id", "source_version", "scope"): _identifier(field, getattr(self, field)) for field in ( "evidence_content_id", "contract_content_id", "projection_content_id", ): if type(getattr(self, field)) is not str or SHA256_ID.fullmatch(getattr(self, field)) is None: raise ValueError(f"{field} must be a content identity") _timestamp("observed_at", self.observed_at) if type(self.token_ids) is not tuple or type(self.representations) is not tuple: raise TypeError("token_ids and representations must be immutable tuples") @classmethod def capture( cls, *, contract: MultiHeadContract, evidence_id: str, source_id: str, source_version: str, observed_at: str, scope: str, token_ids: object, representations: object, ) -> "RepresentationEvidence": if type(contract) is not MultiHeadContract: raise TypeError("contract must be a concrete MultiHeadContract") contract.__post_init__() checked_scope = _identifier("scope", scope) if checked_scope != contract.scope: raise ValueError("evidence scope does not match contract scope") if type(token_ids) is not tuple or not 1 <= len(token_ids) <= contract.maximum_tokens: raise TypeError("token_ids must be a non-empty bounded immutable tuple") checked_ids = tuple(_identifier("token_id", item) for item in token_ids) if len(checked_ids) != len(set(checked_ids)): raise ValueError("token IDs must be unique within the audit fixture") checked_representations = _matrix( "representations", representations, len(checked_ids), contract.model_dimension, contract.minimum_numeric_resolution, contract.maximum_absolute_input ) head_count = len(contract.heads) operations = len(checked_ids) * contract.model_dimension * head_count * ( 2 * contract.head_dimension + contract.value_dimension ) + len(checked_ids) ** 2 * head_count * ( contract.head_dimension + contract.value_dimension ) if operations > contract.maximum_operations: raise ValueError("evidence exceeds maximum_operations") payload = { "contract_content_id": contract.content_id, "projection_content_id": contract.projection_content_id, "evidence_id": _identifier("evidence_id", evidence_id), "source_id": _identifier("source_id", source_id), "source_version": _identifier("source_version", source_version), "observed_at": _timestamp("observed_at", observed_at), "scope": checked_scope, "token_ids": checked_ids, "representations": checked_representations, } return cls( evidence_content_id=_content_id("multi-head-evidence", payload), **payload, ) @dataclass(frozen=True) class HeadResult: head_id: str projection_version: str output_content_id: str attention_weights: tuple[tuple[float, ...], ...] outputs: tuple[tuple[float, ...], ...] @dataclass(frozen=True) class MultiHeadAudit: decision: str contract_content_id: str projection_content_id: str evidence_content_id: str heads: tuple[HeadResult, ...] combined_output: tuple[tuple[float, ...], ...] specialization_conclusion: str required_specialization_evidence: str material_evidence_id: str def _project( rows: tuple[tuple[float, ...], ...], matrix: tuple[tuple[float, ...], ...] ) -> tuple[tuple[float, ...], ...]: return tuple( tuple( math.fsum(row[index] * matrix[index][column] for index in range(len(row))) for column in range(len(matrix[0])) ) for row in rows ) def _softmax(logits: tuple[float, ...], resolution: float) -> tuple[float, ...]: peak = max(logits) exponentials = tuple(math.exp(value - peak) for value in logits) if any(value != 0.0 and value < resolution for value in exponentials): raise ValueError("softmax term is below the contracted numeric resolution") total = math.fsum(exponentials) if not math.isfinite(total) or total <= 0.0: raise ArithmeticError("softmax denominator is invalid") return tuple(value / total for value in exponentials) def _revalidate( contract: MultiHeadContract, evidence: RepresentationEvidence ) -> RepresentationEvidence: if type(contract) is not MultiHeadContract or type(evidence) is not RepresentationEvidence: raise TypeError("audit requires concrete MultiHeadContract and RepresentationEvidence") contract.__post_init__() evidence.__post_init__() rebuilt = RepresentationEvidence.capture( contract=contract, evidence_id=evidence.evidence_id, source_id=evidence.source_id, source_version=evidence.source_version, observed_at=evidence.observed_at, scope=evidence.scope, token_ids=evidence.token_ids, representations=evidence.representations, ) if rebuilt.contract_content_id != evidence.contract_content_id: raise ValueError("evidence belongs to a different multi-head contract") if rebuilt.projection_content_id != evidence.projection_content_id: raise ValueError("evidence binds different projection contents") if rebuilt.evidence_content_id != evidence.evidence_content_id: raise ValueError("evidence content identity does not match its contents") return rebuilt def audit_multi_head_routing( contract: MultiHeadContract, evidence: RepresentationEvidence ) -> MultiHeadAudit: """Validate per-head projections and compute parallel routing without interpretation.""" checked = _revalidate(contract, evidence) head_results: list[HeadResult] = [] head_outputs: list[tuple[tuple[float, ...], ...]] = [] for head in contract.heads: queries = _project(checked.representations, head.query_projection) keys = _project(checked.representations, head.key_projection) values = _project(checked.representations, head.value_projection) weight_rows: list[tuple[float, ...]] = [] output_rows: list[tuple[float, ...]] = [] for query in queries: logits = tuple( math.fsum(left * right for left, right in zip(query, key)) / math.sqrt(contract.head_dimension) for key in keys ) if any(not math.isfinite(item) or abs(item) > contract.maximum_absolute_logit for item in logits): raise ValueError("scaled head logit exceeds the contracted range") weights = _softmax(logits, contract.minimum_numeric_resolution) weight_rows.append(weights) output_rows.append( tuple( math.fsum(weights[index] * values[index][column] for index in range(len(values))) for column in range(contract.value_dimension) ) ) outputs = tuple(output_rows) output_id = _content_id( "head-output", { "head_id": head.head_id, "projection_version": head.projection_version, "projection_content_id": contract.projection_content_id, "evidence_content_id": checked.evidence_content_id, "attention_weights": tuple(weight_rows), "outputs": outputs, }, ) head_results.append( HeadResult( head_id=head.head_id, projection_version=head.projection_version, output_content_id=output_id, attention_weights=tuple(weight_rows), outputs=outputs, ) ) head_outputs.append(outputs) concatenated = tuple( tuple( value for head_output in head_outputs for value in head_output[token_index] ) for token_index in range(len(checked.token_ids)) ) combined = _project(concatenated, contract.output_projection) if any(not math.isfinite(value) for row in combined for value in row): raise ArithmeticError("combined multi-head output is non-finite") material = { "contract_content_id": contract.content_id, "projection_content_id": contract.projection_content_id, "evidence_content_id": checked.evidence_content_id, "head_output_ids": tuple(item.output_content_id for item in head_results), "combined_output": combined, "specialization_conclusion": "NOT_ESTABLISHED_FROM_HEAD_EXISTENCE", } return MultiHeadAudit( decision="PASS", contract_content_id=contract.content_id, projection_content_id=contract.projection_content_id, evidence_content_id=checked.evidence_content_id, heads=tuple(head_results), combined_output=combined, specialization_conclusion="NOT_ESTABLISHED_FROM_HEAD_EXISTENCE", required_specialization_evidence="TASK_METRICS_PLUS_CONTROLLED_HEAD_ABLATION", material_evidence_id=_content_id("multi-head-material", material), ) ILLUSTRATIVE_CONTRACT = MultiHeadContract( contract_version="multi-head-routing-v1", model_id="illustrative-transformer", model_version="model-v1", layer_id="layer-0-self-attention", input_semantics="ordered-token-representations", output_semantics="projected-parallel-routing-result", model_dimension=2, head_dimension=1, value_dimension=1, maximum_tokens=8, maximum_operations=4096, mask_semantics="bidirectional-all-token-pairs-v1", scaling_semantics="divide-dot-by-sqrt-head-dimension-v1", softmax_semantics="rowwise-max-shift-binary64-v1", specialization_claim_policy="requires-task-ablation-evidence-v1", minimum_numeric_resolution=1e-12, maximum_absolute_input=100.0, maximum_absolute_logit=40.0, heads=( HeadProjection("head-0", "projection-v1", ((1.0,), (0.0,)), ((1.0,), (0.0,)), ((1.0,), (0.0,))), HeadProjection("head-1", "projection-v1", ((0.0,), (1.0,)), ((0.0,), (1.0,)), ((0.0,), (1.0,))), ), output_projection=((1.0, 0.0), (0.0, 1.0)), scope="academy:foundation-model-internals", model_owner="team:model", audit_owner="team:model-platform", ) ILLUSTRATIVE_EVIDENCE = RepresentationEvidence.capture( contract=ILLUSTRATIVE_CONTRACT, evidence_id="two-token-fixture", source_id="academy-fixture", source_version="fixture-v1", observed_at="2026-08-25T00:00:00+00:00", scope="academy:foundation-model-internals", token_ids=("token-a", "token-b"), representations=((1.0, 0.0), (0.0, 1.0)), ) def format_example() -> str: report = audit_multi_head_routing(ILLUSTRATIVE_CONTRACT, ILLUSTRATIVE_EVIDENCE) return "\n".join( ( "example=illustrative_only", f"contract_id={report.contract_content_id}", f"projection_id={report.projection_content_id}", f"head_ids={','.join(head.head_id for head in report.heads)}", f"token0_output={','.join(f'{value:.3f}' for value in report.combined_output[0])}", f"specialization={report.specialization_conclusion}", f"required_evidence={report.required_specialization_evidence}", f"material_id={report.material_evidence_id}", f"decision={report.decision}", ) ) if __name__ == "__main__": print(format_example())