code: def _init_weights(
      module,
      n_layer,
      initializer_range=0.02,
      rescale_prenorm_residual=True,
  ):
      if isinstance(module, nn.Linear):
          nn.init.normal_(module.weight, std=initializer_range)
          if module.bias is not None:
              nn.init.zeros_(module.bias)
      elif isinstance(module, nn.Embedding):
          nn.init.normal_(module.weight, std=initializer_range)
  
      if rescale_prenorm_residual:
          # Reinitialize selected weights subject to the OpenAI GPT-2 Paper Scheme:
          #   > A modified initialization which accounts for the accumulation on the residual path with model depth. Scale
          #   > the weights of residual layers at initialization by a factor of 1/√N where N is the # of residual layers.
          #   >   -- GPT-2 :: https://openai.com/blog/better-language-models/
          #
          # Reference (Megatron-LM): https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/model/gpt_model.py
          for name, p in module.named_parameters():
              if name in ["out_proj.weight", "fc2.weight"]:
                  # Special Scaled Initialization --> There are 2 Layer Norms per Transformer Block
                  nn.init.normal_(
                      p,
                      mean=0.0,
                      std=initializer_range / math.sqrt(2 * n_layer),
                  )
  
imports: import math
  from functools import partial
  from collections import namedtuple
  from collections.abc import Sequence
  import torch
  import torch.nn
  import torch.nn.functional
  from transformers.models.gpt2.configuration_gpt2 import GPT2Config
  from src.models.modules.mha import MHA
  from src.models.modules.mlp import Mlp, FusedDenseGeluDense
  from src.models.modules.block import Block
  from src.models.modules.embedding import GPT2Embeddings
label_desc: Use image node_id 4 for calling a global function with example usage: _init_weights(module, n_layer, initializer_range, rescale_prenorm_residual) without return types