Fix direct TransformerBridge config flag assignment

docs/source/content/migrating_to_v3.md

@@ -122,7 +122,7 @@ For the full mapping of legacy → canonical names and the expected tensor shape
 
 Two semantic differences inside `enable_compatibility_mode()` worth knowing if you are porting activation-patching, DLA, or attribution-patching code:
 
-- **`blocks.{i}.hook_mlp_in` fires pre-ln2** (matching legacy `HookedTransformer`). Use `bridge.set_use_hook_mlp_in(True)` to enable it — setting `cfg.use_hook_mlp_in = True` directly is honored when blocks share the bridge's `cfg`, but the setter is the supported entry point. The pre-ln2 placement means cached values from one run can be patched into another and re-flow through `ln2 → mlp` consistently across the bridge and `HookedTransformer`.
+- **`blocks.{i}.hook_mlp_in` fires pre-ln2** (matching legacy `HookedTransformer`). Enable it with `bridge.set_use_hook_mlp_in(True)` or `bridge.cfg.use_hook_mlp_in = True`; direct config assignment routes through the same validation and propagation path as the setter. The pre-ln2 placement means cached values from one run can be patched into another and re-flow through `ln2 → mlp` consistently across the bridge and `HookedTransformer`.
 - **`hook_q_input` / `hook_k_input` / `hook_v_input` / `hook_attn_in`** also fire pre-ln1 in compat mode. On the per-head LN application that follows, the bridge routes through the raw HF norm rather than the `NormalizationBridge` wrapper, so `ln1`'s sub-hooks (`hook_in`, `hook_normalized`, `hook_scale`) do **not** fire once per head the way legacy `LayerNormPre` would. Q/K/V projections downstream still match legacy numerically; only the intermediate LN sub-hook firing is suppressed.
 
 Post-norm architectures (OLMo 2, BERT-style encoders) and MLA blocks (DeepSeek V2/V3/R1) do not participate in the pre-ln1 capture — `MLABlockBridge` does not expose those aliases, and post-norm models would read the post-attention residual instead of the block input.

tests/unit/model_bridge/test_config_flag_assignment.py

@@ -0,0 +1,223 @@
+"""Tests for direct assignment of Bridge-managed hook flags (#1689)."""
+
+from __future__ import annotations
+
+import copy
+import gc
+
+import pytest
+import torch
+from torch import nn
+from transformers import GPT2Config, GPT2LMHeadModel, LlamaConfig, LlamaForCausalLM
+
+from transformer_lens.config import TransformerBridgeConfig
+from transformer_lens.factories.architecture_adapter_factory import (
+    ArchitectureAdapterFactory,
+)
+from transformer_lens.model_bridge import TransformerBridge
+from transformer_lens.model_bridge.sources._bridge_builder import (
+    build_bridge_from_module,
+)
+from transformer_lens.model_bridge.sources.native import NativeModel
+
+
+def _cfg() -> TransformerBridgeConfig:
+    return TransformerBridgeConfig(
+        d_model=32,
+        d_head=16,
+        n_heads=2,
+        n_layers=1,
+        n_ctx=8,
+        d_vocab=16,
+        d_mlp=64,
+        act_fn="gelu",
+        normalization_type="LN",
+        seed=0,
+    )
+
+
+def _tiny_gpt2_bridge() -> TransformerBridge:
+    hf_config = GPT2Config(
+        n_layer=1,
+        n_head=2,
+        n_embd=32,
+        n_positions=8,
+        n_ctx=8,
+        vocab_size=16,
+    )
+    hf_model = GPT2LMHeadModel(hf_config).eval()
+    return build_bridge_from_module(
+        hf_model,
+        "GPT2LMHeadModel",
+        hf_config=hf_config,
+        tokenizer=None,
+        device="cpu",
+    )
+
+
+def _tiny_llama_bridge() -> TransformerBridge:
+    hf_config = LlamaConfig(
+        hidden_size=32,
+        intermediate_size=64,
+        num_hidden_layers=1,
+        num_attention_heads=2,
+        num_key_value_heads=2,
+        vocab_size=16,
+        max_position_embeddings=8,
+    )
+    hf_model = LlamaForCausalLM(hf_config).eval()
+    return build_bridge_from_module(
+        hf_model,
+        "LlamaForCausalLM",
+        hf_config=hf_config,
+        tokenizer=None,
+        device="cpu",
+    )
+
+
+@pytest.fixture(params=["gpt2", "llama"], ids=["shared-config", "cloned-config"])
+def bridge_with_config_mode(request: pytest.FixtureRequest) -> TransformerBridge:
+    bridge = _tiny_gpt2_bridge() if request.param == "gpt2" else _tiny_llama_bridge()
+    attn_config_is_shared = bridge.blocks[0].attn.config is bridge.cfg
+    assert attn_config_is_shared is (request.param == "gpt2")
+    return bridge
+
+
+@pytest.mark.parametrize(
+    ("flag_name", "hook_name"),
+    [
+        ("use_attn_result", "blocks.0.attn.hook_result"),
+        ("use_attn_in", "blocks.0.attn.hook_attn_in"),
+        ("use_hook_mlp_in", "blocks.0.hook_mlp_in"),
+        ("use_split_qkv_input", "blocks.0.attn.hook_q_input"),
+    ],
+)
+def test_direct_assignment_matches_setter_hook_behavior(
+    bridge_with_config_mode: TransformerBridge, flag_name: str, hook_name: str
+) -> None:
+    bridge = bridge_with_config_mode
+    tokens = torch.randint(0, bridge.cfg.d_vocab, (1, 8))
+
+    setattr(bridge.cfg, flag_name, True)
+    _, direct_cache = bridge.run_with_cache(tokens, names_filter=[hook_name])
+
+    setattr(bridge.cfg, flag_name, False)
+    getattr(bridge, f"set_{flag_name}")(True)
+    _, setter_cache = bridge.run_with_cache(tokens, names_filter=[hook_name])
+
+    assert list(direct_cache) == [hook_name]
+    assert list(setter_cache) == [hook_name]
+    assert direct_cache[hook_name].shape == setter_cache[hook_name].shape
+
+
+def test_direct_assignment_preserves_mutual_exclusivity() -> None:
+    bridge = _tiny_gpt2_bridge()
+
+    bridge.cfg.use_split_qkv_input = True
+    with pytest.raises(ValueError, match="mutually exclusive"):
+        bridge.cfg.use_attn_in = True
+    assert bridge.cfg.use_attn_in is False
+
+    bridge.cfg.use_split_qkv_input = False
+    bridge.cfg.use_attn_in = True
+    with pytest.raises(ValueError, match="mutually exclusive"):
+        bridge.cfg.use_split_qkv_input = True
+    assert bridge.cfg.use_split_qkv_input is False
+
+
+@pytest.mark.parametrize("flag_name", ["use_attn_result", "use_attn_in", "use_split_qkv_input"])
+def test_direct_assignment_preserves_unsupported_architecture_errors(
+    monkeypatch: pytest.MonkeyPatch, flag_name: str
+) -> None:
+    bridge = TransformerBridge.boot_native(_cfg())
+
+    class _FakeBlock(nn.Module):
+        def __init__(self) -> None:
+            super().__init__()
+            self.attn = nn.Identity()
+
+    monkeypatch.setattr(bridge, "blocks", nn.ModuleList([_FakeBlock()]), raising=True)
+
+    with pytest.raises(NotImplementedError, match=flag_name):
+        setattr(bridge.cfg, flag_name, True)
+    assert getattr(bridge.cfg, flag_name) is False
+
+
+def test_deepcopied_live_config_is_not_bound_to_original_bridge() -> None:
+    bridge = TransformerBridge.boot_native(_cfg())
+    copied_cfg = copy.deepcopy(bridge.cfg)
+
+    copied_cfg.use_hook_mlp_in = True
+
+    assert copied_cfg.use_hook_mlp_in is True
+    assert bridge.cfg.use_hook_mlp_in is False
+
+
+def test_deepcopied_bridge_rebinds_its_config() -> None:
+    bridge = TransformerBridge.boot_native(_cfg())
+    copied_bridge = copy.deepcopy(bridge)
+
+    copied_bridge.cfg.use_hook_mlp_in = True
+
+    assert copied_bridge.cfg.use_hook_mlp_in is True
+    assert copied_bridge.blocks[0].config.use_hook_mlp_in is True
+    assert bridge.cfg.use_hook_mlp_in is False
+
+
+def test_shallow_copied_bridge_does_not_replace_live_config_binding() -> None:
+    bridge = TransformerBridge.boot_native(_cfg())
+    with pytest.warns(UserWarning, match="already bound to another live"):
+        copied_bridge = copy.copy(bridge)
+
+    assert copied_bridge.cfg is bridge.cfg
+    assert bridge.cfg._bridge_ref() is bridge
+
+    del copied_bridge
+    gc.collect()
+    bridge.cfg.use_hook_mlp_in = True
+
+    assert bridge.blocks[0].config.use_hook_mlp_in is True
+
+
+def test_constructor_warns_when_live_bridge_already_owns_config() -> None:
+    cfg = _cfg()
+    cfg.architecture = "TransformerLensNative"
+    first_model = NativeModel(cfg)
+    second_model = NativeModel(cfg)
+    first_adapter = ArchitectureAdapterFactory.select_architecture_adapter(cfg)
+    second_adapter = ArchitectureAdapterFactory.select_architecture_adapter(cfg)
+    first_adapter.prepare_model(first_model)
+    second_adapter.prepare_model(second_model)
+    first_bridge = TransformerBridge(first_model, first_adapter, tokenizer=None)
+
+    with pytest.warns(UserWarning, match="already bound to another live"):
+        second_bridge = TransformerBridge(second_model, second_adapter, tokenizer=None)
+
+    assert second_bridge.cfg is first_bridge.cfg
+    assert cfg._bridge_ref() is first_bridge
+
+
+def test_attention_flag_propagation_does_not_dispatch_bound_cloned_config() -> None:
+    bridge = _tiny_llama_bridge()
+    cloned_cfg = bridge.blocks[0].attn.config
+    other_bridge = TransformerBridge.boot_native(_cfg())
+    assert cloned_cfg is not bridge.cfg
+    cloned_cfg._bind_bridge(other_bridge)
+
+    bridge.set_use_attn_in(True)
+
+    assert cloned_cfg.use_attn_in is True
+    assert other_bridge.cfg.use_attn_in is False
+
+
+def test_mlp_flag_propagation_does_not_dispatch_bound_cloned_config() -> None:
+    bridge = _tiny_gpt2_bridge()
+    cloned_cfg = bridge.blocks[0].config
+    other_bridge = TransformerBridge.boot_native(_cfg())
+    assert cloned_cfg is not bridge.cfg
+    cloned_cfg._bind_bridge(other_bridge)
+
+    bridge.set_use_hook_mlp_in(True)
+
+    assert cloned_cfg.use_hook_mlp_in is True
+    assert other_bridge.cfg.use_hook_mlp_in is False

transformer_lens/config/transformer_bridge_config.py

@@ -1,6 +1,8 @@
 """Configuration class for TransformerBridge."""
 
-from typing import Optional
+import warnings
+import weakref
+from typing import Any, Optional
 
 import torch
 
@@ -18,6 +20,17 @@ class TransformerBridgeConfig(TransformerLensConfig):
     Also includes all HookedTransformerConfig fields for compatibility.
     """
 
+    __slots__ = ("_bridge_ref",)
+
+    _BRIDGE_MANAGED_HOOK_FLAGS = frozenset(
+        {
+            "use_attn_result",
+            "use_attn_in",
+            "use_hook_mlp_in",
+            "use_split_qkv_input",
+        }
+    )
+
     def __init__(
         self,
         d_model: int,
@@ -109,6 +122,7 @@ def __init__(
         **kwargs,
     ):
         """Initialize TransformerBridgeConfig."""
+        object.__setattr__(self, "_bridge_ref", None)
         super().__init__(
             d_model=d_model,
             d_head=d_head,
@@ -204,9 +218,48 @@ def __init__(
         self.vision_num_layers = vision_num_layers
         self.vision_num_heads = vision_num_heads
         self.mm_tokens_per_image = mm_tokens_per_image
-
         self.__post_init__()
 
+    def __setattr__(self, name: str, value: Any) -> None:
+        """Route live Bridge hook-flag assignments through their public setters."""
+        if name in self._BRIDGE_MANAGED_HOOK_FLAGS:
+            bridge_ref = getattr(self, "_bridge_ref", None)
+            bridge = bridge_ref() if bridge_ref is not None else None
+            if bridge is not None:
+                getattr(bridge, f"set_{name}")(value)
+                return
+        super().__setattr__(name, value)
+
+    def __getstate__(self) -> dict[str, Any]:
+        """Serialize config data without retaining its live Bridge binding."""
+        return self.__dict__.copy()
+
+    def __setstate__(self, state: dict[str, Any]) -> None:
+        """Restore an unbound config copy."""
+        self.__dict__.update(state)
+        object.__setattr__(self, "_bridge_ref", None)
+
+    def _bind_bridge(self, bridge: Any) -> None:
+        """Bind runtime hook-flag assignments to a constructed Bridge."""
+        bridge_ref = getattr(self, "_bridge_ref", None)
+        bound_bridge = bridge_ref() if bridge_ref is not None else None
+        if bound_bridge is None:
+            object.__setattr__(self, "_bridge_ref", weakref.ref(bridge))
+        elif bound_bridge is not bridge:
+            warnings.warn(
+                "TransformerBridgeConfig is already bound to another live "
+                "TransformerBridge; declining to bind it to this instance. "
+                "Direct assignments to Bridge-managed hook flags will continue "
+                "to configure the existing TransformerBridge.",
+                stacklevel=3,
+            )
+
+    def _set_bridge_managed_hook_flag(self, name: str, value: bool) -> None:
+        """Set a managed flag without re-entering the Bridge setter."""
+        if name not in self._BRIDGE_MANAGED_HOOK_FLAGS:
+            raise ValueError(f"Unknown Bridge-managed hook flag: {name}")
+        object.__setattr__(self, name, value)
+
     def __post_init__(self):
         """Post-initialization processing."""
         # dtype is guaranteed to be set at this point