Coverage for transformer_lens/model_bridge/generalized_components/symbolic.py: 100%

1"""Symbolic bridge component. 

2 

3This module contains a bridge component that acts as a structural placeholder 

4when a model doesn't have a corresponding container for grouped subcomponents. 

5""" 

6from typing import Any, Dict, Optional 

7 

8from transformer_lens.model_bridge.generalized_components.base import ( 

9 GeneralizedComponent, 

10) 

11 

12 

13class SymbolicBridge(GeneralizedComponent): 

14 """A placeholder bridge component for maintaining TransformerLens structure. 

15 

16 This bridge is used when a model doesn't have a container component that exists 

17 in the TransformerLens standard structure. For example, OPT has fc1/fc2 layers 

18 directly on the block rather than inside an MLP container. 

19 

20 When the model is set up, the subcomponents defined in this SymbolicBridge 

21 are promoted to the parent component, allowing the TransformerLens structure 

22 to be maintained while correctly mapping to the underlying model's architecture. 

23 

24 Example usage: 

25 # OPT doesn't have an "mlp" container - fc1/fc2 are on the block directly 

26 "mlp": SymbolicBridge( 

27 submodules={ 

28 "in": LinearBridge(name="fc1"), 

29 "out": LinearBridge(name="fc2"), 

30 }, 

31 ) 

32 

33 # During setup, "in" and "out" will be accessible as: 

34 # - blocks[i].mlp.in (pointing to blocks[i].fc1) 

35 # - blocks[i].mlp.out (pointing to blocks[i].fc2) 

36 

37 Attributes: 

38 is_symbolic: Always True, indicates this is a structural placeholder. 

39 """ 

40 

41 is_symbolic: bool = True 

42 

43 def __init__( 

44 self, 

45 submodules: Optional[Dict[str, GeneralizedComponent]] = None, 

46 config: Optional[Any] = None, 

47 ): 

48 """Initialize the SymbolicBridge. 

49 

50 Args: 

51 submodules: Dictionary of submodules to register. These will be set up 

52 using the parent's original_component as their context. 

53 config: Optional configuration object 

54 """ 

55 # SymbolicBridge always has name=None since it doesn't map to a real component 

56 super().__init__(name=None, config=config, submodules=submodules or {}) 

57 

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

59 """Forward pass is not supported for SymbolicBridge. 

60 

61 SymbolicBridge is a structural placeholder and should not be called directly. 

62 The actual computation should go through the subcomponents which are set up 

63 on the parent. 

64 

65 Raises: 

66 RuntimeError: Always, since SymbolicBridge should not be called directly. 

67 """ 

68 raise RuntimeError( 

69 "SymbolicBridge is a structural placeholder and should not be called directly. " 

70 "Use the subcomponents (e.g., 'in', 'out') instead." 

71 )