mechbench-compute 0.10.0__tar.gz

mechbench_compute-0.10.0/PKG-INFO

@@ -0,0 +1,179 @@
+Metadata-Version: 2.4
+Name: mechbench-compute
+Version: 0.10.0
+Summary: Composable mechanistic-interpretability primitives: hook-aware forward, interventions, activation cache, logit lens, probes, and geometry.
+Author-email: Benji Smith <benji@shaxpir.com>
+License: MIT
+Project-URL: Homepage, https://github.com/mechbench/mechbench-compute
+Project-URL: Website, https://mechbench.ai
+Keywords: interpretability,mechanistic-interpretability,mlx,transformers,gemma
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Science/Research
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.11
+Description-Content-Type: text/markdown
+Requires-Dist: huggingface-hub>=0.24
+Requires-Dist: mechbench-schema>=0.13
+Requires-Dist: mlx>=0.20; sys_platform == "darwin" and platform_machine == "arm64"
+Requires-Dist: mlx-lm>=0.20; sys_platform == "darwin" and platform_machine == "arm64"
+Requires-Dist: mlx-vlm>=0.6.1; sys_platform == "darwin" and platform_machine == "arm64"
+Requires-Dist: numpy>=1.26
+Requires-Dist: transformers>=4.40
+Requires-Dist: matplotlib>=3.8
+Requires-Dist: scikit-learn>=1.4
+Provides-Extra: eval
+Requires-Dist: lm-eval>=0.4; extra == "eval"
+Provides-Extra: dev
+Requires-Dist: pytest>=8; extra == "dev"
+Requires-Dist: ruff>=0.5; extra == "dev"
+
+# mechbench-compute
+
+The compute engine for [mechbench](https://mechbench.ai) — composable mechanistic-interpretability primitives built on MLX.
+
+This repo provides:
+
+- **Hook-aware forward pass.** One canonical path through the model; instrumentation via named hook points and TransformerLens-style callbacks.
+- **Declarative interventions.** `Ablate`, `Capture`, `Patch` primitives composable into a single `model.run(..., interventions=[...])` call.
+- **Activation cache.** `ActivationCache` container for collected activations; bf16 throughout, float32 only at the analysis boundary.
+- **Architecture adapter.** `Arch` dataclass that handles per-variant differences (layer count, global-attention pattern, RoPE parameters, etc.). Currently supports Gemma 4 E4B and E2B; the adapter pattern follows TransformerLens 3's `TransformerBridge`.
+- **Analysis helpers.** Logit lens, direct logit attribution (`accumulated_resid`, `decompose_resid`, `head_results`, `logit_attrs`), fact vectors, centroid decoding, probe primitives, head-weight static analysis, geometry metrics.
+- **Plot helpers.** Matplotlib conventions baked in for quick diagnostic figures — not the full visualization surface (that lives in `mechbench-ui`).
+
+See [`PACKAGE_README.md`](PACKAGE_README.md) for the full API tour and worked examples.
+
+## Install
+
+```bash
+pip install mechbench-compute
+```
+
+From source:
+
+```bash
+git clone https://github.com/mechbench/mechbench-compute.git
+cd mechbench-compute
+pip install -e '.[dev]'
+```
+
+Apple Silicon required (MLX is the only supported backend today). A PyTorch backend would live as `mechbench_compute.backends.torch` alongside the MLX one if/when the need arises; splitting repos by backend is explicitly not planned.
+
+## Quick start
+
+```python
+from mechbench_compute import Model, Ablate, Capture
+
+model = Model.load()
+ids = model.tokenize("Complete this sentence with one word: The Eiffel Tower is in")
+
+result = model.run(ids)
+for tok, p in result.top_k(model.tokenizer, k=5):
+    print(f"{tok!r:20s} p={p:.4f}")
+```
+
+## Distributional-target training (`distill` + `lora`)
+
+Primitives for training a model toward a specified *distribution* over
+responses rather than toward example responses (task
+`000114`):
+soft-target cross-entropy at decision tokens has gradient P − T, so the
+adapter learns to *emit the distribution*.
+
+```python
+import mlx.nn as nn
+import mlx.optimizers as optim
+import numpy as np
+from mechbench_compute import Model, distill, lora
+from mechbench_compute.distill import TargetMap
+
+model = Model.load()
+tok = model.tokenizer
+
+# A target is a Map<String, Double> — hardcoded, from JSON, or uniform —
+# with whole-map transforms that each return a new map:
+target = TargetMap.from_json("weights.json").sqrt().normalize()
+target = TargetMap.uniform([str(i) for i in range(1, 7)])   # fair d6
+
+# Compile it against the rendered prompt: items become token paths
+# (multi-token items share trie nodes; a closer appends continuation
+# anchors so the flattening can't leak past the envelope):
+prompt = distill.render_chat(tok, system, "Please roll the die.",
+                             prefill='{ "roll": ')
+trie = target.tokenize(tok, prompt, closer=" }")
+
+n = lora.apply_lora(model.lm)                 # freeze + wrap q/v projections
+step = nn.value_and_grad(model.lm, distill.soft_ce)
+opt = optim.Adam(learning_rate=1e-4)
+rng = np.random.default_rng(7)
+for _ in range(steps):
+    batch = [trie.hard_example(trie.sample(rng)) for _ in range(3)]
+    batch.append(trie.marginal_example())     # exact first-token marginal
+    batch.append(sharp_anchor)                # keeps confident tasks sharp
+    loss, grads = step(model.lm, batch)
+    opt.update(model.lm, grads)
+
+lora.save_adapter(model.lm, "adapter.safetensors")
+# Later, on a fresh model: merge + exact undo
+handle = lora.fuse(model.lm, lora.load_adapter("adapter.safetensors"),
+                   scale=16 / 8)              # alpha / rank from training
+lora.restore(model.lm, handle)
+```
+
+Calibration is measured at item level (`trie.score`, `distill.item_metrics`
+— captured mass, entropy, KL from target) and at the decision token
+(`distill.first_token_metrics`). `python -m mechbench_compute._smoke_distill`
+runs the full lifecycle on E2B.
+
+**Two forward paths.** Training and scoring call `Model.lm` (the text
+decoder, uniform across families) directly — plain module calls,
+differentiable, no instrumentation. `Model.run` remains the hook-aware
+forward for capture/patch/lens work. Adapters bridge the two: `fuse` an
+adapter into the weights and every instrumented run sees the adapted
+model; `restore` flips it back, so base-vs-adapted comparisons run in one
+script.
+
+**Scoring tiers** (task 000227): `score_items` is the sequential
+reference oracle; `score_items_batched` adds length-bucketed batching
+(~1.5×); `score_items_fast` additionally splits the forward via
+`Model.trunk_hidden` / `Model.head_logits` and unembeds only the
+supervised rows (~1.5–2.1× vs oracle, family-dependent — best when
+items share no prefix, e.g. cross-document scoring);
+`score_items_cached` encodes a shared prompt **once** into a KV cache
+(`Model.prompt_cache`) and scores each item's 1–4 suffix tokens against
+a per-item copy — the tier for shared-prompt batteries. Measured
+(flat name batteries, cached vs oracle): E2B 2.8×, Gemma-3-4B 3.4×,
+Qwen-3B 2.6×, Llama-8B 3.0×; positions/attention exact by construction,
+bf16 envelope from decomposed attention: mass-region |ΔlogP| ≤ 0.45,
+renormalized KL ≤ 2.4e-2 bits, idempotent (0.0 across repeat calls).
+Fast-tier envelope: max |ΔlogP| ≤ 0.99 (deep tail) / ≤ 0.24
+(mass region), renorm KL ≤ 6.5e-3 bits. All of it is bf16 matmul
+tiling — same rows, same math, verified bit-exact where shapes match.
+Flat-target KL diagnostics weight the tail, so switch tiers only
+between comparisons, never mid-experiment. Known upstream limit
+(mlx 0.31.2 / mlx-lm 0.31.3 / mlx-vlm 0.6.1): **batched cached decoding
+corrupts every batch row after the first** on both stacks (reproduced
+with natively built B=4 caches and identical rows), which is why the
+cached tier is per-item; batched suffix scoring behind an upstream fix
+is the remaining ~5–10× path.
+
+## Status
+
+The `Arch` adapter supports Gemma 4 E4B and E2B; generalization to other architecture families is ongoing.
+
+The substrate epic that will define how intermediate results are cached and shared across experiments is `000162` (DAG solver + content-addressed memoization). It consumes the canonical-serialization guarantee from `000161` (binary formats) and the path grammar from `000163` (identity scheme).
+
+## Relationship to other mechbench repos
+
+- **`mechbench-schema`** — the typed emission contract. `mechbench-compute` emits records shaped by schema types.
+- **`mechbench-experiments`** — research scripts and findings that consume this package. Uses `mechbench-compute` as its primary dependency.
+- **`mechbench-runner`** — exposes these primitives as agent-callable tools. Imports `mechbench-compute`.
+- **`mechbench-ui`** — TypeScript frontend. Does not import `mechbench-compute` directly; reads bundles produced by it through the `mechbench-schema` contract.
+
+See [mechbench.ai](https://mechbench.ai) for the family overview and the design principles.
+
+## License
+
+MIT.

mechbench_compute-0.10.0/README.md

@@ -0,0 +1,147 @@
+# mechbench-compute
+
+The compute engine for [mechbench](https://mechbench.ai) — composable mechanistic-interpretability primitives built on MLX.
+
+This repo provides:
+
+- **Hook-aware forward pass.** One canonical path through the model; instrumentation via named hook points and TransformerLens-style callbacks.
+- **Declarative interventions.** `Ablate`, `Capture`, `Patch` primitives composable into a single `model.run(..., interventions=[...])` call.
+- **Activation cache.** `ActivationCache` container for collected activations; bf16 throughout, float32 only at the analysis boundary.
+- **Architecture adapter.** `Arch` dataclass that handles per-variant differences (layer count, global-attention pattern, RoPE parameters, etc.). Currently supports Gemma 4 E4B and E2B; the adapter pattern follows TransformerLens 3's `TransformerBridge`.
+- **Analysis helpers.** Logit lens, direct logit attribution (`accumulated_resid`, `decompose_resid`, `head_results`, `logit_attrs`), fact vectors, centroid decoding, probe primitives, head-weight static analysis, geometry metrics.
+- **Plot helpers.** Matplotlib conventions baked in for quick diagnostic figures — not the full visualization surface (that lives in `mechbench-ui`).
+
+See [`PACKAGE_README.md`](PACKAGE_README.md) for the full API tour and worked examples.
+
+## Install
+
+```bash
+pip install mechbench-compute
+```
+
+From source:
+
+```bash
+git clone https://github.com/mechbench/mechbench-compute.git
+cd mechbench-compute
+pip install -e '.[dev]'
+```
+
+Apple Silicon required (MLX is the only supported backend today). A PyTorch backend would live as `mechbench_compute.backends.torch` alongside the MLX one if/when the need arises; splitting repos by backend is explicitly not planned.
+
+## Quick start
+
+```python
+from mechbench_compute import Model, Ablate, Capture
+
+model = Model.load()
+ids = model.tokenize("Complete this sentence with one word: The Eiffel Tower is in")
+
+result = model.run(ids)
+for tok, p in result.top_k(model.tokenizer, k=5):
+    print(f"{tok!r:20s} p={p:.4f}")
+```
+
+## Distributional-target training (`distill` + `lora`)
+
+Primitives for training a model toward a specified *distribution* over
+responses rather than toward example responses (task
+`000114`):
+soft-target cross-entropy at decision tokens has gradient P − T, so the
+adapter learns to *emit the distribution*.
+
+```python
+import mlx.nn as nn
+import mlx.optimizers as optim
+import numpy as np
+from mechbench_compute import Model, distill, lora
+from mechbench_compute.distill import TargetMap
+
+model = Model.load()
+tok = model.tokenizer
+
+# A target is a Map<String, Double> — hardcoded, from JSON, or uniform —
+# with whole-map transforms that each return a new map:
+target = TargetMap.from_json("weights.json").sqrt().normalize()
+target = TargetMap.uniform([str(i) for i in range(1, 7)])   # fair d6
+
+# Compile it against the rendered prompt: items become token paths
+# (multi-token items share trie nodes; a closer appends continuation
+# anchors so the flattening can't leak past the envelope):
+prompt = distill.render_chat(tok, system, "Please roll the die.",
+                             prefill='{ "roll": ')
+trie = target.tokenize(tok, prompt, closer=" }")
+
+n = lora.apply_lora(model.lm)                 # freeze + wrap q/v projections
+step = nn.value_and_grad(model.lm, distill.soft_ce)
+opt = optim.Adam(learning_rate=1e-4)
+rng = np.random.default_rng(7)
+for _ in range(steps):
+    batch = [trie.hard_example(trie.sample(rng)) for _ in range(3)]
+    batch.append(trie.marginal_example())     # exact first-token marginal
+    batch.append(sharp_anchor)                # keeps confident tasks sharp
+    loss, grads = step(model.lm, batch)
+    opt.update(model.lm, grads)
+
+lora.save_adapter(model.lm, "adapter.safetensors")
+# Later, on a fresh model: merge + exact undo
+handle = lora.fuse(model.lm, lora.load_adapter("adapter.safetensors"),
+                   scale=16 / 8)              # alpha / rank from training
+lora.restore(model.lm, handle)
+```
+
+Calibration is measured at item level (`trie.score`, `distill.item_metrics`
+— captured mass, entropy, KL from target) and at the decision token
+(`distill.first_token_metrics`). `python -m mechbench_compute._smoke_distill`
+runs the full lifecycle on E2B.
+
+**Two forward paths.** Training and scoring call `Model.lm` (the text
+decoder, uniform across families) directly — plain module calls,
+differentiable, no instrumentation. `Model.run` remains the hook-aware
+forward for capture/patch/lens work. Adapters bridge the two: `fuse` an
+adapter into the weights and every instrumented run sees the adapted
+model; `restore` flips it back, so base-vs-adapted comparisons run in one
+script.
+
+**Scoring tiers** (task 000227): `score_items` is the sequential
+reference oracle; `score_items_batched` adds length-bucketed batching
+(~1.5×); `score_items_fast` additionally splits the forward via
+`Model.trunk_hidden` / `Model.head_logits` and unembeds only the
+supervised rows (~1.5–2.1× vs oracle, family-dependent — best when
+items share no prefix, e.g. cross-document scoring);
+`score_items_cached` encodes a shared prompt **once** into a KV cache
+(`Model.prompt_cache`) and scores each item's 1–4 suffix tokens against
+a per-item copy — the tier for shared-prompt batteries. Measured
+(flat name batteries, cached vs oracle): E2B 2.8×, Gemma-3-4B 3.4×,
+Qwen-3B 2.6×, Llama-8B 3.0×; positions/attention exact by construction,
+bf16 envelope from decomposed attention: mass-region |ΔlogP| ≤ 0.45,
+renormalized KL ≤ 2.4e-2 bits, idempotent (0.0 across repeat calls).
+Fast-tier envelope: max |ΔlogP| ≤ 0.99 (deep tail) / ≤ 0.24
+(mass region), renorm KL ≤ 6.5e-3 bits. All of it is bf16 matmul
+tiling — same rows, same math, verified bit-exact where shapes match.
+Flat-target KL diagnostics weight the tail, so switch tiers only
+between comparisons, never mid-experiment. Known upstream limit
+(mlx 0.31.2 / mlx-lm 0.31.3 / mlx-vlm 0.6.1): **batched cached decoding
+corrupts every batch row after the first** on both stacks (reproduced
+with natively built B=4 caches and identical rows), which is why the
+cached tier is per-item; batched suffix scoring behind an upstream fix
+is the remaining ~5–10× path.
+
+## Status
+
+The `Arch` adapter supports Gemma 4 E4B and E2B; generalization to other architecture families is ongoing.
+
+The substrate epic that will define how intermediate results are cached and shared across experiments is `000162` (DAG solver + content-addressed memoization). It consumes the canonical-serialization guarantee from `000161` (binary formats) and the path grammar from `000163` (identity scheme).
+
+## Relationship to other mechbench repos
+
+- **`mechbench-schema`** — the typed emission contract. `mechbench-compute` emits records shaped by schema types.
+- **`mechbench-experiments`** — research scripts and findings that consume this package. Uses `mechbench-compute` as its primary dependency.
+- **`mechbench-runner`** — exposes these primitives as agent-callable tools. Imports `mechbench-compute`.
+- **`mechbench-ui`** — TypeScript frontend. Does not import `mechbench-compute` directly; reads bundles produced by it through the `mechbench-schema` contract.
+
+See [mechbench.ai](https://mechbench.ai) for the family overview and the design principles.
+
+## License
+
+MIT.

mechbench_compute-0.10.0/mechbench_compute/__init__.py

@@ -0,0 +1,312 @@
+"""mechbench_compute — a mechanistic-interpretability framework for
+Google's Gemma 4 E4B running locally on Apple Silicon via MLX.
+
+Quick start:
+
+    from mechbench_compute import Model
+
+    model = Model.load("mlx-community/gemma-4-E4B-it-bf16")
+    ids = model.tokenize("Complete this sentence with one word: The Eiffel Tower is in")
+
+    # Forward pass, no instrumentation:
+    result = model.run(ids)
+    for tok, p in result.top_k(model.tokenizer, k=5):
+        print(f'{tok!r:20s} p={p:.4f}')
+
+Capture activations:
+
+    result = model.run(ids, capture=['blocks.23.attn.weights',
+                                      'blocks.14.mlp_out'])
+    weights = result.cache['blocks.23.attn.weights']  # [1, 8, S, S], bf16
+    mlp = result.cache['blocks.14.mlp_out']           # [1, S, 2560], bf16
+
+Modify activations with a hook:
+
+    import mlx.core as mx
+
+    def zero_layer_14_mlp(act, info):
+        return mx.zeros_like(act)
+
+    result = model.run(ids, hooks={'blocks.14.mlp_out': zero_layer_14_mlp})
+
+Declarative interventions (Ablate / Capture / Patch), prompt tooling,
+logit-lens + geometry helpers, and matplotlib plot helpers are all
+re-exported from this module — see README.md for the full API tour.
+
+The full list of hook points is at mechbench_compute.all_hook_names().
+"""
+
+# Everything below the substrate check assumes a substrate. Rather than
+# refusing to import at all, the package now *loads* anywhere and refuses
+# on use — because the modules that report on a machine with no backend
+# (`backends`, `inventory`) live inside this package, and a gate at
+# import time made them reachable only from a machine that did not need
+# them. `doctor` is exactly that machine's tool.
+from .backends import (
+    BACKENDS,
+    Backend,
+    active as active_backend,
+    available as available_backends,
+    describe_platform,
+    require as _require_backend,
+)
+
+if active_backend() is not None:
+    from ._arch import (
+        Arch,
+        D_MODEL,
+        GLOBAL_LAYERS,
+        LAYER_HOOK_POINTS,
+        N_HEADS,
+        N_LAYERS,
+        VOCAB_SIZE,
+        all_hook_names,
+        layer_type,
+    )
+    from .cache import ActivationCache
+    from .errors import (
+        CacheKeyError,
+        InterpError,
+        InvalidHookName,
+        LayerIndexOutOfRange,
+    )
+    from .geometry import (
+        VocabConcentration,
+        centroid_decode,
+        cluster_purity,
+        cohesion,
+        cosine_matrix,
+        effective_vocab_size,
+        entropy_bits,
+        fact_vectors,
+        fact_vectors_at,
+        fact_vectors_at_hook,
+        fact_vectors_pooled,
+        intra_inter_separation,
+        iterate_clusters,
+        nearest_neighbor_purity,
+        orthogonalize_against,
+        silhouette_cosine,
+        top_k_mass,
+        vocab_concentration,
+    )
+    from .generate import generate_labeled_corpus, generate_text
+    from .head_weights import (
+        CircuitAnalysis,
+        CircuitComponent,
+        HeadSpec,
+        PositionWrite,
+        get_head_spec,
+        head_key_tokens,
+        head_ov_actual_writes,
+        head_ov_position_writes,
+        head_read_tokens,
+        ov_circuit,
+        qk_circuit,
+    )
+    from .attribution import (
+        accumulated_resid,
+        decompose_resid,
+        head_results,
+        logit_attrs,
+    )
+    from .probes import Probe
+    from .hooks import HookFn, HookInfo, parse_hook_name
+    from .interventions import Ablate, Capture, Intervention, Patch, compose
+    from .lens import logit_lens_final, logit_lens_per_position
+    from .model import Model, RunResult
+    from .plot import (
+        bar_by_layer,
+        grouped_row_heatmap,
+        head_heatmap,
+        intensity_curve,
+        leaderboard_bar,
+        lens_trajectory,
+        logprob_trajectory,
+        pca_scatter,
+        position_heatmap,
+        probe_diagonal_heatmap,
+        similarity_heatmap,
+    )
+    from .prompts import (
+        Prompt,
+        PromptSet,
+        ValidatedPrompt,
+        ValidatedPromptSet,
+    )
+    from .distill import (
+        Example,
+        TargetMap,
+        TargetTrie,
+        first_token_metrics,
+        item_metrics,
+        render_chat,
+        score_items,
+        score_items_batched,
+        score_items_cached,
+        score_items_fast,
+        soft_ce,
+    )
+    from .lora import (
+        LoRALinear,
+        apply_lora,
+        fuse,
+        load_adapter,
+        restore,
+        save_adapter,
+    )
+
+else:
+
+    def __getattr__(name: str) -> object:
+        """Explain, rather than fail with a missing-module traceback.
+
+        Reached only on a machine with no compute backend, where every
+        name below is genuinely unavailable — so the same explanation is
+        the right answer for all of them.
+
+        Submodules that touch MLX directly (`model`, `protocol`, …) are
+        not covered: importing one still raises out of MLX itself. That
+        is the situation `mechbench-runner doctor` exists to catch first.
+        """
+        if name.startswith("__"):
+            raise AttributeError(name)
+
+        # Submodules that do not touch the substrate stay reachable —
+        # `backends` and `inventory` are precisely what reports on a
+        # machine like this one, and `from mechbench_compute import
+        # inventory` has to keep working. Without this the package
+        # attribute hook shadows the submodule and answers "no backend"
+        # for a module that never needed one.
+        import importlib
+
+        try:
+            return importlib.import_module(f".{name}", __name__)
+        except ImportError:
+            pass
+
+        # Either a name that lives behind the substrate, or a submodule
+        # that failed on MLX itself. Both get the explanation rather than
+        # a traceback out of a package nobody asked for.
+        _require_backend()
+        raise AttributeError(name)  # unreachable; require() always raises
+
+__version__ = "0.10.0"
+
+__all__ = [
+    # Substrate — the part that answers on a machine with no backend
+    "Backend",
+    "BACKENDS",
+    "active_backend",
+    "available_backends",
+    "describe_platform",
+    # Main API
+    "Model",
+    "RunResult",
+    "ActivationCache",
+    # Declarative interventions
+    "Ablate",
+    "Capture",
+    "Patch",
+    "Intervention",
+    "compose",
+    # Prompts (specific prompt collections live in experiments.prompts)
+    "Prompt",
+    "PromptSet",
+    "ValidatedPrompt",
+    "ValidatedPromptSet",
+    # Distributional-target training (distill + lora, task 000114)
+    "TargetMap",
+    "TargetTrie",
+    "Example",
+    "soft_ce",
+    "render_chat",
+    "score_items",
+    "score_items_batched",
+    "score_items_cached",
+    "score_items_fast",
+    "item_metrics",
+    "first_token_metrics",
+    "LoRALinear",
+    "apply_lora",
+    "save_adapter",
+    "load_adapter",
+    "fuse",
+    "restore",
+    # Logit lens
+    "logit_lens_final",
+    "logit_lens_per_position",
+    # Direct logit attribution + residual decomposition
+    "accumulated_resid",
+    "decompose_resid",
+    "head_results",
+    "logit_attrs",
+    # Fact vectors + geometry
+    "fact_vectors",
+    "fact_vectors_at",
+    "fact_vectors_at_hook",
+    "fact_vectors_pooled",
+    "centroid_decode",
+    "cohesion",
+    "cosine_matrix",
+    "intra_inter_separation",
+    "iterate_clusters",
+    "cluster_purity",
+    "silhouette_cosine",
+    "nearest_neighbor_purity",
+    "orthogonalize_against",
+    # Probes (persistent concept vectors)
+    "Probe",
+    # Generation (for corpus-building workflows)
+    "generate_text",
+    "generate_labeled_corpus",
+    # Head-weight analysis (static interp on W_Q/W_K/W_V/W_O)
+    "HeadSpec",
+    "CircuitComponent",
+    "CircuitAnalysis",
+    "PositionWrite",
+    "get_head_spec",
+    "head_read_tokens",
+    "head_key_tokens",
+    "qk_circuit",
+    "ov_circuit",
+    "head_ov_position_writes",
+    "head_ov_actual_writes",
+    # Vocabulary-space concentration
+    "VocabConcentration",
+    "vocab_concentration",
+    "top_k_mass",
+    "entropy_bits",
+    "effective_vocab_size",
+    # Plot helpers
+    "bar_by_layer",
+    "lens_trajectory",
+    "logprob_trajectory",
+    "position_heatmap",
+    "pca_scatter",
+    "similarity_heatmap",
+    "head_heatmap",
+    "probe_diagonal_heatmap",
+    "grouped_row_heatmap",
+    "intensity_curve",
+    "leaderboard_bar",
+    # Hook types (for users writing raw callbacks)
+    "HookInfo",
+    "HookFn",
+    "parse_hook_name",
+    # Architecture facts
+    "Arch",
+    "N_LAYERS",
+    "D_MODEL",
+    "N_HEADS",
+    "VOCAB_SIZE",
+    "GLOBAL_LAYERS",
+    "LAYER_HOOK_POINTS",
+    "layer_type",
+    "all_hook_names",
+    # Errors
+    "InterpError",
+    "InvalidHookName",
+    "LayerIndexOutOfRange",
+    "CacheKeyError",
+]