code: def get_router_linear_layer(config):
      router = torch.nn.Linear(
          config.hidden_size, config.num_moe_experts, bias=False
      )
      with get_cuda_rng_tracker().fork(
          get_data_parallel_rng_tracker_name()
      ):
          config.init_method(router.weight)
      setattr(
          router.weight, "sequence_parallel", config.sequence_parallel
      )
      return router
  
file_path: megatron-lm/megatron/core/transformer/moe/base_moe_layer.py
imports: from abc import ABC, abstractmethod
  import torch
  from megatron.core import parallel_state, tensor_parallel
  from megatron.core.parallel_state import get_tensor_and_expert_parallel_group
  from megatron.core.tensor_parallel import get_cuda_rng_tracker, get_data_parallel_rng_tracker_name
  from megatron.core.transformer.module import MegatronModule
  from megatron.core.transformer.transformer_config import TransformerConfig
label_desc: Use image node_id 2 for calling a global function with example usage: get_router_linear_layer(config) and returns: router