transformer_lens.model_bridge.generalized_components.block module

Block bridge component.

This module contains the bridge component for transformer blocks.

class transformer_lens.model_bridge.generalized_components.block.BlockBridge(name: str, config: Any | None = None, submodules: Dict[str, GeneralizedComponent] | None = None, hook_alias_overrides: Dict[str, str] | None = None, mlp_reads_resid_directly: bool = False)

Bases: GeneralizedComponent

Bridge component for transformer blocks.

This component provides standardized input/output hooks and monkey-patches HuggingFace blocks to insert hooks at positions matching HookedTransformer.

__init__(name: str, config: Any | None = None, submodules: Dict[str, GeneralizedComponent] | None = None, hook_alias_overrides: Dict[str, str] | None = None, mlp_reads_resid_directly: bool = False)

Initialize the block bridge.

Parameters:

  • name – The name of the component in the model

  • config – Optional configuration (unused for BlockBridge)

  • submodules – Dictionary of submodules to register

  • hook_alias_overrides – Optional dictionary to override default hook aliases. For example, {“hook_attn_out”: “ln1_post.hook_out”} will make hook_attn_out point to ln1_post.hook_out instead of the default attn.hook_out.

  • mlp_reads_resid_directly – True for post-norm blocks where the MLP consumes the mid-residual with no pre-MLP norm (OLMo 2 layout). Moves the hook_mlp_in capture from ln2 (whose input there is the raw MLP output) to the MLP itself.

forward(*args: Any, **kwargs: Any) → Any

Forward pass through the block bridge.

Parameters:

  • *args – Input arguments

  • **kwargs – Input keyword arguments

Returns:

The output from the original component

Raises:

StopAtLayerException – If stop_at_layer is set and this block should stop execution

hook_aliases: Dict[str, str | List[str]] = {'hook_attn_in': 'attn.hook_attn_in', 'hook_attn_out': 'attn.hook_out', 'hook_k_input': 'attn.hook_k_input', 'hook_mlp_out': 'mlp.hook_out', 'hook_q_input': 'attn.hook_q_input', 'hook_resid_mid': 'ln2.hook_in', 'hook_resid_post': 'hook_out', 'hook_resid_pre': 'hook_in', 'hook_v_input': 'attn.hook_v_input'}
hook_out_is_single_residual_stream: bool = True
is_list_item: bool = True
real_components: Dict[str, tuple]
training: bool
class transformer_lens.model_bridge.generalized_components.block.DelegatedAttentionBlockBridge(name: str, config: Any | None = None, submodules: Dict[str, GeneralizedComponent] | None = None, hook_alias_overrides: Dict[str, str] | None = None)

Bases: BlockBridge

Block whose attention is delegated wholesale to HF (no split-qkv fork).

For architectures with heterogeneous per-layer attention structure — e.g. Gemma 4, where KV-shared layers have no k_proj/v_proj at all and K==V layers have no v_proj — there is no uniform HookPoint that represents “input that becomes Q/K/V”, so the block-level hook_q_input/ hook_k_input/hook_v_input/hook_attn_in aliases do not apply. Type-level distinction means a reader of the adapter sees DelegatedAttentionBlockBridge and knows those hooks are absent.

maintain_native_attention: bool = True
real_components: Dict[str, tuple]
training: bool
class transformer_lens.model_bridge.generalized_components.block.MLABlockBridge(name: str, config: Any | None = None, submodules: Dict[str, GeneralizedComponent] | None = None, hook_alias_overrides: Dict[str, str] | None = None)

Bases: BlockBridge

Block wrapping Multi-Head Latent Attention (DeepSeek V2/V3/R1).

MLA has no standalone q/k/v projections — Q flows through compressed q_a_proj→q_a_layernorm→q_b_proj, and K/V share a joint kv_a_proj_with_mqa entry point. There is no single HookPoint that represents “input that becomes Q/K/V”, so the block-level hook_q_input/hook_k_input/ hook_v_input/hook_attn_in aliases do not apply. Type-level distinction means a reader of the adapter sees MLABlockBridge and knows those hooks are absent.

class transformer_lens.model_bridge.generalized_components.block.ParallelBlockBridge(name: str, config: Any | None = None, submodules: Dict[str, GeneralizedComponent] | None = None, hook_alias_overrides: Dict[str, str] | None = None)

Bases: BlockBridge

Block where attn and MLP both read the pre-attention residual.

For GPT-J, NeoX, Pythia, Phi, Cohere, CodeGen, and some Falcon variants, output = resid_pre + attn_out + mlp_out — no distinct post-attention residual exists. Matches legacy HookedTransformer which omits hook_resid_mid when cfg.parallel_attn_mlp=True. Type-level distinction means a reader of the adapter sees ParallelBlockBridge and knows the hook is absent.

class transformer_lens.model_bridge.generalized_components.block.ScaledResidualBlockBridge(name: str, config: Any | None = None, submodules: Dict[str, GeneralizedComponent] | None = None, hook_alias_overrides: Dict[str, str] | None = None, mlp_reads_resid_directly: bool = False, residual_contribution_scale: float = 1.0, scaled_attn_submodule: str = 'attn', scaled_mlp_submodule: str | None = 'mlp')

Bases: BlockBridge

Block whose sublayer outputs are scaled before the residual add.

Granite-family HF blocks compute residual + sublayer_out * residual_multiplier inline, so no submodule output equals the tensor added to the residual stream and the legacy aliases cannot be fixed by re-pointing. hook_attn_out / hook_mlp_out become real HookPoints on the block firing on the scaled contribution:

  • no hooks attached: the forward is untouched (bit-exact);

  • read-only hooks: they observe module_out * scale, forward stays bit-exact;

  • a hook that changes the tensor (returned new or mutated in place): the module output is rewritten to hooked / scale so HF’s multiply reconstructs the written value as the contribution (~1-ulp rounding; exact for zero-ablation);

  • any backward hooks: the rewrite always happens so the autograd graph routes through the HookPoint.

__init__(name: str, config: Any | None = None, submodules: Dict[str, GeneralizedComponent] | None = None, hook_alias_overrides: Dict[str, str] | None = None, mlp_reads_resid_directly: bool = False, residual_contribution_scale: float = 1.0, scaled_attn_submodule: str = 'attn', scaled_mlp_submodule: str | None = 'mlp')

scaled_mlp_submodule=None when no single submodule feeds the MLP-side add (GraniteMoeHybrid sums moe + shared_mlp inline) — hook_mlp_out then stays absent rather than firing with a partial tensor.

set_original_component(original_component: Module) → None

Prune contribution HookPoints the bound layer cannot fire.

Heterogeneous blocks (GraniteMoeHybrid mamba layers) lack the attn submodule; a HookPoint that exists but never fires is a silent-no-op intervention trap, so the name must be absent instead.