Coverage for transformer_lens/model_bridge/generalized_components/depthwise_conv1d.py: 42%

1"""Bridge for Mamba-style depthwise causal Conv1d (distinct from GPT-2's Conv1D linear).""" 

2from typing import Any, Dict, Optional, cast 

3 

4import torch 

5 

6from transformer_lens.model_bridge.generalized_components.base import ( 

7 GeneralizedComponent, 

8) 

9 

10 

11class DepthwiseConv1DBridge(GeneralizedComponent): 

12 """Wraps an ``nn.Conv1d`` depthwise causal convolution with input/output hooks. 

13 

14 Hook shapes (channel-first, as HF's MambaMixer transposes before the call): 

15 hook_in: [batch, channels, seq_len] 

16 hook_out: [batch, channels, seq_len + conv_kernel - 1] (pre causal trim) 

17 

18 Decode-step limitation: on stateful generation, HF's Mamba/Mamba-2 mixers 

19 bypass ``self.conv1d(...)`` and read ``self.conv1d.weight`` directly, so the 

20 forward hook never fires on decode steps — only on prefill. For per-step 

21 conv output during decode, compute it manually from the cached conv_states 

22 and ``conv1d.original_component.weight``, or run token-by-token via 

23 ``forward()`` instead of ``generate()``. 

24 """ 

25 

26 def forward(self, input: torch.Tensor, *args: Any, **kwargs: Any) -> torch.Tensor: 

27 if self.original_component is None: 27 ↛ 28line 27 didn't jump to line 28 because the condition on line 27 was never true

28 raise RuntimeError( 

29 f"Original component not set for {self.name}. " 

30 "Call set_original_component() first." 

31 ) 

32 input = self.hook_in(input) 

33 output = self.original_component(input, *args, **kwargs) 

34 output = self.hook_out(output) 

35 return output 

36 

37 def get_random_inputs( 

38 self, 

39 batch_size: int = 2, 

40 seq_len: int = 8, 

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

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

43 ) -> Dict[str, Any]: 

44 if self.original_component is None: 

45 raise RuntimeError( 

46 f"Original component not set for {self.name}. " 

47 "Call set_original_component() first." 

48 ) 

49 device = device or torch.device("cpu") 

50 dtype = dtype or torch.float32 

51 conv = cast(torch.nn.Conv1d, self.original_component) 

52 channels = conv.in_channels # exact, from the wrapped Conv1d 

53 return {"args": (torch.randn(batch_size, channels, seq_len, device=device, dtype=dtype),)} 

54 

55 def __repr__(self) -> str: 

56 if self.original_component is not None: 

57 try: 

58 in_channels = self.original_component.in_channels 

59 out_channels = self.original_component.out_channels 

60 kernel_size = self.original_component.kernel_size 

61 groups = self.original_component.groups 

62 return ( 

63 f"DepthwiseConv1DBridge({in_channels} -> {out_channels}, " 

64 f"kernel_size={kernel_size}, groups={groups})" 

65 ) 

66 except AttributeError: 

67 return ( 

68 f"DepthwiseConv1DBridge(name={self.name}, " 

69 f"original_component={type(self.original_component).__name__})" 

70 ) 

71 return f"DepthwiseConv1DBridge(name={self.name}, original_component=None)"