Attention shape normalization (#997)

* created input output reshape hook functionality

* refactor weight conversion to hook conversion

* ran format

* moved conversion utils

* added auto attention attempt

* added revert function

* ran format

* fixed type checks

* fixed tests

* ran format

* ran format

* fixed test

* replaced config creation with more robust object

* restored init file

* removed type checking

* updated hook to match what it should have been

* made pattern rules a param

* updated jqv component to allow conversion to be passed through

* made c_attn sub module

* fixed type issues

* fixed test

* simplified config flow

* fixed test

* set attention output

* removed output attentions config var

Author: bryce13950

tests/integration/model_bridge/test_bridge_integration.py

@@ -122,5 +122,164 @@ def test_component_access():
     assert hasattr(block, "ln2"), "Block should have second layer norm"
 
 
+def test_joint_qkv_custom_conversion_rule():
+    """Test that custom QKV conversion rules can be passed to JointQKVAttentionBridge."""
+    from transformer_lens.conversion_utils.conversion_steps.rearrange_hook_conversion import (
+        RearrangeHookConversion,
+    )
+    from transformer_lens.model_bridge.generalized_components.joint_qkv_attention import (
+        JointQKVAttentionBridge,
+    )
+    from transformer_lens.model_bridge.generalized_components.linear import LinearBridge
+
+    model_name = "gpt2"  # Use a smaller model for testing
+    bridge = TransformerBridge.boot_transformers(model_name)
+
+    # Create a custom QKV conversion rule
+    custom_qkv_conversion = RearrangeHookConversion(
+        "batch seq (num_attention_heads d_head) -> batch seq num_attention_heads d_head",
+        num_attention_heads=12,  # GPT-2 small has 12 heads
+    )
+
+    # Create QKV config
+    qkv_config = {
+        "split_qkv_matrix": lambda x: (x, x, x),  # Dummy function for test
+    }
+
+    # Create submodules
+    submodules = {
+        "qkv": LinearBridge(name="c_attn"),
+        "o": LinearBridge(name="c_proj"),
+    }
+
+    # This should not raise an error
+    test_bridge = JointQKVAttentionBridge(
+        name="test_joint_qkv",
+        model_config=bridge.cfg,
+        submodules=submodules,
+        qkv_config=qkv_config,
+        qkv_conversion_rule=custom_qkv_conversion,
+    )
+
+    # Verify the custom conversion rule was set on Q, K, V components
+    assert (
+        test_bridge.q.hook_out.hook_conversion is custom_qkv_conversion
+    ), "Custom QKV conversion rule should be set on Q"
+    assert (
+        test_bridge.k.hook_out.hook_conversion is custom_qkv_conversion
+    ), "Custom QKV conversion rule should be set on K"
+    assert (
+        test_bridge.v.hook_out.hook_conversion is custom_qkv_conversion
+    ), "Custom QKV conversion rule should be set on V"
+
+
+def test_attention_pattern_hook_shape_custom_conversion():
+    """Test that custom pattern conversion rules can be passed to attention components."""
+    from transformer_lens.conversion_utils.conversion_steps.rearrange_hook_conversion import (
+        RearrangeHookConversion,
+    )
+
+    model_name = "gpt2"  # Use a smaller model for testing
+    bridge = TransformerBridge.boot_transformers(model_name)
+
+    if bridge.tokenizer.pad_token is None:
+        bridge.tokenizer.pad_token = bridge.tokenizer.eos_token
+
+    # Create a custom conversion rule (this is just for testing the parameter passing)
+    custom_conversion = RearrangeHookConversion(
+        "batch n_heads pos_q pos_k -> batch n_heads pos_q pos_k"  # Same as default but explicitly set
+    )
+
+    # Verify that the attention bridge accepts the custom conversion parameter
+    # We can't easily test this with the existing bridge without recreating it,
+    # but we can at least verify the parameter is accepted without error
+    from transformer_lens.model_bridge.generalized_components.attention import (
+        AttentionBridge,
+    )
+
+    # This should not raise an error
+    test_bridge = AttentionBridge(
+        name="test_attn", config=bridge.cfg, pattern_conversion_rule=custom_conversion
+    )
+
+    # Verify the conversion rule was set
+    assert (
+        test_bridge.hook_pattern.hook_conversion is custom_conversion
+    ), "Custom conversion rule should be set"
+
+
+def test_attention_pattern_hook_shape():
+    """Test that the attention pattern hook produces the correct shape (batch, n_heads, pos, pos)."""
+    model_name = "gpt2"  # Use a smaller model for testing
+    bridge = TransformerBridge.boot_transformers(
+        model_name,
+        hf_config_overrides={
+            "attn_implementation": "eager",
+        },
+    )
+
+    if bridge.tokenizer.pad_token is None:
+        bridge.tokenizer.pad_token = bridge.tokenizer.eos_token
+
+    # Attention output enabled via hf_config_overrides
+
+    # Variable to store captured attention patterns
+    captured_patterns = {}
+
+    def capture_pattern_hook(tensor, hook):
+        """Hook to capture attention patterns."""
+        captured_patterns[hook.name] = tensor.clone()
+        return tensor
+
+    # Add hook to capture attention patterns
+    bridge.blocks[0].attn.hook_pattern.add_hook(capture_pattern_hook)
+
+    try:
+        # Run model with a prompt
+        prompt = "The quick brown fox"
+        tokens = bridge.to_tokens(prompt)
+        batch_size, seq_len = tokens.shape
+
+        # Run forward pass
+        output = bridge(tokens)
+
+        # Verify we captured attention patterns
+        assert len(captured_patterns) > 0, "Should have captured attention patterns"
+
+        # Get the captured pattern tensor
+        pattern_tensor = list(captured_patterns.values())[0]
+
+        # Verify the shape is (batch, n_heads, pos, pos)
+        assert (
+            len(pattern_tensor.shape) == 4
+        ), f"Pattern tensor should be 4D, got {len(pattern_tensor.shape)}D"
+
+        batch_dim, n_heads_dim, pos_q_dim, pos_k_dim = pattern_tensor.shape
+
+        # Verify dimensions make sense
+        assert batch_dim == batch_size, f"Batch dimension should be {batch_size}, got {batch_dim}"
+        assert (
+            n_heads_dim == bridge.cfg.n_heads
+        ), f"Heads dimension should be {bridge.cfg.n_heads}, got {n_heads_dim}"
+        assert (
+            pos_q_dim == seq_len
+        ), f"Query position dimension should be {seq_len}, got {pos_q_dim}"
+        assert pos_k_dim == seq_len, f"Key position dimension should be {seq_len}, got {pos_k_dim}"
+
+        # Verify it's actually attention weights (should be non-negative and roughly sum to 1 along last dim)
+        assert torch.all(pattern_tensor >= 0), "Attention patterns should be non-negative"
+
+        # Check that attention weights roughly sum to 1 along the last dimension (with some tolerance for numerical precision)
+        attention_sums = pattern_tensor.sum(dim=-1)
+        expected_sums = torch.ones_like(attention_sums)
+        assert torch.allclose(
+            attention_sums, expected_sums, atol=1e-5
+        ), "Attention patterns should sum to ~1 along key dimension"
+
+    finally:
+        # Clean up hooks
+        bridge.blocks[0].attn.hook_pattern.remove_hooks()
+
+
 if __name__ == "__main__":
     pytest.main([__file__])

tests/integration/model_bridge/test_bridge_root_module_cache_compatibility.py

@@ -3,7 +3,15 @@
 MODEL = "gpt2"
 
 prompt = "Hello World!"
-bridge = TransformerBridge.boot_transformers(MODEL, device="cpu")
+bridge = TransformerBridge.boot_transformers(
+    MODEL,
+    device="cpu",
+    hf_config_overrides={
+        "attn_implementation": "eager",
+    },
+)
+
+# Attention output enabled via hf_config_overrides
 
 act_names_in_cache = [
     # "hook_embed",

tests/mocks/architecture_adapter.py

@@ -1,4 +1,6 @@
 """Mock architecture adapter for testing."""
+from types import SimpleNamespace
+
 import pytest
 import torch.nn as nn
 
@@ -18,13 +20,11 @@ class MockArchitectureAdapter(ArchitectureAdapter):
     def __init__(self, cfg=None):
         if cfg is None:
             # Create a minimal config for testing
-            cfg = type(
-                "MockConfig",
-                (),
-                {"d_mlp": 512, "intermediate_size": 512, "default_prepend_bos": True},
-            )()
+            cfg = SimpleNamespace(d_mlp=512, intermediate_size=512, default_prepend_bos=True)
         super().__init__(cfg)
         # Use actual bridge instances instead of tuples
+        # Provide minimal config to components that require it
+        attn_cfg = SimpleNamespace(n_heads=1)
         self.component_mapping = {
             "embed": EmbeddingBridge(name="embed"),
             "unembed": EmbeddingBridge(name="unembed"),
@@ -34,7 +34,7 @@ def __init__(self, cfg=None):
                 submodules={
                     "ln1": NormalizationBridge(name="ln1"),
                     "ln2": NormalizationBridge(name="ln2"),
-                    "attn": AttentionBridge(name="attn"),
+                    "attn": AttentionBridge(name="attn", config=attn_cfg),
                     "mlp": MLPBridge(name="mlp"),
                 },
             ),

tests/unit/model_bridge/test_bridge.py

@@ -4,6 +4,7 @@
 and other bridge operations.
 """
 
+from types import SimpleNamespace
 from unittest.mock import MagicMock
 
 import pytest
@@ -42,7 +43,9 @@ def mock_get_component(model, path):
                 comp.set_original_component(model.unembed)
                 return comp
             elif "blocks" in path and "attn" in path:
-                comp = AttentionBridge(name="attn")
+                # Minimal config with n_heads for AttentionBridge
+                attn_cfg = SimpleNamespace(n_heads=1)
+                comp = AttentionBridge(name="attn", config=attn_cfg)
                 comp.set_original_component(model.blocks[0].attn)
                 return comp
             elif "blocks" in path and "mlp" in path:
@@ -99,7 +102,7 @@ def test_format_block_mapping_tuple(self):
                 submodules={
                     "ln1": NormalizationBridge(name="ln1"),
                     "ln2": NormalizationBridge(name="ln2"),
-                    "attn": AttentionBridge(name="attn"),
+                    "attn": AttentionBridge(name="attn", config=SimpleNamespace(n_heads=1)),
                     "mlp": MLPBridge(name="mlp"),
                 },
             )
@@ -124,7 +127,7 @@ def test_format_mixed_mapping(self):
                 name="blocks",
                 submodules={
                     "ln1": NormalizationBridge(name="ln1"),
-                    "attn": AttentionBridge(name="attn"),
+                    "attn": AttentionBridge(name="attn", config=SimpleNamespace(n_heads=1)),
                 },
             ),
             "ln_final": NormalizationBridge(name="ln_final"),
@@ -145,7 +148,7 @@ def test_format_with_prepend_path(self):
         """Test formatting with prepend path parameter."""
         mapping = {
             "ln1": NormalizationBridge(name="ln1"),
-            "attn": AttentionBridge(name="attn"),
+            "attn": AttentionBridge(name="attn", config=SimpleNamespace(n_heads=1)),
         }
         # To test prepending, we need a parent structure in the component mapping
         self.bridge.adapter.component_mapping = {

tests/unit/model_bridge/test_component_setup.py

@@ -1,6 +1,8 @@
 """Tests for component setup utilities."""
 
 
+from types import SimpleNamespace
+
 import pytest
 import torch.nn as nn
 
@@ -55,6 +57,7 @@ def test_setup_submodules_basic(self, mock_model_adapter):
         # Create a component with submodules
         component = AttentionBridge(
             name="self_attn",
+            config=SimpleNamespace(n_heads=1),
             submodules={
                 "q_proj": EmbeddingBridge(name="q_proj"),
                 "k_proj": EmbeddingBridge(name="k_proj"),
@@ -82,10 +85,13 @@ def test_setup_submodules_nested(self):
 
         # Create a component with nested submodules
         inner_component = AttentionBridge(
-            name="q_proj", submodules={}  # This should match a real path
+            name="q_proj",
+            config=SimpleNamespace(n_heads=1),
+            submodules={},  # This should match a real path
         )
         component = AttentionBridge(
             name="attn",
+            config=SimpleNamespace(n_heads=1),
             submodules={
                 "q_proj": inner_component,
             },
@@ -169,7 +175,7 @@ def test_setup_blocks_bridge(self):
             submodules={
                 "ln1": NormalizationBridge(name="ln1"),
                 "ln2": NormalizationBridge(name="ln2"),
-                "attn": AttentionBridge(name="attn"),
+                "attn": AttentionBridge(name="attn", config=SimpleNamespace(n_heads=1)),
                 "mlp": MLPBridge(name="mlp"),
             },
         )
@@ -240,7 +246,7 @@ def __init__(self):
                         submodules={
                             "ln1": NormalizationBridge(name="ln1"),
                             "ln2": NormalizationBridge(name="ln2"),
-                            "attn": AttentionBridge(name="attn"),
+                            "attn": AttentionBridge(name="attn", config=SimpleNamespace(n_heads=1)),
                             "mlp": MLPBridge(name="mlp"),
                         },
                     ),

tests/unit/model_bridge/test_end_to_end_bridge.py

@@ -1,4 +1,5 @@
 """End-to-end tests for the TransformerBridge."""
+from types import SimpleNamespace
 from unittest.mock import MagicMock
 
 import torch.nn as nn
@@ -35,7 +36,7 @@ def test_bridge_creation_and_component_access(self):
                 name="encoder.layers",
                 submodules={
                     "ln1": NormalizationBridge(name="norm1"),
-                    "attn": AttentionBridge(name="self_attn"),
+                    "attn": AttentionBridge(name="self_attn", config=SimpleNamespace(n_heads=1)),
                 },
             ),
         }

tests/unit/test_hook_points.py

@@ -44,8 +44,9 @@ def test_add_hook_with_level(mock_handle):
     assert hook_point.fwd_hooks[0].context_level == 5
 
 
+@mock.patch("transformer_lens.hook_points.LensHandle")
 @mock.patch("torch.utils.hooks.RemovableHandle")
-def test_add_hook_prepend(mock_handle):
+def test_add_hook_prepend(mock_handle, mock_lens_handle):
     mock_handle.id = 0
     mock_handle.next_id = 1
 
@@ -57,6 +58,26 @@ def hook1(activation, hook):
     def hook2(activation, hook):
         return activation
 
+    # Make LensHandle constructor return a simple container capturing the pt_handle ('hook')
+    class _LensHandleBox:
+        def __init__(self, handle, is_permanent, context_level):
+            self.hook = handle
+            self.is_permanent = is_permanent
+            self.context_level = context_level
+
+    mock_lens_handle.side_effect = _LensHandleBox
+
+    # Override register_forward_hook to return mocked handles with incremental ids
+    next_id = {"val": 1}
+
+    def fake_register_forward_hook(fn, prepend=False):
+        handle = mock.MagicMock()
+        handle.id = next_id["val"]
+        next_id["val"] += 1
+        return handle
+
+    hook_point.register_forward_hook = fake_register_forward_hook  # type: ignore[assignment]
+
     hook_point.add_hook(hook1, dir="fwd")
     hook_point.add_hook(hook2, dir="fwd", prepend=True)
 

transformer_lens/__init__.py

@@ -2,6 +2,7 @@
 from . import utilities
 from . import hook_points
 from . import evals
+from . import conversion_utils
 from .past_key_value_caching import (
     HookedTransformerKeyValueCache,
     HookedTransformerKeyValueCacheEntry,