transformer_lens.model_bridge.generalized_components.pos_embed module

Positional embedding bridge component.

This module contains the bridge component for positional embedding layers.

class transformer_lens.model_bridge.generalized_components.pos_embed.PosEmbedBridge(name: str, config: Any | None = None, submodules: Dict[str, GeneralizedComponent] | None = {})

Bases: GeneralizedComponent

Positional embedding bridge that wraps transformer positional embedding layers.

This component provides standardized input/output hooks for positional embeddings.

property W_pos: Tensor

Return the positional embedding weight matrix.

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

Initialize the positional embedding bridge.

Parameters:

  • name – The name of this component

  • config – Optional configuration (unused for PosEmbedBridge)

  • submodules – Dictionary of GeneralizedComponent submodules to register

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

Forward pass through the positional embedding bridge.

This method accepts variable arguments to support different architectures: - Standard models (GPT-2, GPT-Neo): (input_ids, position_ids=None) - OPT models: (attention_mask, past_key_values_length=0, position_ids=None) - Others may have different signatures

Parameters:

  • *args – Positional arguments forwarded to the original component

  • **kwargs – Keyword arguments forwarded to the original component

Returns:

Positional embeddings

property_aliases: Dict[str, str] = {'W_pos': 'weight'}
real_components: Dict[str, tuple]
training: bool