code: def __init__(self, config, weights):
      super().__init__(config)
      self.config = config
  
      self.embed_in = TensorParallelEmbedding(
          prefix="gpt_neox.embed_in", weights=weights
      )
  
      self.layers = nn.ModuleList(
          [
              FlashNeoXLayer(layer_id, config, weights)
              for layer_id in range(config.num_hidden_layers)
          ]
      )
      self.final_layer_norm = FastLayerNorm.load(
          prefix="gpt_neox.final_layer_norm",
          weights=weights,
          eps=config.layer_norm_eps,
      )
  
      self.gradient_checkpointing = False
  
      self.head_size = self.layers[0].attention.head_size
      self.num_heads = self.layers[0].attention.num_heads
  
imports: import torch
  import torch.distributed
  from torch import nn
  from transformers.activations import ACT2FN
  from transformers.modeling_utils import PreTrainedModel
  from transformers.models.gpt_neox import GPTNeoXConfig
  from typing import Optional, List, Tuple
  from text_generation_server.utils import paged_attention, flash_attn
  from text_generation_server.utils.flash_attn import attention
  from text_generation_server.utils.layers import TensorParallelRowLinear, TensorParallelColumnLinear, TensorParallelEmbedding, TensorParallelHead, FastLayerNorm, PositionRotaryEmbedding, get_linear
label_desc: Use image node_id 1 to create a new FlashGPTNeoXModel object from inherited base classes: FlashGPTNeoXPreTrainedModel with example: obj = FlashGPTNeoXModel(config, weights)