transformer_lens.model_bridge.generalized_components.mlp module

MLP bridge component.

This module contains the bridge component for MLP layers.

class transformer_lens.model_bridge.generalized_components.mlp.MLPBridge(name: str | None, config: Any | None = None, submodules: Dict[str, GeneralizedComponent] | None = {}, optional: bool = False)

Bases: GeneralizedComponent

Bridge component for MLP layers.

This component wraps an MLP layer from a remote model and provides a consistent interface for accessing its weights and performing MLP operations.

property W_gate: Tensor | None

MLP gate weight in TL orientation [d_model, d_mlp], or None if ungated.

property W_in: Tensor

MLP input weight in TL orientation [d_model, d_mlp].

property W_out: Tensor

MLP output weight in TL orientation [d_mlp, d_model].

__init__(name: str | None, config: Any | None = None, submodules: Dict[str, GeneralizedComponent] | None = {}, optional: bool = False)

Initialize the MLP bridge.

Parameters:

  • name – The name of the component in the model (None if no container exists)

  • config – Optional configuration (unused for MLPBridge)

  • submodules – Dictionary of submodules to register (e.g., gate_proj, up_proj, down_proj)

  • optional – If True, setup skips this bridge when absent (hybrid architectures).

forward(*args, **kwargs) → Any

Forward pass through the MLP bridge.

Returns a tensor, or the component’s own (hidden, …) tuple re-packed with hooked hidden states for recurrent MLPs.

Parameters:

  • *args – Positional arguments for the original component

  • **kwargs – Keyword arguments for the original component

Returns:

Output hidden states

hook_aliases: Dict[str, str | List[str]] = {'hook_post': 'out.hook_in', 'hook_pre': 'in.hook_out'}
property_aliases: Dict[str, str] = {'b_gate': 'gate.bias', 'b_in': 'in.bias', 'b_out': 'out.bias'}
real_components: Dict[str, tuple]
training: bool
transformer_lens.model_bridge.generalized_components.mlp.normalize_mlp_weight(weight: Tensor, layout: bool | None, proj: Any, pattern: str = 'in') → Tensor

Normalize an MLP projection weight to TL orientation ([d_model, d_mlp] for “in”/W_gate, [d_mlp, d_model] for “out”).

transformer_lens.model_bridge.generalized_components.mlp.weight_layout_in_out(proj: Any) → bool | None

Whether proj’s wrapped module stores its weight as [in, out].

Conv1D (GPT-2 style) stores [in_features, out_features]; nn.Linear stores [out_features, in_features]. Returns None when the wrapped module is neither, so callers fall back to in_features/out_features or a shape heuristic.