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: OpaqueBlockBridge

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

Extends OpaqueBlockBridge with SSM-specific hook aliases:

  • hook_mixer_inmixer.hook_in

  • hook_mixer_outmixer.hook_out

These aliases are only meaningful for architectures whose blocks contain a mixer submodule (SSM or hybrid SSM/attention). For non-SSM architectures use OpaqueBlockBridge directly.

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