transformer_lens.model_bridge.generalized_components.t5gemma_decoder_block module

T5Gemma-specific decoder block bridge.

T5GemmaDecoderLayer uses Gemma-style flat attribute access (not T5’s .layer[] indexing). It has: self-attention + cross-attention + MLP, each with pre/post norms. This bridge monkey-patches the layer forward to insert intermediate hook points.

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

Bases: GeneralizedComponent

Bridge for T5Gemma decoder layers.

Inserts hook points between the three sub-components of each decoder layer: - hook_in (hook_resid_pre): residual before self-attention pre-norm - hook_resid_mid: residual after self-attention + residual add, before cross-attn pre-norm - hook_resid_mid2: residual after cross-attention + residual add, before MLP pre-norm - hook_out (hook_resid_post): residual after MLP + residual add

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

Generic forward pass for bridge components with input/output hooks.

get_expected_parameter_names(prefix: str = '') → list[str]

get_list_size() → int

hook_aliases: Dict[str, str | List[str]] = {'hook_resid_post': 'hook_out', 'hook_resid_pre': 'hook_in'}

is_list_item: bool = True

real_components: Dict[str, tuple]

set_original_component(component: Module) → None

Set the original component that this bridge wraps.

Parameters:

original_component – The original transformer component to wrap

training: bool