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

1"""LiquidAI LFM2 gated short-convolution mixer bridge.""" 

2 

3from transformer_lens.model_bridge.generalized_components.base import ( 

4 GeneralizedComponent, 

5) 

6 

7 

8class Lfm2ShortConvBridge(GeneralizedComponent): 

9 """Wrapper around LFM2's double-gated short-convolution mixer. 

10 

11 Delegates the forward to HF's ``Lfm2ShortConv`` (preserving its fast CUDA / 

12 slow PyTorch dispatch and cache handling) and hooks the residual-stream 

13 input/output. Inner in_proj / conv / out_proj are spliced in as submodules, 

14 so their hooks fire during HF's own forward. 

15 

16 Decode-step caveat: on stateful generation HF's conv path reads 

17 ``conv.weight`` directly instead of calling ``self.conv(...)``, so 

18 ``conv.hook_out`` fires only on prefill — see DepthwiseConv1DBridge. 

19 

20 CUDA caveat: Hooks surrounding the conv1D operation only fire on the hf 

21 "slow path" i.e. if not on CUDA / fast path not available / torch dynamo 

22 compiling. 

23 """ 

24 

25 hook_aliases = { 

26 "hook_in_proj": "in.hook_out", 

27 "hook_conv": "conv.hook_out", 

28 "hook_gated": "out.hook_in", 

29 }