transformer_lens.model_bridge.generalized_components.depthwise_conv1d module

Bridge for Mamba-style depthwise causal Conv1d (distinct from GPT-2’s Conv1D linear).

class transformer_lens.model_bridge.generalized_components.depthwise_conv1d.DepthwiseConv1DBridge(*args: Any, **kwargs: Any)

Bases: GeneralizedComponent

Wraps an nn.Conv1d depthwise causal convolution with input/output hooks.

Hook shapes (channel-first, as HF’s MambaMixer transposes before the call):

hook_in: [batch, channels, seq_len] hook_out: [batch, channels, seq_len + conv_kernel - 1] (pre causal trim)

Because hook_out is captured before the causal trim, negative pos_slice values index the causal padding rather than real tokens — slice positive positions from the left, or trim the last conv_kernel - 1 entries first.

Decode-step limitation: on stateful generation, HF’s Mamba/Mamba-2 mixers bypass self.conv1d(...) and read self.conv1d.weight directly, so the forward hook never fires on decode steps — only on prefill. For per-step conv output during decode, compute it manually from the cached conv_states and conv1d.original_component.weight, or run token-by-token via forward() instead of generate().

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

Generic forward pass for bridge components with input/output hooks.

get_random_inputs(batch_size: int = 2, seq_len: int = 8, device: device | None = None, dtype: dtype | None = None) → Dict[str, Any]