code: def _random_init_weight(self, rng):
      r"""
      Copied from torch.nn.init.kaiming_uniform_
      """
      fan = nn.init._calculate_correct_fan(self.weight[0], "fan_in")
      gain = nn.init.calculate_gain("leaky_relu", math.sqrt(5))
      std = gain / math.sqrt(fan)
      bound = math.sqrt(3.0) * std
      device = self.weight.device
      dtype = self.weight.dtype
      weight = rng.uniform(
          -bound, bound, size=tuple(self.weight.size())
      )
      self.weight.data = torch.from_numpy(weight).to(
          dtype=dtype, device=device
      )
  
      if self.bias is not None:
          fan_in, _ = nn.init._calculate_fan_in_and_fan_out(
              self.weight[0]
          )
          bound = 1 / math.sqrt(fan_in)
          bias = rng.uniform(
              -bound, bound, size=tuple(self.bias.size())
          )
          self.bias.data = torch.from_numpy(bias).to(
              dtype=dtype, device=device
          )
  
imports: import math
  import numpy
  import torch
  import torch.nn
  import torch.nn.functional
  from fmoe.transformer import FMoETransformerMLP
  from fmoe.utils import get_torch_default_comm
  from .balance import reset_gate_hook
  from .balance import generate_megatron_gate_hook
  from .distributed import set_moe_group
label_desc: Use image node_id 2 for calling a global function with example usage: _random_init_weight(rng) without return types