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 

38 def get_random_inputs( 

39 self, 

40 batch_size: int = 2, 

41 seq_len: int = 8, 

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

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

44 ) -> Dict[str, Any]: 

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

46 inputs = super().get_random_inputs( 

47 batch_size=batch_size, seq_len=seq_len, device=device, dtype=dtype 

48 ) 

49 hidden_states = inputs["hidden_states"] 

50 inputs["merged_attention_mask"] = None 

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

52 return inputs 

53 

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

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

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

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

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

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

60 cross_weights = output[2] 

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

62 cross_weights = self.hook_cross_pattern(cross_weights) 

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

64 return output