transformer_lens.model_bridge.generalized_components.ssm_block module¶
Block container for State Space Model (Mamba) layers: norm → mixer → residual.
- class transformer_lens.model_bridge.generalized_components.ssm_block.SSMBlockBridge(name: str, config: Any | None = None, submodules: Dict[str, GeneralizedComponent] | None = None, hook_alias_overrides: Dict[str, str] | None = None)¶
Bases:
OpaqueBlockBridgeBlock bridge for SSM layers (Mamba, Falcon-H1, etc.).
Extends
OpaqueBlockBridgewith SSM-specific hook aliases:hook_mixer_in→mixer.hook_inhook_mixer_out→mixer.hook_out
These aliases are only meaningful for architectures whose blocks contain a
mixersubmodule (SSM or hybrid SSM/attention). For non-SSM architectures useOpaqueBlockBridgedirectly.- hook_aliases: Dict[str, str | List[str]] = {'hook_mixer_in': 'mixer.hook_in', 'hook_mixer_out': 'mixer.hook_out', 'hook_resid_post': 'hook_out', 'hook_resid_pre': 'hook_in'}¶
- real_components: Dict[str, tuple]¶
- training: bool¶