class_docstr: Top-1 Mixture of Experts Layer. Routes input to one of N MLP "experts"
  Curently supports Sinkhorn based expert routing.
code: def __init__(
      self, config: TransformerConfig, submodules: MLPSubmodules
  ):
      super().__init__(config=config)
  
      self.local_experts = torch.nn.ModuleList()
      for _ in range(self.num_local_experts):
          expert = MLP(self.config, submodules, is_expert=True)
          self.local_experts.append(expert)
  
file_path: megatron-lm/megatron/core/transformer/moe/switch_mlp.py
imports: import numpy
  import torch
  from megatron.core.transformer.mlp import MLP, MLPSubmodules
  from megatron.core.transformer.transformer_config import TransformerConfig
  from .base_moe_layer import BaseMoELayer
label_desc: Use image node_id 1 to create a new SwitchMLP object from inherited base classes: BaseMoELayer with example: obj = SwitchMLP(config, submodules)