transformer_lens.model_bridge.generalized_components.position_embedding_hooks_mixin module

Mixin for position embedding hooks (cos/sin) shared across attention bridges.

class transformer_lens.model_bridge.generalized_components.position_embedding_hooks_mixin.PositionEmbeddingHooksMixin

Bases: object

Mixin providing hook_cos/hook_sin and _apply_position_embedding_hooks().

Used by both PositionEmbeddingsAttentionBridge and JointQKVPositionEmbeddingsAttentionBridge to avoid duplicating this logic.

set_rotary_emb(rotary_emb: Any) → None

Set reference to the model’s rotary embedding component.