transformer_lens.model_bridge.supported_architectures.t5gemma2 module

T5Gemma2 architecture adapter (text-only).

T5Gemma2ForConditionalGeneration is a multimodal encoder-decoder model. This adapter bridges the text path only: - Encoder text stack under model.encoder.text_model (the SigLIP vision_tower and

multi_modal_projector are intentionally left unmapped).

  • Decoder stack under model.decoder.

Key differences from T5Gemma: - Encoder text lives at model.encoder.text_model.* (not model.encoder.*). - The decoder uses a single T5Gemma2MergedAttention that fuses self- and

cross-attention with shared q/k/v/o projections; there is no separate cross-attention module and no cross-attention layernorms.

  • Both encoder and decoder attention add Gemma-style QK-norm (q_norm/k_norm).

  • Per-layer sliding/full attention with dual RoPE and per-head QK-norm are all handled natively by HF — the bridge only routes inputs and fires hooks.

class transformer_lens.model_bridge.supported_architectures.t5gemma2.T5Gemma2ArchitectureAdapter(cfg: Any)

Bases: ArchitectureAdapter

Architecture adapter for T5Gemma2ForConditionalGeneration (text-only).

Encoder: BlockBridge over model.encoder.text_model.layers (Gemma-style, QK-norm, no cross-attn) Decoder: T5Gemma2DecoderBlockBridge over model.decoder.layers (merged self+cross attention)

setup_component_testing(hf_model: Any, bridge_model: Any = None) → None

Set up rotary embedding references for T5Gemma2 component testing.

Both the encoder text stack and the decoder carry their own rotary_emb. We set the reference on all PositionEmbeddingsAttentionBridge instances so that component-level forward calls can compute RoPE correctly, force eager attention (so patterns are hookable), and enable native layernorm autograd on QK-norm so the manual encoder path matches HF exactly.