class_docstr: ## Transformer Layer
class_docstr_tok: ["#","#","Transformer","Layer"]
docstr: :param n_hidden: is the embedding size
  :param n_heads: is the number of heads
  :param is_flash_attention: specifies whether to use
      [FlashAttention](https://github.com/HazyResearch/flash-attention)
  
  *Out implementation doesn't include dropout*.
code: def __init__(
      self,
      n_hidden: int = 6_144,
      n_heads: int = 64,
      *,
      is_flash_attention: bool = False,
  ):
      """
      :param n_hidden: is the embedding size
      :param n_heads: is the number of heads
      :param is_flash_attention: specifies whether to use
          [FlashAttention](https://github.com/HazyResearch/flash-attention)
  
      *Out implementation doesn't include dropout*.
      """
      super().__init__()
  
      # Layer normalization before attention
      self.pre_ln_attn = nn.LayerNorm(n_hidden)
      # Layer normalization before FFN
      self.pre_ln_ffn = nn.LayerNorm(n_hidden)
  
      # Attention layer
      self.attention = AttentionLayer(
          n_hidden, n_heads, is_flash_attention=is_flash_attention
      )
      # FFN layer
      self.ffn = FFNLayer(n_hidden)
  
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 TransformerLayer object from inherited base classes: NeoXModule with example: obj = TransformerLayer(n_hidden, n_heads)