transformer_lens.model_bridge.generalized_components.t5gemma2_merged_attention module

Bridge for T5Gemma2’s merged self+cross decoder attention.

T5Gemma2MergedAttention runs decoder self-attention and encoder-decoder cross-attention through one module with shared q/k/v/o projections: it projects encoder_hidden_states through the same k_proj/v_proj, concatenates the encoder K/V onto the decoder K/V, and does a single softmax. In eager mode it returns (attn_output, self_attn_weights, cross_attn_weights), where the self and cross weights are the leading and trailing key-position slices of the merged pattern.

This bridge delegates the math to the native module (the merged/cross logic cannot be reimplemented by the manual attention path) while exposing both pattern slices: hook_pattern for the self-attention slice (handled by the base class) and hook_cross_pattern for the cross-attention slice.

class transformer_lens.model_bridge.generalized_components.t5gemma2_merged_attention.T5Gemma2MergedAttentionBridge(*args: Any, **kwargs: Any)

Bases: AttentionBridge

Native-delegating attention bridge that also hooks the cross-attention pattern.

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

Simplified forward pass - minimal wrapping around original component.

This does minimal wrapping: hook_in → delegate to HF → hook_out. This ensures we match HuggingFace’s exact output without complex intermediate processing.

Parameters:

  • *args – Input arguments to pass to the original component

  • **kwargs – Input keyword arguments to pass to the original component

Returns:

The output from the original component, with only input/output hooks applied

get_random_inputs(batch_size: int = 2, seq_len: int = 8, device: device | None = None, dtype: dtype | None = None) → Dict[str, Any]

Add the merged-attention inputs the native module requires positionally.