transformer_lens.utilities.attention module

Attention.

Utilities for attention components.

transformer_lens.utilities.attention.clamp_qkv(q: Tensor, k: Tensor, v: Tensor, clip: float) → tuple[Tensor, Tensor, Tensor]

Clamp Q/K/V to [-clip, clip] out-of-place (OLMo/OLMoE/MPT clip_qkv).

Out-of-place rather than HF’s clamp_ so tensors wrapped by full backward hooks stay legal to use.

transformer_lens.utilities.attention.complex_attn_linear(input: Float[Tensor, 'batch pos head_index d_model'], w: Float[Tensor, 'head_index d_model d_head'], b: Float[Tensor, 'head_index d_head']) → Float[Tensor, 'batch pos head_index d_head']

Linear layer for attention calculation.

This is almost the same as simple_attn_linear, but the input tensor has an extra head_index dimension, used when calculating the input of each attention head separately.

transformer_lens.utilities.attention.simple_attn_linear(input: Float[Tensor, 'batch pos d_model'], w: Float[Tensor, 'head_index d_model d_head'], b: Float[Tensor, 'head_index d_head']) → Float[Tensor, 'batch pos head_index d_head']

Linear layer for attention calculation.