class_docstr: An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
  models.
class_docstr_tok: ["An","abstract","class","to","handle","weights","initialization","and","a","simple","interface","for","downloading","and","loading","pretrained","models","."]
docstr: Initialize the weights
docstr_tok: ["Initialize","the","weights"]
code: def _init_weights(self, module):
      """Initialize the weights"""
      factor = (
          self.config.initializer_factor
      )  # Used for testing weights initialization
      if isinstance(module, nn.LayerNorm):
          module.weight.data.fill_(factor * 1.0)
          module.bias.data.zero_()
      elif isinstance(module, nn.Linear):
          module.weight.data.normal_(
              mean=0.0, std=factor * ((self.config.d_model) ** -0.5)
          )
          if hasattr(module, "bias") and module.bias is not None:
              module.bias.data.zero_()
      elif isinstance(module, nn.Embedding):
          module.weight.data.normal_(mean=0.0, std=factor * 1.0)
      elif isinstance(module, GPTSanJapaneseModel):
          # Mesh TensorFlow embeddings initialization
          # See https://github.com/tensorflow/mesh/blob/fa19d69eafc9a482aff0b59ddd96b025c0cb207d/mesh_tensorflow/layers.py#L1624
          module.embed_tokens.weight.data.normal_(
              mean=0.0, std=factor * 1.0
          )
          module.position_embeddings.weight.data.normal_(
              mean=0.0, std=factor * 1.0
          )
          if (
              hasattr(module, "extra_position_embeddings")
              and module.extra_position_embeddings is not None
          ):
              module.extra_position_embeddings.weight.data.normal_(
                  mean=0.0, std=factor * 1.0
              )
      elif isinstance(
          module,
          (GPTSanJapaneseModel, GPTSanJapaneseForConditionalGeneration),
      ):
          # Mesh TensorFlow embeddings initialization
          # See https://github.com/tensorflow/mesh/blob/fa19d69eafc9a482aff0b59ddd96b025c0cb207d/mesh_tensorflow/layers.py#L1624
          module.final_logits_bias.data.normal_(
              mean=0.0, std=factor * 1.0
          )
          if (
              hasattr(module, "lm_head")
              and not self.config.tie_word_embeddings
          ):
              module.lm_head.weight.data.normal_(
                  mean=0.0, std=factor * 1.0
              )
      elif isinstance(module, GPTSanJapaneseDenseActDense):
          # Mesh TensorFlow FF initialization
          # See https://github.com/tensorflow/mesh/blob/master/mesh_tensorflow/transformer/transformer_layers.py#L56
          # and https://github.com/tensorflow/mesh/blob/fa19d69eafc9a482aff0b59ddd96b025c0cb207d/mesh_tensorflow/layers.py#L89
          module.wi.weight.data.normal_(
              mean=0.0, std=factor * ((self.config.d_model) ** -0.5)
          )
          if hasattr(module.wi, "bias") and module.wi.bias is not None:
              module.wi.bias.data.zero_()
          module.wo.weight.data.normal_(
              mean=0.0, std=factor * ((self.config.d_ff) ** -0.5)
          )
          if hasattr(module.wo, "bias") and module.wo.bias is not None:
              module.wo.bias.data.zero_()
      elif isinstance(module, GPTSanJapaneseAttention):
          # Multi-headed attention
          d_model = self.config.d_model
          key_value_proj_dim = self.config.d_model
          n_heads = self.config.num_heads
          module.k_proj.weight.data.normal_(
              mean=0.0,
              std=factor * ((d_model * key_value_proj_dim) ** -0.5),
          )
          module.v_proj.weight.data.normal_(
              mean=0.0,
              std=factor * ((d_model * key_value_proj_dim) ** -0.5),
          )
          module.q_proj.weight.data.normal_(
              mean=0.0,
              std=factor * ((d_model * key_value_proj_dim) ** -0.5),
          )
          module.out_proj.weight.data.normal_(
              mean=0.0,
              std=factor * ((n_heads * key_value_proj_dim) ** -0.5),
          )
      elif isinstance(module, GPTSanJapaneseSparseMLP):
          # Mesh TensorFlow attention initialization to avoid scaling before softmax
          # See https://github.com/tensorflow/mesh/blob/fa19d69eafc9a482aff0b59ddd96b025c0cb207d/mesh_tensorflow/transformer/attention.py#L136
          d_model = self.config.d_model
          key_value_proj_dim = self.config.d_model
          n_heads = self.config.num_heads
          module.router.classifier.weight.data.normal_(
              mean=0.0, std=factor * 1
          )
          for idx in range(self.config.num_experts):
              module.experts[f"expert_{idx}"].wi.weight.data.normal_(
                  mean=0.0, std=factor * (d_model**-0.5)
              )
              module.experts[f"expert_{idx}"].wo.weight.data.normal_(
                  mean=0.0, std=factor * (d_model**-0.5)
              )
  
imports: import copy
  from typing import List, Optional, Tuple, Union
  import torch
  import torch.nn
  from ...activations import ACT2FN
  from ...modeling_outputs import MoECausalLMOutputWithPast, MoEModelOutputWithPastAndCrossAttentions
  from ...modeling_utils import PreTrainedModel
  from ...utils import DUMMY_INPUTS, DUMMY_MASK, add_start_docstrings, add_start_docstrings_to_model_forward, is_torch_fx_proxy, logging
  from .configuration_gptsan_japanese import GPTSanJapaneseConfig
label_desc: Use image node_id 2 for calling the GPTSanJapanesePreTrainedModel obj's underlying member method code with example usage: obj._init_weights(module) without return types