transformer_lens.model_bridge.generalized_components.t5gemma2_decoder_block module

T5Gemma2-specific decoder block bridge.

T5Gemma2DecoderLayer replaces T5Gemma’s separate self-attention + cross-attention with a single T5Gemma2MergedAttention module. That module computes decoder self queries/keys/values from hidden_states and cross keys/values from encoder_hidden_states using the same q/k/v/o projections, concatenates the self and cross key/value states, and runs a single softmax. As a result the decoder layer has no separate cross-attention module and no cross-attention layernorms — its structure mirrors the encoder layer plus encoder-state input.

This bridge monkey-patches the layer forward to fire hook points at the canonical HookedTransformer residual-stream positions while delegating all attention math (QK-norm, RoPE, scaling, merged KV) to the native HF module.

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

Bases: GeneralizedComponent

Bridge for T5Gemma2 decoder layers (merged self+cross attention).

Inserts hook points around the two sub-components of each decoder layer: - hook_in (hook_resid_pre): residual before self-attention pre-norm - hook_resid_mid: residual after merged-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