transformer_lens.utilities.bridge_components module

Utilities for traversing and applying functions to every component in a TransformerBridge model.

transformer_lens.utilities.bridge_components.apply_fn_to_all_components(model: TransformerBridge, fn: Callable[[GeneralizedComponent], Any], components: dict | None = None) → dict[str, Any]

Applies a function to all components in the TransformerBridge model. :param model: The TransformerBridge model to apply the function to :param fn: The function to apply to each component :param components: Optional dictionary of components to apply the function to, if None, all components are collected

Returns:

A dictionary mapping component names to the return values of the function

Return type:

return_values

transformer_lens.utilities.bridge_components.collect_all_components(model: TransformerBridge, components: dict) → dict

Collects all components in a TransformerBridge inside a dictionary. The keys are the component names, and the values are the components themselves. :param model: The TransformerBridge model to collect components from :param components: A dictionary to populate with components (modified in-place)

Returns:

Dictionary mapping component names to their respective components

transformer_lens.utilities.bridge_components.collect_all_submodules_of_component(model: TransformerBridge, component: GeneralizedComponent, submodules: dict, block_prefix: str = '') → dict

Recursively collects all submodules of a component in a TransformerBridge model. :param model: The TransformerBridge model to collect submodules from :param component: The component to collect submodules from :param submodules: A dictionary to populate with submodules (modified in-place) :param block_prefix: Prefix for the block name, needed for components that are part of a block bridge

Returns:

Dictionary mapping submodule names to their respective submodules

transformer_lens.utilities.bridge_components.collect_components_of_block_bridge(model: TransformerBridge, component: GeneralizedComponent, components: dict) → dict

Collects all components of a BlockBridge component. :param model: The TransformerBridge model to collect components from :param component: The BlockBridge component to collect components from :param components: A dictionary to populate with components (modified in-place)

Returns:

Dictionary mapping component names to their respective components