Bug/bridge hook gating followups

tests/integration/model_bridge/test_attention_score_sentinel.py

@@ -7,19 +7,27 @@
 
 
 def test_gpt2_compatibility_scores_use_negative_infinity(
-    gpt2_bridge_compat, gpt2_hooked_processed
+    gpt2_bridge_compat, gpt2_goldens_processed
 ) -> None:
-    """GPT-2's direct HF mask is normalized before the compatibility hook."""
-    tokens = gpt2_hooked_processed.to_tokens("The capital of France is")
+    """GPT-2's direct HF mask is normalized before the compatibility hook.
+
+    Anchored on the frozen HookedTransformer goldens rather than a live
+    HookedTransformer, matching the rest of the compatibility suite.
+    """
+    golden = gpt2_goldens_processed
+    tokens = golden.scalars["short_prompt"]
     _, bridge_cache = gpt2_bridge_compat.run_with_cache(tokens, names_filter=[SCORES])
-    _, hooked_cache = gpt2_hooked_processed.run_with_cache(tokens, names_filter=[SCORES])
+    hooked_cache = golden.tensors("activations")
 
     bridge_scores, hooked_scores = bridge_cache[SCORES], hooked_cache[SCORES]
     causal_mask = torch.isneginf(hooked_scores)
     assert causal_mask.any()
     assert torch.isneginf(bridge_scores[causal_mask]).all()
+    # The goldens were captured on different hardware, so the unmasked scores
+    # agree to fp32 accumulation noise rather than bit-exactly. Same tolerance
+    # the sibling golden comparison uses for this hook.
     torch.testing.assert_close(
-        bridge_scores[~causal_mask], hooked_scores[~causal_mask], rtol=0, atol=0
+        bridge_scores[~causal_mask], hooked_scores[~causal_mask], rtol=1e-4, atol=1e-4
     )
 
 

tests/integration/model_bridge/test_audio_frame_entry.py

@@ -0,0 +1,100 @@
+"""Audio frame entry: run the encoder from precomputed frames.
+
+Mirrors HookedAudioEncoder.encoder_output, the audio-path analogue of
+start_at_layer. Deletion evidence for that method: without it the bridge can
+only enter at the waveform, so injecting frames means re-running the conv front
+end. start_at_layer stays refused for audio — this is a separate entry point.
+"""
+
+from __future__ import annotations
+
+import math
+
+import numpy as np
+import pytest
+import torch
+
+from transformer_lens.model_bridge.bridge import TransformerBridge
+
+MODEL = "facebook/hubert-base-ls960"
+SAMPLE_RATE = 16000
+FRAMES_HOOK = "feat_proj.hook_out"
+
+
+@pytest.fixture(scope="module")
+def audio_bridge() -> TransformerBridge:
+    return TransformerBridge.boot_transformers(MODEL, device="cpu")
+
+
+@pytest.fixture(scope="module")
+def waveform() -> torch.Tensor:
+    t = np.linspace(0, 1.0, SAMPLE_RATE, endpoint=False, dtype=np.float32)
+    return torch.tensor(0.1 * np.sin(2 * math.pi * 440.0 * t))[None, :]
+
+
+@pytest.fixture(scope="module")
+def full_run(audio_bridge, waveform):
+    last = f"blocks.{audio_bridge.cfg.n_layers - 1}.hook_out"
+    _, cache = audio_bridge.run_with_cache(
+        waveform, names_filter=[FRAMES_HOOK, "blocks.0.hook_out", last]
+    )
+    return cache, last
+
+
+def test_frame_entry_matches_the_full_waveform_run(audio_bridge, full_run):
+    """Re-entering at the frames reproduces the encoder output exactly."""
+    cache, last = full_run
+    resid = audio_bridge.encoder_output(cache[FRAMES_HOOK])
+    torch.testing.assert_close(resid, cache[last], atol=0.0, rtol=0.0)
+
+
+def test_hooks_fire_from_frame_entry(audio_bridge, full_run):
+    """Block hooks fire on the frame path, so caching composes with it."""
+    cache, last = full_run
+    wanted = {"blocks.0.hook_out", last}
+    cached, fwd_hooks, _ = audio_bridge.get_caching_hooks(names_filter=lambda name: name in wanted)
+    with audio_bridge.hooks(fwd_hooks=fwd_hooks):
+        audio_bridge.encoder_output(cache[FRAMES_HOOK])
+
+    assert set(cached) == wanted
+    for name in wanted:
+        torch.testing.assert_close(cached[name], cache[name], atol=0.0, rtol=0.0)
+
+
+def test_padding_mask_changes_the_encoding(audio_bridge, full_run):
+    """The mask is applied, not ignored."""
+    cache, last = full_run
+    frames = cache[FRAMES_HOOK]
+    mask = torch.ones(frames.shape[:2], dtype=torch.long)
+    mask[:, -10:] = 0
+
+    masked = audio_bridge.encoder_output(frames, one_zero_attention_mask=mask)
+    assert not torch.allclose(masked, cache[last])
+
+
+def test_waveform_shaped_input_is_rejected(audio_bridge, waveform):
+    """A 2D waveform is not frames; say so instead of silently mis-running."""
+    with pytest.raises(ValueError, match=r"\[batch, frames, d_model\]"):
+        audio_bridge.encoder_output(waveform)
+
+
+def test_start_at_layer_remains_refused_for_audio(audio_bridge, waveform):
+    """Frame entry is a separate API; the residual-injection guard is untouched."""
+    with pytest.raises(NotImplementedError, match="audio models"):
+        audio_bridge(waveform, start_at_layer=1)
+
+
+def test_text_models_reject_the_audio_frame_entry():
+    """Non-audio bridges have no conv-frame stage to re-enter."""
+    bridge = TransformerBridge.boot_transformers("gpt2", device="cpu")
+    with pytest.raises(NotImplementedError, match="not an audio model"):
+        bridge.encoder_output(torch.zeros(1, 4, bridge.cfg.d_model))
+
+
+def test_spectrogram_encoders_reject_the_frame_entry():
+    """AST has no conv feature extractor, so there is no frame stage to bypass."""
+    bridge = TransformerBridge.boot_transformers(
+        "MIT/ast-finetuned-audioset-10-10-0.4593", device="cpu"
+    )
+    with pytest.raises(NotImplementedError, match="convolutional front end"):
+        bridge.encoder_output(torch.zeros(1, 4, bridge.cfg.d_model))

tests/integration/model_bridge/test_bert_pooler_hook.py

@@ -0,0 +1,69 @@
+"""The BERT [CLS] pooler is observable through a named bridge hook.
+
+Mirrors HookedEncoder's BertPooler, whose hook_pooler_out carries the
+post-tanh pooled [CLS]. Deletion evidence for that component: without a named
+bridge hook the pooled [CLS] is only reachable coincidentally, via the NSP
+head's unembed.hook_in.
+"""
+
+from __future__ import annotations
+
+import pytest
+import torch
+from transformers import BertForMaskedLM, BertForNextSentencePrediction
+
+from transformer_lens.model_bridge.bridge import TransformerBridge
+
+MODEL = "google-bert/bert-base-cased"
+
+
+@pytest.fixture(scope="module")
+def nsp_bridge() -> TransformerBridge:
+    return TransformerBridge.boot_transformers(
+        MODEL, device="cpu", model_class=BertForNextSentencePrediction
+    )
+
+
+def _tokens(bridge: TransformerBridge) -> torch.Tensor:
+    return bridge.tokenizer("Hello there my friend.", return_tensors="pt")["input_ids"]
+
+
+def test_pooler_hook_matches_huggingfaces_own_pooler(nsp_bridge):
+    """hook_out is the pooled [CLS], checked against HF's pooler directly."""
+    tokens = _tokens(nsp_bridge)
+    _, cache = nsp_bridge.run_with_cache(tokens)
+
+    hf = nsp_bridge.original_model
+    with torch.no_grad():
+        expected = hf.bert.pooler(hf.bert(tokens).last_hidden_state)
+
+    torch.testing.assert_close(cache["pooler.hook_out"], expected, atol=0.0, rtol=0.0)
+
+
+def test_pooler_hook_is_post_activation(nsp_bridge):
+    """HookedEncoder fires hook_pooler_out after tanh; the projection is separate."""
+    tokens = _tokens(nsp_bridge)
+    _, cache = nsp_bridge.run_with_cache(tokens)
+
+    pre_activation = cache["pooler.dense.hook_out"]
+    pooled = cache["pooler.hook_out"]
+    assert not torch.allclose(pre_activation, pooled)
+    torch.testing.assert_close(torch.tanh(pre_activation), pooled)
+
+
+def test_hooked_encoder_hook_name_is_aliased(nsp_bridge):
+    """Code migrated from HookedEncoder asks for hook_pooler_out."""
+    tokens = _tokens(nsp_bridge)
+    _, cache = nsp_bridge.run_with_cache(tokens)
+
+    assert torch.equal(cache["pooler.hook_pooler_out"], cache["pooler.hook_out"])
+
+
+def test_masked_lm_checkpoint_without_a_pooler_still_boots():
+    """BertForMaskedLM leaves bert.pooler as None; the mapping must skip it."""
+    bridge = TransformerBridge.boot_transformers(MODEL, device="cpu", model_class=BertForMaskedLM)
+    tokens = _tokens(bridge)
+    logits, cache = bridge.run_with_cache(tokens)
+
+    assert logits.shape[0] == 1
+    assert not [name for name in cache if "pooler" in name]

tests/integration/model_bridge/test_encdec_weight_stacking_parity.py

@@ -0,0 +1,41 @@
+"""Stacked enc-dec weights on the bridge match HookedEncoderDecoder.
+
+Deletion evidence for HookedEncoderDecoder's stacked-weight properties: the
+bridge must produce the same tensors over chain(encoder, decoder), and the same
+head labels, before the legacy class can go.
+"""
+
+from __future__ import annotations
+
+import pytest
+import torch
+
+from transformer_lens import HookedEncoderDecoder
+from transformer_lens.model_bridge.bridge import TransformerBridge
+
+MODEL = "google-t5/t5-small"
+STACKED = ["W_Q", "W_K", "W_V", "W_O", "W_in", "W_out"]
+
+
+@pytest.fixture(scope="module")
+def hooked() -> HookedEncoderDecoder:
+    return HookedEncoderDecoder.from_pretrained(MODEL, device="cpu")
+
+
+@pytest.fixture(scope="module")
+def bridge() -> TransformerBridge:
+    return TransformerBridge.boot_transformers(MODEL, device="cpu")
+
+
+@pytest.mark.parametrize("name", STACKED)
+def test_stacked_weights_match_hooked_encoder_decoder(name, hooked, bridge):
+    """HookedEncoderDecoder does no weight processing, so these are directly comparable."""
+    expected = getattr(hooked, name)
+    actual = getattr(bridge, name)
+    assert actual.shape == expected.shape
+    torch.testing.assert_close(actual, expected, atol=0.0, rtol=0.0)
+
+
+def test_head_labels_match_hooked_encoder_decoder(hooked, bridge):
+    """all_head_labels is a property on the bridge; HT exposes it as a method."""
+    assert bridge.all_head_labels == hooked.all_head_labels()

tests/integration/model_bridge/test_nsp_sentence_pair_helper.py

@@ -0,0 +1,82 @@
+"""Next-sentence prediction from strings on the bridge.
+
+Mirrors BertNextSentencePrediction's string interface, which cannot be adapted
+onto a bridge (its forward reaches for encoder_output/pooler/nsp_head). This
+helper is where that ergonomics survives the Hooked* removal.
+"""
+
+from __future__ import annotations
+
+import pytest
+import torch
+from transformers import AutoTokenizer, BertForNextSentencePrediction
+
+from transformer_lens.model_bridge.bridge import TransformerBridge
+
+MODEL = "google-bert/bert-base-cased"
+SENTENCE_A = "A man walked into a grocery store."
+SEQUENTIAL_B = "He bought an apple."
+UNRELATED_B = "The Eiffel Tower is in Paris."
+
+
+@pytest.fixture(scope="module")
+def nsp_bridge() -> TransformerBridge:
+    bridge = TransformerBridge.boot_transformers(
+        MODEL, device="cpu", model_class=BertForNextSentencePrediction
+    )
+    bridge.enable_compatibility_mode()
+    return bridge
+
+
+@pytest.fixture(scope="module")
+def hf_tokenizer():
+    return AutoTokenizer.from_pretrained(MODEL)
+
+
+def test_pair_tokenization_matches_huggingface(nsp_bridge, hf_tokenizer):
+    """[CLS] a [SEP] b [SEP] with segment ids, identical to tokenizer(a, b)."""
+    tokens = nsp_bridge.to_sentence_pair_tokens(SENTENCE_A, SEQUENTIAL_B)
+    expected = hf_tokenizer(SENTENCE_A, SEQUENTIAL_B, return_tensors="pt")
+
+    assert torch.equal(tokens["input_ids"], expected["input_ids"])
+    assert torch.equal(tokens["token_type_ids"], expected["token_type_ids"])
+    assert tokens["token_type_ids"].unique().tolist() == [0, 1]
+
+
+def test_logits_match_a_direct_huggingface_nsp_forward(nsp_bridge, hf_tokenizer):
+    encodings = hf_tokenizer(SENTENCE_A, SEQUENTIAL_B, return_tensors="pt")
+    with torch.no_grad():
+        expected = nsp_bridge.original_model(**encodings).logits
+
+    actual = nsp_bridge.predict_next_sentence(SENTENCE_A, SEQUENTIAL_B, return_type="logits")
+    torch.testing.assert_close(actual, expected, atol=0.0, rtol=0.0)
+
+
+def test_predictions_distinguish_sequential_from_unrelated(nsp_bridge):
+    assert nsp_bridge.predict_next_sentence(SENTENCE_A, SEQUENTIAL_B) == (
+        "The sentences are sequential"
+    )
+    assert nsp_bridge.predict_next_sentence(SENTENCE_A, UNRELATED_B) == (
+        "The sentences are NOT sequential"
+    )
+
+
+def test_segment_ids_are_load_bearing(nsp_bridge):
+    """Dropping token_type_ids collapses the NSP margin — why the helper owns them."""
+    tokens = nsp_bridge.to_sentence_pair_tokens(SENTENCE_A, SEQUENTIAL_B)
+    with_segments = nsp_bridge.predict_next_sentence(SENTENCE_A, SEQUENTIAL_B, return_type="logits")
+    without_segments = nsp_bridge(
+        tokens["input_ids"],
+        attention_mask=tokens["attention_mask"],
+        return_type="logits",
+    )
+
+    margin = lambda logits: float((logits[0, 0] - logits[0, 1]).abs())
+    assert margin(with_segments) > 2 * margin(without_segments)
+
+
+def test_helper_rejects_a_model_without_an_nsp_head():
+    """An MLM-headed bridge has no 2-class output; say so instead of decoding noise."""
+    bridge = TransformerBridge.boot_transformers(MODEL, device="cpu")
+    with pytest.raises(ValueError, match="next-sentence-prediction head"):
+        bridge.predict_next_sentence(SENTENCE_A, SEQUENTIAL_B)

tests/integration/model_bridge/test_parent_module_traversal.py

@@ -370,7 +370,11 @@ def test_direct_assign_load_stays_current_after_apply(
 @pytest.mark.parametrize(
     ("case_name", "key_fragment", "expected_keys"),
     (
-        ("bert-nsp", "pooler", {"pooler.weight", "pooler.bias"}),
+        # BERT wraps the pooler module (not its inner Linear) so hook_out carries
+        # the post-tanh pooled [CLS], which nests the weights one level deeper.
+        # Still two keys, so nothing is re-expanded — only renamed. ViT wraps
+        # pooler.dense and keeps the flat names.
+        ("bert-nsp", "pooler", {"pooler.dense.weight", "pooler.dense.bias"}),
         ("vit-bare-pooler", "pooler", {"pooler.weight", "pooler.bias"}),
         (
             "ast-audio-classifier",

tests/integration/model_bridge/test_qwen2_moe_bridge.py

@@ -117,3 +117,13 @@ def test_run_with_cache_captures_moe_hooks(self) -> None:
 
             router_scores_key = f"blocks.{layer_idx}.mlp.hook_router_scores"
             assert router_scores_key not in cache
+
+            # Routing observables mirror HookedTransformer: weights at full
+            # expert width, indices at top-k.
+            weights_key = f"blocks.{layer_idx}.mlp.gate.hook_expert_weights"
+            assert weights_key in cache, f"Missing cache key: {weights_key}"
+            assert cache[weights_key].shape == (flat_tokens, num_experts)
+
+            indices_key = f"blocks.{layer_idx}.mlp.gate.hook_expert_indices"
+            assert indices_key in cache, f"Missing cache key: {indices_key}"
+            assert cache[indices_key].shape == (flat_tokens, bridge.cfg.experts_per_token)

tests/unit/model_bridge/supported_architectures/test_arcee_adapter.py

@@ -97,9 +97,6 @@ def test_attn_not_only_and_eager(self, adapter: ArceeArchitectureAdapter) -> Non
         assert adapter.cfg.attn_only is False
         assert adapter.cfg.attn_implementation == "eager"
 
-    def test_gqa_propagated(self, adapter: ArceeArchitectureAdapter) -> None:
-        assert adapter.cfg.n_key_value_heads == 4
-
 
 class TestArceeAdapterComponentMapping:
     """Component-mapping structure and HF module names. Key contrasts with Llama:

tests/unit/model_bridge/supported_architectures/test_bert_adapter.py

@@ -198,7 +198,10 @@ def test_nsp_only_model_uses_hooked_encoder_names(self) -> None:
 
         adapter.prepare_model(hf_model)
 
-        assert adapter.components["pooler"].name == "bert.pooler.dense"
+        # The pooler module itself is wrapped so hook_out is the post-tanh
+        # pooled [CLS]; the projection stays hookable underneath.
+        assert adapter.components["pooler"].name == "bert.pooler"
+        assert adapter.components["pooler"].submodules["dense"].name == "dense"
         assert adapter.components["unembed"].name == "cls.seq_relationship"
         assert "mlm_head" not in adapter.components
         assert "ln_final" not in adapter.components
@@ -212,7 +215,10 @@ def test_combined_mlm_nsp_model_registers_both_heads(self) -> None:
 
         adapter.prepare_model(hf_model)
 
-        assert adapter.components["pooler"].name == "bert.pooler.dense"
+        # The pooler module itself is wrapped so hook_out is the post-tanh
+        # pooled [CLS]; the projection stays hookable underneath.
+        assert adapter.components["pooler"].name == "bert.pooler"
+        assert adapter.components["pooler"].submodules["dense"].name == "dense"
         assert adapter.components["mlm_head"].name == "cls.predictions.transform.dense"
         assert adapter.components["nsp_head"].name == "cls.seq_relationship"
         assert adapter.components["unembed"].name == "cls.predictions.decoder"

tests/unit/model_bridge/supported_architectures/test_falcon_h1_adapter.py

@@ -123,9 +123,6 @@ def test_not_stateful(self, adapter: FalconH1ArchitectureAdapter) -> None:
     def test_eps_attr_variance_epsilon(self, adapter: FalconH1ArchitectureAdapter) -> None:
         assert adapter.cfg.eps_attr == "variance_epsilon"
 
-    def test_n_key_value_heads_propagated(self, adapter: FalconH1ArchitectureAdapter) -> None:
-        assert adapter.cfg.n_key_value_heads == 2
-
     def test_mamba_intermediate_size_propagated(self, adapter: FalconH1ArchitectureAdapter) -> None:
         # mamba_d_ssm is the inner SSM width directly.
         assert getattr(adapter.cfg, "mamba_intermediate_size", None) == 32