class_docstr: Top-1 MoE
  
  Curently supports Sinkhorn based expert routing.
class_docstr_tok: ["Top-1","MoE","Curently","supports","Sinkhorn","based","expert","routing","."]
code: def __init__(
      self,
      config: ModelParallelConfig,
      num_experts,
      init_method,
      output_layer_init_method,
      hidden_size,
      ffn_hidden_size,
      bias_activation_fusion=True,
      openai_gelu=False,
      onnx_safe=False,
      activation="gelu",
      bias=True,
      transformer_block_type="pre_ln",
      normalization="layernorm",
      layernorm_epsilon=1e-5,
      persist_layer_norm=False,
      sequence_parallel=False,
      dropout=0.0,
  ):
      super(SwitchMLP, self).__init__(config=config)
  
      self.num_experts = num_experts
      self.route_algo = SwitchMLP.sinkhorn
      self.router = tensor_parallel.RowParallelLinear(
          hidden_size,
          num_experts,
          config=config,
          input_is_parallel=False,
          init_method=init_method,
          skip_bias_add=False,
          bias=bias,
      )
  
      mlp_args = {
          "config": config,
          "init_method": init_method,
          "output_layer_init_method": output_layer_init_method,
          "hidden_size": hidden_size,
          "ffn_hidden_size": ffn_hidden_size,
          "bias_activation_fusion": bias_activation_fusion,
          "openai_gelu": openai_gelu,
          "onnx_safe": onnx_safe,
          "activation": activation,
          "bias": bias,
          "transformer_block_type": transformer_block_type,
          "normalization": normalization,
          "layernorm_epsilon": layernorm_epsilon,
          "persist_layer_norm": persist_layer_norm,
          "dropout": dropout,
      }
      self.experts = torch.nn.ModuleList(
          [ParallelMLP(**mlp_args) for _ in range(num_experts)]
      )
  
imports: import torch
  import torch.nn.functional
  from nemo.collections.nlp.modules.common.megatron.adapters.parallel_adapters import AdapterName, MLPInfusedAdapterConfig
  from nemo.collections.nlp.modules.common.megatron.fused_bias_geglu import fused_bias_geglu
  from nemo.collections.nlp.modules.common.megatron.fused_bias_gelu import fused_bias_gelu
  from nemo.collections.nlp.modules.common.megatron.fused_layer_norm import get_layer_norm
  from nemo.collections.nlp.modules.common.megatron.layer_norm_1p import LayerNorm1P
  from nemo.collections.nlp.modules.common.megatron.module import MegatronModule
  from nemo.collections.nlp.modules.common.megatron.utils import ApexGuardDefaults, ApproxGELUActivation, erf_gelu
  from nemo.collections.nlp.modules.common.megatron.utils import openai_gelu
  from nemo.collections.nlp.modules.common.megatron.utils import squared_relu
  from nemo.core import adapter_mixins
label_desc: Use image node_id 1 to create a new SwitchMLP object from inherited base classes: MegatronModule with example: obj = SwitchMLP(config, num_experts, init_method, output_layer_init_method, hidden_size, ffn_hidden_size, bias_activation_fusion, openai_gelu, onnx_safe, activation, bias, transformer_block_type, normalization, layernorm_epsilon, persist_layer_norm, sequence_parallel, dropout)