Coverage for transformer_lens/utilities/matrix.py: 90%

1"""matrix. 

2 

3This module contains utility functions related to the transformer lens implementation of factored 

4matrices. 

5""" 

6from typing import Union 

7 

8import torch 

9from jaxtyping import Float 

10 

11from transformer_lens.FactoredMatrix import FactoredMatrix 

12 

13from .tensors import get_corner 

14 

15 

16def composition_scores( 

17 left: FactoredMatrix, right: FactoredMatrix, broadcast_dims=True 

18) -> Union[ 

19 Float[torch.Tensor, "*leading_dims"], 

20 Float[torch.Tensor, "*leading_dims_left_and_right"], 

21]: 

22 """Composition scores between two factored matrices. 

23 

24 Returns ``||left @ right||_F / (||left||_F * ||right||_F)``, computed from the factored 

25 forms so the full products are never materialized. With ``broadcast_dims``, left and right 

26 leading dims are broadcast against each other (left dims first), scoring every left/right 

27 pair. See ``TransformerBridge.all_composition_scores``. 

28 """ 

29 if broadcast_dims: 29 ↛ 36line 29 didn't jump to line 36 because the condition on line 29 was always true

30 r_leading = right.ndim - 2 

31 l_leading = left.ndim - 2 

32 for i in range(l_leading): 

33 right = right.unsqueeze(i) 

34 for i in range(r_leading): 

35 left = left.unsqueeze(i + l_leading) 

36 assert ( 

37 left.rdim == right.ldim 

38 ), f"Composition scores require left.rdim==right.ldim, shapes were left: {left.shape}, right:{right.shape}" 

39 

40 new_right = right.collapse_r() 

41 new_left = left.collapse_l() 

42 r_norms = new_right.norm(dim=[-2, -1]) 

43 l_norms = new_left.norm(dim=[-2, -1]) 

44 comp_norms = (new_left @ new_right).norm(dim=[-2, -1]) 

45 return comp_norms / r_norms / l_norms 

46 

47 

48def get_matrix_corner(matrix: FactoredMatrix, n=3): 

49 result = get_corner(matrix[tuple(slice(n) for _ in range(matrix.ndim))]) 

50 

51 return result.AB