fix: raise on stop_at_layer when no 'blocks' stack is registered

tests/unit/model_bridge/test_stop_at_layer_guard.py

@@ -0,0 +1,50 @@
+from types import SimpleNamespace
+
+import pytest
+import torch
+import torch.nn as nn
+
+from transformer_lens.config import TransformerBridgeConfig
+from transformer_lens.model_bridge import TransformerBridge
+
+
+def _bare_bridge(**block_lists: nn.Module) -> TransformerBridge:
+    """A TransformerBridge with only the given block lists registered (no HF model)."""
+    bridge = TransformerBridge.__new__(TransformerBridge)
+    nn.Module.__init__(bridge)
+    bridge.cfg = TransformerBridgeConfig(
+        d_model=8,
+        d_head=4,
+        n_layers=1,
+        n_ctx=16,
+        d_vocab=32,
+        d_mlp=16,
+        n_heads=2,
+        architecture="RavenForCausalLM",
+    )
+    for name, module in block_lists.items():
+        bridge.add_module(name, module)
+    return bridge
+
+
+def test_stop_at_layer_raises_without_blocks_stack() -> None:
+    """Raven-style prelude/core_block/coda lists must not silently ignore stop_at_layer."""
+    bridge = _bare_bridge(
+        prelude=nn.ModuleList([nn.Identity()]),
+        core_block=nn.ModuleList([nn.Identity()]),
+        coda=nn.ModuleList([nn.Identity()]),
+    )
+    with pytest.raises(NotImplementedError, match="stop_at_layer requires a 'blocks' stack"):
+        bridge.forward(torch.zeros(1, 3, dtype=torch.long), stop_at_layer=0)
+
+
+def test_blocks_guard_ignores_unregistered_blocks_attribute() -> None:
+    """A wrapped HF model exposing `.blocks` must not satisfy the guard via __getattr__."""
+    bridge = _bare_bridge()
+    bridge.__dict__["original_model"] = SimpleNamespace(blocks=[object()])
+    assert hasattr(bridge, "blocks")  # the trap: __getattr__ falls through to the HF model
+    assert not bridge._has_registered_blocks()
+    with pytest.raises(NotImplementedError, match="stop_at_layer requires a 'blocks' stack"):
+        bridge.forward(torch.zeros(1, 3, dtype=torch.long), stop_at_layer=0)
+    with pytest.raises(NotImplementedError, match="start_at_layer requires a 'blocks' stack"):
+        bridge.forward(torch.zeros(1, 3, 8), start_at_layer=0)

transformer_lens/model_bridge/transformer_bridge.py

@@ -3,6 +3,7 @@
 This module provides the bridge components that wrap remote model components and provide
 a consistent interface for accessing their weights and performing operations.
 """
+
 import inspect
 import logging
 import re
@@ -460,6 +461,16 @@ def n_params_total(self) -> int:
         """
         return self._n_params_total
 
+    def _has_registered_blocks(self) -> bool:
+        """Whether a ``blocks`` stack is registered as a submodule on this bridge.
+
+        Checks ``_modules`` directly rather than ``hasattr``: ``__getattr__`` falls
+        through to the wrapped HF model, so ``hasattr(self, "blocks")`` can be True
+        for a model that merely exposes its own ``.blocks`` attribute.
+        """
+        modules = self.__dict__.get("_modules") or {}
+        return "blocks" in modules
+
     def __getattr__(self, name: str) -> Any:
         """Provide a clear error message for missing attributes."""
         # Re-invoke original_model's property so its descriptive AttributeError
@@ -2095,7 +2106,10 @@ def forward(
                 output is discarded when block k swaps in the residual) but are excluded
                 from ``run_with_cache`` output. Requires an HF model that accepts
                 ``inputs_embeds``; only supported on the standard ``blocks`` stack.
-            stop_at_layer: Layer to stop forward pass at
+            stop_at_layer: Layer to stop forward pass at. Only supported on the
+                standard ``blocks`` stack; architectures that register no ``blocks``
+                (e.g. Raven's ``prelude``/``core_block``/``coda``) raise
+                ``NotImplementedError`` rather than running to completion.
             pixel_values: Optional image tensor for multimodal models (e.g., LLaVA, Gemma3)
                 and vision models (eg. ViT, DeiT).
                 The tensor is passed directly to the underlying HuggingFace model.
@@ -2135,26 +2149,18 @@ def forward(
 
         if start_at_layer is not None:
             input = self._setup_start_at_layer(input, start_at_layer)
-
         # Set stop_at_layer flag on all blocks if requested
         if stop_at_layer is not None:
-            if (
-                hasattr(self, "L_blocks")
-                or hasattr(self, "H_blocks")
-                or hasattr(self, "encoder_blocks")
-                or hasattr(self, "decoder_blocks")
-            ):
+            if not self._has_registered_blocks():
                 raise NotImplementedError(
-                    "stop_at_layer is not supported on non-standard block list "
-                    "names (L_blocks, H_blocks, encoder_blocks, decoder_blocks). "
-                    "The bridge only supports stop_at_layer on 'blocks'."
-                )
-            if hasattr(self, "blocks"):
-                effective_stop_at_layer = (
-                    len(self.blocks) + stop_at_layer if stop_at_layer < 0 else stop_at_layer
+                    "stop_at_layer requires a 'blocks' stack; this architecture "
+                    "does not register one."
                 )
-                for block in self.blocks:
-                    block._stop_at_layer_idx = effective_stop_at_layer
+            effective_stop_at_layer = (
+                len(self.blocks) + stop_at_layer if stop_at_layer < 0 else stop_at_layer
+            )
+            for block in self.blocks:
+                block._stop_at_layer_idx = effective_stop_at_layer
 
         # Map HookedEncoderDecoder-style kwargs to HF-compatible names
         if "decoder_input" in kwargs:
@@ -2449,8 +2455,9 @@ def _setup_start_at_layer(self, input: Any, start_at_layer: int) -> torch.Tensor
                     "start_at_layer is only supported on the standard 'blocks' stack, "
                     f"not {alt!r}."
                 )
-        if not hasattr(self, "blocks"):
+        if not self._has_registered_blocks():
             raise NotImplementedError("start_at_layer requires a 'blocks' stack.")
+
         if not (isinstance(input, torch.Tensor) and input.is_floating_point()):
             raise ValueError(
                 "start_at_layer requires a residual-stream tensor [batch, pos, d_model]; "
@@ -2824,18 +2831,22 @@ def _generate_tokens(
                         temperature=temperature,
                         freq_penalty=freq_penalty,
                         repetition_penalty=repetition_penalty,
-                        tokens=penalty_tokens
-                        if _generate_from_embeds
-                        else (decoder_tokens if is_encoder_decoder else current_tokens),
+                        tokens=(
+                            penalty_tokens
+                            if _generate_from_embeds
+                            else (decoder_tokens if is_encoder_decoder else current_tokens)
+                        ),
                     ).to(self.cfg.device)
                 else:
                     sampled_tokens = utils.sample_logits(
                         final_logits,
                         temperature=0.0,
                         repetition_penalty=repetition_penalty,
-                        tokens=penalty_tokens
-                        if _generate_from_embeds
-                        else (decoder_tokens if is_encoder_decoder else current_tokens),
+                        tokens=(
+                            penalty_tokens
+                            if _generate_from_embeds
+                            else (decoder_tokens if is_encoder_decoder else current_tokens)
+                        ),
                     ).to(self.cfg.device)
 
                 # Freeze rows that finished on an earlier step so they stop emitting