class_docstr: ## Attention layer
class_docstr_tok: ["#","#","Attention","layer"]
docstr: :param n_hidden: the number of features in embeddings
  :param n_heads: the number of attention heads
  :param rope_percentage: percentage of features to add RoPE embeddings
  :param mask_fill: masking fill value for attention matrix
  :param is_flash_attention: specifies whether to use
      [FlashAttention](https://github.com/HazyResearch/flash-attention)
code: def __init__(
      self,
      n_hidden: int = 6_144,
      n_heads: int = 64,
      rope_percentage: float = 0.25,
      mask_fill: float = -10_000.0,
      *,
      is_flash_attention: bool = False,
  ):
      """
      :param n_hidden: the number of features in embeddings
      :param n_heads: the number of attention heads
      :param rope_percentage: percentage of features to add RoPE embeddings
      :param mask_fill: masking fill value for attention matrix
      :param is_flash_attention: specifies whether to use
          [FlashAttention](https://github.com/HazyResearch/flash-attention)
      """
      super().__init__()
  
      self.n_heads = n_heads
      self.mask_fill = mask_fill
  
      # Linear layer for query, key and value
      self.qkv_lin = nn.Linear(n_hidden, n_hidden * 3)
      # Final linear layer
      self.output = nn.Linear(n_hidden, n_hidden)
  
      # Number of features per head
      d_k = n_hidden // n_heads
      # RoPE embedding module
      self.rope = RoPE(int(d_k * rope_percentage))
  
      # Attention scaling factor
      self.scale = 1 / math.sqrt(d_k)
  
      # To cache causal mask
      self.causal_mask = None
  
      # Attention softmax module
      self.softmax = nn.Softmax(dim=-2)
  
      # [FlashAttention](https://github.com/HazyResearch/flash-attention)
      if is_flash_attention:
          try:
              from flash_attn.flash_attention import FlashAttention
  
              self.flash_attention = FlashAttention()
          except ImportError:
              logger.log(
                  "Install flash attention github.com/HazyResearch/flash-attention. "
                  "Falling back to normal attention",
                  Text.warning,
              )
              self.flash_attention = None
      else:
          self.flash_attention = None
  
imports: import copy
  import math
  from typing import Dict, Optional, Set, Callable, Any, Generator, Tuple
  import torch
  from torch import nn
  from torch.cuda.amp import autocast
  from labml import monit, logger
  from labml.logger import Text
  from labml_nn.neox import checkpoint
  from labml_nn.neox.utils.cache import get_cache
label_desc: Use image node_id 1 to create a new AttentionLayer object from inherited base classes: nn with example: obj = AttentionLayer(n_hidden, n_heads, rope_percentage, mask_fill)