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 Because ``hook_out`` is captured before the causal trim, negative ``pos_slice`` 

19 values index the causal padding rather than real tokens — slice positive 

20 positions from the left, or trim the last ``conv_kernel - 1`` entries first. 

21 

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

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

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

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

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

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

28 """ 

29 

30 def __init__(self, *args: Any, **kwargs: Any) -> None: 

31 super().__init__(*args, **kwargs) 

32 # Channel-first: the sequence axis is dim 2 (see the class docstring). 

33 self.hook_in.pos_dim = 2 

34 self.hook_out.pos_dim = 2 

35 

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

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

38 raise RuntimeError( 

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

40 "Call set_original_component() first." 

41 ) 

42 input = self.hook_in(input) 

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

44 output = self.hook_out(output) 

45 return output 

46 

47 def get_random_inputs( 

48 self, 

49 batch_size: int = 2, 

50 seq_len: int = 8, 

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

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

53 ) -> Dict[str, Any]: 

54 if self.original_component is None: 

55 raise RuntimeError( 

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

57 "Call set_original_component() first." 

58 ) 

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

60 dtype = dtype or torch.float32 

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

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

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

64 

65 def __repr__(self) -> str: 

66 if self.original_component is not None: 

67 try: 

68 in_channels = self.original_component.in_channels 

69 out_channels = self.original_component.out_channels 

70 kernel_size = self.original_component.kernel_size 

71 groups = self.original_component.groups 

72 return ( 

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

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

75 ) 

76 except AttributeError: 

77 return ( 

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

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

80 ) 

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