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

Block bridge for SSM layers — direct GeneralizedComponent subclass.

Does not inherit from BlockBridge because BlockBridge’s hook_aliases hardcode transformer-specific names (hook_attn_*, hook_mlp_*, hook_resid_mid).

forward(*args: Any, **kwargs: Any) → Any

Delegate to the HF block with hook_in/hook_out wrapped around it.

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'}

is_list_item: bool = True

real_components: Dict[str, tuple]

training: bool