Coverage for transformer_lens/model_bridge/generalized_components/t5gemma2_merged_attention.py: 93%

1"""Bridge for T5Gemma2's merged self+cross decoder attention. 

2 

3T5Gemma2MergedAttention runs decoder self-attention and encoder-decoder 

4cross-attention through one module with shared q/k/v/o projections: it projects 

5``encoder_hidden_states`` through the same ``k_proj``/``v_proj``, concatenates the 

6encoder K/V onto the decoder K/V, and does a single softmax. In eager mode it 

7returns ``(attn_output, self_attn_weights, cross_attn_weights)``, where the self 

8and cross weights are the leading and trailing key-position slices of the merged 

9pattern. 

10 

11This bridge delegates the math to the native module (the merged/cross logic 

12cannot be reimplemented by the manual attention path) while exposing both pattern 

13slices: ``hook_pattern`` for the self-attention slice (handled by the base class) 

14and ``hook_cross_pattern`` for the cross-attention slice. 

15""" 

16from __future__ import annotations 

17 

18from typing import Any, Dict, Optional 

19 

20import torch 

21 

22from transformer_lens.hook_points import HookPoint 

23from transformer_lens.model_bridge.generalized_components.attention import ( 

24 AttentionBridge, 

25) 

26 

27 

28class T5Gemma2MergedAttentionBridge(AttentionBridge): 

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

30 

31 def __init__(self, *args: Any, **kwargs: Any) -> None: 

32 # The native module unpacks position_embeddings unconditionally, so 

33 # component testing must always supply a (cos, sin) tuple. 

34 kwargs.setdefault("requires_position_embeddings", True) 

35 super().__init__(*args, **kwargs) 

36 self.hook_cross_pattern = HookPoint() 

37 # [batch, heads, q_pos, k_pos]: slice the query position like hook_pattern. 

38 self.hook_cross_pattern.pos_dim = -2 

39 

40 def get_random_inputs( 

41 self, 

42 batch_size: int = 2, 

43 seq_len: int = 8, 

44 device: Optional[torch.device] = None, 

45 dtype: Optional[torch.dtype] = None, 

46 ) -> Dict[str, Any]: 

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

48 inputs = super().get_random_inputs( 

49 batch_size=batch_size, seq_len=seq_len, device=device, dtype=dtype 

50 ) 

51 hidden_states = inputs["hidden_states"] 

52 inputs["merged_attention_mask"] = None 

53 inputs["encoder_hidden_states"] = torch.randn_like(hidden_states) 

54 return inputs 

55 

56 def forward(self, *args: Any, **kwargs: Any) -> Any: 

57 # Base forward delegates to native and hooks output[0] (hook_out) and the 

58 # self-attention weights output[1] (hook_pattern). Native returns a third 

59 # element — the cross-attention weights — which we hook here. 

60 output = super().forward(*args, **kwargs) 

61 if isinstance(output, tuple) and len(output) >= 3: 61 ↛ 66line 61 didn't jump to line 66 because the condition on line 61 was always true

62 cross_weights = output[2] 

63 if isinstance(cross_weights, torch.Tensor) and cross_weights.dim() == 4: 63 ↛ 66line 63 didn't jump to line 66 because the condition on line 63 was always true

64 cross_weights = self.hook_cross_pattern(cross_weights) 

65 output = output[:2] + (cross_weights,) + output[3:] 

66 return output