transformer_lens.components package

Submodules

• transformer_lens.components.abstract_attention module
  • AbstractAttention
    • AbstractAttention.IGNORE
    • AbstractAttention.OV
    • AbstractAttention.QK
    • AbstractAttention.__init__()
    • AbstractAttention.alibi
    • AbstractAttention.apply_causal_mask()
    • AbstractAttention.apply_rotary()
    • AbstractAttention.calculate_attention_scores()
    • AbstractAttention.calculate_qkv_matrices()
    • AbstractAttention.calculate_sin_cos_rotary()
    • AbstractAttention.calculate_z_scores()
    • AbstractAttention.create_alibi_bias()
    • AbstractAttention.create_alibi_multipliers()
    • AbstractAttention.create_alibi_slope()
    • AbstractAttention.forward()
    • AbstractAttention.k_norm
    • AbstractAttention.mask
    • AbstractAttention.q_norm
    • AbstractAttention.rotary_cos
    • AbstractAttention.rotary_sin
    • AbstractAttention.rotate_every_two()
• transformer_lens.components.attention module
  • Attention
    • Attention.__init__()
• transformer_lens.components.bert_block module
  • BertBlock
    • BertBlock.forward()
• transformer_lens.components.bert_embed module
  • BertEmbed
    • BertEmbed.forward()
• transformer_lens.components.bert_mlm_head module
  • BertMLMHead
    • BertMLMHead.forward()
• transformer_lens.components.bert_nsp_head module
  • BertNSPHead
    • BertNSPHead.forward()
• transformer_lens.components.bert_pooler module
  • BertPooler
    • BertPooler.forward()
• transformer_lens.components.embed module
  • Embed
    • Embed.forward()
• transformer_lens.components.grouped_query_attention module
  • GroupedQueryAttention
    • GroupedQueryAttention.W_K
    • GroupedQueryAttention.W_V
    • GroupedQueryAttention.__init__()
    • GroupedQueryAttention.b_K
    • GroupedQueryAttention.b_V
    • GroupedQueryAttention.calculate_attention_scores()
    • GroupedQueryAttention.calculate_qkv_matrices()
    • GroupedQueryAttention.calculate_z_scores()
• transformer_lens.components.layer_norm module
  • LayerNorm
    • LayerNorm.__init__()
    • LayerNorm.forward()
• transformer_lens.components.layer_norm_pre module
  • LayerNormPre
    • LayerNormPre.__init__()
    • LayerNormPre.forward()
• transformer_lens.components.pos_embed module
  • PosEmbed
    • PosEmbed.forward()
• transformer_lens.components.rms_norm module
  • RMSNorm
    • RMSNorm.__init__()
    • RMSNorm.forward()
• transformer_lens.components.rms_norm_pre module
  • RMSNormPre
    • RMSNormPre.__init__()
    • RMSNormPre.forward()
• transformer_lens.components.t5_attention module
  • T5Attention
    • T5Attention.compute_relative_attention_bias()
• transformer_lens.components.t5_block module
  • T5Block
    • T5Block.forward()
• transformer_lens.components.token_typed_embed module
  • TokenTypeEmbed
    • TokenTypeEmbed.forward()
• transformer_lens.components.transformer_block module
  • TransformerBlock
    • TransformerBlock.apply_mlp()
    • TransformerBlock.forward()
    • TransformerBlock.ln1
    • TransformerBlock.ln2
    • TransformerBlock.mlp
• transformer_lens.components.unembed module
  • Unembed
    • Unembed.forward()

Module contents

Hooked Transformer Components.

This module contains all the components (e.g. Attention, MLP, LayerNorm) needed to create many different types of generative language models. They are used by transformer_lens.HookedTransformer.