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

1"""Block container for State Space Model (Mamba) layers: norm → mixer → residual.""" 

2 

3from __future__ import annotations 

4 

5from transformer_lens.model_bridge.generalized_components.opaque_block import ( 

6 OpaqueBlockBridge, 

7) 

8 

9 

10class SSMBlockBridge(OpaqueBlockBridge): 

11 """Block bridge for SSM layers (Mamba, Falcon-H1, etc.). 

12 

13 Extends :class:`OpaqueBlockBridge` with SSM-specific hook aliases: 

14 

15 - ``hook_mixer_in`` → ``mixer.hook_in`` 

16 - ``hook_mixer_out`` → ``mixer.hook_out`` 

17 

18 These aliases are only meaningful for architectures whose blocks contain a 

19 ``mixer`` submodule (SSM or hybrid SSM/attention). For non-SSM architectures 

20 use :class:`OpaqueBlockBridge` directly. 

21 """ 

22 

23 hook_aliases = { 

24 **OpaqueBlockBridge.hook_aliases, 

25 "hook_mixer_in": "mixer.hook_in", 

26 "hook_mixer_out": "mixer.hook_out", 

27 }