code: def __init__(self, config, prefix, weights):
      super().__init__()
      act = config.hidden_act
      self.act = (
          ACT2FN[act]
          if "gelu" not in act
          else lambda x: torch.nn.functional.gelu(
              x,
              approximate="tanh"
              if act in ["gelu_fast", "gelu_pytorch_tanh"]
              else "none",
          )
      )
  
      self.dense_h_to_4h = TensorParallelColumnLinear.load(
          config,
          prefix=f"{prefix}.dense_h_to_4h",
          weights=weights,
          bias=True,
      )
      self.dense_4h_to_h = load_row(
          config,
          prefix=f"{prefix}.dense_4h_to_h",
          weights=weights,
          bias=True,
      )
  
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 FlashMLP object from inherited base classes: nn with example: obj = FlashMLP(config, prefix, weights)