mechbench-schema 0.13.0__tar.gz

mechbench_schema-0.13.0/PKG-INFO

@@ -0,0 +1,157 @@
+Metadata-Version: 2.4
+Name: mechbench-schema
+Version: 0.13.0
+Summary: Typed emission contract for the mechbench family. Pydantic models as single source of truth; generates TypeScript bindings for mechbench-ui.
+Author-email: Benji Smith <benji@shaxpir.com>
+License: MIT
+Project-URL: Homepage, https://github.com/mechbench/mechbench-schema
+Project-URL: Website, https://mechbench.ai
+Keywords: interpretability,schema,pydantic,mechbench
+Classifier: Development Status :: 3 - Alpha
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Topic :: Software Development :: Libraries
+Requires-Python: >=3.11
+Description-Content-Type: text/markdown
+Requires-Dist: pydantic>=2.6
+Requires-Dist: cbor2>=5.6
+Provides-Extra: codegen
+Requires-Dist: datamodel-code-generator[typescript]>=0.25; extra == "codegen"
+Provides-Extra: dev
+Requires-Dist: pytest>=8; extra == "dev"
+Requires-Dist: ruff>=0.5; extra == "dev"
+
+# mechbench-schema
+
+The typed emission contract for the [mechbench](https://mechbench.ai) family.
+
+**Pydantic models are the single source of truth.** TypeScript bindings are generated from them; both halves are published from this one repo to two package registries.
+
+| Consumer | Registry | Package |
+| --- | --- | --- |
+| Python (core, agent, remote, experiments) | PyPI | `pip install mechbench-schema` |
+| TypeScript (ui, skills) | npm | `npm install mechbench-schema` |
+
+The two packages share a name and a version. CI regenerates the TS bindings and fails on drift, so the two halves cannot disagree.
+
+## Organization: by domain axis, not by consumer
+
+Modules are named after the **indexing axis** of the records they carry, not after who renders them. This keeps the schema from becoming "chart data" just because charts are the current biggest consumer; the same records drive CSV exports, agent tool surfaces, the memoization cache, and anything else downstream.
+
+| module | category | archetypal records |
+|---|---|---|
+| `per_layer_data` | indexed by `layer` | ablation damage, DLA diffs, convergence summaries |
+| `per_head_data` | indexed by `(layer, head)` | per-head DLA, OV rank-0 singular values |
+| `attention_trace` | indexed by `(layer, head, pos_from, pos_to)` | attention patterns |
+| `per_layer_per_position_data` | indexed by `(layer, position)` | logit-lens trajectories, causal-trace grids |
+| `vector_data` | atomic: one direction in residual space | captured / steering / probe / centroid vectors |
+| `cluster_data` | collection of vectors with aggregate stats | named clusters, cross-cluster statistics |
+| `identity` | path grammar | `MechbenchPath` validated-string type |
+
+Every payload model is part of a discriminated union keyed by `kind` (or `origin` for vectors), so the UI can dispatch rendering on the tag and the DAG solver can reason about compatible outputs.
+
+## Repo layout
+
+```
+mechbench-schema/
+├── src/mechbench_schema/              # Pydantic models — the source of truth
+│   ├── __init__.py                    # __all__ + __schema_all__
+│   ├── attention_trace.py
+│   ├── cluster_data.py
+│   ├── identity.py
+│   ├── per_head_data.py
+│   ├── per_layer_data.py
+│   ├── per_layer_per_position_data.py
+│   └── vector_data.py
+├── ts/                                # published to npm as "mechbench-schema"
+│   ├── package.json
+│   ├── tsconfig.json
+│   └── src/
+│       ├── index.ts
+│       ├── generated.ts               # codegen output — DO NOT EDIT
+│       └── schema.json                # consolidated JSON Schema — codegen output
+├── scripts/
+│   └── codegen.py                     # Pydantic → JSON Schema → TypeScript
+├── pyproject.toml                     # published to PyPI as "mechbench-schema"
+└── README.md
+```
+
+## Editing the contract
+
+1. Edit an existing module under `src/mechbench_schema/`, or add a new module (following the domain-axis convention).
+2. Export the new type from `src/mechbench_schema/__init__.py`. Add it to `__all__` for Python consumers and to `__schema_all__` for codegen (wire types only).
+3. Run `python scripts/codegen.py`. This writes `ts/src/schema.json` and `ts/src/generated.ts`.
+4. Commit everything in one PR — Python source, generated JSON Schema, generated TypeScript.
+
+CI runs codegen and `git diff --exit-code ts/src/schema.json ts/src/generated.ts`. A drifted PR cannot land.
+
+The codegen uses `json-schema-to-typescript` (via `npx`) for the TS emission. Python → JSON Schema comes from Pydantic's built-in `model_json_schema()` (for `BaseModel` subclasses) and `TypeAdapter(...).json_schema()` (for type aliases like the discriminated unions).
+
+## Install (Python)
+
+```bash
+pip install mechbench-schema
+```
+
+```python
+from mechbench_schema import LayerAblationPayload, DlaSweepPayload
+
+payload = LayerAblationPayload(
+    experiment="step_02_layer_ablation",
+    description="Per-layer ablation damage on FACTUAL_15.",
+    model="mlx-community/gemma-4-E4B-it-bf16",
+    n_layers=42,
+    global_layers=[5, 11, 17, 23, 29, 35, 41],
+    prompts=[...],
+    aggregates=LayerAggregates(mean=[...], median=[...]),
+)
+
+# Emit to JSON — mode="json" canonicalizes datetimes, bytes, etc.
+import json
+path.write_text(json.dumps(payload.model_dump(mode="json"), indent=2))
+```
+
+## Install (TypeScript)
+
+```bash
+npm install mechbench-schema
+```
+
+```ts
+import type { LayerAblationPayload, Vector, MechbenchPath } from "mechbench-schema";
+
+const payload: LayerAblationPayload = JSON.parse(text);
+// All fields, discriminated unions, and nested types come from the generated bindings.
+```
+
+## Identity grammar
+
+`MechbenchPath` is the type used to address every object in the mechbench family — charts, articles, experiments, corpuses, probes, cached DAG intermediates. Five categories:
+
+- `<owner>/<project>/<folders>/<leaf>` — user-named
+- `~canonical/<area>/<path>/<leaf>` — human-readable aliases for globally-shared content
+- `~system/<area>/<path>/<leaf>` — platform / first-party
+- `~hash/<algo>:<digest>` — global content-hashed (deduplicated across users)
+- `<owner>/<project>/~hash/<algo>:<digest>` — workspace-scoped content-hashed
+
+Full spec in the meta repo's IDENTITY_AND_NAMESPACING.md. The Python side (`mechbench_schema.identity`) owns the authoritative validator; a mirror implementation lives in `mechbench-ui/src/lib/mechbenchPath.ts` for UI-side early failure.
+
+## Why one repo, two publications
+
+Schemas that live in two repos drift. The only invariant that matters — "Python and TS agree on the shape" — is enforced by keeping the source in one place and generating the target. Consumers never need the other language's toolchain to install; `pip` and `npm` each resolve to a clean single-language package.
+
+See [the family overview](https://mechbench.ai) for the rationale behind this and other multi-repo decisions.
+
+## Status
+
+Version 0.8.0. Seven modules in active use; the legacy `records.py` holding pen was retired when the domain-axis reorg completed. `mechbench-experiments`' two exporters (`step_02_layer_ablation`, `step_33_dla_factual_sweep`) emit via Pydantic models; `mechbench-ui`'s chart interfaces are one-line aliases over the generated TS types.
+
+Open work lives in the meta repo's `tasks/mechbench-schema/` directory. The two notable open epics downstream of this repo:
+
+- **`000161`** — compact binary formats for records at rest and in transit (safetensors for tensor-bulk, parquet for record-collections).
+- **`000163`** — the identity-and-namespacing epic that produced `MechbenchPath`; Phase 4 (content-addressing grammar) remains open and coordinates with `000162` (the DAG-solver epic in `mechbench-compute`).
+
+## License
+
+MIT.

mechbench_schema-0.13.0/README.md

@@ -0,0 +1,133 @@
+# mechbench-schema
+
+The typed emission contract for the [mechbench](https://mechbench.ai) family.
+
+**Pydantic models are the single source of truth.** TypeScript bindings are generated from them; both halves are published from this one repo to two package registries.
+
+| Consumer | Registry | Package |
+| --- | --- | --- |
+| Python (core, agent, remote, experiments) | PyPI | `pip install mechbench-schema` |
+| TypeScript (ui, skills) | npm | `npm install mechbench-schema` |
+
+The two packages share a name and a version. CI regenerates the TS bindings and fails on drift, so the two halves cannot disagree.
+
+## Organization: by domain axis, not by consumer
+
+Modules are named after the **indexing axis** of the records they carry, not after who renders them. This keeps the schema from becoming "chart data" just because charts are the current biggest consumer; the same records drive CSV exports, agent tool surfaces, the memoization cache, and anything else downstream.
+
+| module | category | archetypal records |
+|---|---|---|
+| `per_layer_data` | indexed by `layer` | ablation damage, DLA diffs, convergence summaries |
+| `per_head_data` | indexed by `(layer, head)` | per-head DLA, OV rank-0 singular values |
+| `attention_trace` | indexed by `(layer, head, pos_from, pos_to)` | attention patterns |
+| `per_layer_per_position_data` | indexed by `(layer, position)` | logit-lens trajectories, causal-trace grids |
+| `vector_data` | atomic: one direction in residual space | captured / steering / probe / centroid vectors |
+| `cluster_data` | collection of vectors with aggregate stats | named clusters, cross-cluster statistics |
+| `identity` | path grammar | `MechbenchPath` validated-string type |
+
+Every payload model is part of a discriminated union keyed by `kind` (or `origin` for vectors), so the UI can dispatch rendering on the tag and the DAG solver can reason about compatible outputs.
+
+## Repo layout
+
+```
+mechbench-schema/
+├── src/mechbench_schema/              # Pydantic models — the source of truth
+│   ├── __init__.py                    # __all__ + __schema_all__
+│   ├── attention_trace.py
+│   ├── cluster_data.py
+│   ├── identity.py
+│   ├── per_head_data.py
+│   ├── per_layer_data.py
+│   ├── per_layer_per_position_data.py
+│   └── vector_data.py
+├── ts/                                # published to npm as "mechbench-schema"
+│   ├── package.json
+│   ├── tsconfig.json
+│   └── src/
+│       ├── index.ts
+│       ├── generated.ts               # codegen output — DO NOT EDIT
+│       └── schema.json                # consolidated JSON Schema — codegen output
+├── scripts/
+│   └── codegen.py                     # Pydantic → JSON Schema → TypeScript
+├── pyproject.toml                     # published to PyPI as "mechbench-schema"
+└── README.md
+```
+
+## Editing the contract
+
+1. Edit an existing module under `src/mechbench_schema/`, or add a new module (following the domain-axis convention).
+2. Export the new type from `src/mechbench_schema/__init__.py`. Add it to `__all__` for Python consumers and to `__schema_all__` for codegen (wire types only).
+3. Run `python scripts/codegen.py`. This writes `ts/src/schema.json` and `ts/src/generated.ts`.
+4. Commit everything in one PR — Python source, generated JSON Schema, generated TypeScript.
+
+CI runs codegen and `git diff --exit-code ts/src/schema.json ts/src/generated.ts`. A drifted PR cannot land.
+
+The codegen uses `json-schema-to-typescript` (via `npx`) for the TS emission. Python → JSON Schema comes from Pydantic's built-in `model_json_schema()` (for `BaseModel` subclasses) and `TypeAdapter(...).json_schema()` (for type aliases like the discriminated unions).
+
+## Install (Python)
+
+```bash
+pip install mechbench-schema
+```
+
+```python
+from mechbench_schema import LayerAblationPayload, DlaSweepPayload
+
+payload = LayerAblationPayload(
+    experiment="step_02_layer_ablation",
+    description="Per-layer ablation damage on FACTUAL_15.",
+    model="mlx-community/gemma-4-E4B-it-bf16",
+    n_layers=42,
+    global_layers=[5, 11, 17, 23, 29, 35, 41],
+    prompts=[...],
+    aggregates=LayerAggregates(mean=[...], median=[...]),
+)
+
+# Emit to JSON — mode="json" canonicalizes datetimes, bytes, etc.
+import json
+path.write_text(json.dumps(payload.model_dump(mode="json"), indent=2))
+```
+
+## Install (TypeScript)
+
+```bash
+npm install mechbench-schema
+```
+
+```ts
+import type { LayerAblationPayload, Vector, MechbenchPath } from "mechbench-schema";
+
+const payload: LayerAblationPayload = JSON.parse(text);
+// All fields, discriminated unions, and nested types come from the generated bindings.
+```
+
+## Identity grammar
+
+`MechbenchPath` is the type used to address every object in the mechbench family — charts, articles, experiments, corpuses, probes, cached DAG intermediates. Five categories:
+
+- `<owner>/<project>/<folders>/<leaf>` — user-named
+- `~canonical/<area>/<path>/<leaf>` — human-readable aliases for globally-shared content
+- `~system/<area>/<path>/<leaf>` — platform / first-party
+- `~hash/<algo>:<digest>` — global content-hashed (deduplicated across users)
+- `<owner>/<project>/~hash/<algo>:<digest>` — workspace-scoped content-hashed
+
+Full spec in the meta repo's IDENTITY_AND_NAMESPACING.md. The Python side (`mechbench_schema.identity`) owns the authoritative validator; a mirror implementation lives in `mechbench-ui/src/lib/mechbenchPath.ts` for UI-side early failure.
+
+## Why one repo, two publications
+
+Schemas that live in two repos drift. The only invariant that matters — "Python and TS agree on the shape" — is enforced by keeping the source in one place and generating the target. Consumers never need the other language's toolchain to install; `pip` and `npm` each resolve to a clean single-language package.
+
+See [the family overview](https://mechbench.ai) for the rationale behind this and other multi-repo decisions.
+
+## Status
+
+Version 0.8.0. Seven modules in active use; the legacy `records.py` holding pen was retired when the domain-axis reorg completed. `mechbench-experiments`' two exporters (`step_02_layer_ablation`, `step_33_dla_factual_sweep`) emit via Pydantic models; `mechbench-ui`'s chart interfaces are one-line aliases over the generated TS types.
+
+Open work lives in the meta repo's `tasks/mechbench-schema/` directory. The two notable open epics downstream of this repo:
+
+- **`000161`** — compact binary formats for records at rest and in transit (safetensors for tensor-bulk, parquet for record-collections).
+- **`000163`** — the identity-and-namespacing epic that produced `MechbenchPath`; Phase 4 (content-addressing grammar) remains open and coordinates with `000162` (the DAG-solver epic in `mechbench-compute`).
+
+## License
+
+MIT.

mechbench_schema-0.13.0/pyproject.toml

@@ -0,0 +1,36 @@
+[build-system]
+requires = ["setuptools>=68", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "mechbench-schema"
+version = "0.13.0"
+description = "Typed emission contract for the mechbench family. Pydantic models as single source of truth; generates TypeScript bindings for mechbench-ui."
+readme = "README.md"
+requires-python = ">=3.11"
+license = { text = "MIT" }
+authors = [{ name = "Benji Smith", email = "benji@shaxpir.com" }]
+keywords = ["interpretability", "schema", "pydantic", "mechbench"]
+classifiers = [
+    "Development Status :: 3 - Alpha",
+    "License :: OSI Approved :: MIT License",
+    "Programming Language :: Python :: 3.11",
+    "Programming Language :: Python :: 3.12",
+    "Topic :: Software Development :: Libraries",
+]
+dependencies = [
+    "pydantic>=2.6",
+    "cbor2>=5.6",
+]
+
+[project.optional-dependencies]
+codegen = ["datamodel-code-generator[typescript]>=0.25"]
+dev = ["pytest>=8", "ruff>=0.5"]
+
+[project.urls]
+Homepage = "https://github.com/mechbench/mechbench-schema"
+Website = "https://mechbench.ai"
+
+[tool.setuptools.packages.find]
+where = ["src"]
+include = ["mechbench_schema*"]

mechbench_schema-0.13.0/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

mechbench_schema-0.13.0/src/mechbench_schema/__init__.py

@@ -0,0 +1,213 @@
+"""mechbench-schema — typed emission contract for the mechbench family.
+
+This package is the single source of truth for every interpretability-record
+shape that crosses a repo boundary. Python consumers (mechbench-compute,
+mechbench-runner, mechbench-experiments) import types from
+here; TypeScript consumers (mechbench-ui) import types from the parallel
+mechbench-schema npm package, which is generated from these Pydantic models
+via scripts/codegen.py.
+
+Rule of thumb: if the shape is ever serialized to disk, sent over the wire,
+or rendered by a non-Python consumer, it belongs here. If the shape is only
+used as an in-memory container within a single repo, it does not.
+"""
+
+from .attention_trace import (
+    AttentionPattern,
+    AttentionTraceData,
+)
+from .codec_cbor import (
+    dump_canonical,
+    load_canonical,
+    load_raw,
+)
+from .identity import (
+    InvalidPathError,
+    MechbenchPath,
+    ParsedPath,
+    PathCategory,
+    make_canonical_path,
+    make_global_hash_path,
+    make_platform_path,
+    make_scoped_hash_path,
+    make_user_path,
+    parse_path,
+)
+from .cluster_data import (
+    Cluster,
+    ClusterSet,
+)
+from .per_head_data import (
+    PerHeadBase,
+    PerHeadData,
+    PerHeadScalarGrid,
+)
+from .document_data import (
+    AnnotationAnchor,
+    AnnotationLayer,
+    AnnotationValue,
+    BASE_KIND_ANNOTATED_TOKENS,
+    BASE_KIND_CONVERSATION,
+    BASE_KIND_TEXT,
+    DocumentCollection,
+    DocumentItem,
+    DocumentPayload,
+    GenerationSpan,
+    KindManifest,
+    RendererBinding,
+    Segment,
+    Segmentation,
+    Trace,
+    Turn,
+)
+from .metric_data import (
+    MetricColumn,
+    MetricTable,
+)
+from .provenance import (
+    Emitted,
+    Fidelity,
+    Provenance,
+    ToolInfo,
+    fidelity_satisfies,
+    fingerprint_params,
+)
+from .per_layer_data import (
+    AblationPrompt,
+    ConvergencePayload,
+    ConvergenceRow,
+    DlaPrompt,
+    DlaSweepPayload,
+    LayerAblationPayload,
+    LayerAggregates,
+    PerLayerBase,
+    PerLayerData,
+)
+from .per_layer_per_position_data import (
+    LogitLensTrajectory,
+    PerLayerPerPositionBase,
+    PerLayerPerPositionData,
+)
+from .vector_data import (
+    CapturedVector,
+    CentroidVector,
+    HookKind,
+    ProbeVector,
+    SteeringVector,
+    Vector,
+)
+
+__version__ = "0.12.0"
+
+__all__ = [
+    # Document collections + annotations + kind manifests (task 000239)
+    "DocumentCollection",
+    "DocumentItem",
+    "DocumentPayload",
+    "Trace",
+    "GenerationSpan",
+    "Segmentation",
+    "Segment",
+    "Turn",
+    "AnnotationLayer",
+    "AnnotationValue",
+    "AnnotationAnchor",
+    "KindManifest",
+    "RendererBinding",
+    "BASE_KIND_TEXT",
+    "BASE_KIND_CONVERSATION",
+    "BASE_KIND_ANNOTATED_TOKENS",
+    # Metric tables (task 000244)
+    "MetricTable",
+    "MetricColumn",
+    # Provenance + fidelity (the emission envelope, task 000237)
+    "Provenance",
+    "ToolInfo",
+    "Emitted",
+    "Fidelity",
+    "fidelity_satisfies",
+    "fingerprint_params",
+    # Per-layer data (records indexed by transformer layer)
+    "PerLayerBase",
+    "LayerAggregates",
+    "AblationPrompt",
+    "LayerAblationPayload",
+    "DlaPrompt",
+    "DlaSweepPayload",
+    "ConvergenceRow",
+    "ConvergencePayload",
+    "PerLayerData",
+    # Per-head data (records indexed by (layer, head))
+    "PerHeadBase",
+    "PerHeadScalarGrid",
+    "PerHeadData",
+    # Attention-trace data (records indexed by (layer, head, position, position))
+    "AttentionPattern",
+    "AttentionTraceData",
+    # Vector data (atomic directions in residual-stream space)
+    "CapturedVector",
+    "SteeringVector",
+    "ProbeVector",
+    "CentroidVector",
+    "HookKind",
+    "Vector",
+    # Cluster data (collections of vectors with aggregate stats)
+    "Cluster",
+    "ClusterSet",
+    # Per-(layer, position) data (residual-stream-axis records)
+    "PerLayerPerPositionBase",
+    "LogitLensTrajectory",
+    "PerLayerPerPositionData",
+    # Identity / namespacing (see docs/IDENTITY_AND_NAMESPACING.md)
+    "MechbenchPath",
+    "ParsedPath",
+    "PathCategory",
+    "InvalidPathError",
+    "parse_path",
+    "make_user_path",
+    "make_canonical_path",
+    "make_platform_path",
+    "make_global_hash_path",
+    "make_scoped_hash_path",
+]
+
+# Names that should be emitted to TS via codegen. Subset of __all__ — excludes
+# Python-side utilities (functions, exceptions, non-Pydantic dataclasses)
+# that don't cross the wire.
+__schema_all__ = [
+    # Per-layer data
+    "PerLayerBase",
+    "LayerAggregates",
+    "AblationPrompt",
+    "LayerAblationPayload",
+    "DlaPrompt",
+    "DlaSweepPayload",
+    "ConvergenceRow",
+    "ConvergencePayload",
+    "PerLayerData",
+    # Per-head data
+    "PerHeadBase",
+    "PerHeadScalarGrid",
+    "PerHeadData",
+    # Attention-trace data
+    "AttentionPattern",
+    "AttentionTraceData",
+    # Vector data
+    "CapturedVector",
+    "SteeringVector",
+    "ProbeVector",
+    "CentroidVector",
+    "HookKind",
+    "Vector",
+    # Cluster data
+    "Cluster",
+    "ClusterSet",
+    # Per-(layer, position) data
+    "PerLayerPerPositionBase",
+    "LogitLensTrajectory",
+    "PerLayerPerPositionData",
+    # Identity (only the validated-string type crosses the wire; parsers and
+    # helpers are Python-side-only)
+    "MechbenchPath",
+    "PathCategory",
+]

mechbench_schema-0.13.0/src/mechbench_schema/attention_trace.py

@@ -0,0 +1,81 @@
+"""Attention-trace data shapes.
+
+Records indexed by `(layer, head, position_from, position_to)` — the
+attention patterns that heads compute during a forward pass. An
+attention trace for a single prompt over a full model is a collection
+of `(layer, head)` weight matrices, each `[seq_len × seq_len]` for
+self-attention.
+
+This module organizes the attention-pattern axis as a first-class
+category, parallel to per_layer_data and per_head_data. A downstream
+consumer can render a single (layer, head) matrix as a heatmap, a
+small-multiples grid over (layer, head), a per-head attention-entropy
+summary, etc.
+
+The kinds today:
+
+  attention_pattern — one (layer, head) attention matrix, possibly
+                      bundled with token labels for axis rendering.
+                      This is the atomic unit; a full trace is a
+                      collection of these.
+
+Future kinds (file when a real consumer needs them):
+- attention_entropy — per-(layer, head, position) scalar summary.
+- attention_topk    — per-(layer, head, position) top-k target
+                      positions by weight.
+- attention_trace   — a sweep containing many attention_patterns plus
+                      shared metadata (prompt_id, token_labels). May
+                      want a `PerHeadOf[AttentionPattern]`-style
+                      container when task 000157's generic work lands.
+
+Migrated from the old records.py as part of the domain-axis reorg
+(task 000156).
+"""
+
+from __future__ import annotations
+
+from typing import Annotated, Literal, Union
+
+from pydantic import BaseModel, Field
+
+
+class AttentionPattern(BaseModel):
+    """Post-softmax attention weights for one (layer, head) over a sequence.
+
+    The weights matrix is stored as a flat row-major `[n_queries * n_keys]`
+    list. Consumers reshape on ingest; this keeps the wire format
+    JSON-native and trivially diffable. `token_labels`, when present,
+    labels both axes (self-attention) — for cross-attention the caller
+    should split into query_labels / key_labels (not yet modeled;
+    add when needed).
+    """
+
+    kind: Literal["attention_pattern"] = "attention_pattern"
+    layer: int = Field(..., ge=0, description="0-indexed transformer layer.")
+    head: int = Field(..., ge=0, description="0-indexed attention head within the layer.")
+    n_queries: int = Field(..., ge=1, description="Query-side sequence length.")
+    n_keys: int = Field(..., ge=1, description="Key-side sequence length.")
+    weights: list[float] = Field(
+        ...,
+        description=(
+            "Row-major flattened post-softmax weights, shape "
+            "[n_queries * n_keys]. weights[q * n_keys + k] is the attention "
+            "weight from query q to key k."
+        ),
+    )
+    token_labels: list[str] | None = Field(
+        None,
+        description=(
+            "Optional per-position decoded tokens, for axis labelling. "
+            "Assumes self-attention (same labels on both axes); if the "
+            "shape is asymmetric, model the split explicitly."
+        ),
+    )
+
+
+# --- Discriminated union over all attention-trace kinds ---------------------
+
+AttentionTraceData = Annotated[
+    Union[AttentionPattern],
+    Field(discriminator="kind"),
+]