transformer_lens.model_bridge.generalized_components.base module

Base class for generalized transformer components.

class transformer_lens.model_bridge.generalized_components.base.CloneOutputUnderGradMixin(*args: Any, **kwargs: Any)

Bases: Module

Clone the forward output so HF’s in-place mutation cannot corrupt it.

Under grad, autograd forbids in-place writes to backward-hook views; under no_grad, cached hook_out tensors alias the storage HF then rewrites. Mix in ahead of a bridge class: class X(CloneOutputUnderGradMixin, LinearBridge).

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

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class transformer_lens.model_bridge.generalized_components.base.GeneralizedComponent(name: str | None, config: Any | None = None, submodules: Dict[str, GeneralizedComponent] | None = None, conversion_rule: BaseTensorConversion | None = None, hook_alias_overrides: Dict[str, str] | None = None, optional: bool = False)

Bases: Module

Base class for generalized transformer components.

This class provides a standardized interface for transformer components and handles hook registration and execution.

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

Initialize the generalized component.

Parameters:

  • name – The name of this component (None if component has no container in remote model)

  • config – Optional configuration object for the component

  • submodules – Dictionary of GeneralizedComponent submodules to register

  • conversion_rule – Optional conversion rule for this component’s hooks

  • hook_alias_overrides – Optional dictionary to override default hook aliases. For example, {“hook_attn_out”: “ln1_post.hook_out”} will make hook_attn_out point to ln1_post.hook_out instead of the default value in self.hook_aliases.

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

add_hook(hook_fn: Callable[[...], Tensor], hook_name: str = 'output') → None

Add a hook function (HookedTransformer-compatible interface).

Parameters:

  • hook_fn – Function to call at this hook point

  • hook_name – Name of the hook point (defaults to “output”)

compatibility_mode: bool = False
disable_warnings: bool = False
forward(*args: Any, **kwargs: Any) → Any

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

get_hooks() → Dict[str, HookPoint]

Get all hooks registered in this component.

hook_aliases: Dict[str, str | List[str]] = {}
hook_out_is_single_residual_stream: bool = False
is_list_item: bool = False
property original_component: Module | None

Get the original component.

property_aliases: Dict[str, str] = {}
remove_hooks(hook_name: str | None = None) → None

Remove hooks (HookedTransformer-compatible interface).

Parameters:

hook_name – Name of the hook point to remove. If None, removes all hooks.

set_original_component(original_component: Module) → None

Set the original component that this bridge wraps.

Parameters:

original_component – The original transformer component to wrap

set_processed_weights(weights: Dict[str, Tensor], verbose: bool = False) → None

Set the processed weights for use in compatibility mode.

This method stores processed weights as attributes on the component so they can be used directly in the forward pass without modifying the original component.

Components should override this method to handle their specific weight structure. The weights dict contains keys like “weight”, “bias”, “W_in”, “W_out”, etc.

If this component has submodules, this method will automatically distribute the weights to those subcomponents using ProcessWeights.distribute_weights_to_components.

Parameters:

  • weights – Dictionary of processed weight tensors

  • verbose – If True, print detailed information about weight setting