mechbench-compute
The compute engine for mechbench — 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,Patchprimitives composable into a singlemodel.run(..., interventions=[...])call. - Activation cache.
ActivationCachecontainer for collected activations; bf16 throughout, float32 only at the analysis boundary. - Architecture adapter.
Archdataclass 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'sTransformerBridge. - 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 for the full API tour and worked examples.
Install
pip install mechbench-compute
From source:
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
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: soft-target cross-entropy at decision tokens has gradient P − T, so the adapter learns to emit the distribution.
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: 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 hook-aware forwards cover Gemma 4, Gemma 3, Qwen 2.x and Llama 3.x.
Relationship to other mechbench repos
mechbench-schema— the typed emission contract.mechbench-computeemits records shaped by schema types.mechbench-experiments— research scripts and findings that consume this package. Usesmechbench-computeas its primary dependency.mechbench-runner— exposes these primitives as agent-callable tools. Importsmechbench-compute.mechbench-ui— TypeScript frontend. Does not importmechbench-computedirectly; reads bundles produced by it through themechbench-schemacontract.
See mechbench.ai for the family overview and the design principles.
License
MIT.
1 maintainer
Metadata
1 maintainer
Release files for mechbench-compute 0.199.0
For a detailed explanation of source distributions (sdists) and built distributions (wheels), please see the package formats documentation.
Source distribution (sdist)
| File | Size | Uploaded | |
|---|---|---|---|
| mechbench_compute-0.199.0.tar.gz | 1.1 MB | Details |
Built distribution (wheel)
| File | Interpreter | ABI | Platform | Reset |
|---|---|---|---|---|
| mechbench_compute-0.199.0-py3-none-any.whl | Python 3 | none | any | Details |
Total release size: 2.1 MB
Release files / mechbench_compute-0.199.0.tar.gz
| Download URL | mechbench_compute-0.199.0.tar.gz |
|---|---|
| Size | 1.1 MB |
| Tags | Source |
|
SHA-256 checksum How to use checksums |
ce40e4ad35a56a01e0ed51012cec08916f3c2f5c924310d329faf2a41a80fa70
|
|
BLAKE2b-256 checksum How to use checksums |
2ec7439a8df3fbf0686b857b66f35f363af9200de99a5184607d947ed25db31c
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
Yes |
| Uploaded via |
twine/7.0.0 CPython/3.13.14
|
Provenance
Provenance describes where a file came from. On PyPI, provenance is shared via attestations, which provide a verifiable record of the build or publishing details. View details, limitations and caveats.
PyPI Publish Attestation
PyPI verified that this artifact, at this checksum, originated from the publisher listed below.
Signed by GitHub Actions, verified by PyPI on Oct 15, 2026.
Transparency logRelease files / mechbench_compute-0.199.0-py3-none-any.whl
| Download URL | mechbench_compute-0.199.0-py3-none-any.whl |
|---|---|
| Size | 987.3 kB |
| Tags | Python 3 |
|
SHA-256 checksum How to use checksums |
bf1dde71cbf4f7302089260b75ffbc51da7dcd61c07121211b999405f5b0dce2
|
|
BLAKE2b-256 checksum How to use checksums |
0e78ea7c74197e00918524614bd7762132e26105014ca33259ce4dba95436907
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
Yes |
| Uploaded via |
twine/7.0.0 CPython/3.13.14
|
Provenance
Provenance describes where a file came from. On PyPI, provenance is shared via attestations, which provide a verifiable record of the build or publishing details. View details, limitations and caveats.
PyPI Publish Attestation
PyPI verified that this artifact, at this checksum, originated from the publisher listed below.
Signed by GitHub Actions, verified by PyPI on Oct 15, 2026.
Transparency log
