Coverage for transformer_lens/model_bridge/composition_scores.py: 86%

1"""Tensor-like container for composition score results with layer-index metadata.""" 

2from typing import Any, List 

3 

4import torch 

5 

6 

7class CompositionScores: 

8 """Composition scores that behave like a tensor but carry layer-index metadata. 

9 

10 Delegates indexing, .shape, arithmetic, and torch.* functions to the 

11 underlying ``scores`` tensor via ``__torch_function__``. On hybrid models 

12 where n_attn_layers < n_layers, ``layer_indices`` maps tensor position i 

13 to the original layer number. 

14 

15 Attributes: 

16 scores: Upper-triangular composition score tensor. 

17 layer_indices: Original layer numbers, e.g. [0, 2, 5]. 

18 head_labels: Labels matching scores dims, e.g. ["L0H0", "L0H1", ...]. 

19 """ 

20 

21 def __init__(self, scores: torch.Tensor, layer_indices: List[int], head_labels: List[str]): 

22 self.scores = scores 

23 self.layer_indices = layer_indices 

24 self.head_labels = head_labels 

25 

26 @classmethod 

27 def __torch_function__(cls, func, types, args=(), kwargs=None): 

28 """Unwrap CompositionScores, including those inside argument containers.""" 

29 if kwargs is None: 29 ↛ 32line 29 didn't jump to line 32 because the condition on line 29 was always true

30 kwargs = {} 

31 

32 def unwrap(value: Any) -> Any: 

33 if isinstance(value, CompositionScores): 

34 return value.scores 

35 if isinstance(value, list): 

36 return [unwrap(item) for item in value] 

37 if isinstance(value, tuple): 

38 return tuple(unwrap(item) for item in value) 

39 if isinstance(value, dict): 

40 return {key: unwrap(item) for key, item in value.items()} 

41 return value 

42 

43 return func(*unwrap(args), **unwrap(kwargs)) 

44 

45 @property 

46 def shape(self) -> torch.Size: 

47 return self.scores.shape 

48 

49 @property 

50 def device(self) -> torch.device: 

51 return self.scores.device 

52 

53 @property 

54 def dtype(self) -> torch.dtype: 

55 return self.scores.dtype 

56 

57 def __getitem__(self, key): 

58 return self.scores[key] 

59 

60 def __getattr__(self, name): 

61 # Guard against recursion during pickle/deepcopy when self.scores isn't set yet 

62 try: 

63 scores = object.__getattribute__(self, "scores") 

64 except AttributeError: 

65 raise AttributeError(name) from None 

66 return getattr(scores, name) 

67 

68 def __gt__(self, other): 

69 return self.scores > other 

70 

71 def __lt__(self, other): 

72 return self.scores < other 

73 

74 def __ge__(self, other): 

75 return self.scores >= other 

76 

77 def __le__(self, other): 

78 return self.scores <= other 

79 

80 def __eq__(self, other): 

81 if isinstance(other, CompositionScores): 81 ↛ 82line 81 didn't jump to line 82 because the condition on line 81 was never true

82 return self.scores == other.scores 

83 return self.scores == other 

84 

85 def __ne__(self, other): 

86 if isinstance(other, CompositionScores): 86 ↛ 87line 86 didn't jump to line 87 because the condition on line 86 was never true

87 return self.scores != other.scores 

88 return self.scores != other 

89 

90 def __repr__(self) -> str: 

91 return ( 

92 f"CompositionScores(shape={self.shape}, " 

93 f"layer_indices={self.layer_indices}, " 

94 f"n_head_labels={len(self.head_labels)})" 

95 )