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: * `h` is the transformer embeddings of shape `[batch_size, seq_len, d_model]`
docstr_tok: ["*","`","h","`","is","the","transformer","embeddings","of","shape","`","[","batch_size",",","seq_len",",","d_model","]","`"]
code: def forward(self, h: torch.Tensor):
      """
      * `h` is the transformer embeddings of shape `[batch_size, seq_len, d_model]`
      """
  
      # Residual connection
      h_res = h
  
      # Pre-normalization
      h = self.norm(h)
  
      # Get query, key, and values and split them in to heads.
      # These will have shapes `[batch_size, seq_len, n_heads, d_k]`
      mh_shape = (*h.shape[:-1], self.n_heads, self.d_k)
      q = self.query(h).view(mh_shape)
      k = self.key(h).view(mh_shape)
      v = self.value(h).view(mh_shape)
  
      # Apply rotary positional embeddings
      q = self.rotary_pe(q)
      k = self.rotary_pe(k)
  
      # Calculate attentions
      attn = torch.einsum("bihd,bjhd->bhij", q, k)
      # Scale it by $\frac{1}{\sqrt{d_k}}$
      attn = attn * self.scale
  
      # Apply masks if it's causal attention
      attn = self.mask_attention(attn)
  
      # Calculate attention probabilities
      attn = self.softmax(attn)
  
      # Get values
      h = torch.einsum("bhij,bjhd->bihd", attn, v)
  
      # Change from shape `[batch_size, seq_len, n_heads, d_k]`
      # to `[batch_size, seq_len, n_heads * d_k]`
      h = h.reshape(*h.shape[:-2], -1)
  
      # Apply final linear layer.
      # The result will have shape `[batch_size, seq_len, d_model]`
      h = self.output(h)
  
      # Add the residual connection
      return h + h_res
  
imports: import math
  from typing import Set
  import torch
  from torch import nn
  from labml.logger import inspect
label_desc: Use image node_id 3 for calling the SelfAttention obj's underlying member method code with example usage: obj.forward(h) and returns: unknown