class_docstr: MOELayer module which implements MixtureOfExperts as described in Gshard_.
  ::
  
      gate = TopKGate(model_dim, num_experts)
      moe = MOELayer(gate, expert)
      output = moe(input)
      l_aux = moe.l_aux
  
  .. Gshard_: https://arxiv.org/pdf/2006.16668.pdf
  
  Args:
      gate (torch.nn.Module):
          gate network
      expert (torch.nn.Module):
          expert network
class_docstr_tok: ["MOELayer","module","which","implements","MixtureOfExperts","as","described","in","Gshard_",".",":",":","gate","=","TopKGate","(","model_dim",",","num_experts",")","moe","=","MOELayer","(","gate",",","expert",")","output","=","moe","(","input",")","l_aux","=","moe.l_aux","..","Gshard_",":","https",":","\/\/arxiv.org\/pdf\/2006.16668.pdf","Args",":","gate","(","torch.nn.Module",")",":","gate","network","expert","(","torch.nn.Module",")",":","expert","network"]
code: def __init__(
      self,
      gate: Module,
      experts: Module,
      ep_group_name,
      ep_size,
      num_local_experts: int,
      use_tutel: bool = False,
  ) -> None:
      super().__init__()
      self.gate = gate
      self.experts = experts
      self.ep_group = None
      self.ep_size = ep_size
      self.ep_group_name = ep_group_name
      self.num_local_experts = num_local_experts
      self.time_falltoall = 0.0
      self.time_salltoall = 0.0
      self.time_moe = 0.0
      self.timers = SynchronizedWallClockTimer()
      self.wall_clock_breakdown = False
  
      self.use_tutel = use_tutel and TUTEL_INSTALLED and gate.k == 1
  
      if self.use_tutel:
          logger.info("Using Tutel optimizations.")
      elif use_tutel and not TUTEL_INSTALLED:
          logger.warning(
              "Tutel optimization requested but not installed. "
              "Proceeding without Tutel."
          )
      elif use_tutel and TUTEL_INSTALLED and gate.k != 1:
          logger.warning(
              "To enable Tutel optimization, use top-1 instead of top-2 gate. "
              "Proceeding without Tutel."
          )
  
imports: from deepspeed.utils.timer import SynchronizedWallClockTimer
  from deepspeed.utils import logger
  from typing import Callable, Dict, TYPE_CHECKING, Any, Optional, Tuple
  import torch
  from torch import Tensor
  from torch.nn import Module
  import torch.nn.functional
  from deepspeed.utils import groups
  from .mappings import drop_tokens, gather_tokens
  from deepspeed import comm
label_desc: Use image node_id 1 to create a new MOELayer object from inherited base classes: Base with example: obj = MOELayer(gate, experts, ep_group_name, ep_size, num_local_experts, use_tutel)