Updating documentation and closing some small gaps from 4.x

devtools/adapter_builder/docs/adapter-specification.md

@@ -249,7 +249,7 @@ def setup_component_testing(self, hf_model, bridge_model=None):
 
 ### `preprocess_weights(state_dict)`
 
-Apply architecture-specific weight transformations before standard processing. Example: Gemma scales embeddings by `sqrt(d_model)`.
+Apply architecture-specific weight transformations before standard processing. Example: Cohere folds `logit_scale` into `unembed.weight`. Do not fold runtime-applied scales such as Gemma's `sqrt(d_model)` embedding scale; HF applies it itself and a fold double-scales.
 
 ### `prepare_loading(model_name, model_kwargs)`
 

docs/source/content/adapter_development/adapter-creation-guide.md

@@ -177,7 +177,7 @@ For combined-QKV models (GPT-2 style), see `gpt2.py`'s `QKVSplitRearrangeConvers
 Implement only the ones you need:
 
 - **`setup_component_testing(hf_model, bridge_model=None)`** — required for RoPE models, to wire the rotary embedding instance through to the attention bridges. Skip for models with standard positional embeddings.
-- **`preprocess_weights(state_dict)`** — for arch-specific weight transforms before standard processing (e.g., Gemma scales embeddings by `sqrt(d_model)`).
+- **`preprocess_weights(state_dict)`** — for arch-specific weight transforms before standard processing (e.g., Cohere folds `logit_scale` into `unembed.weight`; not Gemma's `sqrt(d_model)` embedding scale, which HF applies at runtime).
 - **`prepare_loading(model_name, model_kwargs)`** — patch HF model classes before `from_pretrained()`.
 - **`prepare_model(hf_model)`** — post-load fixups before bridge creation.
 

docs/source/content/adapter_development/adapter-specification.md

@@ -252,7 +252,7 @@ def setup_component_testing(self, hf_model, bridge_model=None):
 
 ### `preprocess_weights(state_dict)`
 
-Apply architecture-specific weight transformations before standard processing. Example: Gemma scales embeddings by `sqrt(d_model)`.
+Apply architecture-specific weight transformations before standard processing. Example: Cohere folds `logit_scale` into `unembed.weight`. Do not fold runtime-applied scales such as Gemma's `sqrt(d_model)` embedding scale; HF applies it itself and a fold double-scales.
 
 ### `prepare_loading(model_name, model_kwargs)`
 

docs/source/content/analysis_tools.md

@@ -34,12 +34,12 @@ available hooks, and the meaning of the selected tensor axes.
 
 | Tool | Requirements to check | Main compute or memory cost |
 |---|---|---|
-| DLA | Targets TransformerBridge in 4.0. Enable compatibility mode and use the standard attention/MLP residual decomposition; Mamba/SSM/Mixer/LinearAttention hybrid layouts are rejected, as are bridges whose readout is not `ln_final -> unembed` on the last block (multi-stream AltUp blocks, a head transform such as BERT's `mlm_head`). When DLA runs the prompt, the model and all submodules must be in evaluation mode, or it raises `ValueError` before the forward; a supplied cache is accepted as is. | One cached forward pass, or reuse of a suitable cache. Head decomposition can require additional per-head results. |
+| DLA | Targets TransformerBridge in 4.0. Enable compatibility mode. Hybrid and SSM stacks are supported, with each sequence mixer's output as its own component (`unit="component"` raises `ValueError` before any forward for a block whose mixer output `decompose_resid` would not emit; a block with no `hook_mlp_out`, such as GraniteMoeHybrid with experts, contributes `resid_post - resid_mid` as its MLP row); `unit="layer"` raises `ValueError` before any forward when a sequential block has no `resid_mid` (mixer-only blocks such as Mamba, NemotronH and Zamba2). Bridges whose readout is not `ln_final -> unembed` on the last block (multi-stream AltUp blocks, a head transform such as BERT's `mlm_head`) are rejected. When DLA runs the prompt, the model and all submodules must be in evaluation mode, or it raises `ValueError` before the forward; a supplied cache is accepted as is. | One cached forward pass, or reuse of a suitable cache. Head decomposition can require additional per-head results. |
 | Activation Patching | Any bridge whose activation names resolve (`resid_pre`, `attn_out`, `mlp_out`, `z`, `q`/`k`/`v`, `pattern`). The metric returns a scalar by default and a `[batch]` tensor under `reduce="mean"` or `"none"`; several metrics can be passed as a mapping and share one sweep. Run the model in evaluation mode. | One forward of the corrupted batch per swept cell (plus one for `baseline=True`); per-example storage is `cells × batch` floats per metric. Statistics are computed afterwards from the stored values and add no forwards. |
 | Attribution Patching | Targets TransformerBridge. Clean/corrupt inputs must be `[batch, seq]` token tensors with matching shapes and aligned positions. The required embedding, head, and MLP hook aliases must exist. Edge granularity additionally requires attention bridges that support per-head results and a pre-Q/K/V residual fork. On GQA/MQA models, where any attention layer's K/V head count (`per_layer_num_key_value_heads` where set, else `n_key_value_heads`) differs from `n_heads`, edge granularity and faithfulness raise `ValueError` before any forward; node granularity remains available. Edge granularity and faithfulness are unsupported on post-norm blocks (such as Gemma 2/3, OLMo 2/3 and GLM-4), scaled-residual blocks (such as Granite) and parallel attention/SSM blocks (such as Falcon-H1), where per-head outputs are normalized, scaled or joined by an unmodeled mixer before the residual add; both raise there, and node granularity remains available. Edge sweeps temporarily enable the required writer and reader hooks, then restore the caller's hook flags. Run the model and all nested stochastic modules in evaluation mode. The metric must return a differentiable scalar, and the current backward-cache path requires an active autograd graph; fully frozen models are not yet supported. | Two forwards and one backward **per prompt pair**; activation and gradient caches. For a uniform-head decoder, the number of edge scores grows as `O(seq_len × (n_layers × n_heads)²)`, while the cache also gains per-head `[batch, seq, n_heads, d_model]` tensors. Pair scores are averaged across the batch. |
 | Direct Path Patching | Requires a TransformerBridge whose attention layers have the pre-Q/K/V residual fork that `use_split_qkv_input` enables (the joint-QKV and position-embedding attention bridges, e.g. GPT-2, Pythia, Llama, Mistral, Qwen2); attention without it, such as GPT-J, GPT-Neo, OPT, BERT or a bridge built with `TransformerBridge.boot_native`, raises `NotImplementedError`. The supplied caches must hold the source layer's `hook_z`. The function turns the fork flag on for the sweep and restores the caller's `use_split_qkv_input` and `use_attn_in`; LayerNorm need not be folded. Q destinations use query heads; K/V destinations use native KV heads, shared by groups of query heads on GQA/MQA. K/V sweeps require a uniform KV count across destination layers (`per_layer_num_key_value_heads` where set, else `n_key_value_heads` or `n_heads`), and raise `ValueError` before any forward if those counts differ. Every block from the source layer on must have an `attn` component; SSM, mixer, linear-attention and AltUp blocks raise `ValueError` before any forward. Source blocks must add the head outputs to the residual unchanged: post-norm blocks (such as Gemma 2/3, OLMo 2/3 and GLM-4), scaled-residual blocks (such as Granite) and parallel attention/SSM blocks (such as Falcon-H1) raise `ValueError` before any forward, since a head's residual write there is not `hook_z @ W_O`. Run the model and all submodules in evaluation mode; both entry points raise `ValueError` before any forward otherwise. | Clean/corrupt caching, then one forward of the corrupted batch per later destination query head or KV group. While the fork is on, every attention layer normalizes `n_heads` Q copies and `n_kv_heads` copies each for K and V, about `n_heads + 2 × n_kv_heads` times the usual attention-input activation memory. Sweeping all source heads adds another sweep dimension. |
-| Logit Lens | Targets TransformerBridge, raw or compatibility mode. Needs a decoder with an `ln_final` component (or no final norm) and `W_U`; encoder-decoder and encoder-only bridges are rejected, as are quantized unembeddings and bridges whose readout is not `ln_final -> unembed` on the last block (multi-stream AltUp blocks, a head transform such as BERT's `mlm_head`; a prompt run with `apply_ln=True` is also checked against the model's own output, a supplied cache only structurally). Hybrid SSM stacks work with `incl_mid=False`. Run the model in evaluation mode. | One cached forward pass over the residual hooks, or reuse of a cache. Logits are produced in chunks of `chunk_size` rows; full-vocabulary output is materialized only when no `targets`, `top_k` or `vocab` selector is given. |
-| Jacobian Lens | Requires a fresh, causal decoder-only Bridge with raw HF weights, without compatibility mode or weight processing. Validate the lens against the model. Run the model and all submodules in evaluation mode; fitting and every call that runs the model raise `ValueError` otherwise. | Loading an existing artifact avoids fitting. The ordinary fitting estimator uses one forward and `ceil(d_model / dim_batch)` backwards per prompt; larger batches increase memory. |
+| Logit Lens | Targets TransformerBridge, raw or compatibility mode. Needs a decoder with an `ln_final` component (or no final norm) and `W_U`; encoder-decoder and encoder-only bridges are rejected, as are quantized unembeddings and bridges whose readout is not `ln_final -> unembed` on the last block (multi-stream AltUp blocks, a head transform such as BERT's `mlm_head`; a prompt run with `apply_ln=True` is also checked against the model's own output, a supplied cache only structurally). `incl_mid=True` needs a `resid_mid` hook on every sequential block, so it raises `ValueError` on mixer-only blocks (Mamba, NemotronH, Zamba2) and works on hybrids whose blocks all have one (Qwen3.5, Jamba, RecurrentGemma). Run the model in evaluation mode. | One cached forward pass over the residual hooks, or reuse of a cache. Logits are produced in chunks of `chunk_size` rows; full-vocabulary output is materialized only when no `targets`, `top_k` or `vocab` selector is given. |
+| Jacobian Lens | Requires a fresh, causal decoder-only Bridge with raw HF weights, without compatibility mode or weight processing; bridges whose readout is not `ln_final -> unembed` on the last block (multi-stream blocks, a head transform such as ModernBERT-decoder's `prediction_head`) are rejected. Validate the lens against the model. Run the model and all submodules in evaluation mode; fitting and every call that runs the model raise `ValueError` otherwise. | Loading an existing artifact avoids fitting. The ordinary fitting estimator uses one forward and `ceil(d_model / dim_batch)` backwards per prompt; larger batches increase memory. |
 | Interchange Interventions | Any model with the hook API (TransformerBridge or `HookedRootModule`). Activations at the hook must be `[batch, d_model]`, or `[batch, pos, d_model]` with one position index per example; base and source rows are paired. | One forward pass per `batch_size` sources to collect activations, then one per `batch_size` bases per intervention with `intervened_logits`; both run without gradients, so memory is bounded by `batch_size`. `run_with_intervention` is a single differentiable pass over all the bases it is given. DAS adds a forward and backward pass per batch per epoch; `micro_batch_size` bounds its memory. |
 | Projection Kernel | The numerical API accepts finite, real floating-point matrices via orthonormal bases. The attention-head wrapper requires Bridge weights with compatible dimensions and ranks. | Basis extraction uses SVD. All-head comparisons allocate basis stacks and a pairwise score grid; they can be large despite requiring no forward pass. |
 
@@ -89,6 +89,12 @@ it is exact for twelve or fewer prompts and a Monte Carlo estimate above that.
 Neither says anything about prompts outside the batch, and neither changes what a
 single cell establishes: the measured effect of one intervention.
 
+On GQA/MQA models the `*_every` helpers pad the `k`/`v` sweeps to `n_heads`. The
+padded cells are zero in `values` and `mean()` (matching the tensor return),
+`NaN` in every statistic, absent from `to_dataframe()`, and `False` in
+`result.valid`; select `p[..., result.valid]` (the mask spans the swept axes, so the
+ellipsis also skips a leading metric axis) before any multiple-comparison correction.
+
 ### Attribution Patching: screen candidates, then measure interventions
 
 At node granularity, the score uses

docs/source/content/contributing.md

@@ -304,7 +304,7 @@ Two test layers:
 ### Common adapter gotchas
 
 - **HF raw config attributes are invisible to TL-side consumers unless explicitly propagated to `self.cfg`.** Walk the HF `config.json` and mirror any non-standard knobs (`sliding_window`, `layer_types`, …) onto `self.cfg` so weight processing and forward passes can see them. Gemma's softcaps are already renamed by the bridge builder (`cfg.output_logits_soft_cap`, `cfg.attn_scores_soft_cap`), and `query_pre_attn_scalar` needs nothing — the wrapped HF attention applies it.
-- **Some config attrs need both surface-on-cfg AND fold-into-weight** via a `preprocess_weights()` override. The trigger: a numerical operation HF's forward applies natively must also be baked into the raw weights, or `bridge.enable_compatibility_mode()` (which calls `process_weights` on raw weights) produces wrong results. Concrete examples in-tree: Cohere `logit_scale` → `unembed.weight`; Gemma embedding scale (`√d_model`) → `embed.weight`. Skip the fold and Phase 3 / Phase 4 of `verify_models` will silently degrade.
+- **Some config attrs need both surface-on-cfg AND fold-into-weight** via a `preprocess_weights()` override. The trigger: a numerical operation HF's forward applies natively must also be baked into the raw weights, or `bridge.enable_compatibility_mode()` (which calls `process_weights` on raw weights) produces wrong results. Concrete examples in-tree: Cohere `logit_scale` → `unembed.weight`. Gemma is the counter-example: HF applies its `√d_model` embedding scale at runtime on both raw and compat paths and `bridge.embed` already captures it, so no Gemma adapter overrides `preprocess_weights`; folding would double-scale. Skip a needed fold and Phase 3 / Phase 4 of `verify_models` will silently degrade.
 - **Tokenizer policy is per-model, not per-architecture.** Sibling models in the same family routinely differ — the chat-instruct variant may prepend BOS where the base does not, padding side can flip, EOS handling can differ. It's worth re-checking `default_prepend_bos`, padding side, and EOS handling against the specific target rather than copying them from a starter adapter. `tokenizer_config.json` is not always reliable on its own — some architectures (Cohere is a notable example) declare `add_bos_token=False` but HF's `__call__` prepends BOS anyway. The most reliable check is to invoke the tokenizer directly:
 
   ```python

docs/source/content/jacobian_lens_fitting.md

@@ -21,7 +21,9 @@ manifest, and set of estimator options.
 ## Requirements
 
 `JacobianLens.fit` requires a freshly booted, causal decoder-only
-`TransformerBridge` with raw Hugging Face weights:
+`TransformerBridge` with raw Hugging Face weights, whose logits are
+`unembed(ln_final(...))` of the last block's single-stream output (a head
+transform such as ModernBERT-decoder's `prediction_head` is rejected):
 
 ```python
 import torch

tests/integration/model_bridge/test_forward_padding_side.py

@@ -184,3 +184,19 @@ def test_right_padded_logits_still_run_unmasked(distilgpt2_bridge) -> None:
         both_logits, _ = distilgpt2_bridge(padded, return_type="both", padding_side="right")
 
     torch.testing.assert_close(both_logits, logits, rtol=0, atol=0)
+
+
+@pytest.mark.parametrize("prepend_bos", [True, False])
+def test_right_padded_loss_keeps_a_real_trailing_eos(distilgpt2_bridge, prepend_bos) -> None:
+    """A real trailing EOS shares the pad id here, yet its transition still counts."""
+    bridge = distilgpt2_bridge
+    eos = bridge.tokenizer.eos_token
+    prompts = ["hello" + eos, "the cat sat on a mat" + eos]
+    with torch.no_grad():
+        per_token = [
+            bridge(prompt, prepend_bos=prepend_bos, return_type="loss", loss_per_token=True)[0]
+            for prompt in prompts
+        ]
+        batched = bridge(prompts, prepend_bos=prepend_bos, padding_side="right", return_type="loss")
+
+    torch.testing.assert_close(batched, torch.cat(per_token).mean())

tests/integration/model_bridge/test_gemma3n_bridge.py

@@ -66,6 +66,21 @@ def test_run_with_cache_text_only(bridge):
     assert len(cache) > 0
 
 
+def test_svd_w_in_uses_hf_ln2_scale(bridge):
+    """W_in is scaled by the per-channel factor HF's ln2 actually applies."""
+    from transformer_lens.SVDInterpreter import SVDInterpreter
+
+    interpreter = SVDInterpreter(bridge)
+    for layer in range(bridge.cfg.n_layers):
+        hf_ln2 = bridge.blocks[layer].ln2.original_component
+        x = torch.randn(8, bridge.cfg.d_model)
+        x_hat = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + hf_ln2.eps)
+        with torch.no_grad():
+            scale = (hf_ln2(x) / x_hat).mean(0)
+        expected = interpreter.params[f"blocks.{layer}.mlp.W_in"].T * scale
+        torch.testing.assert_close(interpreter._get_w_in_matrix(layer), expected)
+
+
 def test_contribution_hooks_are_post_norm_outputs(bridge):
     """hook_attn_out / hook_mlp_out carry the post-sublayer norm outputs, the tensors the
     layer actually adds (Gemma 2/3 contract), checked against an independent HF run."""
@@ -113,3 +128,9 @@ def test_half_converts_unmapped_parameters():
     assert unconverted == []
     with torch.no_grad():
         assert bridge(IDS).dtype == torch.float16
+
+
+def test_start_at_layer_is_refused_on_altup_blocks(bridge):
+    """An AltUp block carries a stack of streams, so there is no single residual to inject."""
+    with pytest.raises(NotImplementedError, match="AltUpBlockBridge"):
+        bridge(torch.zeros(1, 3, bridge.cfg.d_model), start_at_layer=1)

tests/integration/model_bridge/test_post_norm_head_decomposition.py

@@ -26,6 +26,7 @@
 )
 
 from transformer_lens.model_bridge.sources import build_bridge_from_module
+from transformer_lens.utilities.slice import Slice
 
 _TINY = dict(
     hidden_size=64,
@@ -165,6 +166,37 @@ def test_output_bias_is_mapped_through_post_norm(hf_architecture: str, config_cl
             )
 
 
+@pytest.mark.parametrize("hf_architecture,config_cls,config_kwargs", CASES)
+@pytest.mark.parametrize("neuron_slice", [5, Slice(5)], ids=["int", "Slice-int"])
+@pytest.mark.parametrize("pos_slice", [None, -1], ids=["all-pos", "last-pos"])
+@pytest.mark.parametrize("batched", [True, False], ids=["batched", "batchless"])
+@pytest.mark.parametrize("n_outs", [None, 0, 2], ids=["unprojected", "proj-1d", "proj-2d"])
+def test_integer_neuron_slice_matches_list(
+    hf_architecture: str,
+    config_cls: Any,
+    config_kwargs: dict[str, Any],
+    neuron_slice: Any,
+    pos_slice: Any,
+    batched: bool,
+    n_outs: Any,
+) -> None:
+    """An int neuron_slice is the one-element list slice with the neuron axis dropped."""
+    bridge, cache = _run(hf_architecture, config_cls, config_kwargs)
+    if not batched:
+        cache = cache.apply_slice_to_batch_dim(0)
+    d_model = bridge.cfg.d_model
+    project = None
+    if n_outs is not None:
+        shape = (d_model,) if n_outs == 0 else (d_model, n_outs)
+        project = torch.randn(shape, generator=torch.Generator().manual_seed(0))
+    kwargs = dict(layer=0, pos_slice=pos_slice, project_output_onto=project)
+    reference = cache.get_neuron_results(neuron_slice=[5], **kwargs)
+    expected = reference[..., 0] if n_outs == 0 else reference[..., 0, :]
+    result = cache.get_neuron_results(neuron_slice=neuron_slice, **kwargs)
+    assert result.shape == expected.shape
+    torch.testing.assert_close(result, expected)
+
+
 def test_half_precision_stacks_keep_model_dtype() -> None:
     bridge, cache = _run("Gemma2ForCausalLM", Gemma2Config, _TINY, dtype=torch.bfloat16)
     assert cache.stack_head_results().dtype == torch.bfloat16