class_docstr: ## Self-Attention Layer $   ext{A\small{TTN}}$
  
  This applies causal and non-causal [multi-headed self-attention](../mha.html).
class_docstr_tok: ["#","#","Self-Attention","Layer","$","ext","{","A\\small","{","TTN","}","}","$","This","applies","causal","and","non-causal","[","multi-headed","self-attention","]","(","..","\/mha.html",")","."]
docstr: * `d_model` is the number of features in transformer embeddings
  * `n_heads` is the number of attention heads
  * `d_k` is the number of features per head
  * `is_causal` indicates whether this is causal attention (masked)
docstr_tok: ["*","`","d_model","`","is","the","number","of","features","in","transformer","embeddings","*","`","n_heads","`","is","the","number","of","attention","heads","*","`","d_k","`","is","the","number","of","features","per","head","*","`","is_causal","`","indicates","whether","this","is","causal","attention","(","masked",")"]
code: def __init__(
      self, d_model: int, n_heads: int, d_k: int, is_causal: bool
  ):
      """
      * `d_model` is the number of features in transformer embeddings
      * `n_heads` is the number of attention heads
      * `d_k` is the number of features per head
      * `is_causal` indicates whether this is causal attention (masked)
      """
      super().__init__()
  
      self.is_causal = is_causal
      self.n_heads = n_heads
      self.d_k = d_k
  
      # To scale attentions before softmax by $\frac{1}{\sqrt{d_k}}$
      self.scale = 1 / math.sqrt(self.d_k)
  
      # Linear layers for query, key and value heads.
      self.query = nn.Linear(d_model, n_heads * d_k)
      self.key = nn.Linear(d_model, n_heads * d_k)
      self.value = nn.Linear(d_model, n_heads * d_k)
  
      # Pre-norm layer. The paper uses RMSNorm instead.
      self.norm = nn.LayerNorm(d_model)
  
      # Softmax for attention probabilities
      self.softmax = nn.Softmax(dim=-1)
  
      # Rotary positional embeddings
      self.rotary_pe = RotaryPositionalEmbeddings(self.d_k)
  
      # Final linear layer
      self.output = nn.Linear(n_heads * d_k, d_model)
  
imports: import math
  from typing import Set
  import torch
  from torch import nn
  from labml.logger import inspect
label_desc: Use image node_id 1 to create a new SelfAttention object from inherited base classes: nn with example: obj = SelfAttention(d_model, n_heads, d_k, is_causal)