Release v4.1.0

.claude/commands/add-model-support.md

@@ -46,7 +46,7 @@ Each step names the doc to read **when you reach that step** — don't load all
 
 6. **Add the HF repo entry** to [`data/supported_models.json`](../../transformer_lens/tools/model_registry/data/supported_models.json) per [§Adding the HF repo to the registry](../../transformer_lens/model_bridge/supported_architectures/AGENTS.md#adding-the-hf-repo-to-the-registry). Ask the user about adding canonical sibling variants from `CANONICAL_AUTHORS_BY_ARCH[<HFArchClass>]`.
 
-7. **Verify** end-to-end: `/verify-model $ARGUMENTS`. Read both `status` AND per-phase scores. `STATUS_VERIFIED` means hard gates passed (see [§Phase-score thresholds](../../transformer_lens/tools/model_registry/AGENTS.md#phase-score-thresholds)) — but P4's 50% bar is intentionally lenient. P4 well below 100% on a small parity-test model + `status==1` → suspect missing `preprocess_weights` fold or wrong `default_prepend_bos`; investigate before step 8.
+7. **Verify** end-to-end: `/verify-model $ARGUMENTS`. Read both `status` AND per-phase scores. `STATUS_VERIFIED` means hard gates passed (see [§Phase-score thresholds](../../transformer_lens/tools/model_registry/AGENTS.md#phase-score-thresholds)) — but P4's ~54.5% bar (`p4_pass_threshold()`, not a fixed number) is intentionally lenient. P4 well below 100% on a small parity-test model + `status==1` → suspect missing `preprocess_weights` fold or wrong `default_prepend_bos`; investigate before step 8.
 
 8. **Write tests** per [§Required tests](../../transformer_lens/model_bridge/supported_architectures/AGENTS.md#required-tests) (unit + integration). Copy the closest sibling.
 

.claude/commands/verify-model.md

@@ -37,7 +37,7 @@ Full reference: [tools/model_registry/AGENTS.md §Flag reference](../../transfor
 - `--max-memory <gb>` — skip if param estimate exceeds; e.g. `16` on a 24 GB GPU leaves headroom for activations
 - `--phases 1 2 3` — restrict (P4 is slowest; restrict when debugging P1 forward parity)
 - `--dry-run` — see above; always first
-- `--no-hf-reference` / `--no-ht-reference` — skip HF / HT comparison (faster, lower confidence)
+- `--no-hf-reference` — skip the HF reference comparison (faster, lower confidence; the run can only reach `STATUS_PROVISIONAL`, never verified)
 - `--reverify` — re-test `status==1`
 - `--retry-failed` — re-test `status==3` (read existing `note` first)
 
@@ -52,9 +52,10 @@ Hard thresholds (`_MIN_PHASE_SCORES` in `verify_models.py`):
 | 1 | 100% | — | `STATUS_FAILED` |
 | 2 | 75% | `logits_equivalence`, `loss_equivalence` | `STATUS_FAILED` |
 | 3 | 75% | `logits_equivalence`, `loss_equivalence` | `STATUS_FAILED` |
-| 4 | 50% | — | **Non-gating** — adds `"low text quality"` to `note`; never fails. |
+| 4 | ~54.5% (`p4_pass_threshold()`, derived from the judge bake-off noise floor — not a fixed number) | — | **Non-gating** — adds `"low text quality"` to `note`; never fails. |
 | 7 | 75% | `multimodal_forward` | `STATUS_FAILED`. NULL = fail. |
-| 8 | 75% | `audio_forward` | `STATUS_FAILED`. NULL = fail. |
+| 8 | 75% | `audio_forward`, `audio_text_forward` | `STATUS_FAILED`. NULL = fail. |
+| 9 | 75% | `vision_forward`, `vision_cache` | `STATUS_FAILED`. NULL = fail. |
 
 `STATUS_VERIFIED` means hard gates passed. `note` carries quality flags or failure details.
 

.github/workflows/checks.yml

@@ -436,6 +436,7 @@ jobs:
           # - "Grokking_Demo"
           - "Head_Detector_Demo"
           # - "Interactive_Neuroscope"
+          - "Jacobian_Lens_Coordinate_Patch_Benchmark_Demo"
           - "Main_Demo"
           # - "No_Position_Experiment"
           - "Othello_GPT"

AGENTS.md

@@ -58,7 +58,7 @@ make unit-test          # fast, no model loads
 make integration-test   # cross-component
 make acceptance-test    # end-to-end
 make docstring-test     # doctest + doctest-plus
-make notebook-test      # slow; subset run in CI
+make notebook-test      # slow, local-only; CI runs its own per-notebook matrix (notebook-checks in checks.yml), not this target
 make test-pr            # unit + docstring + acceptance + integration (PR-review surface)
 make test               # everything (long; includes benchmarks + notebooks)
 
@@ -105,7 +105,7 @@ Python: **>=3.10, <4.0**. CI tests 3.10, 3.11, 3.12. Format/type/docstring check
 ## 5. Hook naming — HT vs Bridge
 
 - **HT canonical**: uniform across architectures — `hook_embed`, `blocks.{i}.hook_resid_pre`, `blocks.{i}.attn.hook_q`, `blocks.{i}.hook_resid_post`.
-- **Bridge-native**: architecture-shaped — `blocks.{i}.hook_out`, `blocks.{i}.attn.q.hook_out`. HT aliases registered via `build_alias_to_canonical_map()` in [bridge_core.py](transformer_lens/model_bridge/bridge_core.py).
+- **Bridge-native**: architecture-shaped — `blocks.{i}.hook_out`, `blocks.{i}.attn.q.hook_out`. HT aliases come from the `hook_aliases` dicts on the generalized components and `BridgeCore`, registered at boot (not only in compat mode); `build_alias_to_canonical_map()` in [bridge_core.py](transformer_lens/model_bridge/bridge_core.py) reads them.
 
 Prefer Bridge-native names in new code. Raw-HF-forward drivers comparing against `boot_transformers` must match its load configuration (fp32, eager attention) and probe for optional features like `resid_mid` rather than assume.
 
@@ -213,9 +213,9 @@ Load-bearing pins live in [pyproject.toml](pyproject.toml):
 
 | Pin | Where | Why it matters |
 |---|---|---|
-| `transformers>=5.4.0` | `[project] dependencies` | The Bridge adapter contract is written against HF module layouts; every minor HF release can break adapter component-mappings. Bumping is a real test pass. |
+| `transformers>=5.9.0` | `[project] dependencies` | The Bridge adapter contract is written against HF module layouts; every minor HF release can break adapter component-mappings. Bumping is a real test pass. |
 | `torch>=2.6` | `[project] dependencies` | Hook system relies on PyTorch's forward / backward hook semantics; major torch bumps occasionally change ordering. |
-| `accelerate>=0.23.0` | `[project] dependencies` | Required for Llama-family loading. |
+| `accelerate>=1.1.0` | `[project] dependencies` | Required for Llama-family loading and `device_map` disk offload (`align_module_device`). |
 | `numpy>=1.24` / `>=1.26` | `[project] dependencies` (python-version-conditional) | Doctest float formatting can drift across NumPy versions. |
 | `isort==5.8.0` | `[dependency-groups] dev` (exact) | Format check pins to exactly this version; a bump flips the formatting of every file. |
 

CLAUDE.md

@@ -23,7 +23,7 @@
 
 - [AGENTS.md §10](AGENTS.md#10-hard-rules) — hard rules; load-bearing.
 - [AGENTS.md §2](AGENTS.md#2-two-systems-live-in-this-repo) — the single model system (the `Hooked*` model classes were removed in 4.0).
-- [tests/QUARANTINES.md](tests/QUARANTINES.md) — check before debugging any failing test. The macOS-arm64 KV-cache skip is the most common time-sink.
+- [tests/QUARANTINES.md](tests/QUARANTINES.md) — check before debugging any failing test. A skip with a matching reason is not your bug; a failure not on that list is real.
 - [debugging_numerical_divergence.md](docs/source/content/debugging_numerical_divergence.md) — Bridge-vs-HF logit drift bisection.
 - [compatibility_mode.md](docs/source/content/compatibility_mode.md) — `bridge.enable_compatibility_mode()` contract; read before adding tests that use it.
 - [devtools/adapter_builder/](devtools/adapter_builder/README.md) — autonomous adapter builder (contributor tooling; agent-teams mode needs Max, solo mode runs on any tier); manual path is `/add-model-support`.

README.md

@@ -177,14 +177,14 @@ cover:
   `dt_proj`, `out_proj` for Mamba-1; `in_proj`, `conv1d`, `inner_norm`, `out_proj` for
   Mamba-2)
 * Stateful generation with cache-aware decode steps
-* The `compute_effective_attention` utility (in
-  `transformer_lens.model_bridge.supported_architectures.mamba2`) that materializes
-  Mamba-2's SSD-derived attention matrix for comparison with transformer attention
-  patterns
+* `ActivationCache.compute_ssm_effective_attention(layer=...)`, which materializes the
+  SSD-derived attention matrix of a Mamba-2 (or any other SSM / hybrid) layer for
+  comparison with transformer attention patterns; the older
+  `mamba2.compute_effective_attention` remains only as a deprecated alias
 
 Verification lives in the integration tests at
 `tests/integration/model_bridge/test_mamba_adapter.py` and
-`tests/integration/model_bridge/test_mamba2_adapter.py` (81 tests total), and the
+`tests/integration/model_bridge/test_mamba2_adapter.py`, and the
 `verify_models` benchmark suite now covers the SSM and hybrid families. Mamba-1,
 Mamba-2, gated-delta-net (Qwen3.5 / Qwen3-Next), NemotronH, and GraniteMoeHybrid
 all declare `applicable_phases = [1, 2, 3, 4]`, so their forward parity (P1, vs raw

demos/Exploratory_Analysis_Demo.ipynb

@@ -562,7 +562,7 @@
     "\n",
     "The learned translation is folded to `model.unembed.b_U`, a bias added to the logits (note that GPT-2 is not trained with an existing `b_U`). This roughly represents unigram statistics. But we can ignore this because each prompt occurs twice with names in the opposite order, so this perfectly cancels out. \n",
     "\n",
-    "Note that rather than using layernorm scaling we could just study cache[\"ln_final.hook_normalised\"]\n",
+    "Note that rather than using layernorm scaling we could just study cache[\"ln_final.hook_normalized\"]\n",
     "\n",
     "</details>"
    ]
@@ -1815,7 +1815,7 @@
     "\n",
     "Our hook function acts on the attention pattern activation. This has the name\n",
     "\"blocks.{layer}.{layer_type}.hook_{activation_name}\" in general, here it's\n",
-    "\"blocks.{layer}.attn.hook_attn\". And it has shape [batch, head_index, query_pos, token_pos]. Our\n",
+    "\"blocks.{layer}.attn.hook_pattern\". And it has shape [batch, head_index, query_pos, token_pos]. Our\n",
     "hook function takes in the attention pattern activation, calculates the score for the relevant type\n",
     "of head, and write it to an external cache.\n",
     "\n",
@@ -1936,7 +1936,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "The above suggests that it would be a useful bit of infrastructure to have a \"wiki\" for the heads of a model, giving their scores according to some metrics re head functions, like the ones we've seen here. TransformerLens makes this easy to make, as just changing the name input to `HookedTransformer.from_pretrained` gives a different model but in the same architecture, so the same code should work. If you want to make this, I'd love to see it! \n",
+    "The above suggests that it would be a useful bit of infrastructure to have a \"wiki\" for the heads of a model, giving their scores according to some metrics re head functions, like the ones we've seen here. TransformerLens makes this easy to make, as just changing the name input to `TransformerBridge.boot_transformers` gives a different model but in the same architecture, so the same code should work. If you want to make this, I'd love to see it! \n",
     "\n",
     "As a proof of concept, [I made a mosaic of all induction heads across the 40 models then in TransformerLens](https://www.neelnanda.io/mosaic).\n",
     "\n",
@@ -2094,7 +2094,7 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "transformer-lens",
+   "display_name": "transformer-lens (3.12.12)",
    "language": "python",
    "name": "python3"
   },

demos/Inspect_Bridge_Demo.ipynb

@@ -8,7 +8,7 @@
     "# Inspect Driver — Demo\n",
     "\n",
     "Turn a model served through [`inspect_ai`](https://inspect.aisi.org.uk) into a\n",
-    "TransformerLens `HookedTransformer`: `run_with_cache`, named hook points, and\n",
+    "TransformerLens `TransformerBridge`: `run_with_cache`, named hook points, and\n",
     "interventions all work over the Inspect boundary.\n",
     "\n",
     "`boot_inspect` boots the model behind an `inspect_ai` provider and wraps it in a\n",
@@ -679,7 +679,15 @@
     }
    },
    "outputs": [],
-   "source": "# (9) Tool-aware generation: for a tool-capable instruct model the provider renders tool\n# schemas into the chat template and parses tool calls from the output (so Inspect agent\n# loops run). gpt2 has no tool template, so here we just show the tool-call parser:\nfrom transformer_lens.model_bridge.sources.inspect.transformers_provider import _parse_tool_calls\n\ncalls = _parse_tool_calls('<tool_call>{\"name\": \"add\", \"arguments\": {\"a\": 2, \"b\": 2}}</tool_call>')\nprint(\"parsed tool calls:\", [(c.function, c.arguments) for c in calls])"
+   "source": [
+    "# (9) Tool-aware generation: for a tool-capable instruct model the provider renders tool\n",
+    "# schemas into the chat template and parses tool calls from the output (so Inspect agent\n",
+    "# loops run). gpt2 has no tool template, so here we just show the tool-call parser:\n",
+    "from transformer_lens.model_bridge.sources.inspect.transformers_provider import _parse_tool_calls\n",
+    "\n",
+    "calls = _parse_tool_calls('<tool_call>{\"name\": \"add\", \"arguments\": {\"a\": 2, \"b\": 2}}</tool_call>')\n",
+    "print(\"parsed tool calls:\", [(c.function, c.arguments) for c in calls])"
+   ]
   },
   {
    "cell_type": "markdown",
@@ -706,7 +714,7 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "transformer-lens",
+   "display_name": "transformer-lens (3.12.12)",
    "language": "python",
    "name": "python3"
   },
@@ -3112,4 +3120,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 5
-}
+}