Merge upstream/dev-4.x into fix/bridge-adapter-coverage

Author: emerardd

.github/workflows/checks.yml

@@ -94,8 +94,23 @@ jobs:
             ~/.cache/huggingface/hub/models--ArthurConmy--redwood_attn_2l
             ~/.cache/huggingface/hub/models--stas--tiny-random-llama-2
             ~/.cache/huggingface/hub/models--microsoft--DialoGPT-medium
-          key: ${{ runner.os }}-huggingface-models-v4
+            ~/.cache/huggingface/hub/models--hf-internal-testing--tiny-random-gpt2
+            ~/.cache/huggingface/hub/models--hf-internal-testing--tiny-random-Olmo2ForCausalLM
+            ~/.cache/huggingface/hub/models--hf-internal-testing--tiny-random-MptForCausalLM
+            ~/.cache/huggingface/hub/models--hf-internal-testing--tiny-random-Gemma2ForCausalLM
+            ~/.cache/huggingface/hub/models--hf-internal-testing--tiny-random-Gemma3ForCausalLM
+            ~/.cache/huggingface/hub/models--hf-internal-testing--tiny-random-GraniteForCausalLM
+            ~/.cache/huggingface/hub/models--hf-internal-testing--tiny-random-OlmoForCausalLM
+            ~/.cache/huggingface/hub/models--hf-internal-testing--tiny-random-OlmoeForCausalLM
+            ~/.cache/huggingface/hub/models--hf-internal-testing--tiny-random-DeepseekV3ForCausalLM
+            ~/.cache/huggingface/hub/models--katuni4ka--tiny-random-deepseek-v3
+            ~/.cache/huggingface/hub/models--trl-internal-testing--tiny-MistralForCausalLM-0.2
+            ~/.cache/huggingface/hub/models--trl-internal-testing--tiny-BloomForCausalLM
+            ~/.cache/huggingface/hub/models--trl-internal-testing--tiny-Qwen2ForCausalLM-2.5
+          key: ${{ runner.os }}-huggingface-models-v6
           restore-keys: |
+            ${{ runner.os }}-huggingface-models-v5
+            ${{ runner.os }}-huggingface-models-v4
             ${{ runner.os }}-huggingface-models-v2
             ${{ runner.os }}-huggingface-models-v1
             ${{ runner.os }}-huggingface-models
@@ -137,7 +152,22 @@ jobs:
         with:
           path: |
             ~/.cache/huggingface/hub/models--roneneldan--TinyStories-1M*
-          key: ${{ runner.os }}-huggingface-models-mps-v1
+            ~/.cache/huggingface/hub/models--hf-internal-testing--tiny-random-gpt2
+            ~/.cache/huggingface/hub/models--hf-internal-testing--tiny-random-Olmo2ForCausalLM
+            ~/.cache/huggingface/hub/models--hf-internal-testing--tiny-random-MptForCausalLM
+            ~/.cache/huggingface/hub/models--hf-internal-testing--tiny-random-Gemma2ForCausalLM
+            ~/.cache/huggingface/hub/models--hf-internal-testing--tiny-random-Gemma3ForCausalLM
+            ~/.cache/huggingface/hub/models--hf-internal-testing--tiny-random-GraniteForCausalLM
+            ~/.cache/huggingface/hub/models--hf-internal-testing--tiny-random-OlmoForCausalLM
+            ~/.cache/huggingface/hub/models--hf-internal-testing--tiny-random-OlmoeForCausalLM
+            ~/.cache/huggingface/hub/models--hf-internal-testing--tiny-random-DeepseekV3ForCausalLM
+            ~/.cache/huggingface/hub/models--katuni4ka--tiny-random-deepseek-v3
+            ~/.cache/huggingface/hub/models--trl-internal-testing--tiny-MistralForCausalLM-0.2
+            ~/.cache/huggingface/hub/models--trl-internal-testing--tiny-BloomForCausalLM
+            ~/.cache/huggingface/hub/models--trl-internal-testing--tiny-Qwen2ForCausalLM-2.5
+          key: ${{ runner.os }}-huggingface-models-mps-v2
+          restore-keys: |
+            ${{ runner.os }}-huggingface-models-mps-v1
       - name: Install dependencies
         run: |
           uv lock --check
@@ -286,8 +316,23 @@ jobs:
             ~/.cache/huggingface/hub/models--ArthurConmy--redwood_attn_2l
             ~/.cache/huggingface/hub/models--stas--tiny-random-llama-2
             ~/.cache/huggingface/hub/models--microsoft--DialoGPT-medium
-          key: ${{ runner.os }}-huggingface-models-v4
+            ~/.cache/huggingface/hub/models--hf-internal-testing--tiny-random-gpt2
+            ~/.cache/huggingface/hub/models--hf-internal-testing--tiny-random-Olmo2ForCausalLM
+            ~/.cache/huggingface/hub/models--hf-internal-testing--tiny-random-MptForCausalLM
+            ~/.cache/huggingface/hub/models--hf-internal-testing--tiny-random-Gemma2ForCausalLM
+            ~/.cache/huggingface/hub/models--hf-internal-testing--tiny-random-Gemma3ForCausalLM
+            ~/.cache/huggingface/hub/models--hf-internal-testing--tiny-random-GraniteForCausalLM
+            ~/.cache/huggingface/hub/models--hf-internal-testing--tiny-random-OlmoForCausalLM
+            ~/.cache/huggingface/hub/models--hf-internal-testing--tiny-random-OlmoeForCausalLM
+            ~/.cache/huggingface/hub/models--hf-internal-testing--tiny-random-DeepseekV3ForCausalLM
+            ~/.cache/huggingface/hub/models--katuni4ka--tiny-random-deepseek-v3
+            ~/.cache/huggingface/hub/models--trl-internal-testing--tiny-MistralForCausalLM-0.2
+            ~/.cache/huggingface/hub/models--trl-internal-testing--tiny-BloomForCausalLM
+            ~/.cache/huggingface/hub/models--trl-internal-testing--tiny-Qwen2ForCausalLM-2.5
+          key: ${{ runner.os }}-huggingface-models-v6
           restore-keys: |
+            ${{ runner.os }}-huggingface-models-v5
+            ${{ runner.os }}-huggingface-models-v4
             ${{ runner.os }}-huggingface-models-v2
             ${{ runner.os }}-huggingface-models-v1
             ${{ runner.os }}-huggingface-models
@@ -336,11 +381,10 @@ jobs:
         with:
           path: |
             ~/.cache/huggingface/hub/models--gpt2
-          key: ${{ runner.os }}-huggingface-models-v4
-          restore-keys: |
-            ${{ runner.os }}-huggingface-models-v2
-            ${{ runner.os }}-huggingface-models-v1
-            ${{ runner.os }}-huggingface-models
+          # Distinct key: this job needs only gpt2, and sharing the full-suite
+          # key let whichever job saved first pin the cache contents for all
+          # three (a one-path save then serves the suites that need the rest).
+          key: ${{ runner.os }}-huggingface-models-benchmark-v1
       - name: Install dependencies
         run: |
           uv lock --check

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

@@ -276,6 +276,32 @@ Standard positional embeddings + LayerNorm + standard MLP + combined QKV. Used b
 
 Similar to Llama-like but with `MoEBridge` replacing the MLP. Used by: Mixtral, GraniteMoE, OLMoE.
 
+**Interleaved / dense-prefix stacks.** When the config can make some layers a plain
+gated MLP under the same attribute name (`first_k_dense_replace`,
+`mlp_only_layers`/`decoder_sparse_step`, `moe_layer_start_index`,
+`num_dense_layers`, `mlp_layer_types`), also declare the dense projections as
+optional submodules named `dense_gate` / `dense_in` / `dense_out`:
+
+```python
+"mlp": MoEBridge(
+    name="mlp",
+    config=self.cfg,
+    submodules={
+        "gate": MoERouterBridge(name="gate", optional=True),
+        "dense_gate": LinearBridge(name="gate_proj", optional=True),
+        "dense_in": LinearBridge(name="up_proj", optional=True),
+        "dense_out": LinearBridge(name="down_proj", optional=True),
+    },
+),
+```
+
+`MoEBridge` detects that binding per layer and gives the dense layers the
+gated-MLP hook set (`hook_pre`, `hook_pre_linear`, `hook_post`) and weight
+accessors (`W_gate`/`W_in`/`W_out`), while sparse layers keep MoE semantics.
+Without the declaration those layers report `d_model` boundary tensors under
+neuron-hook names (#1645). The keys are `dense_*` rather than `gate`/`in`/`out`
+because `gate` already means the *router* on the sparse layers of the same model.
+
 ### Pattern 4: Multimodal
 
 Extends a text-only pattern with vision encoder and projection bridges. Used by: LLaVA, LLaVA-Next, Gemma3 Multimodal.

docs/source/content/compatibility_mode.md

@@ -64,6 +64,21 @@ After `enable_compatibility_mode()`, these HT hook names fire on the **pre-norm
 - **Post-norm architectures** (OLMo 2, BERT-style) read the **post-attention residual** instead, because the norm semantically lives elsewhere in the block.
 - **MLA blocks** (DeepSeek V2 / V3 / R1) do **not** expose the split-qkv aliases — MLA's compressed K/V doesn't have a clean split.
 
+On post-norm architectures (Gemma 2/3's `ln1_post`/`ln2_post`, OLMo 2/3's `ln1`/`ln2`),
+`blocks.{i}.hook_attn_out` / `hook_mlp_out` fire **after** the post-sublayer norm, so
+that they capture the tensor added to the residual stream and the identities
+`resid_pre + attn_out == resid_mid` and `resid_mid + mlp_out == resid_post` hold
+([issue #1648](https://github.com/TransformerLensOrg/TransformerLens/issues/1648)).
+One consequence for head-level direct logit attribution: per-head contributions from
+`attn.hook_result` no longer sum to `hook_attn_out`, because the norm sits between
+them. This is inherent to post-norm — decompose heads on the pre-norm side
+(`attn.hook_out` for the raw module output) or attribute through the norm explicitly.
+The Granite family has the same raw-vs-contribution split for a different reason:
+HF scales each sublayer output by `residual_multiplier` before the residual add, so
+`hook_attn_out` / `hook_mlp_out` fire on the scaled contribution while
+`attn.hook_out` / `mlp.hook_out` stay raw — `attn.hook_result` sums to the raw
+output, off from the contribution by the multiplier.
+
 An adapter author for a new post-norm or MLA-style architecture must handle these carve-outs in `setup_hook_compatibility`. The Gemma1/Gemma2 adapters are exemplars of when **not** to override `setup_hook_compatibility` — `GemmaTextScaledWordEmbedding` already scales internally, so any added `hook_conversion` would double-scale `embed.hook_out`.
 
 ## The four-quadrant test matrix

docs/source/content/jacobian_lens_fitting.md

@@ -233,6 +233,94 @@ To propose a short-name entry in TransformerLens, open a pull request that adds
 published file to `transformer_lens/tools/analysis/jacobian_lens_registry.json` and
 include the fitting provenance and validation results.
 
+## Sparse decomposition (J-space coordinates)
+
+A fitted lens also decomposes an activation into the concepts it is *disposed to say*.
+`JacobianLens.decompose` writes an activation `x` at layer ℓ as a sparse **nonnegative**
+combination of J-lens vectors `v_t = J_ℓ^T W_U[:, t]` (one direction per vocabulary token),
+selected greedily. `k` is an **upper bound**, not a target: selection stops early once no
+unselected vector is materially positively correlated with the residual (under nonnegativity a
+negatively-correlated vector cannot reduce it), so fewer than `k` vectors may be selected and
+fewer still may be numerically active.
+
+```python
+from transformer_lens.model_bridge import TransformerBridge
+from transformer_lens.tools.analysis import JacobianLens
+
+model = TransformerBridge.boot_transformers("gpt2", device="cpu")
+lens = JacobianLens.from_pretrained("gpt2-small", model=model)
+
+# decompose the activation at a prompt position ...
+result = lens.decompose(model, "The Eiffel Tower is in the city of", layer=6, position=-1, k=8)
+# ... or a raw [d_model] activation you already have (leave position=None):
+#   result = lens.decompose(model, activation, layer=6, k=8)
+
+tokens = [model.to_string(int(t)) for t in result.support]  # the (up to k) *active* J-lens vectors
+coordinates = result.coordinates                            # their nonnegative coefficients
+```
+
+The result exposes **two supports**, because the paper uses two inconsistent operationalizations
+(a main-text sparse nonnegative reconstruction and an appendix projection onto a selected span):
+
+- `support` -- the numerically **active** vectors: the selected vectors whose contribution
+  `coordinates[i] * ||v_t||` is a materially nonzero fraction of `||x||`. `coordinates` is aligned
+  with `support`, and `reconstruction = sum(coordinates * v_t)` over `support`.
+- `selected_support` -- **every** greedily selected vector, including any whose coordinate the
+  nonnegativity constraint drove to zero. It defines the span for `j_space_component`. Hence
+  `len(support) <= len(selected_support) <= k`.
+
+So two vector outputs also need not coincide:
+
+- `reconstruction` -- the nonnegative combination over the active `support`.
+- `j_space_component` (the *J-space component*) -- the orthogonal projection of the activation onto
+  the span of `selected_support` -- with `non_j_space_component = x - j_space_component`, the
+  residual the interventions leave unchanged.
+
+For the default exact NNLS re-solve the `reconstruction` equals the projection onto the *active*
+support (KKT stationarity), so it differs from `j_space_component` exactly when a selected vector
+has a zero coordinate (the projection then uses a strictly larger span).
+
+```
+x  in R^d_model
+   |-- decompose(x, layer, k)
+        |-- support / coordinates  a_t >= 0      (active vectors; reconstruction = sum a_t v_t)
+        |-- selected_support       S             (all selected vectors; defines the span below)
+        |-- j_space_component      Pi_S x        (orthogonal projection onto span of selected v_t)
+        \-- non_j_space_component  x - Pi_S x     (orthogonal to the selected vectors)
+```
+
+Two algorithms are available via `algorithm=`. The default,
+`"nonnegative_orthogonal_matching_pursuit"`, solves a nonnegative least-squares (NNLS) problem
+over the selected atoms in float64 after each step. It checks the result against the KKT
+conditions and raises `RuntimeError` if the check fails. `"gradient_pursuit"` skips that solve
+and uses the directional update from Blumensath & Davies (2008), matching the update used in
+the paper; its projected step is accepted only when it does not increase the residual. The two
+algorithms share the same greedy selection *rule* but, because their coefficient residuals
+differ, may select different vectors at later steps and so return a different `support` and
+`reconstruction`.
+
+### Interpreting the numbers honestly
+
+The quantitative findings below are from Gurnee et al. (2026) and were measured on **closed
+Anthropic models** (Sonnet / Haiku / Opus); on open-weight models the *shape* may hold but the
+exact values will not necessarily transfer.
+
+- The decomposition is **not** a top-k logit-lens readout: because the J-lens vectors are
+  overcomplete and non-orthogonal, it gives "a different (and typically less redundant) set of
+  active concepts than simply taking the top-k by inner product."
+- The J-space is a **small fraction** of the activation: the paper's span projection (the
+  `selected_support` operationalization here) "never [exceeds] more than 10%" of total activation
+  variance, and for concept vectors carries "a median of only 6-7% ... the remaining ~93% lying
+  outside the J-space." Those figures are the paper's own measurements on its models; do not read
+  them off this implementation's `j_space_component` without matching the operationalization.
+- `k` defaults to 25 because the paper "typically choose[s] it to be no more than 25, which we
+  empirically observed to be the number of J-lens vectors that are meaningfully active at a given
+  time." Here `k` is an **upper bound**: `support` returns *at most* `k` active vectors (often
+  fewer), never `k` padded with zero-coefficient slots.
+
+The full-vocabulary dictionary is cached on the model's device and is vocabulary-sized
+(gigabytes for large models); release it with `lens.clear_device_cache()`.
+
 ## Importing an existing lens
 
 `JacobianLens.load()` accepts two file schemas: the standard artifact format (four